mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
feat(gateway): land the chat provider series and full-feature load test on the poc branch (#47388)
* feat(gateway): Azure OpenAI chat completions through the OpenAI passthrough
Adds azure/<deployment> chat completions, streaming and not, on the same
passthrough the openai/ route uses. The request goes to the deployment
path with the api-version query, or to /openai/v1/ for the v1, latest and
preview versions, with api_base, api_key and api_version falling back to
the AZURE_* variables and os.environ references resolved as the proxy
does. A key is sent as api-key, and an azure_ad_token or an Entra ID
token from auth-azure goes as a Bearer token. The body follows the
proxy's gpt, o-series and gpt-5 param rules, picked by base_model or the
deployment name and reading model capabilities from the cost map, so an
unsupported param is refused unless drop_params allows it to be dropped.
Spend is priced as the proxy prices Azure, by azure/<base_model> when one
is set and otherwise by the model the response names, and a content
filter error keeps Azure's inner error for the client
* feat(gateway): Anthropic and Bedrock chat completions streaming
Adds "stream": true to /v1/chat/completions for anthropic/ and bedrock/
deployments, which were refused with 501 before. The request is the
non-streaming one with "stream": true added for Anthropic and sent to
/converse-stream for Bedrock, and the provider's events are decoded the
way the proxy's chunk parsers decode them: text, tool call deltas with
empty arguments sent as "{}", reasoning_content and thinking_blocks with
their index and signature, citations and the mapped finish_reason. A
wrapper in the shape of the proxy's stream wrapper gives every chunk one
id and the client's model name, sends role on the first chunk, adds the
empty finish chunk, strips usage from the chunks and, when
stream_options.include_usage is set, ends with the usage chunk before
[DONE]. The gateway waits for the first chunk before answering, so an
upstream error up to there is retried or falls back to another
deployment, and an error after it ends the stream with the OpenAI error
frame. Usage, cache reads and writes included, is handed to the meter
through the response whether or not the client asked for it, and
dropping the client stream drops the upstream request
* feat(gateway): full openai chat request translation for anthropic and bedrock (#47088)
* feat(gateway): share router cooldowns and follow proxy key invalidations through Redis
Gateway replicas now share state through the coordination Redis the
Python proxy uses (general_settings.coordination_redis, a redis
response cache, or REDIS_*):
- A deployment one replica cools down is written under Python's
deployment:<id>:cooldown key and applied by every other replica's
background poll (one MGET a second), so the fleet stops sending to a
failing deployment together. allowed_fails counts failures across
replicas, as Python's v1 rule does; the v2 failure-rate rule stays
per replica, as in Python. Deployments are named by model_info.id, or
an id derived from the config that is the same on every replica.
- The virtual key cache follows the proxy's
litellm_proxy.auth_cache_invalidation channel, so a key blocked,
deleted or edited through the admin API (or its team or user) is
re-read on every replica at once instead of after up to 60 s.
general_settings.user_api_key_cache_ttl sets the key TTL.
Nothing is added to the path of a successful request. Without Redis,
or while it is down, each replica keeps its own state as before.
Redis Cluster and Sentinel are refused at startup for now.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5
* feat(gateway): hooks and observers around inference requests, with usage and cost
Adds the litellm-gateway-hooks crate: request facts built only when something
reads them, on-path hooks that declare what they read and what they may
change and run in config order, off-path observers on bounded queues, a
usage tap for SSE and JSON responses, and per-deployment prices from the
model cost map overlaid by config prices. The messages and chat handlers run
through it when anything is installed; with nothing installed requests take
the old path.
The key store now also returns the key alias, user, team alias and key and
team metadata, and router attempts carry the deployment id.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5
* feat(gateway): write spend to the proxy's database from a native observer
Adds a spend writer observer that records every gateway request in the
proxy's own Postgres tables with the same statements the Python proxy
uses: spend logs, key/user/team/team member/end user spend, and the
daily user, team and end-user activity tables. Writes are batched on the
proxy's timers and are additive, so replicas and Python proxies can
write side by side. The database pool is now shared between virtual
keys and the writer.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5
* feat(gateway): enforce key, team, user, member and end-user budgets on shared spend counters
Adds a built-in budgets hook that runs the proxy's budget checks, in its
order and with its messages, on the proxy's Redis spend:* counters. Each
request reads and seeds the counters, checks them and reserves its
maximum cost in one Redis script, then settles the reservation to its
real cost off the request path, so replicas and Python proxies sharing
the Redis enforce one budget. Without Redis the counters are kept per
replica. The key store now reads user, team member and end-user budgets,
and auth no longer checks budgets against stored spend.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5
* feat(gateway): enforce key, user, team, member and end-user rate limits on shared counters
Adds a built-in rate limits hook with the proxy's v3 limiter semantics:
requests per window, tokens per window and in-flight requests per key,
on the proxy's own {scope:id}:window|requests|tokens|max_parallel_requests
Redis keys and with its scripts, so replicas and Python proxies sharing
the Redis enforce one limit. Tokens are reserved from the proxy's
estimate and settled to usage when the response ends. Refusals carry the
proxy's message, body and headers on each endpoint, and complete
responses carry its x-ratelimit-* headers. Without Redis the counters are
kept per replica. The key store now reads key, user, team and member
rpm/tpm limits and the key's max_parallel_requests, and end users' budget
limits are read with their budgets; keys with plain rate limits are no
longer refused.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5
* feat(gateway): configurable header, spend metadata and provider metadata hooks
Adds four built-in hooks a config lists under litellm_settings.hooks and
runs in that order: forward_headers copies named client headers to the
provider, strip_headers removes headers from the provider request
(including the gateway's default pass-through), spend_metadata adds
header and key-metadata values to the spend log's spend_logs_metadata,
and provider_metadata sets Messages body fields and headers from
templates over request facts. Hooks' upstream header edits now reach the
provider on both endpoints, and body edits can set nested fields.
Credentials and signer headers are never forwarded or set, and a
misconfigured hook stops the gateway at startup.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5
* feat(gateway): model capacity hook, the proxy's dynamic rate limiter
Adds the model_capacity hook, turned on by listing dynamic_rate_limiter_v3
(or model_capacity) in litellm_settings.callbacks. It caps each model
group at the sum of its deployments' rpm and tpm across every replica,
and with priority_reservation holds each priority's pool to its share
once the group's saturation reaches the threshold, counting it below.
It uses the proxy's model_saturation_check and priority_model keys and
scripts, its refusal messages and headers, and adds a successful
response's tokens to the group and the pool admission chose. The rate
limit counters gain a general check-and-add over request and token
pairs, and the throttling refusal shape is shared.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5
* test(gateway-auth): expect the narrowed unsupported rate limit message
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5
* feat(gateway): prometheus observer with the proxy's metric and label names
Adds the prometheus observer, turned on by listing prometheus in
litellm_settings.callbacks, success_callback or failure_callback. Each
finished request updates the proxy's request, spend, token, latency and
deployment metrics under its names and labels, and /metrics serves them
after the gateway's own. prometheus_metrics_config, the exclude lists
and prometheus_latency_buckets are read as the proxy reads them. The
metrics the gateway already served under the same names come from the
observer alone when it is on. Outcomes now carry when the attempt was
sent upstream, for the LLM API latency and time to first token.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5
* feat(gateway): s3_v2 observer writing the proxy's logging payload to S3
Adds the s3_v2 observer, turned on by listing s3_v2 (or s3_log) in
litellm_settings.callbacks, success_callback or failure_callback. Each
finished request becomes the proxy's StandardLoggingPayload, with the
messages from the request body and the response as the client got it,
streams put back together. It is queued and written every flush interval
under the proxy's key layout, with its Content-MD5, Content-Disposition
and encryption headers, one PUT per request or one NDJSON object per
folder with s3_batch_file_upload. s3_callback_params is read as the
proxy reads it, os.environ references included. Uploads are retried,
dropped and kept for the next flush by the proxy's rules, and what is
queued is written at shutdown. auth-aws gains an S3 signer that signs
the path as sent and every header, as botocore does.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0138knkqKpkoSr6PApenZan5
* feat(gateway): Azure OpenAI chat completions through the OpenAI passthrough
Adds azure/<deployment> chat completions, streaming and not, on the same
passthrough the openai/ route uses. The request goes to the deployment
path with the api-version query, or to /openai/v1/ for the v1, latest and
preview versions, with api_base, api_key and api_version falling back to
the AZURE_* variables and os.environ references resolved as the proxy
does. A key is sent as api-key, and an azure_ad_token or an Entra ID
token from auth-azure goes as a Bearer token. The body follows the
proxy's gpt, o-series and gpt-5 param rules, picked by base_model or the
deployment name and reading model capabilities from the cost map, so an
unsupported param is refused unless drop_params allows it to be dropped.
Spend is priced as the proxy prices Azure, by azure/<base_model> when one
is set and otherwise by the model the response names, and a content
filter error keeps Azure's inner error for the client
* feat(gateway): Anthropic and Bedrock chat completions streaming
Adds "stream": true to /v1/chat/completions for anthropic/ and bedrock/
deployments, which were refused with 501 before. The request is the
non-streaming one with "stream": true added for Anthropic and sent to
/converse-stream for Bedrock, and the provider's events are decoded the
way the proxy's chunk parsers decode them: text, tool call deltas with
empty arguments sent as "{}", reasoning_content and thinking_blocks with
their index and signature, citations and the mapped finish_reason. A
wrapper in the shape of the proxy's stream wrapper gives every chunk one
id and the client's model name, sends role on the first chunk, adds the
empty finish chunk, strips usage from the chunks and, when
stream_options.include_usage is set, ends with the usage chunk before
[DONE]. The gateway waits for the first chunk before answering, so an
upstream error up to there is retried or falls back to another
deployment, and an error after it ends the stream with the OpenAI error
frame. Usage, cache reads and writes included, is handed to the meter
through the response whether or not the client asked for it, and
dropping the client stream drops the upstream request
* feat(gateway): full OpenAI chat request translation for Anthropic and Bedrock
Anthropic and Bedrock Converse deployments now take the whole OpenAI chat request the
Python proxy accepts instead of a narrow allowlist. Messages translate as Python builds
them: system as a string, parts or several messages, text, image_url (URL or data URL),
file and document parts, assistant tool_calls, the tool and function roles, consecutive
role merging, empty content, cache_control on Anthropic and cachePoint on Bedrock.
Params map per provider: max_tokens and max_completion_tokens with Anthropic's model
default, temperature, top_p, stop, tools and tool_choice, parallel_tool_calls,
response_format as a json tool or native output format, reasoning_effort and thinking,
user as metadata, with n and any other unsupported param refused with a 400 unless
drop_params is set globally, on the deployment or on the request, and
allowed_openai_params and additional_drop_params honoured. Responses carry tool calls,
thinking blocks and response_format answers back in OpenAI shape, streamed and not. The
route downloads media URLs a provider cannot take by reference, merges the beta headers
a request needs with forwarded ones, and the gateway hands the cost map to both
providers for the capability rules that read it.
* fix(rust-bridge): read Python's model_cost for chat translation capabilities
Chat translation for Anthropic and Bedrock reads capability flags from the
cost map. The Python bridge passed none, so every SDK chat call through the
Rust path failed. It now reads litellm.model_cost live, so lookups match
Python's, runtime-registered models included. ModelInfo::entry returns a Cow
so a map read from Python can hand over a copy.
Two tests sent temperature != 1 to a model that removed sampling params,
which Python itself refuses; they now send temperature=1.
* fix(model-catalog): read token counts the cost map writes as whole floats
The published cost map now has prompt_cache_min_tokens: 4096.0, which the
typed u64 fields rejected, so the gateway dropped the whole map at startup
and priced every request at zero. Counts accept any whole non-negative
number, as Python reads any JSON number, and a test keeps the checked-in
map parsing.
* feat(gateway): Azure OpenAI chat completions through the OpenAI passthrough (#47074)
Adds azure/<deployment> chat completions, streaming and not, on the same
passthrough the openai/ route uses. The request goes to the deployment
path with the api-version query, or to /openai/v1/ for the v1, latest and
preview versions, with api_base, api_key and api_version falling back to
the AZURE_* variables and os.environ references resolved as the proxy
does. A key is sent as api-key, and an azure_ad_token or an Entra ID
token from auth-azure goes as a Bearer token. The body follows the
proxy's gpt, o-series and gpt-5 param rules, picked by base_model or the
deployment name and reading model capabilities from the cost map, so an
unsupported param is refused unless drop_params allows it to be dropped.
Spend is priced as the proxy prices Azure, by azure/<base_model> when one
is set and otherwise by the model the response names, and a content
filter error keeps Azure's inner error for the client
* feat(gateway): Anthropic and Bedrock chat completions streaming (#47077)
Adds "stream": true to /v1/chat/completions for anthropic/ and bedrock/
deployments, which were refused with 501 before. The request is the
non-streaming one with "stream": true added for Anthropic and sent to
/converse-stream for Bedrock, and the provider's events are decoded the
way the proxy's chunk parsers decode them: text, tool call deltas with
empty arguments sent as "{}", reasoning_content and thinking_blocks with
their index and signature, citations and the mapped finish_reason. A
wrapper in the shape of the proxy's stream wrapper gives every chunk one
id and the client's model name, sends role on the first chunk, adds the
empty finish chunk, strips usage from the chunks and, when
stream_options.include_usage is set, ends with the usage chunk before
[DONE]. The gateway waits for the first chunk before answering, so an
upstream error up to there is retried or falls back to another
deployment, and an error after it ends the stream with the OpenAI error
frame. Usage, cache reads and writes included, is handed to the meter
through the response whether or not the client asked for it, and
dropping the client stream drops the upstream request
---------
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
* feat(gateway): gemini chat completions on google ai studio and vertex ai (#47098)
* feat(gateway): Gemini chat completions on Google AI Studio and Vertex AI
gemini/ deployments and vertex_ai/ deployments of Gemini models now serve
/v1/chat/completions, streamed and not, instead of being refused. Both surfaces share
one translation ported from the Python VertexGeminiConfig. Params map as Python maps
them: temperature, top_p, max_tokens and max_completion_tokens, n, stop, penalties
where the model takes them, seed, logprobs, tools with function declarations and the
built in search, code execution, url context, maps and computer use tools,
tool_choice, response_format as a response schema, reasoning_effort and thinking as a
thinking config, modalities, audio and labels. Messages become contents with a system
instruction when the cost map allows it, images, files, audio, function calls and
responses, and thought signatures carried back on tool call ids
AI Studio sends the key in x-goog-api-key from the deployment or GOOGLE_API_KEY,
GEMINI_API_KEY or PALM_API_KEY, and calls v1alpha for Gemini 3 and v1beta otherwise.
Vertex AI calls the project and location URL with a forwarded bearer or a Google
access token from vertex_credentials, the environment or ADC, and grafts the path onto
a custom api_base. Responses carry text, reasoning, tool calls, grounding annotations,
logprobs and the vertex_ai metadata, a blocked prompt or flagged answer comes back as
content_filter, and usage counts thoughts and cached tokens as Python does so cost is
recorded. The gateway resolves os.environ references in the deployment key, base and
Vertex params, passes the deployment Vertex params to the call, and loads the cost map
for gemini and vertex_ai deployments. Claude and partner models on Vertex AI stay
refused
* test(gateway): gemini cost map test double returns a Cow entry
* feat(gateway): native chat for openai-compatible providers (#47099)
Serve chat completions for hosted_vllm, nvidia_nim, groq, together_ai,
fireworks_ai, xai, mistral, deepseek, openrouter, cerebras and perplexity
through the OpenAI passthrough, with each provider's Python config ported
as a request transform in llms and its response and stream fixups applied
on the way back.
Streams ask for stream_options.include_usage on the client's behalf where
the provider takes it, record the usage through ReportedUsage for the spend
hooks, and keep it from the client unless the client asked for it.
* feat(gateway): claude chat completions on vertex ai (#47101)
vertex_ai/ deployments of Claude models now serve /v1/chat/completions,
streamed and not, instead of being refused. The request is the Anthropic
chat translation built for the vertex_ai host, as Python's
VertexAIAnthropicConfig builds it on top of AnthropicConfig: capabilities
read the vertex_ai/ cost map entry first and native structured output
reads only that entry, service_tier and speed are refused or dropped with
drop_params, and the Anthropic API betas are not added. The body drops
model, carries anthropic_version vertex-2023-10-16, folds output_format
into output_config and keeps effort only for models that take it, and
names its context management and caller betas in anthropic_beta and the
anthropic-beta header, filtered to what Vertex accepts, with the web
search beta alone when there are none.
The call goes to the publisher rawPredict URL for the project and
location, regional or global, or a custom api_base or VERTEXAI_API_BASE,
with a Google access token from vertex_credentials, the environment or
ADC. Streams call streamRawPredict with stream true, alt=sse on a custom
base, decode the Anthropic events as the anthropic/ path does and hand
usage to the hooks through ReportedUsage. Every remote image is
downloaded first, since Vertex takes none by URL, and the response names
the requested model.
* fix(gateway): find anthropic stream usage by event name (#47206)
The stream meter kept a data line only when "message_start" or
"message_delta" appeared in its first 96 bytes. An upstream that puts
"type" after the message object (JSON key order is not fixed) had every
message_start skipped, so input tokens were recorded as 0 and spend,
budgets, rate limits, metrics and S3 logs undercounted.
Anthropic streams name every event, so the event: line now decides; the
byte-window check remains for streams without event lines.
* fix(gateway): stop hooks pinning connection read buffers for a stream's life (#47220)
The s3 logger asks for the response text, and the tap kept every
response chunk as a Bytes slice of the upstream connection's read
buffer, so each open stream held a buffer per chunk until it ended
(about 460 KB per open 10k-token stream locally). Chunks are now copied
into one owned buffer.
The same applies on the request side: the facts kept the raw request
slice and a clone of the client headers, both slices of the client
connection's read buffer. The body is now released once hooks have run
(an owned copy is kept only when an observer reads it) and the headers
are copied.
* feat(loadtest): full-feature mode with recorded-data verification (#47207)
* feat(loadtest): full-feature mode with recorded-data verification
genconfig -features full turns on what both gateways do per request:
owners with rpm/tpm/budgets/max_parallel, deployment capacity, the
dynamic_rate_limiter_v3, prometheus and s3_v2 callbacks, and the four
request hooks with Python stand-ins (custom_callbacks.py). -key-sets
gives each gateway its own keys, users and teams in a shared database.
verify checks spend logs, owner spend, daily tables, Prometheus counters,
S3 logs and the headers the provider received against what loadgen sent,
and compares per-deployment prices across sides. s3stub serves S3 for
the logs. mockllm takes in= for exact input tokens, returns unique
response ids, and counts chosen request headers; loadgen takes -keys
and -header.
* feat(loadtest): s3stub -max-string keeps long runs' logs in memory
* feat(loadtest): bifrost as a third side of the full-feature run (#47236)
genconfig writes a "bifrost" key set as a Bifrost config: a custom provider
per deployment, a routing rule per model group, virtual keys, and with
-features full key/team/customer budgets and rate limits, model configs for
deployment limits, the logs store with S3 payloads, the telemetry plugin and
header forwarding/stripping/logging. A pricing file copies the cost map's
prices per deployment.
mockllm answers Bedrock Converse (streamed and not), serves a self-signed
TLS listener for Bifrost's HTTPS-only Bedrock provider, and echoes the
request's model, which Bifrost prices by. s3stub answers HeadBucket.
loadgen takes -chat-path and -messages-path for Bifrost's
/anthropic/v1/messages.
verify reads a side named bifrost in Bifrost's schema and reports what it
has no counterpart for (users, members, daily tables, provider_metadata,
cross-replica state) as n/a; with three sides it compares every pair.
Claude-Session: https://claude.ai/code/session_01XF7rGf49TWYH5Hqkb14eSL
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
* build(deps): bump aiohttp to 3.14.5 to clear osv-scan (#47369)
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---------
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: Mateo Wang <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Nate Armstrong
GitHub
parent
6dc59535e3
commit
6f345e3eb9
149 changed files with 25167 additions and 1304 deletions
3
litellm-rust/Cargo.lock
generated
3
litellm-rust/Cargo.lock
generated
|
|
@ -4055,6 +4055,7 @@ dependencies = [
|
|||
name = "litellm-inference-chat"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"base64 0.22.1",
|
||||
"bytes",
|
||||
"futures-util",
|
||||
"litellm-auth",
|
||||
|
|
@ -4235,11 +4236,13 @@ dependencies = [
|
|||
"serde_json",
|
||||
"serde_path_to_error",
|
||||
"serde_with",
|
||||
"sha2 0.10.9",
|
||||
"strum",
|
||||
"thiserror 2.0.19",
|
||||
"time",
|
||||
"tokio",
|
||||
"url",
|
||||
"uuid",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
|
|||
|
|
@ -21,6 +21,28 @@ const FINISH_REASONS: &[(&str, &str)] = &[
|
|||
("content_filter", "content_filter"),
|
||||
("stop", "stop"),
|
||||
("length", "length"),
|
||||
("STOP", "stop"),
|
||||
("MAX_TOKENS", "length"),
|
||||
("CONTINUATION", "length"),
|
||||
("SAFETY", "content_filter"),
|
||||
("RECITATION", "content_filter"),
|
||||
("FINISH_REASON_UNSPECIFIED", "stop"),
|
||||
("MALFORMED_FUNCTION_CALL", "stop"),
|
||||
("LANGUAGE", "content_filter"),
|
||||
("OTHER", "content_filter"),
|
||||
("BLOCKLIST", "content_filter"),
|
||||
("PROHIBITED_CONTENT", "content_filter"),
|
||||
("SPII", "content_filter"),
|
||||
("IMAGE_SAFETY", "content_filter"),
|
||||
("IMAGE_PROHIBITED_CONTENT", "content_filter"),
|
||||
("TOO_MANY_TOOL_CALLS", "stop"),
|
||||
("MALFORMED_RESPONSE", "stop"),
|
||||
("NO_IMAGE", "content_filter"),
|
||||
("IMAGE_RECITATION", "content_filter"),
|
||||
("IMAGE_OTHER", "content_filter"),
|
||||
("ESCALATION", "content_filter"),
|
||||
("UNEXPECTED_TOOL_CALL", "stop"),
|
||||
("MISSING_THOUGHT_SIGNATURE", "stop"),
|
||||
];
|
||||
|
||||
pub fn finish_reason_for(provider_reason: &str) -> &'static str {
|
||||
|
|
@ -84,6 +106,11 @@ mod tests {
|
|||
#[case::guardrail_intervened("guardrail_intervened", "content_filter")]
|
||||
#[case::content_filtered("content_filtered", "content_filter")]
|
||||
#[case::content_filter("content_filter", "content_filter")]
|
||||
#[case::gemini_stop("STOP", "stop")]
|
||||
#[case::gemini_max_tokens("MAX_TOKENS", "length")]
|
||||
#[case::gemini_safety("SAFETY", "content_filter")]
|
||||
#[case::gemini_malformed_call("MALFORMED_FUNCTION_CALL", "stop")]
|
||||
#[case::gemini_other("OTHER", "content_filter")]
|
||||
fn maps_every_reason_the_route_can_observe(#[case] reason: &str, #[case] expected: &str) {
|
||||
assert_eq!(finish_reason_for(reason), expected);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -21,28 +21,48 @@ pub enum LlmProviders {
|
|||
Bedrock,
|
||||
#[strum(serialize = "bedrock_mantle")]
|
||||
BedrockMantle,
|
||||
#[strum(serialize = "cerebras")]
|
||||
Cerebras,
|
||||
#[strum(serialize = "cohere")]
|
||||
Cohere,
|
||||
#[strum(serialize = "deepseek")]
|
||||
Deepseek,
|
||||
#[strum(serialize = "edenai")]
|
||||
Edenai,
|
||||
#[strum(serialize = "fireworks_ai")]
|
||||
FireworksAi,
|
||||
/// Google AI Studio (`gemini/<model>`), the Gemini API addressed with an API key.
|
||||
#[strum(serialize = "gemini")]
|
||||
Gemini,
|
||||
#[strum(serialize = "groq")]
|
||||
Groq,
|
||||
/// A vLLM server (`hosted_vllm/<model>`) at the deployment's `api_base`.
|
||||
#[strum(serialize = "hosted_vllm")]
|
||||
HostedVllm,
|
||||
#[strum(serialize = "minimax")]
|
||||
Minimax,
|
||||
#[strum(serialize = "mistral")]
|
||||
Mistral,
|
||||
#[strum(serialize = "nvidia_nim")]
|
||||
NvidiaNim,
|
||||
#[strum(serialize = "openai")]
|
||||
Openai,
|
||||
#[strum(serialize = "openai_like")]
|
||||
OpenaiLike,
|
||||
#[strum(serialize = "openrouter")]
|
||||
Openrouter,
|
||||
#[strum(serialize = "perplexity")]
|
||||
Perplexity,
|
||||
#[strum(serialize = "reducto")]
|
||||
Reducto,
|
||||
#[strum(serialize = "tencent")]
|
||||
Tencent,
|
||||
#[strum(serialize = "together_ai")]
|
||||
TogetherAi,
|
||||
#[strum(serialize = "vertex_ai")]
|
||||
VertexAi,
|
||||
#[strum(serialize = "xai")]
|
||||
Xai,
|
||||
}
|
||||
|
||||
pub fn get_custom_llm_provider<'a>(
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
use std::{
|
||||
ops::{BitOr, BitOrAssign},
|
||||
sync::{Arc, OnceLock},
|
||||
sync::{Arc, Mutex, OnceLock},
|
||||
time::{Duration, SystemTime},
|
||||
};
|
||||
|
||||
|
|
@ -34,7 +34,8 @@ impl Facts {
|
|||
pub const TIMINGS: Self = Self(1 << 7);
|
||||
/// Values earlier hooks attached for observers.
|
||||
pub const ANNOTATIONS: Self = Self(1 << 8);
|
||||
/// The raw request bytes, already held for retries.
|
||||
/// The raw request bytes, already held for retries. Hooks read them while the request is
|
||||
/// admitted; only an observer that asks for them keeps a copy for the rest of the request.
|
||||
pub const REQUEST_BODY: Self = Self(1 << 9);
|
||||
/// The response text. The only fact that makes the gateway keep a stream's bytes.
|
||||
pub const RESPONSE_BODY: Self = Self(1 << 10);
|
||||
|
|
@ -121,7 +122,11 @@ pub struct RequestFacts {
|
|||
/// The output budget the request set, if any.
|
||||
pub max_tokens: Option<u64>,
|
||||
pub started_at: SystemTime,
|
||||
body: Bytes,
|
||||
/// The request bytes while hooks run. Afterwards a compact copy when an observer reads them,
|
||||
/// else empty: the raw bytes are a slice of the connection's read buffer, and holding them
|
||||
/// would hold the whole buffer for as long as the response streams.
|
||||
body: Mutex<Bytes>,
|
||||
body_len: usize,
|
||||
input_tokens: OnceLock<u64>,
|
||||
}
|
||||
|
||||
|
|
@ -148,21 +153,38 @@ impl RequestFacts {
|
|||
stream,
|
||||
max_tokens,
|
||||
started_at: SystemTime::now(),
|
||||
body,
|
||||
body_len: body.len(),
|
||||
body: Mutex::new(body),
|
||||
input_tokens: OnceLock::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// The raw request bytes. Shared, never copied.
|
||||
pub fn body(&self) -> &Bytes {
|
||||
&self.body
|
||||
/// The request bytes: in full while hooks run, afterwards only for observers that read
|
||||
/// [`Facts::REQUEST_BODY`], empty for everyone else.
|
||||
pub fn body(&self) -> Bytes {
|
||||
self.body
|
||||
.lock()
|
||||
.unwrap_or_else(|poison| poison.into_inner())
|
||||
.clone()
|
||||
}
|
||||
|
||||
/// Lets go of the raw bytes once hooks have run, keeping an owned copy when `keep`.
|
||||
pub(crate) fn release_body(&self, keep: bool) {
|
||||
let mut body = self
|
||||
.body
|
||||
.lock()
|
||||
.unwrap_or_else(|poison| poison.into_inner());
|
||||
*body = match keep {
|
||||
true => Bytes::copy_from_slice(&body),
|
||||
false => Bytes::new(),
|
||||
};
|
||||
}
|
||||
|
||||
/// The proxy's cheap input estimate, a token per four bytes of request, made once.
|
||||
pub fn input_tokens(&self) -> u64 {
|
||||
*self
|
||||
.input_tokens
|
||||
.get_or_init(|| (self.body.len() as u64).div_ceil(4).max(1))
|
||||
.get_or_init(|| (self.body_len as u64).div_ceil(4).max(1))
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -63,6 +63,21 @@ impl Pipeline {
|
|||
/// Runs every hook in order. When one refuses, the hooks before it are told the request
|
||||
/// ended with that refusal, and so are the observers.
|
||||
pub async fn admit(&self, request: Arc<RequestFacts>) -> Result<Admission, Refusal> {
|
||||
let admitted = self.run_hooks(request).await;
|
||||
// Refused or not, the request's facts outlive this call (in observer records, or for
|
||||
// the whole response), and should not keep the connection's read buffer with them.
|
||||
let request = match &admitted {
|
||||
Ok(admission) => &admission.request,
|
||||
Err((request, _)) => request,
|
||||
};
|
||||
request.release_body(self.observers.reads().contains(Facts::REQUEST_BODY));
|
||||
admitted.map_err(|(_, refusal)| refusal)
|
||||
}
|
||||
|
||||
async fn run_hooks(
|
||||
&self,
|
||||
request: Arc<RequestFacts>,
|
||||
) -> Result<Admission, (Arc<RequestFacts>, Refusal)> {
|
||||
let mut admission = Admission {
|
||||
request,
|
||||
edit: RequestEdit::default(),
|
||||
|
|
@ -84,8 +99,9 @@ impl Pipeline {
|
|||
ended_at: SystemTime::now(),
|
||||
interrupted: false,
|
||||
};
|
||||
let request = Arc::clone(&admission.request);
|
||||
self.finish(admission, outcome, None);
|
||||
return Err(refusal);
|
||||
return Err((request, refusal));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -304,6 +320,55 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
struct Bodies;
|
||||
|
||||
impl Observer for Bodies {
|
||||
fn name(&self) -> &'static str {
|
||||
"bodies"
|
||||
}
|
||||
|
||||
fn reads(&self) -> Facts {
|
||||
Facts::REQUEST_BODY
|
||||
}
|
||||
|
||||
fn observe<'a>(&'a self, _record: &'a Record) -> BoxFuture<'a, ()> {
|
||||
Box::pin(async {})
|
||||
}
|
||||
}
|
||||
|
||||
/// Hooks see the request as sent; afterwards the facts hold no slice of the read buffer,
|
||||
/// only a copy when an observer reads the body.
|
||||
#[rstest::rstest]
|
||||
#[case::no_observer_reads_it(Vec::new(), b"")]
|
||||
#[case::an_observer_reads_it(vec![Arc::new(Bodies) as Arc<dyn Observer>], b"{\"model\":\"m\"}")]
|
||||
#[tokio::test]
|
||||
async fn the_raw_body_is_released_after_the_hooks(
|
||||
#[case] observers: Vec<Arc<dyn Observer>>,
|
||||
#[case] kept: &[u8],
|
||||
) {
|
||||
let buffer = Bytes::from_static(b"{\"model\":\"m\"} and the rest of the read buffer");
|
||||
let raw = buffer.slice(..13);
|
||||
let request = Arc::new(RequestFacts::new(
|
||||
"r".into(),
|
||||
Route::Messages,
|
||||
None,
|
||||
None,
|
||||
"m".into(),
|
||||
Arc::new(HeaderMap::new()),
|
||||
false,
|
||||
None,
|
||||
raw.clone(),
|
||||
));
|
||||
assert_eq!(request.body().as_ptr(), raw.as_ptr());
|
||||
let pipeline = Pipeline::new(Vec::new(), Observers::start(observers, 8), Arc::default());
|
||||
let admission = pipeline.admit(request).await.unwrap();
|
||||
|
||||
let body = admission.request().body();
|
||||
assert_eq!(&body[..], kept);
|
||||
assert!(body.is_empty() || body.as_ptr() != raw.as_ptr());
|
||||
assert_eq!(admission.request().input_tokens(), 4);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_refusal_settles_the_hooks_before_it() {
|
||||
let log = Log::default();
|
||||
|
|
|
|||
|
|
@ -206,6 +206,15 @@ enum Line {
|
|||
Skipping,
|
||||
}
|
||||
|
||||
/// What the stream's `event:` line named for the event in progress. Anthropic streams name every
|
||||
/// event, which decides a data line whatever order its JSON keys come in.
|
||||
#[derive(Clone, Copy, Debug, PartialEq)]
|
||||
enum Event {
|
||||
Unnamed,
|
||||
Usage,
|
||||
Other,
|
||||
}
|
||||
|
||||
/// Reads usage out of a server-sent-event stream as its chunks pass through. Only the line in
|
||||
/// progress is buffered, and only while it may carry usage.
|
||||
#[derive(Debug)]
|
||||
|
|
@ -213,6 +222,7 @@ pub struct StreamUsage {
|
|||
format: StreamFormat,
|
||||
line: Vec<u8>,
|
||||
state: Line,
|
||||
event: Event,
|
||||
usage: Option<Usage>,
|
||||
model: Option<String>,
|
||||
}
|
||||
|
|
@ -223,6 +233,7 @@ impl StreamUsage {
|
|||
format,
|
||||
line: Vec::new(),
|
||||
state: Line::Reading,
|
||||
event: Event::Unnamed,
|
||||
usage: None,
|
||||
model: None,
|
||||
}
|
||||
|
|
@ -239,8 +250,10 @@ impl StreamUsage {
|
|||
self.take(part);
|
||||
}
|
||||
if end.is_some() {
|
||||
if let Line::Reading = self.state {
|
||||
self.read_line();
|
||||
match self.state {
|
||||
Line::Reading => self.read_line(),
|
||||
// A skipped line is content, never the blank line that ends an event.
|
||||
Line::Skipping => {}
|
||||
}
|
||||
self.line.clear();
|
||||
self.state = Line::Reading;
|
||||
|
|
@ -294,23 +307,43 @@ impl StreamUsage {
|
|||
return false;
|
||||
}
|
||||
match self.format {
|
||||
StreamFormat::Anthropic if line.len() >= DECIDE_BYTES => {
|
||||
let head = &line[..DECIDE_BYTES];
|
||||
contains(head, b"\"message_start\"") || contains(head, b"\"message_delta\"")
|
||||
}
|
||||
StreamFormat::Anthropic if line.len() >= DECIDE_BYTES => match self.event {
|
||||
Event::Usage => true,
|
||||
Event::Other => false,
|
||||
Event::Unnamed => {
|
||||
let head = &line[..DECIDE_BYTES];
|
||||
contains(head, b"\"message_start\"") || contains(head, b"\"message_delta\"")
|
||||
}
|
||||
},
|
||||
StreamFormat::Anthropic | StreamFormat::OpenAi => true,
|
||||
}
|
||||
}
|
||||
|
||||
fn read_line(&mut self) {
|
||||
let line = self.line.strip_suffix(b"\r").unwrap_or(&self.line);
|
||||
if line.is_empty() {
|
||||
self.event = Event::Unnamed;
|
||||
return;
|
||||
}
|
||||
if let Some(name) = line.strip_prefix(b"event:") {
|
||||
let name = name.strip_prefix(b" ").unwrap_or(name);
|
||||
self.event = match name {
|
||||
b"message_start" | b"message_delta" => Event::Usage,
|
||||
_ => Event::Other,
|
||||
};
|
||||
return;
|
||||
}
|
||||
let Some(data) = line.strip_prefix(b"data:") else {
|
||||
return;
|
||||
};
|
||||
let data = data.strip_prefix(b" ").unwrap_or(data);
|
||||
match self.format {
|
||||
StreamFormat::Anthropic => {
|
||||
if !(contains(data, b"\"message_start\"") || contains(data, b"\"message_delta\"")) {
|
||||
let named = self.event == Event::Usage;
|
||||
if !named
|
||||
&& !(contains(data, b"\"message_start\"")
|
||||
|| contains(data, b"\"message_delta\""))
|
||||
{
|
||||
return;
|
||||
}
|
||||
let Ok(event) = serde_json::from_slice::<AnthropicEvent>(data) else {
|
||||
|
|
@ -417,6 +450,32 @@ mod tests {
|
|||
assert_eq!(meter.usage(), None);
|
||||
}
|
||||
|
||||
/// Key order is not part of JSON: an upstream may put `type` after the message object, past
|
||||
/// the decision window, and the event line still marks the usage events.
|
||||
#[rstest]
|
||||
#[case::whole(usize::MAX)]
|
||||
#[case::bytes(1)]
|
||||
#[case::small(7)]
|
||||
fn anthropic_usage_events_are_found_by_event_name(#[case] size: usize) {
|
||||
let stream = concat!(
|
||||
"event: message_start\n",
|
||||
"data: {\"message\":{\"content\":[],\"id\":\"msg_0123456789abcdef0123456789abcdef\",\"model\":\"claude-x\",\"role\":\"assistant\",\"stop_reason\":null,\"type\":\"message\",\"usage\":{\"input_tokens\":1234,\"output_tokens\":1}},\"type\":\"message_start\"}\n",
|
||||
"\n",
|
||||
"event: content_block_delta\n",
|
||||
"data: {\"delta\":{\"text\":\"text that mentions \\\"message_delta\\\" and is longer than the decision window\",\"type\":\"text_delta\"},\"index\":0,\"type\":\"content_block_delta\",\"usage\":{\"output_tokens\":999}}\n",
|
||||
"\n",
|
||||
"event: message_delta\n",
|
||||
"data: {\"delta\":{\"stop_reason\":\"end_turn\",\"stop_sequence\":null},\"type\":\"message_delta\",\"usage\":{\"output_tokens\":42}}\n",
|
||||
"\n",
|
||||
);
|
||||
let mut meter = StreamUsage::new(StreamFormat::Anthropic);
|
||||
for chunk in stream.as_bytes().chunks(size.min(stream.len())) {
|
||||
meter.feed(chunk);
|
||||
}
|
||||
let usage = meter.usage().expect("usage");
|
||||
assert_eq!((usage.input_tokens, usage.output_tokens), (1234, 42));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn crlf_lines_are_read() {
|
||||
let mut meter = StreamUsage::new(StreamFormat::Anthropic);
|
||||
|
|
|
|||
|
|
@ -21,6 +21,8 @@ use litellm_gateway_hooks::Route;
|
|||
|
||||
use crate::{Deployment, Error, Gateway, JsonBody, request};
|
||||
|
||||
const ENV_REFERENCE: &str = "os.environ/";
|
||||
|
||||
pub(crate) async fn create(
|
||||
State(gateway): State<Arc<Gateway>>,
|
||||
identity: AuthenticatedRequest,
|
||||
|
|
@ -135,17 +137,21 @@ async fn handle(
|
|||
) {
|
||||
return openai(gateway, deployment, body, extra_headers.clone()).await;
|
||||
}
|
||||
let (messages, optional_params) = split(body);
|
||||
let (messages, mut optional_params) = split(body);
|
||||
let stream = ChatCompletionsRoute::is_stream(&optional_params);
|
||||
let (api_key, api_base) =
|
||||
connection(gateway, deployment, &mut optional_params).await?;
|
||||
let call = ChatCompletionsCall {
|
||||
model: deployment.model.clone(),
|
||||
messages,
|
||||
optional_params,
|
||||
api_key: deployment.api_key.clone(),
|
||||
api_base: deployment.api_base.clone(),
|
||||
api_key,
|
||||
api_base,
|
||||
custom_llm_provider: deployment.custom_llm_provider.clone(),
|
||||
extra_headers: extra_headers.clone(),
|
||||
timeout: deployment.timeout,
|
||||
drop_params: Some(drop_params(gateway, deployment)),
|
||||
model_info: gateway.model_info(),
|
||||
};
|
||||
if stream {
|
||||
return translated_stream(route, call, name).await;
|
||||
|
|
@ -179,6 +185,60 @@ fn split(mut body: Map<String, Value>) -> (Value, Map<String, Value>) {
|
|||
(messages, body)
|
||||
}
|
||||
|
||||
/// A deployment's connection for a provider chat translates. Python's router reads each
|
||||
/// `os.environ/` setting from the environment when it loads the deployment, and `completion`
|
||||
/// hands the Vertex AI settings to the call, where a request body cannot name them.
|
||||
async fn connection(
|
||||
gateway: &Gateway,
|
||||
deployment: &Deployment,
|
||||
optional_params: &mut Map<String, Value>,
|
||||
) -> Result<(Option<String>, Option<String>), Error> {
|
||||
let vertex = match serde_json::to_value(&deployment.litellm_params.vertex) {
|
||||
Ok(Value::Object(fields)) => fields,
|
||||
_ => Map::new(),
|
||||
};
|
||||
let references: Vec<&str> = [
|
||||
deployment.api_key.as_deref(),
|
||||
deployment.api_base.as_deref(),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.chain(vertex.values().filter_map(Value::as_str))
|
||||
.filter_map(|value| value.strip_prefix(ENV_REFERENCE))
|
||||
.collect();
|
||||
let secrets = match references.is_empty() {
|
||||
true => None,
|
||||
false => Some(
|
||||
gateway
|
||||
.secrets
|
||||
.resolve(&references)
|
||||
.await
|
||||
.map_err(|error| Error::Internal(error.to_string()))?,
|
||||
),
|
||||
};
|
||||
// A reference to an unset variable counts as unset.
|
||||
let resolve = |value: &str| match value.strip_prefix(ENV_REFERENCE) {
|
||||
Some(name) => secrets.as_ref().and_then(|secrets| secrets.get(name)),
|
||||
None => Some(value.to_owned()),
|
||||
};
|
||||
for (name, value) in vertex {
|
||||
let value = match value {
|
||||
Value::String(text) => resolve(&text).map(Value::String),
|
||||
Value::Null => None,
|
||||
other => Some(other),
|
||||
};
|
||||
if let Some(value) = value
|
||||
&& !optional_params.contains_key(&name)
|
||||
{
|
||||
optional_params.insert(name, value);
|
||||
}
|
||||
}
|
||||
Ok((
|
||||
deployment.api_key.as_deref().and_then(resolve),
|
||||
deployment.api_base.as_deref().and_then(resolve),
|
||||
))
|
||||
}
|
||||
|
||||
/// The deployment's `litellm_params` a provider reads beyond its connection: its
|
||||
/// `api_version`, `base_model` and every key the config adds.
|
||||
fn provider_params(deployment: &Deployment) -> Map<String, Value> {
|
||||
|
|
@ -197,12 +257,24 @@ fn provider_params(deployment: &Deployment) -> Map<String, Value> {
|
|||
provider_params
|
||||
}
|
||||
|
||||
/// `litellm.drop_params` or the deployment's own: either one set drops what the provider
|
||||
/// cannot take, as Python's `get_optional_params` reads them.
|
||||
fn drop_params(gateway: &Gateway, deployment: &Deployment) -> bool {
|
||||
gateway.drop_params()
|
||||
|| deployment
|
||||
.litellm_params
|
||||
.drop_params
|
||||
.as_ref()
|
||||
.is_some_and(|flag| flag.is_set())
|
||||
}
|
||||
|
||||
async fn openai(
|
||||
gateway: &Gateway,
|
||||
deployment: &Deployment,
|
||||
body: Map<String, Value>,
|
||||
extra_headers: Option<Map<String, Value>>,
|
||||
) -> Result<Response, Error> {
|
||||
let usage = ReportedUsage::default();
|
||||
let machine = gateway.chat_completions.clone().openai_machine(
|
||||
OpenAiChatCall {
|
||||
model: deployment.model.clone(),
|
||||
|
|
@ -213,13 +285,9 @@ async fn openai(
|
|||
extra_headers,
|
||||
timeout: deployment.timeout,
|
||||
litellm_params: provider_params(deployment),
|
||||
drop_params: gateway.drop_params()
|
||||
|| deployment
|
||||
.litellm_params
|
||||
.drop_params
|
||||
.as_ref()
|
||||
.is_some_and(|flag| flag.is_set()),
|
||||
drop_params: drop_params(gateway, deployment),
|
||||
model_info: gateway.model_info(),
|
||||
usage: usage.clone(),
|
||||
},
|
||||
None,
|
||||
);
|
||||
|
|
@ -227,7 +295,9 @@ async fn openai(
|
|||
|body: Bytes| ([(CONTENT_TYPE, "application/json")], body),
|
||||
|error| Bytes::from(Error::from(error).openai_sse_frame()),
|
||||
);
|
||||
Ok(litellm_host_http::serve(machine, (), (), encoder, None).await?)
|
||||
let mut response = litellm_host_http::serve(machine, (), (), encoder, None).await?;
|
||||
response.extensions_mut().insert(usage);
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
/// A stream from a provider chat translates (Anthropic, Bedrock), sent as OpenAI chunks. The
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ use std::{
|
|||
|
||||
use axum::{
|
||||
body::{Body, BodyDataStream, Bytes, HttpBody},
|
||||
http::{HeaderMap, header::CONTENT_TYPE},
|
||||
http::{HeaderMap, HeaderName, HeaderValue, header::CONTENT_TYPE},
|
||||
response::Response,
|
||||
};
|
||||
use futures_util::Stream;
|
||||
|
|
@ -81,7 +81,7 @@ pub(crate) async fn admit(
|
|||
.then(|| key_facts(hooks, identity))
|
||||
.flatten();
|
||||
let client_headers = match pipeline.reads(Facts::CLIENT_HEADERS) {
|
||||
true => Arc::new(headers.clone()),
|
||||
true => Arc::new(owned_headers(headers)),
|
||||
false => Arc::default(),
|
||||
};
|
||||
let request = RequestFacts::new(
|
||||
|
|
@ -107,6 +107,21 @@ pub(crate) async fn admit(
|
|||
}))
|
||||
}
|
||||
|
||||
/// A copy of the client's headers that owns its bytes. Hyper parses header values, and names it
|
||||
/// does not know, as slices of the connection's read buffer, so a plain clone kept in the facts
|
||||
/// would keep that whole buffer for as long as the response streams.
|
||||
fn owned_headers(headers: &HeaderMap) -> HeaderMap {
|
||||
let mut owned = HeaderMap::with_capacity(headers.len());
|
||||
for (name, value) in headers {
|
||||
let name =
|
||||
HeaderName::from_bytes(name.as_str().as_bytes()).unwrap_or_else(|_| name.clone());
|
||||
let mut copy = HeaderValue::from_bytes(value.as_bytes()).unwrap_or_else(|_| value.clone());
|
||||
copy.set_sensitive(value.is_sensitive());
|
||||
owned.append(name, copy);
|
||||
}
|
||||
owned
|
||||
}
|
||||
|
||||
fn key_facts(hooks: &Hooks, identity: &AuthenticatedRequest) -> Option<Arc<KeyFacts>> {
|
||||
let caller = identity.caller();
|
||||
let authentication = caller.authentication();
|
||||
|
|
@ -277,7 +292,7 @@ enum Meter {
|
|||
Off,
|
||||
Stream(StreamUsage),
|
||||
/// A JSON body, read for its `usage` once complete.
|
||||
Body(Vec<Bytes>),
|
||||
Body(Vec<u8>),
|
||||
}
|
||||
|
||||
/// Records the request when the response body ends or is dropped.
|
||||
|
|
@ -291,7 +306,10 @@ struct Finish {
|
|||
meter: Meter,
|
||||
/// The usage a translated chat stream ended with, sent to the client or not.
|
||||
reported: Option<ReportedUsage>,
|
||||
body: Option<Vec<Bytes>>,
|
||||
/// The response text for observers that read it. Copied, not kept as chunks: each chunk is a
|
||||
/// slice of the upstream connection's read buffer, and holding it would hold that buffer for
|
||||
/// the rest of the stream.
|
||||
body: Option<Vec<u8>>,
|
||||
started: Instant,
|
||||
first_byte: Option<Duration>,
|
||||
ended: bool,
|
||||
|
|
@ -304,10 +322,10 @@ impl Finish {
|
|||
match &mut self.meter {
|
||||
Meter::Off => {}
|
||||
Meter::Stream(meter) => meter.feed(chunk),
|
||||
Meter::Body(chunks) => chunks.push(chunk.clone()),
|
||||
Meter::Body(body) => body.extend_from_slice(chunk),
|
||||
}
|
||||
if let Some(body) = &mut self.body {
|
||||
body.push(chunk.clone());
|
||||
body.extend_from_slice(chunk);
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -322,10 +340,9 @@ impl Finish {
|
|||
Some(usage) => (Some(usage), None),
|
||||
None => (meter.usage(), meter.model().map(str::to_owned)),
|
||||
},
|
||||
Meter::Body(chunks) => {
|
||||
let body = chunks.concat();
|
||||
let model = by_model.then(|| body_model(&body)).flatten();
|
||||
(body_usage(self.route, &body), model)
|
||||
Meter::Body(body) => {
|
||||
let model = by_model.then(|| body_model(body)).flatten();
|
||||
(body_usage(self.route, body), model)
|
||||
}
|
||||
};
|
||||
let cost = by_model
|
||||
|
|
@ -356,7 +373,7 @@ impl Drop for Finish {
|
|||
ended_at: SystemTime::now(),
|
||||
interrupted: !self.ended,
|
||||
};
|
||||
let body = self.body.take().map(|chunks| Bytes::from(chunks.concat()));
|
||||
let body = self.body.take().map(Bytes::from);
|
||||
self.pipeline.finish(admission, outcome, body);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -18,7 +18,10 @@ mod responses;
|
|||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use axum::{Router, routing::post};
|
||||
use litellm_http::{ClientVariant, HttpClientConfig, media::UrlPolicy};
|
||||
use litellm_http::{
|
||||
ClientVariant, HttpClientConfig,
|
||||
media::{MediaFetcher, UrlPolicy},
|
||||
};
|
||||
use litellm_inference::resources::CoreResources;
|
||||
use litellm_inference_chat::ChatCompletionsRoute;
|
||||
use litellm_inference_messages::MessagesRoute;
|
||||
|
|
@ -138,7 +141,12 @@ impl Gateway {
|
|||
provider.clone(),
|
||||
auth.clone(),
|
||||
secrets.clone(),
|
||||
),
|
||||
)
|
||||
.with_media(MediaFetcher::new(
|
||||
&resources.pool,
|
||||
&http,
|
||||
UrlPolicy::default(),
|
||||
)?),
|
||||
messages: MessagesRoute::new(provider.clone(), auth.clone(), secrets.clone()),
|
||||
responses: ResponsesRoute::new(provider, auth.clone(), secrets.clone()),
|
||||
ocr: OcrRoute::new(OcrClient::new(
|
||||
|
|
|
|||
|
|
@ -105,6 +105,8 @@ async fn chat_errors_come_from_core(
|
|||
custom_llm_provider: None,
|
||||
extra_headers: None,
|
||||
timeout: None,
|
||||
drop_params: Some(false),
|
||||
model_info: None,
|
||||
},
|
||||
&(),
|
||||
None,
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
//! Builds the hooks and observers inference requests run through, and the per-deployment prices
|
||||
//! observers that read cost need.
|
||||
|
||||
use std::{collections::HashSet, sync::Arc};
|
||||
use std::{borrow::Cow, collections::HashSet, sync::Arc};
|
||||
|
||||
use litellm_config::{Config, Model};
|
||||
use litellm_gateway_hooks::{
|
||||
|
|
@ -30,19 +30,19 @@ pub async fn install(
|
|||
hooks: Vec<Arc<dyn Hook>>,
|
||||
observers: Vec<Arc<dyn Observer>>,
|
||||
) -> Result<(), Error> {
|
||||
let serves_azure = config.model_list.iter().any(is_azure);
|
||||
if hooks.is_empty() && observers.is_empty() && !serves_azure {
|
||||
let reads_model_info = config.model_list.iter().any(reads_model_info);
|
||||
if hooks.is_empty() && observers.is_empty() && !reads_model_info {
|
||||
return Ok(());
|
||||
}
|
||||
let needs_cost = observers
|
||||
.iter()
|
||||
.any(|observer| observer.reads().contains(Facts::COST))
|
||||
|| hooks.iter().any(|hook| hook.reads().contains(Facts::COST));
|
||||
let catalog = match needs_cost || serves_azure {
|
||||
let catalog = match needs_cost || reads_model_info {
|
||||
true => load_catalog(inference).await.map(Arc::new),
|
||||
false => None,
|
||||
};
|
||||
if let Some(catalog) = catalog.as_ref().filter(|_| serves_azure) {
|
||||
if let Some(catalog) = catalog.as_ref().filter(|_| reads_model_info) {
|
||||
inference.install_model_info(Arc::new(CostMap::new(catalog.clone())));
|
||||
}
|
||||
if hooks.is_empty() && observers.is_empty() {
|
||||
|
|
@ -204,7 +204,39 @@ fn is_azure(model: &Model) -> bool {
|
|||
}
|
||||
}
|
||||
|
||||
/// The model cost map as the capabilities Azure's param rules read.
|
||||
/// Whether the deployment's chat param rules read model capabilities from the cost map:
|
||||
/// Azure, the Anthropic, Bedrock, Gemini and Claude on Vertex AI chat translations, and the
|
||||
/// OpenAI-compatible providers.
|
||||
fn reads_model_info(model: &Model) -> bool {
|
||||
is_azure(model)
|
||||
|| litellm_inference::provider::resolve_llm_provider(
|
||||
&model.litellm_params.model,
|
||||
model.litellm_params.custom_llm_provider.as_deref(),
|
||||
"chat completions",
|
||||
)
|
||||
.is_ok_and(|resolved| {
|
||||
matches!(
|
||||
<&str>::from(resolved.provider),
|
||||
"anthropic"
|
||||
| "bedrock"
|
||||
| "gemini"
|
||||
| "vertex_ai"
|
||||
| "cerebras"
|
||||
| "deepseek"
|
||||
| "fireworks_ai"
|
||||
| "groq"
|
||||
| "hosted_vllm"
|
||||
| "mistral"
|
||||
| "nvidia_nim"
|
||||
| "openrouter"
|
||||
| "perplexity"
|
||||
| "together_ai"
|
||||
| "xai"
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
/// The model cost map as the capabilities the chat param rules read.
|
||||
struct CostMap {
|
||||
catalog: Arc<Catalog>,
|
||||
keys: HashSet<String>,
|
||||
|
|
@ -222,8 +254,10 @@ impl CostMap {
|
|||
}
|
||||
|
||||
impl ModelInfo for CostMap {
|
||||
fn entry(&self, model: &str) -> Option<&Map<String, Value>> {
|
||||
self.catalog.lookup(model).map(|found| found.entry.fields())
|
||||
fn entry(&self, model: &str) -> Option<Cow<'_, Map<String, Value>>> {
|
||||
self.catalog
|
||||
.lookup(model)
|
||||
.map(|found| Cow::Borrowed(found.entry.fields()))
|
||||
}
|
||||
|
||||
fn declares(&self, key: &str) -> bool {
|
||||
|
|
|
|||
|
|
@ -1029,7 +1029,7 @@ impl Hook for RateLimiter {
|
|||
let mut tokens = None;
|
||||
if let Some(min_tpm) = limited.iter().map(|(_, limit)| *limit).min() {
|
||||
let body: Map<String, Value> =
|
||||
serde_json::from_slice(request.body()).unwrap_or_default();
|
||||
serde_json::from_slice(&request.body()).unwrap_or_default();
|
||||
let estimate = estimate(
|
||||
&body,
|
||||
request.key.as_deref(),
|
||||
|
|
|
|||
|
|
@ -82,7 +82,7 @@ pub(super) fn standard_logging_payload(
|
|||
let outcome = &record.outcome;
|
||||
let deployment = outcome.deployment.as_ref();
|
||||
let succeeded = record.succeeded();
|
||||
let body: Map<String, Value> = serde_json::from_slice(request.body()).unwrap_or_default();
|
||||
let body: Map<String, Value> = serde_json::from_slice(&request.body()).unwrap_or_default();
|
||||
let response = match succeeded {
|
||||
true => response(record),
|
||||
false => Value::Object(Map::new()),
|
||||
|
|
|
|||
722
litellm-rust/crates/gateway/tests/chat_translation.rs
Normal file
722
litellm-rust/crates/gateway/tests/chat_translation.rs
Normal file
|
|
@ -0,0 +1,722 @@
|
|||
//! Chat completions for deployments the gateway translates (Anthropic, Bedrock Converse): the
|
||||
//! OpenAI request widened to tools, images and the provider's own limits, the provider body
|
||||
//! the upstream sees, and the OpenAI answer the client gets back, streamed and not.
|
||||
|
||||
use std::{
|
||||
io::Write,
|
||||
sync::{Arc, Mutex, Once},
|
||||
};
|
||||
|
||||
use axum::{
|
||||
Router,
|
||||
body::{Body, Bytes, to_bytes},
|
||||
http::{Request, header::CONTENT_TYPE},
|
||||
response::IntoResponse,
|
||||
};
|
||||
use litellm_config::Config;
|
||||
use litellm_framer::aws_event_stream::{AwsEventStreamCodec, Header, HeaderValue, Message};
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
use tokio::net::TcpListener;
|
||||
use tokio_util::codec::Encoder;
|
||||
use tower::ServiceExt;
|
||||
|
||||
const ANTHROPIC_MODEL: &str = "claude-x";
|
||||
const BEDROCK_MODEL: &str = "anthropic.claude-x-v1:0";
|
||||
/// A one-pixel PNG.
|
||||
const PNG: &str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==";
|
||||
|
||||
/// The gateway reads its cost map from this variable; every test sets it before building a
|
||||
/// gateway, so no test reads the environment while it changes.
|
||||
fn use_local_cost_map() {
|
||||
static ONCE: Once = Once::new();
|
||||
ONCE.call_once(|| {
|
||||
let capabilities = json!({
|
||||
"mode": "chat",
|
||||
"max_tokens": 8192,
|
||||
"max_output_tokens": 8192,
|
||||
"supports_tool_choice": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_vision": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_response_schema": true,
|
||||
});
|
||||
let entry = |provider: &str| {
|
||||
let mut entry = capabilities.clone();
|
||||
entry["litellm_provider"] = json!(provider);
|
||||
entry
|
||||
};
|
||||
let map = json!({
|
||||
ANTHROPIC_MODEL: entry("anthropic"),
|
||||
BEDROCK_MODEL: entry("bedrock_converse"),
|
||||
});
|
||||
let mut file = tempfile::NamedTempFile::new().unwrap();
|
||||
file.write_all(map.to_string().as_bytes()).unwrap();
|
||||
let (_, path) = file.keep().unwrap();
|
||||
// SAFETY: set once, before any test builds a gateway or spawns a thread that reads the
|
||||
// environment.
|
||||
unsafe { std::env::set_var("LITELLM_MODEL_COST_MAP_PATH", path) };
|
||||
});
|
||||
}
|
||||
|
||||
fn anthropic_events(events: &[Value]) -> String {
|
||||
events
|
||||
.iter()
|
||||
.map(|event| {
|
||||
format!(
|
||||
"event: {}\ndata: {event}\n\n",
|
||||
event["type"].as_str().unwrap()
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn anthropic_tool_message() -> Value {
|
||||
json!({
|
||||
"id": "msg_1", "type": "message", "role": "assistant", "model": ANTHROPIC_MODEL,
|
||||
"content": [
|
||||
{"type": "text", "text": "Checking."},
|
||||
{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Paris"}}
|
||||
],
|
||||
"stop_reason": "tool_use", "stop_sequence": null,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5}
|
||||
})
|
||||
}
|
||||
|
||||
/// The answer to a `response_format` request, given as the tool it became.
|
||||
fn anthropic_json_message() -> Value {
|
||||
json!({
|
||||
"id": "msg_3", "type": "message", "role": "assistant", "model": ANTHROPIC_MODEL,
|
||||
"content": [{"type": "tool_use", "id": "toolu_2", "name": "json_tool_call", "input": {"city": "Paris"}}],
|
||||
"stop_reason": "tool_use", "stop_sequence": null,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 5}
|
||||
})
|
||||
}
|
||||
|
||||
fn anthropic_text_message() -> Value {
|
||||
json!({
|
||||
"id": "msg_2", "type": "message", "role": "assistant", "model": ANTHROPIC_MODEL,
|
||||
"content": [{"type": "text", "text": "Sunny."}],
|
||||
"stop_reason": "end_turn", "stop_sequence": null,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 2}
|
||||
})
|
||||
}
|
||||
|
||||
fn anthropic_tool_stream() -> String {
|
||||
anthropic_events(&[
|
||||
json!({"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": ANTHROPIC_MODEL, "content": [], "stop_reason": null, "stop_sequence": null, "usage": {"input_tokens": 10, "output_tokens": 1}}}),
|
||||
json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {}}}),
|
||||
json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{\"city\": "}}),
|
||||
json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "\"Paris\"}"}}),
|
||||
json!({"type": "content_block_stop", "index": 0}),
|
||||
json!({"type": "message_delta", "delta": {"stop_reason": "tool_use", "stop_sequence": null}, "usage": {"output_tokens": 5}}),
|
||||
json!({"type": "message_stop"}),
|
||||
])
|
||||
}
|
||||
|
||||
fn converse_tool_response() -> Value {
|
||||
json!({
|
||||
"output": {"message": {"role": "assistant", "content": [
|
||||
{"text": "Checking."},
|
||||
{"toolUse": {"toolUseId": "tooluse_1", "name": "get_weather", "input": {"city": "Paris"}}}
|
||||
]}},
|
||||
"stopReason": "tool_use",
|
||||
"usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}
|
||||
})
|
||||
}
|
||||
|
||||
fn converse_text_response() -> Value {
|
||||
json!({
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": "Sunny."}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 10, "outputTokens": 2, "totalTokens": 12}
|
||||
})
|
||||
}
|
||||
|
||||
fn bedrock_frame(name: &str, payload: Value, wire: &mut bytes::BytesMut) {
|
||||
let message = Message::new(Bytes::from(payload.to_string()))
|
||||
.add_header(Header::new(
|
||||
":message-type",
|
||||
HeaderValue::String("event".to_owned().into()),
|
||||
))
|
||||
.add_header(Header::new(
|
||||
":event-type",
|
||||
HeaderValue::String(name.to_owned().into()),
|
||||
));
|
||||
AwsEventStreamCodec.encode(message, wire).unwrap();
|
||||
}
|
||||
|
||||
fn converse_tool_stream() -> Bytes {
|
||||
let mut wire = bytes::BytesMut::new();
|
||||
for (name, payload) in [
|
||||
("messageStart", json!({"role": "assistant"})),
|
||||
(
|
||||
"contentBlockStart",
|
||||
json!({"contentBlockIndex": 0, "start": {"toolUse": {"toolUseId": "tooluse_1", "name": "get_weather"}}}),
|
||||
),
|
||||
(
|
||||
"contentBlockDelta",
|
||||
json!({"contentBlockIndex": 0, "delta": {"toolUse": {"input": "{\"city\": "}}}),
|
||||
),
|
||||
(
|
||||
"contentBlockDelta",
|
||||
json!({"contentBlockIndex": 0, "delta": {"toolUse": {"input": "\"Paris\"}"}}}),
|
||||
),
|
||||
("contentBlockStop", json!({"contentBlockIndex": 0})),
|
||||
("messageStop", json!({"stopReason": "tool_use"})),
|
||||
(
|
||||
"metadata",
|
||||
json!({"usage": {"inputTokens": 10, "outputTokens": 5, "totalTokens": 15}, "metrics": {"latencyMs": 9}}),
|
||||
),
|
||||
] {
|
||||
bedrock_frame(name, payload, &mut wire);
|
||||
}
|
||||
wire.freeze()
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct Sent {
|
||||
path: String,
|
||||
headers: Vec<(String, String)>,
|
||||
body: Value,
|
||||
}
|
||||
|
||||
impl Sent {
|
||||
fn header(&self, name: &str) -> Option<&str> {
|
||||
self.headers
|
||||
.iter()
|
||||
.find(|(key, _)| key == name)
|
||||
.map(|(_, value)| value.as_str())
|
||||
}
|
||||
}
|
||||
|
||||
type Calls = Arc<Mutex<Vec<Sent>>>;
|
||||
|
||||
/// Answers each provider by path: a tool call while the conversation has no tool result yet,
|
||||
/// the final answer once it does.
|
||||
async fn upstream(calls: Calls) -> String {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let app = Router::new().fallback(move |request: Request<Body>| {
|
||||
let calls = calls.clone();
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let body: Value =
|
||||
serde_json::from_slice(&to_bytes(body, 1 << 22).await.unwrap()).unwrap();
|
||||
let path = parts.uri.path().to_owned();
|
||||
calls.lock().unwrap().push(Sent {
|
||||
path: path.clone(),
|
||||
headers: parts
|
||||
.headers
|
||||
.iter()
|
||||
.map(|(name, value)| {
|
||||
(
|
||||
name.to_string(),
|
||||
value.to_str().unwrap_or_default().to_owned(),
|
||||
)
|
||||
})
|
||||
.collect(),
|
||||
body: body.clone(),
|
||||
});
|
||||
let answered =
|
||||
body.to_string().contains("tool_result") || body.to_string().contains("toolResult");
|
||||
let json = [(CONTENT_TYPE, "application/json")];
|
||||
if path.ends_with("/converse-stream") {
|
||||
return (
|
||||
[(CONTENT_TYPE, "application/vnd.amazon.eventstream")],
|
||||
converse_tool_stream(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
if path.ends_with("/converse") {
|
||||
let answer = match answered {
|
||||
true => converse_text_response(),
|
||||
false => converse_tool_response(),
|
||||
};
|
||||
return (json, answer.to_string()).into_response();
|
||||
}
|
||||
if body["stream"] == json!(true) {
|
||||
return (
|
||||
[(CONTENT_TYPE, "text/event-stream")],
|
||||
anthropic_tool_stream(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
let answer = match answered {
|
||||
_ if body["tools"][0]["name"] == "json_tool_call" => anthropic_json_message(),
|
||||
true => anthropic_text_message(),
|
||||
false => anthropic_tool_message(),
|
||||
};
|
||||
(json, answer.to_string()).into_response()
|
||||
}
|
||||
});
|
||||
tokio::spawn(async move { axum::serve(listener, app).await });
|
||||
format!("http://{address}")
|
||||
}
|
||||
|
||||
struct Harness {
|
||||
app: Router,
|
||||
calls: Calls,
|
||||
}
|
||||
|
||||
async fn harness() -> Harness {
|
||||
use_local_cost_map();
|
||||
let calls = Calls::default();
|
||||
let upstream = upstream(calls.clone()).await;
|
||||
let config = Config::from_yaml(&format!(
|
||||
"model_list:
|
||||
- model_name: claude
|
||||
litellm_params:
|
||||
model: anthropic/{ANTHROPIC_MODEL}
|
||||
api_base: {upstream}
|
||||
api_key: sk-anthropic
|
||||
- model_name: bedrock
|
||||
litellm_params:
|
||||
model: bedrock/{BEDROCK_MODEL}
|
||||
api_base: {upstream}
|
||||
api_key: bedrock-token
|
||||
- model_name: claude-dropping
|
||||
litellm_params:
|
||||
model: anthropic/{ANTHROPIC_MODEL}
|
||||
api_base: {upstream}
|
||||
api_key: sk-anthropic
|
||||
drop_params: true
|
||||
router_settings:
|
||||
num_retries: 0
|
||||
general_settings:
|
||||
master_key: sk-master
|
||||
"
|
||||
))
|
||||
.unwrap();
|
||||
let inference = litellm_gateway::build_inference(&config).unwrap();
|
||||
litellm_gateway::hooks::install(&inference, &config, Vec::new(), Vec::new())
|
||||
.await
|
||||
.unwrap();
|
||||
let app = litellm_gateway::router(inference, &config, None, None);
|
||||
Harness { app, calls }
|
||||
}
|
||||
|
||||
impl Harness {
|
||||
async fn send(&self, body: Value) -> (u16, String) {
|
||||
let request = Request::post("/v1/chat/completions")
|
||||
.header("content-type", "application/json")
|
||||
.header("authorization", "Bearer sk-master")
|
||||
.body(Body::from(body.to_string()))
|
||||
.unwrap();
|
||||
let response = self.app.clone().oneshot(request).await.unwrap();
|
||||
let status = response.status().as_u16();
|
||||
let body = to_bytes(response.into_body(), 1 << 22).await.unwrap();
|
||||
(status, String::from_utf8(body.to_vec()).unwrap())
|
||||
}
|
||||
|
||||
async fn json(&self, body: Value) -> Value {
|
||||
let (status, text) = self.send(body).await;
|
||||
assert_eq!(status, 200, "{text}");
|
||||
serde_json::from_str(&text).unwrap()
|
||||
}
|
||||
|
||||
fn sent(&self) -> Vec<Sent> {
|
||||
self.calls.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn weather_tool() -> Value {
|
||||
json!({"type": "function", "function": {
|
||||
"name": "get_weather",
|
||||
"description": "The weather in a city",
|
||||
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}
|
||||
}})
|
||||
}
|
||||
|
||||
fn question() -> Value {
|
||||
json!({"role": "user", "content": "Weather in Paris?"})
|
||||
}
|
||||
|
||||
/// The conversation after the model asked for the tool and the tool answered.
|
||||
fn answered(call_id: &str) -> Value {
|
||||
json!([
|
||||
question(),
|
||||
{"role": "assistant", "content": null, "tool_calls": [
|
||||
{"id": call_id, "type": "function", "function": {"name": "get_weather", "arguments": "{\"city\": \"Paris\"}"}}
|
||||
]},
|
||||
{"role": "tool", "tool_call_id": call_id, "content": "sunny, 21C"}
|
||||
])
|
||||
}
|
||||
|
||||
/// The streamed chunks' JSON and whether the stream ended with `[DONE]`.
|
||||
fn frames(body: &str) -> (Vec<Value>, bool) {
|
||||
let mut done = false;
|
||||
let frames = body
|
||||
.split("\n\n")
|
||||
.filter_map(|frame| frame.strip_prefix("data: "))
|
||||
.filter(|data| {
|
||||
let last = *data == "[DONE]";
|
||||
done |= last;
|
||||
!last
|
||||
})
|
||||
.map(|data| serde_json::from_str(data).unwrap())
|
||||
.collect();
|
||||
(frames, done)
|
||||
}
|
||||
|
||||
/// The tool calls the chunks stream, joined: (id, name, arguments) per index.
|
||||
fn streamed_tool_calls(chunks: &[Value]) -> Vec<(String, String, String)> {
|
||||
let mut calls: Vec<(String, String, String)> = Vec::new();
|
||||
for chunk in chunks {
|
||||
let Some(deltas) = chunk["choices"][0]["delta"]["tool_calls"].as_array() else {
|
||||
continue;
|
||||
};
|
||||
for delta in deltas {
|
||||
let index = delta["index"].as_u64().unwrap() as usize;
|
||||
if calls.len() <= index {
|
||||
calls.resize(index + 1, Default::default());
|
||||
}
|
||||
let call = &mut calls[index];
|
||||
if let Some(id) = delta["id"].as_str() {
|
||||
call.0 = id.to_owned();
|
||||
}
|
||||
if let Some(name) = delta["function"]["name"].as_str() {
|
||||
call.1 = name.to_owned();
|
||||
}
|
||||
call.2
|
||||
.push_str(delta["function"]["arguments"].as_str().unwrap_or_default());
|
||||
}
|
||||
}
|
||||
calls
|
||||
}
|
||||
|
||||
fn finish_reason(chunks: &[Value]) -> Option<String> {
|
||||
chunks
|
||||
.iter()
|
||||
.find_map(|chunk| chunk["choices"][0]["finish_reason"].as_str())
|
||||
.map(str::to_owned)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn anthropic_tool_round_trip() {
|
||||
let harness = harness().await;
|
||||
|
||||
let asked = harness
|
||||
.json(json!({
|
||||
"model": "claude",
|
||||
"messages": [question()],
|
||||
"tools": [weather_tool()],
|
||||
"tool_choice": "required",
|
||||
"max_tokens": 64,
|
||||
}))
|
||||
.await;
|
||||
let message = &asked["choices"][0]["message"];
|
||||
assert_eq!(asked["choices"][0]["finish_reason"], "tool_calls");
|
||||
assert_eq!(message["content"], "Checking.");
|
||||
assert_eq!(
|
||||
message["tool_calls"],
|
||||
json!([{"id": "toolu_1", "type": "function", "index": 1,
|
||||
"function": {"name": "get_weather", "arguments": "{\"city\": \"Paris\"}"}}])
|
||||
);
|
||||
|
||||
let done = harness
|
||||
.json(json!({
|
||||
"model": "claude",
|
||||
"messages": answered("toolu_1"),
|
||||
"tools": [weather_tool()],
|
||||
"max_tokens": 64,
|
||||
}))
|
||||
.await;
|
||||
assert_eq!(done["choices"][0]["message"]["content"], "Sunny.");
|
||||
assert_eq!(done["choices"][0]["finish_reason"], "stop");
|
||||
|
||||
let sent = harness.sent();
|
||||
assert_eq!(sent[0].path, "/v1/messages");
|
||||
assert_eq!(
|
||||
sent[0].body,
|
||||
json!({
|
||||
"model": ANTHROPIC_MODEL,
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": "Weather in Paris?"}]}],
|
||||
"max_tokens": 64,
|
||||
"tools": [{"name": "get_weather", "type": "custom", "description": "The weather in a city",
|
||||
"input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}],
|
||||
"tool_choice": {"type": "any"}
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
sent[1].body["messages"],
|
||||
json!([
|
||||
{"role": "user", "content": [{"type": "text", "text": "Weather in Paris?"}]},
|
||||
{"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Paris"}}]},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "sunny, 21C"}]}
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn bedrock_tool_round_trip() {
|
||||
let harness = harness().await;
|
||||
|
||||
let asked = harness
|
||||
.json(json!({
|
||||
"model": "bedrock",
|
||||
"messages": [question()],
|
||||
"tools": [weather_tool()],
|
||||
"tool_choice": {"type": "function", "function": {"name": "get_weather"}},
|
||||
"max_tokens": 64,
|
||||
}))
|
||||
.await;
|
||||
assert_eq!(asked["choices"][0]["finish_reason"], "tool_calls");
|
||||
assert_eq!(
|
||||
asked["choices"][0]["message"]["tool_calls"],
|
||||
json!([{"id": "tooluse_1", "type": "function", "index": 1,
|
||||
"function": {"name": "get_weather", "arguments": "{\"city\": \"Paris\"}"}}])
|
||||
);
|
||||
|
||||
let done = harness
|
||||
.json(json!({
|
||||
"model": "bedrock",
|
||||
"messages": answered("tooluse_1"),
|
||||
"tools": [weather_tool()],
|
||||
"max_tokens": 64,
|
||||
}))
|
||||
.await;
|
||||
assert_eq!(done["choices"][0]["message"]["content"], "Sunny.");
|
||||
assert_eq!(done["choices"][0]["finish_reason"], "stop");
|
||||
|
||||
let sent = harness.sent();
|
||||
assert_eq!(sent[0].path, format!("/model/{BEDROCK_MODEL}/converse"));
|
||||
assert_eq!(
|
||||
sent[0].body["toolConfig"],
|
||||
json!({
|
||||
"tools": [{"toolSpec": {"name": "get_weather", "description": "The weather in a city",
|
||||
"inputSchema": {"json": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}}}}],
|
||||
"toolChoice": {"tool": {"name": "get_weather"}}
|
||||
})
|
||||
);
|
||||
assert_eq!(
|
||||
sent[1].body["messages"],
|
||||
json!([
|
||||
{"role": "user", "content": [{"text": "Weather in Paris?"}]},
|
||||
{"role": "assistant", "content": [{"toolUse": {"toolUseId": "tooluse_1", "name": "get_weather", "input": {"city": "Paris"}}}]},
|
||||
{"role": "user", "content": [{"toolResult": {"toolUseId": "tooluse_1", "content": [{"text": "sunny, 21C"}]}}]}
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::anthropic("claude", "toolu_1")]
|
||||
#[case::bedrock("bedrock", "tooluse_1")]
|
||||
#[tokio::test]
|
||||
async fn a_streamed_tool_call_arrives_as_openai_tool_call_chunks(
|
||||
#[case] model: &str,
|
||||
#[case] id: &str,
|
||||
) {
|
||||
let harness = harness().await;
|
||||
|
||||
let (status, body) = harness
|
||||
.send(json!({
|
||||
"model": model,
|
||||
"messages": [question()],
|
||||
"tools": [weather_tool()],
|
||||
"max_tokens": 64,
|
||||
"stream": true,
|
||||
}))
|
||||
.await;
|
||||
|
||||
assert_eq!(status, 200, "{body}");
|
||||
let (chunks, done) = frames(&body);
|
||||
assert!(done, "{body}");
|
||||
assert_eq!(
|
||||
streamed_tool_calls(&chunks),
|
||||
[(
|
||||
id.to_owned(),
|
||||
"get_weather".to_owned(),
|
||||
"{\"city\": \"Paris\"}".to_owned()
|
||||
)]
|
||||
);
|
||||
assert_eq!(finish_reason(&chunks).as_deref(), Some("tool_calls"));
|
||||
let sent = &harness.sent()[0];
|
||||
let tools = match model {
|
||||
"claude" => sent.body["tools"][0]["name"].clone(),
|
||||
_ => sent.body["toolConfig"]["tools"][0]["toolSpec"]["name"].clone(),
|
||||
};
|
||||
assert_eq!(tools, "get_weather");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::max_tokens("max_tokens")]
|
||||
#[case::max_completion_tokens("max_completion_tokens")]
|
||||
#[tokio::test]
|
||||
async fn bedrock_takes_the_token_limit_in_its_inference_config(#[case] param: &str) {
|
||||
let harness = harness().await;
|
||||
|
||||
harness
|
||||
.json(json!({
|
||||
"model": "bedrock",
|
||||
"messages": [{"role": "system", "content": "Be brief."}, question()],
|
||||
param: 32,
|
||||
"temperature": 0.5,
|
||||
"stop": ["END"],
|
||||
}))
|
||||
.await;
|
||||
|
||||
let sent = &harness.sent()[0];
|
||||
assert_eq!(
|
||||
sent.body["inferenceConfig"],
|
||||
json!({"maxTokens": 32, "temperature": 0.5, "stopSequences": ["END"]})
|
||||
);
|
||||
assert_eq!(sent.body["system"], json!([{"text": "Be brief."}]));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn anthropic_gets_the_model_default_max_tokens_when_the_caller_names_none() {
|
||||
let harness = harness().await;
|
||||
|
||||
harness
|
||||
.json(json!({"model": "claude", "messages": [question()]}))
|
||||
.await;
|
||||
|
||||
assert_eq!(harness.sent()[0].body["max_tokens"], 8192);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn an_image_data_url_reaches_each_provider_inline() {
|
||||
let harness = harness().await;
|
||||
let messages = json!([{"role": "user", "content": [
|
||||
{"type": "text", "text": "What is this?"},
|
||||
{"type": "image_url", "image_url": {"url": format!("data:image/png;base64,{PNG}")}}
|
||||
]}]);
|
||||
|
||||
harness
|
||||
.json(json!({"model": "claude", "messages": messages, "max_tokens": 16}))
|
||||
.await;
|
||||
harness
|
||||
.json(json!({"model": "bedrock", "messages": messages, "max_tokens": 16}))
|
||||
.await;
|
||||
|
||||
let sent = harness.sent();
|
||||
assert_eq!(
|
||||
sent[0].body["messages"][0]["content"][1],
|
||||
json!({"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": PNG}})
|
||||
);
|
||||
assert_eq!(
|
||||
sent[1].body["messages"][0]["content"][1],
|
||||
json!({"image": {"format": "png", "source": {"bytes": PNG}}})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn anthropic_response_format_comes_back_as_the_answer() {
|
||||
let harness = harness().await;
|
||||
let schema =
|
||||
json!({"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]});
|
||||
|
||||
let answer = harness
|
||||
.json(json!({
|
||||
"model": "claude",
|
||||
"messages": [question()],
|
||||
"max_tokens": 64,
|
||||
"response_format": {"type": "json_schema", "json_schema": {"name": "place", "schema": schema}},
|
||||
"user": "end-user-1",
|
||||
}))
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
answer["choices"][0]["message"]["content"],
|
||||
"{\"city\": \"Paris\"}"
|
||||
);
|
||||
assert_eq!(answer["choices"][0]["message"]["tool_calls"], Value::Null);
|
||||
assert_eq!(answer["choices"][0]["finish_reason"], "stop");
|
||||
let sent = &harness.sent()[0];
|
||||
assert_eq!(sent.body["metadata"], json!({"user_id": "end-user-1"}));
|
||||
assert_eq!(sent.body["tools"][0]["name"], "json_tool_call");
|
||||
assert_eq!(sent.body["tools"][0]["input_schema"], schema);
|
||||
assert_eq!(
|
||||
sent.body["tool_choice"],
|
||||
json!({"name": "json_tool_call", "type": "tool"})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::anthropic("claude", json!({"logprobs": true}))]
|
||||
#[case::bedrock("bedrock", json!({"tool_choice": "none", "tools": [weather_tool()]}))]
|
||||
#[case::n_choices("claude", json!({"n": 2}))]
|
||||
#[tokio::test]
|
||||
async fn an_unsupported_param_is_refused_and_dropped_with_drop_params(
|
||||
#[case] model: &str,
|
||||
#[case] extra: Value,
|
||||
) {
|
||||
let harness = harness().await;
|
||||
let mut body = json!({"model": model, "messages": [question()], "max_tokens": 16});
|
||||
body.as_object_mut()
|
||||
.unwrap()
|
||||
.extend(extra.as_object().unwrap().clone());
|
||||
|
||||
let (status, refused) = harness.send(body.clone()).await;
|
||||
assert_eq!(status, 400, "{refused}");
|
||||
assert!(
|
||||
harness.sent().is_empty(),
|
||||
"a refused request never goes out"
|
||||
);
|
||||
|
||||
body["drop_params"] = json!(true);
|
||||
harness.json(body).await;
|
||||
let sent = &harness.sent()[0];
|
||||
for name in extra.as_object().unwrap().keys() {
|
||||
if name == "tools" {
|
||||
continue;
|
||||
}
|
||||
assert!(
|
||||
sent.body.get(name).is_none(),
|
||||
"{name} reached {}",
|
||||
sent.body
|
||||
);
|
||||
}
|
||||
assert!(
|
||||
sent.body
|
||||
.get("toolConfig")
|
||||
.and_then(|config| config.get("toolChoice"))
|
||||
.is_none()
|
||||
);
|
||||
assert!(sent.body.get("drop_params").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_deployment_that_drops_params_drops_them_without_the_request_asking() {
|
||||
let harness = harness().await;
|
||||
|
||||
harness
|
||||
.json(json!({
|
||||
"model": "claude-dropping",
|
||||
"messages": [question()],
|
||||
"max_tokens": 16,
|
||||
"logprobs": true,
|
||||
}))
|
||||
.await;
|
||||
|
||||
assert!(harness.sent()[0].body.get("logprobs").is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_cache_control_part_reaches_anthropic_and_bedrock() {
|
||||
let harness = harness().await;
|
||||
let messages = json!([
|
||||
{"role": "system", "content": [{"type": "text", "text": "Long instructions.", "cache_control": {"type": "ephemeral"}}]},
|
||||
question()
|
||||
]);
|
||||
|
||||
harness
|
||||
.json(json!({"model": "claude", "messages": messages, "max_tokens": 16}))
|
||||
.await;
|
||||
harness
|
||||
.json(json!({"model": "bedrock", "messages": messages, "max_tokens": 16}))
|
||||
.await;
|
||||
|
||||
let sent = harness.sent();
|
||||
assert_eq!(
|
||||
sent[0].body["system"],
|
||||
json!([{"type": "text", "text": "Long instructions.", "cache_control": {"type": "ephemeral"}}])
|
||||
);
|
||||
assert_eq!(
|
||||
sent[1].body["system"],
|
||||
json!([{"text": "Long instructions."}, {"cachePoint": {"type": "default"}}])
|
||||
);
|
||||
assert!(sent[0].header("x-api-key").is_some());
|
||||
}
|
||||
559
litellm-rust/crates/gateway/tests/gemini_chat.rs
Normal file
559
litellm-rust/crates/gateway/tests/gemini_chat.rs
Normal file
|
|
@ -0,0 +1,559 @@
|
|||
//! Gemini chat completions through the gateway, on Google AI Studio (`gemini/<model>`, an API
|
||||
//! key in `x-goog-api-key`) and on Vertex AI (`vertex_ai/<model>`, a service account traded for
|
||||
//! an access token): the request Gemini gets, the OpenAI answer the client gets, and the usage
|
||||
//! and cost the request is recorded with.
|
||||
|
||||
use std::{
|
||||
io::Write,
|
||||
sync::{Arc, Mutex, Once},
|
||||
};
|
||||
|
||||
use axum::{
|
||||
Router,
|
||||
body::{Body, to_bytes},
|
||||
http::{Request, StatusCode, header::CONTENT_TYPE},
|
||||
response::IntoResponse,
|
||||
};
|
||||
use futures_util::future::BoxFuture;
|
||||
use litellm_config::Config;
|
||||
use litellm_gateway_hooks::{Facts, Observer, Record};
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
use tokio::net::TcpListener;
|
||||
use tower::ServiceExt;
|
||||
|
||||
/// Input and output prices of every deployment.
|
||||
const RATES: (f64, f64) = (1.0, 2.0);
|
||||
|
||||
const STUDIO_KEY: &str = "studio-key";
|
||||
const STUDIO_KEY_ENV: &str = "GEMINI_CHAT_TEST_STUDIO_KEY";
|
||||
const ACCESS_TOKEN: &str = "ya29.test-token";
|
||||
const PROJECT: &str = "test-project";
|
||||
const LOCATION: &str = "europe-west4";
|
||||
|
||||
/// A throwaway RSA key the fake service account signs its token request with.
|
||||
const TEST_PRIVATE_KEY: &str = "-----BEGIN PRIVATE KEY-----
|
||||
MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQC95ihJQlfJuaQe
|
||||
ZCwF+EFy/Y5vk7Lx6Nm74He9Oes+6ZbgA2LVdcoFJrNEYHd8S0CCv7cqEp4dpIoR
|
||||
B4HAmQAQF92ZC/S2lcgLAezWQU21TSjDLngdVbn+Jr4syzOwVw/NUigydtQ66Iqv
|
||||
1FpXs/4gQXMgs8FW5XgOhgfKhyMPf71A2qiRaA0vwsrfsmKOY0QVVooXPRfpLIvZ
|
||||
VEOV89ZMQErnYan+m8OADoZ5NJ/Ij6QKb2HkQJNhQC2Zq2T1Y9VLyxY6pXURdVqc
|
||||
XJgFT6YPGjLABLX/oId4g+Fs2YI5mM+g72pYEntnIaE5yicGKBdjyqeVfKZ6/7jg
|
||||
g8evHi6vAgMBAAECggEAGdJjlP6b8Fa5bdaCM/ebcrbuuNZVJVbb0JPHxGfNSLs7
|
||||
pE9hj5QaOdQW2Uviw3h6F61ZCzQH4xD+Iy2po5ZKb2XHYKnDB1bboj+LRGER337T
|
||||
9aJqe9at2VTMVEv3Rdm40NsEk0QcPLxlK16NQFK90gYEUSSQPDAswJDSG2R/zHn+
|
||||
vADI907mW/goEJHeLn8PWGlNlSiR6x+5JJtq+GXCzUzVvJYQSCLGxCSl2x2H+0g7
|
||||
NhFI0zPpdzNmO/h+yhzaFb6Rp5U8+ZsnZ3qYjQ/03gw1myTDKJt1YaO9JvArnNYX
|
||||
hcJQQ8Rt0bHhcrZA16bBOpqZlo5pKCicwI/netgN8QKBgQDcFz7AzdJ26sMSV32V
|
||||
rwrMgIoggt8qDjO1ARwqW35A1TIge0FoW4M4KpsXQGGfT341uU1esXEcyZ/1L/5X
|
||||
3ql2gX4DbOYLZLWYzZGR2hq33oi8HkhN98QrEwL9emSH8NqYX3Xxja3PrmCrSYJe
|
||||
Zbnd9TIm2XkxyMoyXJu6M/QvnwKBgQDc4dzqTbxoGEGa5MuJoGmMwPnqgdG9UM5J
|
||||
eExVnh7osxc2sOdsiPeRjjQTxs9v2kJwctC359OJoo9yGaaJeSghU4LEWJo1sqnA
|
||||
fzSCLammYvtVAtniyNv5Mxk/6Uimi4NNDKaAKB+m4K2uSn3U9AmY7KPYMGaSbS9W
|
||||
XSnobjxm8QKBgC8bPpAvvWs8ZhIn7bY659nLbUT2HeO3dHO6UBf0yzn/J6JyHxbB
|
||||
93zvCZDZc8uQTRgcmCW7XtVlhjoJUqvl+Wlm39zF0xr/LCsPXKfWAb/2/lcdOCaP
|
||||
8Emz4QD10EyUTYUtcWYJB/mafhBLRH8F0Nlj4J8WDu2L51MOJTqeYhZLAoGAWffN
|
||||
icocAbJPlo22sdoa4+/+W5yBF8GAJMDRJtZ+9H1t6SLpQHYRkMIBSETkXUTjZvX9
|
||||
Ocs9iIQkNW9pO/mTdO+VBfCo71JUfknR02xR+6m5gYjlws/ZeYlssXGN2/hbhNiw
|
||||
QOcW7Vv6olFJK6Iy/oz0t6wPO3kpnN3Zogi0paECgYEAwo44M1DdYCtV0snhmYM9
|
||||
5u0mPfYt5P2SVLXyUbr+vFTfrTL/WKnXIJgbsnj3Gvf+GIZv9tKcXhSNmEHQCYX4
|
||||
X3w9iTPddCHuvZ1fpufi2TyArJh0OkoNtLXJHTKrHjf2N+61AQzFiv5WieJrdE+H
|
||||
qr32PTUuVGPyO9LyTY4/RL0=
|
||||
-----END PRIVATE KEY-----
|
||||
";
|
||||
|
||||
/// The gateway reads its cost map from this variable; every test sets it before building a
|
||||
/// gateway, so no test reads the environment while it changes.
|
||||
fn use_local_cost_map() {
|
||||
static ONCE: Once = Once::new();
|
||||
ONCE.call_once(|| {
|
||||
let entry = |provider: &str| {
|
||||
json!({
|
||||
"mode": "chat",
|
||||
"litellm_provider": provider,
|
||||
"max_tokens": 8192,
|
||||
"supports_system_messages": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_reasoning": true,
|
||||
})
|
||||
};
|
||||
let map = json!({
|
||||
"gemini/gemini-2.5-flash": entry("gemini"),
|
||||
"gemini-2.5-flash": entry("vertex_ai-language-models"),
|
||||
});
|
||||
let mut file = tempfile::NamedTempFile::new().unwrap();
|
||||
file.write_all(map.to_string().as_bytes()).unwrap();
|
||||
let (_, path) = file.keep().unwrap();
|
||||
// SAFETY: set once, before any test builds a gateway or spawns a thread that reads the
|
||||
// environment.
|
||||
unsafe {
|
||||
std::env::set_var("LITELLM_MODEL_COST_MAP_PATH", path);
|
||||
std::env::set_var(STUDIO_KEY_ENV, STUDIO_KEY);
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
fn usage_metadata() -> Value {
|
||||
json!({"promptTokenCount": 10, "candidatesTokenCount": 5, "thoughtsTokenCount": 3, "totalTokenCount": 18})
|
||||
}
|
||||
|
||||
fn text_answer(text: &str) -> Value {
|
||||
json!({
|
||||
"candidates": [{"content": {"role": "model", "parts": [{"text": text}]}, "finishReason": "STOP", "index": 0}],
|
||||
"usageMetadata": usage_metadata(),
|
||||
"modelVersion": "gemini-2.5-flash",
|
||||
"responseId": "resp-1",
|
||||
})
|
||||
}
|
||||
|
||||
fn tool_answer() -> Value {
|
||||
json!({
|
||||
"candidates": [{"content": {"role": "model", "parts": [{"functionCall": {"name": "get_weather", "args": {"city": "Paris"}}}]}, "finishReason": "STOP", "index": 0}],
|
||||
"usageMetadata": usage_metadata(),
|
||||
"modelVersion": "gemini-2.5-flash",
|
||||
})
|
||||
}
|
||||
|
||||
fn sse_answer() -> String {
|
||||
[
|
||||
json!({"candidates": [{"content": {"role": "model", "parts": [{"text": "Hello"}]}}], "modelVersion": "gemini-2.5-flash", "responseId": "resp-2"}),
|
||||
json!({"candidates": [{"content": {"role": "model", "parts": [{"text": " world"}]}, "finishReason": "STOP"}], "usageMetadata": usage_metadata(), "modelVersion": "gemini-2.5-flash", "responseId": "resp-2"}),
|
||||
]
|
||||
.iter()
|
||||
.map(|event| format!("data: {event}\r\n\r\n"))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct Sent {
|
||||
path: String,
|
||||
query: Option<String>,
|
||||
api_key: Option<String>,
|
||||
authorization: Option<String>,
|
||||
body: Value,
|
||||
}
|
||||
|
||||
type Calls = Arc<Mutex<Vec<Sent>>>;
|
||||
|
||||
fn has_function_response(body: &Value) -> bool {
|
||||
body["contents"].as_array().is_some_and(|contents| {
|
||||
contents.iter().any(|content| {
|
||||
content["parts"].as_array().is_some_and(|parts| {
|
||||
parts
|
||||
.iter()
|
||||
.any(|part| part.get("function_response").is_some())
|
||||
})
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
/// Google's token endpoint and Gemini's `generateContent`, answering by what the body asks.
|
||||
async fn upstream(calls: Calls) -> String {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let app = Router::new().fallback(move |request: Request<Body>| {
|
||||
let calls = calls.clone();
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let header = |name: &str| {
|
||||
parts
|
||||
.headers
|
||||
.get(name)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_owned)
|
||||
};
|
||||
let bytes = to_bytes(body, 1 << 20).await.unwrap();
|
||||
let path = parts.uri.path().to_owned();
|
||||
if path == "/token" {
|
||||
return (
|
||||
[(CONTENT_TYPE, "application/json")],
|
||||
json!({"access_token": ACCESS_TOKEN, "expires_in": 3600, "token_type": "Bearer"})
|
||||
.to_string(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
let body: Value = serde_json::from_slice(&bytes).unwrap_or(Value::Null);
|
||||
calls.lock().unwrap().push(Sent {
|
||||
path: path.clone(),
|
||||
query: parts.uri.query().map(str::to_owned),
|
||||
api_key: header("x-goog-api-key"),
|
||||
authorization: header("authorization"),
|
||||
body: body.clone(),
|
||||
});
|
||||
match path.ends_with(":streamGenerateContent") {
|
||||
true => ([(CONTENT_TYPE, "text/event-stream")], sse_answer()).into_response(),
|
||||
false => {
|
||||
let answer = match (has_function_response(&body), body.get("tools")) {
|
||||
(true, _) => text_answer("It is sunny in Paris"),
|
||||
(false, Some(_)) => tool_answer(),
|
||||
(false, None) => text_answer("Hello world"),
|
||||
};
|
||||
(
|
||||
StatusCode::OK,
|
||||
[(CONTENT_TYPE, "application/json")],
|
||||
answer.to_string(),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
tokio::spawn(async move { axum::serve(listener, app).await });
|
||||
format!("http://{address}")
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Seen {
|
||||
deployment: Option<String>,
|
||||
usage: Option<litellm_gateway_hooks::Usage>,
|
||||
cost: Option<f64>,
|
||||
}
|
||||
|
||||
struct Collect(Mutex<Vec<Seen>>);
|
||||
|
||||
impl Observer for Collect {
|
||||
fn name(&self) -> &'static str {
|
||||
"collect"
|
||||
}
|
||||
|
||||
fn reads(&self) -> Facts {
|
||||
Facts::DEPLOYMENT | Facts::USAGE | Facts::COST
|
||||
}
|
||||
|
||||
fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
self.0.lock().unwrap().push(Seen {
|
||||
deployment: record
|
||||
.outcome
|
||||
.deployment
|
||||
.as_ref()
|
||||
.map(|deployment| deployment.id.to_string()),
|
||||
usage: record.outcome.usage,
|
||||
cost: record.outcome.cost,
|
||||
});
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct Harness {
|
||||
app: Router,
|
||||
inference: Arc<litellm_gateway_inference::Gateway>,
|
||||
calls: Calls,
|
||||
seen: Arc<Collect>,
|
||||
}
|
||||
|
||||
fn service_account(upstream: &str) -> String {
|
||||
json!({
|
||||
"type": "service_account",
|
||||
"project_id": "credentials-project",
|
||||
"private_key_id": "test-key-id",
|
||||
"private_key": TEST_PRIVATE_KEY,
|
||||
"client_email": "gateway-test@credentials-project.iam.gserviceaccount.com",
|
||||
"client_id": "1",
|
||||
"token_uri": format!("{upstream}/token"),
|
||||
})
|
||||
.to_string()
|
||||
}
|
||||
|
||||
async fn harness() -> Harness {
|
||||
use_local_cost_map();
|
||||
let calls = Calls::default();
|
||||
let upstream = upstream(calls.clone()).await;
|
||||
let (input, output) = RATES;
|
||||
let prices =
|
||||
format!(" input_cost_per_token: {input}\n output_cost_per_token: {output}\n");
|
||||
let credentials = service_account(&upstream);
|
||||
let config = Config::from_yaml(&format!(
|
||||
"model_list:
|
||||
- model_name: studio
|
||||
litellm_params:
|
||||
model: gemini/gemini-2.5-flash
|
||||
api_base: {upstream}/v1beta
|
||||
api_key: {STUDIO_KEY}
|
||||
{prices} model_info:
|
||||
id: studio-1
|
||||
- model_name: studio-env
|
||||
litellm_params:
|
||||
model: gemini/gemini-2.5-flash
|
||||
api_base: {upstream}/v1beta
|
||||
api_key: os.environ/{STUDIO_KEY_ENV}
|
||||
{prices} model_info:
|
||||
id: studio-env-1
|
||||
- model_name: vertex
|
||||
litellm_params:
|
||||
model: vertex_ai/gemini-2.5-flash
|
||||
api_base: {upstream}
|
||||
vertex_project: {PROJECT}
|
||||
vertex_location: {LOCATION}
|
||||
vertex_credentials: '{credentials}'
|
||||
{prices} model_info:
|
||||
id: vertex-1
|
||||
general_settings:
|
||||
master_key: sk-master
|
||||
"
|
||||
))
|
||||
.unwrap();
|
||||
let inference = litellm_gateway::build_inference(&config).unwrap();
|
||||
let seen = Arc::new(Collect(Mutex::new(Vec::new())));
|
||||
litellm_gateway::hooks::install(&inference, &config, Vec::new(), vec![seen.clone()])
|
||||
.await
|
||||
.unwrap();
|
||||
let app = litellm_gateway::router(inference.clone(), &config, None, None);
|
||||
Harness {
|
||||
app,
|
||||
inference,
|
||||
calls,
|
||||
seen,
|
||||
}
|
||||
}
|
||||
|
||||
async fn send(app: &Router, body: Value) -> (u16, String) {
|
||||
let request = Request::post("/v1/chat/completions")
|
||||
.header("content-type", "application/json")
|
||||
.header("authorization", "Bearer sk-master")
|
||||
.body(Body::from(body.to_string()))
|
||||
.unwrap();
|
||||
let response = app.clone().oneshot(request).await.unwrap();
|
||||
let status = response.status().as_u16();
|
||||
let body = to_bytes(response.into_body(), 1 << 20).await.unwrap();
|
||||
(status, String::from_utf8(body.to_vec()).unwrap())
|
||||
}
|
||||
|
||||
async fn seen(harness: &Harness) -> Vec<Seen> {
|
||||
litellm_gateway::hooks::close(&harness.inference).await;
|
||||
std::mem::take(&mut *harness.seen.0.lock().unwrap())
|
||||
}
|
||||
|
||||
fn expected_cost() -> f64 {
|
||||
let (input, output) = RATES;
|
||||
10.0 * input + 8.0 * output
|
||||
}
|
||||
|
||||
fn hi(model: &str) -> Value {
|
||||
json!({
|
||||
"model": model,
|
||||
"messages": [
|
||||
{"role": "system", "content": "Be brief"},
|
||||
{"role": "user", "content": "hi"}
|
||||
],
|
||||
"temperature": 0.2,
|
||||
"max_tokens": 64,
|
||||
})
|
||||
}
|
||||
|
||||
fn sent_body() -> Value {
|
||||
json!({
|
||||
"contents": [{"role": "user", "parts": [{"text": "hi"}]}],
|
||||
"system_instruction": {"parts": [{"text": "Be brief"}]},
|
||||
"generationConfig": {"temperature": 0.2, "max_output_tokens": 64},
|
||||
})
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::studio(
|
||||
"studio",
|
||||
"studio-1",
|
||||
"/v1beta/models/gemini-2.5-flash:generateContent",
|
||||
Some(STUDIO_KEY),
|
||||
None
|
||||
)]
|
||||
#[case::studio_key_from_env(
|
||||
"studio-env",
|
||||
"studio-env-1",
|
||||
"/v1beta/models/gemini-2.5-flash:generateContent",
|
||||
Some(STUDIO_KEY),
|
||||
None
|
||||
)]
|
||||
#[case::vertex(
|
||||
"vertex",
|
||||
"vertex-1",
|
||||
"/v1/projects/test-project/locations/europe-west4/publishers/google/models/gemini-2.5-flash:generateContent",
|
||||
None,
|
||||
Some("Bearer ya29.test-token")
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn a_chat_completion_calls_generate_content_and_records_its_cost(
|
||||
#[case] model: &str,
|
||||
#[case] id: &str,
|
||||
#[case] path: &str,
|
||||
#[case] api_key: Option<&str>,
|
||||
#[case] authorization: Option<&str>,
|
||||
) {
|
||||
let harness = harness().await;
|
||||
let (status, body) = send(&harness.app, hi(model)).await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
let answer: Value = serde_json::from_str(&body).unwrap();
|
||||
assert_eq!(answer["choices"][0]["message"]["content"], "Hello world");
|
||||
assert_eq!(answer["choices"][0]["finish_reason"], "stop");
|
||||
assert_eq!(answer["usage"]["prompt_tokens"], 10);
|
||||
assert_eq!(answer["usage"]["completion_tokens"], 8);
|
||||
assert_eq!(
|
||||
answer["usage"]["completion_tokens_details"]["reasoning_tokens"],
|
||||
3
|
||||
);
|
||||
|
||||
let sent = harness.calls.lock().unwrap()[0].clone();
|
||||
assert_eq!(sent.path, path);
|
||||
assert_eq!(sent.query, None);
|
||||
assert_eq!(sent.api_key.as_deref(), api_key);
|
||||
assert_eq!(sent.authorization.as_deref(), authorization);
|
||||
assert_eq!(sent.body, sent_body());
|
||||
|
||||
let seen = seen(&harness).await;
|
||||
assert_eq!(seen.len(), 1, "{seen:?}");
|
||||
assert_eq!(seen[0].deployment.as_deref(), Some(id));
|
||||
let usage = seen[0].usage.unwrap();
|
||||
assert_eq!((usage.prompt_tokens(), usage.output_tokens), (10, 8));
|
||||
assert_eq!(seen[0].cost, Some(expected_cost()));
|
||||
}
|
||||
|
||||
/// The JSON of each `data:` frame, and whether the stream closed with `[DONE]`.
|
||||
fn frames(body: &str) -> (Vec<Value>, bool) {
|
||||
let mut done = false;
|
||||
let frames = body
|
||||
.split("\n\n")
|
||||
.filter(|frame| !frame.is_empty())
|
||||
.filter_map(|frame| {
|
||||
let data = frame.strip_prefix("data: ").expect("bare data frames");
|
||||
match data {
|
||||
"[DONE]" => {
|
||||
done = true;
|
||||
None
|
||||
}
|
||||
data => Some(serde_json::from_str(data).unwrap()),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
(frames, done)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::studio(
|
||||
"studio",
|
||||
"studio-1",
|
||||
"/v1beta/models/gemini-2.5-flash:streamGenerateContent"
|
||||
)]
|
||||
#[case::vertex(
|
||||
"vertex",
|
||||
"vertex-1",
|
||||
"/v1/projects/test-project/locations/europe-west4/publishers/google/models/gemini-2.5-flash:streamGenerateContent"
|
||||
)]
|
||||
#[tokio::test]
|
||||
async fn a_stream_calls_stream_generate_content_and_records_its_usage(
|
||||
#[case] model: &str,
|
||||
#[case] id: &str,
|
||||
#[case] path: &str,
|
||||
) {
|
||||
let harness = harness().await;
|
||||
let mut request = hi(model);
|
||||
request["stream"] = json!(true);
|
||||
request["stream_options"] = json!({"include_usage": true});
|
||||
let (status, body) = send(&harness.app, request).await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
|
||||
let (frames, done) = frames(&body);
|
||||
assert!(done, "{body}");
|
||||
assert!(
|
||||
frames
|
||||
.iter()
|
||||
.all(|frame| frame["object"] == "chat.completion.chunk")
|
||||
);
|
||||
let content: String = frames
|
||||
.iter()
|
||||
.filter_map(|frame| frame["choices"][0]["delta"]["content"].as_str())
|
||||
.collect();
|
||||
assert_eq!(content, "Hello world", "{body}");
|
||||
let finishes: Vec<&Value> = frames
|
||||
.iter()
|
||||
.filter_map(|frame| frame["choices"][0].get("finish_reason"))
|
||||
.filter(|reason| !reason.is_null())
|
||||
.collect();
|
||||
assert_eq!(finishes, [&json!("stop")], "{body}");
|
||||
let usages: Vec<&Value> = frames
|
||||
.iter()
|
||||
.filter_map(|frame| frame.get("usage"))
|
||||
.filter(|usage| !usage.is_null())
|
||||
.collect();
|
||||
assert_eq!(usages.len(), 1, "{body}");
|
||||
assert_eq!(usages[0]["prompt_tokens"], 10);
|
||||
assert_eq!(usages[0]["completion_tokens"], 8);
|
||||
|
||||
let sent = harness.calls.lock().unwrap()[0].clone();
|
||||
assert_eq!(sent.path, path);
|
||||
assert_eq!(sent.query.as_deref(), Some("alt=sse"));
|
||||
assert_eq!(sent.body, sent_body());
|
||||
|
||||
let seen = seen(&harness).await;
|
||||
assert_eq!(seen.len(), 1, "{seen:?}");
|
||||
assert_eq!(seen[0].deployment.as_deref(), Some(id));
|
||||
assert_eq!(seen[0].cost, Some(expected_cost()));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_tool_call_round_trips_through_function_call_and_response() {
|
||||
let harness = harness().await;
|
||||
let tools = json!([{"type": "function", "function": {
|
||||
"name": "get_weather",
|
||||
"description": "Weather for a city",
|
||||
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"], "additionalProperties": false}
|
||||
}}]);
|
||||
let question = json!({"role": "user", "content": "Weather in Paris?"});
|
||||
let (status, body) = send(
|
||||
&harness.app,
|
||||
json!({"model": "studio", "messages": [question], "tools": tools, "tool_choice": "auto"}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
let answer: Value = serde_json::from_str(&body).unwrap();
|
||||
let message = &answer["choices"][0]["message"];
|
||||
assert_eq!(answer["choices"][0]["finish_reason"], "tool_calls");
|
||||
let call = &message["tool_calls"][0];
|
||||
assert_eq!(call["type"], "function");
|
||||
assert_eq!(call["function"]["name"], "get_weather");
|
||||
assert_eq!(call["function"]["arguments"], "{\"city\": \"Paris\"}");
|
||||
|
||||
let first = harness.calls.lock().unwrap()[0].body.clone();
|
||||
assert_eq!(
|
||||
first["tools"],
|
||||
json!([{"function_declarations": [{
|
||||
"name": "get_weather",
|
||||
"description": "Weather for a city",
|
||||
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}
|
||||
}]}])
|
||||
);
|
||||
assert_eq!(
|
||||
first["toolConfig"],
|
||||
json!({"functionCallingConfig": {"mode": "AUTO"}})
|
||||
);
|
||||
|
||||
let (status, body) = send(
|
||||
&harness.app,
|
||||
json!({"model": "studio", "tools": tools, "messages": [
|
||||
question,
|
||||
{"role": "assistant", "content": null, "tool_calls": [call]},
|
||||
{"role": "tool", "tool_call_id": call["id"], "content": "sunny"}
|
||||
]}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
let answer: Value = serde_json::from_str(&body).unwrap();
|
||||
assert_eq!(
|
||||
answer["choices"][0]["message"]["content"],
|
||||
"It is sunny in Paris"
|
||||
);
|
||||
let second = harness.calls.lock().unwrap()[1].body.clone();
|
||||
assert_eq!(
|
||||
second["contents"],
|
||||
json!([
|
||||
{"role": "user", "parts": [{"text": "Weather in Paris?"}]},
|
||||
{"role": "model", "parts": [{"function_call": {"name": "get_weather", "args": {"city": "Paris"}}}]},
|
||||
{"role": "user", "parts": [{"function_response": {"name": "get_weather", "response": {"content": "sunny"}}}]}
|
||||
])
|
||||
);
|
||||
|
||||
let seen = seen(&harness).await;
|
||||
assert_eq!(seen.len(), 2, "{seen:?}");
|
||||
assert!(seen.iter().all(|seen| seen.cost == Some(expected_cost())));
|
||||
}
|
||||
442
litellm-rust/crates/gateway/tests/openai_compat_chat.rs
Normal file
442
litellm-rust/crates/gateway/tests/openai_compat_chat.rs
Normal file
|
|
@ -0,0 +1,442 @@
|
|||
//! Chat completions for the OpenAI-compatible providers through the gateway: the URL, auth
|
||||
//! and exact body each provider's upstream sees, and the usage and cost a stream is recorded
|
||||
//! with whether or not the client asked to see the usage.
|
||||
|
||||
use std::{
|
||||
collections::BTreeMap,
|
||||
io::Write,
|
||||
sync::{Arc, Mutex, Once},
|
||||
};
|
||||
|
||||
use axum::{
|
||||
Router,
|
||||
body::{Body, to_bytes},
|
||||
http::{Request, header::CONTENT_TYPE},
|
||||
response::IntoResponse,
|
||||
};
|
||||
use futures_util::future::BoxFuture;
|
||||
use litellm_config::Config;
|
||||
use litellm_gateway_hooks::{Facts, Observer, Record};
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
use tokio::net::TcpListener;
|
||||
use tower::ServiceExt;
|
||||
|
||||
const COMPLETION: &str = r#"{"id":"chatcmpl-1","object":"chat.completion","created":1,"model":"m","choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":5,"completion_tokens":6,"total_tokens":11}}"#;
|
||||
const CONTENT_CHUNK: &str = "data: {\"id\":\"c\",\"object\":\"chat.completion.chunk\",\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\n";
|
||||
const USAGE_CHUNK: &str = "data: {\"id\":\"c\",\"object\":\"chat.completion.chunk\",\"model\":\"m\",\"choices\":[],\"usage\":{\"prompt_tokens\":7,\"completion_tokens\":3,\"total_tokens\":10}}\n\n";
|
||||
/// xAI counts reasoning apart from `completion_tokens` but inside `total_tokens`.
|
||||
const XAI_USAGE_CHUNK: &str = "data: {\"id\":\"c\",\"object\":\"chat.completion.chunk\",\"model\":\"m\",\"choices\":[],\"usage\":{\"prompt_tokens\":7,\"completion_tokens\":3,\"total_tokens\":14,\"completion_tokens_details\":{\"reasoning_tokens\":4}}}\n\n";
|
||||
/// Mistral ends its stream with the usage on the last content chunk, unasked.
|
||||
const MISTRAL_LAST_CHUNK: &str = "data: {\"id\":\"c\",\"object\":\"chat.completion.chunk\",\"model\":\"m\",\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":7,\"completion_tokens\":3,\"total_tokens\":10}}\n\n";
|
||||
const DONE: &str = "data: [DONE]\n\n";
|
||||
|
||||
const RATES: (f64, f64) = (1.0, 2.0);
|
||||
|
||||
fn use_local_cost_map() {
|
||||
static ONCE: Once = Once::new();
|
||||
ONCE.call_once(|| {
|
||||
let price = json!({
|
||||
"mode": "chat",
|
||||
"input_cost_per_token": RATES.0,
|
||||
"output_cost_per_token": RATES.1,
|
||||
});
|
||||
let map: serde_json::Map<String, Value> = [
|
||||
("groq/llama-3.1-8b-instant", "groq"),
|
||||
("mistral/mistral-small-latest", "mistral"),
|
||||
("xai/grok-3", "xai"),
|
||||
(
|
||||
"hosted_vllm/meta-llama/Llama-3.1-8B-Instruct",
|
||||
"hosted_vllm",
|
||||
),
|
||||
]
|
||||
.into_iter()
|
||||
.map(|(model, provider)| {
|
||||
let mut entry = price.clone();
|
||||
entry["litellm_provider"] = json!(provider);
|
||||
(model.to_owned(), entry)
|
||||
})
|
||||
.collect();
|
||||
let mut file = tempfile::NamedTempFile::new().unwrap();
|
||||
file.write_all(Value::Object(map).to_string().as_bytes())
|
||||
.unwrap();
|
||||
let (_, path) = file.keep().unwrap();
|
||||
// SAFETY: set once, before any test builds a gateway or spawns a thread that reads the
|
||||
// environment.
|
||||
unsafe { std::env::set_var("LITELLM_MODEL_COST_MAP_PATH", path) };
|
||||
});
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct Sent {
|
||||
path: String,
|
||||
headers: BTreeMap<String, String>,
|
||||
body: Value,
|
||||
}
|
||||
|
||||
type Calls = Arc<Mutex<Vec<Sent>>>;
|
||||
|
||||
fn stream_answer(path: &str) -> String {
|
||||
match path.split('/').nth(1) {
|
||||
Some("xai") => [CONTENT_CHUNK, XAI_USAGE_CHUNK, DONE].concat(),
|
||||
Some("mistral") => [CONTENT_CHUNK, MISTRAL_LAST_CHUNK, DONE].concat(),
|
||||
Some("vllm") => [
|
||||
"data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"reasoning\":\"hmm\"}}]}\n\n",
|
||||
CONTENT_CHUNK,
|
||||
USAGE_CHUNK,
|
||||
DONE,
|
||||
]
|
||||
.concat(),
|
||||
_ => [CONTENT_CHUNK, USAGE_CHUNK, DONE].concat(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn upstream(calls: Calls) -> String {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let app = Router::new().fallback(move |request: Request<Body>| {
|
||||
let calls = calls.clone();
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let headers = parts
|
||||
.headers
|
||||
.iter()
|
||||
.filter_map(|(name, value)| {
|
||||
Some((name.as_str().to_owned(), value.to_str().ok()?.to_owned()))
|
||||
})
|
||||
.collect();
|
||||
let path = parts.uri.path().to_owned();
|
||||
let body: Value =
|
||||
serde_json::from_slice(&to_bytes(body, 1 << 20).await.unwrap()).unwrap();
|
||||
let stream = body["stream"].as_bool().unwrap_or(false);
|
||||
calls.lock().unwrap().push(Sent {
|
||||
path: path.clone(),
|
||||
headers,
|
||||
body,
|
||||
});
|
||||
match stream {
|
||||
true => {
|
||||
([(CONTENT_TYPE, "text/event-stream")], stream_answer(&path)).into_response()
|
||||
}
|
||||
false => ([(CONTENT_TYPE, "application/json")], COMPLETION).into_response(),
|
||||
}
|
||||
}
|
||||
});
|
||||
tokio::spawn(async move { axum::serve(listener, app).await });
|
||||
format!("http://{address}")
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Seen {
|
||||
deployment: Option<String>,
|
||||
usage: Option<(u64, u64)>,
|
||||
cost: Option<f64>,
|
||||
}
|
||||
|
||||
struct Collect(Mutex<Vec<Seen>>);
|
||||
|
||||
impl Observer for Collect {
|
||||
fn name(&self) -> &'static str {
|
||||
"collect"
|
||||
}
|
||||
|
||||
fn reads(&self) -> Facts {
|
||||
Facts::DEPLOYMENT | Facts::USAGE | Facts::COST
|
||||
}
|
||||
|
||||
fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
self.0.lock().unwrap().push(Seen {
|
||||
deployment: record
|
||||
.outcome
|
||||
.deployment
|
||||
.as_ref()
|
||||
.map(|deployment| deployment.id.to_string()),
|
||||
usage: record
|
||||
.outcome
|
||||
.usage
|
||||
.map(|usage| (usage.prompt_tokens(), usage.output_tokens)),
|
||||
cost: record.outcome.cost,
|
||||
});
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct Harness {
|
||||
app: Router,
|
||||
inference: Arc<litellm_gateway_inference::Gateway>,
|
||||
calls: Calls,
|
||||
seen: Arc<Collect>,
|
||||
}
|
||||
|
||||
/// One deployment per provider, each at its own path on the mock upstream.
|
||||
const DEPLOYMENTS: &[(&str, &str, &str, Option<&str>)] = &[
|
||||
(
|
||||
"vllm",
|
||||
"hosted_vllm/meta-llama/Llama-3.1-8B-Instruct",
|
||||
"/vllm/v1",
|
||||
None,
|
||||
),
|
||||
(
|
||||
"nim",
|
||||
"nvidia_nim/meta/llama3-70b-instruct",
|
||||
"/nim/v1",
|
||||
Some("sk-nim"),
|
||||
),
|
||||
(
|
||||
"groq",
|
||||
"groq/llama-3.1-8b-instant",
|
||||
"/groq/openai/v1",
|
||||
Some("sk-groq"),
|
||||
),
|
||||
(
|
||||
"together",
|
||||
"together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo",
|
||||
"/together/v1",
|
||||
Some("sk-together"),
|
||||
),
|
||||
(
|
||||
"fireworks",
|
||||
"fireworks_ai/llama-v3p1-8b-instruct",
|
||||
"/fireworks/inference/v1",
|
||||
Some("sk-fireworks"),
|
||||
),
|
||||
("xai", "xai/grok-3", "/xai/v1", Some("sk-xai")),
|
||||
(
|
||||
"mistral",
|
||||
"mistral/mistral-small-latest",
|
||||
"/mistral",
|
||||
Some("sk-mistral"),
|
||||
),
|
||||
(
|
||||
"deepseek",
|
||||
"deepseek/deepseek-chat",
|
||||
"/deepseek/beta",
|
||||
Some("sk-deepseek"),
|
||||
),
|
||||
(
|
||||
"openrouter",
|
||||
"openrouter/openai/gpt-4o",
|
||||
"/openrouter/api/v1",
|
||||
Some("sk-openrouter"),
|
||||
),
|
||||
(
|
||||
"cerebras",
|
||||
"cerebras/llama3.1-8b",
|
||||
"/cerebras/v1",
|
||||
Some("sk-cerebras"),
|
||||
),
|
||||
(
|
||||
"perplexity",
|
||||
"perplexity/sonar",
|
||||
"/perplexity",
|
||||
Some("sk-perplexity"),
|
||||
),
|
||||
(
|
||||
"custom-openai",
|
||||
"openai/my-served-model",
|
||||
"/custom/v1",
|
||||
Some("sk-custom"),
|
||||
),
|
||||
];
|
||||
|
||||
async fn harness() -> Harness {
|
||||
use_local_cost_map();
|
||||
let calls = Calls::default();
|
||||
let upstream = upstream(calls.clone()).await;
|
||||
let models: String = DEPLOYMENTS
|
||||
.iter()
|
||||
.map(|(name, model, base, key)| {
|
||||
let key = key.map_or_else(String::new, |key| format!("\n api_key: {key}"));
|
||||
format!(
|
||||
" - model_name: {name}
|
||||
litellm_params:
|
||||
model: {model}
|
||||
api_base: {upstream}{base}{key}
|
||||
model_info:
|
||||
id: {name}-1
|
||||
"
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
let config = Config::from_yaml(&format!(
|
||||
"model_list:
|
||||
{models}router_settings:
|
||||
num_retries: 0
|
||||
general_settings:
|
||||
master_key: sk-master
|
||||
"
|
||||
))
|
||||
.unwrap();
|
||||
let inference = litellm_gateway::build_inference(&config).unwrap();
|
||||
let seen = Arc::new(Collect(Mutex::new(Vec::new())));
|
||||
litellm_gateway::hooks::install(&inference, &config, Vec::new(), vec![seen.clone()])
|
||||
.await
|
||||
.unwrap();
|
||||
let app = litellm_gateway::router(inference.clone(), &config, None, None);
|
||||
Harness {
|
||||
app,
|
||||
inference,
|
||||
calls,
|
||||
seen,
|
||||
}
|
||||
}
|
||||
|
||||
async fn send(app: &Router, body: Value) -> (u16, String) {
|
||||
let request = Request::post("/v1/chat/completions")
|
||||
.header("content-type", "application/json")
|
||||
.header("authorization", "Bearer sk-master")
|
||||
.body(Body::from(body.to_string()))
|
||||
.unwrap();
|
||||
let response = app.clone().oneshot(request).await.unwrap();
|
||||
let status = response.status().as_u16();
|
||||
let body = to_bytes(response.into_body(), 1 << 20).await.unwrap();
|
||||
(status, String::from_utf8(body.to_vec()).unwrap())
|
||||
}
|
||||
|
||||
async fn seen(harness: &Harness) -> Vec<Seen> {
|
||||
litellm_gateway::hooks::close(&harness.inference).await;
|
||||
std::mem::take(&mut *harness.seen.0.lock().unwrap())
|
||||
}
|
||||
|
||||
fn hi(model: &str) -> Value {
|
||||
json!({
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"max_completion_tokens": 8,
|
||||
"temperature": 0.5,
|
||||
"metadata": {"trace": "t"},
|
||||
})
|
||||
}
|
||||
|
||||
fn messages() -> Value {
|
||||
json!([{"role": "user", "content": "hi"}])
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::hosted_vllm("vllm", "/vllm/v1/chat/completions", "Bearer fake-api-key", None,
|
||||
json!({"model": "meta-llama/Llama-3.1-8B-Instruct", "messages": messages(), "max_completion_tokens": 8, "temperature": 0.5}))]
|
||||
#[case::nvidia_nim("nim", "/nim/v1/chat/completions", "Bearer sk-nim", None,
|
||||
json!({"model": "meta/llama3-70b-instruct", "messages": messages(), "max_tokens": 8, "temperature": 0.5}))]
|
||||
#[case::groq("groq", "/groq/openai/v1/chat/completions", "Bearer sk-groq", None,
|
||||
json!({"model": "llama-3.1-8b-instant", "messages": messages(), "max_tokens": 8, "temperature": 0.5}))]
|
||||
#[case::together_ai("together", "/together/v1/chat/completions", "Bearer sk-together", None,
|
||||
json!({"model": "meta-llama/Llama-3.3-70B-Instruct-Turbo", "messages": messages(), "max_completion_tokens": 8, "temperature": 0.5}))]
|
||||
#[case::fireworks_ai("fireworks", "/fireworks/inference/v1/chat/completions", "Bearer sk-fireworks", None,
|
||||
json!({"model": "accounts/fireworks/models/llama-v3p1-8b-instruct", "messages": messages(), "max_tokens": 8, "temperature": 0.5}))]
|
||||
#[case::xai("xai", "/xai/v1/chat/completions", "Bearer sk-xai", None,
|
||||
json!({"model": "grok-3", "messages": messages(), "max_tokens": 8, "temperature": 0.5}))]
|
||||
#[case::mistral("mistral", "/mistral/v1/chat/completions", "Bearer sk-mistral", None,
|
||||
json!({"model": "mistral-small-latest", "messages": messages(), "max_tokens": 8, "temperature": 0.5}))]
|
||||
#[case::deepseek("deepseek", "/deepseek/beta/chat/completions", "Bearer sk-deepseek", None,
|
||||
json!({"model": "deepseek-chat", "messages": messages(), "max_completion_tokens": 8, "temperature": 0.5}))]
|
||||
#[case::openrouter("openrouter", "/openrouter/api/v1/chat/completions", "Bearer sk-openrouter",
|
||||
Some(("x-title", "liteLLM")),
|
||||
json!({"model": "openai/gpt-4o", "messages": messages(), "max_completion_tokens": 8, "temperature": 0.5, "usage": {"include": true}}))]
|
||||
#[case::cerebras("cerebras", "/cerebras/v1/chat/completions", "Bearer sk-cerebras", None,
|
||||
json!({"model": "llama3.1-8b", "messages": messages(), "max_tokens": 8, "temperature": 0.5}))]
|
||||
#[case::perplexity("perplexity", "/perplexity/chat/completions", "Bearer sk-perplexity",
|
||||
Some(("x-pplx-integration", "litellm")),
|
||||
json!({"model": "sonar", "messages": messages(), "max_completion_tokens": 8, "temperature": 0.5}))]
|
||||
#[case::openai_custom_base("custom-openai", "/custom/v1/chat/completions", "Bearer sk-custom", None,
|
||||
json!({"model": "my-served-model", "messages": messages(), "max_completion_tokens": 8, "temperature": 0.5}))]
|
||||
#[tokio::test]
|
||||
async fn each_provider_gets_its_url_auth_and_body(
|
||||
#[case] model: &str,
|
||||
#[case] path: &str,
|
||||
#[case] authorization: &str,
|
||||
#[case] header: Option<(&str, &str)>,
|
||||
#[case] body: Value,
|
||||
) {
|
||||
let harness = harness().await;
|
||||
let (status, answer) = send(&harness.app, hi(model)).await;
|
||||
assert_eq!(status, 200, "{answer}");
|
||||
assert_eq!(
|
||||
serde_json::from_str::<Value>(&answer).unwrap()["choices"][0]["message"]["content"],
|
||||
"hi"
|
||||
);
|
||||
|
||||
let sent = harness.calls.lock().unwrap()[0].clone();
|
||||
assert_eq!(sent.path, path);
|
||||
assert_eq!(
|
||||
sent.headers.get("authorization").map(String::as_str),
|
||||
Some(authorization)
|
||||
);
|
||||
if let Some((name, value)) = header {
|
||||
assert_eq!(sent.headers.get(name).map(String::as_str), Some(value));
|
||||
}
|
||||
assert_eq!(sent.body, body);
|
||||
|
||||
let seen = seen(&harness).await;
|
||||
assert_eq!(seen.len(), 1, "{seen:?}");
|
||||
assert_eq!(seen[0].deployment, Some(format!("{model}-1")));
|
||||
assert_eq!(seen[0].usage, Some((5, 6)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::groq_asks_on_the_clients_behalf("groq", false, Some(json!({"include_usage": true})), (7, 3))]
|
||||
#[case::groq_client_asked("groq", true, Some(json!({"include_usage": true})), (7, 3))]
|
||||
#[case::mistral_reports_unasked("mistral", false, None, (7, 3))]
|
||||
#[case::xai_folds_reasoning_into_completion("xai", false, Some(json!({"include_usage": true})), (7, 7))]
|
||||
#[case::hosted_vllm("vllm", false, Some(json!({"include_usage": true})), (7, 3))]
|
||||
#[tokio::test]
|
||||
async fn a_stream_is_metered_and_priced(
|
||||
#[case] model: &str,
|
||||
#[case] client_usage: bool,
|
||||
#[case] stream_options: Option<Value>,
|
||||
#[case] usage: (u64, u64),
|
||||
) {
|
||||
let harness = harness().await;
|
||||
let mut request = hi(model);
|
||||
request["stream"] = json!(true);
|
||||
if client_usage {
|
||||
request["stream_options"] = json!({"include_usage": true});
|
||||
}
|
||||
let (status, body) = send(&harness.app, request).await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
assert!(body.contains("\"content\":\"hi\""), "{body}");
|
||||
assert!(body.ends_with(DONE), "{body}");
|
||||
assert_eq!(body.contains("\"usage\":{"), client_usage, "{body}");
|
||||
|
||||
let sent = harness.calls.lock().unwrap()[0].clone();
|
||||
assert_eq!(sent.body.get("stream_options").cloned(), stream_options);
|
||||
|
||||
let seen = seen(&harness).await;
|
||||
assert_eq!(seen.len(), 1, "{seen:?}");
|
||||
assert_eq!(seen[0].usage, Some(usage));
|
||||
let expected = usage.0 as f64 * RATES.0 + usage.1 as f64 * RATES.1;
|
||||
assert!((seen[0].cost.unwrap() - expected).abs() < 1e-9, "{seen:?}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_vllm_stream_names_its_reasoning_as_litellm_does() {
|
||||
let harness = harness().await;
|
||||
let mut request = hi("vllm");
|
||||
request["stream"] = json!(true);
|
||||
let (status, body) = send(&harness.app, request).await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
assert!(body.contains("\"reasoning_content\":\"hmm\""), "{body}");
|
||||
assert!(!body.contains("\"reasoning\":"), "{body}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::groq_streamed_response_format("groq", json!({"stream": true, "response_format": {"type": "json_object"}}), 501)]
|
||||
#[case::param_the_model_lacks("cerebras", json!({"frequency_penalty": 1}), 400)]
|
||||
#[case::fireworks_file_part("fireworks", json!({"messages": [{"role": "user", "content": [{"type": "file", "file": {"file_id": "f"}}]}]}), 400)]
|
||||
#[tokio::test]
|
||||
async fn requests_a_provider_cannot_take_are_refused_before_the_call(
|
||||
#[case] model: &str,
|
||||
#[case] extra: Value,
|
||||
#[case] status: u16,
|
||||
) {
|
||||
let harness = harness().await;
|
||||
let mut request = hi(model);
|
||||
let (Value::Object(request_fields), Value::Object(extra)) = (&mut request, extra) else {
|
||||
unreachable!()
|
||||
};
|
||||
request_fields.extend(extra);
|
||||
let (sent_status, body) = send(&harness.app, request).await;
|
||||
assert_eq!(sent_status, status, "{body}");
|
||||
assert!(harness.calls.lock().unwrap().is_empty());
|
||||
}
|
||||
624
litellm-rust/crates/gateway/tests/vertex_claude_chat.rs
Normal file
624
litellm-rust/crates/gateway/tests/vertex_claude_chat.rs
Normal file
|
|
@ -0,0 +1,624 @@
|
|||
//! Claude chat completions on Vertex AI through the gateway (`vertex_ai/claude-*`, a service
|
||||
//! account traded for an access token): the `rawPredict` request Vertex gets, the OpenAI answer
|
||||
//! the client gets, and the usage and cost the request is recorded with.
|
||||
|
||||
use std::{
|
||||
io::Write,
|
||||
sync::{Arc, Mutex, Once},
|
||||
};
|
||||
|
||||
use axum::{
|
||||
Router,
|
||||
body::{Body, to_bytes},
|
||||
http::{Request, StatusCode, header::CONTENT_TYPE},
|
||||
response::IntoResponse,
|
||||
};
|
||||
use futures_util::future::BoxFuture;
|
||||
use litellm_config::Config;
|
||||
use litellm_gateway_hooks::{Facts, Observer, Record};
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
use tokio::net::TcpListener;
|
||||
use tower::ServiceExt;
|
||||
|
||||
/// Input and output prices of every deployment.
|
||||
const RATES: (f64, f64) = (1.0, 2.0);
|
||||
|
||||
const ACCESS_TOKEN: &str = "ya29.test-token";
|
||||
const PROJECT: &str = "test-project";
|
||||
const LOCATION: &str = "us-east5";
|
||||
const MODEL: &str = "claude-sonnet-4-5@20250929";
|
||||
|
||||
/// A throwaway RSA key the fake service account signs its token request with.
|
||||
const TEST_PRIVATE_KEY: &str = "-----BEGIN PRIVATE KEY-----
|
||||
MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQC95ihJQlfJuaQe
|
||||
ZCwF+EFy/Y5vk7Lx6Nm74He9Oes+6ZbgA2LVdcoFJrNEYHd8S0CCv7cqEp4dpIoR
|
||||
B4HAmQAQF92ZC/S2lcgLAezWQU21TSjDLngdVbn+Jr4syzOwVw/NUigydtQ66Iqv
|
||||
1FpXs/4gQXMgs8FW5XgOhgfKhyMPf71A2qiRaA0vwsrfsmKOY0QVVooXPRfpLIvZ
|
||||
VEOV89ZMQErnYan+m8OADoZ5NJ/Ij6QKb2HkQJNhQC2Zq2T1Y9VLyxY6pXURdVqc
|
||||
XJgFT6YPGjLABLX/oId4g+Fs2YI5mM+g72pYEntnIaE5yicGKBdjyqeVfKZ6/7jg
|
||||
g8evHi6vAgMBAAECggEAGdJjlP6b8Fa5bdaCM/ebcrbuuNZVJVbb0JPHxGfNSLs7
|
||||
pE9hj5QaOdQW2Uviw3h6F61ZCzQH4xD+Iy2po5ZKb2XHYKnDB1bboj+LRGER337T
|
||||
9aJqe9at2VTMVEv3Rdm40NsEk0QcPLxlK16NQFK90gYEUSSQPDAswJDSG2R/zHn+
|
||||
vADI907mW/goEJHeLn8PWGlNlSiR6x+5JJtq+GXCzUzVvJYQSCLGxCSl2x2H+0g7
|
||||
NhFI0zPpdzNmO/h+yhzaFb6Rp5U8+ZsnZ3qYjQ/03gw1myTDKJt1YaO9JvArnNYX
|
||||
hcJQQ8Rt0bHhcrZA16bBOpqZlo5pKCicwI/netgN8QKBgQDcFz7AzdJ26sMSV32V
|
||||
rwrMgIoggt8qDjO1ARwqW35A1TIge0FoW4M4KpsXQGGfT341uU1esXEcyZ/1L/5X
|
||||
3ql2gX4DbOYLZLWYzZGR2hq33oi8HkhN98QrEwL9emSH8NqYX3Xxja3PrmCrSYJe
|
||||
Zbnd9TIm2XkxyMoyXJu6M/QvnwKBgQDc4dzqTbxoGEGa5MuJoGmMwPnqgdG9UM5J
|
||||
eExVnh7osxc2sOdsiPeRjjQTxs9v2kJwctC359OJoo9yGaaJeSghU4LEWJo1sqnA
|
||||
fzSCLammYvtVAtniyNv5Mxk/6Uimi4NNDKaAKB+m4K2uSn3U9AmY7KPYMGaSbS9W
|
||||
XSnobjxm8QKBgC8bPpAvvWs8ZhIn7bY659nLbUT2HeO3dHO6UBf0yzn/J6JyHxbB
|
||||
93zvCZDZc8uQTRgcmCW7XtVlhjoJUqvl+Wlm39zF0xr/LCsPXKfWAb/2/lcdOCaP
|
||||
8Emz4QD10EyUTYUtcWYJB/mafhBLRH8F0Nlj4J8WDu2L51MOJTqeYhZLAoGAWffN
|
||||
icocAbJPlo22sdoa4+/+W5yBF8GAJMDRJtZ+9H1t6SLpQHYRkMIBSETkXUTjZvX9
|
||||
Ocs9iIQkNW9pO/mTdO+VBfCo71JUfknR02xR+6m5gYjlws/ZeYlssXGN2/hbhNiw
|
||||
QOcW7Vv6olFJK6Iy/oz0t6wPO3kpnN3Zogi0paECgYEAwo44M1DdYCtV0snhmYM9
|
||||
5u0mPfYt5P2SVLXyUbr+vFTfrTL/WKnXIJgbsnj3Gvf+GIZv9tKcXhSNmEHQCYX4
|
||||
X3w9iTPddCHuvZ1fpufi2TyArJh0OkoNtLXJHTKrHjf2N+61AQzFiv5WieJrdE+H
|
||||
qr32PTUuVGPyO9LyTY4/RL0=
|
||||
-----END PRIVATE KEY-----
|
||||
";
|
||||
|
||||
/// The gateway reads its cost map from this variable; every test sets it before building a
|
||||
/// gateway, so no test reads the environment while it changes.
|
||||
fn use_local_cost_map() {
|
||||
static ONCE: Once = Once::new();
|
||||
ONCE.call_once(|| {
|
||||
let map = json!({
|
||||
format!("vertex_ai/{MODEL}"): {
|
||||
"mode": "chat",
|
||||
"litellm_provider": "vertex_ai-anthropic_models",
|
||||
"max_output_tokens": 64000,
|
||||
"supports_function_calling": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_reasoning": true,
|
||||
},
|
||||
});
|
||||
let mut file = tempfile::NamedTempFile::new().unwrap();
|
||||
file.write_all(map.to_string().as_bytes()).unwrap();
|
||||
let (_, path) = file.keep().unwrap();
|
||||
// SAFETY: set once, before any test builds a gateway or spawns a thread that reads the
|
||||
// environment.
|
||||
unsafe { std::env::set_var("LITELLM_MODEL_COST_MAP_PATH", path) };
|
||||
});
|
||||
}
|
||||
|
||||
fn message(content: Value, stop_reason: &str) -> Value {
|
||||
json!({
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": content,
|
||||
"stop_reason": stop_reason,
|
||||
"stop_sequence": null,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 4},
|
||||
})
|
||||
}
|
||||
|
||||
fn text_answer(text: &str) -> Value {
|
||||
message(json!([{"type": "text", "text": text}]), "end_turn")
|
||||
}
|
||||
|
||||
fn tool_answer() -> Value {
|
||||
message(
|
||||
json!([{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Paris"}}]),
|
||||
"tool_use",
|
||||
)
|
||||
}
|
||||
|
||||
fn thinking_answer() -> Value {
|
||||
message(
|
||||
json!([
|
||||
{"type": "thinking", "thinking": "The user says hi.", "signature": "sig-1"},
|
||||
{"type": "text", "text": "Hello"}
|
||||
]),
|
||||
"end_turn",
|
||||
)
|
||||
}
|
||||
|
||||
fn sse_answer() -> String {
|
||||
[
|
||||
json!({"type": "message_start", "message": {"id": "msg_2", "type": "message", "role": "assistant", "model": "claude-sonnet-4-5", "content": [], "stop_reason": null, "stop_sequence": null, "usage": {"input_tokens": 10, "output_tokens": 1}}}),
|
||||
json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
|
||||
json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello"}}),
|
||||
json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": " world"}}),
|
||||
json!({"type": "content_block_stop", "index": 0}),
|
||||
json!({"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": null}, "usage": {"output_tokens": 12}}),
|
||||
json!({"type": "message_stop"}),
|
||||
]
|
||||
.iter()
|
||||
.map(|event| format!("event: {}\ndata: {event}\n\n", event["type"].as_str().unwrap()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct Sent {
|
||||
path: String,
|
||||
query: Option<String>,
|
||||
authorization: Option<String>,
|
||||
api_key: Option<String>,
|
||||
anthropic_version: Option<String>,
|
||||
anthropic_beta: Option<String>,
|
||||
body: Value,
|
||||
}
|
||||
|
||||
type Calls = Arc<Mutex<Vec<Sent>>>;
|
||||
|
||||
fn has_tool_result(body: &Value) -> bool {
|
||||
body["messages"].as_array().is_some_and(|messages| {
|
||||
messages.iter().any(|message| {
|
||||
message["content"]
|
||||
.as_array()
|
||||
.is_some_and(|blocks| blocks.iter().any(|block| block["type"] == "tool_result"))
|
||||
})
|
||||
})
|
||||
}
|
||||
|
||||
/// Google's token endpoint and Vertex's `rawPredict` and `streamRawPredict`, answering by what
|
||||
/// the body asks.
|
||||
async fn upstream(calls: Calls) -> String {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let address = listener.local_addr().unwrap();
|
||||
let app = Router::new().fallback(move |request: Request<Body>| {
|
||||
let calls = calls.clone();
|
||||
async move {
|
||||
let (parts, body) = request.into_parts();
|
||||
let header = |name: &str| {
|
||||
parts
|
||||
.headers
|
||||
.get(name)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_owned)
|
||||
};
|
||||
let bytes = to_bytes(body, 1 << 20).await.unwrap();
|
||||
let path = parts.uri.path().to_owned();
|
||||
if path == "/token" {
|
||||
return (
|
||||
[(CONTENT_TYPE, "application/json")],
|
||||
json!({"access_token": ACCESS_TOKEN, "expires_in": 3600, "token_type": "Bearer"})
|
||||
.to_string(),
|
||||
)
|
||||
.into_response();
|
||||
}
|
||||
let body: Value = serde_json::from_slice(&bytes).unwrap_or(Value::Null);
|
||||
calls.lock().unwrap().push(Sent {
|
||||
path: path.clone(),
|
||||
query: parts.uri.query().map(str::to_owned),
|
||||
authorization: header("authorization"),
|
||||
api_key: header("x-api-key"),
|
||||
anthropic_version: header("anthropic-version"),
|
||||
anthropic_beta: header("anthropic-beta"),
|
||||
body: body.clone(),
|
||||
});
|
||||
match path.ends_with(":streamRawPredict") {
|
||||
true => ([(CONTENT_TYPE, "text/event-stream")], sse_answer()).into_response(),
|
||||
false => {
|
||||
let answer = match (has_tool_result(&body), body.get("tools"), body.get("thinking")) {
|
||||
(true, _, _) => text_answer("It is sunny in Paris"),
|
||||
(false, Some(_), _) => tool_answer(),
|
||||
(false, None, Some(_)) => thinking_answer(),
|
||||
(false, None, None) => text_answer("Hello world"),
|
||||
};
|
||||
(
|
||||
StatusCode::OK,
|
||||
[(CONTENT_TYPE, "application/json")],
|
||||
answer.to_string(),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
tokio::spawn(async move { axum::serve(listener, app).await });
|
||||
format!("http://{address}")
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Seen {
|
||||
deployment: Option<String>,
|
||||
usage: Option<litellm_gateway_hooks::Usage>,
|
||||
cost: Option<f64>,
|
||||
}
|
||||
|
||||
struct Collect(Mutex<Vec<Seen>>);
|
||||
|
||||
impl Observer for Collect {
|
||||
fn name(&self) -> &'static str {
|
||||
"collect"
|
||||
}
|
||||
|
||||
fn reads(&self) -> Facts {
|
||||
Facts::DEPLOYMENT | Facts::USAGE | Facts::COST
|
||||
}
|
||||
|
||||
fn observe<'a>(&'a self, record: &'a Record) -> BoxFuture<'a, ()> {
|
||||
Box::pin(async move {
|
||||
self.0.lock().unwrap().push(Seen {
|
||||
deployment: record
|
||||
.outcome
|
||||
.deployment
|
||||
.as_ref()
|
||||
.map(|deployment| deployment.id.to_string()),
|
||||
usage: record.outcome.usage,
|
||||
cost: record.outcome.cost,
|
||||
});
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct Harness {
|
||||
app: Router,
|
||||
inference: Arc<litellm_gateway_inference::Gateway>,
|
||||
calls: Calls,
|
||||
seen: Arc<Collect>,
|
||||
}
|
||||
|
||||
fn service_account(upstream: &str) -> String {
|
||||
json!({
|
||||
"type": "service_account",
|
||||
"project_id": "credentials-project",
|
||||
"private_key_id": "test-key-id",
|
||||
"private_key": TEST_PRIVATE_KEY,
|
||||
"client_email": "gateway-test@credentials-project.iam.gserviceaccount.com",
|
||||
"client_id": "1",
|
||||
"token_uri": format!("{upstream}/token"),
|
||||
})
|
||||
.to_string()
|
||||
}
|
||||
|
||||
async fn harness() -> Harness {
|
||||
use_local_cost_map();
|
||||
let calls = Calls::default();
|
||||
let upstream = upstream(calls.clone()).await;
|
||||
let (input, output) = RATES;
|
||||
let prices =
|
||||
format!(" input_cost_per_token: {input}\n output_cost_per_token: {output}\n");
|
||||
let credentials = service_account(&upstream);
|
||||
let config = Config::from_yaml(&format!(
|
||||
"model_list:
|
||||
- model_name: claude
|
||||
litellm_params:
|
||||
model: vertex_ai/{MODEL}
|
||||
api_base: {upstream}
|
||||
vertex_project: {PROJECT}
|
||||
vertex_location: {LOCATION}
|
||||
vertex_credentials: '{credentials}'
|
||||
{prices} model_info:
|
||||
id: claude-1
|
||||
- model_name: claude-global
|
||||
litellm_params:
|
||||
model: vertex_ai/{MODEL}
|
||||
api_base: {upstream}
|
||||
vertex_project: {PROJECT}
|
||||
vertex_location: global
|
||||
vertex_credentials: '{credentials}'
|
||||
{prices} model_info:
|
||||
id: claude-global-1
|
||||
general_settings:
|
||||
master_key: sk-master
|
||||
"
|
||||
))
|
||||
.unwrap();
|
||||
let inference = litellm_gateway::build_inference(&config).unwrap();
|
||||
let seen = Arc::new(Collect(Mutex::new(Vec::new())));
|
||||
litellm_gateway::hooks::install(&inference, &config, Vec::new(), vec![seen.clone()])
|
||||
.await
|
||||
.unwrap();
|
||||
let app = litellm_gateway::router(inference.clone(), &config, None, None);
|
||||
Harness {
|
||||
app,
|
||||
inference,
|
||||
calls,
|
||||
seen,
|
||||
}
|
||||
}
|
||||
|
||||
async fn send(app: &Router, body: Value) -> (u16, String) {
|
||||
let request = Request::post("/v1/chat/completions")
|
||||
.header("content-type", "application/json")
|
||||
.header("authorization", "Bearer sk-master")
|
||||
.body(Body::from(body.to_string()))
|
||||
.unwrap();
|
||||
let response = app.clone().oneshot(request).await.unwrap();
|
||||
let status = response.status().as_u16();
|
||||
let body = to_bytes(response.into_body(), 1 << 20).await.unwrap();
|
||||
(status, String::from_utf8(body.to_vec()).unwrap())
|
||||
}
|
||||
|
||||
async fn seen(harness: &Harness) -> Vec<Seen> {
|
||||
litellm_gateway::hooks::close(&harness.inference).await;
|
||||
std::mem::take(&mut *harness.seen.0.lock().unwrap())
|
||||
}
|
||||
|
||||
fn cost(prompt: u64, completion: u64) -> f64 {
|
||||
let (input, output) = RATES;
|
||||
prompt as f64 * input + completion as f64 * output
|
||||
}
|
||||
|
||||
fn hi(model: &str) -> Value {
|
||||
json!({
|
||||
"model": model,
|
||||
"messages": [
|
||||
{"role": "system", "content": "Be brief"},
|
||||
{"role": "user", "content": "hi"}
|
||||
],
|
||||
"temperature": 0.2,
|
||||
"max_tokens": 64,
|
||||
})
|
||||
}
|
||||
|
||||
fn sent_body() -> Value {
|
||||
json!({
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}],
|
||||
"temperature": 0.2,
|
||||
"max_tokens": 64,
|
||||
"anthropic_version": "vertex-2023-10-16",
|
||||
"system": [{"type": "text", "text": "Be brief"}],
|
||||
})
|
||||
}
|
||||
|
||||
fn raw_predict_path(location: &str, method: &str) -> String {
|
||||
format!(
|
||||
"/v1/projects/{PROJECT}/locations/{location}/publishers/anthropic/models/{MODEL}:{method}"
|
||||
)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::regional("claude", "claude-1", LOCATION)]
|
||||
#[case::global("claude-global", "claude-global-1", "global")]
|
||||
#[tokio::test]
|
||||
async fn a_chat_completion_calls_raw_predict_and_records_its_cost(
|
||||
#[case] model: &str,
|
||||
#[case] id: &str,
|
||||
#[case] location: &str,
|
||||
) {
|
||||
let harness = harness().await;
|
||||
let (status, body) = send(&harness.app, hi(model)).await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
let answer: Value = serde_json::from_str(&body).unwrap();
|
||||
// `VertexAIAnthropicConfig.transform_response` names the requested model.
|
||||
assert_eq!(answer["model"], MODEL);
|
||||
assert_eq!(answer["choices"][0]["message"]["role"], "assistant");
|
||||
assert_eq!(answer["choices"][0]["message"]["content"], "Hello world");
|
||||
assert_eq!(answer["choices"][0]["finish_reason"], "stop");
|
||||
assert_eq!(answer["usage"]["prompt_tokens"], 10);
|
||||
assert_eq!(answer["usage"]["completion_tokens"], 4);
|
||||
assert_eq!(answer["usage"]["total_tokens"], 14);
|
||||
|
||||
let sent = harness.calls.lock().unwrap()[0].clone();
|
||||
assert_eq!(sent.path, raw_predict_path(location, "rawPredict"));
|
||||
assert_eq!(sent.query, None);
|
||||
assert_eq!(
|
||||
sent.authorization.as_deref(),
|
||||
Some("Bearer ya29.test-token")
|
||||
);
|
||||
assert_eq!(sent.api_key, None);
|
||||
assert_eq!(sent.anthropic_version.as_deref(), Some("2023-06-01"));
|
||||
assert_eq!(sent.anthropic_beta, None);
|
||||
assert_eq!(sent.body, sent_body());
|
||||
|
||||
let seen = seen(&harness).await;
|
||||
assert_eq!(seen.len(), 1, "{seen:?}");
|
||||
assert_eq!(seen[0].deployment.as_deref(), Some(id));
|
||||
let usage = seen[0].usage.unwrap();
|
||||
assert_eq!((usage.prompt_tokens(), usage.output_tokens), (10, 4));
|
||||
assert_eq!(seen[0].cost, Some(cost(10, 4)));
|
||||
}
|
||||
|
||||
/// The JSON of each `data:` frame, and whether the stream closed with `[DONE]`.
|
||||
fn frames(body: &str) -> (Vec<Value>, bool) {
|
||||
let mut done = false;
|
||||
let frames = body
|
||||
.split("\n\n")
|
||||
.filter(|frame| !frame.is_empty())
|
||||
.filter_map(|frame| {
|
||||
let data = frame.strip_prefix("data: ").expect("bare data frames");
|
||||
match data {
|
||||
"[DONE]" => {
|
||||
done = true;
|
||||
None
|
||||
}
|
||||
data => Some(serde_json::from_str(data).unwrap()),
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
(frames, done)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::regional("claude", "claude-1", LOCATION)]
|
||||
#[case::global("claude-global", "claude-global-1", "global")]
|
||||
#[tokio::test]
|
||||
async fn a_stream_calls_stream_raw_predict_and_records_its_usage(
|
||||
#[case] model: &str,
|
||||
#[case] id: &str,
|
||||
#[case] location: &str,
|
||||
) {
|
||||
let harness = harness().await;
|
||||
let mut request = hi(model);
|
||||
request["stream"] = json!(true);
|
||||
request["stream_options"] = json!({"include_usage": true});
|
||||
let (status, body) = send(&harness.app, request).await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
|
||||
let (frames, done) = frames(&body);
|
||||
assert!(done, "{body}");
|
||||
assert!(
|
||||
frames
|
||||
.iter()
|
||||
.all(|frame| frame["object"] == "chat.completion.chunk" && frame["model"] == model),
|
||||
"{body}"
|
||||
);
|
||||
let content: String = frames
|
||||
.iter()
|
||||
.filter_map(|frame| frame["choices"][0]["delta"]["content"].as_str())
|
||||
.collect();
|
||||
assert_eq!(content, "Hello world", "{body}");
|
||||
let finishes: Vec<&Value> = frames
|
||||
.iter()
|
||||
.filter_map(|frame| frame["choices"][0].get("finish_reason"))
|
||||
.filter(|reason| !reason.is_null())
|
||||
.collect();
|
||||
assert_eq!(finishes, [&json!("stop")], "{body}");
|
||||
let usages: Vec<&Value> = frames
|
||||
.iter()
|
||||
.filter_map(|frame| frame.get("usage"))
|
||||
.filter(|usage| !usage.is_null())
|
||||
.collect();
|
||||
assert_eq!(usages.len(), 1, "{body}");
|
||||
assert_eq!(usages[0]["prompt_tokens"], 10);
|
||||
assert_eq!(usages[0]["completion_tokens"], 12);
|
||||
|
||||
let sent = harness.calls.lock().unwrap()[0].clone();
|
||||
assert_eq!(sent.path, raw_predict_path(location, "streamRawPredict"));
|
||||
// A custom `api_base` asks for SSE, as `_check_custom_proxy` does.
|
||||
assert_eq!(sent.query.as_deref(), Some("alt=sse"));
|
||||
assert_eq!(
|
||||
sent.authorization.as_deref(),
|
||||
Some("Bearer ya29.test-token")
|
||||
);
|
||||
let mut expected = sent_body();
|
||||
expected["stream"] = json!(true);
|
||||
assert_eq!(sent.body, expected);
|
||||
|
||||
let seen = seen(&harness).await;
|
||||
assert_eq!(seen.len(), 1, "{seen:?}");
|
||||
assert_eq!(seen[0].deployment.as_deref(), Some(id));
|
||||
let usage = seen[0].usage.unwrap();
|
||||
assert_eq!((usage.prompt_tokens(), usage.output_tokens), (10, 12));
|
||||
assert_eq!(seen[0].cost, Some(cost(10, 12)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_stream_without_include_usage_still_reports_usage_to_the_hooks() {
|
||||
let harness = harness().await;
|
||||
let mut request = hi("claude");
|
||||
request["stream"] = json!(true);
|
||||
let (status, body) = send(&harness.app, request).await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
let (frames, done) = frames(&body);
|
||||
assert!(done, "{body}");
|
||||
assert!(
|
||||
frames
|
||||
.iter()
|
||||
.all(|frame| frame.get("usage").is_none_or(Value::is_null)),
|
||||
"{body}"
|
||||
);
|
||||
let seen = seen(&harness).await;
|
||||
assert_eq!(seen.len(), 1, "{seen:?}");
|
||||
assert_eq!(seen[0].cost, Some(cost(10, 12)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_tool_call_round_trips_through_tool_use_and_tool_result() {
|
||||
let harness = harness().await;
|
||||
let tools = json!([{"type": "function", "function": {
|
||||
"name": "get_weather",
|
||||
"description": "Weather for a city",
|
||||
"parameters": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}
|
||||
}}]);
|
||||
let question = json!({"role": "user", "content": "Weather in Paris?"});
|
||||
let (status, body) = send(
|
||||
&harness.app,
|
||||
json!({"model": "claude", "messages": [question], "tools": tools, "tool_choice": "auto", "max_tokens": 64}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
let answer: Value = serde_json::from_str(&body).unwrap();
|
||||
assert_eq!(answer["choices"][0]["finish_reason"], "tool_calls");
|
||||
let call = &answer["choices"][0]["message"]["tool_calls"][0];
|
||||
assert_eq!(call["id"], "toolu_1");
|
||||
assert_eq!(call["type"], "function");
|
||||
assert_eq!(call["function"]["name"], "get_weather");
|
||||
assert_eq!(call["function"]["arguments"], "{\"city\": \"Paris\"}");
|
||||
|
||||
let first = harness.calls.lock().unwrap()[0].body.clone();
|
||||
assert_eq!(
|
||||
first,
|
||||
json!({
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": "Weather in Paris?"}]}],
|
||||
"tools": [{
|
||||
"name": "get_weather",
|
||||
"input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
|
||||
"type": "custom",
|
||||
"description": "Weather for a city"
|
||||
}],
|
||||
"tool_choice": {"type": "auto"},
|
||||
"max_tokens": 64,
|
||||
"anthropic_version": "vertex-2023-10-16",
|
||||
})
|
||||
);
|
||||
|
||||
let (status, body) = send(
|
||||
&harness.app,
|
||||
json!({"model": "claude", "tools": tools, "max_tokens": 64, "messages": [
|
||||
question,
|
||||
{"role": "assistant", "content": null, "tool_calls": [call]},
|
||||
{"role": "tool", "tool_call_id": "toolu_1", "content": "sunny"}
|
||||
]}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
let answer: Value = serde_json::from_str(&body).unwrap();
|
||||
assert_eq!(
|
||||
answer["choices"][0]["message"]["content"],
|
||||
"It is sunny in Paris"
|
||||
);
|
||||
let second = harness.calls.lock().unwrap()[1].body.clone();
|
||||
assert_eq!(
|
||||
second["messages"],
|
||||
json!([
|
||||
{"role": "user", "content": [{"type": "text", "text": "Weather in Paris?"}]},
|
||||
{"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Paris"}}]},
|
||||
{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": "sunny"}]}
|
||||
])
|
||||
);
|
||||
|
||||
let seen = seen(&harness).await;
|
||||
assert_eq!(seen.len(), 2, "{seen:?}");
|
||||
assert!(seen.iter().all(|seen| seen.cost == Some(cost(10, 4))));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn reasoning_effort_becomes_a_thinking_budget_and_comes_back_as_reasoning() {
|
||||
let harness = harness().await;
|
||||
let (status, body) = send(
|
||||
&harness.app,
|
||||
json!({"model": "claude", "messages": [{"role": "user", "content": "hi"}], "reasoning_effort": "low"}),
|
||||
)
|
||||
.await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
let answer: Value = serde_json::from_str(&body).unwrap();
|
||||
let message = &answer["choices"][0]["message"];
|
||||
assert_eq!(message["content"], "Hello");
|
||||
assert_eq!(message["reasoning_content"], "The user says hi.");
|
||||
assert_eq!(
|
||||
message["thinking_blocks"],
|
||||
json!([{"type": "thinking", "thinking": "The user says hi.", "signature": "sig-1"}])
|
||||
);
|
||||
|
||||
let sent = harness.calls.lock().unwrap()[0].body.clone();
|
||||
assert_eq!(
|
||||
sent,
|
||||
json!({
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}],
|
||||
"thinking": {"type": "enabled", "budget_tokens": 1024},
|
||||
"max_tokens": 1024 + 4096,
|
||||
"anthropic_version": "vertex-2023-10-16",
|
||||
})
|
||||
);
|
||||
let seen = seen(&harness).await;
|
||||
assert_eq!(seen[0].cost, Some(cost(10, 4)));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn an_anthropic_api_only_param_is_refused() {
|
||||
let harness = harness().await;
|
||||
let mut request = hi("claude");
|
||||
request["service_tier"] = json!("auto");
|
||||
let (status, body) = send(&harness.app, request).await;
|
||||
assert_eq!(status, 400, "{body}");
|
||||
assert!(harness.calls.lock().unwrap().is_empty());
|
||||
}
|
||||
|
|
@ -6,11 +6,13 @@ license.workspace = true
|
|||
repository.workspace = true
|
||||
|
||||
[dependencies]
|
||||
base64.workspace = true
|
||||
bytes.workspace = true
|
||||
futures-util.workspace = true
|
||||
litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] }
|
||||
litellm-cache-response.workspace = true
|
||||
litellm-core-utils.workspace = true
|
||||
litellm-framer.workspace = true
|
||||
litellm-host.workspace = true
|
||||
litellm-http.workspace = true
|
||||
litellm-inference.workspace = true
|
||||
|
|
@ -20,17 +22,16 @@ litellm-secrets.workspace = true
|
|||
reqwest.workspace = true
|
||||
serde_json = { workspace = true, features = ["preserve_order"] }
|
||||
tokio = { workspace = true, features = ["time"] }
|
||||
tokio-util = { version = "0.7", features = ["codec"] }
|
||||
tracing.workspace = true
|
||||
uuid.workspace = true
|
||||
|
||||
[dev-dependencies]
|
||||
litellm-cache-memory.workspace = true
|
||||
litellm-framer.workspace = true
|
||||
litellm-host-native.workspace = true
|
||||
litellm-http = { workspace = true, features = ["test-support"] }
|
||||
litellm-inference-testing.workspace = true
|
||||
litellm-tracing.workspace = true
|
||||
rstest.workspace = true
|
||||
tokio-util = { version = "0.7", features = ["codec"] }
|
||||
tokio = { workspace = true, features = ["io-util", "sync"] }
|
||||
wiremock.workspace = true
|
||||
|
|
|
|||
|
|
@ -129,6 +129,7 @@ impl ChatCompletionsRoute {
|
|||
headers: authenticated.headers,
|
||||
body,
|
||||
api_key,
|
||||
adapter: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -4,7 +4,16 @@ use litellm_llms::{
|
|||
anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG,
|
||||
base_llm::chat::transformation::BaseConfig,
|
||||
bedrock::chat::converse_transformation::BEDROCK_CHAT_COMPLETIONS_CONFIG,
|
||||
gemini::chat::transformation::GOOGLE_AI_STUDIO_GEMINI_CHAT_COMPLETIONS_CONFIG,
|
||||
openai_like::chat::transformation::OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG,
|
||||
vertex_ai::{
|
||||
gemini::vertex_and_google_ai_studio_gemini::{
|
||||
VERTEX_GEMINI_CHAT_COMPLETIONS_CONFIG, is_vertex_gemini_model,
|
||||
},
|
||||
vertex_ai_partner_models::anthropic::transformation::{
|
||||
VERTEX_ANTHROPIC_CHAT_COMPLETIONS_CONFIG, is_vertex_claude_model,
|
||||
},
|
||||
},
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
|
|
@ -15,7 +24,10 @@ const HEADER_CONTEXT: &str = "chat completions";
|
|||
pub(super) enum ChatProvider {
|
||||
Anthropic,
|
||||
Bedrock,
|
||||
Gemini,
|
||||
OpenaiLike,
|
||||
VertexAnthropic,
|
||||
VertexGemini,
|
||||
}
|
||||
|
||||
impl ChatProvider {
|
||||
|
|
@ -23,16 +35,28 @@ impl ChatProvider {
|
|||
match self {
|
||||
Self::Anthropic => &ANTHROPIC_CHAT_COMPLETIONS_CONFIG,
|
||||
Self::Bedrock => &BEDROCK_CHAT_COMPLETIONS_CONFIG,
|
||||
Self::Gemini => &GOOGLE_AI_STUDIO_GEMINI_CHAT_COMPLETIONS_CONFIG,
|
||||
Self::OpenaiLike => &OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG,
|
||||
Self::VertexAnthropic => &VERTEX_ANTHROPIC_CHAT_COMPLETIONS_CONFIG,
|
||||
Self::VertexGemini => &VERTEX_GEMINI_CHAT_COMPLETIONS_CONFIG,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn chat_completions_provider(provider: LlmProviders) -> Option<ChatProvider> {
|
||||
/// Vertex AI serves its Gemini and Claude models here; its other partner models stay refused.
|
||||
pub(super) fn chat_completions_provider(
|
||||
provider: LlmProviders,
|
||||
model: &str,
|
||||
) -> Option<ChatProvider> {
|
||||
match provider {
|
||||
LlmProviders::Anthropic => Some(ChatProvider::Anthropic),
|
||||
LlmProviders::Bedrock => Some(ChatProvider::Bedrock),
|
||||
LlmProviders::Gemini => Some(ChatProvider::Gemini),
|
||||
LlmProviders::OpenaiLike => Some(ChatProvider::OpenaiLike),
|
||||
LlmProviders::VertexAi if is_vertex_gemini_model(model) => Some(ChatProvider::VertexGemini),
|
||||
LlmProviders::VertexAi if is_vertex_claude_model(model) => {
|
||||
Some(ChatProvider::VertexAnthropic)
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
|
@ -42,3 +66,35 @@ pub(super) fn string_headers(
|
|||
) -> Result<Vec<(String, String)>, Error> {
|
||||
shared_string_headers(HEADER_CONTEXT, extra_headers).map_err(Error::from)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[rstest]
|
||||
#[case::studio(LlmProviders::Gemini, "gemini-2.5-flash", Some("gemini"))]
|
||||
#[case::vertex_gemini(LlmProviders::VertexAi, "gemini-2.5-pro", Some("vertex_gemini"))]
|
||||
#[case::vertex_claude(
|
||||
LlmProviders::VertexAi,
|
||||
"claude-sonnet-4@20250514",
|
||||
Some("vertex_claude")
|
||||
)]
|
||||
#[case::vertex_partner(LlmProviders::VertexAi, "meta/llama-3.1-405b-instruct-maas", None)]
|
||||
#[case::vertex_mistral(LlmProviders::VertexAi, "mistral-large@2407", None)]
|
||||
#[case::vertex_gemma(LlmProviders::VertexAi, "gemma/gemma-3-27b-it", None)]
|
||||
fn vertex_routes_its_gemini_and_claude_models(
|
||||
#[case] provider: LlmProviders,
|
||||
#[case] model: &str,
|
||||
#[case] expected: Option<&str>,
|
||||
) {
|
||||
let routed = chat_completions_provider(provider, model).map(|provider| match provider {
|
||||
ChatProvider::Gemini => "gemini",
|
||||
ChatProvider::VertexGemini => "vertex_gemini",
|
||||
ChatProvider::VertexAnthropic => "vertex_claude",
|
||||
ChatProvider::Anthropic | ChatProvider::Bedrock | ChatProvider::OpenaiLike => "other",
|
||||
});
|
||||
assert_eq!(routed, expected);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ pub(super) async fn execute(
|
|||
config,
|
||||
url,
|
||||
body,
|
||||
stream_shape: _,
|
||||
stream_shape,
|
||||
optional_params,
|
||||
environment,
|
||||
secrets,
|
||||
|
|
@ -122,7 +122,7 @@ pub(super) async fn execute(
|
|||
))
|
||||
})?;
|
||||
config
|
||||
.transform_response(&model, ProviderChatResponseData { body })
|
||||
.transform_response(&model, ProviderChatResponseData { body }, &stream_shape)
|
||||
.map_err(Error::from)
|
||||
.map_err(as_response_error)
|
||||
},
|
||||
|
|
@ -229,6 +229,8 @@ mod tests {
|
|||
custom_llm_provider: None,
|
||||
extra_headers: None,
|
||||
timeout: None,
|
||||
drop_params: Some(false),
|
||||
model_info: None,
|
||||
})
|
||||
.unwrap(),
|
||||
std::sync::Arc::new(|_: &str| None),
|
||||
|
|
|
|||
|
|
@ -8,7 +8,9 @@ pub use litellm_inference::RouteError as Error;
|
|||
mod common_utils;
|
||||
pub mod constants;
|
||||
pub(crate) mod handler;
|
||||
mod media;
|
||||
pub mod openai;
|
||||
pub mod openai_compatible;
|
||||
mod prepare;
|
||||
pub mod stream;
|
||||
use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse;
|
||||
|
|
@ -25,6 +27,7 @@ pub struct ChatCompletionsRoute {
|
|||
auth: Arc<AuthServices>,
|
||||
secrets: Arc<dyn SecretSource>,
|
||||
cache: Option<litellm_cache_response::ScopedCache>,
|
||||
media: Option<litellm_http::media::MediaFetcher>,
|
||||
}
|
||||
|
||||
impl ChatCompletionsRoute {
|
||||
|
|
@ -38,6 +41,16 @@ impl ChatCompletionsRoute {
|
|||
auth,
|
||||
secrets,
|
||||
cache: None,
|
||||
media: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Downloads the media URLs a provider cannot take by reference. Without one such a
|
||||
/// request is refused as unsupported.
|
||||
pub fn with_media(self, media: litellm_http::media::MediaFetcher) -> Self {
|
||||
Self {
|
||||
media: Some(media),
|
||||
..self
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -77,7 +90,8 @@ impl ChatCompletionsRoute {
|
|||
interceptors: &impl litellm_host::interceptors::Interceptors<Error>,
|
||||
observers: Option<&ObservationSender>,
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
let resolved = resolve_request(request)?;
|
||||
let mut resolved = resolve_request(request)?;
|
||||
self.inline_remote_media(&mut resolved).await?;
|
||||
let snapshot = self
|
||||
.secrets
|
||||
.resolve(&resolved.config.secret_names())
|
||||
|
|
|
|||
158
litellm-rust/crates/inference-chat/src/media.rs
Normal file
158
litellm-rust/crates/inference-chat/src/media.rs
Normal file
|
|
@ -0,0 +1,158 @@
|
|||
//! The media a translated request needs downloaded first: an image or file URL the provider
|
||||
//! cannot take by reference, which Python fetches while it builds the request.
|
||||
|
||||
use std::{collections::HashMap, time::Duration};
|
||||
|
||||
use base64::Engine;
|
||||
use futures_util::future::try_join_all;
|
||||
use litellm_http::media::DownloadPolicy;
|
||||
use litellm_llms::base_llm::chat::media::{
|
||||
MAX_MEDIA_DOWNLOAD_BYTES, downloaded_media_type, inline_media,
|
||||
};
|
||||
|
||||
use crate::{ChatCompletionsRoute, Error, types::ResolvedChatCompletionsRequest};
|
||||
|
||||
const MEDIA_DOWNLOAD: DownloadPolicy = DownloadPolicy {
|
||||
timeout: Duration::from_secs(60),
|
||||
max_bytes: MAX_MEDIA_DOWNLOAD_BYTES,
|
||||
max_redirects: 10,
|
||||
};
|
||||
|
||||
impl ChatCompletionsRoute {
|
||||
/// Downloads every URL the config names and inlines each as a data URL. A route with no
|
||||
/// fetcher leaves such a request to the host; a URL that cannot be fetched is the
|
||||
/// caller's error, as Python's image fetch failure is.
|
||||
pub(crate) async fn inline_remote_media(
|
||||
&self,
|
||||
request: &mut ResolvedChatCompletionsRequest<'_>,
|
||||
) -> Result<(), Error> {
|
||||
let urls = request.config.remote_media(&request.messages);
|
||||
if urls.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let fetcher = self
|
||||
.media
|
||||
.as_ref()
|
||||
.ok_or(Error::Unsupported("message media URL"))?;
|
||||
let fetch = |url: String| async move {
|
||||
let unable = |reason: String| {
|
||||
Error::InvalidRequest(
|
||||
format!("Unable to fetch media from url {url}: {reason}").into(),
|
||||
)
|
||||
};
|
||||
let parsed = reqwest::Url::parse(&url).map_err(|err| unable(err.to_string()))?;
|
||||
let media = fetcher
|
||||
.fetch(parsed, MEDIA_DOWNLOAD)
|
||||
.await
|
||||
.map_err(|err| unable(err.to_string()))?;
|
||||
let data = format!(
|
||||
"data:{};base64,{}",
|
||||
downloaded_media_type(&url, &media.content_type),
|
||||
base64::engine::general_purpose::STANDARD.encode(media.bytes)
|
||||
);
|
||||
Ok::<_, Error>((url, data))
|
||||
};
|
||||
let downloaded: HashMap<String, String> = try_join_all(urls.into_iter().map(fetch))
|
||||
.await?
|
||||
.into_iter()
|
||||
.collect();
|
||||
inline_media(&mut request.messages, &downloaded);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use litellm_http::{Client, media::MediaFetcher};
|
||||
|
||||
use rstest::rstest;
|
||||
use serde_json::{Value, json};
|
||||
use wiremock::{
|
||||
Mock, MockServer, ResponseTemplate,
|
||||
matchers::{method, path},
|
||||
};
|
||||
|
||||
use crate::{
|
||||
ChatCompletionsRoute, Error, prepare::resolve_request, types::ChatCompletionsRequest,
|
||||
};
|
||||
|
||||
fn route(media: bool) -> ChatCompletionsRoute {
|
||||
let route = crate::test_support::route();
|
||||
match media {
|
||||
true => route.with_media(MediaFetcher::for_test(Client::plain_for_test())),
|
||||
false => route,
|
||||
}
|
||||
}
|
||||
|
||||
fn request(url: &str) -> ChatCompletionsRequest<'static> {
|
||||
ChatCompletionsRequest {
|
||||
model: "anthropic/claude-sonnet-4-5",
|
||||
messages: json!([{"role": "user", "content": [
|
||||
{"type": "image_url", "image_url": {"url": url}}
|
||||
]}]),
|
||||
optional_params: serde_json::Map::new(),
|
||||
api_key: Some("sk-test"),
|
||||
api_base: None,
|
||||
custom_llm_provider: None,
|
||||
extra_headers: None,
|
||||
timeout: None,
|
||||
drop_params: Some(false),
|
||||
model_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn image(request: &crate::types::ResolvedChatCompletionsRequest<'_>) -> Value {
|
||||
serde_json::to_value(&request.messages[0].content).unwrap()[0]["image_url"]["url"].clone()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_remote_image_is_downloaded_and_inlined() {
|
||||
let server = MockServer::start().await;
|
||||
Mock::given(method("GET"))
|
||||
.and(path("/cat.png"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_raw(b"png".to_vec(), "image/png"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
let url = format!("{}/cat.png", server.uri());
|
||||
let url: &'static str = Box::leak(url.into_boxed_str());
|
||||
let mut resolved = resolve_request(request(url)).unwrap();
|
||||
|
||||
route(true)
|
||||
.inline_remote_media(&mut resolved)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(image(&resolved), "data:image/png;base64,cG5n");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_media_url_that_cannot_be_fetched_is_the_callers_error() {
|
||||
let server = MockServer::start().await;
|
||||
let url = format!("{}/missing.png", server.uri());
|
||||
let url: &'static str = Box::leak(url.into_boxed_str());
|
||||
let mut resolved = resolve_request(request(url)).unwrap();
|
||||
|
||||
let error = route(true)
|
||||
.inline_remote_media(&mut resolved)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert!(matches!(error, Error::InvalidRequest(_)), "{error:?}");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn without_a_fetcher_a_media_url_is_unsupported_and_a_data_url_is_left_alone() {
|
||||
let mut remote = resolve_request(request("http://example.com/cat.png")).unwrap();
|
||||
assert!(matches!(
|
||||
route(false).inline_remote_media(&mut remote).await,
|
||||
Err(Error::Unsupported(_))
|
||||
));
|
||||
|
||||
let mut inline = resolve_request(request("data:image/png;base64,cG5n")).unwrap();
|
||||
route(false).inline_remote_media(&mut inline).await.unwrap();
|
||||
assert_eq!(image(&inline), "data:image/png;base64,cG5n");
|
||||
}
|
||||
}
|
||||
|
|
@ -24,8 +24,11 @@ use litellm_llms::{
|
|||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{
|
||||
ChatCompletionsRoute, Error, common_utils::string_headers,
|
||||
ChatCompletionsRoute, Error,
|
||||
common_utils::string_headers,
|
||||
constants::CHAT_COMPLETIONS_TIMEOUT_SECS,
|
||||
openai_compatible::{ResponseAdapter, compatible_chat_config},
|
||||
stream::ReportedUsage,
|
||||
};
|
||||
|
||||
pub struct OpenAiChatCompletions;
|
||||
|
|
@ -61,6 +64,9 @@ pub struct OpenAiChatCall {
|
|||
pub drop_params: bool,
|
||||
/// The model cost map, which Azure's param rules read.
|
||||
pub model_info: Option<Arc<dyn ModelInfo>>,
|
||||
/// Where a stream's usage is recorded for the host when the client did not ask to see
|
||||
/// it.
|
||||
pub usage: ReportedUsage,
|
||||
}
|
||||
|
||||
/// A provider request ready to go out: where, with which headers and body, and who serves it.
|
||||
|
|
@ -71,17 +77,18 @@ pub(crate) struct PreparedCall {
|
|||
pub(crate) headers: Headers,
|
||||
pub(crate) body: Map<String, Value>,
|
||||
pub(crate) api_key: Option<String>,
|
||||
pub(crate) adapter: Option<ResponseAdapter>,
|
||||
}
|
||||
|
||||
impl ChatCompletionsRoute {
|
||||
/// Whether the deployment is served by the OpenAI wire passthrough: OpenAI itself, or
|
||||
/// Azure OpenAI.
|
||||
/// Whether the deployment is served by the OpenAI wire passthrough: OpenAI itself, Azure
|
||||
/// OpenAI, or an OpenAI-compatible provider.
|
||||
pub fn serves_openai(model: &str, custom_llm_provider: Option<&str>) -> bool {
|
||||
resolve_llm_provider(model, custom_llm_provider, "chat completions").is_ok_and(|resolved| {
|
||||
matches!(
|
||||
resolved.provider,
|
||||
LlmProviders::Openai | LlmProviders::Azure
|
||||
)
|
||||
) || compatible_chat_config(resolved.provider).is_some()
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -149,10 +156,14 @@ impl ChatCompletionsRoute {
|
|||
"chat completions",
|
||||
)?;
|
||||
let model = resolved.model.to_string();
|
||||
let prepared = match resolved.provider {
|
||||
LlmProviders::Openai => self.prepare_openai(model, call).await?,
|
||||
LlmProviders::Azure => self.prepare_azure(model, call).await?,
|
||||
other => return Err(Error::InvalidProvider(<&str>::from(other).to_string())),
|
||||
let usage = call.usage.clone();
|
||||
let prepared = match (resolved.provider, compatible_chat_config(resolved.provider)) {
|
||||
(LlmProviders::Openai, _) => self.prepare_openai(model, call).await?,
|
||||
(LlmProviders::Azure, _) => self.prepare_azure(model, call).await?,
|
||||
(_, Some(config)) => self.prepare_openai_compatible(config, model, call).await?,
|
||||
(other, None) => {
|
||||
return Err(Error::InvalidProvider(<&str>::from(other).to_string()));
|
||||
}
|
||||
};
|
||||
let PreparedCall {
|
||||
model,
|
||||
|
|
@ -161,6 +172,7 @@ impl ChatCompletionsRoute {
|
|||
headers,
|
||||
body,
|
||||
api_key,
|
||||
adapter,
|
||||
} = prepared;
|
||||
litellm_inference::diagnostic::provider(&model, provider);
|
||||
let request_context = RequestContext {
|
||||
|
|
@ -192,12 +204,20 @@ impl ChatCompletionsRoute {
|
|||
.unwrap_or(false);
|
||||
let response = send(&self.http, wire, stream, timeout).await?;
|
||||
let output = if stream {
|
||||
let chunks = chunks(response, timeout);
|
||||
CallOutput::Stream {
|
||||
head: (),
|
||||
chunks: chunks(response, timeout),
|
||||
chunks: match adapter {
|
||||
Some(adapter) => adapter.stream(chunks, usage),
|
||||
None => chunks,
|
||||
},
|
||||
}
|
||||
} else {
|
||||
let body = response.bytes().await.map_err(network)?;
|
||||
let body = match &adapter {
|
||||
Some(adapter) => adapter.complete(body),
|
||||
None => body,
|
||||
};
|
||||
context
|
||||
.response_received(&String::from_utf8_lossy(&body))
|
||||
.await?;
|
||||
|
|
@ -246,6 +266,7 @@ impl ChatCompletionsRoute {
|
|||
headers: authenticated.headers,
|
||||
body,
|
||||
api_key: call.api_key,
|
||||
adapter: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
@ -354,6 +375,9 @@ mod tests {
|
|||
#[case::azure("azure/gpt-4o", None, true)]
|
||||
#[case::azure_ai("azure_ai/gpt-4o", None, false)]
|
||||
#[case::openai_like("openai_like/gpt-6-astra", None, false)]
|
||||
#[case::groq("groq/llama-3.1-8b-instant", None, true)]
|
||||
#[case::hosted_vllm("meta-llama/Llama-3.1-8B", Some("hosted_vllm"), true)]
|
||||
#[case::nvidia_nim("nvidia_nim/meta/llama3-70b-instruct", None, true)]
|
||||
#[case::anthropic("anthropic/claude", None, false)]
|
||||
#[case::unprefixed("gpt-6-astra", None, false)]
|
||||
fn only_openai_deployments_take_the_passthrough(
|
||||
|
|
|
|||
416
litellm-rust/crates/inference-chat/src/openai_compatible.rs
Normal file
416
litellm-rust/crates/inference-chat/src/openai_compatible.rs
Normal file
|
|
@ -0,0 +1,416 @@
|
|||
//! The OpenAI-compatible providers (`litellm/llms/<provider>/chat/transformation.py`) through
|
||||
//! the OpenAI passthrough: each provider's config shapes the request, and the response goes
|
||||
//! back as the provider sent it but for the fixups its config makes. A stream is relayed
|
||||
//! event by event, with the usage the call asked for on the client's behalf recorded for the
|
||||
//! host and kept from the client.
|
||||
|
||||
use bytes::{Bytes, BytesMut};
|
||||
use futures_util::{StreamExt, stream::BoxStream};
|
||||
use litellm_auth::{CredentialPlacement, SecretValue};
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use litellm_framer::sse::{SseBlock, SseBlockCodec};
|
||||
use litellm_llms::{
|
||||
ErrorDetail,
|
||||
base_llm::{
|
||||
auth::{AuthScheme, ValidatedEnvironment, resolve_auth},
|
||||
sse::DONE_FRAME,
|
||||
},
|
||||
cerebras::chat::transformation::CEREBRAS_CHAT_CONFIG,
|
||||
deepseek::chat::transformation::DEEPSEEK_CHAT_CONFIG,
|
||||
fireworks_ai::chat::transformation::FIREWORKS_AI_CHAT_CONFIG,
|
||||
groq::chat::transformation::GROQ_CHAT_CONFIG,
|
||||
hosted_vllm::chat::transformation::HOSTED_VLLM_CHAT_CONFIG,
|
||||
mistral::chat::transformation::MISTRAL_CHAT_CONFIG,
|
||||
nvidia_nim::chat::transformation::NVIDIA_NIM_CHAT_CONFIG,
|
||||
openai::chat::compatible::{
|
||||
CompatibleChatBody, CompatibleContext, OpenAiCompatibleChatConfig, ResponseShape,
|
||||
fix_response, transform_request,
|
||||
},
|
||||
openrouter::chat::transformation::OPENROUTER_CHAT_CONFIG,
|
||||
perplexity::chat::transformation::PERPLEXITY_CHAT_CONFIG,
|
||||
together_ai::chat::transformation::TOGETHER_AI_CHAT_CONFIG,
|
||||
xai::chat::transformation::XAI_CHAT_CONFIG,
|
||||
};
|
||||
use litellm_llms_types::formats::chat_completions::{
|
||||
ChatCompletionsUsage, CompletionTokensDetails, PromptTokensDetails,
|
||||
};
|
||||
use serde_json::Value;
|
||||
use tokio_util::codec::Decoder;
|
||||
|
||||
use crate::{
|
||||
ChatCompletionsRoute, Error,
|
||||
common_utils::string_headers,
|
||||
openai::{OpenAiChatCall, PreparedCall},
|
||||
stream::ReportedUsage,
|
||||
};
|
||||
|
||||
const ENV_REFERENCE: &str = "os.environ/";
|
||||
|
||||
/// The config of a provider served through the OpenAI passthrough with its own transforms.
|
||||
pub fn compatible_chat_config(
|
||||
provider: LlmProviders,
|
||||
) -> Option<&'static dyn OpenAiCompatibleChatConfig> {
|
||||
Some(match provider {
|
||||
LlmProviders::Cerebras => &CEREBRAS_CHAT_CONFIG,
|
||||
LlmProviders::Deepseek => &DEEPSEEK_CHAT_CONFIG,
|
||||
LlmProviders::FireworksAi => &FIREWORKS_AI_CHAT_CONFIG,
|
||||
LlmProviders::Groq => &GROQ_CHAT_CONFIG,
|
||||
LlmProviders::HostedVllm => &HOSTED_VLLM_CHAT_CONFIG,
|
||||
LlmProviders::Mistral => &MISTRAL_CHAT_CONFIG,
|
||||
LlmProviders::NvidiaNim => &NVIDIA_NIM_CHAT_CONFIG,
|
||||
LlmProviders::Openrouter => &OPENROUTER_CHAT_CONFIG,
|
||||
LlmProviders::Perplexity => &PERPLEXITY_CHAT_CONFIG,
|
||||
LlmProviders::TogetherAi => &TOGETHER_AI_CHAT_CONFIG,
|
||||
LlmProviders::Xai => &XAI_CHAT_CONFIG,
|
||||
_ => return None,
|
||||
})
|
||||
}
|
||||
|
||||
/// What the response of a call to an OpenAI-compatible provider goes through on its way back.
|
||||
#[derive(Clone)]
|
||||
pub(crate) struct ResponseAdapter {
|
||||
config: &'static dyn OpenAiCompatibleChatConfig,
|
||||
shape: ResponseShape,
|
||||
}
|
||||
|
||||
impl ChatCompletionsRoute {
|
||||
pub(crate) async fn prepare_openai_compatible(
|
||||
&self,
|
||||
config: &'static dyn OpenAiCompatibleChatConfig,
|
||||
model: String,
|
||||
call: OpenAiChatCall,
|
||||
) -> Result<PreparedCall, Error> {
|
||||
let OpenAiChatCall {
|
||||
body,
|
||||
api_key,
|
||||
api_base,
|
||||
extra_headers,
|
||||
drop_params,
|
||||
model_info,
|
||||
..
|
||||
} = call;
|
||||
let references: Vec<&str> = [api_key.as_deref(), api_base.as_deref()]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(|value| value.strip_prefix(ENV_REFERENCE))
|
||||
.collect();
|
||||
let names: Vec<&str> = config
|
||||
.secret_names()
|
||||
.into_iter()
|
||||
.chain(references)
|
||||
.collect();
|
||||
let secrets = self.secrets.resolve(&names).await?;
|
||||
let lookup = |name: &str| secrets.get(name);
|
||||
// Python's router reads every `os.environ/` setting from the environment when it
|
||||
// loads the deployment. A reference to an unset variable counts as unset.
|
||||
let resolve = |value: &str| match value.strip_prefix(ENV_REFERENCE) {
|
||||
Some(name) => lookup(name),
|
||||
None => Some(value.to_owned()),
|
||||
};
|
||||
let api_key = api_key.as_deref().and_then(resolve);
|
||||
let api_base = api_base.as_deref().and_then(resolve);
|
||||
let url = config.complete_url(&config.api_base(api_base.as_deref(), &lookup)?);
|
||||
let key = config.api_key(api_key.as_deref(), &lookup)?;
|
||||
let provider_headers = config.headers(&body, &lookup);
|
||||
let context = CompatibleContext {
|
||||
model: &model,
|
||||
drop_params,
|
||||
model_info: model_info.as_deref(),
|
||||
};
|
||||
let CompatibleChatBody { body, shape } = transform_request(config, &context, body)?;
|
||||
let defaults: Vec<(&str, &str)> = provider_headers
|
||||
.iter()
|
||||
.map(|(name, value)| (name.as_str(), value.as_str()))
|
||||
.chain([("content-type", "application/json")])
|
||||
.collect();
|
||||
let environment = ValidatedEnvironment {
|
||||
headers: litellm_http::request::with_default_headers(
|
||||
string_headers(extra_headers)?,
|
||||
&defaults,
|
||||
),
|
||||
auth: AuthScheme::Credential {
|
||||
placement: CredentialPlacement::Bearer,
|
||||
secret: SecretValue::new(key),
|
||||
},
|
||||
};
|
||||
let authenticated = resolve_auth(&self.auth, environment, &lookup).await?;
|
||||
Ok(PreparedCall {
|
||||
model,
|
||||
provider: config.name(),
|
||||
url,
|
||||
headers: authenticated.headers,
|
||||
body,
|
||||
api_key,
|
||||
adapter: Some(ResponseAdapter { config, shape }),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl ResponseAdapter {
|
||||
/// A complete response, re-encoded only when a fixup changed it.
|
||||
pub(crate) fn complete(&self, body: Bytes) -> Bytes {
|
||||
let Ok(Value::Object(mut response)) = serde_json::from_slice::<Value>(&body) else {
|
||||
return body;
|
||||
};
|
||||
match fix_response(self.config, &mut response, &self.shape) {
|
||||
true => serde_json::to_vec(&response).map_or(body, Bytes::from),
|
||||
false => body,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn stream(
|
||||
self,
|
||||
upstream: BoxStream<'static, Result<Bytes, Error>>,
|
||||
usage: ReportedUsage,
|
||||
) -> BoxStream<'static, Result<Bytes, Error>> {
|
||||
let relay = Relay {
|
||||
upstream,
|
||||
buffer: BytesMut::new(),
|
||||
codec: SseBlockCodec::default(),
|
||||
adapter: self,
|
||||
usage,
|
||||
ended: false,
|
||||
};
|
||||
futures_util::stream::try_unfold(relay, |mut relay| async move {
|
||||
relay
|
||||
.next_frame()
|
||||
.await
|
||||
.map(|frame| frame.map(|frame| (frame, relay)))
|
||||
})
|
||||
.boxed()
|
||||
}
|
||||
|
||||
/// One event's frame for the client, or `None` when the client is not to see it.
|
||||
fn frame(&self, block: SseBlock, usage: &ReportedUsage) -> Option<Bytes> {
|
||||
if block.event.data.trim() == "[DONE]" {
|
||||
return Some(Bytes::from_static(DONE_FRAME));
|
||||
}
|
||||
let Ok(Value::Object(mut chunk)) = serde_json::from_str::<Value>(&block.event.data) else {
|
||||
return Some(block.bytes);
|
||||
};
|
||||
let mut changed = self.config.fix_chunk(&mut chunk, &self.shape);
|
||||
if let Some(reported) = chunk.get("usage").and_then(usage_from_json) {
|
||||
usage.set(reported);
|
||||
}
|
||||
if !self.shape.client_usage {
|
||||
if chunk.get("usage").is_some_and(|usage| !usage.is_null()) {
|
||||
chunk.remove("usage");
|
||||
changed = true;
|
||||
}
|
||||
if chunk
|
||||
.get("choices")
|
||||
.and_then(Value::as_array)
|
||||
.is_some_and(Vec::is_empty)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
}
|
||||
if !changed {
|
||||
return Some(block.bytes);
|
||||
}
|
||||
let mut frame = b"data: ".to_vec();
|
||||
serde_json::to_writer(&mut frame, &chunk).ok()?;
|
||||
frame.extend_from_slice(b"\n\n");
|
||||
Some(Bytes::from(frame))
|
||||
}
|
||||
}
|
||||
|
||||
struct Relay {
|
||||
upstream: BoxStream<'static, Result<Bytes, Error>>,
|
||||
buffer: BytesMut,
|
||||
codec: SseBlockCodec,
|
||||
adapter: ResponseAdapter,
|
||||
usage: ReportedUsage,
|
||||
ended: bool,
|
||||
}
|
||||
|
||||
impl Relay {
|
||||
async fn next_frame(&mut self) -> Result<Option<Bytes>, Error> {
|
||||
loop {
|
||||
let block = match self.ended {
|
||||
false => self.codec.decode(&mut self.buffer),
|
||||
true => self.codec.decode_eof(&mut self.buffer),
|
||||
}
|
||||
.map_err(|error| {
|
||||
Error::InvalidResponse(ErrorDetail::from(format!(
|
||||
"the provider sent a malformed stream: {error}"
|
||||
)))
|
||||
})?;
|
||||
match block {
|
||||
Some(block) => {
|
||||
if let Some(frame) = self.adapter.frame(block, &self.usage) {
|
||||
return Ok(Some(frame));
|
||||
}
|
||||
}
|
||||
None if self.ended => return Ok(None),
|
||||
None => match self.upstream.next().await {
|
||||
Some(bytes) => self.buffer.extend_from_slice(&bytes?),
|
||||
None => self.ended = true,
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The usage a chunk or response reports, read as leniently as Python's `Usage` reads it.
|
||||
fn usage_from_json(usage: &Value) -> Option<ChatCompletionsUsage> {
|
||||
let usage = usage.as_object()?;
|
||||
let count = |value: Option<&Value>| value.and_then(Value::as_u64);
|
||||
let prompt = usage.get("prompt_tokens_details");
|
||||
let completion = usage.get("completion_tokens_details");
|
||||
Some(ChatCompletionsUsage {
|
||||
prompt_tokens: count(usage.get("prompt_tokens")).unwrap_or(0),
|
||||
completion_tokens: count(usage.get("completion_tokens")).unwrap_or(0),
|
||||
total_tokens: count(usage.get("total_tokens")).unwrap_or(0),
|
||||
prompt_tokens_details: PromptTokensDetails {
|
||||
cached_tokens: count(prompt.and_then(|d| d.get("cached_tokens"))).unwrap_or(0),
|
||||
cache_creation_tokens: count(prompt.and_then(|d| d.get("cache_creation_tokens")))
|
||||
.unwrap_or(0),
|
||||
text_tokens: count(prompt.and_then(|d| d.get("text_tokens"))).unwrap_or(0),
|
||||
cache_write_tokens: None,
|
||||
cache_creation_token_details: prompt
|
||||
.and_then(|d| d.get("cache_creation_token_details"))
|
||||
.filter(|details| details.is_object())
|
||||
.cloned(),
|
||||
},
|
||||
completion_tokens_details: completion.filter(|details| details.is_object()).map(
|
||||
|details| CompletionTokensDetails {
|
||||
reasoning_tokens: count(details.get("reasoning_tokens")),
|
||||
text_tokens: count(details.get("text_tokens")),
|
||||
},
|
||||
),
|
||||
server_tool_use: usage
|
||||
.get("server_tool_use")
|
||||
.filter(|tools| tools.is_object())
|
||||
.cloned(),
|
||||
cache_creation_input_tokens: None,
|
||||
cache_read_input_tokens: None,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn adapter(
|
||||
config: &'static dyn OpenAiCompatibleChatConfig,
|
||||
client_usage: bool,
|
||||
) -> ResponseAdapter {
|
||||
ResponseAdapter {
|
||||
config,
|
||||
shape: ResponseShape {
|
||||
stream: true,
|
||||
client_usage,
|
||||
..ResponseShape::default()
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
async fn relay(adapter: ResponseAdapter, parts: &[&'static str]) -> (String, ReportedUsage) {
|
||||
let usage = ReportedUsage::default();
|
||||
let upstream = futures_util::stream::iter(
|
||||
parts
|
||||
.iter()
|
||||
.map(|part| Ok::<_, Error>(Bytes::from_static(part.as_bytes())))
|
||||
.collect::<Vec<_>>(),
|
||||
)
|
||||
.boxed();
|
||||
let frames: Vec<Bytes> = adapter
|
||||
.stream(upstream, usage.clone())
|
||||
.map(Result::unwrap)
|
||||
.collect()
|
||||
.await;
|
||||
(String::from_utf8(frames.concat()).unwrap(), usage)
|
||||
}
|
||||
|
||||
const CONTENT: &str =
|
||||
"data: {\"id\":\"c\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hi\"}}]}\n\n";
|
||||
const USAGE: &str = "data: {\"id\":\"c\",\"choices\":[],\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":2,\"total_tokens\":5}}\n\n";
|
||||
const DONE: &str = "data: [DONE]\n\n";
|
||||
|
||||
#[rstest]
|
||||
#[case::kept_from_the_client(false, format!("{CONTENT}{DONE}"))]
|
||||
#[case::shown_when_asked(true, format!("{CONTENT}{USAGE}{DONE}"))]
|
||||
#[tokio::test]
|
||||
async fn usage_is_recorded_and_shown_only_when_asked(
|
||||
#[case] client_usage: bool,
|
||||
#[case] expected: String,
|
||||
) {
|
||||
let (sent, usage) = relay(
|
||||
adapter(&GROQ_CHAT_CONFIG, client_usage),
|
||||
&[
|
||||
"data: {\"id\":\"c\",\"choices\":[{\"index\":0,",
|
||||
"\"delta\":{\"content\":\"hi\"}}]}\n\n",
|
||||
USAGE,
|
||||
DONE,
|
||||
],
|
||||
)
|
||||
.await;
|
||||
assert_eq!(sent, expected);
|
||||
let usage = usage.get().unwrap();
|
||||
assert_eq!(
|
||||
(
|
||||
usage.prompt_tokens,
|
||||
usage.completion_tokens,
|
||||
usage.total_tokens
|
||||
),
|
||||
(3, 2, 5)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn chunk_fixups_reencode_only_the_chunks_they_change() {
|
||||
let (sent, _) = relay(
|
||||
adapter(&HOSTED_VLLM_CHAT_CONFIG, true),
|
||||
&[
|
||||
CONTENT,
|
||||
"data: {\"choices\":[{\"index\":0,\"delta\":{\"reasoning\":\"r\"}}]}\n\n",
|
||||
DONE,
|
||||
],
|
||||
)
|
||||
.await;
|
||||
assert_eq!(
|
||||
sent,
|
||||
format!(
|
||||
"{CONTENT}data: {{\"choices\":[{{\"index\":0,\"delta\":{{\"reasoning_content\":\"r\"}}}}]}}\n\n{DONE}"
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_usage_chunk_with_choices_loses_only_its_usage() {
|
||||
let (sent, usage) = relay(
|
||||
adapter(&MISTRAL_CHAT_CONFIG, false),
|
||||
&["data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":1,\"completion_tokens\":1,\"total_tokens\":2}}\n\n"],
|
||||
)
|
||||
.await;
|
||||
assert_eq!(
|
||||
sent,
|
||||
"data: {\"choices\":[{\"index\":0,\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n"
|
||||
);
|
||||
assert_eq!(usage.get().unwrap().total_tokens, 2);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_complete_response_is_reencoded_only_when_fixed() {
|
||||
let raw = Bytes::from_static(
|
||||
br#"{"choices": [{"message": {"role": "assistant", "content": "x"}}]}"#,
|
||||
);
|
||||
assert_eq!(
|
||||
adapter(&CEREBRAS_CHAT_CONFIG, false).complete(raw.clone()),
|
||||
raw
|
||||
);
|
||||
let fixed = adapter(&HOSTED_VLLM_CHAT_CONFIG, false).complete(Bytes::from_static(
|
||||
br#"{"choices":[{"message":{"role":"assistant","content":"x","reasoning":"r"}}]}"#,
|
||||
));
|
||||
assert_eq!(
|
||||
serde_json::from_slice::<Value>(&fixed).unwrap(),
|
||||
json!({"choices": [{"message": {"role": "assistant", "content": "x", "reasoning_content": "r"}}]})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,7 +1,10 @@
|
|||
use litellm_auth::SecretValue;
|
||||
use litellm_core_utils::settings::Lookup;
|
||||
use litellm_http::request::with_default_headers;
|
||||
use litellm_llms::base_llm::{auth::ValidatedEnvironment, chat::transformation::BaseConfig};
|
||||
use litellm_llms::base_llm::{
|
||||
auth::{Headers, ValidatedEnvironment},
|
||||
chat::{params::ChatRequestContext, transformation::BaseConfig},
|
||||
};
|
||||
use litellm_llms_types::formats::chat_completions::ChatMessage;
|
||||
use litellm_secrets::source::Secrets;
|
||||
use serde_json::Value;
|
||||
|
|
@ -26,7 +29,7 @@ pub(super) fn resolve_provider_config<'a>(
|
|||
custom_llm_provider: Option<&'a str>,
|
||||
) -> Result<ResolvedProvider, Error> {
|
||||
let provider_info = resolve_llm_provider(model, custom_llm_provider, "chat completions")?;
|
||||
let config = chat_completions_provider(provider_info.provider)
|
||||
let config = chat_completions_provider(provider_info.provider, provider_info.model)
|
||||
.ok_or_else(|| Error::InvalidProvider(<&str>::from(provider_info.provider).to_string()))?
|
||||
.config();
|
||||
Ok(ResolvedProvider {
|
||||
|
|
@ -72,9 +75,38 @@ pub(super) fn resolve_request(
|
|||
api_base: request.api_base,
|
||||
extra_headers: request.extra_headers,
|
||||
timeout: request.timeout,
|
||||
drop_params: request.drop_params,
|
||||
model_info: request.model_info,
|
||||
})
|
||||
}
|
||||
|
||||
/// The headers a transformation asked for, added to the ones already going out. A header
|
||||
/// already there (a forwarded `anthropic-beta`, say) keeps its values and gains the new ones,
|
||||
/// comma separated, the way Python unions beta headers.
|
||||
fn merge_headers(headers: &mut Headers, added: Headers) {
|
||||
for (name, value) in added {
|
||||
match headers
|
||||
.iter_mut()
|
||||
.find(|(existing, _)| existing.eq_ignore_ascii_case(&name))
|
||||
{
|
||||
Some((_, existing)) => {
|
||||
let mut values: Vec<&str> = existing
|
||||
.split(',')
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.collect();
|
||||
for value in value.split(',').map(str::trim) {
|
||||
if !value.is_empty() && !values.contains(&value) {
|
||||
values.push(value);
|
||||
}
|
||||
}
|
||||
*existing = values.join(", ");
|
||||
}
|
||||
None => headers.push((name, value)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn validate_environment(
|
||||
request: &ResolvedChatCompletionsRequest<'_>,
|
||||
model: &str,
|
||||
|
|
@ -101,7 +133,7 @@ pub(super) fn prepare_provider_request(
|
|||
request: ResolvedChatCompletionsRequest<'_>,
|
||||
secrets: Secrets,
|
||||
) -> Result<ProviderChatCompletionsRequest, Error> {
|
||||
let environment =
|
||||
let mut environment =
|
||||
validate_environment(&request, &request.model, request.config, secrets.as_ref())?;
|
||||
let model = request.model;
|
||||
let config = request.config;
|
||||
|
|
@ -112,8 +144,17 @@ pub(super) fn prepare_provider_request(
|
|||
&request.optional_params,
|
||||
&env_lookup,
|
||||
)?;
|
||||
let transformed =
|
||||
config.transform_request(&model, request.messages, request.optional_params.clone())?;
|
||||
let context = ChatRequestContext {
|
||||
drop_params: request.drop_params,
|
||||
model_info: request.model_info.as_deref(),
|
||||
};
|
||||
let transformed = config.transform_request(
|
||||
&model,
|
||||
request.messages,
|
||||
request.optional_params.clone(),
|
||||
&context,
|
||||
)?;
|
||||
merge_headers(&mut environment.headers, transformed.headers);
|
||||
|
||||
Ok(ProviderChatCompletionsRequest {
|
||||
model,
|
||||
|
|
@ -136,7 +177,7 @@ mod tests {
|
|||
use litellm_llms::base_llm::auth::{AuthScheme, resolve_auth};
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use super::{prepare_provider_request, resolve_request};
|
||||
use super::{merge_headers, prepare_provider_request, resolve_request};
|
||||
use crate::{
|
||||
Error,
|
||||
types::{ChatCompletionsRequest, ProviderChatCompletionsRequest},
|
||||
|
|
@ -183,6 +224,8 @@ mod tests {
|
|||
custom_llm_provider: provider,
|
||||
extra_headers: None,
|
||||
timeout: None,
|
||||
drop_params: Some(false),
|
||||
model_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -636,4 +679,27 @@ mod tests {
|
|||
"prepare did not carry the bearer token"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn transformation_headers_union_with_the_forwarded_ones() {
|
||||
let mut headers = vec![
|
||||
("Anthropic-Beta".to_owned(), "a,b".to_owned()),
|
||||
("x-other".to_owned(), "1".to_owned()),
|
||||
];
|
||||
merge_headers(
|
||||
&mut headers,
|
||||
vec![
|
||||
("anthropic-beta".to_owned(), "b, c".to_owned()),
|
||||
("x-new".to_owned(), "2".to_owned()),
|
||||
],
|
||||
);
|
||||
assert_eq!(
|
||||
headers,
|
||||
[
|
||||
("Anthropic-Beta".to_owned(), "a, b, c".to_owned()),
|
||||
("x-other".to_owned(), "1".to_owned()),
|
||||
("x-new".to_owned(), "2".to_owned()),
|
||||
]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -69,6 +69,8 @@ impl ChatCompletionsRoute {
|
|||
custom_llm_provider: call.custom_llm_provider.as_deref(),
|
||||
extra_headers: call.extra_headers,
|
||||
timeout: call.timeout,
|
||||
drop_params: call.drop_params,
|
||||
model_info: call.model_info,
|
||||
};
|
||||
self.run(request, cache_options, interceptors, observers)
|
||||
.await
|
||||
|
|
|
|||
|
|
@ -71,7 +71,7 @@ impl ReportedUsage {
|
|||
self.0.lock().unwrap_or_else(|e| e.into_inner()).clone()
|
||||
}
|
||||
|
||||
fn set(&self, usage: ChatCompletionsUsage) {
|
||||
pub(crate) fn set(&self, usage: ChatCompletionsUsage) {
|
||||
*self.0.lock().unwrap_or_else(|e| e.into_inner()) = Some(usage);
|
||||
}
|
||||
}
|
||||
|
|
@ -140,7 +140,7 @@ impl ChatCompletionsRoute {
|
|||
let timeout = call
|
||||
.timeout
|
||||
.unwrap_or(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS));
|
||||
let resolved = resolve_request(ChatCompletionsRequest {
|
||||
let mut resolved = resolve_request(ChatCompletionsRequest {
|
||||
model: &call.model,
|
||||
messages: call.messages,
|
||||
optional_params,
|
||||
|
|
@ -149,7 +149,10 @@ impl ChatCompletionsRoute {
|
|||
custom_llm_provider: call.custom_llm_provider.as_deref(),
|
||||
extra_headers: call.extra_headers,
|
||||
timeout: call.timeout,
|
||||
drop_params: call.drop_params,
|
||||
model_info: call.model_info,
|
||||
})?;
|
||||
self.inline_remote_media(&mut resolved).await?;
|
||||
let snapshot = self
|
||||
.secrets
|
||||
.resolve(&resolved.config.secret_names())
|
||||
|
|
@ -1019,6 +1022,8 @@ mod tests {
|
|||
custom_llm_provider: None,
|
||||
extra_headers: None,
|
||||
timeout: Some(Duration::from_secs(5)),
|
||||
drop_params: Some(false),
|
||||
model_info: None,
|
||||
},
|
||||
response_model: "public/model".into(),
|
||||
usage: ReportedUsage::default(),
|
||||
|
|
|
|||
|
|
@ -1,20 +1,20 @@
|
|||
use litellm_secrets::source::Secrets;
|
||||
use std::time::Duration;
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use litellm_auth::SecretValue;
|
||||
use litellm_llms::base_llm::{
|
||||
auth::ValidatedEnvironment,
|
||||
chat::{streaming::StreamShape, transformation::BaseConfig},
|
||||
model_info::ModelInfo,
|
||||
};
|
||||
use litellm_llms_types::formats::chat_completions::ChatMessage;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
/// A `/chat/completions` call as it crosses into the core.
|
||||
///
|
||||
/// `optional_params` arrives already mapped to the provider's own parameter
|
||||
/// names by the host, exactly as the messages route receives an already
|
||||
/// Anthropic-shaped body. The core owns the conversation translation, the
|
||||
/// provider call, and the response normalization.
|
||||
/// `optional_params` are the OpenAI request's own params; the provider config maps them,
|
||||
/// refusing what the provider cannot take unless `drop_params` says to drop it. The core
|
||||
/// owns the conversation translation, the provider call, and the response normalization.
|
||||
pub struct ChatCompletionsRequest<'a> {
|
||||
pub model: &'a str,
|
||||
pub messages: Value,
|
||||
|
|
@ -24,6 +24,11 @@ pub struct ChatCompletionsRequest<'a> {
|
|||
pub custom_llm_provider: Option<&'a str>,
|
||||
pub extra_headers: Option<Map<String, Value>>,
|
||||
pub timeout: Option<Duration>,
|
||||
/// `drop_params` from the deployment or the proxy settings, `None` when neither says. A
|
||||
/// request's own `drop_params` rides in `optional_params`.
|
||||
pub drop_params: Option<bool>,
|
||||
/// The model cost map, for the param rules that read model capabilities.
|
||||
pub model_info: Option<Arc<dyn ModelInfo>>,
|
||||
}
|
||||
|
||||
pub struct ChatCompletionsCall {
|
||||
|
|
@ -35,6 +40,11 @@ pub struct ChatCompletionsCall {
|
|||
pub custom_llm_provider: Option<String>,
|
||||
pub extra_headers: Option<Map<String, Value>>,
|
||||
pub timeout: Option<Duration>,
|
||||
/// `drop_params` from the deployment or the proxy settings, `None` when neither says. A
|
||||
/// request's own `drop_params` rides in `optional_params`.
|
||||
pub drop_params: Option<bool>,
|
||||
/// The model cost map, for the param rules that read model capabilities.
|
||||
pub model_info: Option<Arc<dyn ModelInfo>>,
|
||||
}
|
||||
|
||||
impl From<ChatCompletionsRequest<'_>> for ChatCompletionsCall {
|
||||
|
|
@ -48,6 +58,8 @@ impl From<ChatCompletionsRequest<'_>> for ChatCompletionsCall {
|
|||
custom_llm_provider: request.custom_llm_provider.map(str::to_owned),
|
||||
extra_headers: request.extra_headers,
|
||||
timeout: request.timeout,
|
||||
drop_params: request.drop_params,
|
||||
model_info: request.model_info,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -62,6 +74,11 @@ pub struct ResolvedChatCompletionsRequest<'a> {
|
|||
pub api_base: Option<&'a str>,
|
||||
pub extra_headers: Option<Map<String, Value>>,
|
||||
pub timeout: Option<Duration>,
|
||||
/// `drop_params` from the deployment or the proxy settings, `None` when neither says. A
|
||||
/// request's own `drop_params` rides in `optional_params`.
|
||||
pub drop_params: Option<bool>,
|
||||
/// The model cost map, for the param rules that read model capabilities.
|
||||
pub model_info: Option<Arc<dyn ModelInfo>>,
|
||||
}
|
||||
|
||||
pub struct ProviderChatCompletionsRequest {
|
||||
|
|
|
|||
|
|
@ -84,6 +84,8 @@ async fn the_same_route_entrypoint_reports_facts_with_or_without_caching(
|
|||
custom_llm_provider: None,
|
||||
extra_headers: None,
|
||||
timeout: None,
|
||||
drop_params: Some(false),
|
||||
model_info: None,
|
||||
},
|
||||
&(),
|
||||
Some(observer.clone()),
|
||||
|
|
@ -254,6 +256,8 @@ async fn chat_cache_identity_follows_resolved_configuration_and_request_callback
|
|||
custom_llm_provider: None,
|
||||
extra_headers: None,
|
||||
timeout: None,
|
||||
drop_params: Some(false),
|
||||
model_info: None,
|
||||
},
|
||||
&hooks,
|
||||
None,
|
||||
|
|
@ -319,6 +323,8 @@ async fn signed_requests_bypass_response_caching(cache: Arc<dyn ResponseCacheSer
|
|||
messages:json!([{"role":"user","content":"hello"}]),
|
||||
optional_params:json!({"aws_access_key_id":"test-access","aws_secret_access_key":"test-secret","aws_region_name":"eu-west-1"}).as_object().unwrap().clone(),
|
||||
api_key:None,api_base:Some(&upstream.uri()),custom_llm_provider:None,extra_headers:None,timeout:None,
|
||||
drop_params: Some(false),
|
||||
model_info: None,
|
||||
}, &hooks, None).await.unwrap();
|
||||
assert_eq!(
|
||||
serde_json::to_value(response).unwrap()["usage"]["total_tokens"],
|
||||
|
|
|
|||
|
|
@ -47,6 +47,8 @@ fn request() -> ChatCompletionsRequest<'static> {
|
|||
custom_llm_provider: None,
|
||||
extra_headers: None,
|
||||
timeout: Some(Duration::from_secs(10)),
|
||||
drop_params: Some(false),
|
||||
model_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -163,7 +165,7 @@ async fn bedrock_round_trip_is_signed_and_normalized(request: ChatCompletionsReq
|
|||
#[case::missing_usage(
|
||||
r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn"}"#
|
||||
)]
|
||||
#[case::tool_use_block(r#"{"model":"m","content":[{"type":"tool_use","id":"t","name":"f","input":{}}],"stop_reason":"tool_use","usage":{"input_tokens":1,"output_tokens":1}}"#)]
|
||||
#[case::malformed_token_count(r#"{"model":"m","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","usage":{"input_tokens":"1","output_tokens":1}}"#)]
|
||||
#[case::not_json("not json")]
|
||||
#[tokio::test]
|
||||
async fn a_response_it_cannot_normalize_is_reported_as_already_sent(
|
||||
|
|
@ -294,6 +296,8 @@ async fn direct_and_hosted_calls_share_hooks_and_lifecycle(
|
|||
custom_llm_provider: call.custom_llm_provider.as_deref(),
|
||||
extra_headers: call.extra_headers,
|
||||
timeout: call.timeout,
|
||||
drop_params: Some(false),
|
||||
model_info: None,
|
||||
},
|
||||
&host,
|
||||
Some(host.events.0.sender.clone()),
|
||||
|
|
|
|||
|
|
@ -17,6 +17,9 @@ pub struct ChatCompletionToolCallChunk {
|
|||
pub tool_type: String,
|
||||
pub function: ChatCompletionToolCallFunctionChunk,
|
||||
pub index: i64,
|
||||
/// Provider data that rides on the call, such as Gemini's `thought_signature`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_specific_fields: Option<Map<String, Value>>,
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
|
|
@ -87,6 +90,7 @@ pub struct ChatCompletionsUsage {
|
|||
}
|
||||
|
||||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
#[derive(Default)]
|
||||
pub struct ChatCompletionsChoiceMessage {
|
||||
pub role: String,
|
||||
// Whether an empty turn is `None` or `""` is the provider's choice, not a
|
||||
|
|
@ -94,6 +98,17 @@ pub struct ChatCompletionsChoiceMessage {
|
|||
// while Converse assigns the joined string unconditionally. Each config
|
||||
// mirrors its own, so keep this optional and serialize it even when None.
|
||||
pub content: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub tool_calls: Option<Vec<ChatCompletionToolCallChunk>>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub reasoning_content: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub thinking_blocks: Option<Vec<ChatCompletionThinkingBlock>>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_specific_fields: Option<Map<String, Value>>,
|
||||
/// OpenAI `url_citation` annotations, which Gemini's grounding sources become.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub annotations: Option<Vec<Value>>,
|
||||
}
|
||||
|
||||
#[macro_rules_attribute::apply(crate::wire_type)]
|
||||
|
|
@ -101,6 +116,8 @@ pub struct ChatCompletionsChoice {
|
|||
pub index: u64,
|
||||
pub message: ChatCompletionsChoiceMessage,
|
||||
pub finish_reason: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub logprobs: Option<Value>,
|
||||
}
|
||||
|
||||
/// The normalized response handed back to the host.
|
||||
|
|
@ -114,4 +131,8 @@ pub struct ChatCompletionsResponse {
|
|||
pub model: String,
|
||||
pub choices: Vec<ChatCompletionsChoice>,
|
||||
pub usage: ChatCompletionsUsage,
|
||||
/// Top-level fields a provider adds beside the OpenAI ones, such as Vertex AI's
|
||||
/// `vertex_ai_grounding_metadata`.
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
# rules
|
||||
|
||||
- Shared across Vertex AI's Messages and OCR adapters: the global host, the global and default locations, the `rawPredict` methods and the Anthropic API version Vertex expects in the body
|
||||
- Shared by the Gemini chat adapters on Vertex AI and Google AI Studio: the `generateContent` methods, the Gemini API host, its Files API prefix and its key header
|
||||
- Location validation and host selection are in `llms/src/vertex_ai/common_utils.rs`
|
||||
|
||||
# references
|
||||
|
|
|
|||
|
|
@ -5,3 +5,13 @@ pub const DEFAULT_LOCATION: &str = "us-central1";
|
|||
pub const RAW_PREDICT: &str = "rawPredict";
|
||||
pub const STREAM_RAW_PREDICT: &str = "streamRawPredict";
|
||||
pub const ANTHROPIC_VERSION: &str = "vertex-2023-10-16";
|
||||
|
||||
pub const GENERATE_CONTENT: &str = "generateContent";
|
||||
pub const STREAM_GENERATE_CONTENT: &str = "streamGenerateContent";
|
||||
/// The Gemini API host Google AI Studio keys address.
|
||||
pub const GEMINI_API_BASE: &str = "https://generativelanguage.googleapis.com";
|
||||
/// Files uploaded through the Gemini Files API, passed by reference and never fetched.
|
||||
pub const GEMINI_FILES_API_URI_PREFIX: &str =
|
||||
"https://generativelanguage.googleapis.com/v1beta/files/";
|
||||
/// The header a Gemini API key travels in, never the `key=` query parameter.
|
||||
pub const GEMINI_API_KEY_HEADER: &str = "x-goog-api-key";
|
||||
|
|
|
|||
|
|
@ -22,6 +22,8 @@ litellm-secrets.workspace = true
|
|||
litellm-python-compat.workspace = true
|
||||
litellm-router-types.workspace = true
|
||||
base64.workspace = true
|
||||
sha2.workspace = true
|
||||
uuid.workspace = true
|
||||
bytes.workspace = true
|
||||
data-url = "0.3.2"
|
||||
futures-util.workspace = true
|
||||
|
|
|
|||
|
|
@ -2,3 +2,4 @@
|
|||
|
||||
- Chat Completions translated onto Anthropic Messages. Both format specs live under `llms-types/src/formats/`
|
||||
- `top_k` is not forwarded, because Python gates it per model inside `transform_request`
|
||||
- `build_messages_request` serves every host of the translation; `AnthropicHost` picks the host's cost-map reads, `service_tier`/`speed` support and whether the Anthropic API betas are added (Vertex AI sets its own in `vertex_ai/vertex_ai_partner_models/anthropic`)
|
||||
|
|
|
|||
|
|
@ -4,8 +4,9 @@
|
|||
//! Each event becomes one provider chunk, empty ones included; the host's stream wrapper
|
||||
//! (Python's `CustomStreamWrapper`) decides which of them the client sees.
|
||||
//!
|
||||
//! DIVERGENCE: not built here, since the chat route never asks for what produces them:
|
||||
//! `json_mode` / `response_format` tool conversion, `web_search_calls` and
|
||||
//! DIVERGENCE: a `response_format` tool's answer is held until its block ends and sent as
|
||||
//! one piece of content, where Python streams it as it decides the envelope; an answer
|
||||
//! spread over several blocks is sent per block. Not built: `web_search_calls` and
|
||||
//! `code_interpreter_results` summaries, and a tool call's `caller`. The reasoning token
|
||||
//! estimate Python makes from streamed thinking text (with its tokenizer) is not made, so
|
||||
//! such a stream reports `completion_tokens_details` without counts unless Anthropic sends
|
||||
|
|
@ -28,7 +29,14 @@ use serde_json::{Map, Value, json};
|
|||
|
||||
use crate::{
|
||||
Error,
|
||||
base_llm::{base_model_iterator::StreamTransformer, chat::streaming::StreamShape},
|
||||
base_llm::{
|
||||
base_model_iterator::StreamTransformer,
|
||||
chat::{
|
||||
params::{RESPONSE_FORMAT_TOOL_NAME, python_dumps},
|
||||
response_format::{envelope_contents, json_tool_answer},
|
||||
streaming::StreamShape,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
/// Anthropic sends no status with an `error` event; Python raises it as a 500.
|
||||
|
|
@ -51,6 +59,10 @@ pub struct ModelResponseIterator {
|
|||
web_search_results: Vec<Value>,
|
||||
tool_results: Vec<Value>,
|
||||
compaction_blocks: Vec<Value>,
|
||||
/// Whether the tool block in progress is the `response_format` tool.
|
||||
is_response_format_tool: bool,
|
||||
converted_response_format_tool: bool,
|
||||
response_format_buffer: String,
|
||||
}
|
||||
|
||||
/// What one event contributes to its chunk.
|
||||
|
|
@ -79,9 +91,58 @@ impl ModelResponseIterator {
|
|||
web_search_results: Vec::new(),
|
||||
tool_results: Vec::new(),
|
||||
compaction_blocks: Vec::new(),
|
||||
is_response_format_tool: false,
|
||||
converted_response_format_tool: false,
|
||||
response_format_buffer: String::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// `_handle_json_mode_chunk`: the `response_format` tool's arguments become content.
|
||||
fn handle_json_mode(&mut self, parsed: &mut Parsed) {
|
||||
if !self.shape.json_mode {
|
||||
return;
|
||||
}
|
||||
let Some(tool_use) = parsed.tool_use.take() else {
|
||||
return;
|
||||
};
|
||||
if tool_use.id.is_some() {
|
||||
self.is_response_format_tool =
|
||||
tool_use.function.name.as_deref() == Some(RESPONSE_FORMAT_TOOL_NAME);
|
||||
if self.is_response_format_tool {
|
||||
self.tool_index -= 1;
|
||||
}
|
||||
}
|
||||
if !self.is_response_format_tool {
|
||||
parsed.tool_use = Some(tool_use);
|
||||
return;
|
||||
}
|
||||
self.converted_response_format_tool = true;
|
||||
if self.shape.response_format_schema.is_some() {
|
||||
match tool_use.id {
|
||||
Some(_) => self.response_format_buffer.clear(),
|
||||
None => self
|
||||
.response_format_buffer
|
||||
.push_str(&tool_use.function.arguments),
|
||||
}
|
||||
parsed.text = String::new();
|
||||
return;
|
||||
}
|
||||
parsed.text = json_tool_answer(&[tool_use.function.arguments], None);
|
||||
}
|
||||
|
||||
/// `_held_answer`: the held arguments, or what their envelope wraps.
|
||||
fn flush_response_format_buffer(&mut self) -> String {
|
||||
let held = std::mem::take(&mut self.response_format_buffer);
|
||||
let Some(schema) = &self.shape.response_format_schema else {
|
||||
return String::new();
|
||||
};
|
||||
serde_json::from_str::<Value>(&held)
|
||||
.ok()
|
||||
.and_then(|payload| envelope_contents(&payload, schema, RESPONSE_FORMAT_TOOL_NAME))
|
||||
.map(|contents| python_dumps(&contents))
|
||||
.unwrap_or(held)
|
||||
}
|
||||
|
||||
/// Whether the tool block so far streamed only empty argument fragments.
|
||||
pub fn check_empty_tool_call_args(&self) -> bool {
|
||||
match self.content_blocks.first() {
|
||||
|
|
@ -125,6 +186,7 @@ impl ModelResponseIterator {
|
|||
provider_specific_fields: None,
|
||||
},
|
||||
index: self.tool_index,
|
||||
provider_specific_fields: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -273,6 +335,10 @@ impl ModelResponseIterator {
|
|||
if in_tool && self.check_empty_tool_call_args() {
|
||||
parsed.tool_use = Some(self.tool_call(None, None, "{}".into()));
|
||||
}
|
||||
if self.is_response_format_tool {
|
||||
parsed.text = self.flush_response_format_buffer();
|
||||
}
|
||||
self.is_response_format_tool = false;
|
||||
self.current_content_block_type = None;
|
||||
parsed
|
||||
}
|
||||
|
|
@ -282,10 +348,16 @@ impl ModelResponseIterator {
|
|||
delta: MessagesDelta,
|
||||
usage: Option<MessagesStreamUsage>,
|
||||
) -> Parsed {
|
||||
let provider_finish_reason =
|
||||
finish_reason_for(delta.stop_reason.as_deref().unwrap_or("stop")).to_string();
|
||||
let called_only_the_response_format_tool = provider_finish_reason == "tool_calls"
|
||||
&& self.converted_response_format_tool
|
||||
&& self.tool_index < 0;
|
||||
let mut parsed = Parsed {
|
||||
finish_reason: Some(
|
||||
finish_reason_for(delta.stop_reason.as_deref().unwrap_or("stop")).to_string(),
|
||||
),
|
||||
finish_reason: Some(match called_only_the_response_format_tool {
|
||||
true => "stop".to_owned(),
|
||||
false => provider_finish_reason,
|
||||
}),
|
||||
usage: usage.map(|usage| self.handle_usage(&usage)),
|
||||
..Parsed::default()
|
||||
};
|
||||
|
|
@ -314,7 +386,7 @@ impl ModelResponseIterator {
|
|||
&mut self,
|
||||
event: MessagesStreamEvent,
|
||||
) -> Result<ChatCompletionChunk, Error> {
|
||||
let parsed = match event {
|
||||
let mut parsed = match event {
|
||||
MessagesStreamEvent::ContentBlockDelta { delta, .. } => self.content_block_delta(delta),
|
||||
MessagesStreamEvent::ContentBlockStart { content_block, .. } => {
|
||||
self.content_block_start(content_block)
|
||||
|
|
@ -341,6 +413,7 @@ impl ModelResponseIterator {
|
|||
});
|
||||
}
|
||||
};
|
||||
self.handle_json_mode(&mut parsed);
|
||||
Ok(self.chunk(parsed))
|
||||
}
|
||||
|
||||
|
|
|
|||
777
litellm-rust/crates/llms/src/anthropic/chat/messages.rs
Normal file
777
litellm-rust/crates/llms/src/anthropic/chat/messages.rs
Normal file
|
|
@ -0,0 +1,777 @@
|
|||
//! Chat messages as the Messages API takes them: Python's `translate_system_message`,
|
||||
//! `place_mid_conversation_system` and `anthropic_messages_pt`.
|
||||
|
||||
use std::collections::HashSet;
|
||||
|
||||
use base64::Engine;
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::{
|
||||
Error, ErrorDetail,
|
||||
anthropic::common_utils::ENCRYPTED_REASONING_SIGNATURE_PREFIX,
|
||||
base_llm::chat::{
|
||||
media::{data_url, image_url},
|
||||
params::{parse_tool_arguments, truthy},
|
||||
},
|
||||
};
|
||||
|
||||
pub const EMPTY_TEXT_PLACEHOLDER: &str =
|
||||
"[System: Empty message content sanitised to satisfy protocol]";
|
||||
pub const CONVERTED_SYSTEM_NOTE: &str = "Operator note (not from the user): the following was originally a mid-conversation system-role reminder.";
|
||||
const TOOL_USE_ID_FALLBACK: &str = "tool_use_id";
|
||||
|
||||
fn bad_request(message: impl Into<String>) -> Error {
|
||||
Error::InvalidRequest(ErrorDetail::Message(message.into()))
|
||||
}
|
||||
|
||||
pub fn role(message: &Value) -> &str {
|
||||
match message.get("role").and_then(Value::as_str) {
|
||||
Some("developer") => "system",
|
||||
Some(role) => role,
|
||||
None => "",
|
||||
}
|
||||
}
|
||||
|
||||
/// Python's truthiness of a message's `content`.
|
||||
fn content_is_empty(message: &Value) -> bool {
|
||||
match message.get("content") {
|
||||
None | Some(Value::Null) => true,
|
||||
Some(Value::String(text)) => text.is_empty(),
|
||||
Some(Value::Array(items)) => items.is_empty(),
|
||||
Some(_) => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// `[^a-zA-Z0-9_-]` -> `_`, an empty id becoming `tool_use_id`.
|
||||
pub fn sanitize_tool_use_id(id: &str) -> String {
|
||||
let sanitized: String = id
|
||||
.chars()
|
||||
.map(
|
||||
|c| match c.is_ascii_alphanumeric() || c == '_' || c == '-' {
|
||||
true => c,
|
||||
false => '_',
|
||||
},
|
||||
)
|
||||
.collect();
|
||||
match sanitized.is_empty() {
|
||||
true => TOOL_USE_ID_FALLBACK.to_owned(),
|
||||
false => sanitized,
|
||||
}
|
||||
}
|
||||
|
||||
/// `add_cache_control_to_content`: a dict `cache_control` on the source is copied over.
|
||||
fn with_cache_control(mut block: Value, source: &Value) -> Value {
|
||||
if let (Some(cache_control @ Value::Object(_)), Value::Object(fields)) =
|
||||
(source.get("cache_control"), &mut block)
|
||||
{
|
||||
fields.insert("cache_control".into(), cache_control.clone());
|
||||
}
|
||||
block
|
||||
}
|
||||
|
||||
/// `_cache_control` of the mid-conversation module: only an ephemeral one, rebuilt.
|
||||
fn normalized_cache_control(source: &Value) -> Option<Value> {
|
||||
let cache_control = source.get("cache_control")?.as_object()?;
|
||||
if cache_control.get("type").and_then(Value::as_str) != Some("ephemeral") {
|
||||
return None;
|
||||
}
|
||||
Some(match cache_control.get("ttl").and_then(Value::as_str) {
|
||||
Some(ttl @ ("1h" | "5m")) => json!({"type": "ephemeral", "ttl": ttl}),
|
||||
_ => json!({"type": "ephemeral"}),
|
||||
})
|
||||
}
|
||||
|
||||
/// The leading run of system messages, and the rest.
|
||||
pub fn split_leading_system_run(messages: Vec<Value>) -> (Vec<Value>, Vec<Value>) {
|
||||
let leading = messages
|
||||
.iter()
|
||||
.position(|message| role(message) != "system")
|
||||
.unwrap_or(messages.len());
|
||||
let mut messages = messages;
|
||||
let rest = messages.split_off(leading);
|
||||
(messages, rest)
|
||||
}
|
||||
|
||||
/// `translate_system_message` for the leading run.
|
||||
pub fn system_blocks(run: &[Value]) -> Vec<Value> {
|
||||
let mut blocks = Vec::new();
|
||||
for message in run {
|
||||
match message.get("content") {
|
||||
Some(Value::String(text)) => {
|
||||
if text.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let mut block = Map::from_iter([
|
||||
("type".to_owned(), json!("text")),
|
||||
("text".to_owned(), json!(text)),
|
||||
]);
|
||||
if let Some(cache_control) = message.get("cache_control") {
|
||||
block.insert("cache_control".into(), cache_control.clone());
|
||||
}
|
||||
blocks.push(Value::Object(block));
|
||||
}
|
||||
Some(Value::Array(parts)) => {
|
||||
for part in parts {
|
||||
let text = part.get("text").cloned().unwrap_or(Value::Null);
|
||||
let kind = part.get("type").cloned().unwrap_or(Value::Null);
|
||||
let empty = match &text {
|
||||
Value::String(text) => text.is_empty(),
|
||||
Value::Null => true,
|
||||
_ => false,
|
||||
};
|
||||
if kind.as_str() == Some("text") && empty {
|
||||
continue;
|
||||
}
|
||||
let mut block =
|
||||
Map::from_iter([("type".to_owned(), kind), ("text".to_owned(), text)]);
|
||||
if let Some(cache_control) = part.get("cache_control") {
|
||||
block.insert("cache_control".into(), cache_control.clone());
|
||||
}
|
||||
blocks.push(Value::Object(block));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
blocks
|
||||
}
|
||||
|
||||
/// The non-empty text parts of a system message, each with its normalized `cache_control`.
|
||||
fn system_text_parts(message: &Value) -> Vec<Value> {
|
||||
let block = |text: &str, cache_control: Option<Value>| {
|
||||
let mut block = Map::from_iter([
|
||||
("type".to_owned(), json!("text")),
|
||||
("text".to_owned(), json!(text)),
|
||||
]);
|
||||
if let Some(cache_control) = cache_control {
|
||||
block.insert("cache_control".into(), cache_control);
|
||||
}
|
||||
Value::Object(block)
|
||||
};
|
||||
match message.get("content") {
|
||||
Some(Value::String(text)) if !text.is_empty() => {
|
||||
vec![block(text, normalized_cache_control(message))]
|
||||
}
|
||||
Some(Value::Array(parts)) => parts
|
||||
.iter()
|
||||
.filter(|part| part.get("type").and_then(Value::as_str) == Some("text"))
|
||||
.filter_map(|part| {
|
||||
let text = part.get("text")?.as_str().filter(|text| !text.is_empty())?;
|
||||
Some(block(text, normalized_cache_control(part)))
|
||||
})
|
||||
.collect(),
|
||||
_ => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// `system_message_as_user` for each message of a run that carries text.
|
||||
pub fn converted_user_turns(run: &[Value]) -> Vec<Value> {
|
||||
run.iter()
|
||||
.filter_map(|message| {
|
||||
let parts = system_text_parts(message);
|
||||
if parts.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let mut content = vec![json!({"type": "text", "text": CONVERTED_SYSTEM_NOTE})];
|
||||
content.extend(parts);
|
||||
Some(json!({"role": "user", "content": content}))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, PartialEq, Eq)]
|
||||
enum Kind {
|
||||
System,
|
||||
Tool,
|
||||
User,
|
||||
Other,
|
||||
}
|
||||
|
||||
fn kind(message: &Value) -> Kind {
|
||||
match role(message) {
|
||||
"system" => Kind::System,
|
||||
"tool" | "function" => Kind::Tool,
|
||||
"user" => Kind::User,
|
||||
_ => Kind::Other,
|
||||
}
|
||||
}
|
||||
|
||||
/// `_converted_for_unflagged_model`: every later system message becomes a user turn in
|
||||
/// place, a run caught before a tool result moving behind it.
|
||||
pub fn place_mid_conversation_system(messages: Vec<Value>) -> Vec<Value> {
|
||||
if !messages.iter().any(|message| kind(message) == Kind::System) {
|
||||
return messages;
|
||||
}
|
||||
let mut runs: Vec<(Kind, Vec<Value>)> = Vec::new();
|
||||
for message in messages {
|
||||
let message_kind = kind(&message);
|
||||
match runs.last_mut() {
|
||||
Some((last, run)) if *last == message_kind => run.push(message),
|
||||
_ => runs.push((message_kind, vec![message])),
|
||||
}
|
||||
}
|
||||
let mut placed = Vec::new();
|
||||
for index in 0..runs.len() {
|
||||
let (run_kind, run) = &runs[index];
|
||||
let follower = runs.get(index + 1).map(|(kind, _)| *kind);
|
||||
match run_kind {
|
||||
Kind::System => {
|
||||
if follower != Some(Kind::Tool) {
|
||||
placed.extend(converted_user_turns(run));
|
||||
}
|
||||
}
|
||||
Kind::Tool if index > 0 && runs[index - 1].0 == Kind::System => {
|
||||
placed.extend(run.iter().cloned());
|
||||
placed.extend(converted_user_turns(&runs[index - 1].1));
|
||||
}
|
||||
_ => placed.extend(run.iter().cloned()),
|
||||
}
|
||||
}
|
||||
placed
|
||||
}
|
||||
|
||||
/// `_sanitize_empty_text_content`.
|
||||
fn sanitize_empty_text(mut message: Value) -> Value {
|
||||
if !matches!(role(&message), "user" | "assistant") {
|
||||
return message;
|
||||
}
|
||||
match message.get_mut("content") {
|
||||
Some(Value::String(text)) if text.trim().is_empty() => {
|
||||
*text = EMPTY_TEXT_PLACEHOLDER.to_owned();
|
||||
}
|
||||
Some(Value::Array(parts)) => {
|
||||
for part in parts {
|
||||
if part.get("type").and_then(Value::as_str) != Some("text") {
|
||||
continue;
|
||||
}
|
||||
let blank = part
|
||||
.get("text")
|
||||
.and_then(Value::as_str)
|
||||
.is_none_or(|text| text.trim().is_empty());
|
||||
if blank && let Value::Object(fields) = part {
|
||||
fields.insert("text".into(), json!(EMPTY_TEXT_PLACEHOLDER));
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
message
|
||||
}
|
||||
|
||||
/// `convert_to_anthropic_image_obj`: a data URL as its media type and payload.
|
||||
fn image_source(url: &str, format: Option<&str>) -> Result<(String, String), Error> {
|
||||
let (media_type, data) = data_url(url).ok_or_else(|| {
|
||||
bad_request(format!(
|
||||
"Image url not in expected format. Example Expected input - \"image_url\": \"data:image/jpeg;base64,{{base64_image}}\". Supported formats - ['image/jpeg', 'image/png', 'image/gif', 'image/webp']. Error: {url:.40}"
|
||||
))
|
||||
})?;
|
||||
let media_type = match format {
|
||||
Some(format) => format.to_owned(),
|
||||
None => media_type.replace("\\/", "/"),
|
||||
};
|
||||
Ok((media_type, data.to_owned()))
|
||||
}
|
||||
|
||||
/// `create_anthropic_image_param`.
|
||||
fn image_block(part: &Value) -> Result<Value, Error> {
|
||||
let url = image_url(part).unwrap_or_default();
|
||||
let format = part
|
||||
.get("image_url")
|
||||
.and_then(|image| image.get("format"))
|
||||
.and_then(Value::as_str);
|
||||
if url.starts_with("https://") {
|
||||
return Ok(json!({"type": "image", "source": {"type": "url", "url": url}}));
|
||||
}
|
||||
if url.starts_with("http") {
|
||||
return Err(Error::Unsupported(
|
||||
"an http:// image the host did not download",
|
||||
));
|
||||
}
|
||||
let (media_type, data) = image_source(url, format)?;
|
||||
Ok(
|
||||
json!({"type": "image", "source": {"type": "base64", "media_type": media_type, "data": data}}),
|
||||
)
|
||||
}
|
||||
|
||||
/// `anthropic_process_openai_file_message`.
|
||||
fn file_block(part: &Value) -> Result<Value, Error> {
|
||||
let file = part.get("file").ok_or_else(|| {
|
||||
bad_request("Content block has type='file' but is missing the required 'file' field")
|
||||
})?;
|
||||
let format = file.get("format").and_then(Value::as_str);
|
||||
if let Some(data) = file
|
||||
.get("file_data")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|data| !data.is_empty())
|
||||
{
|
||||
if data.starts_with("http") {
|
||||
return Err(Error::Unsupported(
|
||||
"a file_data URL the host did not download",
|
||||
));
|
||||
}
|
||||
let (media_type, data) = image_source(data, format)?;
|
||||
return Ok(
|
||||
json!({"type": "document", "source": {"type": "base64", "media_type": media_type, "data": data}}),
|
||||
);
|
||||
}
|
||||
match file
|
||||
.get("file_id")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|id| !id.is_empty())
|
||||
{
|
||||
Some(id) if format.is_none() && id.starts_with("http") => {
|
||||
Ok(json!({"type": "document", "source": {"type": "url", "url": id}}))
|
||||
}
|
||||
// Uploaded files need the Files API betas and, for a container upload, the code
|
||||
// execution tool.
|
||||
Some(_) => Err(Error::Unsupported("an Anthropic Files API file_id")),
|
||||
None => Err(bad_request(format!(
|
||||
"Either file_data or file_id must be present in the file message: {part}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
/// `_is_anthropic_document_data_uri`.
|
||||
fn is_document_data_uri(url: &str) -> bool {
|
||||
url.strip_prefix("data:")
|
||||
.map(|rest| rest.split([';', ',']).next().unwrap_or_default())
|
||||
.is_some_and(|media_type| matches!(media_type, "application/pdf" | "text/plain"))
|
||||
}
|
||||
|
||||
/// `convert_to_anthropic_tool_result`.
|
||||
fn tool_result(message: &Value) -> Result<Value, Error> {
|
||||
let content = match message.get("content") {
|
||||
Some(Value::String(text)) => json!(text),
|
||||
Some(Value::Array(parts)) => {
|
||||
let mut blocks = Vec::new();
|
||||
for part in parts {
|
||||
match part.get("type").and_then(Value::as_str) {
|
||||
Some("text") => {
|
||||
let mut block = Map::from_iter([
|
||||
("type".to_owned(), json!("text")),
|
||||
(
|
||||
"text".to_owned(),
|
||||
part.get("text").cloned().unwrap_or(Value::Null),
|
||||
),
|
||||
]);
|
||||
if let Some(cache_control) =
|
||||
part.get("cache_control").filter(|value| !value.is_null())
|
||||
{
|
||||
block.insert("cache_control".into(), cache_control.clone());
|
||||
}
|
||||
blocks.push(Value::Object(block));
|
||||
}
|
||||
Some("image_url") => {
|
||||
let url = image_url(part).unwrap_or_default();
|
||||
let block = match is_document_data_uri(url) {
|
||||
true => {
|
||||
file_block(&json!({"type": "file", "file": {"file_data": url}}))?
|
||||
}
|
||||
false => image_block(part)?,
|
||||
};
|
||||
blocks.push(with_cache_control(block, part));
|
||||
}
|
||||
Some("tool_reference") => blocks.push(json!({
|
||||
"type": "tool_reference",
|
||||
"tool_name": part.get("tool_name").cloned().unwrap_or(Value::Null),
|
||||
})),
|
||||
Some("file") => blocks.push(with_cache_control(file_block(part)?, part)),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Value::Array(blocks)
|
||||
}
|
||||
_ => json!(""),
|
||||
};
|
||||
let id = match role(message) {
|
||||
"tool" => message
|
||||
.get("tool_call_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned)
|
||||
.ok_or_else(|| bad_request("A tool message needs a tool_call_id."))?,
|
||||
_ => message
|
||||
.get("tool_call_id")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|id| !id.is_empty())
|
||||
.map(str::to_owned)
|
||||
.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
|
||||
};
|
||||
let mut result = Map::from_iter([
|
||||
("type".to_owned(), json!("tool_result")),
|
||||
("tool_use_id".to_owned(), json!(sanitize_tool_use_id(&id))),
|
||||
("content".to_owned(), content),
|
||||
]);
|
||||
if let Some(cache_control) = message
|
||||
.get("cache_control")
|
||||
.filter(|value| !value.is_null())
|
||||
{
|
||||
result.insert("cache_control".into(), cache_control.clone());
|
||||
}
|
||||
Ok(Value::Object(result))
|
||||
}
|
||||
|
||||
pub fn is_encrypted_reasoning(block: &Value) -> bool {
|
||||
let field = match block.get("type").and_then(Value::as_str) {
|
||||
Some("thinking") => block.get("signature"),
|
||||
Some("redacted_thinking") => block.get("data"),
|
||||
_ => None,
|
||||
};
|
||||
field
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| value.starts_with(ENCRYPTED_REASONING_SIGNATURE_PREFIX))
|
||||
}
|
||||
|
||||
/// `is_unsignable_thinking_block`.
|
||||
pub fn is_unsignable_thinking_block(block: &Value) -> bool {
|
||||
if is_encrypted_reasoning(block) {
|
||||
return true;
|
||||
}
|
||||
block.get("type").and_then(Value::as_str) == Some("thinking")
|
||||
&& block
|
||||
.get("signature")
|
||||
.and_then(Value::as_str)
|
||||
.is_none_or(str::is_empty)
|
||||
}
|
||||
|
||||
/// `native_redacted_thinking_data`: the base64 wrapper Converse adds, removed.
|
||||
pub fn native_redacted_thinking_data(data: &str) -> String {
|
||||
let engine = base64::engine::general_purpose::STANDARD;
|
||||
engine
|
||||
.decode(data)
|
||||
.ok()
|
||||
.and_then(|bytes| String::from_utf8(bytes).ok())
|
||||
.filter(|payload| payload.is_ascii() && engine.decode(payload).is_ok())
|
||||
.unwrap_or_else(|| data.to_owned())
|
||||
}
|
||||
|
||||
/// `_replayable_thinking_blocks`.
|
||||
fn replayable_thinking_blocks(blocks: &[Value]) -> Vec<Value> {
|
||||
blocks
|
||||
.iter()
|
||||
.filter(|block| !is_unsignable_thinking_block(block))
|
||||
.map(|block| {
|
||||
let mut block = block.clone();
|
||||
if let Value::Object(fields) = &mut block {
|
||||
if fields.get("type").and_then(Value::as_str) == Some("redacted_thinking")
|
||||
&& let Some(Value::String(data)) = fields.get("data")
|
||||
{
|
||||
let native = native_redacted_thinking_data(data);
|
||||
fields.insert("data".into(), json!(native));
|
||||
}
|
||||
fields.remove("index");
|
||||
}
|
||||
block
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// The tool calls of an assistant message, as `tool_use` blocks.
|
||||
fn tool_uses(tool_calls: &[Value], seen: &mut HashSet<String>) -> Result<Vec<Value>, Error> {
|
||||
let mut blocks = Vec::new();
|
||||
for call in tool_calls {
|
||||
if call.get("type").and_then(Value::as_str) != Some("function") {
|
||||
continue;
|
||||
}
|
||||
let id = call.get("id").and_then(Value::as_str).unwrap_or_default();
|
||||
if id.starts_with("srvtoolu_") {
|
||||
return Err(Error::Unsupported(
|
||||
"replaying an Anthropic server tool call",
|
||||
));
|
||||
}
|
||||
let function = call.get("function");
|
||||
let name = function
|
||||
.and_then(|function| function.get("name"))
|
||||
.cloned()
|
||||
.unwrap_or(Value::Null);
|
||||
let input = parse_tool_arguments(function.and_then(|function| function.get("arguments")))?;
|
||||
let id = sanitize_tool_use_id(id);
|
||||
let block = with_cache_control(
|
||||
json!({"type": "tool_use", "id": id, "name": name, "input": input}),
|
||||
call,
|
||||
);
|
||||
if !seen.insert(id) {
|
||||
continue;
|
||||
}
|
||||
blocks.push(block);
|
||||
}
|
||||
Ok(blocks)
|
||||
}
|
||||
|
||||
/// `anthropic_messages_pt` over messages whose later system messages are already placed.
|
||||
pub fn anthropic_messages_pt(messages: Vec<Value>) -> Result<Vec<Value>, Error> {
|
||||
let messages: Vec<Value> = messages.into_iter().map(sanitize_empty_text).collect();
|
||||
if messages.is_empty() {
|
||||
return Err(bad_request(
|
||||
"Anthropic requires at least one non-system message. Either provide one, or set `litellm.modify_params = True` // `litellm_settings::modify_params: True` to add the dummy user message - {'role': 'user', 'content': 'Please continue.'}.",
|
||||
));
|
||||
}
|
||||
let user_types = |message: &Value| matches!(role(message), "user" | "tool" | "function");
|
||||
// `_tool_result_later_in_turn`: whether a client tool result follows in the same turn.
|
||||
let mut tool_result_later = vec![false; messages.len()];
|
||||
let mut later = false;
|
||||
for index in (0..messages.len()).rev() {
|
||||
tool_result_later[index] = later;
|
||||
let message = &messages[index];
|
||||
later = user_types(message) && (later || matches!(role(message), "tool" | "function"));
|
||||
}
|
||||
|
||||
let mut out: Vec<Value> = Vec::new();
|
||||
let mut index = 0;
|
||||
while index < messages.len() {
|
||||
let start = index;
|
||||
let mut user_content: Vec<Value> = Vec::new();
|
||||
while index < messages.len() && user_types(&messages[index]) {
|
||||
let message = &messages[index];
|
||||
if role(message) == "user" {
|
||||
match message.get("content") {
|
||||
_ if content_is_empty(message) => {
|
||||
if !tool_result_later[index]
|
||||
&& message.get("cache_control").is_some_and(truthy)
|
||||
{
|
||||
user_content.push(with_cache_control(
|
||||
json!({"type": "text", "text": EMPTY_TEXT_PLACEHOLDER}),
|
||||
message,
|
||||
));
|
||||
}
|
||||
}
|
||||
Some(Value::Array(parts)) => {
|
||||
for part in parts {
|
||||
match part.get("type").and_then(Value::as_str) {
|
||||
Some("image_url") => {
|
||||
user_content.push(with_cache_control(image_block(part)?, part))
|
||||
}
|
||||
Some("text") => user_content.push(with_cache_control(
|
||||
json!({"type": "text", "text": part.get("text").cloned().unwrap_or(Value::Null)}),
|
||||
part,
|
||||
)),
|
||||
Some("document") => {
|
||||
user_content.push(with_cache_control(part.clone(), part))
|
||||
}
|
||||
Some("file") => {
|
||||
user_content.push(with_cache_control(file_block(part)?, part))
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(Value::String(text)) => user_content.push(with_cache_control(
|
||||
json!({"type": "text", "text": text}),
|
||||
message,
|
||||
)),
|
||||
_ => {}
|
||||
}
|
||||
} else {
|
||||
user_content.push(tool_result(message)?);
|
||||
}
|
||||
index += 1;
|
||||
}
|
||||
if !user_content.is_empty() {
|
||||
out.push(json!({"role": "user", "content": user_content}));
|
||||
} else if let Some(empty) = messages[start..index]
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|message| role(message) == "user" && content_is_empty(message))
|
||||
{
|
||||
out.push(json!({"role": "user", "content": [with_cache_control(
|
||||
json!({"type": "text", "text": EMPTY_TEXT_PLACEHOLDER}),
|
||||
empty,
|
||||
)]}));
|
||||
}
|
||||
|
||||
let mut seen_tool_ids: HashSet<String> = HashSet::new();
|
||||
let mut assistant_content: Vec<Value> = Vec::new();
|
||||
while index < messages.len() && role(&messages[index]) == "assistant" {
|
||||
let message = &messages[index];
|
||||
let provider_fields = message.get("provider_specific_fields");
|
||||
if let Some(Value::Array(blocks)) =
|
||||
provider_fields.and_then(|fields| fields.get("compaction_blocks"))
|
||||
{
|
||||
assistant_content.extend(blocks.iter().cloned());
|
||||
}
|
||||
if provider_fields.is_some_and(|fields| {
|
||||
fields
|
||||
.get("web_search_results")
|
||||
.is_some_and(|v| !v.is_null())
|
||||
|| fields.get("tool_results").is_some_and(|v| !v.is_null())
|
||||
}) && message
|
||||
.get("tool_calls")
|
||||
.and_then(Value::as_array)
|
||||
.is_some_and(|calls| !calls.is_empty())
|
||||
{
|
||||
return Err(Error::Unsupported(
|
||||
"replaying Anthropic server tool results",
|
||||
));
|
||||
}
|
||||
let thinking_blocks = match message.get("thinking_blocks") {
|
||||
Some(Value::Array(blocks)) => Some(replayable_thinking_blocks(blocks)),
|
||||
_ => None,
|
||||
};
|
||||
let parts = message.get("content").and_then(Value::as_array);
|
||||
let inline_thinking = parts.is_some_and(|parts| {
|
||||
parts.iter().any(|part| {
|
||||
matches!(
|
||||
part.get("type").and_then(Value::as_str),
|
||||
Some("thinking" | "redacted_thinking")
|
||||
)
|
||||
})
|
||||
});
|
||||
if let Some(blocks) = thinking_blocks
|
||||
&& !inline_thinking
|
||||
{
|
||||
assistant_content.extend(blocks);
|
||||
}
|
||||
match message.get("content") {
|
||||
Some(Value::Array(parts)) => {
|
||||
for part in parts.iter().filter(|part| part.is_object()) {
|
||||
let kind = part.get("type").and_then(Value::as_str).unwrap_or_default();
|
||||
let text = part.get("text").and_then(Value::as_str).unwrap_or_default();
|
||||
if kind == "thinking" && !is_unsignable_thinking_block(part) {
|
||||
assistant_content.push(part.clone());
|
||||
} else if kind == "text" && !text.is_empty() {
|
||||
assistant_content.push(with_cache_control(
|
||||
json!({"type": "text", "text": text}),
|
||||
part,
|
||||
));
|
||||
} else if kind == "server_tool_use" || kind.ends_with("_tool_result") {
|
||||
assistant_content.push(part.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(Value::String(text)) if !text.is_empty() => {
|
||||
assistant_content.push(with_cache_control(
|
||||
json!({"type": "text", "text": text}),
|
||||
message,
|
||||
));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
if let Some(Value::Array(calls)) = message.get("tool_calls") {
|
||||
assistant_content.extend(tool_uses(calls, &mut seen_tool_ids)?);
|
||||
}
|
||||
if let Some(call) = message.get("function_call").filter(|call| !call.is_null()) {
|
||||
let input = parse_tool_arguments(call.get("arguments"))?;
|
||||
assistant_content.push(json!({
|
||||
"type": "tool_use",
|
||||
"id": uuid::Uuid::new_v4().to_string(),
|
||||
"name": call.get("name").filter(|name| !name.is_null()).cloned().unwrap_or_else(|| json!("")),
|
||||
"input": input,
|
||||
}));
|
||||
}
|
||||
index += 1;
|
||||
}
|
||||
if !assistant_content.is_empty() {
|
||||
out.push(json!({"role": "assistant", "content": assistant_content}));
|
||||
}
|
||||
if index == start {
|
||||
return Err(bad_request(format!(
|
||||
"Invalid Message passed in {}",
|
||||
messages[index]
|
||||
)));
|
||||
}
|
||||
}
|
||||
if let Some(last) = out.last_mut()
|
||||
&& last["role"] == "assistant"
|
||||
&& let Some(Value::Array(blocks)) = last.get_mut("content")
|
||||
{
|
||||
for block in blocks {
|
||||
if block.get("type").and_then(Value::as_str) == Some("text")
|
||||
&& let Some(Value::String(text)) = block.get_mut("text")
|
||||
{
|
||||
*text = text.trim_end().to_owned();
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn list(value: Value) -> Vec<Value> {
|
||||
value.as_array().unwrap().clone()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn tool_round_trip_merges_turns() {
|
||||
let messages = anthropic_messages_pt(list(json!([
|
||||
{"role": "user", "content": "weather?"},
|
||||
{"role": "assistant", "content": null, "tool_calls": [{"id": "call.1", "type": "function", "function": {"name": "get", "arguments": "{\"c\": \"P\"}"}}]},
|
||||
{"role": "tool", "tool_call_id": "call.1", "content": "sunny"},
|
||||
{"role": "user", "content": [{"type": "text", "text": "thanks", "cache_control": {"type": "ephemeral"}}]}
|
||||
])))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
messages,
|
||||
list(json!([
|
||||
{"role": "user", "content": [{"type": "text", "text": "weather?"}]},
|
||||
{"role": "assistant", "content": [{"type": "tool_use", "id": "call_1", "name": "get", "input": {"c": "P"}}]},
|
||||
{"role": "user", "content": [
|
||||
{"type": "tool_result", "tool_use_id": "call_1", "content": "sunny"},
|
||||
{"type": "text", "text": "thanks", "cache_control": {"type": "ephemeral"}}
|
||||
]}
|
||||
]))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn images_empty_text_and_final_assistant_whitespace() {
|
||||
let messages = anthropic_messages_pt(list(json!([
|
||||
{"role": "user", "content": [
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBO"}},
|
||||
{"type": "image_url", "image_url": "https://x/y.png"},
|
||||
{"type": "text", "text": " "}
|
||||
]},
|
||||
{"role": "assistant", "content": "Sure "}
|
||||
])))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
messages[0]["content"],
|
||||
json!([
|
||||
{"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "iVBO"}},
|
||||
{"type": "image", "source": {"type": "url", "url": "https://x/y.png"}},
|
||||
{"type": "text", "text": EMPTY_TEXT_PLACEHOLDER}
|
||||
])
|
||||
);
|
||||
assert_eq!(
|
||||
messages[1]["content"],
|
||||
json!([{"type": "text", "text": "Sure"}])
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn later_system_messages_become_operator_notes() {
|
||||
let placed = place_mid_conversation_system(list(json!([
|
||||
{"role": "user", "content": "a"},
|
||||
{"role": "assistant", "content": "b", "tool_calls": [{"id": "t", "type": "function", "function": {"name": "f", "arguments": "{}"}}]},
|
||||
{"role": "system", "content": "remember", "cache_control": {"type": "ephemeral", "ttl": "1h", "x": 1}},
|
||||
{"role": "tool", "tool_call_id": "t", "content": "r"}
|
||||
])));
|
||||
assert_eq!(placed[2]["role"], "tool");
|
||||
assert_eq!(
|
||||
placed[3],
|
||||
json!({"role": "user", "content": [
|
||||
{"type": "text", "text": CONVERTED_SYSTEM_NOTE},
|
||||
{"type": "text", "text": "remember", "cache_control": {"type": "ephemeral", "ttl": "1h"}}
|
||||
]})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn unsigned_thinking_is_dropped_and_signed_thinking_leads() {
|
||||
let messages = anthropic_messages_pt(list(json!([
|
||||
{"role": "user", "content": "q"},
|
||||
{"role": "assistant", "content": "a", "thinking_blocks": [
|
||||
{"type": "thinking", "thinking": "t", "signature": "s", "index": 0},
|
||||
{"type": "thinking", "thinking": "u"}
|
||||
]}
|
||||
])))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
messages[1]["content"],
|
||||
json!([{"type": "thinking", "thinking": "t", "signature": "s"}, {"type": "text", "text": "a"}])
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,2 +1,4 @@
|
|||
pub mod handler;
|
||||
pub mod messages;
|
||||
pub mod params;
|
||||
pub mod transformation;
|
||||
|
|
|
|||
1415
litellm-rust/crates/llms/src/anthropic/chat/params.rs
Normal file
1415
litellm-rust/crates/llms/src/anthropic/chat/params.rs
Normal file
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
|
|
@ -8,56 +8,15 @@ use serde_json::{Map, Value};
|
|||
|
||||
use crate::{
|
||||
Error, ErrorDetail,
|
||||
base_llm::model_info::{self, ModelInfo},
|
||||
base_llm::{
|
||||
chat::params::{
|
||||
OPENAI_PARAMS, UNVALIDATED_PARAMS, ordered, request_drop_params, string_list, truthy,
|
||||
},
|
||||
model_info::{self, ModelInfo},
|
||||
},
|
||||
openai::chat::transformation::{is_litellm_owned, take_extra_body},
|
||||
};
|
||||
|
||||
/// The keys of `DEFAULT_CHAT_COMPLETION_PARAM_VALUES` a request carries as OpenAI params, in
|
||||
/// that table's order. Every other key goes to Azure through `extra_body`.
|
||||
const OPENAI_PARAMS: &[&str] = &[
|
||||
"functions",
|
||||
"function_call",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"n",
|
||||
"stream",
|
||||
"stream_options",
|
||||
"stop",
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"modalities",
|
||||
"prediction",
|
||||
"audio",
|
||||
"presence_penalty",
|
||||
"frequency_penalty",
|
||||
"logit_bias",
|
||||
"user",
|
||||
"response_format",
|
||||
"seed",
|
||||
"tools",
|
||||
"tool_choice",
|
||||
"max_retries",
|
||||
"logprobs",
|
||||
"top_logprobs",
|
||||
"extra_headers",
|
||||
"parallel_tool_calls",
|
||||
"reasoning_effort",
|
||||
"verbosity",
|
||||
"thinking",
|
||||
"web_search_options",
|
||||
"include_server_side_tool_invocations",
|
||||
"service_tier",
|
||||
"safety_identifier",
|
||||
"prompt_cache_key",
|
||||
"prompt_cache_retention",
|
||||
"store",
|
||||
"metadata",
|
||||
"context_management",
|
||||
];
|
||||
|
||||
/// `PROVIDER_UNVALIDATED_PARAMS`: never refused, only left out when unsupported.
|
||||
const UNVALIDATED_PARAMS: &[&str] = &["user", "stream_options", "stream", "max_retries"];
|
||||
|
||||
/// `AzureOpenAIConfig.get_supported_openai_params`.
|
||||
const AZURE_GPT_PARAMS: &[&str] = &[
|
||||
"temperature",
|
||||
|
|
@ -410,6 +369,7 @@ impl AzureChatRequest<'_> {
|
|||
let model = self.deployment;
|
||||
if model.starts_with(RESPONSES_ROUTE)
|
||||
|| model_info::field(self.info(), &format!("azure/{model}"), "mode")
|
||||
.as_ref()
|
||||
.and_then(Value::as_str)
|
||||
== Some("responses")
|
||||
{
|
||||
|
|
@ -830,7 +790,9 @@ fn openai_gpt_params(model: &str, info: Option<&dyn ModelInfo>) -> Vec<&'static
|
|||
.split_once(RESPONSES_ROUTE)
|
||||
.map_or(model, |(_, rest)| rest);
|
||||
if matches!(
|
||||
model_info::field(info, catalog_name, "litellm_provider").and_then(Value::as_str),
|
||||
model_info::field(info, catalog_name, "litellm_provider")
|
||||
.as_ref()
|
||||
.and_then(Value::as_str),
|
||||
Some("openai" | "text-completion-openai")
|
||||
) {
|
||||
params.push("user");
|
||||
|
|
@ -913,14 +875,6 @@ fn azure_messages(messages: Value) -> Value {
|
|||
Value::Array(messages)
|
||||
}
|
||||
|
||||
/// The request's OpenAI params in `DEFAULT_CHAT_COMPLETION_PARAM_VALUES` order, the order
|
||||
/// Python maps them in.
|
||||
fn ordered(params: &Map<String, Value>) -> impl Iterator<Item = (&'static str, &Value)> {
|
||||
OPENAI_PARAMS
|
||||
.iter()
|
||||
.filter_map(|name| params.get(*name).map(|value| (*name, value)))
|
||||
}
|
||||
|
||||
/// `_map_openai_params`: every supported param still unmapped goes as sent.
|
||||
fn copy_supported(
|
||||
non_default: &mut Map<String, Value>,
|
||||
|
|
@ -936,40 +890,6 @@ fn copy_supported(
|
|||
}
|
||||
}
|
||||
|
||||
/// Python's truthiness for a JSON value.
|
||||
fn truthy(value: &Value) -> bool {
|
||||
match value {
|
||||
Value::Null => false,
|
||||
Value::Bool(flag) => *flag,
|
||||
Value::Number(number) => number.as_f64() != Some(0.0),
|
||||
Value::String(text) => !text.is_empty(),
|
||||
Value::Array(items) => !items.is_empty(),
|
||||
Value::Object(object) => !object.is_empty(),
|
||||
}
|
||||
}
|
||||
|
||||
/// `normalize_drop_params`: a boolean, or a config's `"true"`.
|
||||
fn request_drop_params(value: Option<Value>) -> bool {
|
||||
match value {
|
||||
Some(Value::Bool(flag)) => flag,
|
||||
Some(Value::String(text)) => text.trim().eq_ignore_ascii_case("true"),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn string_list(value: Option<Value>) -> Vec<String> {
|
||||
match value {
|
||||
Some(Value::Array(items)) => items
|
||||
.into_iter()
|
||||
.filter_map(|item| match item {
|
||||
Value::String(text) => Some(text),
|
||||
_ => None,
|
||||
})
|
||||
.collect(),
|
||||
_ => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::HashMap;
|
||||
|
|
@ -997,8 +917,8 @@ mod tests {
|
|||
}
|
||||
|
||||
impl ModelInfo for Entries {
|
||||
fn entry(&self, model: &str) -> Option<&Map<String, Value>> {
|
||||
self.0.get(model)
|
||||
fn entry(&self, model: &str) -> Option<std::borrow::Cow<'_, Map<String, Value>>> {
|
||||
self.0.get(model).map(std::borrow::Cow::Borrowed)
|
||||
}
|
||||
|
||||
fn declares(&self, key: &str) -> bool {
|
||||
|
|
|
|||
209
litellm-rust/crates/llms/src/base_llm/chat/media.rs
Normal file
209
litellm-rust/crates/llms/src/base_llm/chat/media.rs
Normal file
|
|
@ -0,0 +1,209 @@
|
|||
//! Media a chat message points at. Python downloads an image or file URL a provider cannot
|
||||
//! take by reference while it builds the request; here the host downloads them first, from
|
||||
//! the list a config names, and hands the messages back with each one inlined as a data URL.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use litellm_llms_types::formats::chat_completions::{ChatMessage, ChatMessageContent};
|
||||
use serde_json::Value;
|
||||
|
||||
/// Python's `MAX_IMAGE_URL_DOWNLOAD_SIZE_MB`.
|
||||
pub const MAX_MEDIA_DOWNLOAD_BYTES: u64 = 50 * 1024 * 1024;
|
||||
|
||||
/// A part's `image_url`, given as a string or as `{"url": ...}`.
|
||||
pub fn image_url(part: &Value) -> Option<&str> {
|
||||
let image_url = part.get("image_url")?;
|
||||
image_url
|
||||
.as_str()
|
||||
.or_else(|| image_url.get("url").and_then(Value::as_str))
|
||||
}
|
||||
|
||||
/// `data:<meta>;base64,<data>` split into its media type and payload, the way Python's
|
||||
/// `split("data:")[1].split(";base64,")` reads it.
|
||||
pub fn data_url(url: &str) -> Option<(&str, &str)> {
|
||||
let (_, rest) = url.split_once("data:")?;
|
||||
let (media_type, data) = rest.split_once(";base64,")?;
|
||||
(!data.contains(";base64,")).then_some((media_type, data))
|
||||
}
|
||||
|
||||
fn parts(message: &ChatMessage) -> &[Value] {
|
||||
match &message.content {
|
||||
Some(ChatMessageContent::Parts(parts)) => parts,
|
||||
_ => &[],
|
||||
}
|
||||
}
|
||||
|
||||
/// The URLs of `image_url` parts (and, with `files`, of `file.file_id` / `file.file_data`)
|
||||
/// in the messages `roles` names that `remote` says must be downloaded, in order, each once.
|
||||
pub fn remote_urls(
|
||||
messages: &[ChatMessage],
|
||||
roles: &[&str],
|
||||
files: bool,
|
||||
remote: impl Fn(&str) -> bool,
|
||||
) -> Vec<String> {
|
||||
let mut urls: Vec<String> = Vec::new();
|
||||
let mut add = |url: &str| {
|
||||
if remote(url) && !urls.iter().any(|seen| seen == url) {
|
||||
urls.push(url.to_owned());
|
||||
}
|
||||
};
|
||||
for message in messages
|
||||
.iter()
|
||||
.filter(|message| roles.contains(&message.role.as_str()))
|
||||
{
|
||||
for part in parts(message) {
|
||||
match part.get("type").and_then(Value::as_str) {
|
||||
Some("image_url") => {
|
||||
if let Some(url) = image_url(part) {
|
||||
add(url);
|
||||
}
|
||||
}
|
||||
Some("file") if files => {
|
||||
let file = part.get("file");
|
||||
let source = ["file_id", "file_data"].into_iter().find_map(|key| {
|
||||
file.and_then(|file| file.get(key)).and_then(Value::as_str)
|
||||
});
|
||||
if let Some(url) = source {
|
||||
add(url);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
urls
|
||||
}
|
||||
|
||||
/// Every reference to a downloaded URL replaced by its data URL.
|
||||
pub fn inline_media(messages: &mut [ChatMessage], downloaded: &HashMap<String, String>) {
|
||||
if downloaded.is_empty() {
|
||||
return;
|
||||
}
|
||||
for message in messages {
|
||||
let Some(ChatMessageContent::Parts(parts)) = &mut message.content else {
|
||||
continue;
|
||||
};
|
||||
for part in parts {
|
||||
match part.get("type").and_then(Value::as_str) {
|
||||
Some("image_url") => match part.get_mut("image_url") {
|
||||
Some(Value::String(url)) => {
|
||||
if let Some(data) = downloaded.get(url.as_str()) {
|
||||
*url = data.clone();
|
||||
}
|
||||
}
|
||||
Some(Value::Object(image)) => {
|
||||
if let Some(Value::String(url)) = image.get_mut("url")
|
||||
&& let Some(data) = downloaded.get(url.as_str())
|
||||
{
|
||||
*url = data.clone();
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
},
|
||||
Some("file") => {
|
||||
if let Some(Value::Object(file)) = part.get_mut("file") {
|
||||
for key in ["file_id", "file_data"] {
|
||||
if let Some(Value::String(url)) = file.get_mut(key)
|
||||
&& let Some(data) = downloaded.get(url.as_str())
|
||||
{
|
||||
*url = data.clone();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The media type a download is inlined with: the response's own, or one guessed from the
|
||||
/// URL's extension the way Python's `infer_content_type_from_url_and_content` falls back.
|
||||
pub fn downloaded_media_type(url: &str, content_type: &str) -> String {
|
||||
let declared = content_type
|
||||
.split(';')
|
||||
.next()
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
if !declared.is_empty() && declared != "application/octet-stream" && declared.contains('/') {
|
||||
return declared;
|
||||
}
|
||||
let path = url
|
||||
.split(['?', '#'])
|
||||
.next()
|
||||
.unwrap_or(url)
|
||||
.to_ascii_lowercase();
|
||||
let extension = path.rsplit('.').next().unwrap_or_default();
|
||||
match extension {
|
||||
"png" => "image/png",
|
||||
"jpg" | "jpeg" => "image/jpeg",
|
||||
"gif" => "image/gif",
|
||||
"webp" => "image/webp",
|
||||
"pdf" => "application/pdf",
|
||||
"txt" => "text/plain",
|
||||
"csv" => "text/csv",
|
||||
"html" | "htm" => "text/html",
|
||||
"md" => "text/markdown",
|
||||
"mp4" => "video/mp4",
|
||||
_ => "image/jpeg",
|
||||
}
|
||||
.to_owned()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn messages(value: Value) -> Vec<ChatMessage> {
|
||||
serde_json::from_value(value).unwrap()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn remote_urls_are_listed_once_and_inlined() {
|
||||
let mut list = messages(json!([
|
||||
{"role": "user", "content": [
|
||||
{"type": "image_url", "image_url": "http://a/x.png"},
|
||||
{"type": "image_url", "image_url": {"url": "http://a/x.png"}},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}},
|
||||
{"type": "file", "file": {"file_id": "https://b/doc.pdf"}}
|
||||
]},
|
||||
{"role": "assistant", "content": [{"type": "image_url", "image_url": "http://c"}]}
|
||||
]));
|
||||
let urls = remote_urls(&list, &["user", "tool"], true, |url| {
|
||||
url.starts_with("http")
|
||||
});
|
||||
assert_eq!(urls, vec!["http://a/x.png", "https://b/doc.pdf"]);
|
||||
let downloaded = HashMap::from([(
|
||||
"http://a/x.png".to_owned(),
|
||||
"data:image/png;base64,Zm8=".to_owned(),
|
||||
)]);
|
||||
inline_media(&mut list, &downloaded);
|
||||
let content = serde_json::to_value(&list[0].content).unwrap();
|
||||
assert_eq!(content[0]["image_url"], "data:image/png;base64,Zm8=");
|
||||
assert_eq!(content[1]["image_url"]["url"], "data:image/png;base64,Zm8=");
|
||||
assert_eq!(content[3]["file"]["file_id"], "https://b/doc.pdf");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("data:image/png;base64,AAAA", Some(("image/png", "AAAA")))]
|
||||
#[case("https://x", None)]
|
||||
fn data_urls_split(#[case] url: &str, #[case] expected: Option<(&str, &str)>) {
|
||||
assert_eq!(data_url(url), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case("https://x/a.PNG?x=1", "", "image/png")]
|
||||
#[case("https://x/a", "image/webp; charset=x", "image/webp")]
|
||||
#[case("https://x/a.pdf", "application/octet-stream", "application/pdf")]
|
||||
fn media_types_fall_back_to_the_extension(
|
||||
#[case] url: &str,
|
||||
#[case] content_type: &str,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
assert_eq!(downloaded_media_type(url, content_type), expected);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,2 +1,5 @@
|
|||
pub mod media;
|
||||
pub mod params;
|
||||
pub mod response_format;
|
||||
pub mod streaming;
|
||||
pub mod transformation;
|
||||
|
|
|
|||
434
litellm-rust/crates/llms/src/base_llm/chat/params.rs
Normal file
434
litellm-rust/crates/llms/src/base_llm/chat/params.rs
Normal file
|
|
@ -0,0 +1,434 @@
|
|||
//! Python's `get_optional_params` for the providers chat translates: the request's keys are
|
||||
//! sorted into OpenAI params, keys LiteLLM consumes itself, and provider-specific keys that
|
||||
//! pass through; an OpenAI param the provider does not take is dropped under `drop_params`
|
||||
//! and refused otherwise.
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{
|
||||
Error, ErrorDetail, base_llm::model_info::ModelInfo,
|
||||
openai::chat::transformation::is_litellm_owned,
|
||||
};
|
||||
|
||||
/// The keys of `DEFAULT_CHAT_COMPLETION_PARAM_VALUES` a request carries as OpenAI params, in
|
||||
/// that table's order, which is the order Python maps them in.
|
||||
pub const OPENAI_PARAMS: &[&str] = &[
|
||||
"functions",
|
||||
"function_call",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"n",
|
||||
"stream",
|
||||
"stream_options",
|
||||
"stop",
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"modalities",
|
||||
"prediction",
|
||||
"audio",
|
||||
"presence_penalty",
|
||||
"frequency_penalty",
|
||||
"logit_bias",
|
||||
"user",
|
||||
"response_format",
|
||||
"seed",
|
||||
"tools",
|
||||
"tool_choice",
|
||||
"max_retries",
|
||||
"logprobs",
|
||||
"top_logprobs",
|
||||
"extra_headers",
|
||||
"parallel_tool_calls",
|
||||
"reasoning_effort",
|
||||
"verbosity",
|
||||
"thinking",
|
||||
"web_search_options",
|
||||
"include_server_side_tool_invocations",
|
||||
"service_tier",
|
||||
"safety_identifier",
|
||||
"prompt_cache_key",
|
||||
"prompt_cache_retention",
|
||||
"store",
|
||||
"metadata",
|
||||
"context_management",
|
||||
];
|
||||
|
||||
/// `PROVIDER_UNVALIDATED_PARAMS`: never refused, only left out when unsupported.
|
||||
pub const UNVALIDATED_PARAMS: &[&str] = &["user", "stream_options", "stream", "max_retries"];
|
||||
|
||||
/// The tool Python adds when `response_format` asks for JSON and the model has no native
|
||||
/// structured output.
|
||||
pub const RESPONSE_FORMAT_TOOL_NAME: &str = "json_tool_call";
|
||||
|
||||
/// Python's `DEFAULT_MAX_TOKENS`, added to a thinking budget when the caller set no limit.
|
||||
pub const DEFAULT_MAX_TOKENS: u64 = 4096;
|
||||
|
||||
/// The request's OpenAI params in `DEFAULT_CHAT_COMPLETION_PARAM_VALUES` order.
|
||||
pub fn ordered(params: &Map<String, Value>) -> impl Iterator<Item = (&'static str, &Value)> {
|
||||
OPENAI_PARAMS
|
||||
.iter()
|
||||
.filter_map(|name| params.get(*name).map(|value| (*name, value)))
|
||||
}
|
||||
|
||||
/// Python's truthiness for a JSON value.
|
||||
pub fn truthy(value: &Value) -> bool {
|
||||
match value {
|
||||
Value::Null => false,
|
||||
Value::Bool(flag) => *flag,
|
||||
Value::Number(number) => number.as_f64() != Some(0.0),
|
||||
Value::String(text) => !text.is_empty(),
|
||||
Value::Array(items) => !items.is_empty(),
|
||||
Value::Object(object) => !object.is_empty(),
|
||||
}
|
||||
}
|
||||
|
||||
/// `normalize_drop_params`: a boolean, or a config's `"true"`.
|
||||
pub fn request_drop_params(value: Option<Value>) -> bool {
|
||||
match value {
|
||||
Some(Value::Bool(flag)) => flag,
|
||||
Some(Value::String(text)) => text.trim().eq_ignore_ascii_case("true"),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn string_list(value: Option<Value>) -> Vec<String> {
|
||||
match value {
|
||||
Some(Value::Array(items)) => items
|
||||
.into_iter()
|
||||
.filter_map(|item| match item {
|
||||
Value::String(text) => Some(text),
|
||||
_ => None,
|
||||
})
|
||||
.collect(),
|
||||
_ => Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// What the host tells a config about the request beyond its params.
|
||||
#[derive(Clone, Copy, Default)]
|
||||
pub struct ChatRequestContext<'a> {
|
||||
/// `drop_params` from the proxy settings or the deployment. `None` when the host cannot
|
||||
/// say, as when Python hosts the call: a request Python would drop or refuse a param of
|
||||
/// then goes back to Python instead.
|
||||
pub drop_params: Option<bool>,
|
||||
/// The model cost map, which several param rules read.
|
||||
pub model_info: Option<&'a dyn ModelInfo>,
|
||||
}
|
||||
|
||||
/// Whether a param the provider does not take is dropped, refused, or left to Python.
|
||||
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
|
||||
pub struct DropPolicy(Option<bool>);
|
||||
|
||||
impl DropPolicy {
|
||||
pub fn new(drop_params: Option<bool>) -> Self {
|
||||
Self(drop_params)
|
||||
}
|
||||
|
||||
/// Whether unsupported params are dropped.
|
||||
pub fn drops(self) -> bool {
|
||||
self.0 == Some(true)
|
||||
}
|
||||
|
||||
/// Whether the host said what `drop_params` is.
|
||||
pub fn known(self) -> bool {
|
||||
self.0.is_some()
|
||||
}
|
||||
|
||||
/// A param Python drops under `drop_params` and refuses otherwise: `Ok` means drop it.
|
||||
pub fn unsupported(self, refusal: impl FnOnce() -> Error) -> Result<(), Error> {
|
||||
match self.0 {
|
||||
Some(true) => Ok(()),
|
||||
Some(false) => Err(refusal()),
|
||||
None => Err(Error::Unsupported(
|
||||
"a param Python drops or refuses by drop_params",
|
||||
)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Python's `UnsupportedParamsError` for one param with its own message.
|
||||
pub fn unsupported_value(message: impl Into<String>) -> Error {
|
||||
Error::InvalidRequest(ErrorDetail::Message(message.into()))
|
||||
}
|
||||
|
||||
/// A request's params as `get_optional_params` sorts them.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct SortedParams {
|
||||
/// The OpenAI params, refused or dropped ones removed.
|
||||
pub openai: Map<String, Value>,
|
||||
/// Every other key, sent to the provider as given.
|
||||
pub provider: Map<String, Value>,
|
||||
/// Keys the config consumes as call configuration.
|
||||
pub config: Map<String, Value>,
|
||||
/// `allowed_openai_params`: sent as given even when the provider does not list them.
|
||||
pub allowed: Vec<String>,
|
||||
pub drop: DropPolicy,
|
||||
}
|
||||
|
||||
/// Sorts `params` and refuses or drops the OpenAI params `supported` leaves out.
|
||||
pub fn sort_params(
|
||||
params: Map<String, Value>,
|
||||
config_params: &[&str],
|
||||
context: &ChatRequestContext<'_>,
|
||||
supported: impl FnOnce(&Map<String, Value>) -> Result<Vec<&'static str>, Error>,
|
||||
provider: &'static str,
|
||||
model: &str,
|
||||
) -> Result<SortedParams, Error> {
|
||||
let mut params = params;
|
||||
let request_drop = request_drop_params(params.remove("drop_params"));
|
||||
let drop = DropPolicy::new(match context.drop_params {
|
||||
Some(configured) => Some(configured || request_drop),
|
||||
None if request_drop => Some(true),
|
||||
None => None,
|
||||
});
|
||||
let allowed = string_list(params.remove("allowed_openai_params"));
|
||||
let additional_drop = string_list(params.remove("additional_drop_params"));
|
||||
let mut sorted = SortedParams {
|
||||
allowed,
|
||||
drop,
|
||||
..SortedParams::default()
|
||||
};
|
||||
for (name, value) in params {
|
||||
if config_params.contains(&name.as_str()) {
|
||||
sorted.config.insert(name, value);
|
||||
continue;
|
||||
}
|
||||
if value.is_null() || is_litellm_owned(&name) || additional_drop.contains(&name) {
|
||||
continue;
|
||||
}
|
||||
match OPENAI_PARAMS.contains(&name.as_str()) {
|
||||
true => sorted.openai.insert(name, value),
|
||||
false => sorted.provider.insert(name, value),
|
||||
};
|
||||
}
|
||||
let supported = supported(&sorted.openai)?;
|
||||
let unsupported: Vec<String> = sorted
|
||||
.openai
|
||||
.iter()
|
||||
.filter(|(name, _)| {
|
||||
!supported.contains(&name.as_str()) && !sorted.allowed.iter().any(|a| a == *name)
|
||||
})
|
||||
.map(|(name, _)| name.clone())
|
||||
.collect();
|
||||
let refused: Vec<String> = unsupported
|
||||
.iter()
|
||||
.filter(|name| {
|
||||
!UNVALIDATED_PARAMS.contains(&name.as_str())
|
||||
&& !(name.as_str() == "n"
|
||||
&& sorted.openai.get(name.as_str()).and_then(Value::as_f64) == Some(1.0))
|
||||
})
|
||||
.cloned()
|
||||
.collect();
|
||||
if !refused.is_empty() {
|
||||
drop.unsupported(|| {
|
||||
Error::InvalidRequest(ErrorDetail::UnsupportedParams {
|
||||
provider,
|
||||
params: refused.clone(),
|
||||
model: model.to_string(),
|
||||
})
|
||||
})?;
|
||||
}
|
||||
for name in unsupported {
|
||||
sorted.openai.remove(&name);
|
||||
}
|
||||
Ok(sorted)
|
||||
}
|
||||
|
||||
/// A capability flag the cost map sets for the first of `candidates` it knows.
|
||||
///
|
||||
/// Without a cost map Python's own lookups have nothing to go on either, so the answer is
|
||||
/// `None`, except when the host could still fall back to Python, which does have one.
|
||||
pub fn capability(
|
||||
context: &ChatRequestContext<'_>,
|
||||
candidates: &[String],
|
||||
key: &str,
|
||||
) -> Result<Option<bool>, Error> {
|
||||
let Some(info) = context.model_info else {
|
||||
return match context.drop_params {
|
||||
None => Err(Error::Unsupported(
|
||||
"a request whose translation reads the model cost map",
|
||||
)),
|
||||
Some(_) => Ok(None),
|
||||
};
|
||||
};
|
||||
Ok(candidates
|
||||
.iter()
|
||||
.filter_map(|candidate| info.entry(candidate))
|
||||
.find_map(|entry| entry.get(key).and_then(Value::as_bool)))
|
||||
}
|
||||
|
||||
/// A numeric field of the first of `candidates` the cost map knows.
|
||||
pub fn model_field(
|
||||
context: &ChatRequestContext<'_>,
|
||||
candidates: &[String],
|
||||
key: &str,
|
||||
) -> Option<Value> {
|
||||
let info = context.model_info?;
|
||||
candidates
|
||||
.iter()
|
||||
.filter_map(|candidate| info.entry(candidate))
|
||||
.find_map(|entry| entry.get(key).filter(|value| !value.is_null()).cloned())
|
||||
}
|
||||
|
||||
/// Whether the cost map has an entry for any of `candidates`.
|
||||
pub fn model_known(context: &ChatRequestContext<'_>, candidates: &[String]) -> bool {
|
||||
context
|
||||
.model_info
|
||||
.is_some_and(|info| candidates.iter().any(|c| info.entry(c).is_some()))
|
||||
}
|
||||
|
||||
/// `json.dumps(value)` with Python's default separators and ASCII escaping.
|
||||
pub fn python_dumps(value: &Value) -> String {
|
||||
litellm_python_compat::json::dumps(&litellm_python_compat::json::from_json(value.clone()))
|
||||
.unwrap_or_else(|_| value.to_string())
|
||||
}
|
||||
|
||||
/// `json.loads` of tool call arguments, as `parse_tool_call_arguments` reads them: empty
|
||||
/// arguments are `{}`.
|
||||
pub fn parse_tool_arguments(arguments: Option<&Value>) -> Result<Value, Error> {
|
||||
match arguments {
|
||||
None | Some(Value::Null) => Ok(Value::Object(Map::new())),
|
||||
Some(Value::String(text)) if text.trim().is_empty() => Ok(Value::Object(Map::new())),
|
||||
Some(Value::String(text)) => serde_json::from_str(text).map_err(|error| {
|
||||
Error::InvalidRequest(ErrorDetail::Message(format!(
|
||||
"tool call arguments are not valid JSON: {error}"
|
||||
)))
|
||||
}),
|
||||
Some(other) => Ok(other.clone()),
|
||||
}
|
||||
}
|
||||
|
||||
/// `is_thinking_enabled` over a params map.
|
||||
pub fn is_thinking_enabled(params: &Map<String, Value>) -> bool {
|
||||
let thinking = params.get("thinking");
|
||||
thinking == Some(&Value::Bool(true))
|
||||
|| thinking
|
||||
.and_then(|thinking| thinking.get("type"))
|
||||
.and_then(Value::as_str)
|
||||
== Some("enabled")
|
||||
|| params
|
||||
.get("reasoning_effort")
|
||||
.is_some_and(|effort| !effort.is_null())
|
||||
}
|
||||
|
||||
/// The JSON schema a `response_format` names: `response_schema`, or `json_schema.schema`.
|
||||
pub fn response_format_schema(value: &Value) -> Option<Value> {
|
||||
if let Some(schema) = value.get("response_schema") {
|
||||
return Some(schema.clone());
|
||||
}
|
||||
value
|
||||
.get("json_schema")
|
||||
.and_then(|json_schema| json_schema.get("schema"))
|
||||
.cloned()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn object(value: Value) -> Map<String, Value> {
|
||||
match value {
|
||||
Value::Object(map) => map,
|
||||
other => panic!("expected an object, got {other}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn sort(params: Value, drop_params: Option<bool>) -> Result<SortedParams, Error> {
|
||||
sort_params(
|
||||
object(params),
|
||||
&["aws_region_name"],
|
||||
&ChatRequestContext {
|
||||
drop_params,
|
||||
model_info: None,
|
||||
},
|
||||
|_| Ok(vec!["temperature", "max_tokens"]),
|
||||
"anthropic",
|
||||
"claude",
|
||||
)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn params_sort_into_openai_provider_config_and_owned() {
|
||||
let sorted = sort(
|
||||
json!({"temperature": 0.5, "top_k": 3, "aws_region_name": "us-east-1",
|
||||
"metadata": {"x": 1}, "mock_response": "hi", "_litellm_x": 1, "user": "u", "seed": null}),
|
||||
Some(false),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(sorted.openai, object(json!({"temperature": 0.5})));
|
||||
assert_eq!(sorted.provider, object(json!({"top_k": 3})));
|
||||
assert_eq!(
|
||||
sorted.config,
|
||||
object(json!({"aws_region_name": "us-east-1"}))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::refused(Some(false), false)]
|
||||
#[case::dropped(Some(true), true)]
|
||||
#[case::request_drop(None, true)]
|
||||
fn an_unsupported_param_is_refused_or_dropped(
|
||||
#[case] drop_params: Option<bool>,
|
||||
#[case] dropped: bool,
|
||||
) {
|
||||
let mut body = json!({"temperature": 0.5, "seed": 3});
|
||||
if drop_params.is_none() {
|
||||
body["drop_params"] = json!(true);
|
||||
}
|
||||
let result = sort(body, drop_params);
|
||||
match dropped {
|
||||
true => assert_eq!(result.unwrap().openai, object(json!({"temperature": 0.5}))),
|
||||
false => assert!(matches!(
|
||||
result,
|
||||
Err(Error::InvalidRequest(ErrorDetail::UnsupportedParams { ref params, .. }))
|
||||
if params == &vec!["seed".to_string()]
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn an_unknown_drop_policy_falls_back() {
|
||||
assert!(matches!(
|
||||
sort(json!({"seed": 3}), None),
|
||||
Err(Error::Unsupported(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn allowed_and_additional_drop_params_are_honoured() {
|
||||
let sorted = sort(
|
||||
json!({"seed": 3, "top_k": 2, "allowed_openai_params": ["seed"], "additional_drop_params": ["top_k"]}),
|
||||
Some(false),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(sorted.openai, object(json!({"seed": 3})));
|
||||
assert!(sorted.provider.is_empty());
|
||||
assert_eq!(sorted.allowed, vec!["seed".to_string()]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::n_of_one(json!({"n": 1}), true)]
|
||||
#[case::n_of_two(json!({"n": 2}), false)]
|
||||
#[case::user(json!({"user": "u"}), true)]
|
||||
fn unvalidated_params_and_one_choice_are_never_refused(
|
||||
#[case] body: Value,
|
||||
#[case] accepted: bool,
|
||||
) {
|
||||
let result = sort(body, Some(false));
|
||||
assert_eq!(result.is_ok(), accepted, "{result:?}");
|
||||
if let Ok(sorted) = result {
|
||||
assert!(sorted.openai.is_empty());
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn python_dumps_matches_json_dumps() {
|
||||
assert_eq!(
|
||||
python_dumps(&json!({"b": "é", "a": [1, 2.5]})),
|
||||
r#"{"b": "\u00e9", "a": [1, 2.5]}"#
|
||||
);
|
||||
}
|
||||
}
|
||||
445
litellm-rust/crates/llms/src/base_llm/chat/response_format.rs
Normal file
445
litellm-rust/crates/llms/src/base_llm/chat/response_format.rs
Normal file
|
|
@ -0,0 +1,445 @@
|
|||
//! `response_format` answers delivered through the `json_tool_call` tool, as Python's
|
||||
//! `base_utils.py` reads them back: the schema's shape, the envelopes a model wraps the
|
||||
//! answer in, and an answer split across several tool calls.
|
||||
|
||||
use std::collections::BTreeSet;
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::base_llm::chat::params::python_dumps;
|
||||
|
||||
/// `ResponseFormatSchema`: what the requested schema declares, enough to tell an answer
|
||||
/// from an envelope around it.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct ResponseFormatSchema {
|
||||
pub property_names: BTreeSet<String>,
|
||||
pub required_names: BTreeSet<String>,
|
||||
pub allows_undeclared_keys: bool,
|
||||
}
|
||||
|
||||
impl ResponseFormatSchema {
|
||||
/// `response_format_schema`: `None` unless `value` is an object whose `properties`, if
|
||||
/// present, is one too.
|
||||
pub fn from_schema(value: &Value) -> Option<Self> {
|
||||
let document = value.as_object()?;
|
||||
let properties = match document.get("properties") {
|
||||
None => Map::new(),
|
||||
Some(Value::Object(properties)) => properties.clone(),
|
||||
Some(_) => return None,
|
||||
};
|
||||
let required_names = match document.get("required") {
|
||||
Some(Value::Array(items)) if items.iter().all(Value::is_string) => items
|
||||
.iter()
|
||||
.filter_map(Value::as_str)
|
||||
.map(str::to_owned)
|
||||
.collect(),
|
||||
_ => BTreeSet::new(),
|
||||
};
|
||||
let additional = document.get("additionalProperties");
|
||||
let allows_undeclared_keys = additional == Some(&Value::Bool(true))
|
||||
|| additional.is_some_and(Value::is_object)
|
||||
|| document
|
||||
.get("patternProperties")
|
||||
.and_then(Value::as_object)
|
||||
.is_some_and(|patterns| !patterns.is_empty());
|
||||
Some(Self {
|
||||
property_names: properties.keys().cloned().collect(),
|
||||
required_names,
|
||||
allows_undeclared_keys,
|
||||
})
|
||||
}
|
||||
|
||||
fn accepts_a_lone_undeclared_key(&self) -> bool {
|
||||
self.allows_undeclared_keys && self.required_names.is_empty()
|
||||
}
|
||||
}
|
||||
|
||||
const ENVELOPE_NAME_KEYS: [&str; 2] = ["name", "tool_name"];
|
||||
|
||||
/// `_json_object_or_encoded_object`.
|
||||
fn object_or_encoded_object(value: &Value) -> Option<Map<String, Value>> {
|
||||
match value {
|
||||
Value::String(text) => match serde_json::from_str(text) {
|
||||
Ok(Value::Object(object)) => Some(object),
|
||||
_ => None,
|
||||
},
|
||||
Value::Object(object) => Some(object.clone()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
struct EnvelopeLayer {
|
||||
inner: Map<String, Value>,
|
||||
names_the_tool: bool,
|
||||
}
|
||||
|
||||
fn envelope_layer(
|
||||
payload: &Map<String, Value>,
|
||||
schema: &ResponseFormatSchema,
|
||||
tool_name: &str,
|
||||
) -> Option<EnvelopeLayer> {
|
||||
let name_keys: BTreeSet<&str> = ENVELOPE_NAME_KEYS
|
||||
.into_iter()
|
||||
.filter(|key| payload.get(*key).and_then(Value::as_str) == Some(tool_name))
|
||||
.collect();
|
||||
let label_keys: BTreeSet<&str> = match (name_keys.is_empty(), schema.required_names.is_empty())
|
||||
{
|
||||
(false, _) => name_keys.clone(),
|
||||
(true, true) => BTreeSet::new(),
|
||||
(true, false) => payload
|
||||
.iter()
|
||||
.filter(|(key, value)| {
|
||||
value.is_string()
|
||||
&& !schema.property_names.contains(*key)
|
||||
&& !schema.required_names.contains(*key)
|
||||
})
|
||||
.map(|(key, _)| key.as_str())
|
||||
.collect(),
|
||||
};
|
||||
let wrapper_keys: Vec<&String> = payload
|
||||
.keys()
|
||||
.filter(|key| !label_keys.contains(key.as_str()))
|
||||
.collect();
|
||||
let [wrapper] = wrapper_keys.as_slice() else {
|
||||
return None;
|
||||
};
|
||||
if schema.property_names.contains(*wrapper) {
|
||||
return None;
|
||||
}
|
||||
let wrapped = &payload[wrapper.as_str()];
|
||||
let inner = match name_keys.is_empty() {
|
||||
false => object_or_encoded_object(wrapped),
|
||||
true => wrapped.as_object().cloned(),
|
||||
}?;
|
||||
Some(EnvelopeLayer {
|
||||
inner,
|
||||
names_the_tool: !name_keys.is_empty(),
|
||||
})
|
||||
}
|
||||
|
||||
fn holds_the_schema_object(layer: &EnvelopeLayer, schema: &ResponseFormatSchema) -> bool {
|
||||
layer.names_the_tool
|
||||
|| (!schema.accepts_a_lone_undeclared_key()
|
||||
&& layer
|
||||
.inner
|
||||
.keys()
|
||||
.any(|key| schema.property_names.contains(key)))
|
||||
}
|
||||
|
||||
fn schema_object_inside_envelope(
|
||||
payload: &Map<String, Value>,
|
||||
schema: &ResponseFormatSchema,
|
||||
tool_name: &str,
|
||||
) -> Option<Map<String, Value>> {
|
||||
let mut layer = envelope_layer(payload, schema, tool_name);
|
||||
let mut first = true;
|
||||
while let Some(current) = layer {
|
||||
if !first && schema.property_names.is_empty() {
|
||||
return None;
|
||||
}
|
||||
first = false;
|
||||
if holds_the_schema_object(¤t, schema) {
|
||||
return Some(current.inner);
|
||||
}
|
||||
layer = envelope_layer(¤t.inner, schema, tool_name);
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// `response_format_envelope_contents`.
|
||||
pub fn envelope_contents(
|
||||
payload: &Value,
|
||||
schema: &ResponseFormatSchema,
|
||||
tool_name: &str,
|
||||
) -> Option<Value> {
|
||||
let object = payload.as_object()?;
|
||||
if let Some(inner) = schema_object_inside_envelope(object, schema, tool_name) {
|
||||
return Some(Value::Object(inner));
|
||||
}
|
||||
if schema.property_names.contains("values") {
|
||||
return None;
|
||||
}
|
||||
object
|
||||
.get("values")
|
||||
.filter(|values| !values.is_null())
|
||||
.cloned()
|
||||
}
|
||||
|
||||
/// `unwrap_response_format_envelope`.
|
||||
pub fn unwrap_envelope(payload: Value, schema: &ResponseFormatSchema, tool_name: &str) -> Value {
|
||||
envelope_contents(&payload, schema, tool_name).unwrap_or(payload)
|
||||
}
|
||||
|
||||
fn field_names(
|
||||
blocks: &[Map<String, Value>],
|
||||
key: &str,
|
||||
schema: &ResponseFormatSchema,
|
||||
) -> Option<Vec<String>> {
|
||||
let names: Vec<String> = blocks
|
||||
.iter()
|
||||
.filter_map(|block| block[key].as_str())
|
||||
.filter(|name| schema.property_names.contains(*name))
|
||||
.map(str::to_owned)
|
||||
.collect();
|
||||
let distinct: BTreeSet<&String> = names.iter().collect();
|
||||
(distinct.len() == blocks.len()).then_some(names)
|
||||
}
|
||||
|
||||
fn is_a_label_every_block_repeats(blocks: &[Map<String, Value>], key: &str) -> bool {
|
||||
blocks[0][key].is_string() && blocks.iter().all(|block| block[key] == blocks[0][key])
|
||||
}
|
||||
|
||||
fn value_key(blocks: &[Map<String, Value>], keys: &[&String]) -> Option<String> {
|
||||
if let [only] = keys {
|
||||
return Some((*only).clone());
|
||||
}
|
||||
let varying: Vec<&&String> = keys
|
||||
.iter()
|
||||
.filter(|key| !is_a_label_every_block_repeats(blocks, key))
|
||||
.collect();
|
||||
match varying.as_slice() {
|
||||
[only] => Some((**only).clone()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn split_answer_fields(
|
||||
blocks: &[Map<String, Value>],
|
||||
schema: &ResponseFormatSchema,
|
||||
) -> Option<Map<String, Value>> {
|
||||
let keys: BTreeSet<&String> = blocks[0].keys().collect();
|
||||
if keys.iter().any(|key| schema.property_names.contains(*key))
|
||||
|| blocks
|
||||
.iter()
|
||||
.any(|block| block.keys().collect::<BTreeSet<_>>() != keys)
|
||||
{
|
||||
return None;
|
||||
}
|
||||
// Python iterates a frozenset here; the outcome only depends on there being one match.
|
||||
let named: Vec<(&String, Vec<String>)> = keys
|
||||
.iter()
|
||||
.filter_map(|key| Some((*key, field_names(blocks, key, schema)?)))
|
||||
.collect();
|
||||
let [(name_key, names)] = named.as_slice() else {
|
||||
return None;
|
||||
};
|
||||
let rest: Vec<&String> = keys.iter().copied().filter(|key| key != name_key).collect();
|
||||
let value_key = value_key(blocks, &rest)?;
|
||||
Some(
|
||||
names
|
||||
.iter()
|
||||
.cloned()
|
||||
.zip(blocks.iter().map(|block| block[value_key.as_str()].clone()))
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
fn merged_partial_answers(
|
||||
blocks: &[Map<String, Value>],
|
||||
schema: &ResponseFormatSchema,
|
||||
) -> Option<Map<String, Value>> {
|
||||
let keys: BTreeSet<&String> = blocks.iter().flat_map(Map::keys).collect();
|
||||
if !keys.iter().all(|key| schema.property_names.contains(*key))
|
||||
|| keys.len() != blocks.iter().map(Map::len).sum::<usize>()
|
||||
{
|
||||
return None;
|
||||
}
|
||||
Some(
|
||||
blocks
|
||||
.iter()
|
||||
.flat_map(|block| block.iter().map(|(k, v)| (k.clone(), v.clone())))
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
/// `split_response_format_answer`: one answer spread over several tool calls.
|
||||
pub fn split_answer(
|
||||
encoded_blocks: &[String],
|
||||
schema: &ResponseFormatSchema,
|
||||
) -> Option<Map<String, Value>> {
|
||||
let blocks: Vec<Map<String, Value>> = encoded_blocks
|
||||
.iter()
|
||||
.map(|encoded| object_or_encoded_object(&Value::String(encoded.clone())))
|
||||
.collect::<Option<_>>()?;
|
||||
if blocks.len() < 2 {
|
||||
return None;
|
||||
}
|
||||
split_answer_fields(&blocks, schema).or_else(|| merged_partial_answers(&blocks, schema))
|
||||
}
|
||||
|
||||
/// Anthropic's `_convert_tool_response_to_message`: the message content the `json_tool_call`
|
||||
/// calls' arguments become.
|
||||
pub fn json_tool_answer(arguments: &[String], schema: Option<&ResponseFormatSchema>) -> String {
|
||||
let first = arguments.first().cloned().unwrap_or_default();
|
||||
if let Some(split) = schema.and_then(|schema| split_answer(arguments, schema)) {
|
||||
return python_dumps(&Value::Object(split));
|
||||
}
|
||||
let Ok(payload) = serde_json::from_str::<Value>(&first) else {
|
||||
return first;
|
||||
};
|
||||
match schema {
|
||||
Some(schema) => python_dumps(&unwrap_envelope(
|
||||
payload,
|
||||
schema,
|
||||
crate::base_llm::chat::params::RESPONSE_FORMAT_TOOL_NAME,
|
||||
)),
|
||||
None => {
|
||||
let values = payload
|
||||
.as_object()
|
||||
.and_then(|object| object.get("values"))
|
||||
.filter(|values| !values.is_null())
|
||||
.cloned();
|
||||
python_dumps(values.as_ref().unwrap_or(&payload))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `unpack_defs`: every `$ref` into `$defs` / `definitions` replaced by a copy of its
|
||||
/// target, nested definitions shadowing outer ones, a circular reference left as it is.
|
||||
pub fn unpack_defs(schema: &mut Value, defs: &Map<String, Value>) {
|
||||
fn scoped(node: &Map<String, Value>, outer: &Map<String, Value>) -> Map<String, Value> {
|
||||
let mut defs = outer.clone();
|
||||
for key in ["$defs", "definitions"] {
|
||||
if let Some(Value::Object(local)) = node.get(key) {
|
||||
defs.extend(local.iter().map(|(k, v)| (k.clone(), v.clone())));
|
||||
}
|
||||
}
|
||||
defs
|
||||
}
|
||||
fn walk(node: &mut Value, defs: &Map<String, Value>, chain: &BTreeSet<String>) {
|
||||
match node {
|
||||
Value::Object(object) => {
|
||||
if let Some(reference) = object.get("$ref").and_then(Value::as_str) {
|
||||
let name = reference.rsplit('/').next().unwrap_or(reference).to_owned();
|
||||
if chain.contains(&name) {
|
||||
return;
|
||||
}
|
||||
let Some(target) = defs.get(&name).cloned() else {
|
||||
return;
|
||||
};
|
||||
let child_defs = match &target {
|
||||
Value::Object(target) => scoped(target, defs),
|
||||
_ => defs.clone(),
|
||||
};
|
||||
*node = target;
|
||||
let mut chain = chain.clone();
|
||||
chain.insert(name);
|
||||
walk(node, &child_defs, &chain);
|
||||
return;
|
||||
}
|
||||
if object.contains_key("$ref") {
|
||||
return;
|
||||
}
|
||||
let current = scoped(object, defs);
|
||||
for value in object.values_mut() {
|
||||
walk(value, ¤t, chain);
|
||||
}
|
||||
}
|
||||
Value::Array(items) => {
|
||||
for item in items {
|
||||
walk(item, defs, chain);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
let root = match schema {
|
||||
Value::Object(object) => scoped(object, defs),
|
||||
_ => defs.clone(),
|
||||
};
|
||||
walk(schema, &root, &BTreeSet::new());
|
||||
}
|
||||
|
||||
/// `unpack_legacy_defs`: refs into `definitions` or OpenAPI `components.schemas` inlined,
|
||||
/// those blocks removed. `$defs` is left alone.
|
||||
pub fn unpack_legacy_defs(schema: &mut Map<String, Value>) {
|
||||
let has_components = schema
|
||||
.get("components")
|
||||
.and_then(|components| components.get("schemas"))
|
||||
.is_some();
|
||||
if !schema.contains_key("definitions") && !has_components {
|
||||
return;
|
||||
}
|
||||
let mut defs = match schema.remove("components") {
|
||||
Some(Value::Object(mut components)) => match components.remove("schemas") {
|
||||
Some(Value::Object(schemas)) => schemas,
|
||||
_ => Map::new(),
|
||||
},
|
||||
_ => Map::new(),
|
||||
};
|
||||
if let Some(Value::Object(definitions)) = schema.remove("definitions") {
|
||||
defs.extend(definitions);
|
||||
}
|
||||
let mut value = Value::Object(std::mem::take(schema));
|
||||
unpack_defs(&mut value, &defs);
|
||||
if let Value::Object(object) = value {
|
||||
*schema = object;
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn schema(value: Value) -> ResponseFormatSchema {
|
||||
ResponseFormatSchema::from_schema(&value).unwrap()
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::bare(json!({"name": "a"}), json!({"name": "a"}))]
|
||||
#[case::named_envelope(json!({"name": "json_tool_call", "arguments": {"name": "a"}}), json!({"name": "a"}))]
|
||||
#[case::encoded_envelope(json!({"tool_name": "json_tool_call", "parameters": "{\"name\": \"a\"}"}), json!({"name": "a"}))]
|
||||
#[case::wrapper(json!({"answer": {"name": "a"}}), json!({"name": "a"}))]
|
||||
#[case::values(json!({"values": {"x": 1}}), json!({"x": 1}))]
|
||||
fn envelopes_are_unwrapped(#[case] payload: Value, #[case] expected: Value) {
|
||||
let schema = schema(
|
||||
json!({"type": "object", "properties": {"name": {"type": "string"}}, "required": ["name"]}),
|
||||
);
|
||||
assert_eq!(
|
||||
unwrap_envelope(payload, &schema, "json_tool_call"),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn an_answer_split_over_named_fields_is_joined() {
|
||||
let schema = schema(json!({"properties": {"a": {}, "b": {}}}));
|
||||
let blocks = vec![
|
||||
r#"{"field": "a", "value": 1}"#.to_owned(),
|
||||
r#"{"field": "b", "value": 2}"#.to_owned(),
|
||||
];
|
||||
assert_eq!(
|
||||
split_answer(&blocks, &schema),
|
||||
Some(json!({"a": 1, "b": 2}).as_object().unwrap().clone())
|
||||
);
|
||||
let partial = vec![r#"{"a": 1}"#.to_owned(), r#"{"b": 2}"#.to_owned()];
|
||||
assert_eq!(
|
||||
json_tool_answer(&partial, Some(&schema)),
|
||||
r#"{"a": 1, "b": 2}"#
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn an_undecodable_answer_is_returned_as_sent() {
|
||||
assert_eq!(json_tool_answer(&["{bad".to_owned()], None), "{bad");
|
||||
assert_eq!(
|
||||
json_tool_answer(&[r#"{"values": [1]}"#.to_owned()], None),
|
||||
"[1]"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn refs_are_inlined_and_cycles_left_alone() {
|
||||
let mut value =
|
||||
json!({"properties": {"a": {"$ref": "#/$defs/A"}, "n": {"$ref": "#/$defs/Node"}}});
|
||||
let defs = json!({"A": {"type": "string"}, "Node": {"properties": {"next": {"$ref": "#/$defs/Node"}}}});
|
||||
unpack_defs(&mut value, defs.as_object().unwrap());
|
||||
assert_eq!(
|
||||
value,
|
||||
json!({"properties": {"a": {"type": "string"}, "n": {"properties": {"next": {"$ref": "#/$defs/Node"}}}}})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,5 +1,7 @@
|
|||
use std::collections::HashMap;
|
||||
|
||||
use crate::base_llm::chat::response_format::ResponseFormatSchema;
|
||||
|
||||
use bytes::Bytes;
|
||||
use futures_util::{StreamExt, stream::BoxStream};
|
||||
use litellm_llms_types::formats::chat_completions::ChatCompletionChunk;
|
||||
|
|
@ -35,6 +37,12 @@ pub struct StreamShape {
|
|||
pub json_mode: bool,
|
||||
pub speed: Option<String>,
|
||||
pub tool_name_reverse_map: HashMap<String, String>,
|
||||
/// The schema of the `json_tool_call` tool a `response_format` added, which tells its
|
||||
/// answer apart from an envelope around it.
|
||||
pub response_format_schema: Option<ResponseFormatSchema>,
|
||||
/// The text of a trailing assistant message sent with `prefix: true`, which Anthropic
|
||||
/// continues and Python puts back in front of the answer.
|
||||
pub prefix_prompt: Option<String>,
|
||||
}
|
||||
|
||||
/// A wire decoder paired with the iterator that turns its events into chat chunks.
|
||||
|
|
|
|||
|
|
@ -5,7 +5,10 @@ use serde_json::{Map, Value};
|
|||
|
||||
use crate::{
|
||||
Error,
|
||||
base_llm::chat::streaming::{ChatStream, StreamShape},
|
||||
base_llm::chat::{
|
||||
params::ChatRequestContext,
|
||||
streaming::{ChatStream, StreamShape},
|
||||
},
|
||||
};
|
||||
|
||||
/// The provider-shaped request body a config produces. Named rather than a bare
|
||||
|
|
@ -14,6 +17,8 @@ use crate::{
|
|||
pub struct ProviderChatRequestData {
|
||||
pub body: Value,
|
||||
pub stream_shape: StreamShape,
|
||||
/// Headers the request needs beyond the config's defaults, such as `anthropic-beta`.
|
||||
pub headers: Vec<(String, String)>,
|
||||
}
|
||||
|
||||
/// The raw provider response body handed back to a config for normalization.
|
||||
|
|
@ -42,8 +47,11 @@ pub struct Unsupported(pub &'static str);
|
|||
pub trait BaseConfig: Sync {
|
||||
fn secret_names(&self) -> Vec<&'static str>;
|
||||
|
||||
/// Supported OpenAI parameter names paired with their provider names.
|
||||
fn supported_openai_param_mappings(&self) -> &'static [(&'static str, &'static str)];
|
||||
/// Supported OpenAI parameter names paired with their provider names, for a config
|
||||
/// whose requests the host checks against an allowlist.
|
||||
fn supported_openai_param_mappings(&self) -> &'static [(&'static str, &'static str)] {
|
||||
&[]
|
||||
}
|
||||
|
||||
fn get_complete_url(
|
||||
&self,
|
||||
|
|
@ -58,14 +66,23 @@ pub trait BaseConfig: Sync {
|
|||
model: &str,
|
||||
messages: Vec<ChatMessage>,
|
||||
optional_params: Map<String, Value>,
|
||||
context: &ChatRequestContext<'_>,
|
||||
) -> Result<ProviderChatRequestData, Error>;
|
||||
|
||||
/// `shape` is what the request decided about its answer, the same one a stream reads.
|
||||
fn transform_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response: ProviderChatResponseData,
|
||||
shape: &StreamShape,
|
||||
) -> Result<ChatCompletionsResponse, Error>;
|
||||
|
||||
/// The media URLs the messages point at that Python downloads while building this
|
||||
/// provider's request. The host fetches them and inlines each as a data URL first.
|
||||
fn remote_media(&self, _messages: &[ChatMessage]) -> Vec<String> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
/// `None` means this config has no streaming path yet, so the host keeps the request.
|
||||
fn model_response_iterator(&self, _shape: StreamShape) -> Option<ChatStream> {
|
||||
None
|
||||
|
|
|
|||
|
|
@ -1,11 +1,14 @@
|
|||
//! What the model cost map says about a model, for configs whose request rules depend on it
|
||||
//! the way Python's `supports_*` helpers read `litellm.model_cost`.
|
||||
|
||||
use std::borrow::Cow;
|
||||
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub trait ModelInfo: Send + Sync {
|
||||
/// The cost map entry `model` names, if there is one.
|
||||
fn entry(&self, model: &str) -> Option<&Map<String, Value>>;
|
||||
/// The cost map entry `model` names, if there is one. A map held in memory lends it; one
|
||||
/// read from elsewhere, such as Python's `litellm.model_cost`, hands over a copy.
|
||||
fn entry(&self, model: &str) -> Option<Cow<'_, Map<String, Value>>>;
|
||||
|
||||
/// Whether any entry of the map sets `key`.
|
||||
fn declares(&self, key: &str) -> bool;
|
||||
|
|
@ -14,21 +17,22 @@ pub trait ModelInfo: Send + Sync {
|
|||
/// Python's `_supports_factory`: the entry sets `key` to `true`. An unknown model supports
|
||||
/// nothing.
|
||||
pub fn supports(info: Option<&dyn ModelInfo>, model: &str, key: &str) -> bool {
|
||||
field(info, model, key).and_then(Value::as_bool) == Some(true)
|
||||
field(info, model, key).and_then(|value| value.as_bool()) == Some(true)
|
||||
}
|
||||
|
||||
/// Python's `is_explicitly_disabled_factory`: the entry sets `key` to `false`. An unknown
|
||||
/// model or an unset key is not disabled.
|
||||
pub fn explicitly_disabled(info: Option<&dyn ModelInfo>, model: &str, key: &str) -> bool {
|
||||
field(info, model, key).and_then(Value::as_bool) == Some(false)
|
||||
field(info, model, key).and_then(|value| value.as_bool()) == Some(false)
|
||||
}
|
||||
|
||||
/// The value the entry for `model` sets for `key`.
|
||||
pub fn field<'a>(info: Option<&'a dyn ModelInfo>, model: &str, key: &str) -> Option<&'a Value> {
|
||||
pub fn field(info: Option<&dyn ModelInfo>, model: &str, key: &str) -> Option<Value> {
|
||||
info?
|
||||
.entry(model)?
|
||||
.get(key)
|
||||
.filter(|value| !value.is_null())
|
||||
.cloned()
|
||||
}
|
||||
|
||||
/// Whether the map has an entry for `model`.
|
||||
|
|
|
|||
1343
litellm-rust/crates/llms/src/bedrock/chat/converse_messages.rs
Normal file
1343
litellm-rust/crates/llms/src/bedrock/chat/converse_messages.rs
Normal file
File diff suppressed because it is too large
Load diff
1039
litellm-rust/crates/llms/src/bedrock/chat/converse_params.rs
Normal file
1039
litellm-rust/crates/llms/src/bedrock/chat/converse_params.rs
Normal file
File diff suppressed because it is too large
Load diff
File diff suppressed because it is too large
Load diff
|
|
@ -23,9 +23,13 @@ use crate::{
|
|||
anthropic::chat::handler::ModelResponseIterator,
|
||||
base_llm::{
|
||||
base_model_iterator::StreamTransformer,
|
||||
chat::streaming::{ChatStream, StreamShape},
|
||||
chat::{
|
||||
params::RESPONSE_FORMAT_TOOL_NAME,
|
||||
streaming::{ChatStream, StreamShape},
|
||||
},
|
||||
messages::streaming::{ByteStream, EventStream},
|
||||
},
|
||||
bedrock::chat::converse_transformation::unwrap_bedrock_properties,
|
||||
};
|
||||
|
||||
#[derive(Deserialize)]
|
||||
|
|
@ -200,7 +204,13 @@ pub fn converse_event_stream(bytes: ByteStream) -> BoxStream<'static, Result<Val
|
|||
/// ConverseStream events to Chat Completions chunks, as Python's `AWSEventStreamDecoder`
|
||||
/// `converse_chunk_parser` builds them.
|
||||
///
|
||||
/// DIVERGENCE: `json_mode` tool conversion is not done (the route never asks for it),
|
||||
/// With `json_mode` the `json_tool_call` tool a `response_format` became is streamed back as
|
||||
/// text: its argument fragments are held and the whole answer (with a lone `properties`
|
||||
/// envelope unwrapped) goes out as content when the block stops, and a `tool_use` stop with
|
||||
/// no other tool call finishes as `stop`.
|
||||
///
|
||||
/// DIVERGENCE: Python fake-streams a json_mode Converse call (one non-streaming request
|
||||
/// replayed as chunks), so its chunk boundaries differ; the content and finish reason match.
|
||||
/// `redactedContent` is passed through without unwrapping, and the reasoning token estimate
|
||||
/// Python makes from streamed thinking text is not made.
|
||||
pub struct ConverseStreamIterator {
|
||||
|
|
@ -211,6 +221,10 @@ pub struct ConverseStreamIterator {
|
|||
thinking_ran: bool,
|
||||
thinking_block_index: i64,
|
||||
thinking_content_block_index: Option<u64>,
|
||||
/// The block in progress is the `json_tool_call` tool a `response_format` became.
|
||||
in_response_format_tool: bool,
|
||||
/// Some `json_tool_call` block has been streamed back as text.
|
||||
converted_response_format_tool: bool,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
|
|
@ -233,6 +247,8 @@ impl ConverseStreamIterator {
|
|||
thinking_ran: false,
|
||||
thinking_block_index: -1,
|
||||
thinking_content_block_index: None,
|
||||
in_response_format_tool: false,
|
||||
converted_response_format_tool: false,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -254,6 +270,7 @@ impl ConverseStreamIterator {
|
|||
index: self
|
||||
.tool_calls_index
|
||||
.unwrap_or_else(|| i64::try_from(index).unwrap_or(i64::MAX)),
|
||||
provider_specific_fields: None,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -295,7 +312,15 @@ impl ConverseStreamIterator {
|
|||
fn start(&mut self, start: &Value, index: u64) -> ConverseParsed {
|
||||
let mut parsed = ConverseParsed::default();
|
||||
self.content_blocks.clear();
|
||||
self.in_response_format_tool = false;
|
||||
if let Some(tool) = start.get("toolUse").filter(|tool| !tool.is_null()) {
|
||||
if self.shape.json_mode
|
||||
&& tool.get("name").and_then(Value::as_str) == Some(RESPONSE_FORMAT_TOOL_NAME)
|
||||
{
|
||||
self.in_response_format_tool = true;
|
||||
self.converted_response_format_tool = true;
|
||||
return parsed;
|
||||
}
|
||||
self.tool_calls_index = Some(self.tool_calls_index.map_or(0, |index| index + 1));
|
||||
let name = tool.get("name").and_then(Value::as_str).map(|name| {
|
||||
self.shape
|
||||
|
|
@ -329,6 +354,9 @@ impl ConverseStreamIterator {
|
|||
.unwrap_or_default()
|
||||
.to_owned();
|
||||
self.content_blocks.push(input.clone());
|
||||
if self.in_response_format_tool {
|
||||
return parsed;
|
||||
}
|
||||
parsed.tool_use = Some(self.tool_call(None, None, input, index));
|
||||
} else if let Some(reasoning) = delta.get("reasoningContent") {
|
||||
parsed
|
||||
|
|
@ -359,6 +387,13 @@ impl ConverseStreamIterator {
|
|||
self.start(start, index)
|
||||
} else if let Some(delta) = data.get("delta") {
|
||||
self.delta(delta, index)
|
||||
} else if data.get("contentBlockIndex").is_some() && self.in_response_format_tool {
|
||||
self.in_response_format_tool = false;
|
||||
let answer = std::mem::take(&mut self.content_blocks).concat();
|
||||
ConverseParsed {
|
||||
text: unwrap_bedrock_properties(if answer.is_empty() { "{}" } else { &answer }),
|
||||
..ConverseParsed::default()
|
||||
}
|
||||
} else if data.get("contentBlockIndex").is_some() {
|
||||
// A block's stop: a tool that streamed only empty fragments closes with `{}`.
|
||||
let empty =
|
||||
|
|
@ -368,10 +403,15 @@ impl ConverseStreamIterator {
|
|||
..ConverseParsed::default()
|
||||
}
|
||||
} else if let Some(reason) = data.get("stopReason") {
|
||||
let mut finish = finish_reason_for(reason.as_str().unwrap_or("stop"));
|
||||
if finish == "tool_calls"
|
||||
&& self.converted_response_format_tool
|
||||
&& self.tool_calls_index.is_none()
|
||||
{
|
||||
finish = "stop";
|
||||
}
|
||||
ConverseParsed {
|
||||
finish_reason: Some(
|
||||
finish_reason_for(reason.as_str().unwrap_or("stop")).to_owned(),
|
||||
),
|
||||
finish_reason: Some(finish.to_owned()),
|
||||
..ConverseParsed::default()
|
||||
}
|
||||
} else if let Some(usage) = data.get("usage") {
|
||||
|
|
@ -735,6 +775,39 @@ mod tests {
|
|||
assert_eq!(chunks[7]["choices"][0]["finish_reason"], "tool_calls");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_response_format_tool_streams_back_as_text() {
|
||||
let mut iterator = ConverseStreamIterator::new(StreamShape {
|
||||
json_mode: true,
|
||||
..StreamShape::default()
|
||||
});
|
||||
let chunks: Vec<serde_json::Value> = serde_json::json!([
|
||||
{"contentBlockIndex": 0, "start": {"toolUse": {"toolUseId": "t1", "name": "json_tool_call"}}},
|
||||
{"contentBlockIndex": 0, "delta": {"toolUse": {"input": "{\"properties\": {\"a\""}}},
|
||||
{"contentBlockIndex": 0, "delta": {"toolUse": {"input": ": 1}}"}}},
|
||||
{"contentBlockIndex": 0},
|
||||
{"stopReason": "tool_use"}
|
||||
])
|
||||
.as_array()
|
||||
.unwrap()
|
||||
.iter()
|
||||
.flat_map(|event| iterator.transform(event.clone()).unwrap())
|
||||
.map(|chunk| serde_json::to_value(chunk).unwrap())
|
||||
.collect();
|
||||
for chunk in &chunks {
|
||||
assert_eq!(
|
||||
chunk["choices"][0]["delta"]["tool_calls"],
|
||||
serde_json::Value::Null
|
||||
);
|
||||
}
|
||||
let content: String = chunks
|
||||
.iter()
|
||||
.filter_map(|chunk| chunk["choices"][0]["delta"]["content"].as_str())
|
||||
.collect();
|
||||
assert_eq!(content, "{\"a\": 1}");
|
||||
assert_eq!(chunks[4]["choices"][0]["finish_reason"], "stop");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn converse_reasoning_streams_as_thinking_blocks() {
|
||||
let chunks = converse_chunks(serde_json::json!([
|
||||
|
|
|
|||
|
|
@ -1,2 +1,4 @@
|
|||
pub mod converse_messages;
|
||||
pub mod converse_params;
|
||||
pub mod converse_transformation;
|
||||
pub mod invoke_handler;
|
||||
|
|
|
|||
10
litellm-rust/crates/llms/src/cerebras/chat/AGENTS.md
Normal file
10
litellm-rust/crates/llms/src/cerebras/chat/AGENTS.md
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
# rules
|
||||
|
||||
- Cerebras through the OpenAI SDK URL rule at `https://api.cerebras.ai/v1`
|
||||
- `reasoning_effort` only for models the cost map marks as reasoning; `max_completion_tokens` goes as `max_tokens`
|
||||
- Shared flow and response handling: `../../openai/chat/compatible.rs`
|
||||
|
||||
# references
|
||||
|
||||
- https://inference-docs.cerebras.ai/api-reference/chat-completions
|
||||
- litellm/llms/cerebras/chat.py
|
||||
1
litellm-rust/crates/llms/src/cerebras/chat/mod.rs
Normal file
1
litellm-rust/crates/llms/src/cerebras/chat/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
163
litellm-rust/crates/llms/src/cerebras/chat/transformation.rs
Normal file
163
litellm-rust/crates/llms/src/cerebras/chat/transformation.rs
Normal file
|
|
@ -0,0 +1,163 @@
|
|||
//! `litellm/llms/cerebras/chat.py`: Cerebras through the OpenAI SDK at
|
||||
//! `https://api.cerebras.ai/v1`, with `reasoning_effort` for the models the cost map marks as
|
||||
//! reasoning ones and `max_completion_tokens` sent as `max_tokens`.
|
||||
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{
|
||||
Error,
|
||||
openai::chat::compatible::{CompatibleContext, OpenAiCompatibleChatConfig, sdk_complete_url},
|
||||
};
|
||||
|
||||
const SUPPORTED_PARAMS: &[&str] = &[
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"response_format",
|
||||
"seed",
|
||||
"stop",
|
||||
"stream",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"tool_choice",
|
||||
"tools",
|
||||
"user",
|
||||
"max_retries",
|
||||
"extra_headers",
|
||||
];
|
||||
|
||||
pub struct CerebrasChatConfig;
|
||||
|
||||
pub const CEREBRAS_CHAT_CONFIG: CerebrasChatConfig = CerebrasChatConfig;
|
||||
|
||||
impl OpenAiCompatibleChatConfig for CerebrasChatConfig {
|
||||
fn provider(&self) -> LlmProviders {
|
||||
LlmProviders::Cerebras
|
||||
}
|
||||
|
||||
fn label(&self) -> &'static str {
|
||||
"Cerebras"
|
||||
}
|
||||
|
||||
fn api_base_env(&self) -> &'static [&'static str] {
|
||||
&["CEREBRAS_API_BASE"]
|
||||
}
|
||||
|
||||
fn default_api_base(&self) -> Option<&'static str> {
|
||||
Some("https://api.cerebras.ai/v1")
|
||||
}
|
||||
|
||||
fn api_key_env(&self) -> &'static [&'static str] {
|
||||
&["CEREBRAS_API_KEY"]
|
||||
}
|
||||
|
||||
fn complete_url(&self, api_base: &str) -> String {
|
||||
sdk_complete_url(api_base)
|
||||
}
|
||||
|
||||
fn supported_params(&self, context: &CompatibleContext<'_>) -> Vec<&'static str> {
|
||||
let mut params = SUPPORTED_PARAMS.to_vec();
|
||||
if context.supports(self.name(), "supports_reasoning") {
|
||||
params.push("reasoning_effort");
|
||||
}
|
||||
params
|
||||
}
|
||||
|
||||
fn map_params(
|
||||
&self,
|
||||
non_default: &Map<String, Value>,
|
||||
optional: &mut Map<String, Value>,
|
||||
context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
let supported = self.supported_params(context);
|
||||
for (name, value) in non_default {
|
||||
if name == "max_completion_tokens" {
|
||||
optional.insert("max_tokens".into(), value.clone());
|
||||
} else if supported.contains(&name.as_str()) {
|
||||
optional.insert(name.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::borrow::Cow;
|
||||
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::{
|
||||
ErrorDetail, base_llm::model_info::ModelInfo, openai::chat::compatible::transform_request,
|
||||
};
|
||||
|
||||
struct Reasoning;
|
||||
|
||||
impl ModelInfo for Reasoning {
|
||||
fn entry(&self, model: &str) -> Option<Cow<'_, Map<String, Value>>> {
|
||||
(model == "cerebras/gpt-oss-120b").then(|| {
|
||||
Cow::Owned(Map::from_iter([(
|
||||
"supports_reasoning".to_owned(),
|
||||
json!(true),
|
||||
)]))
|
||||
})
|
||||
}
|
||||
|
||||
fn declares(&self, _key: &str) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
fn send(model: &str, params: Value) -> Result<Map<String, Value>, Error> {
|
||||
let Value::Object(body) = params else {
|
||||
panic!("an object")
|
||||
};
|
||||
transform_request(
|
||||
&CEREBRAS_CHAT_CONFIG,
|
||||
&CompatibleContext {
|
||||
model,
|
||||
drop_params: false,
|
||||
model_info: Some(&Reasoning),
|
||||
},
|
||||
body,
|
||||
)
|
||||
.map(|sent| sent.body)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn reasoning_effort_goes_to_reasoning_models() {
|
||||
let sent = send(
|
||||
"gpt-oss-120b",
|
||||
json!({"messages": [], "reasoning_effort": "low", "max_completion_tokens": 5}),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
Value::Object(sent),
|
||||
json!({"model": "gpt-oss-120b", "messages": [], "reasoning_effort": "low", "max_tokens": 5})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn reasoning_effort_is_refused_elsewhere() {
|
||||
assert!(matches!(
|
||||
send(
|
||||
"llama3.1-8b",
|
||||
json!({"messages": [], "reasoning_effort": "low"})
|
||||
),
|
||||
Err(Error::InvalidRequest(ErrorDetail::UnsupportedParams { .. }))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn frequency_penalty_is_refused() {
|
||||
assert!(
|
||||
send(
|
||||
"llama3.1-8b",
|
||||
json!({"messages": [], "frequency_penalty": 1})
|
||||
)
|
||||
.is_err()
|
||||
);
|
||||
}
|
||||
}
|
||||
1
litellm-rust/crates/llms/src/cerebras/mod.rs
Normal file
1
litellm-rust/crates/llms/src/cerebras/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod chat;
|
||||
11
litellm-rust/crates/llms/src/deepseek/chat/AGENTS.md
Normal file
11
litellm-rust/crates/llms/src/deepseek/chat/AGENTS.md
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
# rules
|
||||
|
||||
- DeepSeek at `https://api.deepseek.com/beta`; `/chat/completions` is appended unless the base ends with it
|
||||
- `thinking` and `reasoning_effort` become the `thinking` toggle, only function tools are sent
|
||||
- A thinking turn replays `reasoning_content` on every assistant message
|
||||
- Shared flow and response handling: `../../openai/chat/compatible.rs`
|
||||
|
||||
# references
|
||||
|
||||
- https://api-docs.deepseek.com/api/create-chat-completion
|
||||
- litellm/llms/deepseek/chat/transformation.py
|
||||
1
litellm-rust/crates/llms/src/deepseek/chat/mod.rs
Normal file
1
litellm-rust/crates/llms/src/deepseek/chat/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
352
litellm-rust/crates/llms/src/deepseek/chat/transformation.rs
Normal file
352
litellm-rust/crates/llms/src/deepseek/chat/transformation.rs
Normal file
|
|
@ -0,0 +1,352 @@
|
|||
//! `litellm/llms/deepseek/chat/transformation.py`: DeepSeek at `https://api.deepseek.com/beta`.
|
||||
//! `thinking` and `reasoning_effort` become DeepSeek's thinking toggle, only function tools go,
|
||||
//! content lists collapse to strings unless a vision model can take them, and a thinking turn
|
||||
//! replays `reasoning_content` on every assistant message.
|
||||
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::{
|
||||
Error,
|
||||
base_llm::chat::params::truthy,
|
||||
openai::chat::{
|
||||
compatible::{
|
||||
CompatibleContext, GPT_BASE_PARAMS, OpenAiCompatibleChatConfig, content_list_to_str,
|
||||
copy_supported, gpt_messages, role, search_results_text,
|
||||
},
|
||||
transformation::OPENAI_CHAT_COMPLETIONS_PATH,
|
||||
},
|
||||
};
|
||||
|
||||
pub struct DeepSeekChatConfig;
|
||||
|
||||
pub const DEEPSEEK_CHAT_CONFIG: DeepSeekChatConfig = DeepSeekChatConfig;
|
||||
|
||||
impl OpenAiCompatibleChatConfig for DeepSeekChatConfig {
|
||||
fn provider(&self) -> LlmProviders {
|
||||
LlmProviders::Deepseek
|
||||
}
|
||||
|
||||
fn label(&self) -> &'static str {
|
||||
"DeepSeek"
|
||||
}
|
||||
|
||||
fn api_base_env(&self) -> &'static [&'static str] {
|
||||
&["DEEPSEEK_API_BASE"]
|
||||
}
|
||||
|
||||
fn default_api_base(&self) -> Option<&'static str> {
|
||||
Some("https://api.deepseek.com/beta")
|
||||
}
|
||||
|
||||
fn api_key_env(&self) -> &'static [&'static str] {
|
||||
&["DEEPSEEK_API_KEY"]
|
||||
}
|
||||
|
||||
fn complete_url(&self, api_base: &str) -> String {
|
||||
let base = api_base.trim_end_matches('/');
|
||||
match base.ends_with(OPENAI_CHAT_COMPLETIONS_PATH) {
|
||||
true => base.to_owned(),
|
||||
false => format!("{base}{OPENAI_CHAT_COMPLETIONS_PATH}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn supported_params(&self, _context: &CompatibleContext<'_>) -> Vec<&'static str> {
|
||||
GPT_BASE_PARAMS
|
||||
.iter()
|
||||
.copied()
|
||||
.chain(["response_format", "thinking", "reasoning_effort"])
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn map_params(
|
||||
&self,
|
||||
non_default: &Map<String, Value>,
|
||||
optional: &mut Map<String, Value>,
|
||||
context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
copy_supported(non_default, optional, &self.supported_params(context));
|
||||
let thinking = optional.remove("thinking");
|
||||
let effort = optional.remove("reasoning_effort");
|
||||
let thinking_type = thinking
|
||||
.as_ref()
|
||||
.and_then(|thinking| thinking.get("type"))
|
||||
.and_then(Value::as_str)
|
||||
.filter(|kind| ["enabled", "disabled"].contains(kind));
|
||||
if let Some(kind) = thinking_type {
|
||||
optional.insert("thinking".into(), json!({"type": kind}));
|
||||
} else if let Some(effort) = effort.filter(|effort| !effort.is_null()) {
|
||||
let kind = match effort.as_str() {
|
||||
Some("none") => "disabled",
|
||||
_ => "enabled",
|
||||
};
|
||||
optional.insert("thinking".into(), json!({"type": kind}));
|
||||
}
|
||||
drop_unsupported_tools(optional);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn transform_messages(
|
||||
&self,
|
||||
messages: Vec<Value>,
|
||||
optional: &Map<String, Value>,
|
||||
context: &CompatibleContext<'_>,
|
||||
) -> Result<Vec<Value>, Error> {
|
||||
let thinking = context.supports(self.name(), "supports_reasoning")
|
||||
&& optional
|
||||
.get("thinking")
|
||||
.and_then(|thinking| thinking.get("type"))
|
||||
== Some(&json!("enabled"));
|
||||
let messages: Vec<Value> = match thinking {
|
||||
true => messages.into_iter().map(fill_reasoning_content).collect(),
|
||||
false => messages,
|
||||
};
|
||||
let forward_images = messages
|
||||
.iter()
|
||||
.any(|message| message.get("content").is_some_and(Value::is_array))
|
||||
&& context.supports(self.name(), "supports_vision");
|
||||
gpt_messages(
|
||||
messages
|
||||
.into_iter()
|
||||
.map(|message| forward_or_collapse_content(message, forward_images))
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// `_drop_unsupported_tools`: DeepSeek takes only function tools.
|
||||
fn drop_unsupported_tools(optional: &mut Map<String, Value>) {
|
||||
let Some(Value::Array(tools)) = optional.get("tools") else {
|
||||
return;
|
||||
};
|
||||
let is_function = |tool: &Value| tool.get("type").and_then(Value::as_str) == Some("function");
|
||||
if tools.is_empty() || tools.iter().all(is_function) {
|
||||
return;
|
||||
}
|
||||
let functions: Vec<Value> = tools
|
||||
.iter()
|
||||
.filter(|tool| is_function(tool))
|
||||
.cloned()
|
||||
.collect();
|
||||
if functions.is_empty() {
|
||||
for key in ["tools", "tool_choice", "parallel_tool_calls"] {
|
||||
optional.remove(key);
|
||||
}
|
||||
return;
|
||||
}
|
||||
let names: Vec<&str> = functions
|
||||
.iter()
|
||||
.filter_map(|tool| tool.get("function")?.get("name")?.as_str())
|
||||
.collect();
|
||||
let choice_matches = match optional.get("tool_choice") {
|
||||
Some(Value::Object(choice)) => {
|
||||
choice.get("type").and_then(Value::as_str) == Some("function")
|
||||
&& choice
|
||||
.get("function")
|
||||
.and_then(|function| function.get("name"))
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|name| names.contains(&name))
|
||||
}
|
||||
_ => true,
|
||||
};
|
||||
if !choice_matches {
|
||||
optional.remove("tool_choice");
|
||||
}
|
||||
optional.insert("tools".into(), Value::Array(functions));
|
||||
}
|
||||
|
||||
/// `_fill_reasoning_content`.
|
||||
fn fill_reasoning_content(mut message: Value) -> Value {
|
||||
if role(&message) != Some("assistant") || message.get("reasoning_content").is_some_and(truthy) {
|
||||
return message;
|
||||
}
|
||||
let Some(object) = message.as_object_mut() else {
|
||||
return message;
|
||||
};
|
||||
let stored = object
|
||||
.get("provider_specific_fields")
|
||||
.and_then(|fields| fields.get("reasoning_content"))
|
||||
.filter(|stored| truthy(stored))
|
||||
.cloned();
|
||||
match stored {
|
||||
Some(stored) => {
|
||||
object.insert("reasoning_content".into(), stored);
|
||||
if let Some(Value::Object(fields)) = object.get_mut("provider_specific_fields") {
|
||||
fields.remove("reasoning_content");
|
||||
}
|
||||
}
|
||||
None => {
|
||||
object.insert("reasoning_content".into(), json!(" "));
|
||||
}
|
||||
}
|
||||
message
|
||||
}
|
||||
|
||||
fn forward_or_collapse_content(message: Value, forward_images: bool) -> Value {
|
||||
if forward_images
|
||||
&& let Some(Value::Array(content)) = message.get("content")
|
||||
&& is_vision_forwardable(&message, content)
|
||||
{
|
||||
return with_search_results_text_block(message);
|
||||
}
|
||||
let collapsed = content_list_to_str(&message);
|
||||
if collapsed.is_empty() || message.get("content") == Some(&Value::String(collapsed.clone())) {
|
||||
return message;
|
||||
}
|
||||
let mut message = message;
|
||||
if let Some(object) = message.as_object_mut() {
|
||||
object.insert("content".into(), Value::String(collapsed));
|
||||
}
|
||||
message
|
||||
}
|
||||
|
||||
fn is_vision_forwardable(message: &Value, content: &[Value]) -> bool {
|
||||
let forwardable = |block: &Value| match block.get("type").and_then(Value::as_str) {
|
||||
Some("image_url") => match block.get("image_url") {
|
||||
Some(Value::String(url)) => !url.is_empty(),
|
||||
Some(Value::Object(image)) => image
|
||||
.get("url")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|url| !url.is_empty()),
|
||||
_ => false,
|
||||
},
|
||||
Some("text") => block.get("text").is_some_and(Value::is_string),
|
||||
_ => false,
|
||||
};
|
||||
role(message) == Some("user")
|
||||
&& content.iter().all(forwardable)
|
||||
&& content
|
||||
.iter()
|
||||
.any(|block| block.get("type").and_then(Value::as_str) == Some("image_url"))
|
||||
}
|
||||
|
||||
fn with_search_results_text_block(mut message: Value) -> Value {
|
||||
let text = search_results_text(message.get("search_results"));
|
||||
if text.is_empty() {
|
||||
return message;
|
||||
}
|
||||
if let Some(object) = message.as_object_mut() {
|
||||
object.remove("search_results");
|
||||
if let Some(Value::Array(content)) = object.get_mut("content") {
|
||||
content.push(json!({"type": "text", "text": text}));
|
||||
}
|
||||
}
|
||||
message
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::borrow::Cow;
|
||||
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::{base_llm::model_info::ModelInfo, openai::chat::compatible::transform_request};
|
||||
|
||||
struct Capable;
|
||||
|
||||
impl ModelInfo for Capable {
|
||||
fn entry(&self, model: &str) -> Option<Cow<'_, Map<String, Value>>> {
|
||||
let entry = match model {
|
||||
"deepseek/deepseek-reasoner" => json!({"supports_reasoning": true}),
|
||||
"deepseek/deepseek-vl" => json!({"supports_vision": true}),
|
||||
_ => return None,
|
||||
};
|
||||
let Value::Object(entry) = entry else {
|
||||
unreachable!()
|
||||
};
|
||||
Some(Cow::Owned(entry))
|
||||
}
|
||||
|
||||
fn declares(&self, _key: &str) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
fn send(model: &str, params: Value) -> Map<String, Value> {
|
||||
let Value::Object(body) = params else {
|
||||
panic!("an object")
|
||||
};
|
||||
transform_request(
|
||||
&DEEPSEEK_CHAT_CONFIG,
|
||||
&CompatibleContext {
|
||||
model,
|
||||
drop_params: false,
|
||||
model_info: Some(&Capable),
|
||||
},
|
||||
body,
|
||||
)
|
||||
.unwrap()
|
||||
.body
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::base(
|
||||
"https://api.deepseek.com/beta",
|
||||
"https://api.deepseek.com/beta/chat/completions"
|
||||
)]
|
||||
#[case::named("https://h/chat/completions/", "https://h/chat/completions")]
|
||||
fn the_url_ends_in_chat_completions_once(#[case] base: &str, #[case] expected: &str) {
|
||||
assert_eq!(DEEPSEEK_CHAT_CONFIG.complete_url(base), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::thinking_wins(json!({"thinking": {"type": "disabled", "budget_tokens": 9}, "reasoning_effort": "high"}), json!({"type": "disabled"}))]
|
||||
#[case::effort_none(json!({"reasoning_effort": "none"}), json!({"type": "disabled"}))]
|
||||
#[case::effort_on(json!({"reasoning_effort": "low"}), json!({"type": "enabled"}))]
|
||||
fn reasoning_becomes_the_thinking_toggle(#[case] params: Value, #[case] thinking: Value) {
|
||||
let Value::Object(mut params) = params else {
|
||||
unreachable!()
|
||||
};
|
||||
params.insert("messages".into(), json!([]));
|
||||
let sent = send("deepseek-chat", Value::Object(params));
|
||||
assert_eq!(sent["thinking"], thinking);
|
||||
assert!(!sent.contains_key("reasoning_effort"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn only_function_tools_go_and_a_choice_naming_a_dropped_tool_goes_too() {
|
||||
let sent = send(
|
||||
"deepseek-chat",
|
||||
json!({"messages": [],
|
||||
"tools": [{"type": "function", "function": {"name": "f"}}, {"type": "namespace", "name": "n"}],
|
||||
"tool_choice": {"type": "function", "function": {"name": "n"}}}),
|
||||
);
|
||||
assert_eq!(
|
||||
sent["tools"],
|
||||
json!([{"type": "function", "function": {"name": "f"}}])
|
||||
);
|
||||
assert!(!sent.contains_key("tool_choice"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_thinking_turn_replays_reasoning_on_assistant_messages() {
|
||||
let sent = send(
|
||||
"deepseek-reasoner",
|
||||
json!({"thinking": {"type": "enabled"}, "messages": [
|
||||
{"role": "assistant", "content": "a", "provider_specific_fields": {"reasoning_content": "r"}},
|
||||
{"role": "assistant", "content": "b"}]}),
|
||||
);
|
||||
assert_eq!(
|
||||
sent["messages"],
|
||||
json!([
|
||||
{"role": "assistant", "content": "a", "provider_specific_fields": {}, "reasoning_content": "r"},
|
||||
{"role": "assistant", "content": "b", "reasoning_content": " "}])
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::collapsed("deepseek-chat", json!("look"))]
|
||||
#[case::forwarded("deepseek-vl", json!([{"type": "text", "text": "look"}, {"type": "image_url", "image_url": {"url": "https://i"}}]))]
|
||||
fn content_lists_collapse_unless_vision_takes_them(
|
||||
#[case] model: &str,
|
||||
#[case] content: Value,
|
||||
) {
|
||||
let sent = send(
|
||||
model,
|
||||
json!({"messages": [{"role": "user", "content": [
|
||||
{"type": "text", "text": "look"}, {"type": "image_url", "image_url": "https://i"}]}]}),
|
||||
);
|
||||
assert_eq!(sent["messages"][0]["content"], content);
|
||||
}
|
||||
}
|
||||
|
|
@ -1 +1,2 @@
|
|||
pub mod chat;
|
||||
pub mod messages;
|
||||
|
|
|
|||
12
litellm-rust/crates/llms/src/fireworks_ai/chat/AGENTS.md
Normal file
12
litellm-rust/crates/llms/src/fireworks_ai/chat/AGENTS.md
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
# rules
|
||||
|
||||
- Fireworks AI at `https://api.fireworks.ai/inference/v1`, the model named by its full resource (`resolve_fireworks_resource_name`)
|
||||
- Tool, tool choice and reasoning params are gated by the cost map, file parts and images on non-vision models are rejected
|
||||
- vLLM/NIM `extra_body` keys are translated (`truncate_prompt_tokens`, `chat_template_kwargs`, `guided_*`) or dropped
|
||||
- Tool calls Fireworks writes into the content go back as `tool_calls`; the answer's model gets the `fireworks_ai/` prefix
|
||||
- Shared flow and response handling: `../../openai/chat/compatible.rs`
|
||||
|
||||
# references
|
||||
|
||||
- https://docs.fireworks.ai/api-reference/post-chatcompletions
|
||||
- litellm/llms/fireworks_ai/chat/transformation.py
|
||||
1
litellm-rust/crates/llms/src/fireworks_ai/chat/mod.rs
Normal file
1
litellm-rust/crates/llms/src/fireworks_ai/chat/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
680
litellm-rust/crates/llms/src/fireworks_ai/chat/transformation.rs
Normal file
680
litellm-rust/crates/llms/src/fireworks_ai/chat/transformation.rs
Normal file
|
|
@ -0,0 +1,680 @@
|
|||
//! `litellm/llms/fireworks_ai/chat/transformation.py`: Fireworks AI at
|
||||
//! `https://api.fireworks.ai/inference/v1`, with models named by their Fireworks resource,
|
||||
//! capability gated tool and reasoning params, vLLM style `extra_body` keys translated or
|
||||
//! dropped, and tool calls Fireworks writes into the content moved back to `tool_calls`.
|
||||
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::{
|
||||
Error, ErrorDetail,
|
||||
base_llm::chat::{
|
||||
params::{python_dumps, truthy},
|
||||
response_format::unpack_legacy_defs,
|
||||
},
|
||||
openai::chat::compatible::{
|
||||
CompatibleContext, OpenAiCompatibleChatConfig, RESPONSE_FORMAT_TOOL_NAME, ResponseShape,
|
||||
gpt_messages, messages_mut, role, strip_response_only_fields,
|
||||
},
|
||||
};
|
||||
|
||||
const SUPPORTED_PARAMS: &[&str] = &[
|
||||
"stream",
|
||||
"max_completion_tokens",
|
||||
"max_tokens",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"top_k",
|
||||
"frequency_penalty",
|
||||
"presence_penalty",
|
||||
"n",
|
||||
"stop",
|
||||
"response_format",
|
||||
"user",
|
||||
"logprobs",
|
||||
"prompt_truncate_len",
|
||||
"context_length_exceeded_behavior",
|
||||
"seed",
|
||||
"top_logprobs",
|
||||
"min_p",
|
||||
"typical_p",
|
||||
"repetition_penalty",
|
||||
"mirostat_target",
|
||||
"mirostat_lr",
|
||||
"logit_bias",
|
||||
"echo",
|
||||
"echo_last",
|
||||
"ignore_eos",
|
||||
"prompt_cache_key",
|
||||
"prompt_cache_isolation_key",
|
||||
"raw_output",
|
||||
"perf_metrics_in_response",
|
||||
"return_token_ids",
|
||||
"safe_tokenization",
|
||||
"service_tier",
|
||||
"speculation",
|
||||
"prediction",
|
||||
"stream_options",
|
||||
"sampling_mask",
|
||||
];
|
||||
const NIM_VLLM_STRIP_PARAMS: &[&str] = &[
|
||||
"stop_token_ids",
|
||||
"include_stop_str_in_output",
|
||||
"skip_special_tokens",
|
||||
"spaces_between_special_tokens",
|
||||
"best_of",
|
||||
"use_beam_search",
|
||||
"guided_decoding_backend",
|
||||
"guided_regex",
|
||||
"add_generation_prompt",
|
||||
"continue_final_message",
|
||||
"add_special_tokens",
|
||||
"detokenize",
|
||||
"allowed_token_ids",
|
||||
"bad_words",
|
||||
"include_reasoning",
|
||||
"nvext",
|
||||
];
|
||||
const EXTRA_BODY_CONSUMED_PARAMS: &[&str] = &[
|
||||
"truncate_prompt_tokens",
|
||||
"chat_template_kwargs",
|
||||
"guided_json",
|
||||
"guided_grammar",
|
||||
"guided_choice",
|
||||
];
|
||||
const FIREROUTER: &str = "firerouter";
|
||||
const ROUTER_SHORT_NAMES: &[&str] = &[FIREROUTER, "auto", "auto-instant"];
|
||||
const SESSION_AFFINITY_HEADER: &str = "x-session-affinity";
|
||||
|
||||
pub struct FireworksAiChatConfig;
|
||||
|
||||
pub const FIREWORKS_AI_CHAT_CONFIG: FireworksAiChatConfig = FireworksAiChatConfig;
|
||||
|
||||
impl OpenAiCompatibleChatConfig for FireworksAiChatConfig {
|
||||
fn provider(&self) -> LlmProviders {
|
||||
LlmProviders::FireworksAi
|
||||
}
|
||||
|
||||
fn label(&self) -> &'static str {
|
||||
"Fireworks AI"
|
||||
}
|
||||
|
||||
fn api_base_env(&self) -> &'static [&'static str] {
|
||||
&["FIREWORKS_API_BASE"]
|
||||
}
|
||||
|
||||
fn default_api_base(&self) -> Option<&'static str> {
|
||||
Some("https://api.fireworks.ai/inference/v1")
|
||||
}
|
||||
|
||||
fn api_key_env(&self) -> &'static [&'static str] {
|
||||
&[
|
||||
"FIREWORKS_API_KEY",
|
||||
"FIREWORKS_AI_API_KEY",
|
||||
"FIREWORKSAI_API_KEY",
|
||||
"FIREWORKS_AI_TOKEN",
|
||||
]
|
||||
}
|
||||
|
||||
/// `with_fireworks_session_affinity`: the caller's session pins its calls to one node so
|
||||
/// prompt caching hits.
|
||||
fn headers(
|
||||
&self,
|
||||
body: &Map<String, Value>,
|
||||
_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Vec<(String, String)> {
|
||||
let session = body
|
||||
.get("litellm_session_id")
|
||||
.filter(|id| truthy(id))
|
||||
.or_else(|| {
|
||||
body.get("metadata")
|
||||
.and_then(|metadata| metadata.get("session_id"))
|
||||
.filter(|id| truthy(id))
|
||||
});
|
||||
session
|
||||
.map(|id| match id {
|
||||
Value::String(id) => id.clone(),
|
||||
other => other.to_string(),
|
||||
})
|
||||
.map(|id| vec![(SESSION_AFFINITY_HEADER.to_owned(), id)])
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
fn wire_model(&self, model: &str) -> String {
|
||||
resolve_fireworks_resource_name(model)
|
||||
}
|
||||
|
||||
fn supported_params(&self, context: &CompatibleContext<'_>) -> Vec<&'static str> {
|
||||
let mut params = SUPPORTED_PARAMS.to_vec();
|
||||
if exact_capability(context, "supports_function_calling") != Some(false) {
|
||||
params.extend(["tools", "parallel_tool_calls"]);
|
||||
}
|
||||
if capability(context, "supports_tool_choice") {
|
||||
params.push("tool_choice");
|
||||
}
|
||||
if capability(context, "supports_reasoning") {
|
||||
params.extend(["reasoning_effort", "reasoning_history", "thinking"]);
|
||||
}
|
||||
params
|
||||
}
|
||||
|
||||
fn map_params(
|
||||
&self,
|
||||
non_default: &Map<String, Value>,
|
||||
optional: &mut Map<String, Value>,
|
||||
context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
if non_default.get("thinking").is_some_and(|v| !v.is_null())
|
||||
&& non_default
|
||||
.get("reasoning_effort")
|
||||
.is_some_and(|v| !v.is_null())
|
||||
{
|
||||
return Err(Error::InvalidRequest(ErrorDetail::Message(
|
||||
"Fireworks AI chat completions does not support specifying both `thinking` and \
|
||||
`reasoning_effort` in the same request."
|
||||
.into(),
|
||||
)));
|
||||
}
|
||||
let supported = self.supported_params(context);
|
||||
let tools_set = non_default
|
||||
.get("tools")
|
||||
.is_some_and(|tools| !tools.is_null());
|
||||
let mut response_format_tool = None;
|
||||
for (name, value) in non_default {
|
||||
match name.as_str() {
|
||||
"tool_choice" => {
|
||||
let choice = match value.as_str() {
|
||||
Some("required") => json!("any"),
|
||||
_ => value.clone(),
|
||||
};
|
||||
optional.insert("tool_choice".into(), choice);
|
||||
}
|
||||
"response_format" if tools_set => {
|
||||
if let Some(schema) = response_schema(value) {
|
||||
response_format_tool = Some(json!({"type": "function",
|
||||
"function": {"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": schema}}));
|
||||
optional.insert("json_mode".into(), Value::Bool(true));
|
||||
}
|
||||
}
|
||||
"max_completion_tokens" => {
|
||||
optional.insert("max_tokens".into(), value.clone());
|
||||
}
|
||||
"reasoning_effort" => {
|
||||
if let Some(effort) = map_reasoning_effort(value) {
|
||||
optional.insert("reasoning_effort".into(), effort);
|
||||
}
|
||||
}
|
||||
name if supported.contains(&name) && !value.is_null() => {
|
||||
optional.insert(name.to_owned(), value.clone());
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
if let Some(tool) = response_format_tool
|
||||
&& let Some(Value::Array(tools)) = optional.get_mut("tools")
|
||||
{
|
||||
tools.push(tool);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `map_extra_body_params`.
|
||||
fn map_extra_body(
|
||||
&self,
|
||||
extra_body: &mut Map<String, Value>,
|
||||
optional: &mut Map<String, Value>,
|
||||
context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
let mut promoted = Map::new();
|
||||
if let Some(tokens) = extra_body
|
||||
.get("truncate_prompt_tokens")
|
||||
.filter(|v| !v.is_null())
|
||||
&& !extra_body.contains_key("prompt_truncate_len")
|
||||
&& !optional.contains_key("prompt_truncate_len")
|
||||
{
|
||||
promoted.insert("prompt_truncate_len".into(), tokens.clone());
|
||||
}
|
||||
let explicit_reasoning = ["reasoning_effort", "thinking"]
|
||||
.iter()
|
||||
.any(|key| optional.contains_key(*key) || extra_body.contains_key(*key));
|
||||
if let Some(Value::Object(kwargs)) = extra_body.get("chat_template_kwargs")
|
||||
&& !explicit_reasoning
|
||||
&& let Some(effort) = effort_from_chat_template_kwargs(kwargs)
|
||||
&& context.supports("fireworks_ai", "supports_reasoning")
|
||||
{
|
||||
promoted.insert("reasoning_effort".into(), effort);
|
||||
}
|
||||
let has_response_format =
|
||||
optional.contains_key("response_format") || extra_body.contains_key("response_format");
|
||||
if !has_response_format && let Some(format) = guided_response_format(extra_body) {
|
||||
promoted.insert("response_format".into(), format);
|
||||
}
|
||||
extra_body.retain(|key, _| {
|
||||
!NIM_VLLM_STRIP_PARAMS.contains(&key.as_str())
|
||||
&& !EXTRA_BODY_CONSUMED_PARAMS.contains(&key.as_str())
|
||||
&& (key != "response_format" || !optional.contains_key("response_format"))
|
||||
});
|
||||
optional.extend(promoted);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `_transform_messages_helper`, then the base transform.
|
||||
fn transform_messages(
|
||||
&self,
|
||||
messages: Vec<Value>,
|
||||
_optional: &Map<String, Value>,
|
||||
context: &CompatibleContext<'_>,
|
||||
) -> Result<Vec<Value>, Error> {
|
||||
let vision = exact_capability(context, "supports_vision");
|
||||
for message in messages
|
||||
.iter()
|
||||
.filter(|message| role(message) == Some("user"))
|
||||
{
|
||||
for kind in message
|
||||
.get("content")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(|part| part.get("type").and_then(Value::as_str))
|
||||
{
|
||||
if kind == "file" {
|
||||
return Err(Error::InvalidRequest(ErrorDetail::Message(
|
||||
"Fireworks AI chat completions does not support file content blocks. For \
|
||||
PDFs, convert pages to images and send image_url blocks to a Fireworks \
|
||||
vision model, or extract text before calling a text-only model."
|
||||
.into(),
|
||||
)));
|
||||
}
|
||||
if kind == "image_url" && vision == Some(false) {
|
||||
return Err(Error::InvalidRequest(ErrorDetail::Message(format!(
|
||||
"Fireworks AI model {} does not support image inputs. Use a Fireworks \
|
||||
vision model or remove image_url content blocks.",
|
||||
resolve_fireworks_resource_name(context.model)
|
||||
))));
|
||||
}
|
||||
}
|
||||
}
|
||||
gpt_messages(
|
||||
messages
|
||||
.into_iter()
|
||||
.map(|mut message| {
|
||||
strip_response_only_fields(&mut message);
|
||||
message
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
fn finish_request(
|
||||
&self,
|
||||
body: &mut Map<String, Value>,
|
||||
_context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
if let Some(Value::Array(tools)) = body.get_mut("tools") {
|
||||
for function in tools
|
||||
.iter_mut()
|
||||
.filter(|tool| tool.get("type").and_then(Value::as_str) == Some("function"))
|
||||
.filter_map(|tool| tool.get_mut("function")?.as_object_mut())
|
||||
{
|
||||
function.remove("strict");
|
||||
if let Some(Value::Object(parameters)) = function.get_mut("parameters") {
|
||||
unpack_legacy_defs(parameters);
|
||||
}
|
||||
}
|
||||
}
|
||||
if body.get("stream").is_some_and(truthy) {
|
||||
match body.get_mut("stream_options") {
|
||||
Some(Value::Object(options)) => {
|
||||
if options.get("include_usage") != Some(&Value::Bool(false)) {
|
||||
options.insert("include_usage".into(), Value::Bool(true));
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
body.insert("stream_options".into(), json!({"include_usage": true}));
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn fix_response(&self, body: &mut Map<String, Value>, shape: &ResponseShape) -> bool {
|
||||
let mut changed = false;
|
||||
if let Some(Value::String(model)) = body.get_mut("model") {
|
||||
*model = format!("fireworks_ai/{model}");
|
||||
changed = true;
|
||||
}
|
||||
if !shape.tools_sent {
|
||||
return changed;
|
||||
}
|
||||
for message in messages_mut(body) {
|
||||
changed |= content_tool_call_to_tool_calls(message, &shape.tool_names);
|
||||
}
|
||||
changed
|
||||
}
|
||||
}
|
||||
|
||||
/// `resolve_fireworks_resource_name`: the full Fireworks resource a short model name means.
|
||||
pub fn resolve_fireworks_resource_name(model: &str) -> String {
|
||||
let stripped = model.strip_prefix("fireworks_ai/").unwrap_or(model);
|
||||
if stripped.starts_with("accounts/") || stripped.starts_with("FW-") || stripped.contains('#') {
|
||||
return stripped.to_owned();
|
||||
}
|
||||
if stripped.starts_with("routers/") || stripped.starts_with("models/") {
|
||||
return format!("accounts/fireworks/{stripped}");
|
||||
}
|
||||
if stripped.ends_with("-fast")
|
||||
|| ROUTER_SHORT_NAMES.contains(&stripped)
|
||||
|| stripped.starts_with(&format!("{FIREROUTER}/"))
|
||||
{
|
||||
return format!("accounts/fireworks/routers/{stripped}");
|
||||
}
|
||||
format!("accounts/fireworks/models/{stripped}")
|
||||
}
|
||||
|
||||
/// `_get_model_cost_capability_exact`.
|
||||
fn exact_capability(context: &CompatibleContext<'_>, key: &str) -> Option<bool> {
|
||||
let short = context
|
||||
.model
|
||||
.strip_prefix("fireworks_ai/")
|
||||
.unwrap_or(context.model);
|
||||
let short = short
|
||||
.strip_prefix("accounts/fireworks/models/")
|
||||
.unwrap_or(short);
|
||||
let firerouter = format!("accounts/fireworks/routers/{FIREROUTER}");
|
||||
let mut candidates = vec![
|
||||
context.model.to_owned(),
|
||||
format!("fireworks_ai/{short}"),
|
||||
format!("fireworks_ai/accounts/fireworks/models/{short}"),
|
||||
];
|
||||
if resolve_fireworks_resource_name(context.model).starts_with(&format!("{firerouter}/")) {
|
||||
candidates.push(format!("fireworks_ai/{firerouter}"));
|
||||
}
|
||||
context.capability_of(&candidates, key)
|
||||
}
|
||||
|
||||
fn capability(context: &CompatibleContext<'_>, key: &str) -> bool {
|
||||
exact_capability(context, key) == Some(true) || context.supports("fireworks_ai", key)
|
||||
}
|
||||
|
||||
/// `_map_reasoning_effort`: booleans as levels, `auto` left to Fireworks.
|
||||
fn map_reasoning_effort(value: &Value) -> Option<Value> {
|
||||
let effort = match value {
|
||||
Value::Object(object) => object.get("effort").cloned().unwrap_or(Value::Null),
|
||||
other => other.clone(),
|
||||
};
|
||||
match effort {
|
||||
Value::Bool(true) => Some(json!("medium")),
|
||||
Value::Bool(false) => Some(json!("none")),
|
||||
Value::String(effort) if effort == "auto" => None,
|
||||
Value::Null => None,
|
||||
other => Some(other),
|
||||
}
|
||||
}
|
||||
|
||||
fn response_schema(value: &Value) -> Option<Value> {
|
||||
let schema = match value.get("response_schema") {
|
||||
Some(schema) => schema.clone(),
|
||||
None => value.get("json_schema")?.get("schema")?.clone(),
|
||||
};
|
||||
truthy(&schema).then_some(schema)
|
||||
}
|
||||
|
||||
/// `effort_from_chat_template_kwargs`.
|
||||
fn effort_from_chat_template_kwargs(kwargs: &Map<String, Value>) -> Option<Value> {
|
||||
let flag = |keys: &[&str]| keys.iter().find_map(|key| kwargs.get(*key)?.as_bool());
|
||||
if flag(&["enable_thinking", "thinking"]) == Some(false) {
|
||||
return Some(json!("none"));
|
||||
}
|
||||
if let Some(budget) = kwargs.get("reasoning_budget").and_then(Value::as_f64)
|
||||
&& budget > 0.0
|
||||
{
|
||||
return Some(json!(budget as i64));
|
||||
}
|
||||
(flag(&["low_effort"]) == Some(true)).then(|| json!("low"))
|
||||
}
|
||||
|
||||
/// `translate_guided_params`.
|
||||
fn guided_response_format(extra_body: &Map<String, Value>) -> Option<Value> {
|
||||
let present = |key: &str| extra_body.get(key).filter(|value| !value.is_null());
|
||||
let schema = |schema: &Value, name: &str| json!({"type": "json_schema", "json_schema": {"name": name, "schema": schema}});
|
||||
if let Some(guided) = present("guided_json") {
|
||||
return Some(schema(guided, "response"));
|
||||
}
|
||||
if let Some(grammar) = present("guided_grammar") {
|
||||
return Some(json!({"type": "grammar", "grammar": grammar}));
|
||||
}
|
||||
present("guided_choice")
|
||||
.map(|choices| schema(&json!({"type": "string", "enum": choices}), "choice"))
|
||||
}
|
||||
|
||||
/// `_handle_message_content_with_tool_calls`: a call to a sent tool, written as the content.
|
||||
fn content_tool_call_to_tool_calls(
|
||||
message: &mut Map<String, Value>,
|
||||
tool_names: &[String],
|
||||
) -> bool {
|
||||
if message
|
||||
.get("tool_calls")
|
||||
.is_some_and(|calls| !calls.is_null())
|
||||
{
|
||||
return false;
|
||||
}
|
||||
let Some(Value::Object(call)) = message
|
||||
.get("content")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(|content| serde_json::from_str::<Value>(content).ok())
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
let Some(name) = call.get("name").and_then(Value::as_str) else {
|
||||
return false;
|
||||
};
|
||||
if name == RESPONSE_FORMAT_TOOL_NAME || !tool_names.iter().any(|tool| tool == name) {
|
||||
return false;
|
||||
}
|
||||
let arguments = match call.get("arguments") {
|
||||
Some(Value::String(arguments)) => arguments.clone(),
|
||||
Some(Value::Object(arguments)) => python_dumps(&Value::Object(arguments.clone())),
|
||||
Some(Value::Null) | None => match call.get("parameters") {
|
||||
Some(parameters @ Value::Object(_)) => python_dumps(parameters),
|
||||
_ => String::new(),
|
||||
},
|
||||
Some(_) => return false,
|
||||
};
|
||||
message.insert(
|
||||
"tool_calls".into(),
|
||||
json!([{"id": uuid::Uuid::new_v4().to_string(), "type": "function",
|
||||
"function": {"name": name, "arguments": arguments}}]),
|
||||
);
|
||||
message.insert("content".into(), Value::Null);
|
||||
true
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::borrow::Cow;
|
||||
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::{
|
||||
base_llm::model_info::ModelInfo,
|
||||
openai::chat::compatible::{CompatibleChatBody, fix_response, transform_request},
|
||||
};
|
||||
|
||||
struct Registry;
|
||||
|
||||
impl ModelInfo for Registry {
|
||||
fn entry(&self, model: &str) -> Option<Cow<'_, Map<String, Value>>> {
|
||||
let entry = match model {
|
||||
"fireworks_ai/accounts/fireworks/models/thinker" => {
|
||||
json!({"supports_reasoning": true, "supports_tool_choice": true})
|
||||
}
|
||||
"fireworks_ai/text-only" => {
|
||||
json!({"supports_vision": false, "supports_function_calling": false})
|
||||
}
|
||||
_ => return None,
|
||||
};
|
||||
let Value::Object(entry) = entry else {
|
||||
unreachable!()
|
||||
};
|
||||
Some(Cow::Owned(entry))
|
||||
}
|
||||
|
||||
fn declares(&self, _key: &str) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
fn send(model: &str, params: Value) -> Result<CompatibleChatBody, Error> {
|
||||
let Value::Object(body) = params else {
|
||||
panic!("an object")
|
||||
};
|
||||
transform_request(
|
||||
&FIREWORKS_AI_CHAT_CONFIG,
|
||||
&CompatibleContext {
|
||||
model,
|
||||
drop_params: false,
|
||||
model_info: Some(&Registry),
|
||||
},
|
||||
body,
|
||||
)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::short(
|
||||
"llama-v3p1-8b-instruct",
|
||||
"accounts/fireworks/models/llama-v3p1-8b-instruct"
|
||||
)]
|
||||
#[case::router("glm-4p6-fast", "accounts/fireworks/routers/glm-4p6-fast")]
|
||||
#[case::models_prefix("models/x", "accounts/fireworks/models/x")]
|
||||
#[case::full("accounts/me/models/x", "accounts/me/models/x")]
|
||||
#[case::deployment("x#accounts/me/deployments/d", "x#accounts/me/deployments/d")]
|
||||
fn models_are_named_by_their_resource(#[case] model: &str, #[case] resource: &str) {
|
||||
assert_eq!(resolve_fireworks_resource_name(model), resource);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn reasoning_models_take_tool_choice_and_mapped_efforts() {
|
||||
let sent = send("thinker", json!({"messages": [], "tool_choice": "required",
|
||||
"tools": [{"type": "function", "function": {"name": "f", "strict": true,
|
||||
"parameters": {"type": "object", "properties": {"a": {"$ref": "#/definitions/A"}}, "definitions": {"A": {"type": "string"}}}}}],
|
||||
"reasoning_effort": true, "max_completion_tokens": 8}))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
Value::Object(sent.body),
|
||||
json!({"model": "accounts/fireworks/models/thinker", "messages": [],
|
||||
"tool_choice": "any",
|
||||
"tools": [{"type": "function", "function": {"name": "f",
|
||||
"parameters": {"type": "object", "properties": {"a": {"type": "string"}}}}}],
|
||||
"reasoning_effort": "medium", "max_tokens": 8})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn thinking_and_effort_together_are_rejected() {
|
||||
assert!(matches!(
|
||||
send(
|
||||
"thinker",
|
||||
json!({"messages": [], "thinking": {"type": "enabled"}, "reasoning_effort": "low"})
|
||||
),
|
||||
Err(Error::InvalidRequest(ErrorDetail::Message(_)))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::file(json!({"type": "file", "file": {"file_id": "f"}}))]
|
||||
#[case::image_on_text_model(json!({"type": "image_url", "image_url": {"url": "https://i"}}))]
|
||||
fn content_fireworks_cannot_read_is_rejected(#[case] part: Value) {
|
||||
assert!(matches!(
|
||||
send(
|
||||
"text-only",
|
||||
json!({"messages": [{"role": "user", "content": [part]}]})
|
||||
),
|
||||
Err(Error::InvalidRequest(ErrorDetail::Message(_)))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn tools_are_refused_where_the_registry_rules_out_function_calling() {
|
||||
assert!(matches!(
|
||||
send(
|
||||
"text-only",
|
||||
json!({"messages": [], "tools": [{"type": "function", "function": {"name": "f"}}]})
|
||||
),
|
||||
Err(Error::InvalidRequest(ErrorDetail::UnsupportedParams { .. }))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_response_format_beside_tools_becomes_one_more_tool() {
|
||||
let sent = send("unknown", json!({"messages": [],
|
||||
"tools": [{"type": "function", "function": {"name": "f"}}],
|
||||
"response_format": {"type": "json_schema", "json_schema": {"schema": {"type": "object"}}}}))
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
sent.body["tools"],
|
||||
json!([{"type": "function", "function": {"name": "f"}},
|
||||
{"type": "function", "function": {"name": "json_tool_call", "parameters": {"type": "object"}}}])
|
||||
);
|
||||
assert!(!sent.body.contains_key("response_format"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn vllm_extra_body_keys_are_translated_or_dropped() {
|
||||
let sent = send(
|
||||
"unknown",
|
||||
json!({"messages": [], "truncate_prompt_tokens": 100,
|
||||
"guided_choice": ["a", "b"], "stop_token_ids": [1], "custom": 1}),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
Value::Object(sent.body),
|
||||
json!({"model": "accounts/fireworks/models/unknown", "messages": [],
|
||||
"prompt_truncate_len": 100,
|
||||
"response_format": {"type": "json_schema", "json_schema": {"name": "choice", "schema": {"type": "string", "enum": ["a", "b"]}}},
|
||||
"custom": 1})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_session_pins_the_call() {
|
||||
let body = Map::from_iter([("metadata".to_owned(), json!({"session_id": "s1"}))]);
|
||||
assert_eq!(
|
||||
FIREWORKS_AI_CHAT_CONFIG.headers(&body, &|_| None),
|
||||
vec![("x-session-affinity".to_owned(), "s1".to_owned())]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::sent_tool("f", true)]
|
||||
#[case::json_tool("json_tool_call", false)]
|
||||
#[case::unknown_tool("g", false)]
|
||||
fn tool_calls_in_the_content_move_to_tool_calls(#[case] name: &str, #[case] moved: bool) {
|
||||
let content = json!({"name": name, "arguments": {"a": 1}}).to_string();
|
||||
let mut body = Map::from_iter([
|
||||
("model".to_owned(), json!("accounts/fireworks/models/x")),
|
||||
(
|
||||
"choices".to_owned(),
|
||||
json!([{"message": {"role": "assistant", "content": content}}]),
|
||||
),
|
||||
]);
|
||||
let shape = ResponseShape {
|
||||
tools_sent: true,
|
||||
tool_names: vec!["f".into(), "json_tool_call".into()],
|
||||
..ResponseShape::default()
|
||||
};
|
||||
fix_response(&FIREWORKS_AI_CHAT_CONFIG, &mut body, &shape);
|
||||
assert_eq!(
|
||||
body["model"],
|
||||
json!("fireworks_ai/accounts/fireworks/models/x")
|
||||
);
|
||||
let message = &body["choices"][0]["message"];
|
||||
assert_eq!(message["content"].is_null(), moved);
|
||||
if moved {
|
||||
assert_eq!(
|
||||
message["tool_calls"][0]["function"],
|
||||
json!({"name": "f", "arguments": "{\"a\": 1}"})
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
1
litellm-rust/crates/llms/src/fireworks_ai/mod.rs
Normal file
1
litellm-rust/crates/llms/src/fireworks_ai/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod chat;
|
||||
11
litellm-rust/crates/llms/src/gemini/chat/AGENTS.md
Normal file
11
litellm-rust/crates/llms/src/gemini/chat/AGENTS.md
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
# rules
|
||||
|
||||
- Gemini on Google AI Studio: reuses `vertex_ai/gemini` for params, the body, the response and the stream, with `GeminiSurface::AiStudio`
|
||||
- The API key goes in the `x-goog-api-key` header, never in the URL
|
||||
- Gemini 3 and newer call `v1alpha`, older models `v1beta`; a custom `api_base` is called at `{api_base}/models/{model}:generateContent`
|
||||
|
||||
# references
|
||||
|
||||
- https://ai.google.dev/api/generate-content
|
||||
- litellm/llms/gemini/chat/transformation.py
|
||||
- litellm/llms/vertex_ai/vertex_llm_base.py
|
||||
1
litellm-rust/crates/llms/src/gemini/chat/mod.rs
Normal file
1
litellm-rust/crates/llms/src/gemini/chat/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
252
litellm-rust/crates/llms/src/gemini/chat/transformation.rs
Normal file
252
litellm-rust/crates/llms/src/gemini/chat/transformation.rs
Normal file
|
|
@ -0,0 +1,252 @@
|
|||
//! Python's `GoogleAIStudioGeminiConfig` for `/chat/completions`, with the key and URL
|
||||
//! `_get_token_and_url` and `_check_custom_proxy` build for `custom_llm_provider="gemini"`.
|
||||
|
||||
use litellm_auth::{CredentialPlacement, SecretValue};
|
||||
use litellm_llms_types::{
|
||||
formats::chat_completions::{ChatCompletionsResponse, ChatMessage},
|
||||
providers::vertex_ai::{GEMINI_API_BASE, GEMINI_API_KEY_HEADER, GENERATE_CONTENT},
|
||||
};
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{
|
||||
Error,
|
||||
base_llm::{
|
||||
auth::{AuthScheme, Headers, ValidatedEnvironment},
|
||||
chat::{
|
||||
params::ChatRequestContext,
|
||||
streaming::{ChatStream, StreamShape},
|
||||
transformation::{
|
||||
BaseConfig, ProviderChatRequestData, ProviderChatResponseData, Unsupported,
|
||||
},
|
||||
},
|
||||
},
|
||||
vertex_ai::gemini::{
|
||||
handler::gemini_chat_stream,
|
||||
transformation::remote_media,
|
||||
vertex_and_google_ai_studio_gemini::{
|
||||
CONFIG_PARAMS, GeminiSurface, build_request, is_gemini_3_or_newer, streaming_url,
|
||||
transform_generate_content_response,
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
pub const GEMINI_API_BASE_ENV: &str = "GEMINI_API_BASE";
|
||||
pub const GEMINI_API_KEY_ENV: &str = "GEMINI_API_KEY";
|
||||
|
||||
/// `get_api_key_from_env`, then the older PaLM key.
|
||||
const API_KEY_ENVS: &[&str] = &["GOOGLE_API_KEY", GEMINI_API_KEY_ENV, "PALM_API_KEY"];
|
||||
|
||||
pub struct GoogleAiStudioGeminiConfig;
|
||||
|
||||
pub const GOOGLE_AI_STUDIO_GEMINI_CHAT_COMPLETIONS_CONFIG: GoogleAiStudioGeminiConfig =
|
||||
GoogleAiStudioGeminiConfig;
|
||||
|
||||
/// The id AI Studio is called with: Python strips the provider prefix.
|
||||
fn studio_model(model: &str) -> &str {
|
||||
model.strip_prefix("gemini/").unwrap_or(model)
|
||||
}
|
||||
|
||||
fn api_key(
|
||||
api_key: Option<&str>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<String, Error> {
|
||||
api_key
|
||||
.filter(|key| !key.is_empty())
|
||||
.map(str::to_owned)
|
||||
.or_else(|| {
|
||||
API_KEY_ENVS
|
||||
.iter()
|
||||
.find_map(|name| env_lookup(name).filter(|key| !key.is_empty()))
|
||||
})
|
||||
.ok_or(Error::Auth(litellm_auth::Error::MissingApiKey {
|
||||
provider: "Gemini",
|
||||
environment_variable: GEMINI_API_KEY_ENV,
|
||||
}))
|
||||
}
|
||||
|
||||
impl BaseConfig for GoogleAiStudioGeminiConfig {
|
||||
fn secret_names(&self) -> Vec<&'static str> {
|
||||
API_KEY_ENVS
|
||||
.iter()
|
||||
.copied()
|
||||
.chain([GEMINI_API_BASE_ENV])
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// `get_gemini_url`, or `{api_base}/models/{model}:generateContent` for a custom base.
|
||||
fn get_complete_url(
|
||||
&self,
|
||||
api_base: Option<&str>,
|
||||
model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<String, Error> {
|
||||
let model = studio_model(model);
|
||||
let custom = api_base
|
||||
.map(str::to_owned)
|
||||
.or_else(|| env_lookup(GEMINI_API_BASE_ENV))
|
||||
.filter(|base| !base.is_empty());
|
||||
Ok(match custom {
|
||||
Some(base) => format!("{base}/models/{model}:{GENERATE_CONTENT}"),
|
||||
None => {
|
||||
let version = match is_gemini_3_or_newer(model) {
|
||||
true => "v1alpha",
|
||||
false => "v1beta",
|
||||
};
|
||||
format!("{GEMINI_API_BASE}/{version}/models/{model}:{GENERATE_CONTENT}")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
fn transform_request(
|
||||
&self,
|
||||
model: &str,
|
||||
messages: Vec<ChatMessage>,
|
||||
optional_params: Map<String, Value>,
|
||||
context: &ChatRequestContext<'_>,
|
||||
) -> Result<ProviderChatRequestData, Error> {
|
||||
build_request(
|
||||
GeminiSurface::AiStudio,
|
||||
model,
|
||||
messages,
|
||||
optional_params,
|
||||
context,
|
||||
)
|
||||
}
|
||||
|
||||
fn transform_response(
|
||||
&self,
|
||||
model: &str,
|
||||
response: ProviderChatResponseData,
|
||||
_shape: &StreamShape,
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
transform_generate_content_response(model, response.body)
|
||||
}
|
||||
|
||||
fn remote_media(&self, messages: &[ChatMessage]) -> Vec<String> {
|
||||
remote_media(messages, GeminiSurface::AiStudio)
|
||||
}
|
||||
|
||||
fn model_response_iterator(&self, _shape: StreamShape) -> Option<ChatStream> {
|
||||
Some(gemini_chat_stream())
|
||||
}
|
||||
|
||||
fn streaming_request(&self, url: String, body: Value) -> Result<(String, Value), Error> {
|
||||
Ok((streaming_url(&url), body))
|
||||
}
|
||||
|
||||
/// The key always wins over a forwarded header, as Python writes `x-goog-api-key` last.
|
||||
fn validate_environment(
|
||||
&self,
|
||||
headers: Headers,
|
||||
api_key_param: Option<&str>,
|
||||
_model: &str,
|
||||
_optional_params: &Map<String, Value>,
|
||||
env_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<ValidatedEnvironment, Error> {
|
||||
Ok(ValidatedEnvironment {
|
||||
headers,
|
||||
auth: AuthScheme::Credential {
|
||||
placement: CredentialPlacement::Header(GEMINI_API_KEY_HEADER),
|
||||
secret: SecretValue::new(api_key(api_key_param, env_lookup)?),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
fn config_params(&self) -> &'static [&'static str] {
|
||||
CONFIG_PARAMS
|
||||
}
|
||||
|
||||
/// Every request is translated; a param or message Python refuses is refused here too.
|
||||
fn unsupported_reason(
|
||||
&self,
|
||||
_messages: &[ChatMessage],
|
||||
_optional_params: &Map<String, Value>,
|
||||
) -> Option<Unsupported> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn no_env(_: &str) -> Option<String> {
|
||||
None
|
||||
}
|
||||
|
||||
fn base_env(name: &str) -> Option<String> {
|
||||
(name == GEMINI_API_BASE_ENV).then(|| "https://proxy.example/gemini".to_owned())
|
||||
}
|
||||
|
||||
fn key_env(name: &str) -> Option<String> {
|
||||
match name {
|
||||
"GEMINI_API_KEY" => Some("gemini-key".into()),
|
||||
"PALM_API_KEY" => Some("palm-key".into()),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::v1beta(None, "gemini/gemini-2.5-flash", no_env as fn(&str) -> Option<String>, "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent")]
|
||||
#[case::v1alpha(None, "gemini-3-pro-preview", no_env as fn(&str) -> Option<String>, "https://generativelanguage.googleapis.com/v1alpha/models/gemini-3-pro-preview:generateContent")]
|
||||
#[case::custom(Some("https://gw.example/v1beta"), "gemini-2.5-flash", no_env as fn(&str) -> Option<String>, "https://gw.example/v1beta/models/gemini-2.5-flash:generateContent")]
|
||||
#[case::env_base(None, "gemini-2.5-flash", base_env as fn(&str) -> Option<String>, "https://proxy.example/gemini/models/gemini-2.5-flash:generateContent")]
|
||||
fn urls_follow_pythons_gemini_url(
|
||||
#[case] api_base: Option<&str>,
|
||||
#[case] model: &str,
|
||||
#[case] env: fn(&str) -> Option<String>,
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let url = GOOGLE_AI_STUDIO_GEMINI_CHAT_COMPLETIONS_CONFIG
|
||||
.get_complete_url(api_base, model, &Map::new(), &env)
|
||||
.unwrap();
|
||||
assert_eq!(url, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::explicit(Some("explicit"), "explicit")]
|
||||
#[case::from_env(None, "gemini-key")]
|
||||
fn the_key_goes_in_the_goog_header(#[case] key: Option<&str>, #[case] expected: &str) {
|
||||
let validated = GOOGLE_AI_STUDIO_GEMINI_CHAT_COMPLETIONS_CONFIG
|
||||
.validate_environment(Vec::new(), key, "gemini-2.5-flash", &Map::new(), &key_env)
|
||||
.unwrap();
|
||||
assert!(matches!(
|
||||
validated.auth,
|
||||
AuthScheme::Credential {
|
||||
placement: CredentialPlacement::Header("x-goog-api-key"),
|
||||
ref secret
|
||||
} if secret.expose() == expected
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_missing_key_is_an_auth_error() {
|
||||
let error = GOOGLE_AI_STUDIO_GEMINI_CHAT_COMPLETIONS_CONFIG
|
||||
.validate_environment(Vec::new(), None, "gemini-2.5-flash", &Map::new(), &no_env)
|
||||
.unwrap_err();
|
||||
assert!(matches!(
|
||||
error,
|
||||
Error::Auth(litellm_auth::Error::MissingApiKey {
|
||||
provider: "Gemini",
|
||||
..
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn streams_call_stream_generate_content_with_sse() {
|
||||
let (url, _) = GOOGLE_AI_STUDIO_GEMINI_CHAT_COMPLETIONS_CONFIG
|
||||
.streaming_request(
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/m:generateContent".into(),
|
||||
Value::Null,
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
url,
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/m:streamGenerateContent?alt=sse"
|
||||
);
|
||||
}
|
||||
}
|
||||
1
litellm-rust/crates/llms/src/gemini/mod.rs
Normal file
1
litellm-rust/crates/llms/src/gemini/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod chat;
|
||||
12
litellm-rust/crates/llms/src/groq/chat/AGENTS.md
Normal file
12
litellm-rust/crates/llms/src/groq/chat/AGENTS.md
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
# rules
|
||||
|
||||
- Groq at `https://api.groq.com/openai/v1`
|
||||
- A JSON schema on a model without native structured outputs goes as a forced `json_tool_call` tool, refused beside caller tools
|
||||
- A streamed `response_format` is refused: Python fakes that stream from a complete call
|
||||
- `web_search_options` adds the `browser_search` tool; `service_tier` in the answer is held to OpenAI's values
|
||||
- Shared flow and response handling: `../../openai/chat/compatible.rs`
|
||||
|
||||
# references
|
||||
|
||||
- https://console.groq.com/docs/api-reference#chat
|
||||
- litellm/llms/groq/chat/transformation.py
|
||||
1
litellm-rust/crates/llms/src/groq/chat/mod.rs
Normal file
1
litellm-rust/crates/llms/src/groq/chat/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
368
litellm-rust/crates/llms/src/groq/chat/transformation.rs
Normal file
368
litellm-rust/crates/llms/src/groq/chat/transformation.rs
Normal file
|
|
@ -0,0 +1,368 @@
|
|||
//! `litellm/llms/groq/chat/transformation.py`: Groq at `https://api.groq.com/openai/v1`. A
|
||||
//! JSON schema a model cannot take natively goes as a forced `json_tool_call` tool whose call
|
||||
//! comes back as content, `web_search_options` turns on the `browser_search` tool, and the
|
||||
//! response's `service_tier` is held to the values OpenAI defines.
|
||||
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::{
|
||||
Error, ErrorDetail,
|
||||
openai::chat::compatible::{
|
||||
CompatibleContext, GPT_BASE_PARAMS, OpenAiCompatibleChatConfig, RESPONSE_FORMAT_TOOL_NAME,
|
||||
ResponseShape, copy_supported, gpt_messages, json_mode_tool_call_to_content, messages_mut,
|
||||
role, sanitize_usage, strip_response_only_fields,
|
||||
},
|
||||
};
|
||||
|
||||
const COMPOUND_MODELS: &[&str] = &["compound", "compound-mini"];
|
||||
const SERVICE_TIERS: &[&str] = &["auto", "default", "flex"];
|
||||
|
||||
pub struct GroqChatConfig;
|
||||
|
||||
pub const GROQ_CHAT_CONFIG: GroqChatConfig = GroqChatConfig;
|
||||
|
||||
impl OpenAiCompatibleChatConfig for GroqChatConfig {
|
||||
fn provider(&self) -> LlmProviders {
|
||||
LlmProviders::Groq
|
||||
}
|
||||
|
||||
fn label(&self) -> &'static str {
|
||||
"Groq"
|
||||
}
|
||||
|
||||
fn api_base_env(&self) -> &'static [&'static str] {
|
||||
&["GROQ_API_BASE"]
|
||||
}
|
||||
|
||||
fn default_api_base(&self) -> Option<&'static str> {
|
||||
Some("https://api.groq.com/openai/v1")
|
||||
}
|
||||
|
||||
fn api_key_env(&self) -> &'static [&'static str] {
|
||||
&["GROQ_API_KEY"]
|
||||
}
|
||||
|
||||
fn supported_params(&self, context: &CompatibleContext<'_>) -> Vec<&'static str> {
|
||||
let web_search =
|
||||
is_compound(context.model) || context.supports(self.name(), "supports_web_search");
|
||||
let mut params: Vec<&'static str> = GPT_BASE_PARAMS
|
||||
.iter()
|
||||
.copied()
|
||||
.filter(|param| *param != "max_retries")
|
||||
.filter(|param| web_search || *param != "web_search_options")
|
||||
.chain(["response_format"])
|
||||
.collect();
|
||||
if context.supports(self.name(), "supports_reasoning") {
|
||||
params.push("reasoning_effort");
|
||||
}
|
||||
params
|
||||
}
|
||||
|
||||
fn map_params(
|
||||
&self,
|
||||
non_default: &Map<String, Value>,
|
||||
optional: &mut Map<String, Value>,
|
||||
context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
let mut params = non_default.clone();
|
||||
if params.contains_key("response_format")
|
||||
&& params.get("stream") == Some(&Value::Bool(true))
|
||||
{
|
||||
return Err(Error::Unsupported(
|
||||
"streaming a groq response_format, which Python fakes from a complete call",
|
||||
));
|
||||
}
|
||||
if let Some(schema) = response_schema(params.get("response_format"))
|
||||
&& context.capability(self.name(), "supports_response_schema") != Some(true)
|
||||
{
|
||||
if params.contains_key("tools") {
|
||||
return Err(Error::InvalidRequest(ErrorDetail::Message(format!(
|
||||
"Groq model '{}' does not support native structured outputs. LiteLLM uses a \
|
||||
tool-calling workaround for structured outputs on this model, which is \
|
||||
incompatible with user-provided tools. Either use a model that supports \
|
||||
native structured outputs (e.g., gpt-oss-120b, llama-4, kimi-k2), or remove \
|
||||
the tools parameter. See: \
|
||||
https://console.groq.com/docs/structured-outputs#supported-models",
|
||||
context.model
|
||||
))));
|
||||
}
|
||||
optional.insert(
|
||||
"tools".into(),
|
||||
json!([{"type": "function", "function": {"name": RESPONSE_FORMAT_TOOL_NAME, "parameters": schema}}]),
|
||||
);
|
||||
optional.insert(
|
||||
"tool_choice".into(),
|
||||
json!({"type": "function", "function": {"name": RESPONSE_FORMAT_TOOL_NAME}}),
|
||||
);
|
||||
optional.insert("json_mode".into(), Value::Bool(true));
|
||||
params.remove("response_format");
|
||||
}
|
||||
let web_search = params.remove("web_search_options");
|
||||
copy_supported(¶ms, optional, &self.supported_params(context));
|
||||
if let Some(max_tokens) = optional.remove("max_completion_tokens") {
|
||||
optional.insert("max_tokens".into(), max_tokens);
|
||||
}
|
||||
if web_search.is_none() || is_compound(context.model) {
|
||||
return Ok(());
|
||||
}
|
||||
let tools = optional
|
||||
.entry("tools")
|
||||
.or_insert_with(|| Value::Array(Vec::new()));
|
||||
if tools.is_null() {
|
||||
*tools = Value::Array(Vec::new());
|
||||
}
|
||||
if let Value::Array(tools) = tools
|
||||
&& !tools
|
||||
.iter()
|
||||
.any(|tool| tool.get("type").and_then(Value::as_str) == Some("browser_search"))
|
||||
{
|
||||
tools.push(json!({"type": "browser_search"}));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn transform_messages(
|
||||
&self,
|
||||
messages: Vec<Value>,
|
||||
_optional: &Map<String, Value>,
|
||||
_context: &CompatibleContext<'_>,
|
||||
) -> Result<Vec<Value>, Error> {
|
||||
gpt_messages(
|
||||
messages
|
||||
.into_iter()
|
||||
.map(|mut message| {
|
||||
if role(&message) == Some("assistant")
|
||||
&& let Value::Object(object) = &mut message
|
||||
{
|
||||
object.retain(|_, value| !value.is_null());
|
||||
strip_response_only_fields(&mut message);
|
||||
}
|
||||
message
|
||||
})
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
fn fix_response(&self, body: &mut Map<String, Value>, shape: &ResponseShape) -> bool {
|
||||
let mut changed = sanitize_usage(body);
|
||||
if shape.json_mode {
|
||||
changed |= json_mode_tool_call_to_content(body);
|
||||
}
|
||||
let tier = body.get("service_tier").and_then(Value::as_str);
|
||||
if !tier.is_some_and(|tier| SERVICE_TIERS.contains(&tier)) {
|
||||
body.insert("service_tier".into(), json!("auto"));
|
||||
changed = true;
|
||||
}
|
||||
changed | add_web_search_usage(body)
|
||||
}
|
||||
}
|
||||
|
||||
fn is_compound(model: &str) -> bool {
|
||||
COMPOUND_MODELS.contains(&model.strip_prefix("groq/").unwrap_or(model))
|
||||
}
|
||||
|
||||
fn response_schema(response_format: Option<&Value>) -> Option<Value> {
|
||||
let format = response_format?.as_object()?;
|
||||
match format.get("response_schema") {
|
||||
Some(schema) => Some(schema.clone()),
|
||||
None => format.get("json_schema")?.get("schema").cloned(),
|
||||
}
|
||||
.filter(|schema| !schema.is_null())
|
||||
}
|
||||
|
||||
/// `_add_web_search_usage`: the searches and page opens Groq ran, from each message's
|
||||
/// `executed_tools`.
|
||||
fn add_web_search_usage(body: &mut Map<String, Value>) -> bool {
|
||||
if !body.get("usage").is_some_and(Value::is_object) {
|
||||
return false;
|
||||
}
|
||||
let actions: Vec<String> = messages_mut(body)
|
||||
.filter_map(|message| message.get("executed_tools")?.as_array())
|
||||
.flatten()
|
||||
.filter_map(|tool| {
|
||||
tool.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|name| !name.is_empty())
|
||||
.or_else(|| tool.get("type").and_then(Value::as_str))
|
||||
.map(str::to_owned)
|
||||
})
|
||||
.collect();
|
||||
let count = |names: &[&str]| {
|
||||
actions
|
||||
.iter()
|
||||
.filter(|a| names.contains(&a.as_str()))
|
||||
.count()
|
||||
};
|
||||
let searches = count(&["browser.search", "browser_search"]);
|
||||
let opens = count(&["browser.open"]);
|
||||
if searches == 0 && opens == 0 {
|
||||
return false;
|
||||
}
|
||||
if let Some(Value::Object(usage)) = body.get_mut("usage") {
|
||||
usage.insert(
|
||||
"server_tool_use".into(),
|
||||
json!({"web_search_requests": searches, "browser_open_requests": opens}),
|
||||
);
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::borrow::Cow;
|
||||
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::{
|
||||
base_llm::model_info::ModelInfo,
|
||||
openai::chat::compatible::{fix_response, transform_request},
|
||||
};
|
||||
|
||||
fn object(value: Value) -> Map<String, Value> {
|
||||
let Value::Object(map) = value else {
|
||||
panic!("an object")
|
||||
};
|
||||
map
|
||||
}
|
||||
|
||||
fn send(model: &str, params: Value) -> Result<Map<String, Value>, Error> {
|
||||
transform_request(
|
||||
&GROQ_CHAT_CONFIG,
|
||||
&CompatibleContext {
|
||||
model,
|
||||
drop_params: false,
|
||||
model_info: None,
|
||||
},
|
||||
object(params),
|
||||
)
|
||||
.map(|sent| sent.body)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_schema_without_native_support_goes_as_a_forced_tool() {
|
||||
let sent = transform_request(
|
||||
&GROQ_CHAT_CONFIG,
|
||||
&CompatibleContext {
|
||||
model: "llama-3.1-8b-instant",
|
||||
drop_params: false,
|
||||
model_info: None,
|
||||
},
|
||||
object(json!({"messages": [], "max_completion_tokens": 3,
|
||||
"response_format": {"type": "json_schema", "json_schema": {"name": "x", "schema": {"type": "object"}}}})),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
Value::Object(sent.body),
|
||||
json!({"model": "llama-3.1-8b-instant", "messages": [],
|
||||
"tools": [{"type": "function", "function": {"name": "json_tool_call", "parameters": {"type": "object"}}}],
|
||||
"tool_choice": {"type": "function", "function": {"name": "json_tool_call"}},
|
||||
"max_tokens": 3})
|
||||
);
|
||||
assert!(sent.shape.json_mode);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_schema_workaround_with_user_tools_is_rejected() {
|
||||
assert!(matches!(
|
||||
send(
|
||||
"llama-3.1-8b-instant",
|
||||
json!({"messages": [], "tools": [{"type": "function", "function": {"name": "f"}}],
|
||||
"response_format": {"type": "json_schema", "json_schema": {"schema": {}}}})
|
||||
),
|
||||
Err(Error::InvalidRequest(ErrorDetail::Message(_)))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_streamed_response_format_is_refused() {
|
||||
assert!(matches!(
|
||||
send(
|
||||
"llama-3.1-8b-instant",
|
||||
json!({"messages": [], "stream": true, "response_format": {"type": "json_object"}})
|
||||
),
|
||||
Err(Error::Unsupported(_))
|
||||
));
|
||||
}
|
||||
|
||||
struct WebSearch;
|
||||
|
||||
impl ModelInfo for WebSearch {
|
||||
fn entry(&self, model: &str) -> Option<Cow<'_, Map<String, Value>>> {
|
||||
(model == "groq/openai/gpt-oss-20b").then(|| {
|
||||
Cow::Owned(Map::from_iter([(
|
||||
"supports_web_search".to_owned(),
|
||||
json!(true),
|
||||
)]))
|
||||
})
|
||||
}
|
||||
|
||||
fn declares(&self, _key: &str) -> bool {
|
||||
true
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::adds_browser_search("openai/gpt-oss-20b", json!([{"type": "browser_search"}]))]
|
||||
#[case::compound_searches_itself("compound", Value::Null)]
|
||||
fn web_search_turns_on_browser_search(#[case] model: &str, #[case] tools: Value) {
|
||||
let sent = transform_request(
|
||||
&GROQ_CHAT_CONFIG,
|
||||
&CompatibleContext {
|
||||
model,
|
||||
drop_params: false,
|
||||
model_info: Some(&WebSearch),
|
||||
},
|
||||
object(json!({"messages": [], "web_search_options": {"search_context_size": "low"}})),
|
||||
)
|
||||
.unwrap()
|
||||
.body;
|
||||
assert_eq!(sent.get("tools").cloned().unwrap_or(Value::Null), tools);
|
||||
assert!(!sent.contains_key("web_search_options"));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn assistant_turns_lose_nulls_and_response_only_fields() {
|
||||
let sent = send(
|
||||
"llama-3.1-8b-instant",
|
||||
json!({"messages": [
|
||||
{"role": "assistant", "content": "a", "function_call": null, "thinking_blocks": []}]}),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
sent["messages"],
|
||||
json!([{"role": "assistant", "content": "a"}])
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::unknown(json!("on_demand"), json!("auto"))]
|
||||
#[case::missing(Value::Null, json!("auto"))]
|
||||
#[case::kept(json!("flex"), json!("flex"))]
|
||||
fn service_tier_is_held_to_openai_values(#[case] tier: Value, #[case] expected: Value) {
|
||||
let mut body =
|
||||
object(json!({"choices": [], "service_tier": tier, "usage": {"prompt_tokens": null}}));
|
||||
fix_response(&GROQ_CHAT_CONFIG, &mut body, &ResponseShape::default());
|
||||
assert_eq!(body["service_tier"], expected);
|
||||
assert_eq!(body["usage"]["prompt_tokens"], json!(0));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn executed_browser_tools_reach_the_usage() {
|
||||
let mut body = object(
|
||||
json!({"service_tier": "auto", "usage": {"prompt_tokens": 1},
|
||||
"choices": [{"message": {"role": "assistant", "content": "x",
|
||||
"executed_tools": [{"type": "browser_search"}, {"name": "browser.open"}, {"name": "browser.search"}]}}]}),
|
||||
);
|
||||
assert!(fix_response(
|
||||
&GROQ_CHAT_CONFIG,
|
||||
&mut body,
|
||||
&ResponseShape::default()
|
||||
));
|
||||
assert_eq!(
|
||||
body["usage"]["server_tool_use"],
|
||||
json!({"web_search_requests": 2, "browser_open_requests": 1})
|
||||
);
|
||||
}
|
||||
}
|
||||
1
litellm-rust/crates/llms/src/groq/mod.rs
Normal file
1
litellm-rust/crates/llms/src/groq/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod chat;
|
||||
11
litellm-rust/crates/llms/src/hosted_vllm/chat/AGENTS.md
Normal file
11
litellm-rust/crates/llms/src/hosted_vllm/chat/AGENTS.md
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
# rules
|
||||
|
||||
- vLLM's OpenAI-compatible server at the deployment's `api_base`; there is no default host, and an unset key sends `fake-api-key`
|
||||
- Tools lose `additionalProperties: false` and `strict`, custom tools go as function tools, `thinking` becomes `reasoning_effort`
|
||||
- Replayed assistant turns lose `thinking_blocks` and non-string `reasoning_content`, content lists flatten to text
|
||||
- Shared flow and response handling: `../../openai/chat/compatible.rs`
|
||||
|
||||
# references
|
||||
|
||||
- https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html
|
||||
- litellm/llms/hosted_vllm/chat/transformation.py
|
||||
1
litellm-rust/crates/llms/src/hosted_vllm/chat/mod.rs
Normal file
1
litellm-rust/crates/llms/src/hosted_vllm/chat/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
385
litellm-rust/crates/llms/src/hosted_vllm/chat/transformation.rs
Normal file
385
litellm-rust/crates/llms/src/hosted_vllm/chat/transformation.rs
Normal file
|
|
@ -0,0 +1,385 @@
|
|||
//! `litellm/llms/hosted_vllm/chat/transformation.py`: a vLLM server at the deployment's
|
||||
//! `api_base`. Tools lose `additionalProperties: false` and `strict` and custom tools become
|
||||
//! function tools, `thinking` becomes a `reasoning_effort`, replayed assistant turns lose the
|
||||
//! fields vLLM rejects with their content lists flattened, and video files go as `video_url`.
|
||||
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::{
|
||||
Error,
|
||||
base_llm::chat::params::python_dumps,
|
||||
openai::chat::compatible::{
|
||||
CompatibleContext, GPT_BASE_PARAMS, OpenAiCompatibleChatConfig, copy_supported,
|
||||
gpt_messages, reasoning_effort_from_thinking_budget, role,
|
||||
},
|
||||
vertex_ai::gemini::schema::{remove_additional_properties, remove_strict_from_schema},
|
||||
};
|
||||
|
||||
pub struct HostedVllmChatConfig;
|
||||
|
||||
pub const HOSTED_VLLM_CHAT_CONFIG: HostedVllmChatConfig = HostedVllmChatConfig;
|
||||
|
||||
impl OpenAiCompatibleChatConfig for HostedVllmChatConfig {
|
||||
fn provider(&self) -> LlmProviders {
|
||||
LlmProviders::HostedVllm
|
||||
}
|
||||
|
||||
fn label(&self) -> &'static str {
|
||||
"hosted vLLM"
|
||||
}
|
||||
|
||||
fn api_base_env(&self) -> &'static [&'static str] {
|
||||
&["HOSTED_VLLM_API_BASE"]
|
||||
}
|
||||
|
||||
fn default_api_base(&self) -> Option<&'static str> {
|
||||
None
|
||||
}
|
||||
|
||||
fn api_key_env(&self) -> &'static [&'static str] {
|
||||
&["HOSTED_VLLM_API_KEY"]
|
||||
}
|
||||
|
||||
/// vLLM needs no key, and Python sends this one when none is set.
|
||||
fn default_api_key(&self) -> Option<&'static str> {
|
||||
Some("fake-api-key")
|
||||
}
|
||||
|
||||
fn supported_params(&self, _context: &CompatibleContext<'_>) -> Vec<&'static str> {
|
||||
GPT_BASE_PARAMS
|
||||
.iter()
|
||||
.copied()
|
||||
.chain(["response_format", "reasoning_effort", "thinking"])
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn map_params(
|
||||
&self,
|
||||
non_default: &Map<String, Value>,
|
||||
optional: &mut Map<String, Value>,
|
||||
context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
let mut params = non_default.clone();
|
||||
if let Some(mut tools) = params.remove("tools").filter(|tools| !tools.is_null()) {
|
||||
remove_additional_properties(&mut tools);
|
||||
remove_strict_from_schema(&mut tools);
|
||||
if let Value::Array(items) = tools {
|
||||
tools = Value::Array(
|
||||
items
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(index, tool)| custom_tool_as_function(index, tool))
|
||||
.collect(),
|
||||
);
|
||||
}
|
||||
params.insert("tools".into(), tools);
|
||||
}
|
||||
if let Some(Value::Object(thinking)) = params.remove("thinking")
|
||||
&& thinking.get("type").and_then(Value::as_str) == Some("enabled")
|
||||
&& !params.contains_key("reasoning_effort")
|
||||
{
|
||||
let budget = thinking
|
||||
.get("budget_tokens")
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or(0);
|
||||
params.insert(
|
||||
"reasoning_effort".into(),
|
||||
json!(reasoning_effort_from_thinking_budget(budget)),
|
||||
);
|
||||
}
|
||||
copy_supported(¶ms, optional, &self.supported_params(context));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn transform_messages(
|
||||
&self,
|
||||
messages: Vec<Value>,
|
||||
_optional: &Map<String, Value>,
|
||||
_context: &CompatibleContext<'_>,
|
||||
) -> Result<Vec<Value>, Error> {
|
||||
let messages = messages
|
||||
.into_iter()
|
||||
.map(|mut message| {
|
||||
match role(&message) {
|
||||
Some("assistant") => assistant_for_vllm(&mut message),
|
||||
Some("user") => video_files_as_video_urls(&mut message),
|
||||
_ => {}
|
||||
}
|
||||
message
|
||||
})
|
||||
.collect();
|
||||
gpt_messages(messages)
|
||||
}
|
||||
}
|
||||
|
||||
/// `_convert_custom_tools_to_function_tools` for one tool.
|
||||
fn custom_tool_as_function(index: usize, tool: Value) -> Value {
|
||||
if tool.get("type").and_then(Value::as_str) != Some("custom") {
|
||||
return tool;
|
||||
}
|
||||
let custom = tool
|
||||
.get("custom")
|
||||
.filter(|custom| custom.is_object())
|
||||
.cloned()
|
||||
.unwrap_or_else(|| json!({}));
|
||||
let pick = |key: &str| {
|
||||
custom
|
||||
.get(key)
|
||||
.filter(|value| crate::base_llm::chat::params::truthy(value))
|
||||
.or_else(|| tool.get(key).filter(|value| !value.is_null()))
|
||||
.cloned()
|
||||
};
|
||||
let name = match pick("name") {
|
||||
Some(Value::String(name)) => name,
|
||||
Some(other) => python_str(&other),
|
||||
None => format!("custom_tool_{index}"),
|
||||
};
|
||||
let parameters = match pick("input_schema") {
|
||||
Some(schema @ Value::Object(_)) => schema,
|
||||
_ => json!({
|
||||
"type": "object",
|
||||
"properties": {"input": {"type": "string", "description": "Raw tool input payload."}},
|
||||
"required": ["input"],
|
||||
}),
|
||||
};
|
||||
let mut function = Map::from_iter([
|
||||
("name".to_owned(), Value::String(name)),
|
||||
("parameters".to_owned(), parameters),
|
||||
]);
|
||||
if let Some(Value::String(description)) = pick("description") {
|
||||
function.insert("description".into(), Value::String(description));
|
||||
}
|
||||
json!({"type": "function", "function": function})
|
||||
}
|
||||
|
||||
fn python_str(value: &Value) -> String {
|
||||
match value {
|
||||
Value::String(text) => text.clone(),
|
||||
other => python_dumps(other),
|
||||
}
|
||||
}
|
||||
|
||||
/// The assistant half of `HostedVLLMChatConfig._transform_messages`.
|
||||
fn assistant_for_vllm(message: &mut Value) {
|
||||
let Some(object) = message.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
object.remove("thinking_blocks");
|
||||
if !object
|
||||
.get("reasoning_content")
|
||||
.is_some_and(Value::is_string)
|
||||
{
|
||||
object.remove("reasoning_content");
|
||||
}
|
||||
let Some(Value::Array(parts)) = object.get("content").cloned() else {
|
||||
return;
|
||||
};
|
||||
let mut texts = Vec::new();
|
||||
let mut tool_calls = Vec::new();
|
||||
let mut blocks = Vec::new();
|
||||
let mut structured = false;
|
||||
for part in parts {
|
||||
match part.get("type").and_then(Value::as_str) {
|
||||
Some("text") if part.is_object() => {
|
||||
texts.push(
|
||||
part.get("text")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_owned(),
|
||||
);
|
||||
blocks.push(part);
|
||||
}
|
||||
Some("tool_use") if part.is_object() => {
|
||||
let input = part.get("input").cloned().unwrap_or_else(|| json!({}));
|
||||
let arguments = match input {
|
||||
Value::String(text) => text,
|
||||
other => python_dumps(&other),
|
||||
};
|
||||
tool_calls.push(json!({
|
||||
"id": part.get("id").cloned().unwrap_or(Value::Null),
|
||||
"type": "function",
|
||||
"function": {"name": part.get("name").cloned().unwrap_or(Value::Null), "arguments": arguments},
|
||||
}));
|
||||
}
|
||||
_ => {
|
||||
blocks.push(part);
|
||||
structured = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
if !tool_calls.is_empty() {
|
||||
match object.get_mut("tool_calls") {
|
||||
Some(Value::Array(existing)) => {
|
||||
let known: Vec<Value> = existing
|
||||
.iter()
|
||||
.filter_map(|call| call.get("id").filter(|id| !id.is_null()).cloned())
|
||||
.collect();
|
||||
existing.extend(
|
||||
tool_calls
|
||||
.into_iter()
|
||||
.filter(|call| !known.contains(&call["id"])),
|
||||
);
|
||||
}
|
||||
_ => {
|
||||
object.insert("tool_calls".into(), Value::Array(tool_calls));
|
||||
}
|
||||
}
|
||||
}
|
||||
let content = match structured {
|
||||
true => Value::Array(blocks),
|
||||
false => Value::String(texts.join("\n")),
|
||||
};
|
||||
object.insert("content".into(), content);
|
||||
}
|
||||
|
||||
/// `_is_video_file`: a `file` part whose format, data URL or file id names a video.
|
||||
fn is_video_file(part: &Value) -> bool {
|
||||
if part.get("type").and_then(Value::as_str) != Some("file") {
|
||||
return false;
|
||||
}
|
||||
let file = part.get("file");
|
||||
let field = |key: &str| {
|
||||
file.and_then(|file| file.get(key))
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.is_empty())
|
||||
};
|
||||
if field("format").is_some_and(|format| format.starts_with("video/")) {
|
||||
return true;
|
||||
}
|
||||
if let Some(data) = field("file_data") {
|
||||
return data_url_mime(data).is_some_and(|mime| mime.starts_with("video/"));
|
||||
}
|
||||
if let Some(id) = field("file_id") {
|
||||
return url_video_mime(id);
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// `parse_mime_type`: the media type of a `data:` URL.
|
||||
fn data_url_mime(data: &str) -> Option<&str> {
|
||||
let rest = data.strip_prefix("data:")?;
|
||||
let end = rest.find([';', ','])?;
|
||||
Some(&rest[..end])
|
||||
}
|
||||
|
||||
/// `get_image_mime_type_from_url` for the video extensions it knows.
|
||||
fn url_video_mime(url: &str) -> bool {
|
||||
let path = url.split(['?', '#']).next().unwrap_or(url).to_lowercase();
|
||||
[
|
||||
".mp4", ".mov", ".avi", ".flv", ".mpeg", ".mpg", ".webm", ".wmv", ".3gpp",
|
||||
]
|
||||
.iter()
|
||||
.any(|extension| path.ends_with(extension))
|
||||
}
|
||||
|
||||
fn video_files_as_video_urls(message: &mut Value) {
|
||||
let Some(Value::Array(parts)) = message.get_mut("content") else {
|
||||
return;
|
||||
};
|
||||
for part in parts.iter_mut() {
|
||||
if !is_video_file(part) {
|
||||
continue;
|
||||
}
|
||||
let file = part.get("file").cloned().unwrap_or_default();
|
||||
let url = ["file_id", "file_data"]
|
||||
.iter()
|
||||
.find_map(|key| {
|
||||
file.get(*key)
|
||||
.filter(|value| crate::base_llm::chat::params::truthy(value))
|
||||
})
|
||||
.cloned();
|
||||
if let Some(url) = url {
|
||||
*part = json!({"type": "video_url", "video_url": {"url": url}});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::openai::chat::compatible::transform_request;
|
||||
|
||||
fn send(params: Value) -> Map<String, Value> {
|
||||
let Value::Object(body) = params else {
|
||||
panic!("an object")
|
||||
};
|
||||
transform_request(
|
||||
&HOSTED_VLLM_CHAT_CONFIG,
|
||||
&CompatibleContext {
|
||||
model: "Qwen/Qwen3-8B",
|
||||
drop_params: false,
|
||||
model_info: None,
|
||||
},
|
||||
body,
|
||||
)
|
||||
.unwrap()
|
||||
.body
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn tools_lose_strict_and_closed_objects_and_custom_tools_become_functions() {
|
||||
let sent = send(json!({"messages": [], "tools": [
|
||||
{"type": "function", "function": {"name": "f", "strict": true,
|
||||
"parameters": {"type": "object", "additionalProperties": false, "properties": {}}}},
|
||||
{"type": "custom", "custom": {"name": "grep", "description": "search"}},
|
||||
]}));
|
||||
assert_eq!(
|
||||
sent["tools"],
|
||||
json!([
|
||||
{"type": "function", "function": {"name": "f", "parameters": {"type": "object", "properties": {}}}},
|
||||
{"type": "function", "function": {"name": "grep", "description": "search", "parameters": {
|
||||
"type": "object",
|
||||
"properties": {"input": {"type": "string", "description": "Raw tool input payload."}},
|
||||
"required": ["input"]}}},
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::budget(json!({"thinking": {"type": "enabled", "budget_tokens": 3000}}), Some("medium"))]
|
||||
#[case::effort_wins(json!({"thinking": {"type": "enabled", "budget_tokens": 3000}, "reasoning_effort": "low"}), Some("low"))]
|
||||
#[case::disabled(json!({"thinking": {"type": "disabled"}}), None)]
|
||||
fn thinking_becomes_a_reasoning_effort(#[case] params: Value, #[case] effort: Option<&str>) {
|
||||
let mut body = json!({"messages": []});
|
||||
body.as_object_mut()
|
||||
.unwrap()
|
||||
.extend(params.as_object().unwrap().clone());
|
||||
let sent = send(body);
|
||||
assert!(!sent.contains_key("thinking"));
|
||||
assert_eq!(sent.get("reasoning_effort").and_then(Value::as_str), effort);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn assistant_turns_are_flattened_for_vllm() {
|
||||
let sent = send(json!({"messages": [
|
||||
{"role": "assistant", "thinking_blocks": [{"type": "thinking"}], "reasoning_content": {"x": 1},
|
||||
"content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"},
|
||||
{"type": "tool_use", "id": "t1", "name": "f", "input": {"q": 1}}]},
|
||||
{"role": "user", "content": [{"type": "file", "file": {"file_data": "data:video/mp4;base64,AA"}}]},
|
||||
]}));
|
||||
assert_eq!(
|
||||
sent["messages"],
|
||||
json!([
|
||||
{"role": "assistant", "content": "a\nb", "tool_calls": [
|
||||
{"id": "t1", "type": "function", "function": {"name": "f", "arguments": "{\"q\": 1}"}}]},
|
||||
{"role": "user", "content": [{"type": "video_url", "video_url": {"url": "data:video/mp4;base64,AA"}}]},
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn the_key_defaults_for_a_server_that_takes_none() {
|
||||
assert_eq!(
|
||||
HOSTED_VLLM_CHAT_CONFIG.api_key(None, &|_| None).unwrap(),
|
||||
"fake-api-key"
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn an_api_base_is_required() {
|
||||
assert!(HOSTED_VLLM_CHAT_CONFIG.api_base(None, &|_| None).is_err());
|
||||
}
|
||||
}
|
||||
1
litellm-rust/crates/llms/src/hosted_vllm/mod.rs
Normal file
1
litellm-rust/crates/llms/src/hosted_vllm/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod chat;
|
||||
|
|
@ -5,17 +5,26 @@ pub mod azure_ai;
|
|||
pub mod base_llm;
|
||||
pub mod bedrock;
|
||||
pub mod bedrock_mantle;
|
||||
pub mod cerebras;
|
||||
pub mod cohere;
|
||||
pub mod deepseek;
|
||||
pub mod edenai;
|
||||
mod error;
|
||||
pub mod fireworks_ai;
|
||||
pub mod gemini;
|
||||
pub mod groq;
|
||||
pub mod hosted_vllm;
|
||||
pub mod minimax;
|
||||
pub mod mistral;
|
||||
pub mod nvidia_nim;
|
||||
pub mod openai;
|
||||
pub mod openai_like;
|
||||
pub mod openrouter;
|
||||
pub mod perplexity;
|
||||
pub mod reducto;
|
||||
pub mod tencent;
|
||||
pub mod together_ai;
|
||||
pub mod vertex_ai;
|
||||
pub mod xai;
|
||||
|
||||
pub use error::{Error, ErrorDetail, ErrorSource};
|
||||
|
|
|
|||
11
litellm-rust/crates/llms/src/mistral/chat/AGENTS.md
Normal file
11
litellm-rust/crates/llms/src/mistral/chat/AGENTS.md
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
# rules
|
||||
|
||||
- Mistral at `https://api.mistral.ai/v1`; a base without `/v1` gets it appended
|
||||
- `tool_choice` narrows to `auto`/`any`/`none`, `seed` goes as `random_seed`, magistral reasoning adds Mistral's system prompt
|
||||
- Content blocks in answers and chunks become text plus `reasoning_content`
|
||||
- Shared flow and response handling: `../../openai/chat/compatible.rs`
|
||||
|
||||
# references
|
||||
|
||||
- https://docs.mistral.ai/api/#tag/chat
|
||||
- litellm/llms/mistral/chat/transformation.py
|
||||
1
litellm-rust/crates/llms/src/mistral/chat/mod.rs
Normal file
1
litellm-rust/crates/llms/src/mistral/chat/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
605
litellm-rust/crates/llms/src/mistral/chat/transformation.rs
Normal file
605
litellm-rust/crates/llms/src/mistral/chat/transformation.rs
Normal file
|
|
@ -0,0 +1,605 @@
|
|||
//! `litellm/llms/mistral/chat/transformation.py`: Mistral at `https://api.mistral.ai/v1`.
|
||||
//! Tool choices narrow to Mistral's `auto`/`any`/`none`, `seed` goes as `random_seed`,
|
||||
//! magistral models get Mistral's reasoning system prompt, replayed messages lose what
|
||||
//! Mistral's schema forbids, and content lists in the answer become text plus
|
||||
//! `reasoning_content`.
|
||||
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::{
|
||||
Error,
|
||||
openai::chat::compatible::{
|
||||
CompatibleContext, OpenAiCompatibleChatConfig, ResponseShape, content_list_to_str,
|
||||
gpt_messages, messages_mut, nearest_declared_reasoning_effort, role, strip_internal_fields,
|
||||
},
|
||||
};
|
||||
|
||||
const REASONING_PROMPT_FLAG: &str = "_add_reasoning_prompt";
|
||||
/// `_get_mistral_reasoning_system_prompt`, verbatim with its indentation.
|
||||
const REASONING_SYSTEM_PROMPT: &str = r"
|
||||
<s>[SYSTEM_PROMPT]system_prompt
|
||||
A user will ask you to solve a task. You should first draft your thinking process (inner monologue) until you have derived the final answer. Afterwards, write a self-contained summary of your thoughts (i.e. your summary should be succinct but contain all the critical steps you needed to reach the conclusion). You should use Markdown to format your response. Write both your thoughts and summary in the same language as the task posed by the user. NEVER use \boxed{} in your response.
|
||||
|
||||
Your thinking process must follow the template below:
|
||||
<think>
|
||||
Your thoughts or/and draft, like working through an exercise on scratch paper. Be as casual and as long as you want until you are confident to generate a correct answer.
|
||||
</think>
|
||||
|
||||
Here, provide a concise summary that reflects your reasoning and presents a clear final answer to the user. Don't mention that this is a summary.
|
||||
|
||||
Problem:
|
||||
|
||||
[/SYSTEM_PROMPT][INST]user_message[/INST]<think>
|
||||
reasoning_traces
|
||||
</think>
|
||||
assistant_response</s>[INST]user_message[/INST]
|
||||
";
|
||||
/// The keys of `ChatCompletionAssistantMessage` besides `role`.
|
||||
const ASSISTANT_MESSAGE_KEYS: &[&str] = &[
|
||||
"content",
|
||||
"name",
|
||||
"tool_calls",
|
||||
"function_call",
|
||||
"reasoning_content",
|
||||
"cache_control",
|
||||
"thinking_blocks",
|
||||
"reasoning_items",
|
||||
];
|
||||
|
||||
pub struct MistralChatConfig;
|
||||
|
||||
pub const MISTRAL_CHAT_CONFIG: MistralChatConfig = MistralChatConfig;
|
||||
|
||||
impl OpenAiCompatibleChatConfig for MistralChatConfig {
|
||||
fn provider(&self) -> LlmProviders {
|
||||
LlmProviders::Mistral
|
||||
}
|
||||
|
||||
fn label(&self) -> &'static str {
|
||||
"Mistral"
|
||||
}
|
||||
|
||||
fn api_base_env(&self) -> &'static [&'static str] {
|
||||
&["MISTRAL_AZURE_API_BASE"]
|
||||
}
|
||||
|
||||
fn default_api_base(&self) -> Option<&'static str> {
|
||||
Some("https://api.mistral.ai/v1")
|
||||
}
|
||||
|
||||
fn api_key_env(&self) -> &'static [&'static str] {
|
||||
&["MISTRAL_AZURE_API_KEY", "MISTRAL_API_KEY"]
|
||||
}
|
||||
|
||||
/// Mistral always needs `/v1` at the end of the base.
|
||||
fn api_base(
|
||||
&self,
|
||||
configured: Option<&str>,
|
||||
lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<String, Error> {
|
||||
let base = configured
|
||||
.filter(|base| !base.is_empty())
|
||||
.map(str::to_owned)
|
||||
.or_else(|| lookup("MISTRAL_AZURE_API_BASE"))
|
||||
.unwrap_or_else(|| "https://api.mistral.ai/v1".to_owned());
|
||||
Ok(match base.ends_with("/v1") {
|
||||
true => base,
|
||||
false => format!("{base}/v1"),
|
||||
})
|
||||
}
|
||||
|
||||
fn supported_params(&self, context: &CompatibleContext<'_>) -> Vec<&'static str> {
|
||||
let magistral = is_magistral(context.model);
|
||||
let mut params = vec![
|
||||
"stream",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"tools",
|
||||
"tool_choice",
|
||||
"seed",
|
||||
"stop",
|
||||
"response_format",
|
||||
"parallel_tool_calls",
|
||||
];
|
||||
if magistral {
|
||||
params.push("thinking");
|
||||
}
|
||||
if magistral || context.supports(self.name(), "supports_reasoning") {
|
||||
params.push("reasoning_effort");
|
||||
}
|
||||
params
|
||||
}
|
||||
|
||||
fn map_params(
|
||||
&self,
|
||||
non_default: &Map<String, Value>,
|
||||
optional: &mut Map<String, Value>,
|
||||
context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
let magistral = is_magistral(context.model);
|
||||
for (name, value) in non_default {
|
||||
match name.as_str() {
|
||||
"max_tokens" if !non_default.contains_key("max_completion_tokens") => {
|
||||
optional.insert("max_tokens".into(), value.clone());
|
||||
}
|
||||
"max_completion_tokens" => {
|
||||
optional.insert("max_tokens".into(), value.clone());
|
||||
}
|
||||
"tools" => {
|
||||
let mut tools = value.clone();
|
||||
remove_json_schema_refs(&mut tools, 10);
|
||||
optional.insert("tools".into(), tools);
|
||||
}
|
||||
"stream" if value == &Value::Bool(true) => {
|
||||
optional.insert("stream".into(), value.clone());
|
||||
}
|
||||
"temperature" | "top_p" | "stop" | "response_format" | "parallel_tool_calls" => {
|
||||
optional.insert(name.clone(), value.clone());
|
||||
}
|
||||
"tool_choice" => {
|
||||
if let Some(choice) = value.as_str() {
|
||||
let choice = match choice {
|
||||
"auto" | "none" => choice,
|
||||
_ => "any",
|
||||
};
|
||||
optional.insert("tool_choice".into(), json!(choice));
|
||||
}
|
||||
}
|
||||
"seed" => {
|
||||
optional.insert("extra_body".into(), json!({"random_seed": value}));
|
||||
}
|
||||
"reasoning_effort" | "thinking" if magistral => {
|
||||
optional.insert(REASONING_PROMPT_FLAG.into(), Value::Bool(true));
|
||||
}
|
||||
"reasoning_effort" => {
|
||||
let effort = match (
|
||||
value.as_str(),
|
||||
context.declared_reasoning_efforts(self.name()),
|
||||
) {
|
||||
(Some(requested), Some(declared)) => {
|
||||
json!(nearest_declared_reasoning_effort(requested, &declared))
|
||||
}
|
||||
_ => value.clone(),
|
||||
};
|
||||
optional.insert("reasoning_effort".into(), effort);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn map_extra_body(
|
||||
&self,
|
||||
extra_body: &mut Map<String, Value>,
|
||||
_optional: &mut Map<String, Value>,
|
||||
_context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
extra_body.remove("client_metadata");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn transform_messages(
|
||||
&self,
|
||||
messages: Vec<Value>,
|
||||
optional: &Map<String, Value>,
|
||||
context: &CompatibleContext<'_>,
|
||||
) -> Result<Vec<Value>, Error> {
|
||||
let mut messages = messages;
|
||||
if is_magistral(context.model) && optional.contains_key(REASONING_PROMPT_FLAG) {
|
||||
add_reasoning_system_prompt(&mut messages);
|
||||
}
|
||||
for message in messages.iter_mut() {
|
||||
if role(message) == Some("assistant") {
|
||||
strip_internal_fields(message);
|
||||
}
|
||||
}
|
||||
let has_media = messages.iter().any(|message| {
|
||||
message
|
||||
.get("content")
|
||||
.and_then(Value::as_array)
|
||||
.is_some_and(|parts| {
|
||||
parts.iter().any(|part| {
|
||||
matches!(
|
||||
part.get("type").and_then(Value::as_str),
|
||||
Some("image_url" | "file")
|
||||
)
|
||||
})
|
||||
})
|
||||
});
|
||||
if has_media {
|
||||
let mut messages = gpt_messages(messages)?;
|
||||
messages.iter_mut().for_each(file_parts_to_file_ids);
|
||||
return Ok(messages);
|
||||
}
|
||||
let messages = messages
|
||||
.into_iter()
|
||||
.map(|mut message| {
|
||||
let text = content_list_to_str(&message);
|
||||
if let Some(object) = message.as_object_mut() {
|
||||
if !text.is_empty() {
|
||||
object.insert("content".into(), Value::String(text));
|
||||
}
|
||||
keep_name_on_tool_messages(object);
|
||||
rebuild_tool_calls(object);
|
||||
}
|
||||
message
|
||||
})
|
||||
.filter(|message| !is_empty_assistant_message(message))
|
||||
.map(|mut message| {
|
||||
if let Some(object) = message.as_object_mut() {
|
||||
object.retain(|_, value| !value.is_null());
|
||||
}
|
||||
message
|
||||
})
|
||||
.collect();
|
||||
gpt_messages(messages)
|
||||
}
|
||||
|
||||
fn finish_request(
|
||||
&self,
|
||||
body: &mut Map<String, Value>,
|
||||
_context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
body.remove(REASONING_PROMPT_FLAG);
|
||||
body.remove("client_metadata");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn fix_response(&self, body: &mut Map<String, Value>, _shape: &ResponseShape) -> bool {
|
||||
let mut changed = false;
|
||||
for message in messages_mut(body) {
|
||||
match message.get("content") {
|
||||
Some(Value::String(content)) if content.is_empty() => {
|
||||
message.insert("content".into(), Value::Null);
|
||||
changed = true;
|
||||
}
|
||||
Some(Value::Array(blocks)) if !blocks.is_empty() => {
|
||||
let (text, thinking) = split_content_blocks(blocks);
|
||||
message.insert("content".into(), Value::String(text));
|
||||
if !thinking.is_empty() {
|
||||
message.insert("reasoning_content".into(), Value::String(thinking));
|
||||
}
|
||||
changed = true;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
changed
|
||||
}
|
||||
|
||||
fn fix_chunk(&self, chunk: &mut Map<String, Value>, _shape: &ResponseShape) -> bool {
|
||||
let mut changed = false;
|
||||
let Some(Value::Array(choices)) = chunk.get_mut("choices") else {
|
||||
return false;
|
||||
};
|
||||
for delta in choices
|
||||
.iter_mut()
|
||||
.filter_map(|choice| choice.get_mut("delta")?.as_object_mut())
|
||||
{
|
||||
let Some(Value::Array(blocks)) = delta.get("content") else {
|
||||
continue;
|
||||
};
|
||||
let (text, thinking_blocks, reasoning) = normalize_stream_blocks(blocks);
|
||||
delta.insert("content".into(), text.map_or(Value::Null, Value::String));
|
||||
match reasoning {
|
||||
Some(reasoning) if !thinking_blocks.is_empty() => {
|
||||
delta.insert("thinking_blocks".into(), Value::Array(thinking_blocks));
|
||||
delta.insert("reasoning_content".into(), Value::String(reasoning));
|
||||
}
|
||||
_ => {
|
||||
delta.remove("thinking_blocks");
|
||||
delta.remove("reasoning_content");
|
||||
}
|
||||
}
|
||||
changed = true;
|
||||
}
|
||||
changed
|
||||
}
|
||||
}
|
||||
|
||||
fn is_magistral(model: &str) -> bool {
|
||||
model.to_lowercase().contains("magistral")
|
||||
}
|
||||
|
||||
/// `remove_json_schema_refs`.
|
||||
fn remove_json_schema_refs(value: &mut Value, depth: usize) {
|
||||
if depth == 0 {
|
||||
return;
|
||||
}
|
||||
match value {
|
||||
Value::Object(object) => {
|
||||
object.remove("$id");
|
||||
object.remove("$schema");
|
||||
object
|
||||
.values_mut()
|
||||
.for_each(|nested| remove_json_schema_refs(nested, depth - 1));
|
||||
}
|
||||
Value::Array(items) => items
|
||||
.iter_mut()
|
||||
.for_each(|item| remove_json_schema_refs(item, depth - 1)),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
/// `_add_reasoning_system_prompt_if_needed`.
|
||||
fn add_reasoning_system_prompt(messages: &mut Vec<Value>) {
|
||||
let Some(system) = messages
|
||||
.iter_mut()
|
||||
.find(|message| role(message) == Some("system"))
|
||||
.and_then(Value::as_object_mut)
|
||||
else {
|
||||
messages.insert(
|
||||
0,
|
||||
json!({"role": "system", "content": REASONING_SYSTEM_PROMPT}),
|
||||
);
|
||||
return;
|
||||
};
|
||||
let content = match system.remove("content") {
|
||||
Some(Value::String(existing)) => json!(format!("{REASONING_SYSTEM_PROMPT}\n\n{existing}")),
|
||||
Some(Value::Array(existing)) => Value::Array(
|
||||
std::iter::once(
|
||||
json!({"type": "text", "text": format!("{REASONING_SYSTEM_PROMPT}\n\n")}),
|
||||
)
|
||||
.chain(existing)
|
||||
.collect(),
|
||||
),
|
||||
Some(Value::Null) => json!(format!("{REASONING_SYSTEM_PROMPT}\n\nNone")),
|
||||
Some(other) => json!(format!("{REASONING_SYSTEM_PROMPT}\n\n{other}")),
|
||||
None => json!(format!("{REASONING_SYSTEM_PROMPT}\n\n")),
|
||||
};
|
||||
system.insert("content".into(), content);
|
||||
}
|
||||
|
||||
/// `_handle_message_with_file`: Mistral takes a file part by its `file_id`.
|
||||
fn file_parts_to_file_ids(message: &mut Value) {
|
||||
let Some(Value::Array(parts)) = message.get_mut("content") else {
|
||||
return;
|
||||
};
|
||||
for part in parts
|
||||
.iter_mut()
|
||||
.filter_map(Value::as_object_mut)
|
||||
.filter(|part| part.get("type").and_then(Value::as_str) == Some("file"))
|
||||
{
|
||||
let file_id = part
|
||||
.get("file")
|
||||
.and_then(|file| file.get("file_id"))
|
||||
.filter(|id| crate::base_llm::chat::params::truthy(id))
|
||||
.cloned();
|
||||
if let Some(file_id) = file_id {
|
||||
part.insert("file_id".into(), file_id);
|
||||
part.remove("file");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `_handle_name_in_message`: Mistral takes `name` only on tool messages.
|
||||
fn keep_name_on_tool_messages(message: &mut Map<String, Value>) {
|
||||
let Some(name) = message.get("name").filter(|name| !name.is_null()) else {
|
||||
return;
|
||||
};
|
||||
let blank = name.as_str().is_some_and(|name| name.trim().is_empty());
|
||||
if message.get("role").and_then(Value::as_str) != Some("tool") || blank {
|
||||
message.remove("name");
|
||||
}
|
||||
}
|
||||
|
||||
/// `_handle_tool_call_message`: each call as `{id, type: "function", function}`.
|
||||
fn rebuild_tool_calls(message: &mut Map<String, Value>) {
|
||||
let Some(Value::Array(calls)) = message.get_mut("tool_calls") else {
|
||||
return;
|
||||
};
|
||||
for call in calls.iter_mut() {
|
||||
*call = json!({
|
||||
"id": call.get("id").cloned().unwrap_or(Value::Null),
|
||||
"type": "function",
|
||||
"function": call.get("function").cloned().unwrap_or(Value::Null),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/// `_is_empty_assistant_message`.
|
||||
fn is_empty_assistant_message(message: &Value) -> bool {
|
||||
role(message) == Some("assistant")
|
||||
&& ASSISTANT_MESSAGE_KEYS
|
||||
.iter()
|
||||
.all(|key| match message.get(*key) {
|
||||
None | Some(Value::Null) => true,
|
||||
Some(Value::String(content)) => *key == "content" && content.is_empty(),
|
||||
Some(_) => false,
|
||||
})
|
||||
}
|
||||
|
||||
/// `_handle_content_list_to_str_conversion` for one message: the last text block, and the
|
||||
/// last thinking block's text parts joined by newlines.
|
||||
fn split_content_blocks(blocks: &[Value]) -> (String, String) {
|
||||
let mut text = String::new();
|
||||
let mut thinking = String::new();
|
||||
for block in blocks {
|
||||
match block.get("type").and_then(Value::as_str) {
|
||||
Some("thinking") => thinking = thinking_texts(block).join("\n"),
|
||||
Some("text") => {
|
||||
text = block
|
||||
.get("text")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_owned();
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
(text, thinking)
|
||||
}
|
||||
|
||||
fn thinking_texts(block: &Value) -> Vec<&str> {
|
||||
block
|
||||
.get("thinking")
|
||||
.and_then(Value::as_array)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter(|part| part.get("type").and_then(Value::as_str) == Some("text"))
|
||||
.map(|part| part.get("text").and_then(Value::as_str).unwrap_or_default())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// `MistralChatResponseIterator._normalize_content_blocks`.
|
||||
fn normalize_stream_blocks(blocks: &[Value]) -> (Option<String>, Vec<Value>, Option<String>) {
|
||||
let mut text: Option<String> = None;
|
||||
let mut thinking_blocks = Vec::new();
|
||||
let mut reasoning = Vec::new();
|
||||
for block in blocks {
|
||||
match block.get("type").and_then(Value::as_str) {
|
||||
Some("thinking") => {
|
||||
let thought = thinking_texts(block).concat();
|
||||
if !thought.is_empty() {
|
||||
thinking_blocks.push(
|
||||
json!({"type": "thinking", "thinking": thought, "signature": "mistral"}),
|
||||
);
|
||||
reasoning.push(thought);
|
||||
}
|
||||
}
|
||||
Some("text") => text.get_or_insert_with(String::new).push_str(
|
||||
block
|
||||
.get("text")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default(),
|
||||
),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
let reasoning = (!reasoning.is_empty()).then(|| reasoning.join("\n"));
|
||||
(text, thinking_blocks, reasoning)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::openai::chat::compatible::{fix_response, transform_request};
|
||||
|
||||
fn object(value: Value) -> Map<String, Value> {
|
||||
let Value::Object(map) = value else {
|
||||
panic!("an object")
|
||||
};
|
||||
map
|
||||
}
|
||||
|
||||
fn send(model: &str, params: Value) -> Map<String, Value> {
|
||||
transform_request(
|
||||
&MISTRAL_CHAT_CONFIG,
|
||||
&CompatibleContext {
|
||||
model,
|
||||
drop_params: false,
|
||||
model_info: None,
|
||||
},
|
||||
object(params),
|
||||
)
|
||||
.unwrap()
|
||||
.body
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::bare("https://example.ai", "https://example.ai/v1")]
|
||||
#[case::versioned("https://example.ai/v1", "https://example.ai/v1")]
|
||||
fn the_base_always_ends_in_v1(#[case] configured: &str, #[case] expected: &str) {
|
||||
assert_eq!(
|
||||
MISTRAL_CHAT_CONFIG
|
||||
.api_base(Some(configured), &|_| None)
|
||||
.unwrap(),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn params_take_mistral_names() {
|
||||
let sent = send(
|
||||
"mistral-large-latest",
|
||||
json!({"messages": [], "max_tokens": 1,
|
||||
"max_completion_tokens": 9, "seed": 3, "stream": false, "tool_choice": "required",
|
||||
"tools": [{"type": "function", "function": {"name": "f", "parameters": {"$schema": "s", "$id": "i", "type": "object"}}}]}),
|
||||
);
|
||||
assert_eq!(
|
||||
Value::Object(sent),
|
||||
json!({"model": "mistral-large-latest", "messages": [], "max_tokens": 9,
|
||||
"tools": [{"type": "function", "function": {"name": "f", "parameters": {"type": "object"}}}],
|
||||
"tool_choice": "any", "random_seed": 3})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn magistral_reasoning_adds_the_system_prompt_instead_of_the_param() {
|
||||
let sent = send(
|
||||
"magistral-medium-latest",
|
||||
json!({"reasoning_effort": "high",
|
||||
"messages": [{"role": "system", "content": "be brief"}]}),
|
||||
);
|
||||
assert_eq!(
|
||||
sent["messages"],
|
||||
json!([{"role": "system", "content": format!("{REASONING_SYSTEM_PROMPT}\n\nbe brief")}])
|
||||
);
|
||||
assert!(!sent.contains_key("reasoning_effort"));
|
||||
assert!(!sent.contains_key(REASONING_PROMPT_FLAG));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn replayed_messages_lose_what_mistral_forbids() {
|
||||
let sent = send(
|
||||
"mistral-small",
|
||||
json!({"messages": [
|
||||
{"role": "user", "name": "u", "content": [{"type": "text", "text": "hi"}]},
|
||||
{"role": "assistant", "content": "", "reasoning_content": "r", "thinking_blocks": []},
|
||||
{"role": "assistant", "content": null, "tool_calls": [{"index": 0, "id": "c", "type": "function", "function": {"name": "f", "arguments": "{}"}}]},
|
||||
{"role": "tool", "name": "f", "tool_call_id": "c", "content": "ok"}]}),
|
||||
);
|
||||
assert_eq!(
|
||||
sent["messages"],
|
||||
json!([
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "tool_calls": [{"id": "c", "type": "function", "function": {"name": "f", "arguments": "{}"}}]},
|
||||
{"role": "tool", "name": "f", "tool_call_id": "c", "content": "ok"}])
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn file_parts_go_by_file_id() {
|
||||
let sent = send(
|
||||
"mistral-small",
|
||||
json!({"messages": [{"role": "user", "content": [
|
||||
{"type": "text", "text": "read"}, {"type": "file", "file": {"file_id": "abc"}}]}]}),
|
||||
);
|
||||
assert_eq!(
|
||||
sent["messages"][0]["content"],
|
||||
json!([{"type": "text", "text": "read"}, {"type": "file", "file_id": "abc"}])
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty(json!(""), json!({"role": "assistant", "content": null}))]
|
||||
#[case::blocks(json!([{"type": "thinking", "thinking": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}]}, {"type": "text", "text": "done"}]),
|
||||
json!({"role": "assistant", "content": "done", "reasoning_content": "a\nb"}))]
|
||||
fn answers_become_text_and_reasoning(#[case] content: Value, #[case] message: Value) {
|
||||
let mut body =
|
||||
object(json!({"choices": [{"message": {"role": "assistant", "content": content}}]}));
|
||||
assert!(fix_response(
|
||||
&MISTRAL_CHAT_CONFIG,
|
||||
&mut body,
|
||||
&ResponseShape::default()
|
||||
));
|
||||
assert_eq!(body["choices"][0]["message"], message);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn streamed_blocks_become_text_and_thinking() {
|
||||
let mut chunk = object(json!({"choices": [{"index": 0, "delta": {"content": [
|
||||
{"type": "thinking", "thinking": [{"type": "text", "text": "hm"}]}, {"type": "text", "text": "ok"}]}}]}));
|
||||
assert!(MISTRAL_CHAT_CONFIG.fix_chunk(&mut chunk, &ResponseShape::default()));
|
||||
assert_eq!(
|
||||
chunk["choices"][0]["delta"],
|
||||
json!({"content": "ok", "thinking_blocks": [{"type": "thinking", "thinking": "hm", "signature": "mistral"}], "reasoning_content": "hm"})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1 +1,2 @@
|
|||
pub mod chat;
|
||||
pub mod ocr;
|
||||
|
|
|
|||
11
litellm-rust/crates/llms/src/nvidia_nim/chat/AGENTS.md
Normal file
11
litellm-rust/crates/llms/src/nvidia_nim/chat/AGENTS.md
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
# rules
|
||||
|
||||
- NIM's hosted API through the OpenAI SDK URL rule: `{api_base}/chat/completions`
|
||||
- The supported params depend on the model family; `max_completion_tokens` goes as `max_tokens`
|
||||
- No `stream_options`, so stream usage is only what NIM sends unasked
|
||||
- Shared flow and response handling: `../../openai/chat/compatible.rs`
|
||||
|
||||
# references
|
||||
|
||||
- https://docs.api.nvidia.com/nim/reference/llm-apis
|
||||
- litellm/llms/nvidia_nim/chat/transformation.py
|
||||
1
litellm-rust/crates/llms/src/nvidia_nim/chat/mod.rs
Normal file
1
litellm-rust/crates/llms/src/nvidia_nim/chat/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
196
litellm-rust/crates/llms/src/nvidia_nim/chat/transformation.rs
Normal file
196
litellm-rust/crates/llms/src/nvidia_nim/chat/transformation.rs
Normal file
|
|
@ -0,0 +1,196 @@
|
|||
//! `litellm/llms/nvidia_nim/chat/transformation.py`: NVIDIA NIM through the OpenAI SDK at
|
||||
//! `https://integrate.api.nvidia.com/v1`, with the params each model family takes and
|
||||
//! `max_completion_tokens` sent as `max_tokens`.
|
||||
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
use crate::{
|
||||
Error,
|
||||
openai::chat::compatible::{CompatibleContext, OpenAiCompatibleChatConfig, sdk_complete_url},
|
||||
};
|
||||
|
||||
const GEMMA_MODELS: &[&str] = &[
|
||||
"google/recurrentgemma-2b",
|
||||
"google/gemma-2-27b-it",
|
||||
"google/gemma-2-9b-it",
|
||||
"gemma-2-9b-it",
|
||||
];
|
||||
|
||||
pub struct NvidiaNimChatConfig;
|
||||
|
||||
pub const NVIDIA_NIM_CHAT_CONFIG: NvidiaNimChatConfig = NvidiaNimChatConfig;
|
||||
|
||||
impl OpenAiCompatibleChatConfig for NvidiaNimChatConfig {
|
||||
fn provider(&self) -> LlmProviders {
|
||||
LlmProviders::NvidiaNim
|
||||
}
|
||||
|
||||
fn label(&self) -> &'static str {
|
||||
"NVIDIA NIM"
|
||||
}
|
||||
|
||||
fn api_base_env(&self) -> &'static [&'static str] {
|
||||
&["NVIDIA_NIM_API_BASE"]
|
||||
}
|
||||
|
||||
fn default_api_base(&self) -> Option<&'static str> {
|
||||
Some("https://integrate.api.nvidia.com/v1")
|
||||
}
|
||||
|
||||
fn api_key_env(&self) -> &'static [&'static str] {
|
||||
&["NVIDIA_NIM_API_KEY"]
|
||||
}
|
||||
|
||||
fn complete_url(&self, api_base: &str) -> String {
|
||||
sdk_complete_url(api_base)
|
||||
}
|
||||
|
||||
fn supported_params(&self, context: &CompatibleContext<'_>) -> Vec<&'static str> {
|
||||
let model = context.model;
|
||||
if GEMMA_MODELS.contains(&model) {
|
||||
return vec![
|
||||
"stream",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"max_tokens",
|
||||
"stop",
|
||||
"seed",
|
||||
];
|
||||
}
|
||||
match model {
|
||||
"nvidia/nemotron-4-340b-instruct" => vec![
|
||||
"stream",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
],
|
||||
"nvidia/nemotron-4-340b-reward" => vec!["stream"],
|
||||
"google/codegemma-1.1-7b" => vec![
|
||||
"stream",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"frequency_penalty",
|
||||
"presence_penalty",
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"stop",
|
||||
],
|
||||
_ => vec![
|
||||
"stream",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"frequency_penalty",
|
||||
"presence_penalty",
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"stop",
|
||||
"seed",
|
||||
"tools",
|
||||
"tool_choice",
|
||||
"parallel_tool_calls",
|
||||
"response_format",
|
||||
],
|
||||
}
|
||||
}
|
||||
|
||||
fn map_params(
|
||||
&self,
|
||||
non_default: &Map<String, Value>,
|
||||
optional: &mut Map<String, Value>,
|
||||
context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
let supported = self.supported_params(context);
|
||||
for (name, value) in non_default {
|
||||
if name == "max_completion_tokens" {
|
||||
optional.insert("max_tokens".into(), value.clone());
|
||||
} else if supported.contains(&name.as_str()) {
|
||||
optional.insert(name.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::{ErrorDetail, openai::chat::compatible::transform_request};
|
||||
|
||||
fn context(model: &str, drop_params: bool) -> CompatibleContext<'_> {
|
||||
CompatibleContext {
|
||||
model,
|
||||
drop_params,
|
||||
model_info: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn body(value: Value) -> Map<String, Value> {
|
||||
let Value::Object(map) = value else {
|
||||
panic!("an object")
|
||||
};
|
||||
map
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn max_completion_tokens_goes_as_max_tokens() {
|
||||
let sent = transform_request(
|
||||
&NVIDIA_NIM_CHAT_CONFIG,
|
||||
&context("meta/llama3-70b-instruct", false),
|
||||
body(
|
||||
json!({"model": "x", "messages": [{"role": "user", "content": "hi"}],
|
||||
"max_completion_tokens": 9, "seed": 1, "top_k": 4}),
|
||||
),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
Value::Object(sent.body),
|
||||
json!({"model": "meta/llama3-70b-instruct", "messages": [{"role": "user", "content": "hi"}],
|
||||
"max_tokens": 9, "seed": 1, "top_k": 4})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::refused(false, true)]
|
||||
#[case::dropped(true, false)]
|
||||
fn a_param_the_model_family_lacks_is_refused_or_dropped(
|
||||
#[case] drop_params: bool,
|
||||
#[case] refused: bool,
|
||||
) {
|
||||
let result = transform_request(
|
||||
&NVIDIA_NIM_CHAT_CONFIG,
|
||||
&context("google/gemma-2-9b-it", drop_params),
|
||||
body(
|
||||
json!({"messages": [], "tools": [{"type": "function", "function": {"name": "f"}}]}),
|
||||
),
|
||||
);
|
||||
match refused {
|
||||
true => assert!(matches!(
|
||||
result,
|
||||
Err(Error::InvalidRequest(ErrorDetail::UnsupportedParams {
|
||||
provider: "nvidia_nim",
|
||||
..
|
||||
}))
|
||||
)),
|
||||
false => assert!(!result.unwrap().body.contains_key("tools")),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn stream_options_are_not_sent_since_nim_does_not_list_them() {
|
||||
let sent = transform_request(
|
||||
&NVIDIA_NIM_CHAT_CONFIG,
|
||||
&context("meta/llama3-70b-instruct", false),
|
||||
body(
|
||||
json!({"messages": [], "stream": true, "stream_options": {"include_usage": true}}),
|
||||
),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(!sent.body.contains_key("stream_options"));
|
||||
assert!(sent.shape.client_usage);
|
||||
}
|
||||
}
|
||||
1
litellm-rust/crates/llms/src/nvidia_nim/mod.rs
Normal file
1
litellm-rust/crates/llms/src/nvidia_nim/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod chat;
|
||||
|
|
@ -4,3 +4,4 @@
|
|||
- Passthrough, not a `BaseConfig`: that contract parses messages and normalizes the response into `ChatCompletionsResponse`, which drops tool calls, ids and every field it does not type. OpenAI's own body and SSE bytes reach the caller as sent
|
||||
- The request changes only by the deployment model name, the `extra_body` merge the OpenAI SDK does, and dropping the LiteLLM-owned keys Python consumes before the call
|
||||
- Endpoint and credential fallbacks are shared with Responses in `../common_utils.rs`
|
||||
- `compatible.rs` is the request flow and response fixups the OpenAI-compatible providers share. Each provider's config lives in `src/<provider>/chat/`, and its body and SSE bytes change only where its config's fixups change them
|
||||
|
|
|
|||
985
litellm-rust/crates/llms/src/openai/chat/compatible.rs
Normal file
985
litellm-rust/crates/llms/src/openai/chat/compatible.rs
Normal file
|
|
@ -0,0 +1,985 @@
|
|||
//! The chat flow every OpenAI-compatible provider shares, `litellm/llms/openai/chat/gpt_transformation.py`
|
||||
//! (`OpenAIGPTConfig`) as the base each provider config overrides, run over the caller's
|
||||
//! OpenAI-shaped body: Python's `get_optional_params` sorts the request's OpenAI params from the
|
||||
//! rest, refuses or drops what the provider does not take, maps the rest by the provider's
|
||||
//! rules, and passes every other key through as `extra_body` does. The messages go through the
|
||||
//! provider's `_transform_messages`, and the response reaches the caller as the provider sent
|
||||
//! it, changed only by the fixups each provider's `transform_response` and stream iterator make.
|
||||
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::{
|
||||
Error, ErrorDetail,
|
||||
base_llm::{
|
||||
chat::params::{ChatRequestContext, sort_params, truthy},
|
||||
model_info::ModelInfo,
|
||||
},
|
||||
openai::chat::transformation::{OPENAI_CHAT_COMPLETIONS_PATH, is_stream, take_extra_body},
|
||||
};
|
||||
|
||||
/// The `base_params` of `OpenAIGPTConfig.get_supported_openai_params`.
|
||||
pub const GPT_BASE_PARAMS: &[&str] = &[
|
||||
"frequency_penalty",
|
||||
"logit_bias",
|
||||
"logprobs",
|
||||
"top_logprobs",
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"modalities",
|
||||
"prediction",
|
||||
"n",
|
||||
"presence_penalty",
|
||||
"seed",
|
||||
"stop",
|
||||
"stream",
|
||||
"stream_options",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"tools",
|
||||
"tool_choice",
|
||||
"function_call",
|
||||
"functions",
|
||||
"max_retries",
|
||||
"extra_headers",
|
||||
"parallel_tool_calls",
|
||||
"audio",
|
||||
"web_search_options",
|
||||
"service_tier",
|
||||
"safety_identifier",
|
||||
"prompt_cache_key",
|
||||
"prompt_cache_retention",
|
||||
"store",
|
||||
];
|
||||
|
||||
/// `LITELLM_RESPONSE_ONLY_MESSAGE_FIELDS`: what LiteLLM adds to a returned message that no
|
||||
/// OpenAI-compatible endpoint takes back.
|
||||
const RESPONSE_ONLY_MESSAGE_FIELDS: &[&str] = &[
|
||||
"thinking_blocks",
|
||||
"reasoning_items",
|
||||
"provider_specific_fields",
|
||||
];
|
||||
/// `LITELLM_INTERNAL_MESSAGE_FIELDS`: the response-only fields and `reasoning_content`.
|
||||
const INTERNAL_MESSAGE_FIELDS: &[&str] = &[
|
||||
"thinking_blocks",
|
||||
"reasoning_items",
|
||||
"provider_specific_fields",
|
||||
"reasoning_content",
|
||||
];
|
||||
/// `DEFAULT_MAX_RECURSE_DEPTH`.
|
||||
const MAX_RECURSE_DEPTH: usize = 100;
|
||||
const STREAM_OPTIONS: &str = "stream_options";
|
||||
const INCLUDE_USAGE: &str = "include_usage";
|
||||
/// `RESPONSE_FORMAT_TOOL_NAME`.
|
||||
pub const RESPONSE_FORMAT_TOOL_NAME: &str = "json_tool_call";
|
||||
const POTENTIAL_PDF_URL_STARTS: &[&str] = &["https://", "http://", "www."];
|
||||
/// The content part types `_content_parts_contain_image` looks for.
|
||||
const IMAGE_PART_TYPES: &[&str] = &["image_url", "input_image", "image"];
|
||||
/// `DEFAULT_REASONING_EFFORT_{LOW,MEDIUM,HIGH}_THINKING_BUDGET` at their defaults.
|
||||
const THINKING_BUDGETS: (u64, u64, u64) = (1024, 2048, 4096);
|
||||
/// `REASONING_EFFORT_STRENGTH_ORDER`.
|
||||
const EFFORT_STRENGTH_ORDER: &[&str] = &["minimal", "low", "medium", "high", "xhigh", "max"];
|
||||
|
||||
/// What a provider's param rules read about the call.
|
||||
#[derive(Clone, Copy)]
|
||||
pub struct CompatibleContext<'a> {
|
||||
/// The model as the provider names it, without the route prefix.
|
||||
pub model: &'a str,
|
||||
/// `drop_params` from the deployment, the proxy settings or the request.
|
||||
pub drop_params: bool,
|
||||
pub model_info: Option<&'a dyn ModelInfo>,
|
||||
}
|
||||
|
||||
impl CompatibleContext<'_> {
|
||||
/// `supports_factory`: the `<provider>/<model>` entry's flag, else the bare model's.
|
||||
pub fn capability(&self, provider: &str, key: &str) -> Option<bool> {
|
||||
self.capability_of(
|
||||
&[format!("{provider}/{}", self.model), self.model.to_string()],
|
||||
key,
|
||||
)
|
||||
}
|
||||
|
||||
/// The first of `candidates` whose cost map entry sets `key`.
|
||||
pub fn capability_of(&self, candidates: &[String], key: &str) -> Option<bool> {
|
||||
let info = self.model_info?;
|
||||
candidates
|
||||
.iter()
|
||||
.filter_map(|candidate| info.entry(candidate))
|
||||
.find_map(|entry| entry.get(key).and_then(Value::as_bool))
|
||||
}
|
||||
|
||||
pub fn supports(&self, provider: &str, key: &str) -> bool {
|
||||
self.capability(provider, key) == Some(true)
|
||||
}
|
||||
|
||||
/// `declared_reasoning_efforts_for_model`: the `reasoning_effort_levels` the entry for the
|
||||
/// model declares, read through its bare twin. `None` when the map says nothing.
|
||||
pub fn declared_reasoning_efforts(&self, provider: &str) -> Option<Vec<String>> {
|
||||
let info = self.model_info?;
|
||||
let prefixed = format!("{provider}/{}", self.model);
|
||||
let (entry, prefixed_hit) = match info.entry(&prefixed) {
|
||||
Some(entry) => (entry, true),
|
||||
None => (info.entry(self.model)?, false),
|
||||
};
|
||||
let levels = match entry.get("reasoning_effort_levels") {
|
||||
Some(levels) if !levels.is_null() => levels.clone(),
|
||||
_ if prefixed_hit => info
|
||||
.entry(self.model)
|
||||
.and_then(|bare| bare.get("reasoning_effort_levels").cloned())
|
||||
.unwrap_or(Value::Null),
|
||||
_ => Value::Null,
|
||||
};
|
||||
let Value::Array(levels) = levels else {
|
||||
return None;
|
||||
};
|
||||
Some(
|
||||
levels
|
||||
.into_iter()
|
||||
.filter_map(|level| level.as_str().map(str::to_owned))
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// `nearest_declared_reasoning_effort`: the weakest declared level at least as strong as the
|
||||
/// request, else the strongest declared one. `none` and unknown levels go as asked.
|
||||
pub fn nearest_declared_reasoning_effort(requested: &str, declared: &[String]) -> String {
|
||||
let rank = |effort: &str| EFFORT_STRENGTH_ORDER.iter().position(|e| *e == effort);
|
||||
let mut ranked: Vec<&str> = declared
|
||||
.iter()
|
||||
.map(String::as_str)
|
||||
.filter(|effort| rank(effort).is_some())
|
||||
.collect();
|
||||
ranked.sort_by_key(|effort| rank(effort));
|
||||
let Some(requested_rank) = rank(requested) else {
|
||||
return requested.to_owned();
|
||||
};
|
||||
if ranked.is_empty() || ranked.contains(&requested) {
|
||||
return requested.to_owned();
|
||||
}
|
||||
ranked
|
||||
.iter()
|
||||
.find(|effort| rank(effort) >= Some(requested_rank))
|
||||
.or(ranked.last())
|
||||
.map_or_else(|| requested.to_owned(), |effort| (*effort).to_owned())
|
||||
}
|
||||
|
||||
/// `reasoning_effort_from_thinking_budget` at the default budget thresholds.
|
||||
pub fn reasoning_effort_from_thinking_budget(budget_tokens: i64) -> &'static str {
|
||||
let (low, medium, high) = THINKING_BUDGETS;
|
||||
match budget_tokens {
|
||||
tokens if tokens >= high as i64 => "high",
|
||||
tokens if tokens >= medium as i64 => "medium",
|
||||
tokens if tokens >= low as i64 => "low",
|
||||
_ => "minimal",
|
||||
}
|
||||
}
|
||||
|
||||
/// What a request decided about its answer, which the response and stream fixups read.
|
||||
#[derive(Clone, Debug, Default, PartialEq, Eq)]
|
||||
pub struct ResponseShape {
|
||||
pub stream: bool,
|
||||
/// The client asked to see the usage on a stream. When it did not, the usage the call asked
|
||||
/// for on its behalf is recorded and kept from the client.
|
||||
pub client_usage: bool,
|
||||
/// The provider answers a `response_format` with a `json_tool_call` tool call, which goes
|
||||
/// back as content.
|
||||
pub json_mode: bool,
|
||||
/// The names of the function tools the request sent.
|
||||
pub tool_names: Vec<String>,
|
||||
/// Whether the request sent `tools` at all.
|
||||
pub tools_sent: bool,
|
||||
}
|
||||
|
||||
/// A request ready for an OpenAI-compatible provider.
|
||||
#[derive(Debug, PartialEq)]
|
||||
pub struct CompatibleChatBody {
|
||||
pub body: Map<String, Value>,
|
||||
pub shape: ResponseShape,
|
||||
}
|
||||
|
||||
/// One OpenAI-compatible provider's config: its endpoint and credential fallbacks, and the
|
||||
/// overrides it makes to `OpenAIGPTConfig`.
|
||||
pub trait OpenAiCompatibleChatConfig: Sync {
|
||||
fn provider(&self) -> LlmProviders;
|
||||
|
||||
/// The provider as a missing key error names it.
|
||||
fn label(&self) -> &'static str;
|
||||
|
||||
/// The variables an unset `api_base` falls back to, in order.
|
||||
fn api_base_env(&self) -> &'static [&'static str];
|
||||
|
||||
fn default_api_base(&self) -> Option<&'static str>;
|
||||
|
||||
/// The variables an unset key falls back to, in order.
|
||||
fn api_key_env(&self) -> &'static [&'static str];
|
||||
|
||||
/// The key sent when neither the deployment nor the environment has one; `None` means the
|
||||
/// call needs one.
|
||||
fn default_api_key(&self) -> Option<&'static str> {
|
||||
None
|
||||
}
|
||||
|
||||
/// Other variables the provider's headers read.
|
||||
fn header_env(&self) -> &'static [&'static str] {
|
||||
&[]
|
||||
}
|
||||
|
||||
fn name(&self) -> &'static str {
|
||||
self.provider().into()
|
||||
}
|
||||
|
||||
/// Every variable a call may read.
|
||||
fn secret_names(&self) -> Vec<&'static str> {
|
||||
self.api_base_env()
|
||||
.iter()
|
||||
.chain(self.api_key_env())
|
||||
.chain(self.header_env())
|
||||
.copied()
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// The deployment's `api_base`, else the environment's, else the provider's own.
|
||||
fn api_base(
|
||||
&self,
|
||||
configured: Option<&str>,
|
||||
lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<String, Error> {
|
||||
configured
|
||||
.filter(|base| !base.is_empty())
|
||||
.map(str::to_owned)
|
||||
.or_else(|| self.api_base_env().iter().find_map(|name| lookup(name)))
|
||||
.or_else(|| self.default_api_base().map(str::to_owned))
|
||||
.ok_or_else(|| {
|
||||
Error::InvalidRequest(ErrorDetail::Message(format!(
|
||||
"Missing API Base - set api_base on the {} deployment or {}",
|
||||
self.name(),
|
||||
self.api_base_env().join(" or ")
|
||||
)))
|
||||
})
|
||||
}
|
||||
|
||||
/// `OpenAIGPTConfig.get_complete_url`: `/chat/completions` appended unless the base already
|
||||
/// names it.
|
||||
fn complete_url(&self, api_base: &str) -> String {
|
||||
gpt_complete_url(api_base)
|
||||
}
|
||||
|
||||
/// The deployment's key, else the environment's, else the provider's default.
|
||||
fn api_key(
|
||||
&self,
|
||||
configured: Option<&str>,
|
||||
lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Result<String, Error> {
|
||||
configured
|
||||
.filter(|key| !key.is_empty())
|
||||
.map(str::to_owned)
|
||||
.or_else(|| {
|
||||
self.api_key_env()
|
||||
.iter()
|
||||
.find_map(|name| lookup(name).filter(|key| !key.is_empty()))
|
||||
})
|
||||
.or_else(|| self.default_api_key().map(str::to_owned))
|
||||
.ok_or(Error::Auth(litellm_auth::Error::MissingApiKey {
|
||||
provider: self.label(),
|
||||
environment_variable: self.api_key_env().first().copied().unwrap_or_default(),
|
||||
}))
|
||||
}
|
||||
|
||||
/// Headers the provider adds to every call, read off the caller's body before it is
|
||||
/// sorted.
|
||||
fn headers(
|
||||
&self,
|
||||
_body: &Map<String, Value>,
|
||||
_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Vec<(String, String)> {
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
/// The model the body names.
|
||||
fn wire_model(&self, model: &str) -> String {
|
||||
model.to_owned()
|
||||
}
|
||||
|
||||
/// `get_supported_openai_params`.
|
||||
fn supported_params(&self, context: &CompatibleContext<'_>) -> Vec<&'static str>;
|
||||
|
||||
/// `map_openai_params`: `non_default` holds only supported params.
|
||||
fn map_params(
|
||||
&self,
|
||||
non_default: &Map<String, Value>,
|
||||
optional: &mut Map<String, Value>,
|
||||
context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
copy_supported(non_default, optional, &self.supported_params(context));
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// The keys sent as `extra_body`, which a provider may rewrite.
|
||||
fn map_extra_body(
|
||||
&self,
|
||||
_extra_body: &mut Map<String, Value>,
|
||||
_optional: &mut Map<String, Value>,
|
||||
_context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `_transform_messages`.
|
||||
fn transform_messages(
|
||||
&self,
|
||||
messages: Vec<Value>,
|
||||
_optional: &Map<String, Value>,
|
||||
_context: &CompatibleContext<'_>,
|
||||
) -> Result<Vec<Value>, Error> {
|
||||
gpt_messages(messages)
|
||||
}
|
||||
|
||||
/// Whether `cache_control` marks reach the provider; `OpenAIGPTConfig` strips them.
|
||||
fn keeps_cache_control(&self, _context: &CompatibleContext<'_>) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// The provider's `transform_request` changes over the assembled body, before the
|
||||
/// `extra_body` keys are merged over it.
|
||||
fn finish_request(
|
||||
&self,
|
||||
_body: &mut Map<String, Value>,
|
||||
_context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `transform_response`'s fixups to a complete response. Returns whether it changed.
|
||||
fn fix_response(&self, _body: &mut Map<String, Value>, _shape: &ResponseShape) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// The stream iterator's fixups to one chunk. Returns whether it changed.
|
||||
fn fix_chunk(&self, chunk: &mut Map<String, Value>, _shape: &ResponseShape) -> bool {
|
||||
rename_reasoning_in_deltas(chunk)
|
||||
}
|
||||
}
|
||||
|
||||
/// `OpenAIGPTConfig.get_complete_url`.
|
||||
pub fn gpt_complete_url(api_base: &str) -> String {
|
||||
let base = api_base.trim_end_matches('/');
|
||||
match base.contains("chat/completions") {
|
||||
true => base.to_owned(),
|
||||
false => format!("{base}{OPENAI_CHAT_COMPLETIONS_PATH}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// The URL the OpenAI SDK calls for a `base_url`.
|
||||
pub fn sdk_complete_url(api_base: &str) -> String {
|
||||
format!(
|
||||
"{}{OPENAI_CHAT_COMPLETIONS_PATH}",
|
||||
api_base.trim_end_matches('/')
|
||||
)
|
||||
}
|
||||
|
||||
/// `OpenAIGPTConfig._map_openai_params`: every supported param as sent.
|
||||
pub fn copy_supported(
|
||||
non_default: &Map<String, Value>,
|
||||
optional: &mut Map<String, Value>,
|
||||
supported: &[&str],
|
||||
) {
|
||||
for (name, value) in non_default {
|
||||
if supported.contains(&name.as_str()) {
|
||||
optional.insert(name.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The request for `config`, built from the caller's body.
|
||||
pub fn transform_request(
|
||||
config: &dyn OpenAiCompatibleChatConfig,
|
||||
context: &CompatibleContext<'_>,
|
||||
body: Map<String, Value>,
|
||||
) -> Result<CompatibleChatBody, Error> {
|
||||
let mut body = body;
|
||||
body.remove("model");
|
||||
let messages = match body.remove("messages") {
|
||||
None | Some(Value::Null) => Vec::new(),
|
||||
Some(Value::Array(messages)) => messages,
|
||||
Some(other) => {
|
||||
return Err(Error::InvalidRequest(ErrorDetail::InvalidValue {
|
||||
field: "messages",
|
||||
expected: "an array",
|
||||
actual: other,
|
||||
}));
|
||||
}
|
||||
};
|
||||
let mut user_extra_body = take_extra_body(&mut body)?;
|
||||
user_extra_body.remove("model");
|
||||
let stream = is_stream(&body)?;
|
||||
let client_usage = stream
|
||||
&& body
|
||||
.get(STREAM_OPTIONS)
|
||||
.and_then(|options| options.get(INCLUDE_USAGE))
|
||||
== Some(&Value::Bool(true));
|
||||
let supported = config.supported_params(context);
|
||||
let sorted = sort_params(
|
||||
body,
|
||||
&[],
|
||||
&ChatRequestContext {
|
||||
drop_params: Some(context.drop_params),
|
||||
model_info: context.model_info,
|
||||
},
|
||||
|_| Ok(supported.clone()),
|
||||
config.name(),
|
||||
context.model,
|
||||
)?;
|
||||
let mut non_default = sorted.openai;
|
||||
// The proxy's stream usage tracking: a stream the provider can report usage on asks for
|
||||
// it, so the call is accounted for, and the client sees it only if it asked.
|
||||
if stream && !client_usage && supported.contains(&STREAM_OPTIONS) {
|
||||
let mut options = match non_default.remove(STREAM_OPTIONS) {
|
||||
Some(Value::Object(options)) => options,
|
||||
_ => Map::new(),
|
||||
};
|
||||
options.insert(INCLUDE_USAGE.into(), Value::Bool(true));
|
||||
non_default.insert(STREAM_OPTIONS.into(), Value::Object(options));
|
||||
}
|
||||
|
||||
let mut optional = Map::new();
|
||||
config.map_params(&non_default, &mut optional, context)?;
|
||||
// `apply_openai_param_overrides`: what the caller allowed goes as sent.
|
||||
for name in &sorted.allowed {
|
||||
if !optional.contains_key(name)
|
||||
&& let Some(value) = non_default.get(name)
|
||||
{
|
||||
optional.insert(name.clone(), value.clone());
|
||||
}
|
||||
}
|
||||
let json_mode = optional
|
||||
.remove("json_mode")
|
||||
.is_some_and(|value| value == Value::Bool(true));
|
||||
let mut extra_body = match optional.remove("extra_body") {
|
||||
Some(Value::Object(extra)) => extra,
|
||||
_ => Map::new(),
|
||||
};
|
||||
extra_body.extend(user_extra_body);
|
||||
extra_body.extend(sorted.provider);
|
||||
config.map_extra_body(&mut extra_body, &mut optional, context)?;
|
||||
|
||||
let mut messages = config.transform_messages(messages, &optional, context)?;
|
||||
if !config.keeps_cache_control(context) {
|
||||
for message in &mut messages {
|
||||
filter_value(message, "cache_control", 0);
|
||||
}
|
||||
if let Some(Value::Array(tools)) = optional.get_mut("tools") {
|
||||
for tool in tools {
|
||||
filter_value(tool, "cache_control", 0);
|
||||
}
|
||||
}
|
||||
}
|
||||
if context.drop_params {
|
||||
for message in &mut messages {
|
||||
without_prompt_cache_breakpoints(message);
|
||||
}
|
||||
if let Some(Value::Array(tools)) = optional.get_mut("tools") {
|
||||
for tool in tools {
|
||||
without_prompt_cache_breakpoints(tool);
|
||||
}
|
||||
}
|
||||
}
|
||||
optional.remove("max_retries");
|
||||
if !optional.get("tools").is_some_and(truthy) && !optional.get("functions").is_some_and(truthy)
|
||||
{
|
||||
optional.remove("tool_choice");
|
||||
}
|
||||
|
||||
let mut sent = Map::new();
|
||||
sent.insert(
|
||||
"model".into(),
|
||||
Value::String(config.wire_model(context.model)),
|
||||
);
|
||||
sent.insert("messages".into(), Value::Array(messages));
|
||||
sent.extend(optional);
|
||||
config.finish_request(&mut sent, context)?;
|
||||
extra_body.remove("model");
|
||||
sent.extend(extra_body);
|
||||
|
||||
let tools = sent.get("tools").and_then(Value::as_array);
|
||||
let shape = ResponseShape {
|
||||
stream,
|
||||
client_usage,
|
||||
json_mode,
|
||||
tool_names: tools
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(|tool| tool.get("function")?.get("name")?.as_str())
|
||||
.map(str::to_owned)
|
||||
.collect(),
|
||||
tools_sent: sent.get("tools").is_some_and(|tools| !tools.is_null()),
|
||||
};
|
||||
Ok(CompatibleChatBody { body: sent, shape })
|
||||
}
|
||||
|
||||
/// A complete response as the caller receives it. Returns whether it changed.
|
||||
pub fn fix_response(
|
||||
config: &dyn OpenAiCompatibleChatConfig,
|
||||
body: &mut Map<String, Value>,
|
||||
shape: &ResponseShape,
|
||||
) -> bool {
|
||||
let renamed = rename_reasoning_in_messages(body);
|
||||
config.fix_response(body, shape) || renamed
|
||||
}
|
||||
|
||||
/// `OpenAIGPTConfig._transform_messages`: `tool_reference` parts dropped from tool messages,
|
||||
/// and each user content part normalized.
|
||||
pub fn gpt_messages(messages: Vec<Value>) -> Result<Vec<Value>, Error> {
|
||||
if messages.iter().any(tool_message_carries_image) {
|
||||
// `hoist_images_from_tool_messages` moves them into an inserted user message.
|
||||
return Err(Error::Unsupported(
|
||||
"images in tool messages, which Python moves into a user message",
|
||||
));
|
||||
}
|
||||
messages
|
||||
.into_iter()
|
||||
.map(drop_tool_reference_parts)
|
||||
.map(|mut message| {
|
||||
if role(&message) == Some("user")
|
||||
&& let Some(Value::Array(parts)) = message.get_mut("content")
|
||||
{
|
||||
for part in parts.iter_mut() {
|
||||
transform_content_item(part)?;
|
||||
}
|
||||
}
|
||||
Ok(message)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn role(message: &Value) -> Option<&str> {
|
||||
message.get("role").and_then(Value::as_str)
|
||||
}
|
||||
|
||||
fn is_tool_reference(part: &Value) -> bool {
|
||||
part.get("type").and_then(Value::as_str) == Some("tool_reference")
|
||||
}
|
||||
|
||||
/// `drop_tool_reference_parts_from_tool_messages` for one message.
|
||||
fn drop_tool_reference_parts(mut message: Value) -> Value {
|
||||
if role(&message) != Some("tool") {
|
||||
return message;
|
||||
}
|
||||
if let Some(content) = message.get_mut("content")
|
||||
&& let Value::Array(parts) = content
|
||||
&& parts.iter().any(is_tool_reference)
|
||||
{
|
||||
parts.retain(|part| !is_tool_reference(part));
|
||||
if parts.is_empty() {
|
||||
*content = Value::String(String::new());
|
||||
}
|
||||
}
|
||||
message
|
||||
}
|
||||
|
||||
fn tool_message_carries_image(message: &Value) -> bool {
|
||||
fn has_image(parts: &[Value], depth: usize) -> bool {
|
||||
parts.iter().any(|part| {
|
||||
part.get("type")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|kind| IMAGE_PART_TYPES.contains(&kind))
|
||||
|| (depth < 4
|
||||
&& part
|
||||
.get("content")
|
||||
.and_then(Value::as_array)
|
||||
.is_some_and(|nested| has_image(nested, depth + 1)))
|
||||
})
|
||||
}
|
||||
role(message) == Some("tool")
|
||||
&& message
|
||||
.get("content")
|
||||
.and_then(Value::as_array)
|
||||
.is_some_and(|parts| has_image(parts, 1))
|
||||
}
|
||||
|
||||
/// `OpenAIGPTConfig._transform_content_item`.
|
||||
fn transform_content_item(part: &mut Value) -> Result<(), Error> {
|
||||
let Some(object) = part.as_object_mut() else {
|
||||
return Ok(());
|
||||
};
|
||||
match object.get("type").and_then(Value::as_str) {
|
||||
Some("image_url") => match object.get_mut("image_url") {
|
||||
Some(Value::String(url)) => {
|
||||
let url = std::mem::take(url);
|
||||
object.insert("image_url".into(), json!({"url": url}));
|
||||
}
|
||||
Some(Value::Object(image)) => {
|
||||
image.remove("format");
|
||||
}
|
||||
_ => {}
|
||||
},
|
||||
Some("file") => {
|
||||
let Some(Value::Object(file)) = object.get_mut("file") else {
|
||||
return Err(Error::InvalidRequest(ErrorDetail::Message(
|
||||
"Content block has type='file' but is missing the required 'file' field".into(),
|
||||
)));
|
||||
};
|
||||
file.remove("format");
|
||||
if file
|
||||
.get("file_id")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|id| POTENTIAL_PDF_URL_STARTS.iter().any(|s| id.starts_with(s)))
|
||||
{
|
||||
// `_handle_pdf_url` downloads it into `file_data`.
|
||||
return Err(Error::Unsupported(
|
||||
"a file_id URL Python downloads into file_data",
|
||||
));
|
||||
}
|
||||
if file.get("file_data").is_some_and(|data| !data.is_null())
|
||||
&& file.get("filename").is_none_or(Value::is_null)
|
||||
{
|
||||
file.insert("filename".into(), Value::String("my_file.pdf".into()));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `filter_value_from_dict`: `key` removed from every object, at any depth.
|
||||
pub fn filter_value(value: &mut Value, key: &str, depth: usize) {
|
||||
if depth > MAX_RECURSE_DEPTH {
|
||||
return;
|
||||
}
|
||||
if let Value::Object(object) = value {
|
||||
object.remove(key);
|
||||
for nested in object.values_mut() {
|
||||
match nested {
|
||||
Value::Object(_) => filter_value(nested, key, depth + 1),
|
||||
Value::Array(items) => items
|
||||
.iter_mut()
|
||||
.filter(|item| item.is_object())
|
||||
.for_each(|item| filter_value(item, key, depth + 1)),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `without_prompt_cache_breakpoints`: the marker removed from a message or tool, its content
|
||||
/// parts, its tool calls and its `function`.
|
||||
fn without_prompt_cache_breakpoints(block: &mut Value) {
|
||||
const MARKER: &str = "prompt_cache_breakpoint";
|
||||
let Some(object) = block.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
object.remove(MARKER);
|
||||
for key in ["content", "tool_calls"] {
|
||||
if let Some(Value::Array(items)) = object.get_mut(key) {
|
||||
for item in items.iter_mut().filter_map(Value::as_object_mut) {
|
||||
item.remove(MARKER);
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(Value::Object(function)) = object.get_mut("function") {
|
||||
function.remove(MARKER);
|
||||
}
|
||||
}
|
||||
|
||||
fn without_fields(message: &mut Value, fields: &[&str]) {
|
||||
let Some(object) = message.as_object_mut() else {
|
||||
return;
|
||||
};
|
||||
for field in fields {
|
||||
object.remove(*field);
|
||||
}
|
||||
if let Some(Value::Array(calls)) = object.get_mut("tool_calls") {
|
||||
for call in calls.iter_mut().filter_map(Value::as_object_mut) {
|
||||
call.remove("provider_specific_fields");
|
||||
if let Some(Value::Object(function)) = call.get_mut("function") {
|
||||
function.remove("provider_specific_fields");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `strip_litellm_response_only_fields`.
|
||||
pub fn strip_response_only_fields(message: &mut Value) {
|
||||
without_fields(message, RESPONSE_ONLY_MESSAGE_FIELDS);
|
||||
}
|
||||
|
||||
/// `strip_litellm_internal_message_fields`.
|
||||
pub fn strip_internal_fields(message: &mut Value) {
|
||||
without_fields(message, INTERNAL_MESSAGE_FIELDS);
|
||||
}
|
||||
|
||||
/// `extract_search_results_text`.
|
||||
pub fn search_results_text(search_results: Option<&Value>) -> String {
|
||||
let Some(Value::Array(results)) = search_results else {
|
||||
return String::new();
|
||||
};
|
||||
let mut text = String::new();
|
||||
for result in results.iter().filter_map(Value::as_object) {
|
||||
for key in ["source", "title"] {
|
||||
if let Some(value) = result.get(key).and_then(Value::as_str) {
|
||||
text.push_str(value);
|
||||
}
|
||||
}
|
||||
if let Some(Value::Array(blocks)) = result.get("content") {
|
||||
for block in blocks {
|
||||
if let Some(value) = block.get("text").and_then(Value::as_str) {
|
||||
text.push_str(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(citations) = result.get("citations").filter(|c| !c.is_null()) {
|
||||
text.push_str(&serde_json::to_string(citations).unwrap_or_default());
|
||||
}
|
||||
}
|
||||
text
|
||||
}
|
||||
|
||||
/// `convert_content_list_to_str`: the text parts joined, then any `search_results` text.
|
||||
pub fn content_list_to_str(message: &Value) -> String {
|
||||
let mut text = String::new();
|
||||
match message.get("content") {
|
||||
Some(Value::Array(parts)) => {
|
||||
for part in parts {
|
||||
if let Some(value) = part.get("text").and_then(Value::as_str) {
|
||||
text.push_str(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
Some(Value::String(content)) => text.push_str(content),
|
||||
_ => {}
|
||||
}
|
||||
text.push_str(&search_results_text(message.get("search_results")));
|
||||
text
|
||||
}
|
||||
|
||||
/// The `reasoning` some OpenAI-compatible servers send, under LiteLLM's `reasoning_content`.
|
||||
pub fn rename_reasoning(message: &mut Map<String, Value>) -> bool {
|
||||
if message.contains_key("reasoning_content") || !message.contains_key("reasoning") {
|
||||
return false;
|
||||
}
|
||||
if let Some(reasoning) = message.remove("reasoning") {
|
||||
message.insert("reasoning_content".into(), reasoning);
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
/// `_map_reasoning_to_reasoning_content` over a chunk's deltas.
|
||||
pub fn rename_reasoning_in_deltas(chunk: &mut Map<String, Value>) -> bool {
|
||||
let mut changed = false;
|
||||
if let Some(Value::Array(choices)) = chunk.get_mut("choices") {
|
||||
for delta in choices
|
||||
.iter_mut()
|
||||
.filter_map(|choice| choice.get_mut("delta")?.as_object_mut())
|
||||
{
|
||||
changed |= rename_reasoning(delta);
|
||||
}
|
||||
}
|
||||
changed
|
||||
}
|
||||
|
||||
/// `extract_reasoning_content` over a complete response's messages.
|
||||
pub fn rename_reasoning_in_messages(body: &mut Map<String, Value>) -> bool {
|
||||
let mut changed = false;
|
||||
for message in messages_mut(body) {
|
||||
changed |= rename_reasoning(message);
|
||||
}
|
||||
changed
|
||||
}
|
||||
|
||||
/// Each choice's message in a complete response.
|
||||
pub fn messages_mut(
|
||||
body: &mut Map<String, Value>,
|
||||
) -> impl Iterator<Item = &mut Map<String, Value>> {
|
||||
body.get_mut("choices")
|
||||
.and_then(Value::as_array_mut)
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.filter_map(|choice| choice.get_mut("message")?.as_object_mut())
|
||||
}
|
||||
|
||||
/// `OpenAILikeChatConfig._sanitize_usage_obj`: a null `*_tokens` usage count becomes 0.
|
||||
pub fn sanitize_usage(body: &mut Map<String, Value>) -> bool {
|
||||
let mut changed = false;
|
||||
if let Some(Value::Object(usage)) = body.get_mut("usage") {
|
||||
for (key, value) in usage.iter_mut() {
|
||||
if key.ends_with("_tokens") && value.is_null() {
|
||||
*value = json!(0);
|
||||
changed = true;
|
||||
}
|
||||
}
|
||||
}
|
||||
changed
|
||||
}
|
||||
|
||||
/// `_json_mode_convert_tool_response_to_message`: a lone tool call's arguments become the
|
||||
/// content.
|
||||
pub fn json_mode_tool_call_to_content(body: &mut Map<String, Value>) -> bool {
|
||||
let mut changed = false;
|
||||
for message in messages_mut(body) {
|
||||
let arguments = match message.get("tool_calls") {
|
||||
Some(Value::Array(calls)) if calls.len() == 1 => calls[0]
|
||||
.get("function")
|
||||
.and_then(|function| function.get("arguments"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_owned(),
|
||||
_ => continue,
|
||||
};
|
||||
message.insert("content".into(), Value::String(arguments));
|
||||
message.insert("tool_calls".into(), Value::Null);
|
||||
changed = true;
|
||||
}
|
||||
changed
|
||||
}
|
||||
|
||||
/// Python's `str` indices of `[n]` citation markers in `text`: `(number, start, end)` in code
|
||||
/// points, as `re.finditer(r"\[(\d+)\]")` reports them.
|
||||
pub fn citation_markers(text: &str) -> Vec<(usize, usize, usize)> {
|
||||
let chars: Vec<char> = text.chars().collect();
|
||||
let mut markers = Vec::new();
|
||||
let mut index = 0;
|
||||
while index < chars.len() {
|
||||
if chars[index] == '[' {
|
||||
let digits = chars[index + 1..]
|
||||
.iter()
|
||||
.take_while(|c| c.is_ascii_digit())
|
||||
.count();
|
||||
let close = index + 1 + digits;
|
||||
if digits > 0 && chars.get(close) == Some(&']') {
|
||||
let number: String = chars[index + 1..close].iter().collect();
|
||||
if let Ok(number) = number.parse() {
|
||||
markers.push((number, index, close + 1));
|
||||
}
|
||||
index = close + 1;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
index += 1;
|
||||
}
|
||||
markers
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn object(value: Value) -> Map<String, Value> {
|
||||
match value {
|
||||
Value::Object(map) => map,
|
||||
other => panic!("expected an object, got {other}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::bare("https://h/v1", "https://h/v1/chat/completions")]
|
||||
#[case::slash("https://h/v1/", "https://h/v1/chat/completions")]
|
||||
#[case::named("https://h/v1/chat/completions", "https://h/v1/chat/completions")]
|
||||
fn the_gpt_url_appends_the_path_once(#[case] base: &str, #[case] expected: &str) {
|
||||
assert_eq!(gpt_complete_url(base), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn user_content_parts_are_normalized() {
|
||||
let messages = gpt_messages(vec![json!({"role": "user", "content": [
|
||||
{"type": "image_url", "image_url": "https://img"},
|
||||
{"type": "image_url", "image_url": {"url": "u", "format": "image/png", "detail": "low"}},
|
||||
{"type": "file", "file": {"file_data": "data:application/pdf;base64,AA", "format": "application/pdf"}},
|
||||
]})])
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
messages[0]["content"],
|
||||
json!([
|
||||
{"type": "image_url", "image_url": {"url": "https://img"}},
|
||||
{"type": "image_url", "image_url": {"url": "u", "detail": "low"}},
|
||||
{"type": "file", "file": {"file_data": "data:application/pdf;base64,AA", "filename": "my_file.pdf"}},
|
||||
])
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::pdf_url(json!({"role": "user", "content": [{"type": "file", "file": {"file_id": "https://x/a.pdf"}}]}))]
|
||||
#[case::tool_image(json!({"role": "tool", "tool_call_id": "c", "content": [{"type": "image_url", "image_url": {"url": "u"}}]}))]
|
||||
fn shapes_python_rewrites_with_io_are_refused(#[case] message: Value) {
|
||||
assert!(matches!(
|
||||
gpt_messages(vec![message]),
|
||||
Err(Error::Unsupported(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_file_part_without_its_file_is_rejected() {
|
||||
assert!(matches!(
|
||||
gpt_messages(vec![json!({"role": "user", "content": [{"type": "file"}]})]),
|
||||
Err(Error::InvalidRequest(_))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn tool_reference_parts_leave_tool_messages() {
|
||||
let messages = gpt_messages(vec![json!({"role": "tool", "tool_call_id": "c",
|
||||
"content": [{"type": "tool_reference", "tool_name": "f"}]})])
|
||||
.unwrap();
|
||||
assert_eq!(messages[0]["content"], json!(""));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn cache_control_is_filtered_at_any_depth() {
|
||||
let mut message = json!({"role": "user", "cache_control": {"type": "ephemeral"},
|
||||
"content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}]});
|
||||
filter_value(&mut message, "cache_control", 0);
|
||||
assert_eq!(
|
||||
message,
|
||||
json!({"role": "user", "content": [{"type": "text", "text": "hi"}]})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::exact(1024, "low")]
|
||||
#[case::medium(3000, "medium")]
|
||||
#[case::high(5000, "high")]
|
||||
#[case::minimal(10, "minimal")]
|
||||
fn thinking_budgets_bucket_into_efforts(#[case] budget: i64, #[case] effort: &str) {
|
||||
assert_eq!(reasoning_effort_from_thinking_budget(budget), effort);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::declared("high", &["low", "high"], "high")]
|
||||
#[case::rounds_up("medium", &["low", "high"], "high")]
|
||||
#[case::rounds_down("max", &["low", "high"], "high")]
|
||||
#[case::none("none", &["low"], "none")]
|
||||
fn efforts_round_to_a_declared_level(
|
||||
#[case] requested: &str,
|
||||
#[case] declared: &[&str],
|
||||
#[case] expected: &str,
|
||||
) {
|
||||
let declared: Vec<String> = declared.iter().map(|d| (*d).to_owned()).collect();
|
||||
assert_eq!(
|
||||
nearest_declared_reasoning_effort(requested, &declared),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn citation_markers_count_code_points() {
|
||||
assert_eq!(
|
||||
citation_markers("é [1] and [22][x]"),
|
||||
vec![(1, 2, 5), (22, 10, 14)]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn json_mode_moves_a_lone_tool_call_into_the_content() {
|
||||
let mut body = object(
|
||||
json!({"choices": [{"message": {"role": "assistant", "content": null,
|
||||
"tool_calls": [{"id": "c", "type": "function", "function": {"name": "json_tool_call", "arguments": "{\"a\":1}"}}]},
|
||||
"finish_reason": "tool_calls"}]}),
|
||||
);
|
||||
assert!(json_mode_tool_call_to_content(&mut body));
|
||||
assert_eq!(
|
||||
body["choices"][0]["message"],
|
||||
json!({"role": "assistant", "content": "{\"a\":1}", "tool_calls": null})
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1 +1,2 @@
|
|||
pub mod compatible;
|
||||
pub mod transformation;
|
||||
|
|
|
|||
|
|
@ -16,9 +16,13 @@ use crate::{
|
|||
Error,
|
||||
base_llm::{
|
||||
auth::AuthScheme,
|
||||
chat::transformation::{
|
||||
BaseConfig, Headers, ProviderChatRequestData, ProviderChatResponseData,
|
||||
ValidatedEnvironment,
|
||||
chat::{
|
||||
params::ChatRequestContext,
|
||||
streaming::StreamShape,
|
||||
transformation::{
|
||||
BaseConfig, Headers, ProviderChatRequestData, ProviderChatResponseData,
|
||||
ValidatedEnvironment,
|
||||
},
|
||||
},
|
||||
},
|
||||
openai_like::common_utils::{complete_openai_like_url, openai_compatible_provider_info},
|
||||
|
|
@ -90,6 +94,7 @@ impl BaseConfig for OpenAILikeChatConfig {
|
|||
model: &str,
|
||||
messages: Vec<ChatMessage>,
|
||||
optional_params: Map<String, Value>,
|
||||
_context: &ChatRequestContext<'_>,
|
||||
) -> Result<ProviderChatRequestData, Error> {
|
||||
let mut params = Map::from_iter(
|
||||
optional_params
|
||||
|
|
@ -113,6 +118,7 @@ impl BaseConfig for OpenAILikeChatConfig {
|
|||
Ok(ProviderChatRequestData {
|
||||
body: Value::Object(body),
|
||||
stream_shape: Default::default(),
|
||||
headers: Vec::new(),
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -120,6 +126,7 @@ impl BaseConfig for OpenAILikeChatConfig {
|
|||
&self,
|
||||
model: &str,
|
||||
response: ProviderChatResponseData,
|
||||
_shape: &StreamShape,
|
||||
) -> Result<ChatCompletionsResponse, Error> {
|
||||
let mut body = response.body;
|
||||
sanitize_usage(&mut body);
|
||||
|
|
@ -177,6 +184,7 @@ impl BaseConfig for OpenAILikeChatConfig {
|
|||
},
|
||||
..ChatCompletionsUsage::default()
|
||||
},
|
||||
extra: Map::new(),
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -266,12 +274,14 @@ fn normalize_choice(position: usize, choice: &Value) -> Result<ChatCompletionsCh
|
|||
.unwrap_or("assistant")
|
||||
.to_string(),
|
||||
content: content.and_then(Value::as_str).map(str::to_string),
|
||||
..ChatCompletionsChoiceMessage::default()
|
||||
},
|
||||
finish_reason: choice
|
||||
.get("finish_reason")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("")
|
||||
.to_string(),
|
||||
logprobs: None,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
@ -310,14 +320,17 @@ mod tests {
|
|||
|
||||
fn transform(model: &str, msgs: Value, opts: Value) -> Value {
|
||||
OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG
|
||||
.transform_request(model, messages(msgs), params(opts))
|
||||
.transform_request(model, messages(msgs), params(opts), &Default::default())
|
||||
.expect("request transforms")
|
||||
.body
|
||||
}
|
||||
|
||||
fn transform_response(body: Value) -> Result<ChatCompletionsResponse, Error> {
|
||||
OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG
|
||||
.transform_response("some-model", ProviderChatResponseData { body })
|
||||
OPENAI_LIKE_CHAT_COMPLETIONS_CONFIG.transform_response(
|
||||
"some-model",
|
||||
ProviderChatResponseData { body },
|
||||
&Default::default(),
|
||||
)
|
||||
}
|
||||
|
||||
fn reason(msgs: Value, opts: Value) -> Option<Unsupported> {
|
||||
|
|
|
|||
11
litellm-rust/crates/llms/src/openrouter/chat/AGENTS.md
Normal file
11
litellm-rust/crates/llms/src/openrouter/chat/AGENTS.md
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
# rules
|
||||
|
||||
- OpenRouter at `https://openrouter.ai/api/v1` with the `HTTP-Referer` and `X-Title` attribution headers
|
||||
- `cache_control` is kept for caching model families and moved into the last content block
|
||||
- `usage: {include: true}` is always asked for unless the caller set `usage`
|
||||
- Shared flow and response handling: `../../openai/chat/compatible.rs`
|
||||
|
||||
# references
|
||||
|
||||
- https://openrouter.ai/docs/api-reference/chat-completion
|
||||
- litellm/llms/openrouter/chat/transformation.py
|
||||
1
litellm-rust/crates/llms/src/openrouter/chat/mod.rs
Normal file
1
litellm-rust/crates/llms/src/openrouter/chat/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
231
litellm-rust/crates/llms/src/openrouter/chat/transformation.rs
Normal file
231
litellm-rust/crates/llms/src/openrouter/chat/transformation.rs
Normal file
|
|
@ -0,0 +1,231 @@
|
|||
//! `litellm/llms/openrouter/chat/transformation.py`: OpenRouter at
|
||||
//! `https://openrouter.ai/api/v1`, with LiteLLM's attribution headers, `cache_control` kept
|
||||
//! for the model families that cache, moved into the last content block, and the usage
|
||||
//! accounting OpenRouter reports cost in always asked for.
|
||||
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::{
|
||||
Error,
|
||||
openai::chat::compatible::{
|
||||
CompatibleContext, GPT_BASE_PARAMS, OpenAiCompatibleChatConfig, copy_supported,
|
||||
gpt_messages,
|
||||
},
|
||||
};
|
||||
|
||||
const CACHE_CONTROL_MODEL_FAMILIES: &[&str] = &["claude", "gemini", "minimax", "glm", "z-ai"];
|
||||
|
||||
pub struct OpenRouterChatConfig;
|
||||
|
||||
pub const OPENROUTER_CHAT_CONFIG: OpenRouterChatConfig = OpenRouterChatConfig;
|
||||
|
||||
impl OpenAiCompatibleChatConfig for OpenRouterChatConfig {
|
||||
fn provider(&self) -> LlmProviders {
|
||||
LlmProviders::Openrouter
|
||||
}
|
||||
|
||||
fn label(&self) -> &'static str {
|
||||
"OpenRouter"
|
||||
}
|
||||
|
||||
fn api_base_env(&self) -> &'static [&'static str] {
|
||||
&["OPENROUTER_API_BASE"]
|
||||
}
|
||||
|
||||
fn default_api_base(&self) -> Option<&'static str> {
|
||||
Some("https://openrouter.ai/api/v1")
|
||||
}
|
||||
|
||||
fn api_key_env(&self) -> &'static [&'static str] {
|
||||
&["OPENROUTER_API_KEY", "OR_API_KEY"]
|
||||
}
|
||||
|
||||
fn header_env(&self) -> &'static [&'static str] {
|
||||
&["OR_SITE_URL", "OR_APP_NAME"]
|
||||
}
|
||||
|
||||
fn headers(
|
||||
&self,
|
||||
_body: &Map<String, Value>,
|
||||
lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Vec<(String, String)> {
|
||||
let setting = |name: &str, default: &str| {
|
||||
lookup(name)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or_else(|| default.to_owned())
|
||||
};
|
||||
vec![
|
||||
(
|
||||
"HTTP-Referer".into(),
|
||||
setting("OR_SITE_URL", "https://litellm.ai"),
|
||||
),
|
||||
("X-Title".into(), setting("OR_APP_NAME", "liteLLM")),
|
||||
]
|
||||
}
|
||||
|
||||
fn supported_params(&self, _context: &CompatibleContext<'_>) -> Vec<&'static str> {
|
||||
GPT_BASE_PARAMS
|
||||
.iter()
|
||||
.copied()
|
||||
.chain(["response_format", "reasoning_effort", "thinking"])
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn map_params(
|
||||
&self,
|
||||
non_default: &Map<String, Value>,
|
||||
optional: &mut Map<String, Value>,
|
||||
context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
copy_supported(non_default, optional, &self.supported_params(context));
|
||||
if optional.get("reasoning_effort") == Some(&json!("max")) {
|
||||
optional.insert("reasoning_effort".into(), json!("xhigh"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn keeps_cache_control(&self, context: &CompatibleContext<'_>) -> bool {
|
||||
caches(context.model)
|
||||
}
|
||||
|
||||
fn transform_messages(
|
||||
&self,
|
||||
messages: Vec<Value>,
|
||||
_optional: &Map<String, Value>,
|
||||
context: &CompatibleContext<'_>,
|
||||
) -> Result<Vec<Value>, Error> {
|
||||
match caches(context.model) {
|
||||
true => gpt_messages(
|
||||
messages
|
||||
.into_iter()
|
||||
.map(move_cache_control_to_content)
|
||||
.collect(),
|
||||
),
|
||||
false => gpt_messages(messages),
|
||||
}
|
||||
}
|
||||
|
||||
fn finish_request(
|
||||
&self,
|
||||
body: &mut Map<String, Value>,
|
||||
_context: &CompatibleContext<'_>,
|
||||
) -> Result<(), Error> {
|
||||
body.entry("usage")
|
||||
.or_insert_with(|| json!({"include": true}));
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn caches(model: &str) -> bool {
|
||||
let model = model.to_lowercase();
|
||||
CACHE_CONTROL_MODEL_FAMILIES
|
||||
.iter()
|
||||
.any(|family| model.contains(family))
|
||||
}
|
||||
|
||||
/// `_move_cache_control_to_content`: only the last block carries it, to stay within
|
||||
/// Anthropic's four cache breakpoints.
|
||||
fn move_cache_control_to_content(mut message: Value) -> Value {
|
||||
let Some(object) = message.as_object_mut() else {
|
||||
return message;
|
||||
};
|
||||
let Some(cache_control) = object.remove("cache_control").filter(|c| !c.is_null()) else {
|
||||
return message;
|
||||
};
|
||||
match object.get_mut("content") {
|
||||
Some(Value::Array(blocks)) => {
|
||||
if let Some(Value::Object(last)) = blocks.last_mut() {
|
||||
last.insert("cache_control".into(), cache_control);
|
||||
}
|
||||
}
|
||||
content => {
|
||||
let text = content.map(Value::take).unwrap_or(Value::Null);
|
||||
object.insert(
|
||||
"content".into(),
|
||||
json!([{"type": "text", "text": text, "cache_control": cache_control}]),
|
||||
);
|
||||
}
|
||||
}
|
||||
message
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::openai::chat::compatible::transform_request;
|
||||
|
||||
fn send(model: &str, params: Value) -> Map<String, Value> {
|
||||
let Value::Object(body) = params else {
|
||||
panic!("an object")
|
||||
};
|
||||
transform_request(
|
||||
&OPENROUTER_CHAT_CONFIG,
|
||||
&CompatibleContext {
|
||||
model,
|
||||
drop_params: false,
|
||||
model_info: None,
|
||||
},
|
||||
body,
|
||||
)
|
||||
.unwrap()
|
||||
.body
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn openrouter_params_and_usage_accounting_go_along() {
|
||||
let sent = send(
|
||||
"openai/gpt-4o",
|
||||
json!({"messages": [], "reasoning_effort": "max",
|
||||
"transforms": ["middle-out"], "route": "fallback"}),
|
||||
);
|
||||
assert_eq!(
|
||||
Value::Object(sent),
|
||||
json!({"model": "openai/gpt-4o", "messages": [], "reasoning_effort": "xhigh",
|
||||
"usage": {"include": true}, "transforms": ["middle-out"], "route": "fallback"})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_callers_usage_setting_wins() {
|
||||
let sent = send(
|
||||
"openai/gpt-4o",
|
||||
json!({"messages": [], "usage": {"include": false}}),
|
||||
);
|
||||
assert_eq!(sent["usage"], json!({"include": false}));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::caching_family("anthropic/claude-sonnet-4", json!([
|
||||
{"role": "system", "content": [{"type": "text", "text": "s", "cache_control": {"type": "ephemeral"}}]},
|
||||
{"role": "user", "content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b", "cache_control": {"type": "ephemeral"}}]}]))]
|
||||
#[case::other_family("openai/gpt-4o", json!([
|
||||
{"role": "system", "content": "s"},
|
||||
{"role": "user", "content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}]}]))]
|
||||
fn cache_control_reaches_only_families_that_cache(
|
||||
#[case] model: &str,
|
||||
#[case] messages: Value,
|
||||
) {
|
||||
let sent = send(
|
||||
model,
|
||||
json!({"messages": [
|
||||
{"role": "system", "content": "s", "cache_control": {"type": "ephemeral"}},
|
||||
{"role": "user", "cache_control": {"type": "ephemeral"}, "content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}]}]}),
|
||||
);
|
||||
assert_eq!(sent["messages"], messages);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn attribution_headers_default_to_litellm() {
|
||||
assert_eq!(
|
||||
OPENROUTER_CHAT_CONFIG.headers(&Map::new(), &|name| (name == "OR_APP_NAME")
|
||||
.then(|| "app".to_owned())),
|
||||
vec![
|
||||
("HTTP-Referer".to_owned(), "https://litellm.ai".to_owned()),
|
||||
("X-Title".to_owned(), "app".to_owned()),
|
||||
]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,2 +1,3 @@
|
|||
pub mod chat;
|
||||
pub mod messages;
|
||||
pub mod responses;
|
||||
|
|
|
|||
10
litellm-rust/crates/llms/src/perplexity/chat/AGENTS.md
Normal file
10
litellm-rust/crates/llms/src/perplexity/chat/AGENTS.md
Normal file
|
|
@ -0,0 +1,10 @@
|
|||
# rules
|
||||
|
||||
- Perplexity through the OpenAI SDK URL rule at `https://api.perplexity.ai`, with `X-Pplx-Integration: litellm`
|
||||
- Citations become `url_citation` annotations on the first choice, search counts reach `prompt_tokens_details`
|
||||
- Shared flow and response handling: `../../openai/chat/compatible.rs`
|
||||
|
||||
# references
|
||||
|
||||
- https://docs.perplexity.ai/api-reference/chat-completions-post
|
||||
- litellm/llms/perplexity/chat/transformation.py
|
||||
1
litellm-rust/crates/llms/src/perplexity/chat/mod.rs
Normal file
1
litellm-rust/crates/llms/src/perplexity/chat/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod transformation;
|
||||
283
litellm-rust/crates/llms/src/perplexity/chat/transformation.rs
Normal file
283
litellm-rust/crates/llms/src/perplexity/chat/transformation.rs
Normal file
|
|
@ -0,0 +1,283 @@
|
|||
//! `litellm/llms/perplexity/chat/transformation.py`: Perplexity through the OpenAI SDK at
|
||||
//! `https://api.perplexity.ai`, with the subset of OpenAI params it takes, LiteLLM's
|
||||
//! `X-Pplx-Integration` attribution header, and the citations of an answer added to its
|
||||
//! message as `url_citation` annotations with the search counts in the usage.
|
||||
|
||||
use litellm_core_utils::get_llm_provider_logic::LlmProviders;
|
||||
use serde_json::{Map, Value, json};
|
||||
|
||||
use crate::openai::chat::compatible::{
|
||||
CompatibleContext, OpenAiCompatibleChatConfig, ResponseShape, citation_markers,
|
||||
sdk_complete_url,
|
||||
};
|
||||
|
||||
const SUPPORTED_PARAMS: &[&str] = &[
|
||||
"frequency_penalty",
|
||||
"max_tokens",
|
||||
"max_completion_tokens",
|
||||
"presence_penalty",
|
||||
"response_format",
|
||||
"stream",
|
||||
"temperature",
|
||||
"top_p",
|
||||
"max_retries",
|
||||
"extra_headers",
|
||||
];
|
||||
const ATTRIBUTION_HEADER: (&str, &str) = ("X-Pplx-Integration", "litellm");
|
||||
|
||||
pub struct PerplexityChatConfig;
|
||||
|
||||
pub const PERPLEXITY_CHAT_CONFIG: PerplexityChatConfig = PerplexityChatConfig;
|
||||
|
||||
impl OpenAiCompatibleChatConfig for PerplexityChatConfig {
|
||||
fn provider(&self) -> LlmProviders {
|
||||
LlmProviders::Perplexity
|
||||
}
|
||||
|
||||
fn label(&self) -> &'static str {
|
||||
"Perplexity"
|
||||
}
|
||||
|
||||
fn api_base_env(&self) -> &'static [&'static str] {
|
||||
&["PERPLEXITY_API_BASE"]
|
||||
}
|
||||
|
||||
fn default_api_base(&self) -> Option<&'static str> {
|
||||
Some("https://api.perplexity.ai")
|
||||
}
|
||||
|
||||
fn api_key_env(&self) -> &'static [&'static str] {
|
||||
&["PERPLEXITYAI_API_KEY", "PERPLEXITY_API_KEY"]
|
||||
}
|
||||
|
||||
fn complete_url(&self, api_base: &str) -> String {
|
||||
sdk_complete_url(api_base)
|
||||
}
|
||||
|
||||
fn headers(
|
||||
&self,
|
||||
_body: &Map<String, Value>,
|
||||
_lookup: &dyn Fn(&str) -> Option<String>,
|
||||
) -> Vec<(String, String)> {
|
||||
vec![(ATTRIBUTION_HEADER.0.into(), ATTRIBUTION_HEADER.1.into())]
|
||||
}
|
||||
|
||||
fn supported_params(&self, context: &CompatibleContext<'_>) -> Vec<&'static str> {
|
||||
let mut params = SUPPORTED_PARAMS.to_vec();
|
||||
if context.supports(self.name(), "supports_reasoning") {
|
||||
params.push("reasoning_effort");
|
||||
}
|
||||
if context.supports(self.name(), "supports_web_search") {
|
||||
params.push("web_search_options");
|
||||
}
|
||||
params
|
||||
}
|
||||
|
||||
/// `_enhance_usage_with_perplexity_fields` and `_add_citations_as_annotations`.
|
||||
fn fix_response(&self, body: &mut Map<String, Value>, _shape: &ResponseShape) -> bool {
|
||||
let usage_changed = enhance_usage(body);
|
||||
add_citations_as_annotations(body) || usage_changed
|
||||
}
|
||||
}
|
||||
|
||||
/// The search count Perplexity reports, wherever it put it.
|
||||
fn search_queries(body: &Map<String, Value>) -> Option<u64> {
|
||||
let usage = body.get("usage");
|
||||
let field = |value: Option<&Value>, key: &str| value.and_then(|v| v.get(key)).cloned();
|
||||
[
|
||||
field(usage, "num_search_queries"),
|
||||
body.get("num_search_queries").cloned(),
|
||||
field(usage, "search_queries"),
|
||||
body.get("search_queries").cloned(),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.find(|value| !value.is_null())
|
||||
.and_then(|value| value.as_u64())
|
||||
}
|
||||
|
||||
fn enhance_usage(body: &mut Map<String, Value>) -> bool {
|
||||
let citation_chars: usize = match body.get("citations") {
|
||||
Some(Value::Array(citations)) => citations
|
||||
.iter()
|
||||
.filter(|citation| crate::base_llm::chat::params::truthy(citation))
|
||||
.map(|citation| match citation {
|
||||
Value::String(text) => text.chars().count(),
|
||||
other => other.to_string().chars().count(),
|
||||
})
|
||||
.sum(),
|
||||
_ => 0,
|
||||
};
|
||||
let citation_tokens = match citation_chars {
|
||||
0 => 0,
|
||||
chars => (chars / 4).max(1),
|
||||
};
|
||||
let searches = search_queries(body).filter(|count| *count > 0);
|
||||
if citation_tokens == 0 && searches.is_none() {
|
||||
return false;
|
||||
}
|
||||
let usage = body
|
||||
.entry("usage")
|
||||
.or_insert_with(|| json!({"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}));
|
||||
if usage.is_null() {
|
||||
*usage = json!({"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0});
|
||||
}
|
||||
let Some(usage) = usage.as_object_mut() else {
|
||||
return false;
|
||||
};
|
||||
if citation_tokens > 0 {
|
||||
usage.insert("citation_tokens".into(), json!(citation_tokens));
|
||||
}
|
||||
if let Some(searches) = searches {
|
||||
let details = usage
|
||||
.entry("prompt_tokens_details")
|
||||
.or_insert_with(|| json!({}));
|
||||
if !details.is_object() {
|
||||
*details = json!({});
|
||||
}
|
||||
if let Some(details) = details.as_object_mut() {
|
||||
details.insert("web_search_requests".into(), json!(searches));
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
fn add_citations_as_annotations(body: &mut Map<String, Value>) -> bool {
|
||||
let citations: Vec<Option<String>> = match body.get("citations") {
|
||||
Some(Value::Array(citations)) => citations
|
||||
.iter()
|
||||
.map(|citation| citation.as_str().map(str::to_owned))
|
||||
.collect(),
|
||||
_ => Vec::new(),
|
||||
};
|
||||
let titles: Vec<(String, String)> = match body.get("search_results") {
|
||||
Some(Value::Array(results)) => results
|
||||
.iter()
|
||||
.filter_map(|result| {
|
||||
Some((
|
||||
result.get("url")?.as_str()?.to_owned(),
|
||||
result.get("title")?.as_str()?.to_owned(),
|
||||
))
|
||||
})
|
||||
.collect(),
|
||||
_ => Vec::new(),
|
||||
};
|
||||
let Some(message) = body
|
||||
.get_mut("choices")
|
||||
.and_then(Value::as_array_mut)
|
||||
.and_then(|choices| choices.first_mut())
|
||||
.and_then(|choice| choice.get_mut("message"))
|
||||
.and_then(Value::as_object_mut)
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
let Some(content) = message
|
||||
.get("content")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|content| !content.is_empty())
|
||||
else {
|
||||
return false;
|
||||
};
|
||||
let annotations: Vec<Value> = citation_markers(content)
|
||||
.into_iter()
|
||||
.filter_map(|(number, start, end)| {
|
||||
let url = citations.get(number.checked_sub(1)?)?.as_ref()?;
|
||||
let title = titles
|
||||
.iter()
|
||||
.rev()
|
||||
.find(|(known, _)| known == url)
|
||||
.map_or("", |(_, title)| title.as_str());
|
||||
Some(json!({"type": "url_citation", "url_citation": {
|
||||
"url": url, "title": title, "start_index": start, "end_index": end,
|
||||
}}))
|
||||
})
|
||||
.collect();
|
||||
if annotations.is_empty() {
|
||||
return false;
|
||||
}
|
||||
match message.get_mut("annotations") {
|
||||
Some(Value::Array(existing)) => existing.extend(annotations),
|
||||
_ => {
|
||||
message.insert("annotations".into(), Value::Array(annotations));
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::*;
|
||||
use crate::{
|
||||
Error, ErrorDetail,
|
||||
openai::chat::compatible::{fix_response, transform_request},
|
||||
};
|
||||
|
||||
fn object(value: Value) -> Map<String, Value> {
|
||||
let Value::Object(map) = value else {
|
||||
panic!("an object")
|
||||
};
|
||||
map
|
||||
}
|
||||
|
||||
fn send(params: Value) -> Result<Map<String, Value>, Error> {
|
||||
transform_request(
|
||||
&PERPLEXITY_CHAT_CONFIG,
|
||||
&CompatibleContext {
|
||||
model: "sonar",
|
||||
drop_params: false,
|
||||
model_info: None,
|
||||
},
|
||||
object(params),
|
||||
)
|
||||
.map(|sent| sent.body)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn max_completion_tokens_goes_as_sent() {
|
||||
let sent = send(
|
||||
json!({"messages": [], "max_completion_tokens": 7, "search_domain_filter": ["a.com"]}),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
Value::Object(sent),
|
||||
json!({"model": "sonar", "messages": [], "max_completion_tokens": 7, "search_domain_filter": ["a.com"]})
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::tools(json!({"messages": [], "tools": [{"type": "function", "function": {"name": "f"}}]}))]
|
||||
#[case::seed(json!({"messages": [], "seed": 1}))]
|
||||
fn params_perplexity_lacks_are_refused(#[case] params: Value) {
|
||||
assert!(matches!(
|
||||
send(params),
|
||||
Err(Error::InvalidRequest(ErrorDetail::UnsupportedParams { .. }))
|
||||
));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn citations_become_annotations_and_searches_reach_the_usage() {
|
||||
let mut body = object(json!({
|
||||
"model": "sonar",
|
||||
"citations": ["https://a.example", "https://b.example"],
|
||||
"search_results": [{"url": "https://b.example", "title": "B"}],
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "é see [2] and [3]"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 3, "completion_tokens": 4, "total_tokens": 7, "num_search_queries": 2},
|
||||
}));
|
||||
assert!(fix_response(
|
||||
&PERPLEXITY_CHAT_CONFIG,
|
||||
&mut body,
|
||||
&ResponseShape::default()
|
||||
));
|
||||
assert_eq!(
|
||||
body["choices"][0]["message"]["annotations"],
|
||||
json!([{"type": "url_citation", "url_citation": {"url": "https://b.example", "title": "B", "start_index": 6, "end_index": 9}}])
|
||||
);
|
||||
assert_eq!(body["usage"]["citation_tokens"], json!(8));
|
||||
assert_eq!(
|
||||
body["usage"]["prompt_tokens_details"],
|
||||
json!({"web_search_requests": 2})
|
||||
);
|
||||
}
|
||||
}
|
||||
1
litellm-rust/crates/llms/src/perplexity/mod.rs
Normal file
1
litellm-rust/crates/llms/src/perplexity/mod.rs
Normal file
|
|
@ -0,0 +1 @@
|
|||
pub mod chat;
|
||||
11
litellm-rust/crates/llms/src/together_ai/chat/AGENTS.md
Normal file
11
litellm-rust/crates/llms/src/together_ai/chat/AGENTS.md
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
# rules
|
||||
|
||||
- Together AI at `https://api.together.ai/v1`
|
||||
- Tool and `response_format` params are refused or dropped only when the `together_ai/<model>` cost map entry rules them out
|
||||
- `reasoning_effort` is translated per model family, `none` turns reasoning off with `reasoning.enabled: false`
|
||||
- Shared flow and response handling: `../../openai/chat/compatible.rs`
|
||||
|
||||
# references
|
||||
|
||||
- https://docs.together.ai/docs/chat-overview
|
||||
- litellm/llms/together_ai/chat/transformation.py
|
||||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue