mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_scx_ai_provider
# Conflicts: # ui/litellm-dashboard/src/components/provider_info_helpers.tsx
This commit is contained in:
commit
d499e68e55
1826 changed files with 106214 additions and 36368 deletions
|
|
@ -2744,84 +2744,6 @@ jobs:
|
|||
file: ./coverage.xml
|
||||
flags: circleci
|
||||
|
||||
ui_build:
|
||||
docker:
|
||||
- image: cimg/node:24.19@sha256:8966565f07189a67d64d6808a2b127f31dafae566508e3547f55640e1070bfad
|
||||
auth:
|
||||
username: ${DOCKERHUB_USERNAME}
|
||||
password: ${DOCKERHUB_PASSWORD}
|
||||
resource_class: medium+
|
||||
working_directory: ~/project
|
||||
steps:
|
||||
- checkout
|
||||
- skip_if_unrelated_changes:
|
||||
category: client
|
||||
- setup_google_dns
|
||||
- restore_cache:
|
||||
keys:
|
||||
- ui-build-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
|
||||
- ui-build-deps-v1-
|
||||
- restore_cache:
|
||||
keys:
|
||||
- ui-nextjs-cache-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
|
||||
- ui-nextjs-cache-v1-
|
||||
- run:
|
||||
name: Install dependencies
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
npm ci
|
||||
- save_cache:
|
||||
key: ui-build-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
|
||||
paths:
|
||||
- ui/litellm-dashboard/node_modules
|
||||
- run:
|
||||
name: Build UI
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
source ./build_ui.sh
|
||||
- save_cache:
|
||||
key: ui-nextjs-cache-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
|
||||
paths:
|
||||
- ui/litellm-dashboard/.next/cache
|
||||
- persist_to_workspace:
|
||||
root: .
|
||||
paths:
|
||||
- litellm/proxy/_experimental/out
|
||||
|
||||
ui_unit_tests:
|
||||
docker:
|
||||
- image: cimg/node:24.19@sha256:8966565f07189a67d64d6808a2b127f31dafae566508e3547f55640e1070bfad
|
||||
auth:
|
||||
username: ${DOCKERHUB_USERNAME}
|
||||
password: ${DOCKERHUB_PASSWORD}
|
||||
resource_class: xlarge
|
||||
working_directory: ~/project
|
||||
steps:
|
||||
- checkout
|
||||
- skip_if_unrelated_changes:
|
||||
category: client
|
||||
- setup_google_dns
|
||||
- restore_cache:
|
||||
keys:
|
||||
- ui-unit-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
|
||||
- ui-unit-deps-v1-
|
||||
- run:
|
||||
name: Install dependencies
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
npm ci
|
||||
- save_cache:
|
||||
key: ui-unit-deps-v1-{{ checksum "ui/litellm-dashboard/package-lock.json" }}
|
||||
paths:
|
||||
- ui/litellm-dashboard/node_modules
|
||||
- run:
|
||||
name: Run UI unit tests (Vitest)
|
||||
command: |
|
||||
cd ui/litellm-dashboard
|
||||
|
||||
CI=true npm run test -- --run \
|
||||
--pool forks --poolOptions.forks.maxForks=6
|
||||
|
||||
e2e_ui_testing:
|
||||
docker:
|
||||
- image: cimg/python:3.12-browsers@sha256:b432899af01c9a311bf74f4f22e9ada2e5306d4b1b4383f8d29e1228a5844ef2
|
||||
|
|
@ -3181,12 +3103,6 @@ workflows:
|
|||
filters: *main_branches
|
||||
- litellm_router_unit_testing:
|
||||
filters: *main_branches
|
||||
- ui_build:
|
||||
filters: *main_branches
|
||||
- ui_unit_tests:
|
||||
requires:
|
||||
- ui_build
|
||||
filters: *main_branches
|
||||
- auth_ui_unit_tests:
|
||||
filters: *main_branches
|
||||
- proxy_behavior_tests:
|
||||
|
|
|
|||
56
.github/ISSUE_TEMPLATE/bug_report.yml
vendored
56
.github/ISSUE_TEMPLATE/bug_report.yml
vendored
|
|
@ -23,30 +23,56 @@ body:
|
|||
label: What happened?
|
||||
description: Also tell us, what did you expect to happen?
|
||||
placeholder: Tell us what you see!
|
||||
value: "A bug happened!"
|
||||
validations:
|
||||
required: true
|
||||
- type: textarea
|
||||
id: steps-to-reproduce
|
||||
id: user-flow
|
||||
attributes:
|
||||
label: Steps to Reproduce
|
||||
description: Please provide a numbered list of the exact steps to reproduce this bug (include a curl/python snippet to reproduce it). Number each step (1., 2., 3., ...) in the order you performed them.
|
||||
label: User Flow
|
||||
description: |
|
||||
Two ordered lists, "Before a (hypothetical) fix" and "After a (hypothetical) fix", walking the same end user through the same task, written strictly from that user's seat. Every rule below applies.
|
||||
|
||||
- Describe the real application and the routes its users actually hit, not a generic scenario
|
||||
- Lead each list with one plain sentence saying where the flow fails (before) or would succeed (after), then number the steps
|
||||
- Every step is something the user does or observes: the HTTP method and full URL they hit, what they sent, and what visibly came back (status code, error text, the shape of an ID). UI steps name the page URL and what is on screen
|
||||
- No LiteLLM internals: never name functions, files, DB tables, config classes, hooks, callbacks, or code paths. "The upload hands back an ID that looks like OpenAI's own `file-abc123` instead of the scrambled one the gateway returned" is right, "no managed-file row was registered" is wrong
|
||||
- Keep the two lists step-for-step identical until they diverge, so the broken step is obvious
|
||||
- If the bug has a security or authorization consequence, end each list with what another user can do that they shouldn't be able to, and what they could no longer do after a fix
|
||||
placeholder: |
|
||||
1. config.yaml file/ .env file/ etc.
|
||||
2. Run the following code...
|
||||
3. Observe the error...
|
||||
value: |
|
||||
1.
|
||||
2.
|
||||
3.
|
||||
Before a (hypothetical) fix: a developer whose app streams chat completions gets no token counts back, so their cost dashboard reads zero
|
||||
|
||||
1. They send POST https://litellm-domain/v1/chat/completions with "stream": true and no stream_options
|
||||
2. The last SSE chunk arrives with "usage": null, so their app records 0 prompt and 0 completion tokens
|
||||
3. They open https://litellm-domain/ui/?page=logs and see the request logged at $0 spend
|
||||
|
||||
After a (hypothetical) fix: the same request comes back with real token counts, so the dashboard shows real spend
|
||||
|
||||
1. The proxy admin sets always_include_stream_usage: true and restarts the proxy
|
||||
2. The developer sends the same POST https://litellm-domain/v1/chat/completions with "stream": true and no stream_options
|
||||
3. The last SSE chunk now carries a usage object with real prompt and completion token counts
|
||||
4. https://litellm-domain/ui/?page=logs shows that request at non-zero spend
|
||||
validations:
|
||||
required: true
|
||||
- type: textarea
|
||||
id: logs
|
||||
id: proof-of-bug
|
||||
attributes:
|
||||
label: Relevant log output
|
||||
description: Please copy and paste any relevant log output. This will be automatically formatted into code, so no need for backticks.
|
||||
render: shell
|
||||
label: Proof the bug occurs
|
||||
description: |
|
||||
The commands (e.g., curl) and their full output, screenshots, or a screen recording demonstrating that the bug happens. Every rule below applies.
|
||||
|
||||
- The proof must be completely e2e with no mocks, against a live proxy you ran yourself (e.g., `litellm --config config.yaml --detailed_debug` on localhost:4000), hitting real LLM provider APIs, costing real $ if needed, where the bug involves a provider call. `pytest` commands are not enough
|
||||
- Show exactly what the end user sees or does, matching the User Flow above step for step
|
||||
- Start with the config.yaml (or SDK setup) and any env vars the proxy ran with, then the exact version or commit hash the proof was captured at, so a maintainer can stand up the same proxy before running your commands. Keep the real values for env vars that aren't sensitive, they are often the reason the bug happens, and redact only the secrets: never paste a real API key, virtual key, database URL, or other credential, here or anywhere else in the issue
|
||||
- If the bug applies to more than one of the LLM endpoints (/v1/responses, /v1/chat/completions, /v1/messages), include proof for every one of them, not just one
|
||||
- For UI bugs: include screenshots and the page URLs you were on. Scrub keys and tokens out of screenshots too (for example, the virtual key is briefly shown in the panel right after you create a virtual key)
|
||||
placeholder: |
|
||||
Config / setup the proxy ran with:
|
||||
|
||||
Version or commit:
|
||||
|
||||
Commands and their full output:
|
||||
validations:
|
||||
required: true
|
||||
- type: dropdown
|
||||
id: component
|
||||
attributes:
|
||||
|
|
|
|||
49
.github/ISSUE_TEMPLATE/feature_request.yml
vendored
49
.github/ISSUE_TEMPLATE/feature_request.yml
vendored
|
|
@ -24,10 +24,53 @@ body:
|
|||
validations:
|
||||
required: true
|
||||
- type: textarea
|
||||
id: motivation
|
||||
id: user-flow
|
||||
attributes:
|
||||
label: Motivation, pitch
|
||||
description: Please outline the motivation for the proposal. Is your feature request related to a specific problem? e.g., "I'm working on X and would like Y to be possible". If this is related to another GitHub issue, please link here too.
|
||||
label: User Flow
|
||||
description: |
|
||||
Two ordered lists, "Before this feature (today)" and "After this feature (ideal user flow)", walking the same end user through the same task, written strictly from that user's seat. Every rule below applies.
|
||||
|
||||
- Describe the real application and the routes its users actually hit, not a generic scenario. Link any related GitHub issue or provider API docs
|
||||
- Lead each list with one plain sentence saying where the flow dead-ends today and what it would let them do instead, then number the steps
|
||||
- Every step is something the user does or observes: the HTTP method and full URL they hit, what they sent, and what visibly came back (status code, error text, the shape of an ID). UI steps name the page URL and what is on screen
|
||||
- No LiteLLM internals: never name functions, files, DB tables, config classes, hooks, callbacks, or code paths. Ask for the behavior you need, not the implementation you imagine
|
||||
- Keep the two lists step-for-step identical until they diverge, so the missing capability is obvious
|
||||
- "Before this feature" is also where you show the workaround you're living with, which is what tells us how badly this is needed
|
||||
placeholder: |
|
||||
Before this feature (today): a developer batching nightly summaries has no way to mark those calls as low priority, so they compete with live traffic for the same rate limit
|
||||
|
||||
1. They send POST https://litellm-domain/v1/chat/completions for 500 documents in a loop
|
||||
2. Around document 120 they start getting 429s naming the rpm limit, and their user-facing chat app starts getting them too
|
||||
3. Their workaround is a hand-rolled sleep between calls, which stretches the batch to 3 hours and still collides at peak
|
||||
|
||||
After this feature (ideal user flow): the same batch runs as background work that yields to live traffic
|
||||
|
||||
1. The developer sends the same POST with "service_tier": "flex"
|
||||
2. Batch calls queue behind interactive ones instead of 429ing, and the response comes back with the tier it was served at
|
||||
3. The live chat app keeps returning 200s throughout the batch
|
||||
4. https://litellm-domain/ui/?page=logs shows the batch requests tagged with that tier
|
||||
validations:
|
||||
required: true
|
||||
- type: textarea
|
||||
id: how-far-you-got
|
||||
attributes:
|
||||
label: How far you got
|
||||
description: |
|
||||
Run as many steps of the "After this feature (ideal user flow)" list as you can against a live proxy you ran yourself (e.g., `litellm --config config.yaml --detailed_debug` on localhost:4000), then paste the commands (e.g., curl) and their full output, ending at the step that dead-ends. Every rule below applies.
|
||||
|
||||
- Say plainly what stopped you there, in user terms: the option you passed came back ignored, the response 400'd naming an unsupported field, there is no button on the page for it. This is what proves the feature is genuinely missing rather than undocumented
|
||||
- No mocks. Where the flow involves a provider call, hit the real provider API, even if it costs real $. `pytest` commands are not enough
|
||||
- Include the config.yaml (or SDK setup) and env vars the proxy ran with, plus the version or commit you were on. Keep the real values for env vars that aren't sensitive, and redact only the secrets: never paste a real API key, virtual key, database URL, or other credential, here or anywhere else in the issue
|
||||
- If the provider already supports this, link their API docs and paste a direct call to them succeeding, so we can see the shape LiteLLM should be sending
|
||||
- For UI asks: include screenshots of the page you got stuck on and its URL. Scrub keys and tokens out of screenshots too (for example, the virtual key is briefly shown in the panel right after you create a virtual key)
|
||||
placeholder: |
|
||||
Config / setup the proxy ran with:
|
||||
|
||||
Version or commit:
|
||||
|
||||
Commands and their full output, up to the step that dead-ends:
|
||||
|
||||
What stopped me there:
|
||||
validations:
|
||||
required: true
|
||||
- type: dropdown
|
||||
|
|
|
|||
40
.github/actions/cache-prisma-binaries/action.yml
vendored
Normal file
40
.github/actions/cache-prisma-binaries/action.yml
vendored
Normal file
|
|
@ -0,0 +1,40 @@
|
|||
name: "Cache Prisma binaries"
|
||||
description: >-
|
||||
Cache the Prisma CLI and engine binaries that `prisma generate` downloads, so
|
||||
only the first job on a given prisma-client-py version pays for the download.
|
||||
|
||||
prisma-client-py shells out to `npm install prisma@<version>` whenever its
|
||||
binary cache directory has no CLI entrypoint, which pulls ~85 MB of query and
|
||||
schema engines over the network. That normally takes a few seconds, but it is
|
||||
unbounded: one shard of a proxy-db run took 5m18s on that single step versus
|
||||
3.8s on its eleven siblings, which pushed the job past its timeout and got a
|
||||
fully passing test run cancelled.
|
||||
|
||||
Callers must not set PRISMA_BINARY_CACHE_DIR. The prisma-client-py default
|
||||
(~/.cache/prisma-python/binaries/<prisma-version>/<engine-version>) is already
|
||||
keyed by both versions, so a cache entry can never be served to a run that
|
||||
expects different binaries.
|
||||
|
||||
runs:
|
||||
using: composite
|
||||
steps:
|
||||
- name: Resolve prisma-client-py version
|
||||
id: version
|
||||
shell: bash
|
||||
run: |
|
||||
version="$(grep -A1 '^name = "prisma"$' uv.lock | sed -n 's/^version = "\(.*\)"$/\1/p' | head -1)"
|
||||
if [ -z "${version}" ]; then
|
||||
echo "could not resolve the prisma package version from uv.lock" >&2
|
||||
exit 1
|
||||
fi
|
||||
echo "version=${version}" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Restore Prisma binaries
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
# ~/.cache/prisma-python holds the npm install tree prisma-client-py
|
||||
# drives; ~/.cache/prisma is where @prisma/engines stages its downloads.
|
||||
path: |
|
||||
~/.cache/prisma-python
|
||||
~/.cache/prisma
|
||||
key: ${{ runner.os }}-prisma-binaries-${{ steps.version.outputs.version }}
|
||||
7
.github/ci-coverage-allowlist.yml
vendored
7
.github/ci-coverage-allowlist.yml
vendored
|
|
@ -136,13 +136,6 @@ test_paths:
|
|||
- tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py
|
||||
|
||||
dockerfiles:
|
||||
- reason: >-
|
||||
The componentized images the microservices chart deploys are built by no job; wiring both into
|
||||
the scan workflow costs a full image build each and is deferred to a change that prices the
|
||||
whole set
|
||||
paths:
|
||||
- backend/Dockerfile
|
||||
- gateway/Dockerfile
|
||||
- reason: >-
|
||||
The dashboard container is a static Next.js export served by nginx, and the dashboard build
|
||||
and lint workflows already exercise that output, so building the image adds no signal about it
|
||||
|
|
|
|||
33
.github/pull_request_template.md
vendored
33
.github/pull_request_template.md
vendored
|
|
@ -13,6 +13,33 @@ How it solves it:
|
|||
- <blah>
|
||||
- ...
|
||||
|
||||
## User Flow
|
||||
|
||||
<!-- Two ordered lists, Before and After, walking the same end user through the same task, written strictly from that user's seat
|
||||
Read the linked issue, ticket, or customer thread first so the flow reflects the real application and the routes its users actually hit; don't invent a generic scenario
|
||||
Lead each list with one plain sentence saying where the flow fails (Before) or succeeds (After), then number the steps
|
||||
Every step is something the user does or observes: the HTTP method and full URL they hit, what they sent, and what visibly came back (status code, error text, the shape of an ID). UI steps name the page URL and what is on screen
|
||||
No LiteLLM internals: never name functions, files, DB tables, config classes, hooks, callbacks, or code paths. "The upload hands back an ID that looks like OpenAI's own `file-abc123` instead of the scrambled one the gateway returned" is right, "no managed-file row was registered" is wrong
|
||||
Keep the two lists step-for-step identical until they diverge, so the changed step is obvious
|
||||
If the bug had a security or authorization consequence, end each list with what another user could or could no longer do
|
||||
Regenerate this section whenever new commits change the PR's behavior, so it never describes an older revision
|
||||
|
||||
Example:
|
||||
|
||||
Before: a developer whose app streams chat completions gets no token counts back, so their cost dashboard reads zero
|
||||
|
||||
1. They send POST https://litellm-domain/v1/chat/completions with `"stream": true` and no `stream_options`
|
||||
2. The last SSE chunk arrives with `"usage": null`, so their app records 0 prompt and 0 completion tokens
|
||||
3. They open https://litellm-domain/ui/?page=logs and see the request logged at $0 spend
|
||||
|
||||
After: the same request comes back with real token counts, so the dashboard shows real spend
|
||||
|
||||
1. The proxy admin sets `always_include_stream_usage: true` and restarts the proxy
|
||||
2. The developer sends the same POST https://litellm-domain/v1/chat/completions with `"stream": true` and no `stream_options`
|
||||
3. The last SSE chunk now carries a `usage` object with real prompt and completion token counts
|
||||
4. https://litellm-domain/ui/?page=logs shows that request at non-zero spend
|
||||
-->
|
||||
|
||||
## Relevant issues
|
||||
|
||||
<!-- e.g., "Fixes #000" -->
|
||||
|
|
@ -56,7 +83,11 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
|
|||
🚄 Infrastructure
|
||||
✅ Test
|
||||
|
||||
## Changes
|
||||
## Caveats (if any)
|
||||
|
||||
<!-- Short bullet points, just like the TLDR: one line per bullet, roughly 10 words max
|
||||
Call out known limitations, follow-up work, or anything a reviewer should watch out for
|
||||
Leave this section empty if there are none -->
|
||||
|
||||
## QA runbook
|
||||
|
||||
|
|
|
|||
29
.github/scripts/triage_with_llm.py
vendored
29
.github/scripts/triage_with_llm.py
vendored
|
|
@ -582,7 +582,9 @@ def build_issue_prompt(*, title: str, body: str) -> str:
|
|||
Commands whose external dependencies (LLM provider, DB,
|
||||
network) are mocked or stubbed do NOT count.
|
||||
Prose-only "steps to reproduce" with no run output, video, or
|
||||
screenshot do NOT satisfy (1).
|
||||
screenshot do NOT satisfy (1). An unfilled template scaffold
|
||||
(bare headings such as "Version or commit:" with nothing under
|
||||
them, empty numbered lists) counts as absent, not as evidence.
|
||||
(2) Expected vs. actual behavior (`has_expected_vs_actual`).
|
||||
|
||||
FAIL the bug report if either (1) or (2) is missing. Do not bias
|
||||
|
|
@ -595,6 +597,13 @@ def build_issue_prompt(*, title: str, body: str) -> str:
|
|||
that it does not today).
|
||||
- Motivation / use case with a concrete example (config, API call,
|
||||
UI flow, or scenario showing what's blocked today).
|
||||
- END-TO-END EVIDENCE OF THE DEAD-END (set
|
||||
`has_dead_end_evidence=true` only when this is present): a video,
|
||||
a screenshot, or the exact command(s) actually run paired with
|
||||
their real output, showing the point where the flow stops today.
|
||||
Mocked or stubbed dependencies do NOT count, and an unfilled
|
||||
template scaffold (bare headings, empty numbered lists) counts as
|
||||
absent.
|
||||
|
||||
For an issue that is neither a bug report nor a feature request (a
|
||||
question, support request, or discussion), PASS as long as it has a
|
||||
|
|
@ -608,6 +617,7 @@ def build_issue_prompt(*, title: str, body: str) -> str:
|
|||
"has_repro": boolean,
|
||||
"has_expected_vs_actual": boolean,
|
||||
"has_motivation_example": boolean,
|
||||
"has_dead_end_evidence": boolean,
|
||||
"missing": ["plain-english strings naming what is missing"],
|
||||
"explanation": "1-2 sentence reasoning for the team to skim"
|
||||
}}
|
||||
|
|
@ -705,6 +715,10 @@ _ISSUE_BUG_LABELS: tuple[tuple[str, str], ...] = (
|
|||
)
|
||||
_ISSUE_FEATURE_LABELS: tuple[tuple[str, str], ...] = (
|
||||
("has_motivation_example", "Motivation and concrete example"),
|
||||
(
|
||||
"has_dead_end_evidence",
|
||||
"End-to-end evidence of the dead-end (video, screenshot, or command + real output)",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -836,8 +850,11 @@ def format_issue_close_comment(verdict: dict) -> str:
|
|||
"video, a screenshot, or the exact commands you ran with their real output / "
|
||||
"traceback) plus expected vs. actual behavior. Written steps with no run output, "
|
||||
"video, or screenshot don't count, and mocked or stubbed runs don't count.\n"
|
||||
" - For **feature requests**: a concrete description of what should change, plus a "
|
||||
"use case and example (config / API call / UI flow).\n"
|
||||
" - For **feature requests**: a concrete description of what should change, a "
|
||||
"use case and example (config / API call / UI flow), plus end-to-end evidence of "
|
||||
"the dead-end (a video, a screenshot, or the exact commands you ran with their "
|
||||
"real output showing where the flow stops today). Mocked or stubbed runs don't "
|
||||
"count.\n"
|
||||
"2. Comment `@agent-shin reconsider`. I'll re-run triage and reopen the issue if it "
|
||||
"now meets the bar. (GitHub doesn't let external authors reopen an issue a maintainer "
|
||||
"or bot closed, so the comment-based reconsider is the reliable path.)\n"
|
||||
|
|
@ -943,8 +960,10 @@ def format_grace_warning_issue_comment(verdict: dict) -> str:
|
|||
"screenshot, or the exact commands you ran with their real output / traceback) plus "
|
||||
"expected vs. actual behavior. Written steps with no run output don't count, and "
|
||||
"mocked or stubbed runs don't count.\n"
|
||||
"- For **feature requests**: a concrete description of what should change, plus a use "
|
||||
"case and example (config / API call / UI flow).\n"
|
||||
"- For **feature requests**: a concrete description of what should change, a use "
|
||||
"case and example (config / API call / UI flow), plus end-to-end evidence of the "
|
||||
"dead-end (a video, a screenshot, or the exact commands you ran with their real "
|
||||
"output showing where the flow stops today). Mocked or stubbed runs don't count.\n"
|
||||
"\n"
|
||||
"**If the issue does get auto-closed in 2 hours**, comment `@agent-shin reconsider` "
|
||||
"and I'll re-evaluate. If it now meets the bar, I'll reopen the issue.\n"
|
||||
|
|
|
|||
34
.github/workflows/_test-unit-base.yml
vendored
34
.github/workflows/_test-unit-base.yml
vendored
|
|
@ -18,10 +18,25 @@ on:
|
|||
type: number
|
||||
default: 2
|
||||
timeout-minutes:
|
||||
description: "Job timeout in minutes"
|
||||
description: >-
|
||||
Timeout for the test step alone. Setup (checkout, dependency install,
|
||||
Prisma client generation) gets its own allowance on top, so a slow
|
||||
runner or a cold binary download can never cancel passing tests.
|
||||
required: false
|
||||
type: number
|
||||
default: 20
|
||||
job-timeout-minutes:
|
||||
description: >-
|
||||
Backstop for the whole job. Keep it >= `timeout-minutes` plus 35: 30 for
|
||||
the per-step ceilings on the setup steps below, and 5 for the runner
|
||||
overhead the job clock charges but no step owns (job init, step
|
||||
transitions, post-job cleanup). That headroom is what makes the test
|
||||
budget a floor rather than a hope, since setup cannot overrun into it
|
||||
without failing its own step first. GitHub expressions have no
|
||||
arithmetic, so the sum is passed in rather than computed.
|
||||
required: false
|
||||
type: number
|
||||
default: 55
|
||||
max-failures:
|
||||
description: "Stop after this many failures"
|
||||
required: false
|
||||
|
|
@ -44,30 +59,35 @@ jobs:
|
|||
run:
|
||||
name: Run tests
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: ${{ inputs.timeout-minutes }}
|
||||
timeout-minutes: ${{ inputs.job-timeout-minutes }}
|
||||
outputs:
|
||||
decision: ${{ steps.changes.outputs.decision }}
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
timeout-minutes: 3
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Detect backend-relevant changes
|
||||
id: changes
|
||||
timeout-minutes: 2
|
||||
uses: ./.github/actions/detect-backend-changes
|
||||
|
||||
- name: Set up Python
|
||||
timeout-minutes: 3
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
timeout-minutes: 3
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
timeout-minutes: 5
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
|
|
@ -79,18 +99,24 @@ jobs:
|
|||
|
||||
- name: Install dependencies
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 8
|
||||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: 3
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Generate Prisma client
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
timeout-minutes: 3
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Run tests
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
timeout-minutes: ${{ inputs.timeout-minutes }}
|
||||
env:
|
||||
TEST_PATH: ${{ inputs.test-path }}
|
||||
MAX_FAILURES: ${{ inputs.max-failures }}
|
||||
|
|
|
|||
4
.github/workflows/check-schema-sync.yml
vendored
4
.github/workflows/check-schema-sync.yml
vendored
|
|
@ -10,6 +10,10 @@ on:
|
|||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
check-sync:
|
||||
name: Verify schema.prisma copies match root
|
||||
|
|
|
|||
50
.github/workflows/check-ui-api-types.yml
vendored
50
.github/workflows/check-ui-api-types.yml
vendored
|
|
@ -2,18 +2,19 @@ name: Check UI API Types Sync
|
|||
|
||||
on:
|
||||
pull_request:
|
||||
paths:
|
||||
- "litellm/proxy/**"
|
||||
- "litellm/types/**"
|
||||
- "ui/litellm-dashboard/src/lib/http/schema.d.ts"
|
||||
- "ui/litellm-dashboard/scripts/gen-api-types.mjs"
|
||||
- "ui/litellm-dashboard/package.json"
|
||||
- "ui/litellm-dashboard/package-lock.json"
|
||||
- ".github/workflows/check-ui-api-types.yml"
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
check-sync:
|
||||
name: Verify schema.d.ts matches the proxy OpenAPI spec
|
||||
|
|
@ -24,18 +25,39 @@ jobs:
|
|||
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
fetch-depth: 2
|
||||
|
||||
- name: Detect changes that can affect the generated types
|
||||
id: changes
|
||||
run: |
|
||||
set -euo pipefail
|
||||
if ! base="$(git rev-parse --verify --quiet HEAD^2 >/dev/null && git rev-parse HEAD^1)"; then
|
||||
echo "Not a pull request merge commit, running the full check."
|
||||
echo "relevant=true" >> "$GITHUB_OUTPUT"
|
||||
exit 0
|
||||
fi
|
||||
files="$(git diff --name-only "$base" HEAD)"
|
||||
if grep -Eq '^(litellm/(proxy|types)/|ui/litellm-dashboard/(src/lib/http/schema\.d\.ts|scripts/gen-api-types\.mjs|package(-lock)?\.json)$|\.github/workflows/check-ui-api-types\.yml$)' <<< "$files"; then
|
||||
echo "relevant=true" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "No proxy, types or generator changes in this pull request, nothing to verify."
|
||||
echo "relevant=false" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
- name: Set up Python
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Set up uv
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
uses: ./.github/actions/setup-uv-with-retries
|
||||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache uv dependencies
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.3.0
|
||||
with:
|
||||
path: |
|
||||
|
|
@ -46,14 +68,19 @@ jobs:
|
|||
${{ runner.os }}-uv-
|
||||
|
||||
- name: Install backend dependencies
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
run: .github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
run: uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
- name: Set up Node.js
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0.0
|
||||
with:
|
||||
node-version-file: ui/litellm-dashboard/.nvmrc
|
||||
|
|
@ -61,16 +88,19 @@ jobs:
|
|||
cache-dependency-path: ui/litellm-dashboard/package-lock.json
|
||||
|
||||
- name: Install dashboard dependencies
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
working-directory: ui/litellm-dashboard
|
||||
run: npm ci
|
||||
|
||||
- name: Regenerate types from the live spec
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
working-directory: ui/litellm-dashboard
|
||||
env:
|
||||
LITELLM_PYTHON: "uv run --no-sync python"
|
||||
run: npm run gen:api
|
||||
|
||||
- name: Fail if types are stale
|
||||
if: steps.changes.outputs.relevant == 'true'
|
||||
run: |
|
||||
if ! git diff --exit-code -- ui/litellm-dashboard/src/lib/http/schema.d.ts; then
|
||||
echo "::error file=ui/litellm-dashboard/src/lib/http/schema.d.ts::Generated API types are out of sync with the proxy OpenAPI spec."
|
||||
|
|
|
|||
4
.github/workflows/conventional-commits.yml
vendored
4
.github/workflows/conventional-commits.yml
vendored
|
|
@ -14,6 +14,10 @@ on:
|
|||
permissions:
|
||||
pull-requests: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
lint-pr-title:
|
||||
name: Validate PR title
|
||||
|
|
|
|||
|
|
@ -15,6 +15,10 @@ on:
|
|||
|
||||
permissions: {}
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
guard:
|
||||
name: Block fork dependency changes
|
||||
|
|
|
|||
4
.github/workflows/helm_unit_test.yml
vendored
4
.github/workflows/helm_unit_test.yml
vendored
|
|
@ -9,6 +9,10 @@ on:
|
|||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
unit-test:
|
||||
runs-on: ubuntu-latest
|
||||
|
|
|
|||
65
.github/workflows/image-scan.yml
vendored
65
.github/workflows/image-scan.yml
vendored
|
|
@ -12,6 +12,11 @@ on:
|
|||
- docker/Dockerfile.non_root
|
||||
- migrations/Dockerfile
|
||||
- migrations/run.py
|
||||
- gateway/Dockerfile
|
||||
- gateway/main.py
|
||||
- backend/Dockerfile
|
||||
- backend/main.py
|
||||
- docker/component_entrypoint.sh
|
||||
- litellm-proxy-extras/**
|
||||
- tests/proxy_migration_tests/**
|
||||
- uv.lock
|
||||
|
|
@ -147,3 +152,63 @@ jobs:
|
|||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_offline_image_migration.py -v
|
||||
|
||||
gateway-image:
|
||||
name: gateway-image
|
||||
runs-on: ubuntu-latest
|
||||
if: >-
|
||||
github.event_name != 'pull_request' ||
|
||||
github.event.pull_request.head.repo.full_name == github.repository
|
||||
timeout-minutes: 30
|
||||
permissions:
|
||||
contents: read
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Build gateway image
|
||||
run: docker build -f gateway/Dockerfile -t litellm-gateway-scan:${{ github.sha }} .
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Verify the gateway serves offline as a non-root uid
|
||||
env:
|
||||
LITELLM_IMAGE: litellm-gateway-scan:${{ github.sha }}
|
||||
LITELLM_COMPONENT_PORT: "4000"
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v
|
||||
|
||||
backend-image:
|
||||
name: backend-image
|
||||
runs-on: ubuntu-latest
|
||||
if: >-
|
||||
github.event_name != 'pull_request' ||
|
||||
github.event.pull_request.head.repo.full_name == github.repository
|
||||
timeout-minutes: 30
|
||||
permissions:
|
||||
contents: read
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- name: Build backend image
|
||||
run: docker build -f backend/Dockerfile -t litellm-backend-scan:${{ github.sha }} .
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
with:
|
||||
python-version: "3.12"
|
||||
|
||||
- name: Verify the backend serves offline as a non-root uid
|
||||
env:
|
||||
LITELLM_IMAGE: litellm-backend-scan:${{ github.sha }}
|
||||
LITELLM_COMPONENT_PORT: "4001"
|
||||
run: |
|
||||
python -m pip install "pytest==9.0.3"
|
||||
python -m pytest tests/proxy_migration_tests/test_component_image_serves_offline.py -v
|
||||
|
|
|
|||
5
.github/workflows/mutation-test.yml
vendored
5
.github/workflows/mutation-test.yml
vendored
|
|
@ -57,9 +57,10 @@ jobs:
|
|||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router --extra saml
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
|
|
|
|||
|
|
@ -43,12 +43,13 @@ jobs:
|
|||
with:
|
||||
version: "0.10.9"
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
# The gate provisions its own measurement env (.venv-typecheck: a frozen
|
||||
# uv sync of its canonical dependency groups plus a generated Prisma
|
||||
# client), so no install step here can drift from what local runs measure.
|
||||
- name: Emit basedpyright counts for HEAD
|
||||
env:
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
run: |
|
||||
python scripts/type_check_gate.py --emit-counts-dir "$RUNNER_TEMP/basedpyright-counts"
|
||||
counts_file=$(ls "$RUNNER_TEMP"/basedpyright-counts/basedpyright-counts-*.json)
|
||||
|
|
|
|||
6
.github/workflows/test-code-quality.yml
vendored
6
.github/workflows/test-code-quality.yml
vendored
|
|
@ -65,6 +65,12 @@ jobs:
|
|||
- name: check_provider_folders_documented
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_provider_folders_documented.py
|
||||
|
||||
- name: check_prisma_binary_cache
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_prisma_binary_cache.py
|
||||
|
||||
- name: check_workflow_startup_safety
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/check_workflow_startup_safety.py
|
||||
|
||||
- name: router_code_coverage
|
||||
run: uv run --no-sync python ./tests/code_coverage_tests/router_code_coverage.py
|
||||
|
||||
|
|
|
|||
21
.github/workflows/test-linting.yml
vendored
21
.github/workflows/test-linting.yml
vendored
|
|
@ -11,6 +11,10 @@ on:
|
|||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
lint:
|
||||
runs-on: ubuntu-latest
|
||||
|
|
@ -39,9 +43,10 @@ jobs:
|
|||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
HEAD_SHA: ${{ github.event.pull_request.head.sha }}
|
||||
run: |
|
||||
MERGE_BASE=$(gh api "repos/${{ github.repository }}/compare/${BASE_SHA}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha')
|
||||
retry() { "$@" || { sleep 15; "$@"; } || { sleep 30; "$@"; }; }
|
||||
MERGE_BASE=$(retry gh api "repos/${{ github.repository }}/compare/${BASE_SHA}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha')
|
||||
test -n "$MERGE_BASE"
|
||||
git fetch --no-tags --depth=1 origin "$MERGE_BASE"
|
||||
retry git fetch --no-tags --depth=1 origin "$MERGE_BASE"
|
||||
echo "GATE_BASE_SHA=$MERGE_BASE" >> "$GITHUB_ENV"
|
||||
|
||||
- name: Set up Python
|
||||
|
|
@ -67,12 +72,13 @@ jobs:
|
|||
run: |
|
||||
uv sync --frozen --group proxy-dev --group e2e-dev
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
# basedpyright resolves Prisma's generated client (litellm/proxy/schema.prisma)
|
||||
# only after `prisma generate` writes prisma/client.py et al. Without this the
|
||||
# DB wrappers typed against the generated client would degrade to Unknown.
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
|
|
@ -115,7 +121,6 @@ jobs:
|
|||
- name: Check basedpyright budget (delta vs base)
|
||||
env:
|
||||
GH_TOKEN: ${{ github.token }}
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
run: |
|
||||
uv run --no-sync python scripts/type_check_gate.py --base "$GATE_BASE_SHA"
|
||||
|
||||
|
|
@ -157,7 +162,8 @@ jobs:
|
|||
env:
|
||||
BASE_SHA: ${{ github.event.pull_request.base.sha }}
|
||||
run: |
|
||||
git fetch --no-tags --depth=1 origin "$BASE_SHA"
|
||||
retry() { "$@" || { sleep 15; "$@"; } || { sleep 30; "$@"; }; }
|
||||
retry git fetch --no-tags --depth=1 origin "$BASE_SHA"
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
|
||||
|
|
@ -201,7 +207,8 @@ jobs:
|
|||
GITGUARDIAN_API_KEY: ${{ secrets.GITGUARDIAN_API_KEY }}
|
||||
run: |
|
||||
if [ -n "$GITGUARDIAN_API_KEY" ]; then
|
||||
git fetch --no-tags --unshallow origin
|
||||
retry() { "$@" || { sleep 15; "$@"; } || { sleep 30; "$@"; }; }
|
||||
retry git fetch --no-tags --unshallow origin
|
||||
uv tool run --from 'ggshield==1.48.0' ggshield secret scan repo .
|
||||
else
|
||||
echo "GITGUARDIAN_API_KEY not set, skipping ggshield scan"
|
||||
|
|
|
|||
4
.github/workflows/test-litellm-ui-build.yml
vendored
4
.github/workflows/test-litellm-ui-build.yml
vendored
|
|
@ -10,6 +10,10 @@ on:
|
|||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
build-ui:
|
||||
runs-on: ubuntu-latest
|
||||
|
|
|
|||
4
.github/workflows/test-litellm-ui-lint.yml
vendored
4
.github/workflows/test-litellm-ui-lint.yml
vendored
|
|
@ -10,6 +10,10 @@ on:
|
|||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
frontend-lint:
|
||||
runs-on: ubuntu-latest
|
||||
|
|
|
|||
5
.github/workflows/test-litellm-ui-unit.yml
vendored
5
.github/workflows/test-litellm-ui-unit.yml
vendored
|
|
@ -42,6 +42,11 @@ jobs:
|
|||
- name: Install dependencies
|
||||
run: npm ci
|
||||
|
||||
- name: Run UI type tests (Vitest)
|
||||
env:
|
||||
CI: "true"
|
||||
run: npm run test:types
|
||||
|
||||
- name: Run UI unit tests (Vitest)
|
||||
env:
|
||||
CI: "true"
|
||||
|
|
|
|||
4
.github/workflows/test-mcp.yml
vendored
4
.github/workflows/test-mcp.yml
vendored
|
|
@ -11,6 +11,10 @@ on:
|
|||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
test:
|
||||
runs-on: ubuntu-latest
|
||||
|
|
|
|||
4
.github/workflows/test-model-map.yaml
vendored
4
.github/workflows/test-model-map.yaml
vendored
|
|
@ -11,6 +11,10 @@ on:
|
|||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
|
||||
cancel-in-progress: ${{ github.event_name == 'pull_request' }}
|
||||
|
||||
jobs:
|
||||
validate-model-prices-json:
|
||||
runs-on: ubuntu-latest
|
||||
|
|
|
|||
54
.github/workflows/test-terraform-modules.yml
vendored
Normal file
54
.github/workflows/test-terraform-modules.yml
vendored
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
name: Terraform Modules
|
||||
|
||||
on:
|
||||
push:
|
||||
paths:
|
||||
- "terraform/litellm/aws/**"
|
||||
- ".github/workflows/test-terraform-modules.yml"
|
||||
pull_request:
|
||||
branches:
|
||||
- main
|
||||
- litellm_internal_staging
|
||||
- litellm_oss_staging
|
||||
- "litellm_**"
|
||||
paths:
|
||||
- "terraform/litellm/aws/**"
|
||||
- ".github/workflows/test-terraform-modules.yml"
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
|
||||
concurrency:
|
||||
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
aws-module:
|
||||
name: fmt, validate, test (aws)
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
defaults:
|
||||
run:
|
||||
working-directory: terraform/litellm/aws
|
||||
steps:
|
||||
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
|
||||
with:
|
||||
persist-credentials: false
|
||||
|
||||
- uses: hashicorp/setup-terraform@b9cd54a3c349d3f38e8881555d616ced269862dd # v3.1.2
|
||||
with:
|
||||
terraform_version: 1.13.3
|
||||
terraform_wrapper: false
|
||||
|
||||
- name: fmt
|
||||
run: terraform fmt -recursive -check -diff
|
||||
|
||||
- name: init
|
||||
run: terraform init -backend=false -input=false
|
||||
|
||||
- name: validate
|
||||
run: terraform validate
|
||||
|
||||
# Plan-only, mock_provider-backed: no AWS credentials, no API calls.
|
||||
- name: test
|
||||
run: terraform test
|
||||
|
|
@ -92,9 +92,10 @@ jobs:
|
|||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
|
|
|
|||
|
|
@ -65,10 +65,12 @@ jobs:
|
|||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Generate Prisma client
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
|
|
|
|||
6
.github/workflows/test-unit-proxy-db.yml
vendored
6
.github/workflows/test-unit-proxy-db.yml
vendored
|
|
@ -28,6 +28,10 @@ concurrency:
|
|||
# Most of a shard's time is pytest plugin load + xdist worker imports +
|
||||
# pytest-cov instrumentation, not the tests themselves. Keeping per-shard
|
||||
# work low and matching worker count to runner cores is what controls it.
|
||||
# * `timeout` bounds the pytest step only. Checkout, dependency install, and
|
||||
# Prisma client generation draw on a separate allowance in the base
|
||||
# workflow, so slow setup shows up as a slow job rather than as a
|
||||
# cancelled shard whose tests were passing.
|
||||
# * workers: 4 matches the 4-core ubuntu-latest runner. -n 8 on 4 cores
|
||||
# oversubscribes 2x and workers fight for CPU during their cold-start
|
||||
# imports (measured ~441% CPU for -n 8 locally, i.e. ~55% effective).
|
||||
|
|
@ -131,8 +135,6 @@ jobs:
|
|||
test-path: >-
|
||||
tests/proxy_unit_tests/test_proxy_server.py
|
||||
tests/proxy_unit_tests/test_proxy_server_keys.py
|
||||
tests/proxy_unit_tests/test_proxy_server_caching.py
|
||||
tests/proxy_unit_tests/test_proxy_server_langfuse.py
|
||||
tests/proxy_unit_tests/test_proxy_server_spend.py
|
||||
tests/proxy_unit_tests/test_aproxy_startup.py
|
||||
workers: 4
|
||||
|
|
|
|||
|
|
@ -38,12 +38,15 @@ jobs:
|
|||
tests/test_litellm/proxy/google_endpoints
|
||||
tests/test_litellm/proxy/openai_files_endpoint
|
||||
tests/test_litellm/proxy/batches_endpoints
|
||||
tests/test_litellm/proxy/fine_tuning_endpoints
|
||||
tests/test_litellm/proxy/vector_store_files_endpoints
|
||||
tests/test_litellm/proxy/video_endpoints
|
||||
tests/test_litellm/proxy/response_api_endpoints
|
||||
tests/test_litellm/proxy/image_endpoints
|
||||
tests/test_litellm/proxy/vector_store_endpoints
|
||||
tests/test_litellm/proxy/agent_endpoints
|
||||
tests/test_litellm/proxy/a2a
|
||||
tests/test_litellm/proxy/credential_endpoints
|
||||
tests/test_litellm/proxy/discovery_endpoints
|
||||
tests/test_litellm/proxy/health_endpoints
|
||||
tests/test_litellm/proxy/shutdown
|
||||
|
|
@ -73,4 +76,5 @@ jobs:
|
|||
workers: 4
|
||||
reruns: 2
|
||||
timeout-minutes: 60
|
||||
job-timeout-minutes: 95
|
||||
artifact-name: proxy-server
|
||||
|
|
|
|||
6
.github/workflows/test-unit-proxy-legacy.yml
vendored
6
.github/workflows/test-unit-proxy-legacy.yml
vendored
|
|
@ -82,10 +82,12 @@ jobs:
|
|||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra google --extra proxy --extra semantic-router
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Generate Prisma client
|
||||
if: steps.changes.outputs.decision != 'skip'
|
||||
env:
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
|
|
|
|||
5
.github/workflows/weekly_load_anomaly.yml
vendored
5
.github/workflows/weekly_load_anomaly.yml
vendored
|
|
@ -51,9 +51,10 @@ jobs:
|
|||
run: |
|
||||
.github/scripts/uv_sync_with_retries.sh --frozen --group ci --group proxy-dev --extra proxy
|
||||
|
||||
- name: Cache Prisma binaries
|
||||
uses: ./.github/actions/cache-prisma-binaries
|
||||
|
||||
- name: Generate Prisma client
|
||||
env:
|
||||
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
|
||||
run: |
|
||||
uv run --no-sync prisma generate --schema litellm/proxy/schema.prisma
|
||||
|
||||
|
|
|
|||
18
CLAUDE.md
18
CLAUDE.md
|
|
@ -1,4 +1,12 @@
|
|||
Do not write any comments (existing comments can stay) unless explicitly asked to in a user (not system) prompt
|
||||
Do not write comments unless they are any of:
|
||||
- absolutely necessary to explain some very complex business logic (in which case, keep it concise and clear)
|
||||
- used as an input for tools to read and act on. For example:
|
||||
- entries in `.git-blame-ignore-revs` saying which commit is excluded from git blame
|
||||
- a lint or type checker suppression like `# mutable-ok` or `# pyright: ignore[reportArgumentType] # <reason>` when introducing a truly unavoidable violation
|
||||
- a TODO or FIXME
|
||||
- Not great to have those, but if it's unavoidable, make sure to include a strong, concise reason for why it's there or, better yet, link to a GitHub issue for the follow-up work
|
||||
|
||||
Explanation: The point of this rule is to keep out AI slop comments. AI writes way too many and way too verbose comments. Code comments are, in a way, a violation of DRY code. You must update logic in two locations to change the code, and "hard to change" is literally the definition of tech debt. We should instead aim to write code that is intuitive and clear, even at a glance, to the reader, being both easy to maintain and high performance
|
||||
|
||||
Don't assume that the existing code is correct or the right way of doing things / good coding patterns. In fact, there are a lot of bad coding practices, overly complex code, code smells, etc. If something doesn't look right, speak up. Feel free to break existing patterns or question weird existing code to make new code high quality, as in:
|
||||
|
||||
|
|
@ -9,7 +17,7 @@ Don't assume that the existing code is correct or the right way of doing things
|
|||
- easy to maintain/change
|
||||
- modern
|
||||
|
||||
In that order of importance
|
||||
In descending order of importance
|
||||
|
||||
When adding new features, add meaningful tests. Don't add tests that don't check anything substantial and is there just to make the code coverage pass. Yes, code coverage is important, but I'd rather have no signal whether the code is working than tests that don't fail when code is broken. The goal is to have tests that would fail before the feature was added/if the code was mutated in a way that breaks the feature and succeed only when the feature is fully working. I should run mutation testing and see > 90% kill rate
|
||||
|
||||
|
|
@ -21,7 +29,9 @@ End-to-end tests belong in `tests/e2e/` and must follow the harness conventions
|
|||
|
||||
When creating PRs, don't set base to `main`. `litellm_internal_staging` is the default base branch and serves that purpose for both internal and external / OSS contributions
|
||||
|
||||
When writing a PR body, treat the comments and imperative instructions inside @.github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
|
||||
When writing a PR body, treat the comments and imperative instructions inside .github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule
|
||||
|
||||
Same applies for filing bug reports and feature requests, with .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml, respectively
|
||||
|
||||
If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank
|
||||
|
||||
|
|
@ -41,7 +51,7 @@ Python max line length is 120, not 88
|
|||
|
||||
When you fix violations gated by `ruff-strict-budget.json`, `type-discipline-budget.json`, or `basedpyright-code-budget.json`, run `make lint-budget-update` and commit the lowered limits so the ceilings ratchet down instead of leaving stale headroom. It measures the working tree, so it must contain exactly the fixes you're committing
|
||||
|
||||
`make pre-commit` saves its complete output to a log file in .git (overwriting previous pre-commit logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
|
||||
`make check` (f.k.a. `make pre-commit`, which still works identically as an alias) saves its complete output to a log file in .git (overwriting previous logs) and prints that path as its first and last output lines. To inspect a run, read or grep that log instead of re-running the multi-minute checks just to see a different slice
|
||||
|
||||
If you're trying to create a new function that relies on untyped stuff, instead of adding more Any's and pushing `reportAny` / `reportExplicitAny` closer to their basedpyright ceilings, just validate it in the caller with Pydantic (a model or `TypeAdapter` that returns the typed thing or raises will do) and then pass the now typed variable in
|
||||
|
||||
|
|
|
|||
18
Makefile
18
Makefile
|
|
@ -8,7 +8,7 @@
|
|||
lint-basedpyright lint-e2e-basedpyright lint-basedpyright-budget-update lint-type-discipline lint-type-discipline-budget-update \
|
||||
lint-ruff-budget lint-ruff-budget-update lint-budget-update lint-gate \
|
||||
install-dev install-proxy-dev install-test-deps install-hooks \
|
||||
install-helm-unittest check-circular-imports check-import-safety pre-commit \
|
||||
install-helm-unittest check-circular-imports check-import-safety check pre-commit \
|
||||
lint-install lint-fetch-base bootstrap
|
||||
|
||||
# Default target
|
||||
|
|
@ -22,7 +22,8 @@ help:
|
|||
@echo " make install-test-deps - Install the full local test environment"
|
||||
@echo " make install-helm-unittest - Install helm unittest plugin"
|
||||
@echo " make install-hooks - Install git hooks (Conventional Commits + Branches)"
|
||||
@echo " make pre-commit - Run CI-equivalent lint on staged files (run before committing)"
|
||||
@echo " make check - Run CI-equivalent lint on staged files, or on the diff vs the base branch when nothing is staged"
|
||||
@echo " make pre-commit - Legacy alias for make check"
|
||||
@echo " make format - Apply ruff format code formatting"
|
||||
@echo " make format-check - Check ruff format code formatting (matches CI)"
|
||||
@echo " make lint - Run all linting (Ruff, basedpyright, format check, circular imports, import safety)"
|
||||
|
|
@ -236,13 +237,20 @@ lint-checks: lint-format-check-changed lint-ruff lint-gate lint-type-discipline
|
|||
# Faster linting for local development (only checks changed code)
|
||||
lint-dev: lint-format-changed check-circular-imports check-import-safety
|
||||
|
||||
# Run the gating CI checks against your staged files right before committing. Mirrors
|
||||
# Run the gating CI checks against your changes. Scopes to staged files when anything
|
||||
# is staged (warning about changed files left unstaged); with nothing staged it falls
|
||||
# back to the working tree's diff against the merge base with the base branch, so a
|
||||
# fresh merge commit or an unstaged working tree still gets checked. Mirrors
|
||||
# test-linting.yml (Python), test-litellm-ui-build.yml's frontend-lint (dashboard), and
|
||||
# check-ui-api-types.yml (API-type drift), skipping any whose files you didn't stage.
|
||||
# check-ui-api-types.yml (API-type drift), skipping any whose files aren't in scope.
|
||||
# Not auto-installed as a git hook so it never slows an unrelated human commit.
|
||||
pre-commit: bootstrap
|
||||
check: bootstrap
|
||||
./scripts/pre_commit_lint.sh
|
||||
|
||||
pre-commit:
|
||||
@echo "make pre-commit is a legacy alias; use make check" >&2
|
||||
@$(MAKE) check
|
||||
|
||||
# Testing targets
|
||||
test: install-test-deps
|
||||
$(UV_RUN) pytest tests/
|
||||
|
|
|
|||
|
|
@ -1,30 +1,30 @@
|
|||
{
|
||||
"reportAny": {
|
||||
"limit": 28842
|
||||
"limit": 22947
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2634
|
||||
"limit": 2579
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"limit": 329
|
||||
"limit": 323
|
||||
},
|
||||
"reportAttributeAccessIssue": {
|
||||
"limit": 514
|
||||
"limit": 488
|
||||
},
|
||||
"reportCallIssue": {
|
||||
"limit": 117
|
||||
"limit": 114
|
||||
},
|
||||
"reportConstantRedefinition": {
|
||||
"limit": 40
|
||||
},
|
||||
"reportDeprecated": {
|
||||
"limit": 215
|
||||
"limit": 213
|
||||
},
|
||||
"reportDuplicateImport": {
|
||||
"limit": 19
|
||||
},
|
||||
"reportExplicitAny": {
|
||||
"limit": 9103
|
||||
"limit": 7312
|
||||
},
|
||||
"reportFunctionMemberAccess": {
|
||||
"limit": 7
|
||||
|
|
@ -54,10 +54,10 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportMissingParameterType": {
|
||||
"limit": 5843
|
||||
"limit": 5707
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15816
|
||||
"limit": 15642
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 40
|
||||
|
|
@ -72,7 +72,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportOptionalMemberAccess": {
|
||||
"limit": 1078
|
||||
"limit": 1069
|
||||
},
|
||||
"reportOptionalOperand": {
|
||||
"limit": 0
|
||||
|
|
@ -84,52 +84,52 @@
|
|||
"limit": 56
|
||||
},
|
||||
"reportPrivateUsage": {
|
||||
"limit": 1825
|
||||
"limit": 1824
|
||||
},
|
||||
"reportRedeclaration": {
|
||||
"limit": 8
|
||||
},
|
||||
"reportReturnType": {
|
||||
"limit": 218
|
||||
"limit": 213
|
||||
},
|
||||
"reportTypedDictNotRequiredAccess": {
|
||||
"limit": 27
|
||||
"limit": 26
|
||||
},
|
||||
"reportUndefinedVariable": {
|
||||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 45110
|
||||
"limit": 44776
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 113
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 39838
|
||||
"limit": 39237
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 20237
|
||||
"limit": 19969
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 31383
|
||||
"limit": 30881
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 122
|
||||
"limit": 117
|
||||
},
|
||||
"reportUnnecessaryComparison": {
|
||||
"limit": 701
|
||||
"limit": 699
|
||||
},
|
||||
"reportUnnecessaryContains": {
|
||||
"limit": 5
|
||||
},
|
||||
"reportUnnecessaryIsInstance": {
|
||||
"limit": 864
|
||||
"limit": 853
|
||||
},
|
||||
"reportUntypedBaseClass": {
|
||||
"limit": 0
|
||||
},
|
||||
"reportUntypedFunctionDecorator": {
|
||||
"limit": 33
|
||||
"limit": 27
|
||||
},
|
||||
"reportUnusedClass": {
|
||||
"limit": 23
|
||||
|
|
@ -138,7 +138,7 @@
|
|||
"limit": 139
|
||||
},
|
||||
"reportUnusedImport": {
|
||||
"limit": 555
|
||||
"limit": 545
|
||||
},
|
||||
"reportUnusedVariable": {
|
||||
"limit": 146
|
||||
|
|
|
|||
|
|
@ -99,6 +99,7 @@ class BaseEmailLogger(CustomLogger):
|
|||
email_html_content = USER_INVITATION_EMAIL_TEMPLATE.format(
|
||||
email_logo_url=email_params.logo_url,
|
||||
recipient_email=email_params.recipient_email,
|
||||
invitation_link=email_params.base_url,
|
||||
base_url=email_params.base_url,
|
||||
email_support_contact=email_params.support_contact,
|
||||
email_footer=email_params.signature,
|
||||
|
|
@ -826,10 +827,15 @@ class BaseEmailLogger(CustomLogger):
|
|||
"""
|
||||
# Early validation
|
||||
if not user_id:
|
||||
verbose_proxy_logger.debug("No user_id provided for invitation link")
|
||||
verbose_proxy_logger.warning(
|
||||
"No user_id provided for invitation link. Email will link to base URL instead of onboarding page"
|
||||
)
|
||||
return base_url
|
||||
|
||||
if not await self._is_prisma_client_available():
|
||||
verbose_proxy_logger.warning(
|
||||
"Prisma client not available. Email will link to base URL instead of onboarding page"
|
||||
)
|
||||
return base_url
|
||||
|
||||
# Wait for any concurrent invitation creation to complete
|
||||
|
|
@ -839,11 +845,15 @@ class BaseEmailLogger(CustomLogger):
|
|||
invitation = await self._get_or_create_invitation(user_id)
|
||||
if not invitation:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Failed to get/create invitation for user_id: {user_id}"
|
||||
f"Failed to get/create invitation for user_id: {user_id}. Email will link to base URL instead of onboarding page"
|
||||
)
|
||||
return base_url
|
||||
|
||||
return self._construct_invitation_link(invitation.id, base_url)
|
||||
invitation_link = self._construct_invitation_link(invitation.id, base_url)
|
||||
verbose_proxy_logger.info(
|
||||
f"Successfully created invitation link for user_id: {user_id}"
|
||||
)
|
||||
return invitation_link
|
||||
|
||||
async def _is_prisma_client_available(self) -> bool:
|
||||
"""Check if Prisma client is available"""
|
||||
|
|
@ -921,7 +931,9 @@ class BaseEmailLogger(CustomLogger):
|
|||
|
||||
# http://localhost:4000/ui/onboarding?invitation_id=7a096b3a-37c6-440f-9dd1-ba22e8043f6b
|
||||
"""
|
||||
return f"{base_url}/ui/onboarding?invitation_id={invitation_id}"
|
||||
base_url = base_url.rstrip("/")
|
||||
invitation_link = f"{base_url}/ui/onboarding?invitation_id={invitation_id}"
|
||||
return invitation_link
|
||||
|
||||
async def send_email(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
|
|||
"""
|
||||
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, List, Optional, Tuple
|
||||
from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Tuple
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -23,6 +23,15 @@ if TYPE_CHECKING:
|
|||
|
||||
CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
|
||||
|
||||
TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
|
||||
"completed",
|
||||
"complete",
|
||||
"failed",
|
||||
"expired",
|
||||
"cancelled",
|
||||
"stale_expired",
|
||||
)
|
||||
|
||||
|
||||
class CheckBatchCost:
|
||||
def __init__(
|
||||
|
|
@ -43,11 +52,15 @@ class CheckBatchCost:
|
|||
# the guaranteed-failing primary query on every subsequent cycle.
|
||||
self._has_batch_processed_column: bool = True
|
||||
|
||||
async def _get_user_info(self, batch_id, user_id) -> dict:
|
||||
async def _get_user_info(self, batch_id: str, user_id: Optional[str]) -> Dict[str, Any]:
|
||||
"""
|
||||
Look up user email and key alias by user_id for enriching the S3 callback metadata.
|
||||
Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None).
|
||||
Returns an empty dict when user_id is None: batches created by a team or service
|
||||
account key carry no user id, and find_unique(where={"user_id": None}) raises.
|
||||
"""
|
||||
if not user_id:
|
||||
return {}
|
||||
try:
|
||||
user_row = await self.prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": user_id}
|
||||
|
|
@ -62,17 +75,77 @@ class CheckBatchCost:
|
|||
verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
|
||||
return {}
|
||||
|
||||
async def _get_key_alias(self, batch_id: str, api_key: str | None) -> str | None:
|
||||
"""Resolve the creating virtual key's alias from its hashed token."""
|
||||
if not api_key:
|
||||
return None
|
||||
try:
|
||||
key_row = await self.prisma_client.db.litellm_verificationtoken.find_unique(
|
||||
where={"token": api_key}
|
||||
)
|
||||
return getattr(key_row, "key_alias", None) if key_row is not None else None
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}")
|
||||
return None
|
||||
|
||||
async def _get_team_alias(self, team_id: str | None) -> str | None:
|
||||
"""Resolve a team's alias from its id."""
|
||||
if not team_id:
|
||||
return None
|
||||
try:
|
||||
team_row = await self.prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": team_id}
|
||||
)
|
||||
return getattr(team_row, "team_alias", None) if team_row is not None else None
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}")
|
||||
return None
|
||||
|
||||
async def _build_creator_attribution_metadata(
|
||||
self, job: "LiteLLM_ManagedObjectTable", batch_id: str
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Rebuild the spend-tracking metadata for the key, team, and tags that created the
|
||||
batch so the batch-cost spend log is attributed the same way a non-batch request
|
||||
is. Rows created before api_key and request_tags were persisted carry only
|
||||
created_by and team_id, and fall back to those. A named creating key owns
|
||||
user_api_key_alias; when it has no alias, or the key has since been rotated or
|
||||
deleted, the field keeps the creating user's alias that _get_user_info filled in,
|
||||
because a resolvable name is more useful on the spend row than a null.
|
||||
"""
|
||||
api_key = getattr(job, "api_key", None)
|
||||
team_id = getattr(job, "team_id", None)
|
||||
request_tags = getattr(job, "request_tags", None)
|
||||
|
||||
metadata: Dict[str, Any] = {
|
||||
"user_api_key_user_id": job.created_by,
|
||||
"user_api_key": api_key,
|
||||
"user_api_key_team_id": team_id,
|
||||
**(await self._get_user_info(batch_id, job.created_by)),
|
||||
}
|
||||
|
||||
key_alias = await self._get_key_alias(batch_id, api_key)
|
||||
if key_alias is not None:
|
||||
metadata["user_api_key_alias"] = key_alias
|
||||
team_alias = await self._get_team_alias(team_id)
|
||||
if team_alias is not None:
|
||||
metadata["user_api_key_team_alias"] = team_alias
|
||||
if isinstance(request_tags, list) and request_tags:
|
||||
metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)]
|
||||
|
||||
return metadata
|
||||
|
||||
async def _cleanup_stale_managed_objects(self) -> None:
|
||||
"""
|
||||
Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days
|
||||
in non-terminal states as 'stale_expired'. These will never complete and
|
||||
should not be polled.
|
||||
"""
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
|
||||
result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
cutoff: Final = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
|
||||
result: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
where={
|
||||
"file_purpose": "batch",
|
||||
"status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
|
||||
"status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)},
|
||||
"created_at": {"lt": cutoff},
|
||||
},
|
||||
data={"status": "stale_expired"},
|
||||
|
|
@ -83,6 +156,26 @@ class CheckBatchCost:
|
|||
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
|
||||
)
|
||||
|
||||
if not self._has_batch_processed_column:
|
||||
return
|
||||
|
||||
# A row already in a terminal status is never rewritten by the sweep above, so
|
||||
# without this it keeps a poll-page slot forever and starves newer batches.
|
||||
retired: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
where={
|
||||
"file_purpose": "batch",
|
||||
"batch_processed": False,
|
||||
"status": {"in": ["complete", "completed"]},
|
||||
"created_at": {"lt": cutoff},
|
||||
},
|
||||
data={"batch_processed": True},
|
||||
)
|
||||
if retired > 0:
|
||||
verbose_proxy_logger.warning(
|
||||
f"CheckBatchCost: gave up on {retired} completed managed objects older than "
|
||||
f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
|
||||
)
|
||||
|
||||
async def _fallback_find_jobs(self) -> list:
|
||||
"""Query batch jobs without the batch_processed filter (for older schemas)."""
|
||||
return await self.prisma_client.db.litellm_managedobjecttable.find_many(
|
||||
|
|
@ -103,6 +196,68 @@ class CheckBatchCost:
|
|||
order={"created_at": "asc"},
|
||||
)
|
||||
|
||||
async def _retire_job(self, job: "LiteLLM_ManagedObjectTable", reason: str) -> None:
|
||||
"""
|
||||
Take a row that can never be costed out of the poll page. Leaving it selectable
|
||||
would burn one of the MAX_OBJECTS_PER_POLL_CYCLE slots on every future cycle, and
|
||||
once enough such rows accumulate no newer batch is ever reached. Older schemas
|
||||
without batch_processed can only be excluded through the status filter.
|
||||
"""
|
||||
data: Final = (
|
||||
{"batch_processed": True}
|
||||
if self._has_batch_processed_column
|
||||
else {"status": "stale_expired"}
|
||||
)
|
||||
try:
|
||||
await self.prisma_client.db.litellm_managedobjecttable.update(
|
||||
where={"id": job.id},
|
||||
data=data,
|
||||
)
|
||||
except Exception as db_err:
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}"
|
||||
)
|
||||
return
|
||||
verbose_proxy_logger.warning(
|
||||
f"CheckBatchCost: job {job.id} can never be costed ({reason}), "
|
||||
"so it will no longer be polled"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _has_unified_id_without_model(job: "LiteLLM_ManagedObjectTable") -> bool:
|
||||
"""A unified id that decodes but carries no model_id can never be routed."""
|
||||
from litellm.proxy.openai_files_endpoints.common_utils import (
|
||||
convert_b64_uid_to_unified_uid,
|
||||
get_model_id_from_unified_batch_id,
|
||||
)
|
||||
|
||||
decoded: Final = convert_b64_uid_to_unified_uid(job.unified_object_id)
|
||||
return (
|
||||
decoded != job.unified_object_id
|
||||
and get_model_id_from_unified_batch_id(decoded) is None
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_batch_gone_at_provider(error: Exception, batch_id: str) -> bool:
|
||||
"""
|
||||
A 404 naming the batch means the provider dropped its record of it, so no later
|
||||
retrieve can ever succeed. A 404 about anything else, a renamed Azure deployment
|
||||
or a fallback deployment that never saw this batch, is still fixable in config, so
|
||||
it keeps retrying.
|
||||
"""
|
||||
import openai
|
||||
|
||||
from litellm.exceptions import NotFoundError
|
||||
|
||||
return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error)
|
||||
|
||||
def _batch_deployment_exists(self, model_id: str) -> bool:
|
||||
"""A 404 only proves the batch is gone when it came from the batch's own
|
||||
deployment. Once that deployment leaves the router, default fallbacks can
|
||||
silently send the retrieve to a provider that never saw the batch, so its
|
||||
404 must not retire the row; the staleness sweep bounds it instead."""
|
||||
return self.llm_router.get_deployment(model_id=model_id) is not None
|
||||
|
||||
@staticmethod
|
||||
def _record_error(
|
||||
prom_logger: Optional["PrometheusLogger"], error_type: str
|
||||
|
|
@ -485,9 +640,6 @@ class CheckBatchCost:
|
|||
function_id=str(uuid.uuid4()),
|
||||
)
|
||||
|
||||
creator_user_id = job.created_by
|
||||
user_info = await self._get_user_info(batch_id, job.created_by)
|
||||
|
||||
logging_obj.update_environment_variables(
|
||||
litellm_params={
|
||||
# set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks
|
||||
|
|
@ -496,11 +648,7 @@ class CheckBatchCost:
|
|||
"user-agent": CHECK_BATCH_COST_USER_AGENT,
|
||||
}
|
||||
},
|
||||
"metadata": {
|
||||
"user_api_key_user_id": creator_user_id,
|
||||
"user_api_key_team_id": getattr(job, "team_id", None),
|
||||
**user_info,
|
||||
},
|
||||
"metadata": await self._build_creator_attribution_metadata(job, batch_id),
|
||||
},
|
||||
optional_params={},
|
||||
)
|
||||
|
|
@ -588,6 +736,8 @@ class CheckBatchCost:
|
|||
for job in jobs:
|
||||
routing = self._resolve_job_routing(job, prom_logger)
|
||||
if routing is None:
|
||||
if self._has_unified_id_without_model(job):
|
||||
await self._retire_job(job, "unified object id has no model id")
|
||||
continue
|
||||
model_id, batch_id = routing
|
||||
|
||||
|
|
@ -610,6 +760,8 @@ class CheckBatchCost:
|
|||
)
|
||||
if prom_logger:
|
||||
prom_logger.record_check_batch_cost_error("provider_retrieval_error")
|
||||
if self._is_batch_gone_at_provider(e, batch_id) and self._batch_deployment_exists(model_id):
|
||||
await self._retire_job(job, f"batch {batch_id} no longer exists at the provider")
|
||||
continue
|
||||
|
||||
## RETRIEVE THE BATCH JOB OUTPUT FILE
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -19,6 +19,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.auth.auth_checks import delete_cached_project_object
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field
|
||||
from litellm.proxy.management_helpers.utils import (
|
||||
|
|
@ -514,6 +515,7 @@ async def update_project(
|
|||
litellm_proxy_admin_name,
|
||||
premium_user,
|
||||
prisma_client,
|
||||
user_api_key_cache,
|
||||
)
|
||||
|
||||
try:
|
||||
|
|
@ -672,6 +674,11 @@ async def update_project(
|
|||
include={"litellm_budget_table": True, "object_permission": True},
|
||||
)
|
||||
|
||||
await delete_cached_project_object(
|
||||
project_id=data.project_id,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
return updated_project
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
|
|
@ -710,7 +717,7 @@ async def delete_project(
|
|||
}'
|
||||
```
|
||||
"""
|
||||
from litellm.proxy.proxy_server import premium_user, prisma_client
|
||||
from litellm.proxy.proxy_server import premium_user, prisma_client, user_api_key_cache
|
||||
|
||||
try:
|
||||
if not premium_user:
|
||||
|
|
@ -773,6 +780,11 @@ async def delete_project(
|
|||
prisma_models.LiteLLM_ProjectTable | None
|
||||
) = await prisma_client.db.litellm_projecttable.delete(where={"project_id": project_id})
|
||||
|
||||
await delete_cached_project_object(
|
||||
project_id=project_id,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
deleted_projects.append(deleted_project)
|
||||
|
||||
return deleted_projects
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.54"
|
||||
version = "0.1.55"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.54"
|
||||
version = "0.1.55"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -105,6 +105,10 @@ spec:
|
|||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
restartPolicy: OnFailure
|
||||
{{- with .Values.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.affinity }}
|
||||
affinity:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
|
|
|
|||
|
|
@ -290,3 +290,27 @@ tests:
|
|||
value:
|
||||
allowPrivilegeEscalation: false
|
||||
readOnlyRootFilesystem: true
|
||||
- it: should schedule onto the same nodes as the gateway
|
||||
template: migrations-job.yaml
|
||||
set:
|
||||
migrationJob:
|
||||
enabled: true
|
||||
nodeSelector:
|
||||
karpenter.sh/nodepool: litellm-e2e
|
||||
tolerations:
|
||||
- key: workload
|
||||
operator: Equal
|
||||
value: litellm-e2e
|
||||
effect: NoSchedule
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.nodeSelector
|
||||
value:
|
||||
karpenter.sh/nodepool: litellm-e2e
|
||||
- equal:
|
||||
path: spec.template.spec.tolerations
|
||||
value:
|
||||
- key: workload
|
||||
operator: Equal
|
||||
value: litellm-e2e
|
||||
effect: NoSchedule
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DailyTeamSpend" ADD COLUMN IF NOT EXISTS "ptu_flat_cost" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
-- Add api_key and request_tags columns to LiteLLM_ManagedObjectTable
|
||||
-- Captured at batch-create time so CheckBatchCost can attribute batch-cost spend
|
||||
-- back to the creating virtual key (and its tags) even when created_by is null.
|
||||
ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "api_key" TEXT;
|
||||
ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN IF NOT EXISTS "request_tags" JSONB DEFAULT '[]';
|
||||
|
|
@ -0,0 +1 @@
|
|||
ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "tier_turns" JSONB NOT NULL DEFAULT '{}';
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "settings_updated_at" TIMESTAMP(3);
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "settings_updated_at" TIMESTAMP(3);
|
||||
|
|
@ -0,0 +1,49 @@
|
|||
-- CreateTable
|
||||
CREATE TABLE "LiteLLM_ShadowEvalJob" (
|
||||
"id" TEXT NOT NULL,
|
||||
"api_key_id" TEXT NOT NULL,
|
||||
"router_name" TEXT NOT NULL,
|
||||
"judge_model" TEXT NOT NULL,
|
||||
"shadow_percentage" DOUBLE PRECISION NOT NULL,
|
||||
"max_turns" INTEGER NOT NULL,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
"created_by" TEXT,
|
||||
"ends_at" TIMESTAMP(3) NOT NULL,
|
||||
"stopped_at" TIMESTAMP(3),
|
||||
|
||||
CONSTRAINT "LiteLLM_ShadowEvalJob_pkey" PRIMARY KEY ("id")
|
||||
);
|
||||
|
||||
-- CreateTable
|
||||
CREATE TABLE "LiteLLM_ShadowEvalAttempt" (
|
||||
"id" TEXT NOT NULL,
|
||||
"job_id" TEXT NOT NULL,
|
||||
"request_id" TEXT NOT NULL,
|
||||
"outcome" TEXT NOT NULL,
|
||||
"tier" TEXT,
|
||||
"real_model" TEXT,
|
||||
"shadow_model" TEXT,
|
||||
"confidence" DOUBLE PRECISION,
|
||||
"judge_cost" DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
"error" TEXT,
|
||||
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
|
||||
CONSTRAINT "LiteLLM_ShadowEvalAttempt_pkey" PRIMARY KEY ("id")
|
||||
);
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_ShadowEvalJob_api_key_id_idx" ON "LiteLLM_ShadowEvalJob"("api_key_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_ShadowEvalJob_created_at_idx" ON "LiteLLM_ShadowEvalJob"("created_at");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX "LiteLLM_ShadowEvalAttempt_job_id_idx" ON "LiteLLM_ShadowEvalAttempt"("job_id");
|
||||
|
||||
|
||||
-- One active job per key, enforced by the database rather than a read-then-create in the
|
||||
-- start endpoint, which races against a concurrent start on another pod. Partial indexes
|
||||
-- are not expressible in schema.prisma, so this lives here only. Active means not yet
|
||||
-- stopped; the start endpoint stamps stopped_at on expired jobs before creating.
|
||||
CREATE UNIQUE INDEX "LiteLLM_ShadowEvalJob_one_active_per_key"
|
||||
ON "LiteLLM_ShadowEvalJob"("api_key_id") WHERE "stopped_at" IS NULL;
|
||||
|
|
@ -30,7 +30,7 @@ model LiteLLM_BudgetTable {
|
|||
end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget
|
||||
tags LiteLLM_TagTable[] // multiple tags can have the same budget
|
||||
team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team
|
||||
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
|
||||
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
|
||||
}
|
||||
|
||||
// Models on proxy
|
||||
|
|
@ -452,6 +452,7 @@ model LiteLLM_VerificationToken {
|
|||
created_by String?
|
||||
updated_at DateTime? @default(now()) @updatedAt @map("updated_at")
|
||||
updated_by String?
|
||||
settings_updated_at DateTime? @map("settings_updated_at")
|
||||
last_active DateTime? // When this key was last used
|
||||
rotation_count Int? @default(0) // Number of times key has been rotated
|
||||
auto_rotate Boolean? @default(false) // Whether this key should be auto-rotated
|
||||
|
|
@ -548,6 +549,7 @@ model LiteLLM_DeletedVerificationToken {
|
|||
created_by String? // Original creator
|
||||
updated_at DateTime? // Last update timestamp before deletion
|
||||
updated_by String? // Last user who updated before deletion
|
||||
settings_updated_at DateTime? // Last configuration change before deletion
|
||||
last_active DateTime? // When this key was last used before deletion
|
||||
rotation_count Int? @default(0)
|
||||
auto_rotate Boolean? @default(false)
|
||||
|
|
@ -893,6 +895,7 @@ model LiteLLM_DailyTeamSpend {
|
|||
api_requests BigInt @default(0)
|
||||
successful_requests BigInt @default(0)
|
||||
failed_requests BigInt @default(0)
|
||||
ptu_flat_cost Float @default(0.0)
|
||||
created_at DateTime @default(now())
|
||||
updated_at DateTime @updatedAt
|
||||
|
||||
|
|
@ -985,6 +988,8 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
|
|||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
team_id String?
|
||||
api_key String?
|
||||
request_tags Json? @default("[]")
|
||||
updated_at DateTime @updatedAt
|
||||
updated_by String?
|
||||
|
||||
|
|
@ -1439,11 +1444,50 @@ model LiteLLM_AutoRouterSession {
|
|||
total_tokens BigInt @default(0)
|
||||
spend Float @default(0)
|
||||
saved_spend Float @default(0)
|
||||
tier_turns Json @default("{}")
|
||||
|
||||
@@id([api_key, session_id, router_name])
|
||||
@@index([last_turn_at], map: "idx_autorouter_session_last_turn")
|
||||
}
|
||||
|
||||
// Shadow eval: pre-adoption evaluation of an auto-router against a key's live traffic.
|
||||
// A sampled slice of requests is duplicated through the router in a detached task and an
|
||||
// LLM judge compares real vs shadow responses blind. The job row is immutable config plus
|
||||
// stopped_at; every count, status, and spend figure is derived from the append-only
|
||||
// attempt rows, so nothing can disagree across pods or stop races.
|
||||
model LiteLLM_ShadowEvalJob {
|
||||
id String @id @default(cuid())
|
||||
api_key_id String // hashed virtual key whose traffic is shadowed
|
||||
router_name String
|
||||
judge_model String
|
||||
shadow_percentage Float
|
||||
max_turns Int // sample budget: judge at most this many turns
|
||||
created_at DateTime @default(now())
|
||||
created_by String?
|
||||
ends_at DateTime
|
||||
stopped_at DateTime?
|
||||
|
||||
@@index([api_key_id])
|
||||
@@index([created_at])
|
||||
}
|
||||
|
||||
// One row per sampled pipeline: a blind verdict (real | shadow | tie) or an error.
|
||||
model LiteLLM_ShadowEvalAttempt {
|
||||
id String @id @default(cuid())
|
||||
job_id String
|
||||
request_id String // the judged real request
|
||||
outcome String // real | shadow | tie | error
|
||||
tier String? // router's tier for the prompt, when classified
|
||||
real_model String?
|
||||
shadow_model String?
|
||||
confidence Float?
|
||||
judge_cost Float @default(0)
|
||||
error String?
|
||||
created_at DateTime @default(now())
|
||||
|
||||
@@index([job_id])
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Workflow Run Tracking
|
||||
//
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-proxy-extras"
|
||||
version = "0.4.84"
|
||||
version = "0.4.85"
|
||||
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.4.84"
|
||||
version = "0.4.85"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-proxy-extras==",
|
||||
|
|
|
|||
|
|
@ -197,6 +197,7 @@ standard_logging_payload_excluded_fields: Optional[List[str]] = (
|
|||
None # Fields to exclude from StandardLoggingPayload before callbacks receive it
|
||||
)
|
||||
log_raw_request_response: bool = False
|
||||
request_correlation_in_logs: bool = False
|
||||
redact_messages_in_exceptions: Optional[bool] = False
|
||||
redact_user_api_key_info: Optional[bool] = False
|
||||
# When True (default — preserves historical behavior), the Router appends
|
||||
|
|
@ -245,6 +246,7 @@ use_chat_completions_url_for_anthropic_messages: bool = bool(
|
|||
# config.yaml.
|
||||
strip_anthropic_total_tokens: bool = False
|
||||
anthropic_sse_ping_interval_seconds: float = 15.0
|
||||
sse_keepalive_ping_interval_seconds: float | None = None
|
||||
route_all_chat_openai_to_responses: bool = (
|
||||
os.getenv("LITELLM_ROUTE_ALL_CHAT_OPENAI_TO_RESPONSES", "false").lower() == "true"
|
||||
) # When True, routes all OpenAI /chat/completions requests through the Responses API bridge
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import ast
|
||||
import contextvars
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
|
@ -6,12 +7,44 @@ from datetime import datetime
|
|||
from logging import Formatter
|
||||
from typing import Any, Final
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.litellm_core_utils.secret_redaction import redact_string
|
||||
|
||||
set_verbose = False
|
||||
|
||||
session_id_var: Final[contextvars.ContextVar[str]] = contextvars.ContextVar("session_id", default="")
|
||||
trace_id_var: Final[contextvars.ContextVar[str]] = contextvars.ContextVar("trace_id", default="")
|
||||
|
||||
_MAX_CORRELATION_ID_LENGTH: Final = 256
|
||||
|
||||
|
||||
def _sanitize_correlation_id(value: str) -> str:
|
||||
"""Strip control characters, bound length, and redact credential-shaped
|
||||
content before a caller-controlled trace_id/session_id (e.g.
|
||||
litellm_session_id, x-litellm-trace-id) is stamped into log lines.
|
||||
|
||||
Without the first two, a caller could embed \\r/\\n or terminal escape
|
||||
sequences to forge fake log entries, or submit an oversized value repeated
|
||||
across every log line for the request. Without the redaction, a caller
|
||||
could smuggle a real credential (e.g. an sk-... key) through this field:
|
||||
CorrelationContextFilter stamps trace_id/session_id onto the record after
|
||||
SecretRedactionFilter has already run, so those two fields never otherwise
|
||||
pass through credential redaction.
|
||||
"""
|
||||
stripped: Final = "".join(ch for ch in value if ch.isprintable())
|
||||
return _redact_string(stripped[:_MAX_CORRELATION_ID_LENGTH])
|
||||
|
||||
|
||||
def set_session_id(session_id: str) -> "contextvars.Token[str]":
|
||||
return session_id_var.set(_sanitize_correlation_id(session_id))
|
||||
|
||||
|
||||
def set_trace_id(trace_id: str) -> "contextvars.Token[str]":
|
||||
return trace_id_var.set(_sanitize_correlation_id(trace_id))
|
||||
|
||||
|
||||
if set_verbose is True:
|
||||
logging.warning(
|
||||
"`litellm.set_verbose` is deprecated. Please set `os.environ['LITELLM_LOG'] = 'DEBUG'` for debug logs."
|
||||
|
|
@ -77,6 +110,28 @@ class SecretRedactionFilter(logging.Filter):
|
|||
_secret_filter: Final = SecretRedactionFilter()
|
||||
|
||||
|
||||
class CorrelationContextFilter(logging.Filter):
|
||||
"""Stamps each log record with the current request's trace_id and session_id from contextvars.
|
||||
|
||||
Works in tandem with JsonFormatter: the formatter's record.__dict__ loop picks up these
|
||||
attributes as first-class JSON fields without any formatter-level code.
|
||||
"""
|
||||
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
if not litellm.request_correlation_in_logs:
|
||||
return True
|
||||
trace_id: Final = trace_id_var.get()
|
||||
if trace_id:
|
||||
record.trace_id = trace_id # rebind-ok: stamping the LogRecord is the Filter interface's contract
|
||||
session_id: Final = session_id_var.get()
|
||||
if session_id:
|
||||
record.session_id = session_id # rebind-ok: stamping the LogRecord is the Filter interface's contract
|
||||
return True
|
||||
|
||||
|
||||
_correlation_filter: Final = CorrelationContextFilter()
|
||||
|
||||
|
||||
json_logs = bool(os.getenv("JSON_LOGS", False))
|
||||
# Create a handler for the logger (you may need to adapt this based on your needs)
|
||||
log_level: Final = os.getenv("LITELLM_LOG", "DEBUG")
|
||||
|
|
@ -84,6 +139,7 @@ numeric_level: Final[str] = getattr(logging, log_level.upper())
|
|||
handler: Final = logging.StreamHandler()
|
||||
handler.setLevel(numeric_level)
|
||||
handler.addFilter(_secret_filter)
|
||||
handler.addFilter(_correlation_filter)
|
||||
|
||||
|
||||
def _try_parse_json_message(message: str) -> dict[str, Any] | None:
|
||||
|
|
@ -146,6 +202,11 @@ def _get_standard_record_attrs() -> frozenset:
|
|||
|
||||
_STANDARD_RECORD_ATTRS: Final = _get_standard_record_attrs()
|
||||
|
||||
# CorrelationContextFilter is the only legitimate source for these two JSON fields;
|
||||
# see JsonFormatter.format() for why they're excluded from the generic message-content
|
||||
# and extra-attribute promotion paths.
|
||||
_RESERVED_CORRELATION_FIELDS: Final = frozenset(("trace_id", "session_id"))
|
||||
|
||||
|
||||
class JsonFormatter(Formatter):
|
||||
def __init__(self):
|
||||
|
|
@ -164,13 +225,18 @@ class JsonFormatter(Formatter):
|
|||
"timestamp": self.formatTime(record),
|
||||
}
|
||||
|
||||
# Parse embedded JSON or Python dict repr in message so sub-fields become first-class properties
|
||||
# Parse embedded JSON or Python dict repr in message so sub-fields become first-class properties.
|
||||
# trace_id/session_id are excluded here unconditionally (not just "if not already
|
||||
# set") - CorrelationContextFilter is the only legitimate source for these two
|
||||
# fields, and a message that merely happens to parse as JSON/dict (e.g. a proxy
|
||||
# log line dumping raw request headers) must never be able to claim them, even on
|
||||
# a record the filter hasn't stamped yet (no correlation context active for it).
|
||||
parsed = _try_parse_json_message(message_str)
|
||||
if parsed is None:
|
||||
parsed = _try_parse_embedded_python_dict(message_str)
|
||||
if parsed is not None:
|
||||
for key, value in parsed.items():
|
||||
if key not in json_record:
|
||||
if key not in json_record and key not in _RESERVED_CORRELATION_FIELDS:
|
||||
json_record[key] = value
|
||||
|
||||
# Include extra attributes passed via logger.debug("msg", extra={...})
|
||||
|
|
@ -178,6 +244,18 @@ class JsonFormatter(Formatter):
|
|||
if key not in _STANDARD_RECORD_ATTRS and key not in json_record:
|
||||
json_record[key] = value
|
||||
|
||||
# trace_id/session_id are reserved: CorrelationContextFilter is the only
|
||||
# legitimate source for these two fields. Without this, a message string
|
||||
# that happens to parse as JSON/dict (e.g. a proxy log line dumping raw
|
||||
# request headers) with a "trace_id"/"session_id" key would have already
|
||||
# claimed the key at the parsed-message step above, and the extra-attributes
|
||||
# loop's "key not in json_record" guard would then skip the real value -
|
||||
# letting a caller-supplied header spoof another request's correlation ids.
|
||||
for reserved_key in _RESERVED_CORRELATION_FIELDS:
|
||||
value = getattr(record, reserved_key, None)
|
||||
if value:
|
||||
json_record[reserved_key] = value
|
||||
|
||||
# Set component/logger only if not already supplied via extra={...}
|
||||
if "component" not in json_record:
|
||||
json_record["component"] = record.name
|
||||
|
|
@ -190,12 +268,34 @@ class JsonFormatter(Formatter):
|
|||
return safe_dumps(json_record)
|
||||
|
||||
|
||||
class CorrelationPlainFormatter(logging.Formatter):
|
||||
"""Appends trace_id/session_id to plain-text log lines stamped by CorrelationContextFilter.
|
||||
|
||||
Mirrors JsonFormatter's handling of these two fields so request_correlation_in_logs
|
||||
behaves the same whether or not json_logs is enabled.
|
||||
"""
|
||||
|
||||
def format(self, record: logging.LogRecord) -> str:
|
||||
formatted: Final = super().format(record)
|
||||
trace_id: Final = getattr(record, "trace_id", None)
|
||||
session_id: Final = getattr(record, "session_id", None)
|
||||
if not trace_id and not session_id:
|
||||
return formatted
|
||||
parts: Final = tuple(
|
||||
p
|
||||
for p in (f"trace_id={trace_id}" if trace_id else None, f"session_id={session_id}" if session_id else None)
|
||||
if p
|
||||
)
|
||||
return f"{formatted} [{' '.join(parts)}]"
|
||||
|
||||
|
||||
# Function to set up exception handlers for JSON logging
|
||||
def _setup_json_exception_handlers(formatter):
|
||||
# Create a handler with JSON formatting for exceptions
|
||||
error_handler: Final = logging.StreamHandler()
|
||||
error_handler.setFormatter(formatter)
|
||||
error_handler.addFilter(_secret_filter)
|
||||
error_handler.addFilter(_correlation_filter)
|
||||
|
||||
# Setup excepthook for uncaught exceptions
|
||||
def json_excepthook(exc_type, exc_value, exc_traceback):
|
||||
|
|
@ -243,7 +343,7 @@ if json_logs:
|
|||
handler.setFormatter(JsonFormatter())
|
||||
_setup_json_exception_handlers(JsonFormatter())
|
||||
else:
|
||||
formatter: Final = logging.Formatter(
|
||||
formatter: Final = CorrelationPlainFormatter(
|
||||
"\033[92m%(asctime)s - %(name)s:%(levelname)s\033[0m: %(filename)s:%(lineno)s - %(message)s",
|
||||
datefmt="%H:%M:%S",
|
||||
)
|
||||
|
|
@ -346,6 +446,7 @@ def _initialize_loggers_with_handler(handler: logging.Handler):
|
|||
- Prevents bubbling to parent/root (critical to prevent duplicate JSON logs)
|
||||
"""
|
||||
handler.addFilter(_secret_filter)
|
||||
handler.addFilter(_correlation_filter)
|
||||
for lg in _get_loggers_to_initialize():
|
||||
lg.handlers.clear() # remove any existing handlers
|
||||
lg.addHandler(handler) # add JSON formatter handler
|
||||
|
|
|
|||
|
|
@ -202,9 +202,9 @@ async def handle_a2a_localhost_retry(
|
|||
# Fix the agent card URL
|
||||
set_agent_card_url(agent_card, error.base_url)
|
||||
|
||||
# Reuse the httpx client LiteLLM attached at creation. It carries this agent's
|
||||
# trace-id and auth headers, so a fresh client would drop them. Only clients built
|
||||
# by ``create_a2a_client`` have it; an externally-supplied client cannot be retried.
|
||||
# Reuse the httpx client and call context LiteLLM attached at creation, since the
|
||||
# context carries this agent's trace-id/auth headers. Only clients built by
|
||||
# ``create_a2a_client`` have them; an externally-supplied client cannot be retried.
|
||||
httpx_client: Final = getattr(a2a_client, "_litellm_httpx_client", None)
|
||||
if httpx_client is None:
|
||||
raise RuntimeError(
|
||||
|
|
@ -220,5 +220,8 @@ async def handle_a2a_localhost_retry(
|
|||
),
|
||||
)
|
||||
new_client._litellm_httpx_client = httpx_client
|
||||
new_client._litellm_call_context = getattr( # pyright: ignore[reportAttributeAccessIssue] # LiteLLM-owned stash
|
||||
a2a_client, "_litellm_call_context", None
|
||||
)
|
||||
new_client._litellm_agent_card = agent_card
|
||||
return new_client
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ A2A Streaming Events (in order):
|
|||
4. Status update (kind: "status-update") - Final status "completed" with final=true
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from collections.abc import AsyncIterator, Callable, Coroutine, Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
import litellm
|
||||
|
|
@ -54,7 +54,7 @@ class A2ACompletionBridgeHandler:
|
|||
agent_extra_headers: Mapping[str, str] | None,
|
||||
*,
|
||||
stream: bool,
|
||||
) -> Mapping[str, Any]:
|
||||
) -> Mapping[str, object]:
|
||||
# Extract message from params
|
||||
message: Final = params.get("message", {})
|
||||
|
||||
|
|
@ -63,7 +63,7 @@ class A2ACompletionBridgeHandler:
|
|||
|
||||
# Get completion params
|
||||
custom_llm_provider: Final = litellm_params.get("custom_llm_provider")
|
||||
model: Final = litellm_params.get("model", "agent")
|
||||
model: Final[str] = litellm_params.get("model", "agent")
|
||||
|
||||
# Build full model string if provider specified
|
||||
# Skip prepending if model already starts with the provider prefix
|
||||
|
|
@ -109,13 +109,16 @@ class A2ACompletionBridgeHandler:
|
|||
return completion_params
|
||||
|
||||
@staticmethod
|
||||
async def _acompletion(completion_params: Mapping[str, Any]) -> ModelResponse | CustomStreamWrapper:
|
||||
return await litellm.acompletion(**completion_params)
|
||||
async def _acompletion(completion_params: Mapping[str, object]) -> ModelResponse | CustomStreamWrapper:
|
||||
acompletion_fn: Final[Callable[..., Coroutine[object, object, ModelResponse | CustomStreamWrapper]]] = vars(
|
||||
litellm
|
||||
)["acompletion"]
|
||||
return await acompletion_fn(**completion_params)
|
||||
|
||||
@staticmethod
|
||||
async def handle_non_streaming(
|
||||
request_id: str,
|
||||
params: dict[str, Any],
|
||||
params: dict[str, object],
|
||||
litellm_params: dict[str, Any],
|
||||
api_base: str | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
|
|
@ -296,8 +299,8 @@ class A2ACompletionBridgeHandler:
|
|||
# Convenience functions that delegate to the class methods
|
||||
async def handle_a2a_completion(
|
||||
request_id: str,
|
||||
params: dict[str, Any],
|
||||
litellm_params: dict[str, Any],
|
||||
params: dict[str, object],
|
||||
litellm_params: dict[str, object],
|
||||
api_base: str | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
) -> dict[str, object]:
|
||||
|
|
@ -313,8 +316,8 @@ async def handle_a2a_completion(
|
|||
|
||||
async def handle_a2a_completion_streaming(
|
||||
request_id: str,
|
||||
params: dict[str, Any],
|
||||
litellm_params: dict[str, Any],
|
||||
params: dict[str, object],
|
||||
litellm_params: dict[str, object],
|
||||
api_base: str | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
) -> AsyncIterator[dict[str, object]]:
|
||||
|
|
|
|||
|
|
@ -12,7 +12,8 @@ Provides standalone functions with @client decorator for LiteLLM logging integra
|
|||
import asyncio
|
||||
import datetime
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator, Coroutine
|
||||
from collections.abc import AsyncIterator, Coroutine, Mapping
|
||||
from types import ModuleType
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional, cast
|
||||
|
||||
import litellm
|
||||
|
|
@ -30,6 +31,7 @@ from litellm.utils import client
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from a2a.client import Client as A2AClientType
|
||||
from a2a.client import ClientCallContext as A2ACallContextType
|
||||
from a2a.compat.v0_3.types import (
|
||||
AgentCard,
|
||||
Message,
|
||||
|
|
@ -37,15 +39,18 @@ if TYPE_CHECKING:
|
|||
SendMessageResponse,
|
||||
SendStreamingMessageRequest,
|
||||
SendStreamingMessageResponse,
|
||||
SendStreamingMessageSuccessResponse,
|
||||
Task,
|
||||
)
|
||||
from a2a.types.a2a_pb2 import SendMessageRequest as CoreSendMessageRequest
|
||||
from a2a.types.a2a_pb2 import StreamResponse as CoreStreamResponse
|
||||
|
||||
# Runtime imports — requires a2a-sdk>=1.1.0
|
||||
A2A_SDK_AVAILABLE = False
|
||||
_a2a_conversions: Any = None
|
||||
_a2a_conversions: ModuleType | None = None
|
||||
|
||||
try:
|
||||
from a2a.client import Client, ClientConfig, create_client
|
||||
from a2a.client import Client, ClientCallContext, ClientConfig, create_client
|
||||
from a2a.compat.v0_3 import conversions as _a2a_conversions
|
||||
from a2a.compat.v0_3.types import (
|
||||
Message,
|
||||
|
|
@ -60,6 +65,7 @@ try:
|
|||
A2A_SDK_AVAILABLE = True
|
||||
except ImportError:
|
||||
Client = None
|
||||
ClientCallContext = None
|
||||
ClientConfig = None
|
||||
create_client = None
|
||||
|
||||
|
|
@ -126,7 +132,7 @@ _A2A_COST_PARAM_KEYS: Final = ("cost_per_query", "input_cost_per_token", "output
|
|||
|
||||
def _set_litellm_params_on_logging_obj(
|
||||
kwargs: dict[str, Any],
|
||||
litellm_params: dict[str, Any],
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> None:
|
||||
"""
|
||||
Merge the agent's pricing params into model_call_details["litellm_params"]
|
||||
|
|
@ -148,7 +154,7 @@ def _set_litellm_params_on_logging_obj(
|
|||
logging_obj.model_call_details["litellm_params"] = {**existing, **cost_params}
|
||||
|
||||
|
||||
def _get_a2a_model_info(a2a_client: Any, kwargs: dict[str, Any]) -> str:
|
||||
def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: dict[str, Any]) -> str:
|
||||
"""
|
||||
Extract agent info and set model/custom_llm_provider for cost tracking.
|
||||
|
||||
|
|
@ -177,7 +183,7 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: dict[str, Any]) -> str:
|
|||
return agent_name
|
||||
|
||||
|
||||
def _get_a2a_client_agent_card(a2a_client: Any) -> Optional["AgentCard"]:
|
||||
def _get_a2a_client_agent_card(a2a_client: "A2AClientType") -> Optional["AgentCard"]:
|
||||
agent_card = cast(Optional["AgentCard"], getattr(a2a_client, "_litellm_agent_card", None))
|
||||
if agent_card is not None:
|
||||
return agent_card
|
||||
|
|
@ -189,9 +195,9 @@ def _get_a2a_client_agent_card(a2a_client: Any) -> Optional["AgentCard"]:
|
|||
|
||||
async def _send_message_via_completion_bridge(
|
||||
request: "SendMessageRequest",
|
||||
custom_llm_provider: str,
|
||||
custom_llm_provider: object,
|
||||
api_base: str | None,
|
||||
litellm_params: dict[str, Any],
|
||||
litellm_params: dict[str, object],
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
) -> LiteLLMSendMessageResponse:
|
||||
"""
|
||||
|
|
@ -218,6 +224,24 @@ async def _send_message_via_completion_bridge(
|
|||
return LiteLLMSendMessageResponse.from_dict(response_dict, request_id=str(request.id))
|
||||
|
||||
|
||||
def _get_a2a_call_context(a2a_client: "A2AClientType") -> Optional["A2ACallContextType"]:
|
||||
return getattr(a2a_client, "_litellm_call_context", None)
|
||||
|
||||
|
||||
def _to_core_send_message_request(request: "SendMessageRequest") -> "CoreSendMessageRequest":
|
||||
from a2a.compat.v0_3 import conversions
|
||||
|
||||
return conversions.to_core_send_message_request(request)
|
||||
|
||||
|
||||
def _to_compat_stream_response(
|
||||
event: "CoreStreamResponse", request_id: str | int
|
||||
) -> "SendStreamingMessageSuccessResponse":
|
||||
from a2a.compat.v0_3 import conversions
|
||||
|
||||
return conversions.to_compat_stream_response(event, request_id=request_id)
|
||||
|
||||
|
||||
async def _send_message(a2a_client: "A2AClientType", request: "SendMessageRequest") -> "SendMessageResponse":
|
||||
"""Send a non-streaming message via a2a-sdk 1.x and return JSON-RPC response."""
|
||||
if _a2a_conversions is None:
|
||||
|
|
@ -225,17 +249,14 @@ async def _send_message(a2a_client: "A2AClientType", request: "SendMessageReques
|
|||
"The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk"
|
||||
)
|
||||
|
||||
pb_request: Final = _a2a_conversions.to_core_send_message_request(request)
|
||||
pb_request: Final = _to_core_send_message_request(request)
|
||||
last_event = None
|
||||
async for event in a2a_client.send_message(pb_request):
|
||||
async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)):
|
||||
last_event = event
|
||||
if last_event is None:
|
||||
raise RuntimeError("A2A send_message failed: no response received from agent.")
|
||||
|
||||
stream_compat: Final = _a2a_conversions.to_compat_stream_response(
|
||||
last_event,
|
||||
request_id=request.id,
|
||||
)
|
||||
stream_compat: Final = _to_compat_stream_response(last_event, request_id=request.id)
|
||||
result: Final = stream_compat.result
|
||||
if not isinstance(result, (Message, Task)):
|
||||
raise RuntimeError(
|
||||
|
|
@ -300,12 +321,9 @@ async def _stream_messages(
|
|||
"The 'a2a' package is required for A2A agent invocation. Install it with: pip install a2a-sdk"
|
||||
)
|
||||
|
||||
pb_request: Final = _a2a_conversions.to_core_send_message_request(request)
|
||||
async for event in a2a_client.send_message(pb_request):
|
||||
compat_chunk = _a2a_conversions.to_compat_stream_response(
|
||||
event,
|
||||
request_id=request.id,
|
||||
)
|
||||
pb_request: Final[CoreSendMessageRequest] = _a2a_conversions.to_core_send_message_request(request)
|
||||
async for event in a2a_client.send_message(pb_request, context=_get_a2a_call_context(a2a_client)):
|
||||
compat_chunk = _to_compat_stream_response(event, request_id=request.id)
|
||||
yield SendStreamingMessageResponse(root=compat_chunk)
|
||||
|
||||
|
||||
|
|
@ -362,10 +380,10 @@ async def asend_message(
|
|||
a2a_client: Optional["A2AClientType"] = None,
|
||||
request: Optional["SendMessageRequest"] = None,
|
||||
api_base: str | None = None,
|
||||
litellm_params: dict[str, Any] | None = None,
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
agent_id: str | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
**kwargs: Any,
|
||||
**kwargs: object,
|
||||
) -> LiteLLMSendMessageResponse:
|
||||
"""
|
||||
Async: Send a message to an A2A agent.
|
||||
|
|
@ -479,7 +497,7 @@ async def asend_message(
|
|||
response: Final = LiteLLMSendMessageResponse.from_a2a_response(a2a_response, request_id=str(request.id))
|
||||
|
||||
# Calculate token usage from request and response
|
||||
response_dict: Final = a2a_response.model_dump(mode="json", exclude_none=True)
|
||||
response_dict: Final[dict[str, object]] = a2a_response.model_dump(mode="json", exclude_none=True)
|
||||
(
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
|
|
@ -510,7 +528,7 @@ def send_message(
|
|||
a2a_client: "A2AClientType",
|
||||
request: "SendMessageRequest",
|
||||
**kwargs: Any,
|
||||
) -> LiteLLMSendMessageResponse | Coroutine[Any, Any, LiteLLMSendMessageResponse]:
|
||||
) -> LiteLLMSendMessageResponse | Coroutine[object, object, LiteLLMSendMessageResponse]:
|
||||
"""
|
||||
Sync: Send a message to an A2A agent.
|
||||
|
||||
|
|
@ -539,9 +557,9 @@ def _build_streaming_logging_obj(
|
|||
request: "SendStreamingMessageRequest",
|
||||
agent_name: str,
|
||||
agent_id: str | None,
|
||||
litellm_params: dict[str, Any] | None,
|
||||
metadata: dict[str, Any] | None,
|
||||
proxy_server_request: dict[str, Any] | None,
|
||||
litellm_params: dict[str, object] | None,
|
||||
metadata: dict[str, object] | None,
|
||||
proxy_server_request: dict[str, object] | None,
|
||||
) -> Logging:
|
||||
"""Build logging object for streaming A2A requests."""
|
||||
start_time: Final = datetime.datetime.now()
|
||||
|
|
@ -582,10 +600,10 @@ async def asend_message_streaming(
|
|||
a2a_client: Optional["A2AClientType"] = None,
|
||||
request: Optional["SendStreamingMessageRequest"] = None,
|
||||
api_base: str | None = None,
|
||||
litellm_params: dict[str, Any] | None = None,
|
||||
litellm_params: dict[str, object] | None = None,
|
||||
agent_id: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
proxy_server_request: dict[str, Any] | None = None,
|
||||
metadata: dict[str, object] | None = None,
|
||||
proxy_server_request: dict[str, object] | None = None,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
**kwargs: object,
|
||||
) -> AsyncIterator[Any]:
|
||||
|
|
@ -756,26 +774,12 @@ async def create_a2a_client(
|
|||
|
||||
verbose_logger.info("Creating A2A client for %s", base_url)
|
||||
|
||||
# Use get_async_httpx_client with per-agent params so that different agents
|
||||
# (with different extra_headers) get separate cached clients. The params
|
||||
# dict is hashed into the cache key, keeping agent auth isolated while
|
||||
# still reusing connections within the same agent.
|
||||
#
|
||||
# Only pass params that AsyncHTTPHandler.__init__ accepts (e.g. timeout).
|
||||
# Use "disable_aiohttp_transport" key for cache-key-only data (it's
|
||||
# filtered out before reaching the constructor).
|
||||
_client_params: Final[dict] = {"timeout": timeout}
|
||||
if extra_headers:
|
||||
# Encode headers into a cache-key-only param so each unique header
|
||||
# set produces a distinct cache key.
|
||||
_client_params["disable_aiohttp_transport"] = str(sorted(extra_headers.items()))
|
||||
_async_handler: Final = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.A2AProvider,
|
||||
params=_client_params,
|
||||
params={"timeout": timeout},
|
||||
)
|
||||
httpx_client: Final = _async_handler.client
|
||||
if extra_headers:
|
||||
httpx_client.headers.update(extra_headers)
|
||||
verbose_proxy_logger.debug("A2A client created with extra_headers=%s", list(extra_headers.keys()))
|
||||
|
||||
a2a_client: Final = await create_client( # pyright: ignore[reportOptionalCall]
|
||||
|
|
@ -784,11 +788,17 @@ async def create_a2a_client(
|
|||
httpx_client=httpx_client,
|
||||
streaming=streaming,
|
||||
),
|
||||
resolver_http_kwargs={"headers": extra_headers} if extra_headers else None,
|
||||
)
|
||||
# Stash LiteLLM-owned handles on the client so the localhost-retry path can reuse
|
||||
# the configured httpx client (with this agent's trace-id/auth headers) without
|
||||
# excavating a2a-sdk private internals.
|
||||
# the configured httpx client and this agent's headers without excavating
|
||||
# a2a-sdk private internals.
|
||||
a2a_client._litellm_httpx_client = httpx_client
|
||||
a2a_client._litellm_call_context = ( # pyright: ignore[reportAttributeAccessIssue] # LiteLLM-owned stash
|
||||
ClientCallContext(service_parameters=extra_headers) # pyright: ignore[reportOptionalCall] # SDK checked above
|
||||
if extra_headers
|
||||
else None
|
||||
)
|
||||
agent_card: Final = getattr(a2a_client, "_card", None)
|
||||
if agent_card is not None:
|
||||
a2a_client._litellm_agent_card = agent_card
|
||||
|
|
|
|||
|
|
@ -6,16 +6,33 @@ This module provides fake streaming by converting non-streaming responses into s
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any, Final, cast
|
||||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from typing import Any, Final, Protocol, cast, runtime_checkable
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
get_async_httpx_client,
|
||||
)
|
||||
|
||||
_ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object])
|
||||
_STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
_LIST_ADAPTER: Final = TypeAdapter(list[object])
|
||||
_TEXT_ADAPTER: Final = TypeAdapter(str)
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _SupportsModelDump(Protocol):
|
||||
def model_dump(self, *, mode: str, exclude_none: bool) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class _SupportsPydanticDict(Protocol):
|
||||
def dict(self, *, exclude_none: bool) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
class PydanticAITransformation:
|
||||
"""
|
||||
|
|
@ -28,7 +45,7 @@ class PydanticAITransformation:
|
|||
"""
|
||||
|
||||
@staticmethod
|
||||
def _remove_none_values(obj: Any) -> Any:
|
||||
def _remove_none_values(obj: object) -> object:
|
||||
"""
|
||||
Recursively remove None values from a dict/list structure.
|
||||
|
||||
|
|
@ -42,14 +59,18 @@ class PydanticAITransformation:
|
|||
Cleaned object with None values removed
|
||||
"""
|
||||
if isinstance(obj, dict):
|
||||
return {k: PydanticAITransformation._remove_none_values(v) for k, v in obj.items() if v is not None}
|
||||
typed_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python(obj)
|
||||
return {k: PydanticAITransformation._remove_none_values(v) for k, v in typed_dict.items() if v is not None}
|
||||
elif isinstance(obj, list):
|
||||
return [PydanticAITransformation._remove_none_values(item) for item in obj if item is not None]
|
||||
typed_list: Final = _LIST_ADAPTER.validate_python(obj)
|
||||
return [PydanticAITransformation._remove_none_values(item) for item in typed_list if item is not None]
|
||||
else:
|
||||
return obj
|
||||
|
||||
@staticmethod
|
||||
def _params_to_dict(params: Any) -> dict[str, Any]:
|
||||
def _params_to_dict(
|
||||
params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
|
||||
) -> Mapping[str, object]:
|
||||
"""
|
||||
Convert params to a dict, handling Pydantic models.
|
||||
|
||||
|
|
@ -59,10 +80,10 @@ class PydanticAITransformation:
|
|||
Returns:
|
||||
Dict representation of params
|
||||
"""
|
||||
if hasattr(params, "model_dump"):
|
||||
if isinstance(params, _SupportsModelDump):
|
||||
# Pydantic v2 model
|
||||
return params.model_dump(mode="python", exclude_none=True)
|
||||
elif hasattr(params, "dict"):
|
||||
elif isinstance(params, _SupportsPydanticDict):
|
||||
# Pydantic v1 model
|
||||
return params.dict(exclude_none=True)
|
||||
elif isinstance(params, dict):
|
||||
|
|
@ -75,12 +96,12 @@ class PydanticAITransformation:
|
|||
async def _poll_for_completion(
|
||||
client: AsyncHTTPHandler,
|
||||
endpoint: str,
|
||||
task_id: str,
|
||||
task_id: object,
|
||||
request_id: str,
|
||||
max_attempts: int = 30,
|
||||
poll_interval: float = 0.5,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Poll for task completion using tasks/get method.
|
||||
|
||||
|
|
@ -112,10 +133,10 @@ class PydanticAITransformation:
|
|||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
poll_data = response.json()
|
||||
poll_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json())
|
||||
|
||||
result = poll_data.get("result", {})
|
||||
status = result.get("status", {})
|
||||
result = _STR_KEY_DICT_ADAPTER.validate_python(poll_data.get("result", {}))
|
||||
status = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {}))
|
||||
state = status.get("state", "")
|
||||
|
||||
verbose_logger.debug("Pydantic AI: Poll attempt %s/%s, state=%s", attempt + 1, max_attempts, state)
|
||||
|
|
@ -133,10 +154,10 @@ class PydanticAITransformation:
|
|||
async def _send_and_poll_raw(
|
||||
api_base: str,
|
||||
request_id: str,
|
||||
params: Any,
|
||||
params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
|
||||
timeout: float = 60.0,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Send a request to Pydantic AI agent and return the raw task response.
|
||||
|
||||
|
|
@ -153,14 +174,16 @@ class PydanticAITransformation:
|
|||
Raw Pydantic AI task response (with history/artifacts)
|
||||
"""
|
||||
# Convert params to dict if it's a Pydantic model
|
||||
params_dict = PydanticAITransformation._params_to_dict(params)
|
||||
|
||||
# Remove None values - FastA2A doesn't accept null for optional fields
|
||||
params_dict = PydanticAITransformation._remove_none_values(params_dict)
|
||||
params_dict: Final = _ANY_KEY_DICT_ADAPTER.validate_python(
|
||||
PydanticAITransformation._remove_none_values(PydanticAITransformation._params_to_dict(params))
|
||||
)
|
||||
|
||||
# Ensure the message has 'kind': 'message' as required by FastA2A/Pydantic AI
|
||||
if "message" in params_dict:
|
||||
params_dict["message"]["kind"] = "message"
|
||||
message_value: Final = _ANY_KEY_DICT_ADAPTER.validate_python(params_dict["message"])
|
||||
message_value["kind"] = "message"
|
||||
params_dict["message"] = message_value
|
||||
|
||||
# Build A2A JSON-RPC request using message/send method for FastA2A compatibility
|
||||
a2a_request: Final = {
|
||||
|
|
@ -189,11 +212,11 @@ class PydanticAITransformation:
|
|||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
response_data = response.json()
|
||||
response_data = _STR_KEY_DICT_ADAPTER.validate_python(response.json())
|
||||
|
||||
# Check if task is already completed
|
||||
result: Final = response_data.get("result", {})
|
||||
status: Final = result.get("status", {})
|
||||
result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {}))
|
||||
status: Final = _STR_KEY_DICT_ADAPTER.validate_python(result.get("status", {}))
|
||||
state: Final = status.get("state", "")
|
||||
|
||||
if state != "completed":
|
||||
|
|
@ -217,10 +240,10 @@ class PydanticAITransformation:
|
|||
async def send_non_streaming_request(
|
||||
api_base: str,
|
||||
request_id: str,
|
||||
params: Any,
|
||||
params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
|
||||
timeout: float = 60.0,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Send a non-streaming A2A request to Pydantic AI agent and wait for completion.
|
||||
|
||||
|
|
@ -253,10 +276,10 @@ class PydanticAITransformation:
|
|||
async def send_and_get_raw_response(
|
||||
api_base: str,
|
||||
request_id: str,
|
||||
params: Any,
|
||||
params: "_SupportsModelDump | _SupportsPydanticDict | Mapping[str, object]",
|
||||
timeout: float = 60.0,
|
||||
agent_extra_headers: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Send a request to Pydantic AI agent and return the raw task response.
|
||||
|
||||
|
|
@ -282,9 +305,9 @@ class PydanticAITransformation:
|
|||
|
||||
@staticmethod
|
||||
def _transform_to_a2a_response(
|
||||
response_data: dict[str, Any],
|
||||
response_data: Mapping[str, object],
|
||||
request_id: str,
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Transform Pydantic AI task response to standard A2A non-streaming format.
|
||||
|
||||
|
|
@ -328,7 +351,7 @@ class PydanticAITransformation:
|
|||
}
|
||||
|
||||
@staticmethod
|
||||
def _extract_response_text(response_data: dict[str, Any]) -> tuple[str, str, list]:
|
||||
def _extract_response_text(response_data: Mapping[str, object]) -> tuple[object, object, Sequence[object]]:
|
||||
"""
|
||||
Extract response text from completed task response.
|
||||
|
||||
|
|
@ -342,52 +365,53 @@ class PydanticAITransformation:
|
|||
Returns:
|
||||
Tuple of (full_text, message_id, parts)
|
||||
"""
|
||||
result: Final = response_data.get("result", {})
|
||||
result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {}))
|
||||
|
||||
# Try to extract from artifacts first (preferred for results)
|
||||
artifacts: Final = result.get("artifacts", [])
|
||||
if artifacts:
|
||||
for artifact in artifacts:
|
||||
parts = artifact.get("parts", [])
|
||||
for artifact in _LIST_ADAPTER.validate_python(artifacts):
|
||||
parts = _LIST_ADAPTER.validate_python(_STR_KEY_DICT_ADAPTER.validate_python(artifact).get("parts", []))
|
||||
for part in parts:
|
||||
if part.get("kind") == "text":
|
||||
text = part.get("text", "")
|
||||
if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text":
|
||||
text = part_dict.get("text", "")
|
||||
if text:
|
||||
return text, str(uuid4()), parts
|
||||
|
||||
# Fall back to history - get the last agent message
|
||||
history: Final = result.get("history", [])
|
||||
history: Final = _LIST_ADAPTER.validate_python(result.get("history", []))
|
||||
for msg in reversed(history):
|
||||
if msg.get("role") == "agent":
|
||||
parts = msg.get("parts", [])
|
||||
message_id = msg.get("messageId", str(uuid4()))
|
||||
if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "agent":
|
||||
parts = _LIST_ADAPTER.validate_python(msg_dict.get("parts", []))
|
||||
message_id = msg_dict.get("messageId", str(uuid4()))
|
||||
full_text = ""
|
||||
for part in parts:
|
||||
if part.get("kind") == "text":
|
||||
full_text += part.get("text", "")
|
||||
if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text":
|
||||
full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", ""))
|
||||
if full_text:
|
||||
return full_text, message_id, parts
|
||||
|
||||
# Fall back to message field (original format)
|
||||
message: Final = result.get("message", {})
|
||||
if message:
|
||||
parts = message.get("parts", [])
|
||||
message_id = message.get("messageId", str(uuid4()))
|
||||
message_dict: Final = _STR_KEY_DICT_ADAPTER.validate_python(message)
|
||||
parts = _LIST_ADAPTER.validate_python(message_dict.get("parts", []))
|
||||
message_id = message_dict.get("messageId", str(uuid4()))
|
||||
full_text = ""
|
||||
for part in parts:
|
||||
if part.get("kind") == "text":
|
||||
full_text += part.get("text", "")
|
||||
if (part_dict := _STR_KEY_DICT_ADAPTER.validate_python(part)).get("kind") == "text":
|
||||
full_text += _TEXT_ADAPTER.validate_python(part_dict.get("text", ""))
|
||||
return full_text, message_id, parts
|
||||
|
||||
return "", str(uuid4()), []
|
||||
|
||||
@staticmethod
|
||||
async def fake_streaming_from_response(
|
||||
response_data: dict[str, Any],
|
||||
response_data: Mapping[str, object],
|
||||
request_id: str,
|
||||
chunk_size: int = 50,
|
||||
delay_ms: int = 10,
|
||||
) -> AsyncIterator[dict[str, Any]]:
|
||||
) -> AsyncIterator[dict[str, object]]:
|
||||
"""
|
||||
Convert a non-streaming A2A response into fake streaming chunks.
|
||||
|
||||
|
|
@ -410,12 +434,12 @@ class PydanticAITransformation:
|
|||
full_text, message_id, parts = PydanticAITransformation._extract_response_text(response_data)
|
||||
|
||||
# Extract input message from raw response for history
|
||||
result: Final = response_data.get("result", {})
|
||||
history: Final = result.get("history", [])
|
||||
input_message = {}
|
||||
result: Final = _STR_KEY_DICT_ADAPTER.validate_python(response_data.get("result", {}))
|
||||
history: Final = _LIST_ADAPTER.validate_python(result.get("history", []))
|
||||
input_message = _STR_KEY_DICT_ADAPTER.validate_python({})
|
||||
for msg in history:
|
||||
if msg.get("role") == "user":
|
||||
input_message = msg
|
||||
if (msg_dict := _STR_KEY_DICT_ADAPTER.validate_python(msg)).get("role") == "user":
|
||||
input_message = msg_dict
|
||||
break
|
||||
|
||||
# Generate IDs for streaming events
|
||||
|
|
@ -426,45 +450,49 @@ class PydanticAITransformation:
|
|||
|
||||
# 1. Emit initial task event (kind: "task", status: "submitted")
|
||||
# Format matches A2ACompletionBridgeTransformation.create_task_event
|
||||
task_event: Final = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {
|
||||
"contextId": context_id,
|
||||
"history": [
|
||||
{
|
||||
"contextId": context_id,
|
||||
"kind": "message",
|
||||
"messageId": input_message_id,
|
||||
"parts": input_message.get("parts", [{"kind": "text", "text": ""}]),
|
||||
"role": "user",
|
||||
"taskId": task_id,
|
||||
}
|
||||
],
|
||||
"id": task_id,
|
||||
"kind": "task",
|
||||
"status": {
|
||||
"state": "submitted",
|
||||
task_event: Final = _STR_KEY_DICT_ADAPTER.validate_python(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {
|
||||
"contextId": context_id,
|
||||
"history": [
|
||||
{
|
||||
"contextId": context_id,
|
||||
"kind": "message",
|
||||
"messageId": input_message_id,
|
||||
"parts": input_message.get("parts", [{"kind": "text", "text": ""}]),
|
||||
"role": "user",
|
||||
"taskId": task_id,
|
||||
}
|
||||
],
|
||||
"id": task_id,
|
||||
"kind": "task",
|
||||
"status": {
|
||||
"state": "submitted",
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
)
|
||||
yield task_event
|
||||
|
||||
# 2. Emit status update (kind: "status-update", status: "working")
|
||||
# Format matches A2ACompletionBridgeTransformation.create_status_update_event
|
||||
working_event: Final = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {
|
||||
"contextId": context_id,
|
||||
"final": False,
|
||||
"kind": "status-update",
|
||||
"status": {
|
||||
"state": "working",
|
||||
working_event: Final = _STR_KEY_DICT_ADAPTER.validate_python(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {
|
||||
"contextId": context_id,
|
||||
"final": False,
|
||||
"kind": "status-update",
|
||||
"status": {
|
||||
"state": "working",
|
||||
},
|
||||
"taskId": task_id,
|
||||
},
|
||||
"taskId": task_id,
|
||||
},
|
||||
}
|
||||
}
|
||||
)
|
||||
yield working_event
|
||||
|
||||
# Small delay to simulate processing
|
||||
|
|
@ -473,29 +501,32 @@ class PydanticAITransformation:
|
|||
# 3. Emit artifact update chunks (kind: "artifact-update")
|
||||
# Format matches A2ACompletionBridgeTransformation.create_artifact_update_event
|
||||
if full_text:
|
||||
full_text_str: Final = _TEXT_ADAPTER.validate_python(full_text)
|
||||
# Split text into chunks
|
||||
for i in range(0, len(full_text), chunk_size):
|
||||
chunk_text = full_text[i : i + chunk_size]
|
||||
is_last_chunk = (i + chunk_size) >= len(full_text)
|
||||
for i in range(0, len(full_text_str), chunk_size):
|
||||
chunk_text = full_text_str[i : i + chunk_size]
|
||||
is_last_chunk = (i + chunk_size) >= len(full_text_str)
|
||||
|
||||
artifact_event = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {
|
||||
"contextId": context_id,
|
||||
"kind": "artifact-update",
|
||||
"taskId": task_id,
|
||||
"artifact": {
|
||||
"artifactId": artifact_id,
|
||||
"parts": [
|
||||
{
|
||||
"kind": "text",
|
||||
"text": chunk_text,
|
||||
}
|
||||
],
|
||||
artifact_event = _STR_KEY_DICT_ADAPTER.validate_python(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {
|
||||
"contextId": context_id,
|
||||
"kind": "artifact-update",
|
||||
"taskId": task_id,
|
||||
"artifact": {
|
||||
"artifactId": artifact_id,
|
||||
"parts": [
|
||||
{
|
||||
"kind": "text",
|
||||
"text": chunk_text,
|
||||
}
|
||||
],
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
)
|
||||
yield artifact_event
|
||||
|
||||
# Add delay between chunks (except for last chunk)
|
||||
|
|
@ -503,19 +534,21 @@ class PydanticAITransformation:
|
|||
await asyncio.sleep(delay_ms / 1000.0)
|
||||
|
||||
# 4. Emit final status update (kind: "status-update", status: "completed", final: true)
|
||||
completed_event: Final = {
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {
|
||||
"contextId": context_id,
|
||||
"final": True,
|
||||
"kind": "status-update",
|
||||
"status": {
|
||||
"state": "completed",
|
||||
completed_event: Final = _STR_KEY_DICT_ADAPTER.validate_python(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": request_id,
|
||||
"result": {
|
||||
"contextId": context_id,
|
||||
"final": True,
|
||||
"kind": "status-update",
|
||||
"status": {
|
||||
"state": "completed",
|
||||
},
|
||||
"taskId": task_id,
|
||||
},
|
||||
"taskId": task_id,
|
||||
},
|
||||
}
|
||||
}
|
||||
)
|
||||
yield completed_event
|
||||
|
||||
verbose_logger.info("Pydantic AI: Fake streaming completed for request_id=%s", request_id)
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ class ResponsesToCompletionBridgeHandlerInputKwargs(TypedDict):
|
|||
model_response: "ModelResponse"
|
||||
logging_obj: "LiteLLMLoggingObj"
|
||||
custom_llm_provider: str
|
||||
encoding: object
|
||||
|
||||
|
||||
class ResponsesToCompletionBridgeHandler:
|
||||
|
|
@ -102,35 +103,37 @@ class ResponsesToCompletionBridgeHandler:
|
|||
from litellm import LiteLLMLoggingObj
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
model: Final = kwargs.get("model")
|
||||
typed_kwargs: Final[dict[str, object]] = kwargs
|
||||
|
||||
model: Final = typed_kwargs.get("model")
|
||||
if model is None or not isinstance(model, str):
|
||||
raise ValueError("model is required")
|
||||
|
||||
custom_llm_provider: Final = kwargs.get("custom_llm_provider")
|
||||
custom_llm_provider: Final = typed_kwargs.get("custom_llm_provider")
|
||||
if custom_llm_provider is None or not isinstance(custom_llm_provider, str):
|
||||
raise ValueError("custom_llm_provider is required")
|
||||
|
||||
messages: Final = kwargs.get("messages")
|
||||
messages: Final = typed_kwargs.get("messages")
|
||||
if messages is None or not isinstance(messages, list):
|
||||
raise ValueError("messages is required")
|
||||
|
||||
optional_params: Final = kwargs.get("optional_params")
|
||||
optional_params: Final = typed_kwargs.get("optional_params")
|
||||
if optional_params is None or not isinstance(optional_params, dict):
|
||||
raise ValueError("optional_params is required")
|
||||
|
||||
litellm_params: Final = kwargs.get("litellm_params")
|
||||
litellm_params: Final = typed_kwargs.get("litellm_params")
|
||||
if litellm_params is None or not isinstance(litellm_params, dict):
|
||||
raise ValueError("litellm_params is required")
|
||||
|
||||
headers: Final = kwargs.get("headers")
|
||||
headers: Final = typed_kwargs.get("headers")
|
||||
if headers is None or not isinstance(headers, dict):
|
||||
raise ValueError("headers is required")
|
||||
|
||||
model_response: Final = kwargs.get("model_response")
|
||||
model_response: Final = typed_kwargs.get("model_response")
|
||||
if model_response is None or not isinstance(model_response, ModelResponse):
|
||||
raise ValueError("model_response is required")
|
||||
|
||||
logging_obj: Final = kwargs.get("logging_obj")
|
||||
logging_obj: Final = typed_kwargs.get("logging_obj")
|
||||
if logging_obj is None or not isinstance(logging_obj, LiteLLMLoggingObj):
|
||||
raise ValueError("logging_obj is required")
|
||||
|
||||
|
|
@ -143,6 +146,7 @@ class ResponsesToCompletionBridgeHandler:
|
|||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
encoding=typed_kwargs.get("encoding"),
|
||||
)
|
||||
|
||||
def completion(
|
||||
|
|
@ -205,7 +209,7 @@ class ResponsesToCompletionBridgeHandler:
|
|||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
encoding=kwargs.get("encoding"),
|
||||
encoding=validated_kwargs["encoding"],
|
||||
api_key=kwargs.get("api_key"),
|
||||
json_mode=kwargs.get("json_mode"),
|
||||
)
|
||||
|
|
@ -230,7 +234,7 @@ class ResponsesToCompletionBridgeHandler:
|
|||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
encoding=kwargs.get("encoding"),
|
||||
encoding=validated_kwargs["encoding"],
|
||||
api_key=kwargs.get("api_key"),
|
||||
json_mode=kwargs.get("json_mode"),
|
||||
)
|
||||
|
|
@ -303,7 +307,7 @@ class ResponsesToCompletionBridgeHandler:
|
|||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
encoding=kwargs.get("encoding"),
|
||||
encoding=validated_kwargs["encoding"],
|
||||
api_key=kwargs.get("api_key"),
|
||||
json_mode=kwargs.get("json_mode"),
|
||||
)
|
||||
|
|
@ -328,7 +332,7 @@ class ResponsesToCompletionBridgeHandler:
|
|||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
encoding=kwargs.get("encoding"),
|
||||
encoding=validated_kwargs["encoding"],
|
||||
api_key=kwargs.get("api_key"),
|
||||
json_mode=kwargs.get("json_mode"),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -4,8 +4,8 @@ Handler for transforming /chat/completions api requests to litellm.responses req
|
|||
|
||||
import json
|
||||
import os
|
||||
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Union, cast
|
||||
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast
|
||||
|
||||
from openai.types.responses.custom_tool_param import CustomToolParam
|
||||
from openai.types.responses.response_input_param import (
|
||||
|
|
@ -45,6 +45,9 @@ from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from openai.types.responses import ResponseInputImageParam
|
||||
from openai.types.responses.response_text_config_param import (
|
||||
ResponseTextConfigParam as ResponseText,
|
||||
)
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm import LiteLLMLoggingObj, ModelResponse
|
||||
|
|
@ -57,6 +60,19 @@ if TYPE_CHECKING:
|
|||
ChatCompletionThinkingBlock,
|
||||
OpenAIMessageContentListBlock,
|
||||
)
|
||||
from litellm.types.utils import Choices
|
||||
|
||||
|
||||
class _ReasoningSummaryText(TypedDict):
|
||||
type: str
|
||||
text: str
|
||||
|
||||
|
||||
class _BuiltReasoningItem(TypedDict):
|
||||
type: Literal["reasoning"]
|
||||
id: str
|
||||
encrypted_content: str | None
|
||||
summary: Sequence[_ReasoningSummaryText]
|
||||
|
||||
|
||||
def _get_reasoning_items(
|
||||
|
|
@ -72,13 +88,13 @@ def _get_reasoning_items(
|
|||
def _build_reasoning_item(
|
||||
item_id: str,
|
||||
encrypted_content: str | None,
|
||||
summary_raw: Any,
|
||||
) -> dict[str, Any]:
|
||||
summary_raw: Iterable[object] | None,
|
||||
) -> _BuiltReasoningItem:
|
||||
"""Build a ChatCompletionReasoningItem-shaped dict from raw response data.
|
||||
|
||||
Handles both pydantic objects (attribute access) and plain dicts.
|
||||
"""
|
||||
summary: Final[list[dict[str, Any]]] = []
|
||||
summary: Final[list[_ReasoningSummaryText]] = []
|
||||
for s in summary_raw or []:
|
||||
if isinstance(s, dict):
|
||||
summary.append({"type": s.get("type", "summary_text"), "text": s.get("text", "")})
|
||||
|
|
@ -98,7 +114,7 @@ def _build_reasoning_item(
|
|||
|
||||
|
||||
class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False):
|
||||
provider_specific_fields: Mapping[str, Any]
|
||||
provider_specific_fields: Mapping[str, object]
|
||||
|
||||
|
||||
def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict:
|
||||
|
|
@ -142,10 +158,10 @@ def _flat_responses_tool_choice(choice_type: str, name: str) -> ToolChoiceFuncti
|
|||
|
||||
|
||||
def _reasoning_item_to_response_input(
|
||||
r_item: ChatCompletionReasoningItem | dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
r_item: ChatCompletionReasoningItem,
|
||||
) -> dict[str, object]:
|
||||
"""Convert a stored ChatCompletionReasoningItem back to a Responses API input item."""
|
||||
r_input: Final[dict[str, Any]] = {
|
||||
r_input: Final[dict[str, object]] = {
|
||||
"type": "reasoning",
|
||||
"id": r_item.get("id") or f"rs_{id(r_item)}",
|
||||
# summary is always required by the Responses API, even when empty
|
||||
|
|
@ -181,7 +197,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
return _flat_responses_tool_choice(choice_type, nested_name)
|
||||
return tool_choice
|
||||
|
||||
def _handle_raw_dict_response_item(self, item: dict[str, Any], index: int) -> tuple[Any | None, int]:
|
||||
def _handle_raw_dict_response_item(self, item: dict[str, Any], index: int) -> tuple["Choices | None", int]:
|
||||
"""
|
||||
Handle raw dict response items from Responses API (e.g., GPT-5 Codex format).
|
||||
|
||||
|
|
@ -228,8 +244,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
|
||||
def convert_chat_completion_messages_to_responses_api(
|
||||
self, messages: list["AllMessageValues"]
|
||||
) -> tuple[list[Any], str | None]:
|
||||
input_items: Final[list[Any]] = []
|
||||
) -> tuple[list[object], str | None]:
|
||||
input_items: Final[list[object]] = []
|
||||
instructions: str | None = None
|
||||
custom_tool_call_ids: Final = frozenset(
|
||||
tool_call["id"]
|
||||
|
|
@ -270,7 +286,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
# Convert tool message to function call output format
|
||||
# The Responses API expects 'output' to be a list with input_text/input_image types
|
||||
# Using list format for consistency across text and multimodal content
|
||||
tool_output: list[dict[str, Any]]
|
||||
tool_output: list[dict[str, object]]
|
||||
if content is None:
|
||||
tool_output = []
|
||||
elif isinstance(content, str):
|
||||
|
|
@ -308,7 +324,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
function = tool_call.get("function")
|
||||
custom = tool_call.get("custom")
|
||||
if function:
|
||||
input_tool_call: dict[str, Any] = {
|
||||
input_tool_call: dict[str, object] = {
|
||||
"type": "function_call",
|
||||
"call_id": tool_call["id"],
|
||||
}
|
||||
|
|
@ -376,15 +392,15 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
elif key == "web_search_options":
|
||||
self._add_web_search_tool(responses_api_request, value)
|
||||
|
||||
def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, Any]:
|
||||
def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, object]:
|
||||
"""Build sanitized litellm_params with merged metadata."""
|
||||
responses_optional_param_keys: Final = set(ResponsesAPIOptionalRequestParams.__annotations__.keys())
|
||||
sanitized: Final[dict[str, Any]] = {
|
||||
sanitized: Final[dict[str, object]] = {
|
||||
key: value for key, value in litellm_params.items() if key not in responses_optional_param_keys
|
||||
}
|
||||
legacy_metadata: Final = litellm_params.get("metadata")
|
||||
existing_litellm_metadata: Final = litellm_params.get("litellm_metadata")
|
||||
merged_litellm_metadata: Final[dict[str, Any]] = {}
|
||||
merged_litellm_metadata: Final[dict[str, object]] = {}
|
||||
if isinstance(legacy_metadata, dict):
|
||||
merged_litellm_metadata.update(legacy_metadata)
|
||||
if isinstance(existing_litellm_metadata, dict):
|
||||
|
|
@ -424,7 +440,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
litellm_params: dict,
|
||||
headers: dict,
|
||||
litellm_logging_obj: "LiteLLMLoggingObj",
|
||||
client: Any | None = None,
|
||||
client: object | None = None,
|
||||
) -> dict:
|
||||
(
|
||||
input_items,
|
||||
|
|
@ -498,9 +514,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
|
||||
@staticmethod
|
||||
def _convert_response_output_to_choices(
|
||||
output_items: list[Any],
|
||||
handle_raw_dict_callback: Callable | None = None,
|
||||
) -> list[Any]:
|
||||
output_items: Sequence[object],
|
||||
handle_raw_dict_callback: Callable[..., tuple["Choices | None", int]] | None = None,
|
||||
) -> list["Choices"]:
|
||||
"""
|
||||
Convert Responses API output items to chat completion choices.
|
||||
|
||||
|
|
@ -529,11 +545,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
choices: Final[list[Choices]] = []
|
||||
index = 0
|
||||
reasoning_content: str | None = None
|
||||
pending_reasoning_item: dict[str, Any] | None = None
|
||||
pending_reasoning_item: _BuiltReasoningItem | None = None
|
||||
|
||||
# Collect all tool calls to put them in a single choice
|
||||
# (Chat Completions API expects all tool calls in one message)
|
||||
accumulated_tool_calls: Final[list[dict[str, Any]]] = []
|
||||
accumulated_tool_calls: Final[list[Mapping[str, object]]] = []
|
||||
tool_call_index = 0
|
||||
|
||||
for item in output_items:
|
||||
|
|
@ -640,7 +656,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
return choices
|
||||
|
||||
@classmethod
|
||||
def _extract_output_from_completed_event(cls, parsed_chunk: dict[str, Any]) -> list[dict[str, Any]] | None:
|
||||
def _extract_output_from_completed_event(cls, parsed_chunk: Mapping[str, object]) -> list[dict[str, object]] | None:
|
||||
response_payload: Final = parsed_chunk.get("response")
|
||||
if not isinstance(response_payload, dict):
|
||||
return None
|
||||
|
|
@ -650,12 +666,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
return cast(list[dict[str, Any]], response_output)
|
||||
|
||||
@classmethod
|
||||
def _recover_output_items_from_raw_sse(cls, raw_sse: str | None) -> list[dict[str, Any]]:
|
||||
def _recover_output_items_from_raw_sse(cls, raw_sse: str | None) -> list[dict[str, object]]:
|
||||
if not raw_sse or not isinstance(raw_sse, str):
|
||||
return []
|
||||
|
||||
recovered_output_items: Final[dict[int, dict[str, Any]]] = {}
|
||||
recovered_text_only_items: Final[dict[int, dict[str, Any]]] = {}
|
||||
recovered_output_items: Final[dict[int, dict[str, object]]] = {}
|
||||
recovered_text_only_items: Final[dict[int, dict[str, object]]] = {}
|
||||
|
||||
for chunk in raw_sse.splitlines():
|
||||
parsed_chunk = parse_sse_json_chunk(chunk)
|
||||
|
|
@ -690,7 +706,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
# but text-only items at indices without a matching OUTPUT_ITEM_DONE
|
||||
# must still be preserved (e.g. multi-output responses where some
|
||||
# indices only emitted OUTPUT_TEXT_DONE).
|
||||
merged_items: Final[dict[int, dict[str, Any]]] = {**recovered_text_only_items}
|
||||
merged_items: Final[dict[int, dict[str, object]]] = {**recovered_text_only_items}
|
||||
merged_items.update(recovered_output_items)
|
||||
|
||||
if merged_items:
|
||||
|
|
@ -699,7 +715,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
return []
|
||||
|
||||
@classmethod
|
||||
def _recover_output_items_from_logging(cls, logging_obj: "LiteLLMLoggingObj") -> list[dict[str, Any]]:
|
||||
def _recover_output_items_from_logging(cls, logging_obj: "LiteLLMLoggingObj") -> list[dict[str, object]]:
|
||||
model_call_details: Final = getattr(logging_obj, "model_call_details", {}) or {}
|
||||
original_response: Final = model_call_details.get("original_response")
|
||||
return cls._recover_output_items_from_raw_sse(original_response)
|
||||
|
|
@ -714,7 +730,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
messages: list["AllMessageValues"],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
encoding: Any,
|
||||
encoding: object,
|
||||
api_key: str | None = None,
|
||||
json_mode: bool | None = None,
|
||||
) -> "ModelResponse":
|
||||
|
|
@ -788,7 +804,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
) -> BaseModelResponseIterator:
|
||||
return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode)
|
||||
|
||||
def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, Any]:
|
||||
def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, object]:
|
||||
if role == "user" or role == "system" or role == "tool":
|
||||
return {"type": "input_text", "text": content}
|
||||
else:
|
||||
|
|
@ -825,13 +841,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
def _convert_content_to_responses_format(
|
||||
self,
|
||||
content: str
|
||||
| list[Any]
|
||||
| list[object]
|
||||
| Iterable[
|
||||
Union["OpenAIMessageContentListBlock", "ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]
|
||||
]
|
||||
| None,
|
||||
role: str,
|
||||
) -> list[dict[str, Any]]:
|
||||
) -> list[dict[str, object]]:
|
||||
"""Convert chat completion content to responses API format"""
|
||||
from litellm.types.llms.openai import ChatCompletionImageObject
|
||||
|
||||
|
|
@ -973,7 +989,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
|
||||
return optional_params
|
||||
|
||||
def _map_reasoning_effort(self, reasoning_effort: str | dict[str, Any]) -> Reasoning | None:
|
||||
def _map_reasoning_effort(self, reasoning_effort: str | Reasoning) -> Reasoning | None:
|
||||
# If dict is passed, convert it directly to Reasoning object
|
||||
if isinstance(reasoning_effort, dict):
|
||||
return Reasoning(**reasoning_effort)
|
||||
|
|
@ -1006,7 +1022,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
def _add_web_search_tool(
|
||||
self,
|
||||
responses_api_request: ResponsesAPIOptionalRequestParams,
|
||||
web_search_options: Any,
|
||||
web_search_options: object,
|
||||
) -> None:
|
||||
"""
|
||||
Add web search tool to responses API request.
|
||||
|
|
@ -1024,14 +1040,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
tools = []
|
||||
responses_api_request["tools"] = tools
|
||||
|
||||
web_search_tool: Final[dict[str, Any]] = {"type": "web_search"}
|
||||
web_search_tool: Final[dict[str, object]] = {"type": "web_search"}
|
||||
if isinstance(web_search_options, dict):
|
||||
web_search_tool.update(web_search_options)
|
||||
|
||||
# Cast to Any to match the expected union type for tools list items
|
||||
tools.append(cast(Any, web_search_tool))
|
||||
|
||||
def _transform_response_format_to_text_format(self, response_format: dict[str, Any] | Any) -> dict[str, Any] | None:
|
||||
def _transform_response_format_to_text_format(self, response_format: object) -> "ResponseText | None":
|
||||
"""
|
||||
Transform Chat Completion response_format parameter to Responses API text.format parameter.
|
||||
|
||||
|
|
@ -1130,7 +1146,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
|
|||
|
||||
|
||||
class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
|
||||
def __init__(self, streaming_response, sync_stream: bool, json_mode: bool | None = False):
|
||||
def __init__(
|
||||
self,
|
||||
streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel"],
|
||||
sync_stream: bool,
|
||||
json_mode: bool | None = False,
|
||||
):
|
||||
super().__init__(streaming_response, sync_stream, json_mode)
|
||||
self._chat_completion_id: str | None = None
|
||||
self._tool_call_index_map: dict[int, int] = {} # mutable-ok: per-stream accumulator state
|
||||
|
|
@ -1387,7 +1408,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
|
|||
finish_reason: Final = "tool_calls" if has_function_calls else "stop"
|
||||
|
||||
# Extract reasoning items with encrypted_content for round-tripping
|
||||
completed_reasoning_items: list[dict[str, Any]] | None = None
|
||||
completed_reasoning_items: list[_BuiltReasoningItem] | None = None
|
||||
for item in output_items:
|
||||
if not isinstance(item, dict) or item.get("type") != "reasoning":
|
||||
continue
|
||||
|
|
@ -1439,7 +1460,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
|
|||
]
|
||||
)
|
||||
|
||||
def chunk_parser(self, chunk: dict) -> "ModelResponseStream":
|
||||
def chunk_parser(self, chunk: dict[str, object]) -> "ModelResponseStream":
|
||||
"""
|
||||
Parse a Responses API streaming chunk and convert to OpenAI format.
|
||||
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_non
|
|||
DEFAULT_HEALTH_CHECK_PROMPT: Final = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"))
|
||||
AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))
|
||||
ROUTER_MAX_FALLBACKS: Final = int(os.getenv("ROUTER_MAX_FALLBACKS", 5))
|
||||
ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000
|
||||
DEFAULT_BATCH_SIZE: Final = int(os.getenv("DEFAULT_BATCH_SIZE", 512))
|
||||
DEFAULT_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_FLUSH_INTERVAL_SECONDS", 5))
|
||||
DEFAULT_S3_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_S3_FLUSH_INTERVAL_SECONDS", 10))
|
||||
|
|
@ -279,6 +280,7 @@ TOOL_POLICY_CACHE_TTL_SECONDS: Final = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECO
|
|||
GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS: Final = int(
|
||||
os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60)
|
||||
)
|
||||
BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS: Final = 25_000
|
||||
# Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger.
|
||||
# Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire.
|
||||
MAX_SIZE_IN_MEMORY_QUEUE: Final = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8)))
|
||||
|
|
@ -470,6 +472,8 @@ EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE: Final = float(
|
|||
### ANTHROPIC CONSTANTS ###
|
||||
ANTHROPIC_TOKEN_COUNTING_BETA_VERSION = os.getenv("ANTHROPIC_TOKEN_COUNTING_BETA_VERSION", "token-counting-2024-11-01")
|
||||
ANTHROPIC_SKILLS_API_BETA_VERSION: Final = "skills-2025-10-02"
|
||||
ANTHROPIC_BATCHES_ROUTE: Final = "/v1/messages/batches"
|
||||
VERTEX_BATCH_PREDICTION_JOBS_ROUTE: Final = "batchPredictionJobs"
|
||||
ANTHROPIC_WEB_SEARCH_TOOL_MAX_USES: Final = {
|
||||
"low": 1,
|
||||
"medium": 5,
|
||||
|
|
@ -1322,6 +1326,8 @@ X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks"
|
|||
LITELLM_METADATA_FIELD: Final = "litellm_metadata"
|
||||
OLD_LITELLM_METADATA_FIELD: Final = "metadata"
|
||||
RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name"
|
||||
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
|
||||
CONSUMED_REQUEST_TAGS_METADATA_KEY: Final = "_consumed_request_tags"
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
|
||||
LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
|
||||
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (
|
||||
|
|
@ -1477,12 +1483,19 @@ CLOUDZERO_MAX_FETCHED_DATA_RECORDS: Final = int(os.getenv("CLOUDZERO_MAX_FETCHED
|
|||
SPEND_LOG_CLEANUP_JOB_NAME: Final = "spend_log_cleanup"
|
||||
KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job"
|
||||
EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job"
|
||||
WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job"
|
||||
MONTHLY_SPEND_REPORT_JOB_ID: Final = "monthly_spend_report_job"
|
||||
PROMETHEUS_FALLBACK_STATS_JOB_ID: Final = "prometheus_fallback_stats_job"
|
||||
SLACK_DAILY_REPORT_LOCK_ID: Final = "slack_daily_report"
|
||||
SPEND_LOG_RUN_LOOPS: Final = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500))
|
||||
SPEND_LOG_CLEANUP_BATCH_SIZE: Final = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000))
|
||||
SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3))
|
||||
SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = float(
|
||||
os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5)
|
||||
)
|
||||
SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", "300"))
|
||||
SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", "30"))
|
||||
SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = int(os.getenv("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", "100000"))
|
||||
TOOL_SPEND_TOP_TOOLS: Final = 100
|
||||
SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day")
|
||||
SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7))
|
||||
|
|
@ -1492,6 +1505,8 @@ SPEND_LOG_QUEUE_POLL_INTERVAL: Final = float(os.getenv("SPEND_LOG_QUEUE_POLL_INT
|
|||
SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE: Final = int(os.getenv("SPEND_COUNTER_RESEED_LOCKS_MAX_SIZE", 10000))
|
||||
DEFAULT_CRON_JOB_LOCK_TTL_SECONDS: Final = int(os.getenv("DEFAULT_CRON_JOB_LOCK_TTL_SECONDS", 60)) # 1 minute
|
||||
PROXY_BUDGET_RESCHEDULER_MIN_TIME: Final = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597))
|
||||
RESET_BUDGET_JOB_BATCH_SIZE: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_BATCH_SIZE", "500")))
|
||||
RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN: Final = max(1, int(os.getenv("RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN", "100")))
|
||||
PROXY_BATCH_POLLING_INTERVAL: Final = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600))
|
||||
MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50)))
|
||||
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)))
|
||||
|
|
@ -1520,6 +1535,10 @@ APSCHEDULER_REPLACE_EXISTING: Final = os.getenv("APSCHEDULER_REPLACE_EXISTING",
|
|||
"1",
|
||||
] # always replace existing jobs
|
||||
|
||||
# Width of the window scheduled background jobs are spread across, so they do not all fire
|
||||
# on one instant on every replica. Tunable per deployment via general_settings.
|
||||
DEFAULT_STAGGER_WINDOW_SECONDS: Final = 300
|
||||
|
||||
# The number of tag entries are higher than number of user, team entries. This leads to a higher QPS.
|
||||
# This will run tag spcific tasks at a later time to smooth QPS
|
||||
DAILY_TAG_SPEND_BATCH_MULTIPLIER: Final = 2.3
|
||||
|
|
@ -1718,3 +1737,21 @@ BROWSER_SECURITY_HEADERS: Final[frozenset[str]] = frozenset(
|
|||
)
|
||||
|
||||
UNSAFE_PROXY_RESPONSE_HEADERS: Final[frozenset[str]] = HTTP_FRAMING_HEADERS | BROWSER_SECURITY_HEADERS
|
||||
|
||||
# PTU reservation rollup writes rows to LiteLLM_DailyTeamSpend with this
|
||||
# sentinel api_key so PTU flat cost stays distinguishable from real per-request
|
||||
# spend under the table's composite unique constraint.
|
||||
PTU_SENTINEL_API_KEY: Final[str] = "__ptu_flat_cost__"
|
||||
PTU_ROLLUP_JOB_ID: Final[str] = "ptu_flat_cost_rollup_job"
|
||||
PTU_ROLLUP_LOCK_TTL_SECONDS: Final[int] = 900
|
||||
# Furthest back the catch-up pass looks for unpriced PTU days when a deployment
|
||||
# declares no ptu_effective_from, bounding the scan for an open-ended window.
|
||||
PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90
|
||||
# Deployments named in the lapsed-window alert before it is truncated, so a fleet-wide
|
||||
# expiry cannot produce an alert too large for the channel delivering it.
|
||||
PTU_LAPSED_ALERT_LIMIT: Final[int] = 10
|
||||
# Slack allowed when deciding a sentinel row is stale. The row's updated_at and the
|
||||
# run's cutoff are stamped by different hosts, so clock skew between them must not let
|
||||
# one run delete a charge another just wrote. A stale row is hours old and a concurrent
|
||||
# one is seconds old, so a few minutes separates them.
|
||||
PTU_PRUNE_SKEW_GRACE_SECONDS: Final[int] = 300
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ import asyncio
|
|||
import contextvars
|
||||
from collections.abc import Coroutine
|
||||
from functools import partial
|
||||
from typing import Any, Final
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -21,8 +21,10 @@ from litellm.types.llms.openai_evals import (
|
|||
CancelRunResponse,
|
||||
CreateEvalRequest,
|
||||
CreateRunRequest,
|
||||
DataSourceConfig,
|
||||
DeleteEvalResponse,
|
||||
Eval,
|
||||
GraderConfig,
|
||||
ListEvalsParams,
|
||||
ListEvalsResponse,
|
||||
ListRunsParams,
|
||||
|
|
@ -41,13 +43,13 @@ DEFAULT_OPENAI_API_BASE: Final = "https://api.openai.com"
|
|||
|
||||
@client
|
||||
async def acreate_eval(
|
||||
data_source_config: dict[str, Any],
|
||||
testing_criteria: list[dict[str, Any]],
|
||||
data_source_config: DataSourceConfig,
|
||||
testing_criteria: list[GraderConfig],
|
||||
name: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
metadata: dict[str, object] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
|
|
@ -110,17 +112,17 @@ async def acreate_eval(
|
|||
|
||||
@client
|
||||
def create_eval(
|
||||
data_source_config: dict[str, Any],
|
||||
testing_criteria: list[dict[str, Any]],
|
||||
data_source_config: DataSourceConfig,
|
||||
testing_criteria: list[GraderConfig],
|
||||
name: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
metadata: dict[str, object] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
) -> Eval | Coroutine[Any, Any, Eval]:
|
||||
) -> Eval | Coroutine[object, object, Eval]:
|
||||
"""
|
||||
Create a new evaluation
|
||||
|
||||
|
|
@ -231,8 +233,8 @@ async def alist_evals(
|
|||
before: str | None = None,
|
||||
order: str | None = None,
|
||||
order_by: str | None = None,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
|
|
@ -300,12 +302,12 @@ def list_evals(
|
|||
before: str | None = None,
|
||||
order: str | None = None,
|
||||
order_by: str | None = None,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
) -> ListEvalsResponse | Coroutine[Any, Any, ListEvalsResponse]:
|
||||
) -> ListEvalsResponse | Coroutine[object, object, ListEvalsResponse]:
|
||||
"""
|
||||
List all evaluations
|
||||
|
||||
|
|
@ -413,8 +415,8 @@ def list_evals(
|
|||
@client
|
||||
async def aget_eval(
|
||||
eval_id: str,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
|
|
@ -470,12 +472,12 @@ async def aget_eval(
|
|||
@client
|
||||
def get_eval(
|
||||
eval_id: str,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
) -> Eval | Coroutine[Any, Any, Eval]:
|
||||
) -> Eval | Coroutine[object, object, Eval]:
|
||||
"""
|
||||
Get an evaluation by ID
|
||||
|
||||
|
|
@ -564,10 +566,10 @@ def get_eval(
|
|||
async def aupdate_eval(
|
||||
eval_id: str,
|
||||
name: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
metadata: dict[str, object] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
|
|
@ -630,14 +632,14 @@ async def aupdate_eval(
|
|||
def update_eval(
|
||||
eval_id: str,
|
||||
name: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
metadata: dict[str, object] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
) -> Eval | Coroutine[Any, Any, Eval]:
|
||||
) -> Eval | Coroutine[object, object, Eval]:
|
||||
"""
|
||||
Update an evaluation
|
||||
|
||||
|
|
@ -783,8 +785,8 @@ def update_eval(
|
|||
@client
|
||||
async def adelete_eval(
|
||||
eval_id: str,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
|
|
@ -840,12 +842,12 @@ async def adelete_eval(
|
|||
@client
|
||||
def delete_eval(
|
||||
eval_id: str,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
) -> DeleteEvalResponse | Coroutine[Any, Any, DeleteEvalResponse]:
|
||||
) -> DeleteEvalResponse | Coroutine[object, object, DeleteEvalResponse]:
|
||||
"""
|
||||
Delete an evaluation
|
||||
|
||||
|
|
@ -933,8 +935,8 @@ def delete_eval(
|
|||
@client
|
||||
async def acancel_eval(
|
||||
eval_id: str,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
|
|
@ -990,12 +992,12 @@ async def acancel_eval(
|
|||
@client
|
||||
def cancel_eval(
|
||||
eval_id: str,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
) -> CancelEvalResponse | Coroutine[Any, Any, CancelEvalResponse]:
|
||||
) -> CancelEvalResponse | Coroutine[object, object, CancelEvalResponse]:
|
||||
"""
|
||||
Cancel a running evaluation
|
||||
|
||||
|
|
@ -1092,12 +1094,12 @@ def cancel_eval(
|
|||
@client
|
||||
async def acreate_run(
|
||||
eval_id: str,
|
||||
data_source: dict[str, Any],
|
||||
data_source: dict[str, object],
|
||||
name: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
metadata: dict[str, object] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
|
|
@ -1161,16 +1163,16 @@ async def acreate_run(
|
|||
@client
|
||||
def create_run(
|
||||
eval_id: str,
|
||||
data_source: dict[str, Any],
|
||||
data_source: dict[str, object],
|
||||
name: str | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_body: dict[str, Any] | None = None,
|
||||
metadata: dict[str, object] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
extra_body: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
) -> Run | Coroutine[Any, Any, Run]:
|
||||
) -> Run | Coroutine[object, object, Run]:
|
||||
"""
|
||||
Create a new run for an evaluation
|
||||
|
||||
|
|
@ -1280,8 +1282,8 @@ async def alist_runs(
|
|||
after: str | None = None,
|
||||
before: str | None = None,
|
||||
order: str | None = None,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
|
|
@ -1349,12 +1351,12 @@ def list_runs(
|
|||
after: str | None = None,
|
||||
before: str | None = None,
|
||||
order: str | None = None,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
) -> ListRunsResponse | Coroutine[Any, Any, ListRunsResponse]:
|
||||
) -> ListRunsResponse | Coroutine[object, object, ListRunsResponse]:
|
||||
"""
|
||||
List all runs for an evaluation
|
||||
|
||||
|
|
@ -1462,8 +1464,8 @@ def list_runs(
|
|||
async def aget_run(
|
||||
eval_id: str,
|
||||
run_id: str,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
|
|
@ -1522,12 +1524,12 @@ async def aget_run(
|
|||
def get_run(
|
||||
eval_id: str,
|
||||
run_id: str,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
) -> Run | Coroutine[Any, Any, Run]:
|
||||
) -> Run | Coroutine[object, object, Run]:
|
||||
"""
|
||||
Get a specific run
|
||||
|
||||
|
|
@ -1618,8 +1620,8 @@ def get_run(
|
|||
async def acancel_run(
|
||||
eval_id: str,
|
||||
run_id: str,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
|
|
@ -1678,12 +1680,12 @@ async def acancel_run(
|
|||
def cancel_run(
|
||||
eval_id: str,
|
||||
run_id: str,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
) -> CancelRunResponse | Coroutine[Any, Any, CancelRunResponse]:
|
||||
) -> CancelRunResponse | Coroutine[object, object, CancelRunResponse]:
|
||||
"""
|
||||
Cancel a running run
|
||||
|
||||
|
|
@ -1783,8 +1785,8 @@ def cancel_run(
|
|||
async def adelete_run(
|
||||
eval_id: str,
|
||||
run_id: str,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
|
|
@ -1843,12 +1845,12 @@ async def adelete_run(
|
|||
def delete_run(
|
||||
eval_id: str,
|
||||
run_id: str,
|
||||
extra_headers: dict[str, Any] | None = None,
|
||||
extra_query: dict[str, Any] | None = None,
|
||||
extra_headers: dict[str, object] | None = None,
|
||||
extra_query: dict[str, object] | None = None,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
custom_llm_provider: str | None = None,
|
||||
**kwargs,
|
||||
) -> RunDeleteResponse | Coroutine[Any, Any, RunDeleteResponse]:
|
||||
) -> RunDeleteResponse | Coroutine[object, object, RunDeleteResponse]:
|
||||
"""
|
||||
Delete a run
|
||||
|
||||
|
|
|
|||
|
|
@ -6,10 +6,11 @@ import asyncio
|
|||
import base64
|
||||
import os
|
||||
from collections.abc import Awaitable, Callable, Generator
|
||||
from datetime import timedelta
|
||||
from typing import Any, Final, TypeVar
|
||||
|
||||
import httpx
|
||||
from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters
|
||||
from mcp import ClientSession, McpError, ReadResourceResult, Resource, StdioServerParameters
|
||||
from mcp.client.sse import sse_client
|
||||
from mcp.client.stdio import stdio_client
|
||||
|
||||
|
|
@ -69,6 +70,29 @@ def _first_non_cancelled_cause(exc: BaseException) -> BaseException | None:
|
|||
return None
|
||||
|
||||
|
||||
_SDK_READ_TIMEOUT_CODE: Final = int(httpx.codes.REQUEST_TIMEOUT)
|
||||
"""The code the MCP SDK puts on its own elapsed read timeout, an HTTP status in a field that
|
||||
otherwise carries JSON-RPC error codes."""
|
||||
|
||||
|
||||
def _as_read_timeout(exc: BaseException) -> TimeoutError | None:
|
||||
"""The session read timeout elapsing, re-expressed as a ``TimeoutError``, or ``None``.
|
||||
|
||||
The SDK reports its own elapsed read timeout as ``McpError`` carrying an HTTP status code in a
|
||||
field that otherwise holds JSON-RPC error codes, and it relays an upstream's JSON-RPC error
|
||||
through that same class and field. The numeric code alone therefore cannot separate the two, and
|
||||
an upstream answering with application code 408 would be reported as a gateway timeout it never
|
||||
caused. The SDK raises its own from inside an ``except TimeoutError``, so the elapsed timeout is
|
||||
on the context chain, while a relayed error is built from a received message and has no such
|
||||
chain; that is the discriminator.
|
||||
"""
|
||||
if not isinstance(exc, McpError) or exc.error.code != _SDK_READ_TIMEOUT_CODE:
|
||||
return None
|
||||
if not isinstance(exc.__context__, TimeoutError):
|
||||
return None
|
||||
return TimeoutError(exc.error.message)
|
||||
|
||||
|
||||
TSessionResult = TypeVar("TSessionResult")
|
||||
|
||||
|
||||
|
|
@ -347,7 +371,14 @@ class MCPClient:
|
|||
session_kwargs["elicitation_callback"] = self._elicitation_callback
|
||||
if self._logging_callback is not None:
|
||||
session_kwargs["logging_callback"] = self._logging_callback
|
||||
session_ctx: Final = ClientSession(read_stream, write_stream, **session_kwargs)
|
||||
# The SDK drops a response stream that ends without a JSON-RPC reply, so nothing else
|
||||
# ever fails the request.
|
||||
session_ctx: Final = ClientSession(
|
||||
read_stream,
|
||||
write_stream,
|
||||
read_timeout_seconds=timedelta(seconds=self.timeout),
|
||||
**session_kwargs,
|
||||
)
|
||||
session: Final = await session_ctx.__aenter__()
|
||||
try:
|
||||
init_result: Final = await session.initialize()
|
||||
|
|
@ -390,7 +421,16 @@ class MCPClient:
|
|||
self._last_initialize_instructions = None
|
||||
transport_ctx, http_client = self._create_transport_context()
|
||||
return await self._execute_session_operation(transport_ctx, operation)
|
||||
except Exception:
|
||||
except Exception as e:
|
||||
read_timeout: Final = _as_read_timeout(e)
|
||||
if read_timeout is not None:
|
||||
verbose_logger.warning(
|
||||
"MCP client timed out after %ss waiting for %s to answer; the server accepted the "
|
||||
"request and ended its response stream without a JSON-RPC reply",
|
||||
self.timeout,
|
||||
self.server_url or "stdio",
|
||||
)
|
||||
raise read_timeout from e
|
||||
_log: Final = verbose_logger.debug if quiet_on_error else verbose_logger.warning
|
||||
_log("MCP client run_with_session failed for %s", self.server_url or "stdio")
|
||||
raise
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
import json
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from typing import Any, Final, cast
|
||||
from typing import Any, Final, TypedDict, cast
|
||||
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
from litellm import verbose_logger
|
||||
from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema
|
||||
|
|
@ -28,6 +30,19 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
|
||||
class _GenAITextPart(TypedDict, total=False):
|
||||
text: ReadOnly[str]
|
||||
|
||||
|
||||
class _GenAISystemInstruction(TypedDict, total=False):
|
||||
parts: ReadOnly[list[_GenAITextPart]]
|
||||
|
||||
|
||||
class _GenAIPart(TypedDict, total=False):
|
||||
text: ReadOnly[str]
|
||||
functionCall: ReadOnly[dict[str, object]]
|
||||
|
||||
|
||||
class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
||||
"""
|
||||
Wrapper for streaming Google GenAI generate_content responses.
|
||||
|
|
@ -36,9 +51,9 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
|
||||
sent_first_chunk: bool = False
|
||||
# State tracking for accumulating partial tool calls
|
||||
accumulated_tool_calls: dict[str, dict[str, Any]]
|
||||
accumulated_tool_calls: dict[str, dict[str, str]]
|
||||
|
||||
def __init__(self, completion_stream: Any):
|
||||
def __init__(self, completion_stream: object):
|
||||
self.sent_first_chunk = False
|
||||
self.accumulated_tool_calls = {}
|
||||
self._returned_response = False
|
||||
|
|
@ -85,7 +100,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
# After the stream is exhausted, check for any remaining accumulated tool calls
|
||||
if self.accumulated_tool_calls:
|
||||
try:
|
||||
parts: Final = []
|
||||
parts: Final[list[_GenAIPart]] = []
|
||||
for (
|
||||
tool_call_index,
|
||||
tool_call_data,
|
||||
|
|
@ -94,7 +109,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
# For tool calls with no arguments, accumulated_args will be "", which is not valid JSON.
|
||||
# We default to an empty JSON object in this case.
|
||||
parsed_args = json.loads(tool_call_data["arguments"] or "{}")
|
||||
function_call_part = {
|
||||
function_call_part: _GenAIPart = {
|
||||
"functionCall": {
|
||||
"name": tool_call_data["name"] or "undefined_tool_name",
|
||||
"args": parsed_args,
|
||||
|
|
@ -110,7 +125,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper):
|
|||
tool_call_data["arguments"],
|
||||
)
|
||||
if parts:
|
||||
final_chunk: Final = {
|
||||
final_chunk: Final[dict[str, object]] = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"parts": parts, "role": "model"},
|
||||
|
|
@ -273,9 +288,9 @@ class GoogleGenAIAdapter:
|
|||
|
||||
def _add_generic_litellm_params_to_request(
|
||||
self,
|
||||
completion_request_dict: dict[str, Any],
|
||||
completion_request_dict: dict[str, object],
|
||||
litellm_params: GenericLiteLLMParams | None = None,
|
||||
) -> dict:
|
||||
) -> dict[str, object]:
|
||||
"""Add generic litellm params to request. e.g add api_base, api_key, api_version, etc.
|
||||
|
||||
Args:
|
||||
|
|
@ -295,7 +310,7 @@ class GoogleGenAIAdapter:
|
|||
|
||||
def translate_completion_output_params_streaming(
|
||||
self,
|
||||
completion_stream: Any,
|
||||
completion_stream: object,
|
||||
) -> AsyncIterator[bytes] | None:
|
||||
"""Transform streaming completion output to Google GenAI format"""
|
||||
google_genai_wrapper: Final = GoogleGenAIStreamWrapper(completion_stream=completion_stream)
|
||||
|
|
@ -307,12 +322,12 @@ class GoogleGenAIAdapter:
|
|||
tools: list[dict[str, Any]],
|
||||
) -> list[ChatCompletionToolParam]:
|
||||
"""Transform Google GenAI tools to OpenAI tools format"""
|
||||
openai_tools: Final[list[dict[str, Any]]] = []
|
||||
openai_tools: Final[list[dict[str, object]]] = []
|
||||
|
||||
for tool in tools:
|
||||
if "functionDeclarations" in tool:
|
||||
for func_decl in tool["functionDeclarations"]:
|
||||
function_chunk: dict[str, Any] = {
|
||||
function_chunk: dict[str, object] = {
|
||||
"name": func_decl.get("name", ""),
|
||||
}
|
||||
|
||||
|
|
@ -321,7 +336,7 @@ class GoogleGenAIAdapter:
|
|||
if "parametersJsonSchema" in func_decl:
|
||||
function_chunk["parameters"] = func_decl["parametersJsonSchema"]
|
||||
|
||||
openai_tool = {"type": "function", "function": function_chunk}
|
||||
openai_tool: dict[str, object] = {"type": "function", "function": function_chunk}
|
||||
openai_tools.append(openai_tool)
|
||||
|
||||
# normalize the tool schemas
|
||||
|
|
@ -345,7 +360,7 @@ class GoogleGenAIAdapter:
|
|||
def _transform_contents_to_messages(
|
||||
self,
|
||||
contents: list[dict[str, Any]],
|
||||
system_instruction: dict[str, Any] | None = None,
|
||||
system_instruction: _GenAISystemInstruction | None = None,
|
||||
) -> list[AllMessageValues]:
|
||||
"""Transform Google GenAI contents to OpenAI messages format"""
|
||||
messages: Final[list[AllMessageValues]] = []
|
||||
|
|
@ -461,7 +476,7 @@ class GoogleGenAIAdapter:
|
|||
def translate_completion_to_generate_content(
|
||||
self,
|
||||
response: ModelResponse,
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Transform litellm completion response to Google GenAI generate_content format
|
||||
|
||||
|
|
@ -490,7 +505,7 @@ class GoogleGenAIAdapter:
|
|||
parts = [{"text": message_content}] if message_content else []
|
||||
|
||||
# Create Google GenAI format response
|
||||
generate_content_response: Final[dict[str, Any]] = {
|
||||
generate_content_response: Final[dict[str, object]] = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"parts": parts, "role": "model"},
|
||||
|
|
@ -524,7 +539,7 @@ class GoogleGenAIAdapter:
|
|||
self,
|
||||
response: ModelResponse | ModelResponseStream,
|
||||
wrapper: GoogleGenAIStreamWrapper,
|
||||
) -> dict[str, Any] | None:
|
||||
) -> dict[str, object] | None:
|
||||
"""
|
||||
Transform streaming litellm completion chunk to Google GenAI generate_content format
|
||||
|
||||
|
|
@ -560,7 +575,7 @@ class GoogleGenAIAdapter:
|
|||
return None
|
||||
|
||||
# Create Google GenAI streaming format response
|
||||
streaming_chunk: Final[dict[str, Any]] = {
|
||||
streaming_chunk: Final[dict[str, object]] = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"parts": parts, "role": "model"},
|
||||
|
|
@ -597,9 +612,9 @@ class GoogleGenAIAdapter:
|
|||
def _transform_openai_message_to_google_genai_parts(
|
||||
self,
|
||||
message: Any,
|
||||
) -> list[dict[str, Any]]:
|
||||
) -> list[_GenAIPart]:
|
||||
"""Transform OpenAI message to Google GenAI parts format"""
|
||||
parts: Final[list[dict[str, Any]]] = []
|
||||
parts: Final[list[_GenAIPart]] = []
|
||||
|
||||
# Add text content if present
|
||||
if hasattr(message, "content") and message.content:
|
||||
|
|
@ -614,7 +629,7 @@ class GoogleGenAIAdapter:
|
|||
except json.JSONDecodeError:
|
||||
args = {}
|
||||
|
||||
function_call_part = {
|
||||
function_call_part: _GenAIPart = {
|
||||
"functionCall": {
|
||||
"name": tool_call.function.name or "undefined_tool_name",
|
||||
"args": args,
|
||||
|
|
@ -626,14 +641,14 @@ class GoogleGenAIAdapter:
|
|||
|
||||
def _transform_openai_delta_to_google_genai_parts_with_accumulation(
|
||||
self, delta: Any, wrapper: GoogleGenAIStreamWrapper
|
||||
) -> list[dict[str, Any]]:
|
||||
) -> list[_GenAIPart]:
|
||||
"""Transforms OpenAI delta to Google GenAI parts, accumulating streaming tool calls."""
|
||||
|
||||
# 1. Initialize wrapper state if it doesn't exist
|
||||
if not hasattr(wrapper, "accumulated_tool_calls"):
|
||||
wrapper.accumulated_tool_calls = {}
|
||||
|
||||
parts: Final[list[dict[str, Any]]] = []
|
||||
parts: Final[list[_GenAIPart]] = []
|
||||
|
||||
if hasattr(delta, "content") and delta.content:
|
||||
parts.append({"text": delta.content})
|
||||
|
|
@ -686,7 +701,7 @@ class GoogleGenAIAdapter:
|
|||
# The part will be created by a later chunk that brings the name.
|
||||
if accumulated_name:
|
||||
# If successful, create the part and clean up
|
||||
function_call_part = {"functionCall": {"name": accumulated_name, "args": parsed_args}}
|
||||
function_call_part: _GenAIPart = {"functionCall": {"name": accumulated_name, "args": parsed_args}}
|
||||
parts.append(function_call_part)
|
||||
|
||||
# Remove the completed tool call from the accumulator
|
||||
|
|
|
|||
|
|
@ -315,7 +315,12 @@ def image_generation(
|
|||
or get_secret_str("AZURE_API_KEY")
|
||||
)
|
||||
|
||||
azure_ad_token: Final = optional_params.pop("azure_ad_token", None) or get_secret_str("AZURE_AD_TOKEN")
|
||||
azure_ad_token_param: Final = optional_params.pop("azure_ad_token", None)
|
||||
azure_ad_token: Final = (
|
||||
azure_ad_token_param
|
||||
if isinstance(azure_ad_token_param, str) and azure_ad_token_param
|
||||
else get_secret_str("AZURE_AD_TOKEN")
|
||||
)
|
||||
|
||||
# Create azure_ad_token_provider from tenant_id, client_id, client_secret if not already provided
|
||||
if azure_ad_token_provider is None:
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from datetime import timedelta
|
|||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
|
||||
from openai import APIError
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
import litellm.litellm_core_utils
|
||||
|
|
@ -16,7 +17,7 @@ import litellm.litellm_core_utils.litellm_logging
|
|||
import litellm.types
|
||||
from litellm._logging import verbose_logger, verbose_proxy_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.constants import HOURS_IN_A_DAY
|
||||
from litellm.constants import HOURS_IN_A_DAY, SLACK_DAILY_REPORT_LOCK_ID
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
from litellm.integrations.SlackAlerting.budget_alert_types import get_budget_alert_type
|
||||
from litellm.integrations.SlackAlerting.hanging_request_check import (
|
||||
|
|
@ -33,10 +34,14 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
from litellm.proxy._types import (
|
||||
AlertType,
|
||||
CallInfo,
|
||||
InvitationModel,
|
||||
InvitationNew,
|
||||
Litellm_EntityType,
|
||||
UserAPIKeyAuth,
|
||||
VirtualKeyEvent,
|
||||
WebhookEvent,
|
||||
)
|
||||
from litellm.repositories.table_repositories import InvitationLinkRepository
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
from litellm.types.integrations.slack_alerting import *
|
||||
|
|
@ -46,6 +51,7 @@ from .batching_handler import send_to_webhook, squash_payloads
|
|||
from .utils import process_slack_alerting_variables
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
|
||||
from litellm.router import Router as _Router
|
||||
|
||||
Router = _Router
|
||||
|
|
@ -1081,6 +1087,44 @@ Model Info:
|
|||
if email_logo_url is not None or email_support_contact is not None:
|
||||
raise ValueError(f"Trying to Customize Email Alerting\n {CommonProxyErrors.not_premium_user.value}")
|
||||
|
||||
async def _construct_user_invitation_link(self, recipient_user_id: str | None, base_url: str) -> str:
|
||||
from litellm.proxy.management_helpers.user_invitation import (
|
||||
create_invitation_for_user,
|
||||
)
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
if recipient_user_id is None or prisma_client is None:
|
||||
return base_url
|
||||
|
||||
try:
|
||||
existing_invitations: Final = TypeAdapter(list[InvitationModel]).validate_python(
|
||||
await InvitationLinkRepository(prisma_client).table.find_many( # pyright: ignore[reportAny] # untyped prisma boundary (any-ok), result validated by TypeAdapter
|
||||
where={"user_id": recipient_user_id}, # mutable-ok: prisma find_many requires a dict where filter
|
||||
order={"created_at": "desc"}, # mutable-ok: prisma find_many requires a dict order arg
|
||||
),
|
||||
from_attributes=True,
|
||||
)
|
||||
invitation: Final = (
|
||||
existing_invitations[0]
|
||||
if existing_invitations
|
||||
else TypeAdapter(InvitationModel).validate_python(
|
||||
await create_invitation_for_user(
|
||||
data=InvitationNew(user_id=recipient_user_id),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_id=recipient_user_id),
|
||||
),
|
||||
from_attributes=True,
|
||||
)
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # best-effort link build; any DB/creation failure falls back to base_url
|
||||
verbose_proxy_logger.error(
|
||||
"Error creating invitation link for user_id %s: %s",
|
||||
recipient_user_id,
|
||||
str(e),
|
||||
)
|
||||
return base_url
|
||||
|
||||
return f"{base_url.rstrip('/')}/ui/onboarding?invitation_id={invitation.id}"
|
||||
|
||||
async def send_key_created_or_user_invited_email(self, webhook_event: WebhookEvent) -> bool:
|
||||
try:
|
||||
from litellm.proxy.utils import send_email
|
||||
|
|
@ -1139,11 +1183,14 @@ Model Info:
|
|||
team_row: Final = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id})
|
||||
if team_row is not None:
|
||||
team_name = team_row.team_alias or "-"
|
||||
invitation_link: Final = await self._construct_user_invitation_link(
|
||||
recipient_user_id=recipient_user_id, base_url=base_url
|
||||
)
|
||||
email_html_content = USER_INVITED_EMAIL_TEMPLATE.format(
|
||||
email_logo_url=email_logo_url,
|
||||
recipient_email=recipient_email,
|
||||
team_name=team_name,
|
||||
base_url=base_url,
|
||||
base_url=invitation_link,
|
||||
email_support_contact=email_support_contact,
|
||||
)
|
||||
else:
|
||||
|
|
@ -1530,7 +1577,11 @@ Model Info:
|
|||
except Exception:
|
||||
pass
|
||||
|
||||
async def _run_scheduler_helper(self, llm_router) -> bool:
|
||||
async def _run_scheduler_helper(
|
||||
self,
|
||||
llm_router,
|
||||
pod_lock_manager: "PodLockManager | None" = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Returns:
|
||||
- True -> report sent
|
||||
|
|
@ -1555,6 +1606,16 @@ Model Info:
|
|||
interval_seconds: Final = self.alerting_args.daily_report_frequency
|
||||
|
||||
if current_time - report_sent >= interval_seconds:
|
||||
if (
|
||||
pod_lock_manager is not None
|
||||
and (
|
||||
await pod_lock_manager.acquire_lock(
|
||||
cronjob_id=SLACK_DAILY_REPORT_LOCK_ID, ttl=interval_seconds, allow_reentrant=False
|
||||
)
|
||||
)
|
||||
is False
|
||||
):
|
||||
return False
|
||||
# Sneak in the reporting logic here
|
||||
await self.send_daily_reports(router=llm_router)
|
||||
# Also, don't forget to update the report_sent time after sending the report!
|
||||
|
|
@ -1566,7 +1627,11 @@ Model Info:
|
|||
|
||||
return report_sent_bool
|
||||
|
||||
async def _run_scheduled_daily_report(self, llm_router: Any | None = None):
|
||||
async def _run_scheduled_daily_report(
|
||||
self,
|
||||
llm_router: Any | None = None,
|
||||
pod_lock_manager: "PodLockManager | None" = None,
|
||||
):
|
||||
"""
|
||||
If 'daily_reports' enabled
|
||||
|
||||
|
|
@ -1579,7 +1644,7 @@ Model Info:
|
|||
|
||||
if "daily_reports" in self.alert_types:
|
||||
while True:
|
||||
await self._run_scheduler_helper(llm_router=llm_router)
|
||||
await self._run_scheduler_helper(llm_router=llm_router, pod_lock_manager=pod_lock_manager)
|
||||
interval = random.randint(
|
||||
self.alerting_args.report_check_interval - 3,
|
||||
self.alerting_args.report_check_interval + 3,
|
||||
|
|
|
|||
|
|
@ -382,19 +382,23 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
model: str,
|
||||
custom_llm_provider: str | None,
|
||||
tools: list | None = None,
|
||||
enable_prompt_caching: bool | None = None,
|
||||
) -> list[CacheControlInjectionPoint]:
|
||||
"""Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on.
|
||||
|
||||
Caches the system prompt and the trailing turn, so the stable prefix
|
||||
(system + tools + history) is reused while the breakpoint advances with
|
||||
the conversation. Returns [] (stand down) when the flag is off, the
|
||||
provider does not consume cache_control breakpoints (only anthropic /
|
||||
bedrock do), the model lacks prompt-caching support, or the request
|
||||
already carries client-supplied cache_control.
|
||||
``enable_prompt_caching`` is the per-request override (stamped from key
|
||||
metadata by the proxy); True turns auto-injection on for this request
|
||||
even when the global flag is off. Caches the system prompt and the
|
||||
trailing turn, so the stable prefix (system + tools + history) is
|
||||
reused while the breakpoint advances with the conversation. Returns []
|
||||
(stand down) when neither flag is on, the provider does not consume
|
||||
cache_control breakpoints (only anthropic / bedrock do), the model
|
||||
lacks prompt-caching support, or the request already carries
|
||||
client-supplied cache_control.
|
||||
"""
|
||||
import litellm
|
||||
|
||||
if litellm.enable_anthropic_prompt_caching is not True:
|
||||
if litellm.enable_anthropic_prompt_caching is not True and enable_prompt_caching is not True:
|
||||
return []
|
||||
|
||||
provider = custom_llm_provider
|
||||
|
|
@ -433,6 +437,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
model: str,
|
||||
custom_llm_provider: str | None,
|
||||
tools: list | None = None,
|
||||
enable_prompt_caching: bool | None = None,
|
||||
) -> None:
|
||||
"""For /chat/completions: resolve the injection points the request should carry.
|
||||
|
||||
|
|
@ -458,6 +463,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
tools=tools,
|
||||
enable_prompt_caching=enable_prompt_caching,
|
||||
)
|
||||
if points:
|
||||
non_default_params["cache_control_injection_points"] = points
|
||||
|
|
@ -478,12 +484,17 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
judgment happens once per request; points a prior pass wrote back
|
||||
carry the judged stamp and are never re-judged (see
|
||||
``_should_stand_down``). When none are configured but
|
||||
``litellm.enable_anthropic_prompt_caching`` is on, synthesize default
|
||||
breakpoints for the native /v1/messages path. Pops the key from kwargs;
|
||||
``litellm.enable_anthropic_prompt_caching`` or the per-request
|
||||
``enable_prompt_caching`` kwarg (stamped from key metadata) is on,
|
||||
synthesize default breakpoints for the native /v1/messages path. Pops
|
||||
both keys from kwargs;
|
||||
if remaining (non-message) points exist they are written back so
|
||||
downstream transforms can handle them.
|
||||
"""
|
||||
typed_messages = cast(list[AllMessageValues], messages) # cast-ok: Anthropic-shaped dicts from v1/messages
|
||||
enable_prompt_caching: Final = cast( # cast-ok: kwargs is untyped; key stamped as bool by the proxy
|
||||
bool | None, kwargs.pop("enable_prompt_caching", None)
|
||||
)
|
||||
configured: Final = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list
|
||||
list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None)
|
||||
)
|
||||
|
|
@ -497,6 +508,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
tools=tools,
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
enable_prompt_caching=enable_prompt_caching,
|
||||
)
|
||||
if not injection_points:
|
||||
return messages, system
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from typing_extensions import override
|
||||
|
|
@ -12,7 +13,7 @@ from litellm.litellm_core_utils.redact_messages import (
|
|||
should_redact_message_logging,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall, StandardLoggingPayload
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.trace import Span
|
||||
|
|
@ -22,6 +23,7 @@ from litellm.integrations._types.open_inference import (
|
|||
ImageAttributes,
|
||||
MessageAttributes,
|
||||
MessageContentAttributes,
|
||||
OpenInferenceMimeTypeValues,
|
||||
OpenInferenceSpanKindValues,
|
||||
SpanAttributes,
|
||||
ToolCallAttributes,
|
||||
|
|
@ -480,6 +482,7 @@ def set_attributes(span: "Span", kwargs, response_obj, attributes: type[BaseLLMO
|
|||
response_obj_for_attrs,
|
||||
slp,
|
||||
)
|
||||
_safe_emit("mcp tool attrs", _maybe_set_mcp_tool_attrs, span, kwargs, slp, response_obj_for_attrs)
|
||||
|
||||
|
||||
def _sanitize_optional_params(optional_params: dict | None) -> dict:
|
||||
|
|
@ -538,9 +541,12 @@ def _set_request_attributes(
|
|||
if optional_params.get("user"):
|
||||
safe_set_attribute(span, "llm.user", optional_params.get("user"))
|
||||
|
||||
if response_obj and response_obj.get("id"):
|
||||
if not hasattr(response_obj, "get"):
|
||||
return
|
||||
|
||||
if response_obj.get("id"):
|
||||
safe_set_attribute(span, "llm.response.id", response_obj.get("id"))
|
||||
if response_obj and response_obj.get("model"):
|
||||
if response_obj.get("model"):
|
||||
safe_set_attribute(span, "llm.response.model", response_obj.get("model"))
|
||||
|
||||
|
||||
|
|
@ -588,6 +594,8 @@ def _coerce_response_obj_for_attrs(response_obj):
|
|||
- dicts and Pydantic models that already expose `.get` are returned
|
||||
unchanged (preserves all current behavior, including the Responses API
|
||||
flow which relies on Pydantic attribute access).
|
||||
- Pydantic models without `.get` (e.g. the MCP SDK's `CallToolResult`,
|
||||
logged for `call_mcp_tool` spans) are dumped to a dict.
|
||||
- `httpx.Response` and other text-only responses (passthrough routes)
|
||||
are JSON-decoded so the standard extraction paths can read fields like
|
||||
`id`, `model`, and `usage`. On failure the original object is returned
|
||||
|
|
@ -595,6 +603,9 @@ def _coerce_response_obj_for_attrs(response_obj):
|
|||
"""
|
||||
if response_obj is None or hasattr(response_obj, "get"):
|
||||
return response_obj
|
||||
dumped: Final = _to_plain_dict(response_obj)
|
||||
if isinstance(dumped, dict):
|
||||
return dumped
|
||||
text: Final = getattr(response_obj, "text", None)
|
||||
if isinstance(text, str) and text:
|
||||
try:
|
||||
|
|
@ -1058,3 +1069,65 @@ def _parse_passthrough_response(raw_response_obj, coerced_response_obj, kwargs):
|
|||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _maybe_set_mcp_tool_attrs(
|
||||
span: "Span",
|
||||
kwargs: Mapping[str, object],
|
||||
standard_logging_payload: StandardLoggingPayload | None,
|
||||
coerced_response_obj: object,
|
||||
) -> None:
|
||||
"""Render `call_mcp_tool` spans as OpenInference TOOL spans.
|
||||
|
||||
MCP tool calls carry neither `messages` nor `choices`, so the generic
|
||||
extraction paths leave Input/Output blank. The tool name and arguments live
|
||||
in `metadata.mcp_tool_call_metadata`; the result is an MCP `CallToolResult`
|
||||
whose `content` is a list of typed parts.
|
||||
"""
|
||||
if standard_logging_payload is None:
|
||||
return
|
||||
if standard_logging_payload.get("call_type") != CallTypes.call_mcp_tool.value:
|
||||
return
|
||||
|
||||
metadata: Final = standard_logging_payload.get("metadata")
|
||||
mcp_meta: Final[StandardLoggingMCPToolCall | None] = metadata.get("mcp_tool_call_metadata") if metadata else None
|
||||
if mcp_meta is None:
|
||||
return
|
||||
|
||||
tool_name: Final = mcp_meta.get("name") or mcp_meta.get("namespaced_tool_name")
|
||||
if tool_name:
|
||||
safe_set_attribute(span, SpanAttributes.TOOL_NAME, tool_name)
|
||||
|
||||
if should_redact_message_logging(kwargs): # pyright: ignore[reportArgumentType] # reads, never mutates
|
||||
return
|
||||
|
||||
arguments: Final[object] = mcp_meta.get("arguments")
|
||||
if arguments is not None:
|
||||
safe_set_attribute(span, SpanAttributes.INPUT_VALUE, safe_dumps(arguments))
|
||||
safe_set_attribute(span, SpanAttributes.INPUT_MIME_TYPE, OpenInferenceMimeTypeValues.JSON.value)
|
||||
|
||||
_set_mcp_tool_output(span, coerced_response_obj)
|
||||
|
||||
|
||||
def _has_only_text_parts(content: object) -> bool:
|
||||
return not isinstance(content, list) or all(_coerce_text([part]) is not None for part in content)
|
||||
|
||||
|
||||
def _set_mcp_tool_output(span: "Span", coerced_response_obj: object) -> None:
|
||||
if not isinstance(coerced_response_obj, Mapping):
|
||||
return
|
||||
|
||||
content: Final[object] = coerced_response_obj.get("content")
|
||||
text: Final[str | None] = _coerce_text(content)
|
||||
if text and _has_only_text_parts(content):
|
||||
safe_set_attribute(span, SpanAttributes.OUTPUT_VALUE, text)
|
||||
safe_set_attribute(span, SpanAttributes.OUTPUT_MIME_TYPE, OpenInferenceMimeTypeValues.TEXT.value)
|
||||
return
|
||||
|
||||
structured: Final[object] = coerced_response_obj.get("structuredContent")
|
||||
payload: Final[object] = content if content else structured if structured is not None else content
|
||||
if payload is None:
|
||||
return
|
||||
|
||||
safe_set_attribute(span, SpanAttributes.OUTPUT_VALUE, safe_dumps(payload))
|
||||
safe_set_attribute(span, SpanAttributes.OUTPUT_MIME_TYPE, OpenInferenceMimeTypeValues.JSON.value)
|
||||
|
|
|
|||
|
|
@ -16,7 +16,10 @@ import asyncio
|
|||
import os
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
|
|
@ -27,6 +30,16 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
)
|
||||
from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
|
||||
|
||||
DEFAULT_AZURE_AUTHORITY_HOST: Final = "https://login.microsoftonline.com"
|
||||
DEFAULT_AZURE_MONITOR_SCOPE: Final = "https://monitor.azure.com/.default"
|
||||
|
||||
MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType(
|
||||
{
|
||||
"login.microsoftonline.com": DEFAULT_AZURE_MONITOR_SCOPE,
|
||||
"login.microsoftonline.us": "https://monitor.azure.us/.default",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class AzureSentinelLogger(CustomBatchLogger):
|
||||
"""
|
||||
|
|
@ -42,6 +55,7 @@ class AzureSentinelLogger(CustomBatchLogger):
|
|||
client_id: str | None = None,
|
||||
client_secret: str | None = None,
|
||||
audit_stream_name: str | None = None,
|
||||
authority_host: str | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
|
|
@ -62,6 +76,10 @@ class AzureSentinelLogger(CustomBatchLogger):
|
|||
If not provided, will use AZURE_SENTINEL_CLIENT_SECRET or AZURE_CLIENT_SECRET env var.
|
||||
audit_stream_name (str, optional): Stream name from DCR for audit logs.
|
||||
If not provided, will use AZURE_SENTINEL_AUDIT_STREAM_NAME env var or the standard stream name.
|
||||
authority_host (str, optional): Microsoft Entra authority host that issues the OAuth2 token,
|
||||
e.g. "https://login.microsoftonline.us" for Azure Government. If not provided, will use
|
||||
AZURE_SENTINEL_AUTHORITY_HOST or AZURE_AUTHORITY_HOST env vars, or default to the Azure
|
||||
Public Cloud authority. The Azure Monitor audience is derived from it.
|
||||
"""
|
||||
self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
|
||||
|
||||
|
|
@ -76,6 +94,12 @@ class AzureSentinelLogger(CustomBatchLogger):
|
|||
resolved_client_secret: Final = (
|
||||
client_secret or os.getenv("AZURE_SENTINEL_CLIENT_SECRET") or os.getenv("AZURE_CLIENT_SECRET")
|
||||
)
|
||||
resolved_authority_host: Final = self._normalize_authority_host(
|
||||
authority_host
|
||||
or os.getenv("AZURE_SENTINEL_AUTHORITY_HOST")
|
||||
or os.getenv("AZURE_AUTHORITY_HOST")
|
||||
or DEFAULT_AZURE_AUTHORITY_HOST
|
||||
)
|
||||
|
||||
if not resolved_dcr_immutable_id:
|
||||
raise ValueError(
|
||||
|
|
@ -119,7 +143,8 @@ class AzureSentinelLogger(CustomBatchLogger):
|
|||
)
|
||||
|
||||
# OAuth2 scope for Azure Monitor
|
||||
self.oauth_scope = "https://monitor.azure.com/.default"
|
||||
self.authority_host = resolved_authority_host
|
||||
self.oauth_scope = self._resolve_oauth_scope(authority_host=resolved_authority_host)
|
||||
self.oauth_token: str | None = None
|
||||
self.oauth_token_expires_at: float | None = None
|
||||
|
||||
|
|
@ -129,6 +154,26 @@ class AzureSentinelLogger(CustomBatchLogger):
|
|||
self.log_queue: list[StandardLoggingPayload] = []
|
||||
self.audit_log_queue: list[StandardAuditLogPayload] = []
|
||||
|
||||
@staticmethod
|
||||
def _normalize_authority_host(authority_host: str) -> str:
|
||||
"""
|
||||
Normalize an authority host into an absolute URL with no trailing slash.
|
||||
|
||||
Accepts the scheme-qualified form litellm documents ("https://login.microsoftonline.us")
|
||||
and the bare-host form the azure-identity AzureAuthorityHosts constants use.
|
||||
"""
|
||||
stripped: Final = authority_host.strip().rstrip("/")
|
||||
return stripped if "://" in stripped else f"https://{stripped}"
|
||||
|
||||
@staticmethod
|
||||
def _resolve_oauth_scope(authority_host: str) -> str:
|
||||
"""
|
||||
Map an authority host to the Azure Monitor Logs Ingestion audience for the same cloud,
|
||||
falling back to the Azure Public Cloud audience for an unrecognized host.
|
||||
"""
|
||||
host: Final = urlparse(authority_host).hostname or ""
|
||||
return MONITOR_SCOPE_BY_AUTHORITY_HOST.get(host, DEFAULT_AZURE_MONITOR_SCOPE)
|
||||
|
||||
@staticmethod
|
||||
def _build_api_endpoint(endpoint: str, dcr_immutable_id: str, stream_name: str) -> str:
|
||||
return f"{endpoint.rstrip('/')}/dataCollectionRules/{dcr_immutable_id}/streams/{stream_name}?api-version=2023-01-01"
|
||||
|
|
@ -150,7 +195,7 @@ class AzureSentinelLogger(CustomBatchLogger):
|
|||
assert self.client_id is not None, "client_id is required"
|
||||
assert self.client_secret is not None, "client_secret is required"
|
||||
|
||||
token_url: Final = f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token"
|
||||
token_url: Final = f"{self.authority_host}/{self.tenant_id}/oauth2/v2.0/token"
|
||||
|
||||
token_data: Final = {
|
||||
"client_id": self.client_id,
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ Usage:
|
|||
|
||||
import os
|
||||
import time
|
||||
from collections.abc import AsyncIterable, Iterable
|
||||
from typing import Final
|
||||
from urllib.parse import urlparse
|
||||
|
||||
|
|
@ -84,7 +85,7 @@ def _mock_http_handler_post(
|
|||
timeout=None,
|
||||
stream=False,
|
||||
files=None,
|
||||
content=None,
|
||||
content: str | bytes | Iterable[bytes] | AsyncIterable[bytes] | None = None,
|
||||
logging_obj=None,
|
||||
):
|
||||
"""Monkey-patched HTTPHandler.post that intercepts Braintrust calls with endpoint-specific responses."""
|
||||
|
|
|
|||
|
|
@ -54,7 +54,7 @@ USER_INVITED_EMAIL_TEMPLATE: Final = """
|
|||
|
||||
You were invited to use OpenAI Proxy API for team {team_name} <br /> <br />
|
||||
|
||||
<a href="{base_url}" style="display: inline-block; padding: 10px 20px; background-color: #87ceeb; color: #fff; text-decoration: none; border-radius: 20px;">Get Started here</a> <br /> <br />
|
||||
<a href="{base_url}" style="display: inline-block; padding: 10px 20px; background-color: #87ceeb; color: #fff; text-decoration: none; border-radius: 20px;">Accept Invitation</a> <br /> <br />
|
||||
|
||||
|
||||
If you have any questions, please send an email to {email_support_contact} <br /> <br />
|
||||
|
|
|
|||
|
|
@ -131,7 +131,7 @@ USER_INVITATION_EMAIL_TEMPLATE: Final = """
|
|||
</div>
|
||||
|
||||
<div class="btn-container">
|
||||
<a href="{base_url}" class="btn">Accept Invitation</a>
|
||||
<a href="{invitation_link}" class="btn">Accept Invitation</a>
|
||||
</div>
|
||||
|
||||
<div class="quickstart">
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ Usage:
|
|||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterable, Iterable
|
||||
from typing import Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -113,7 +114,7 @@ async def _mock_async_handler_delete(
|
|||
headers=None,
|
||||
timeout=None,
|
||||
stream=False,
|
||||
content=None,
|
||||
content: str | bytes | Iterable[bytes] | AsyncIterable[bytes] | None = None,
|
||||
):
|
||||
"""Monkey-patched AsyncHTTPHandler.delete that intercepts GCS calls."""
|
||||
# Only mock GCS API calls
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ import json
|
|||
import os
|
||||
import re
|
||||
import traceback
|
||||
from typing import Any, Final, Literal
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -158,7 +158,7 @@ class GenericAPILogger(CustomBatchLogger):
|
|||
"endpoint not set for GenericAPILogger, GENERIC_LOGGER_ENDPOINT not found in environment variables"
|
||||
)
|
||||
|
||||
self.headers: dict = self._get_headers(headers)
|
||||
self.headers: dict[str, str] = self._get_headers(headers)
|
||||
self.endpoint: str = endpoint
|
||||
self.event_types: list[API_EVENT_TYPES] | None = event_types
|
||||
self.callback_name: str | None = callback_name
|
||||
|
|
@ -248,18 +248,15 @@ class GenericAPILogger(CustomBatchLogger):
|
|||
await asyncio.sleep(delay)
|
||||
|
||||
async def _post_with_retries(self, data: str) -> httpx.Response:
|
||||
post_kwargs: Final[dict[str, Any]] = {
|
||||
"url": self.endpoint,
|
||||
"headers": self.headers,
|
||||
"data": data,
|
||||
}
|
||||
if self.timeout is not None:
|
||||
post_kwargs["timeout"] = self.timeout
|
||||
|
||||
total_attempts: Final = self.max_retries + 1
|
||||
for attempt in range(total_attempts):
|
||||
try:
|
||||
return await self.async_httpx_client.post(**post_kwargs)
|
||||
return await self.async_httpx_client.post(
|
||||
url=self.endpoint,
|
||||
headers=self.headers,
|
||||
data=data,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
except Exception as e:
|
||||
is_last_attempt = attempt == self.max_retries
|
||||
should_retry = self._should_retry_exception(e)
|
||||
|
|
|
|||
|
|
@ -2,8 +2,9 @@
|
|||
# On success, logs events to Langfuse
|
||||
import os
|
||||
import traceback
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Iterable, Mapping
|
||||
from datetime import datetime
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
from packaging.version import Version
|
||||
|
|
@ -30,6 +31,7 @@ from litellm.types.utils import (
|
|||
ImageResponse,
|
||||
ModelResponse,
|
||||
RerankResponse,
|
||||
StandardLoggingMetadata,
|
||||
StandardLoggingPayload,
|
||||
StandardLoggingPromptManagementMetadata,
|
||||
TextCompletionResponse,
|
||||
|
|
@ -46,6 +48,11 @@ else:
|
|||
Langfuse = Any
|
||||
|
||||
|
||||
_DENIED_STEERING_KEYS: Final = frozenset({"headers", "endpoint", "caching_groups", "previous_models"})
|
||||
_NO_METADATA: Final[Mapping[str, Any]] = MappingProxyType({})
|
||||
_REDACTED_PROXY_HEADERS: Final[frozenset[str]] = frozenset({"authorization", "cookie", "referer"})
|
||||
|
||||
|
||||
def _extract_cache_read_input_tokens(usage_obj) -> int:
|
||||
"""
|
||||
Extract cache_read_input_tokens from usage object.
|
||||
|
|
@ -75,6 +82,22 @@ def _extract_cache_read_input_tokens(usage_obj) -> int:
|
|||
return cache_read_input_tokens
|
||||
|
||||
|
||||
def _as_steering_flag(value: object) -> bool:
|
||||
"""A string ``str_to_bool`` does not recognise falls back to its truthiness."""
|
||||
if isinstance(value, str):
|
||||
parsed: Final = str_to_bool(value)
|
||||
return bool(value) if parsed is None else parsed
|
||||
return bool(value)
|
||||
|
||||
|
||||
def _as_steering_key_sequence(value: object) -> tuple[str, ...]:
|
||||
if isinstance(value, str):
|
||||
return tuple(key.strip() for key in value.split(",") if key.strip())
|
||||
if isinstance(value, Iterable):
|
||||
return tuple(str(key) for key in value)
|
||||
return ()
|
||||
|
||||
|
||||
def resolve_langfuse_credentials(
|
||||
langfuse_public_key=None,
|
||||
langfuse_secret=None,
|
||||
|
|
@ -496,16 +519,14 @@ class LangFuseLogger:
|
|||
else []
|
||||
)
|
||||
|
||||
if standard_logging_object is None:
|
||||
end_user_id = None
|
||||
prompt_management_metadata: StandardLoggingPromptManagementMetadata | None = None
|
||||
else:
|
||||
end_user_id = standard_logging_object["metadata"].get("user_api_key_end_user_id", None)
|
||||
|
||||
prompt_management_metadata = cast(
|
||||
StandardLoggingPromptManagementMetadata | None,
|
||||
standard_logging_object["metadata"].get("prompt_management_metadata", None),
|
||||
)
|
||||
allowlisted_metadata: Final[StandardLoggingMetadata | dict[str, Any]] = (
|
||||
standard_logging_object["metadata"] if standard_logging_object is not None else _NO_METADATA
|
||||
)
|
||||
end_user_id: Final = allowlisted_metadata.get("user_api_key_end_user_id", None)
|
||||
prompt_management_metadata: Final[StandardLoggingPromptManagementMetadata | None] = cast(
|
||||
StandardLoggingPromptManagementMetadata | None,
|
||||
allowlisted_metadata.get("prompt_management_metadata", None),
|
||||
)
|
||||
|
||||
# Clean Metadata before logging - never log raw metadata
|
||||
# the raw metadata can contain circular references which leads to infinite recursion
|
||||
|
|
@ -524,12 +545,7 @@ class LangFuseLogger:
|
|||
tags.append(f"{key}:{value}")
|
||||
|
||||
# clean litellm metadata before logging
|
||||
if key in [
|
||||
"headers",
|
||||
"endpoint",
|
||||
"caching_groups",
|
||||
"previous_models",
|
||||
]:
|
||||
if key in _DENIED_STEERING_KEYS:
|
||||
continue
|
||||
else:
|
||||
clean_metadata[key] = value
|
||||
|
|
@ -552,10 +568,10 @@ class LangFuseLogger:
|
|||
# This allows continuing an existing trace while still returning the correct trace_id
|
||||
if existing_trace_id is not None:
|
||||
trace_id = existing_trace_id
|
||||
update_trace_keys: Final = cast(list, clean_metadata.pop("update_trace_keys", []))
|
||||
update_trace_keys: Final = _as_steering_key_sequence(clean_metadata.pop("update_trace_keys", ()))
|
||||
debug: Final = clean_metadata.pop("debug_langfuse", None)
|
||||
mask_input: Final = clean_metadata.pop("mask_input", False)
|
||||
mask_output: Final = clean_metadata.pop("mask_output", False)
|
||||
mask_input: Final = _as_steering_flag(clean_metadata.pop("mask_input", False))
|
||||
mask_output: Final = _as_steering_flag(clean_metadata.pop("mask_output", False))
|
||||
# Look for masking function in the dedicated location first (set by scrub_sensitive_keys_in_metadata)
|
||||
# Fall back to metadata for backwards compatibility
|
||||
masking_function: Final = litellm_params.get("_langfuse_masking_function") or clean_metadata.pop(
|
||||
|
|
@ -614,19 +630,18 @@ class LangFuseLogger:
|
|||
trace_params["output"] = output if not mask_output else "redacted-by-litellm"
|
||||
|
||||
if debug is True or (isinstance(debug, str) and debug.lower() == "true"):
|
||||
if "metadata" in trace_params:
|
||||
# log the raw_metadata in the trace
|
||||
trace_params["metadata"]["metadata_passed_to_litellm"] = metadata
|
||||
else:
|
||||
trace_params["metadata"] = {"metadata_passed_to_litellm": metadata}
|
||||
debug_metadata: Final = {
|
||||
key: value for key, value in metadata.items() if isinstance(value, (str, int, float, bool))
|
||||
}
|
||||
trace_params["metadata"] = {
|
||||
**(trace_params.get("metadata") or _NO_METADATA),
|
||||
"metadata_passed_to_litellm": debug_metadata,
|
||||
}
|
||||
|
||||
cost: Final = kwargs.get("response_cost", None)
|
||||
verbose_logger.debug("trace: %s", cost)
|
||||
|
||||
clean_metadata["litellm_response_cost"] = cost
|
||||
if standard_logging_object is not None:
|
||||
hidden_params: Final = standard_logging_object.get("hidden_params", {})
|
||||
clean_metadata["hidden_params"] = filter_exceptions_from_params(hidden_params)
|
||||
hidden_params: Final = standard_logging_object.get("hidden_params") if standard_logging_object else None
|
||||
|
||||
if (
|
||||
litellm.langfuse_default_tags is not None
|
||||
|
|
@ -638,22 +653,24 @@ class LangFuseLogger:
|
|||
tags.append(f"proxy_base_url:{proxy_base_url}")
|
||||
|
||||
api_base: Final = litellm_params.get("api_base", None)
|
||||
if api_base:
|
||||
clean_metadata["api_base"] = api_base
|
||||
|
||||
vertex_location: Final = kwargs.get("vertex_location", None)
|
||||
if vertex_location:
|
||||
clean_metadata["vertex_location"] = vertex_location
|
||||
|
||||
aws_region_name: Final = kwargs.get("aws_region_name", None)
|
||||
if aws_region_name:
|
||||
clean_metadata["aws_region_name"] = aws_region_name
|
||||
|
||||
candidate_enrichments: Final = (
|
||||
("litellm_response_cost", cost, True),
|
||||
("hidden_params", filter_exceptions_from_params(hidden_params), hidden_params is not None),
|
||||
("api_base", api_base, bool(api_base)),
|
||||
("vertex_location", vertex_location, bool(vertex_location)),
|
||||
("aws_region_name", aws_region_name, bool(aws_region_name)),
|
||||
("cache_hit", kwargs.get("cache_hit") or False, self._supports_tags() and "cache_hit" in kwargs),
|
||||
)
|
||||
enrichments: Final[Mapping[str, Any]] = {
|
||||
key: value for key, value, include in candidate_enrichments if include
|
||||
}
|
||||
|
||||
if self._supports_tags():
|
||||
if "cache_hit" in kwargs:
|
||||
if kwargs["cache_hit"] is None:
|
||||
kwargs["cache_hit"] = False
|
||||
clean_metadata["cache_hit"] = kwargs["cache_hit"]
|
||||
if "cache_hit" in kwargs and kwargs["cache_hit"] is None:
|
||||
kwargs["cache_hit"] = False # rebind-ok: pre-existing normalization other integrations rely on
|
||||
if existing_trace_id is None:
|
||||
trace_params.update({"tags": tags})
|
||||
|
||||
|
|
@ -666,13 +683,13 @@ class LangFuseLogger:
|
|||
if headers:
|
||||
for key, value in headers.items():
|
||||
# these headers can leak our API keys and/or JWT tokens
|
||||
if key.lower() not in ["authorization", "cookie", "referer"]:
|
||||
if key.lower() not in _REDACTED_PROXY_HEADERS:
|
||||
clean_headers[key] = value
|
||||
|
||||
trace: Final[StatefulTraceClient] = self.Langfuse.trace(**trace_params)
|
||||
|
||||
# Log provider specific information as a span
|
||||
log_provider_specific_information_as_span(trace, clean_metadata)
|
||||
log_provider_specific_information_as_span(trace, enrichments)
|
||||
|
||||
# Log guardrail information as a span
|
||||
self._log_guardrail_information_as_span(
|
||||
|
|
@ -745,7 +762,10 @@ class LangFuseLogger:
|
|||
"output": output if not mask_output else "redacted-by-litellm",
|
||||
"usage": usage,
|
||||
"usage_details": usage_details,
|
||||
"metadata": log_requester_metadata(clean_metadata),
|
||||
"metadata": {
|
||||
**log_requester_metadata(redact_user_api_key_info(metadata=allowlisted_metadata)),
|
||||
**enrichments,
|
||||
},
|
||||
"level": level,
|
||||
"version": clean_metadata.pop("version", None),
|
||||
}
|
||||
|
|
@ -1042,7 +1062,7 @@ def _add_prompt_to_generation_params(
|
|||
|
||||
def log_provider_specific_information_as_span(
|
||||
trace,
|
||||
clean_metadata,
|
||||
clean_metadata: Mapping[str, Any],
|
||||
):
|
||||
"""
|
||||
Logs provider-specific information as spans.
|
||||
|
|
@ -1082,7 +1102,7 @@ def log_provider_specific_information_as_span(
|
|||
)
|
||||
|
||||
|
||||
def log_requester_metadata(clean_metadata: dict):
|
||||
def log_requester_metadata(clean_metadata: Mapping[str, Any]):
|
||||
returned_metadata: Final = {}
|
||||
requester_metadata: Final = clean_metadata.get("requester_metadata") or {}
|
||||
for k, v in clean_metadata.items():
|
||||
|
|
|
|||
|
|
@ -90,7 +90,6 @@ class LangfuseOtelLogger(OpenTelemetry):
|
|||
"generation_name": LangfuseSpanAttributes.GENERATION_NAME,
|
||||
"generation_id": LangfuseSpanAttributes.GENERATION_ID,
|
||||
"parent_observation_id": LangfuseSpanAttributes.PARENT_OBSERVATION_ID,
|
||||
"version": LangfuseSpanAttributes.GENERATION_VERSION,
|
||||
"mask_input": LangfuseSpanAttributes.MASK_INPUT,
|
||||
"mask_output": LangfuseSpanAttributes.MASK_OUTPUT,
|
||||
"trace_user_id": LangfuseSpanAttributes.TRACE_USER_ID,
|
||||
|
|
@ -99,13 +98,18 @@ class LangfuseOtelLogger(OpenTelemetry):
|
|||
"trace_name": LangfuseSpanAttributes.TRACE_NAME,
|
||||
"trace_id": LangfuseSpanAttributes.TRACE_ID,
|
||||
"trace_metadata": LangfuseSpanAttributes.TRACE_METADATA,
|
||||
"trace_version": LangfuseSpanAttributes.TRACE_VERSION,
|
||||
"trace_release": LangfuseSpanAttributes.TRACE_RELEASE,
|
||||
"trace_release": LangfuseSpanAttributes.RELEASE,
|
||||
"existing_trace_id": LangfuseSpanAttributes.EXISTING_TRACE_ID,
|
||||
"update_trace_keys": LangfuseSpanAttributes.UPDATE_TRACE_KEYS,
|
||||
"debug_langfuse": LangfuseSpanAttributes.DEBUG_LANGFUSE,
|
||||
}
|
||||
|
||||
version: Final = (
|
||||
metadata.get("trace_version") if metadata.get("trace_version") is not None else metadata.get("version")
|
||||
)
|
||||
if version is not None:
|
||||
safe_set_attribute(span, LangfuseSpanAttributes.VERSION.value, version)
|
||||
|
||||
for key, enum_attr in mapping.items():
|
||||
if key in metadata and metadata[key] is not None:
|
||||
value = metadata[key]
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ making actual network calls.
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncIterable, Iterable
|
||||
from dataclasses import dataclass
|
||||
from datetime import timedelta
|
||||
from typing import Final, cast
|
||||
|
|
@ -140,7 +141,7 @@ def create_mock_client_factory(config: MockClientConfig):
|
|||
stream=False,
|
||||
logging_obj=None,
|
||||
files=None,
|
||||
content=None,
|
||||
content: str | bytes | Iterable[bytes] | AsyncIterable[bytes] | None = None,
|
||||
):
|
||||
"""Monkey-patched AsyncHTTPHandler.post that intercepts API calls."""
|
||||
if isinstance(url, str) and _is_mock_url(url):
|
||||
|
|
@ -193,7 +194,7 @@ def create_mock_client_factory(config: MockClientConfig):
|
|||
timeout=None,
|
||||
stream=False,
|
||||
files=None,
|
||||
content=None,
|
||||
content: str | bytes | Iterable[bytes] | AsyncIterable[bytes] | None = None,
|
||||
logging_obj=None,
|
||||
):
|
||||
"""Monkey-patched HTTPHandler.post that intercepts API calls."""
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
import os
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -37,9 +38,11 @@ from litellm.types.utils import (
|
|||
# OpenTelemetry imports moved to individual functions to avoid import errors when not installed
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from opentelemetry.sdk.trace import TracerProvider as _SDKTracerProvider
|
||||
from opentelemetry.sdk.trace.export import SpanExporter as _SpanExporter
|
||||
from opentelemetry.trace import Context as _Context
|
||||
from opentelemetry.trace import Span as _Span
|
||||
from opentelemetry.trace import SpanKind as _SpanKind
|
||||
from opentelemetry.trace import Tracer as _Tracer
|
||||
|
||||
from litellm.proxy._types import (
|
||||
|
|
@ -61,6 +64,25 @@ else:
|
|||
ManagementEndpointLoggingPayload = Any
|
||||
Context = Any
|
||||
|
||||
|
||||
class _StartSpanRequiredKwargs(TypedDict):
|
||||
name: str
|
||||
start_time: int
|
||||
context: "Context | None"
|
||||
|
||||
|
||||
class _StartSpanKwargs(_StartSpanRequiredKwargs, total=False):
|
||||
kind: "_SpanKind"
|
||||
|
||||
|
||||
class _UsageCompletionTokensView(TypedDict, total=False):
|
||||
completion_tokens: int
|
||||
|
||||
|
||||
class _ResponseWithUsageView(TypedDict, total=False):
|
||||
usage: "_UsageCompletionTokensView | None"
|
||||
|
||||
|
||||
LITELLM_TRACER_NAME: Final = os.getenv("OTEL_TRACER_NAME", "litellm")
|
||||
LITELLM_METER_NAME: Final = os.getenv("LITELLM_METER_NAME", "litellm")
|
||||
LITELLM_LOGGER_NAME: Final = os.getenv("LITELLM_LOGGER_NAME", "litellm")
|
||||
|
|
@ -297,9 +319,9 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
config: OpenTelemetryConfig | None = None,
|
||||
callback_name: str | None = None,
|
||||
# injection points for testing
|
||||
tracer_provider: Any | None = None,
|
||||
logger_provider: Any | None = None,
|
||||
meter_provider: Any | None = None,
|
||||
tracer_provider: object | None = None,
|
||||
logger_provider: object | None = None,
|
||||
meter_provider: object | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
team_metadata_keys_override: Final = kwargs.pop("baggage_team_metadata_keys", None)
|
||||
|
|
@ -325,7 +347,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
self.OTEL_EXPORTER = self.config.exporter
|
||||
self.OTEL_ENDPOINT = self.config.endpoint
|
||||
self.OTEL_HEADERS = self.config.headers
|
||||
self._tracer_provider_cache: dict[str, Any] = {}
|
||||
self._tracer_provider_cache: dict[str, _SDKTracerProvider] = {}
|
||||
self._init_tracing(tracer_provider)
|
||||
|
||||
_debug_otel: Final = str(os.getenv("DEBUG_OTEL", "False")).lower()
|
||||
|
|
@ -870,7 +892,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
def _emit_guardrail_spans_from_request_data(
|
||||
self,
|
||||
request_data: dict,
|
||||
parent_span: Any | None,
|
||||
parent_span: "Span | None",
|
||||
) -> None:
|
||||
"""Emit ``guardrail`` spans from the request's proxy-internal metadata bucket
|
||||
(``standard_logging_guardrail_information``).
|
||||
|
|
@ -896,7 +918,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
# kwargs["litellm_params"]["metadata"]["_otel_internal"]. Pass the
|
||||
# SAME metadata dict the proxy populated so _handle_failure and
|
||||
# this hook see the same dedupe markers.
|
||||
kwargs: Final[dict[str, Any]] = {
|
||||
kwargs: Final[dict[str, object]] = {
|
||||
"litellm_params": {"metadata": metadata},
|
||||
"standard_logging_object": {
|
||||
"guardrail_information": guardrail_information,
|
||||
|
|
@ -1257,13 +1279,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
response_obj,
|
||||
start_time,
|
||||
end_time,
|
||||
context,
|
||||
context: "Context | None",
|
||||
):
|
||||
from opentelemetry.trace import Status, StatusCode
|
||||
|
||||
otel_tracer: Final[Tracer] = self.get_tracer_to_use_for_request(kwargs)
|
||||
|
||||
span_kwargs: Final[dict[str, Any]] = {
|
||||
span_kwargs: Final[_StartSpanKwargs] = {
|
||||
"name": self._get_span_name(kwargs),
|
||||
"start_time": self._to_ns(start_time),
|
||||
"context": context,
|
||||
|
|
@ -1454,7 +1476,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
) = _resolve_metric_attribute_filter(attributes)
|
||||
self._metric_attr_filter_resolved = True
|
||||
|
||||
def _filter_metric_attributes(self, attrs: dict[str, Any]) -> dict[str, Any]:
|
||||
def _filter_metric_attributes(self, attrs: dict[str, str]) -> dict[str, str]:
|
||||
if not self._metric_attr_filter_resolved:
|
||||
self._ensure_metric_attribute_filter()
|
||||
if self._metric_attr_include is not None:
|
||||
|
|
@ -1559,7 +1581,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
def _record_time_per_output_token_metric(
|
||||
self,
|
||||
kwargs: dict,
|
||||
response_obj: Any | None,
|
||||
response_obj: "_ResponseWithUsageView | None",
|
||||
end_time: datetime,
|
||||
duration_s: float,
|
||||
common_attrs: dict,
|
||||
|
|
@ -1775,10 +1797,10 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
|
||||
@staticmethod
|
||||
def _resolve_guardrail_context(
|
||||
span: Any | None,
|
||||
parent_span: Any | None,
|
||||
fallback_ctx: Any | None,
|
||||
) -> Any | None:
|
||||
span: "Span | None",
|
||||
parent_span: "Span | None",
|
||||
fallback_ctx: "Context | None",
|
||||
) -> "Context | None":
|
||||
"""
|
||||
Return a valid OTEL context for guardrail child spans so they are
|
||||
never orphaned (Issue #5). Priority:
|
||||
|
|
@ -1945,7 +1967,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
if should_create_primary_span:
|
||||
# Span 1: Request sent to litellm SDK
|
||||
otel_tracer: Final[Tracer] = self.get_tracer_to_use_for_request(kwargs)
|
||||
span_kwargs: Final[dict[str, Any]] = {
|
||||
span_kwargs: Final[_StartSpanKwargs] = {
|
||||
"name": self._get_span_name(kwargs),
|
||||
"start_time": self._to_ns(start_time),
|
||||
"context": _parent_context,
|
||||
|
|
@ -2131,10 +2153,10 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
@staticmethod
|
||||
def _tool_calls_kv_pair(
|
||||
tool_calls: list[ChatCompletionMessageToolCall],
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
from litellm.proxy._types import SpanAttributes
|
||||
|
||||
kv_pairs: Final[dict[str, Any]] = {}
|
||||
kv_pairs: Final[dict[str, object]] = {}
|
||||
for idx, tool_call in enumerate(tool_calls):
|
||||
_function = tool_call.get("function")
|
||||
if not _function:
|
||||
|
|
@ -2691,8 +2713,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
import json
|
||||
|
||||
try:
|
||||
_raw_response = json.loads(_raw_response)
|
||||
for param, val in _raw_response.items():
|
||||
_parsed: Final[Mapping[str, object]] = json.loads(_raw_response)
|
||||
for param, val in _parsed.items():
|
||||
self.safe_set_attribute(
|
||||
span=span,
|
||||
key=f"llm.{custom_llm_provider}.{param}",
|
||||
|
|
@ -2722,7 +2744,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
return int(dt * 1e9)
|
||||
return int(dt.timestamp() * 1e9)
|
||||
|
||||
def _get_span_name(self, kwargs):
|
||||
def _get_span_name(self, kwargs) -> str:
|
||||
litellm_params: Final = kwargs.get("litellm_params", {})
|
||||
metadata: Final = litellm_params.get("metadata") or {}
|
||||
generation_name: Final = metadata.get("generation_name")
|
||||
|
|
|
|||
|
|
@ -57,6 +57,7 @@ from litellm.integrations.otel.model.semconv import (
|
|||
Metric,
|
||||
Network,
|
||||
NetworkTransport,
|
||||
RpcSystem,
|
||||
Server,
|
||||
resolve_operation,
|
||||
resolve_provider,
|
||||
|
|
@ -102,6 +103,7 @@ __all__ = [
|
|||
"ProxyRequestSpanData",
|
||||
"RequestContext",
|
||||
"RequestIdentity",
|
||||
"RpcSystem",
|
||||
"Server",
|
||||
"ServerInfo",
|
||||
"ServiceSpanData",
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ from litellm.integrations.otel.model.payloads import (
|
|||
is_mcp_list_tools,
|
||||
is_mcp_tool_call,
|
||||
)
|
||||
from litellm.integrations.otel.model.semconv import Error
|
||||
from litellm.integrations.otel.model.spans import SpanRole, span_role_for_service
|
||||
from litellm.integrations.otel.model.utils import to_ns
|
||||
from litellm.integrations.otel.plumbing.context import (
|
||||
|
|
@ -634,18 +635,23 @@ class OpenTelemetryV2(CustomLogger):
|
|||
"""Stamp the v2 error.* attributes on the FastAPI-owned SERVER span for a
|
||||
failure that dies before any LLM-call span exists (malformed body, auth /
|
||||
validation rejection). Called from the proxy's global exception handler via
|
||||
``_close_dangling_otel_server_span``. The instrumentor still owns the span's
|
||||
status and lifecycle, so this only decorates it — never sets status, never
|
||||
ends it — and emits no exception event, matching v1's SERVER-span behavior
|
||||
and avoiding a duplicate of the event ``async_post_call_failure_hook`` or
|
||||
the ``auth`` phase span already records."""
|
||||
``_close_dangling_otel_server_span``, which swallows the exception into a
|
||||
``JSONResponse`` so the instrumentor never sees it and leaves the span
|
||||
``UNSET``; the status is set here instead (v1 did the same from the handler)
|
||||
so a failed request reads as failed and not merely as a span carrying an
|
||||
error message. The instrumentor still owns the span's lifecycle, so this
|
||||
never ends it. The exception event is recorded only when nothing stamped
|
||||
this span already — ``async_post_call_failure_hook`` and the ``auth`` phase
|
||||
span record their own, and a second event would duplicate it — while the
|
||||
attributes are always restamped so ``error.code`` stays pinned to the real
|
||||
response status."""
|
||||
if span is None or not is_recordable_span(span):
|
||||
return
|
||||
already_stamped: Final = Error.TYPE in (getattr(span, "attributes", None) or ())
|
||||
stamp_error(
|
||||
span,
|
||||
_span_error_from_exception(exception, status_code=status_code),
|
||||
record_event=False,
|
||||
set_status=False,
|
||||
record_event=not already_stamped,
|
||||
)
|
||||
|
||||
async def async_post_call_failure_hook(
|
||||
|
|
|
|||
|
|
@ -31,7 +31,9 @@ from litellm.integrations.otel.model.semconv import (
|
|||
MCP,
|
||||
Error,
|
||||
GenAI,
|
||||
JsonRpc,
|
||||
LiteLLM,
|
||||
RpcSystem,
|
||||
Server,
|
||||
)
|
||||
from litellm.integrations.otel.model.spans import db_system
|
||||
|
|
@ -94,11 +96,14 @@ class GenAIMapper:
|
|||
|
||||
_MCP_ATTRS: dict[str, Callable[[MCPToolCallSpanData], AttrValue | None]] = {
|
||||
GenAI.OPERATION_NAME: lambda d: d.operation.value,
|
||||
JsonRpc.SYSTEM: lambda d: RpcSystem.JSONRPC.value if d.server_address and d.server_port else None,
|
||||
MCP.METHOD_NAME: lambda d: d.method,
|
||||
MCP.SESSION_ID: lambda d: d.session_id,
|
||||
GenAI.TOOL_NAME: lambda d: d.tool_name or None,
|
||||
GenAI.TOOL_CALL_ARGUMENTS: lambda d: d.arguments_json,
|
||||
GenAI.TOOL_CALL_RESULT: lambda d: d.result_json,
|
||||
Server.ADDRESS: lambda d: d.server_address,
|
||||
Server.PORT: lambda d: d.server_port,
|
||||
LiteLLM.MCP_SERVER_NAME: lambda d: d.server_name,
|
||||
LiteLLM.CALL_ID: lambda d: d.identity.call_id or None,
|
||||
f"{LiteLLM.COST_PREFIX}total": lambda d: d.response_cost,
|
||||
|
|
|
|||
|
|
@ -364,6 +364,34 @@ class LLMCallSpanData:
|
|||
# --- the MCP tool-call model ------------------------------------------------- #
|
||||
|
||||
|
||||
def _upstream_address_port(resource: str | None) -> tuple[str | None, int | None]:
|
||||
"""Split a redacted MCP server origin into ``server.address`` / ``server.port``.
|
||||
|
||||
``mcp_server_resource`` is a scheme + host + port origin with userinfo, path,
|
||||
query and fragment already stripped. The port falls back to the scheme default
|
||||
when the origin omits it, because a consumer that keys a downstream dependency
|
||||
off the address renders a missing port as ``0``.
|
||||
|
||||
The origin is rebuilt without its IPv6 brackets upstream, so reading the port can
|
||||
raise on an address the host check still admits: a zone-scoped ``fe80::1%25eth0``
|
||||
leaves a truthy hostname of ``fe80`` behind. Both halves are read inside the guard
|
||||
so an unparseable origin yields no address rather than propagating out of span
|
||||
construction, matching how the redactor guards the same split.
|
||||
"""
|
||||
if not resource:
|
||||
return None, None
|
||||
try:
|
||||
parsed: Final = urlsplit(resource)
|
||||
hostname: Final = parsed.hostname
|
||||
port: Final = parsed.port
|
||||
except ValueError:
|
||||
return None, None
|
||||
if not hostname:
|
||||
return None, None
|
||||
default_port: Final = 443 if parsed.scheme == "https" else 80 if parsed.scheme == "http" else None
|
||||
return hostname, port or default_port
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MCPToolCallSpanData:
|
||||
"""One MCP ``tools/call`` execution, parsed from a closed request's payload.
|
||||
|
|
@ -378,6 +406,8 @@ class MCPToolCallSpanData:
|
|||
method: str
|
||||
tool_name: str
|
||||
server_name: str | None
|
||||
server_address: str | None
|
||||
server_port: int | None
|
||||
session_id: str | None
|
||||
arguments_json: str | None
|
||||
result_json: str | None
|
||||
|
|
@ -390,11 +420,14 @@ class MCPToolCallSpanData:
|
|||
cls, payload: StandardLoggingPayload, capture_content: bool = False
|
||||
) -> MCPToolCallSpanData:
|
||||
meta: Final = _mcp_tool_call_metadata(cast(Mapping[str, object], payload))
|
||||
address, port = _upstream_address_port(as_str(meta.get("mcp_server_resource")) or None)
|
||||
return cls(
|
||||
operation=resolve_operation(as_str(payload.get("call_type"))),
|
||||
method=MCPMethod.TOOLS_CALL.value,
|
||||
tool_name=as_str(meta.get("name")) or "",
|
||||
server_name=as_str(meta.get("mcp_server_name")),
|
||||
server_address=address,
|
||||
server_port=port,
|
||||
session_id=as_str(meta.get("mcp_session_id")),
|
||||
arguments_json=(
|
||||
_json_or_none(meta.get("arguments")) if capture_content and meta.get("arguments") is not None else None
|
||||
|
|
|
|||
|
|
@ -130,11 +130,25 @@ class JsonRpc:
|
|||
"""JSON-RPC keys carried on MCP spans. The error/status code lives in the
|
||||
``rpc.*`` namespace per semconv, not ``jsonrpc.*``."""
|
||||
|
||||
SYSTEM: Final = "rpc.system"
|
||||
REQUEST_ID: Final = "jsonrpc.request.id"
|
||||
PROTOCOL_VERSION: Final = "jsonrpc.protocol.version"
|
||||
RESPONSE_STATUS_CODE: Final = "rpc.response.status_code"
|
||||
|
||||
|
||||
class RpcSystem(str, Enum):
|
||||
"""Well-known values for ``rpc.system``. MCP frames every message as JSON-RPC 2.0.
|
||||
|
||||
Naming the system also classifies the span: a CLIENT span carrying none of the
|
||||
``rpc.*``/``http.*``/``db.*``/``messaging.*`` families records no span type or
|
||||
subtype in backends that derive those from the attribute family. It is emitted
|
||||
only alongside ``server.address``/``server.port``, since a backend that reads it
|
||||
as a downstream dependency names that dependency from the server address.
|
||||
"""
|
||||
|
||||
JSONRPC = "jsonrpc"
|
||||
|
||||
|
||||
class NetworkTransport(str, Enum):
|
||||
"""Well-known values for ``network.transport``."""
|
||||
|
||||
|
|
|
|||
|
|
@ -42,9 +42,7 @@ def langfuse_dynamic_headers(params: StandardCallbackDynamicParams) -> dict[str,
|
|||
public_key: Final = params.get("langfuse_public_key")
|
||||
secret_key: Final = params.get("langfuse_secret_key")
|
||||
if public_key and secret_key:
|
||||
return {
|
||||
"Authorization": _V1Langfuse._get_langfuse_authorization_header(
|
||||
public_key=public_key, secret_key=secret_key
|
||||
)
|
||||
}
|
||||
return _V1Langfuse._build_langfuse_otel_headers(
|
||||
_V1Langfuse._get_langfuse_authorization_header(public_key=public_key, secret_key=secret_key)
|
||||
)
|
||||
return {}
|
||||
|
|
|
|||
|
|
@ -6,12 +6,13 @@ import random
|
|||
import time
|
||||
import uuid
|
||||
from collections import Counter
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict
|
||||
|
||||
import httpx
|
||||
from typing_extensions import Never, ReadOnly
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_batch_logger import CustomBatchLogger
|
||||
|
|
@ -48,7 +49,20 @@ _WEBHOOK_PATH_PROMPT_MODERATION: Final = "/v1/before_prompt/openai/v1"
|
|||
_WEBHOOK_PATH_LOGGING_BATCH: Final = "/v1/litellm/batch"
|
||||
_MAX_QUEUE_SIZE: Final = 10_000
|
||||
_DROP_WARNING_INTERVAL_SECONDS: Final = 60.0
|
||||
_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({})
|
||||
_EMPTY_MAPPING: Final[Mapping[str, Never]] = MappingProxyType({})
|
||||
|
||||
|
||||
class _ServiceToolCall(TypedDict):
|
||||
id: ReadOnly[str]
|
||||
|
||||
|
||||
class _ServiceMessage(TypedDict, total=False):
|
||||
content: ReadOnly[str]
|
||||
tool_calls: ReadOnly[Sequence[_ServiceToolCall]]
|
||||
|
||||
|
||||
class _ServiceChoice(TypedDict, total=False):
|
||||
message: ReadOnly[_ServiceMessage]
|
||||
|
||||
|
||||
class _MalformedToolBlockingResponseError(Exception):
|
||||
|
|
@ -143,7 +157,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
else {"Content-Type": "application/json"}
|
||||
)
|
||||
|
||||
self._periodic_flush_task: asyncio.Task[Any] | None = self._start_periodic_flush_task()
|
||||
self._periodic_flush_task: asyncio.Task[None] | None = self._start_periodic_flush_task()
|
||||
|
||||
@classmethod
|
||||
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
|
||||
|
|
@ -191,7 +205,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
params={"timeout": httpx.Timeout(5.0, connect=2.0)},
|
||||
)
|
||||
|
||||
def _start_periodic_flush_task(self) -> asyncio.Task[Any] | None:
|
||||
def _start_periodic_flush_task(self) -> asyncio.Task[None] | None:
|
||||
"""Start the periodic flush task only when an event loop is already running."""
|
||||
try:
|
||||
loop: Final = asyncio.get_running_loop()
|
||||
|
|
@ -212,7 +226,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
Closing them here would close the shared connection pool for every
|
||||
other logger instance; let LiteLLM manage their lifecycle instead.
|
||||
"""
|
||||
task: Final = getattr(self, "_periodic_flush_task", None)
|
||||
task: Final[asyncio.Task[None] | None] = getattr(self, "_periodic_flush_task", None)
|
||||
if task is not None:
|
||||
task.cancel()
|
||||
|
||||
|
|
@ -253,7 +267,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
|
||||
@staticmethod
|
||||
async def _guarded(
|
||||
coro: Any,
|
||||
coro: Awaitable[GenericGuardrailAPIInputs],
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
label: str,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
|
|
@ -400,7 +414,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
request_data["_rubrik_logging_obj"] = logging_obj
|
||||
|
||||
@staticmethod
|
||||
def _normalize_tool_calls(tool_calls: Any) -> tuple[ChatCompletionMessageToolCall, ...]:
|
||||
def _normalize_tool_calls(tool_calls: Sequence[object]) -> tuple[ChatCompletionMessageToolCall, ...]:
|
||||
"""Convert tool_calls from inputs to ChatCompletionMessageToolCall objects."""
|
||||
return tuple(RubrikLogger._normalize_tool_call(tc) for tc in tool_calls)
|
||||
|
||||
|
|
@ -427,7 +441,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
raise TypeError(f"Cannot normalize tool_call of type {type(tc).__name__}: {tc!r}")
|
||||
|
||||
@staticmethod
|
||||
def _join_texts(texts: Any) -> str:
|
||||
def _join_texts(texts: Sequence[str] | None) -> str:
|
||||
"""Join response text segments into the single content string the
|
||||
webhook evaluates. Empty when there is no assistant text."""
|
||||
if not texts:
|
||||
|
|
@ -439,14 +453,14 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
tool_calls: Sequence[ChatCompletionMessageToolCall],
|
||||
content: str,
|
||||
request_id: str | None,
|
||||
) -> Mapping[str, Any]:
|
||||
) -> Mapping[str, object]:
|
||||
"""Build an OpenAI ChatCompletion-format dict (assistant text + tool
|
||||
calls) for the after_completion webhook.
|
||||
|
||||
``content`` is sent so the webhook can moderate the response text;
|
||||
``None`` when the assistant produced no text (tool-call-only response).
|
||||
"""
|
||||
message: Final[dict[str, Any]] = {
|
||||
message: Final[dict[str, object]] = {
|
||||
"role": "assistant",
|
||||
"content": content or None,
|
||||
}
|
||||
|
|
@ -467,7 +481,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
}
|
||||
|
||||
@staticmethod
|
||||
def _flatten_messages_for_moderation(messages: Any) -> tuple[Mapping[str, Any], ...]:
|
||||
def _flatten_messages_for_moderation(messages: Sequence[object] | None) -> tuple[Mapping[str, Any], ...]:
|
||||
"""Collapse each message's content to a plain string for the webhook.
|
||||
|
||||
litellm normalizes Anthropic ``/v1/messages`` requests to OpenAI shape,
|
||||
|
|
@ -506,8 +520,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
@staticmethod
|
||||
def _build_prompt_moderation_payload(
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: Mapping[str, Any],
|
||||
) -> Mapping[str, Any]:
|
||||
request_data: Mapping[str, object],
|
||||
) -> Mapping[str, object]:
|
||||
"""Build the bare OpenAI request the before_prompt webhook consumes.
|
||||
|
||||
Unlike the after_completion envelope, this endpoint takes a raw OpenAI
|
||||
|
|
@ -516,7 +530,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
``/v1/messages`` requests too. Optional fields are sent only when
|
||||
present so the payload stays clean.
|
||||
"""
|
||||
payload: Final[dict[str, Any]] = {
|
||||
payload: Final[dict[str, object]] = {
|
||||
"model": inputs.get("model") or request_data.get("model") or "",
|
||||
"messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")),
|
||||
}
|
||||
|
|
@ -540,8 +554,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
@staticmethod
|
||||
def _extract_request_data(
|
||||
call_details: Mapping[str, Any],
|
||||
request_data: Mapping[str, Any] | None,
|
||||
) -> Mapping[str, Any]:
|
||||
request_data: Mapping[str, object] | None,
|
||||
) -> Mapping[str, object]:
|
||||
"""Extract original request data from model_call_details for the
|
||||
response moderation service envelope.
|
||||
|
||||
|
|
@ -576,7 +590,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
}
|
||||
|
||||
@staticmethod
|
||||
def _sanitize_proxy_server_request(proxy_server_request: Any) -> Any:
|
||||
def _sanitize_proxy_server_request(proxy_server_request: object) -> object:
|
||||
"""Allowlist only routing fields (``url``, ``method``) when forwarding
|
||||
``proxy_server_request`` to an external webhook, dropping inbound
|
||||
``headers`` (Authorization, Cookie, x-api-key, ...) and the raw
|
||||
|
|
@ -586,17 +600,18 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
return {key: proxy_server_request[key] for key in ("url", "method") if key in proxy_server_request}
|
||||
|
||||
@staticmethod
|
||||
def _resolve_model(request_data: Mapping[str, Any], call_details: Mapping[str, Any]) -> str:
|
||||
def _resolve_model(request_data: Mapping[str, object], call_details: Mapping[str, str]) -> str:
|
||||
"""Get the model name for the ModifyResponseException."""
|
||||
response: Final = request_data.get("response")
|
||||
if response and hasattr(response, "model"):
|
||||
return response.model or "unknown"
|
||||
response_model: Final[str | None] = getattr(response, "model", None)
|
||||
return response_model or "unknown"
|
||||
return call_details.get("model", "unknown")
|
||||
|
||||
# -- Logging hooks ---------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _correlation_id(call_details: Mapping[str, Any], request_data: Mapping[str, Any] | None = None) -> str | None:
|
||||
def _correlation_id(call_details: Mapping[str, str], request_data: Mapping[str, str] | None = None) -> str | None:
|
||||
"""The id that joins a blocked request's two S3 logs by filename: the
|
||||
moderation (``_blocking``) log and the failure (response) log.
|
||||
|
||||
|
|
@ -610,7 +625,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
return call_details.get("litellm_call_id") or (request_data or _EMPTY_MAPPING).get("litellm_call_id")
|
||||
|
||||
@classmethod
|
||||
def _apply_correlation_id(cls, payload: dict[str, Any], source: Mapping[str, Any]) -> None:
|
||||
def _apply_correlation_id(cls, payload: dict[str, object], source: Mapping[str, str]) -> None:
|
||||
"""Pin ``payload["id"]`` to ``litellm_call_id`` in place so this log
|
||||
shares its S3 filename id with the moderation (``_blocking``) and
|
||||
failure logs for the same request -- for every provider.
|
||||
|
|
@ -630,7 +645,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
payload["id"] = correlated
|
||||
|
||||
@staticmethod
|
||||
def _prepend_system_prompt(payload: dict[str, Any], source: Mapping[str, Any]) -> None:
|
||||
def _prepend_system_prompt(payload: dict[str, object], source: Mapping[str, object]) -> None:
|
||||
"""Prepend ``source["system"]`` onto ``payload["messages"]``.
|
||||
|
||||
Builds a NEW messages list rather than mutating ``payload["messages"]``
|
||||
|
|
@ -658,7 +673,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
exc_info=True,
|
||||
)
|
||||
|
||||
async def _prepare_log_payload(self, kwargs: Mapping[str, Any], event_type: str) -> StandardLoggingPayload | None:
|
||||
async def _prepare_log_payload(
|
||||
self, kwargs: Mapping[str, object], event_type: str
|
||||
) -> StandardLoggingPayload | None:
|
||||
"""Shared logic for success logging (sampled)."""
|
||||
if random.random() > self.sampling_rate:
|
||||
verbose_logger.debug("Skipping Rubrik %s logging (sampling_rate=%s)", event_type, self.sampling_rate)
|
||||
|
|
@ -697,7 +714,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
self._dropped_since_warning = 0
|
||||
self._last_drop_warning_time = now
|
||||
|
||||
async def _enqueue_log_event(self, kwargs: Mapping[str, Any], event_type: str):
|
||||
async def _enqueue_log_event(self, kwargs: Mapping[str, object], event_type: str):
|
||||
try:
|
||||
payload: Final = await self._prepare_log_payload(kwargs, event_type)
|
||||
if payload is None:
|
||||
|
|
@ -862,7 +879,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
|
||||
base: Final = call_details.get("standard_logging_object")
|
||||
if base is not None:
|
||||
payload: dict = safe_deep_copy(base)
|
||||
payload: dict[str, object] = safe_deep_copy(base)
|
||||
else:
|
||||
verbose_logger.debug(
|
||||
"Rubrik: standard_logging_object not yet on model_call_details "
|
||||
|
|
@ -908,7 +925,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
cls,
|
||||
call_details: Mapping[str, Any],
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
) -> dict[str, Any]:
|
||||
) -> dict[str, object]:
|
||||
# Convert datetime to a Unix float so json.dumps can serialize it.
|
||||
# httpx's json= parameter uses stdlib json.dumps with no custom encoder.
|
||||
_raw_start: Final = call_details.get("start_time")
|
||||
|
|
@ -996,7 +1013,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
|
||||
# -- Webhook services ------------------------------------------------------
|
||||
|
||||
async def _post_json(self, endpoint: str, payload: Mapping[str, Any], service_name: str) -> Mapping[str, Any]:
|
||||
async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> Mapping[str, Any]:
|
||||
"""POST ``payload`` to a Rubrik webhook and return its dict response.
|
||||
|
||||
Raises:
|
||||
|
|
@ -1010,7 +1027,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
headers=self._headers,
|
||||
)
|
||||
http_response.raise_for_status()
|
||||
result: Final = http_response.json()
|
||||
result: Final[object] = http_response.json()
|
||||
if not isinstance(result, dict):
|
||||
raise TypeError(
|
||||
f"{service_name} returned non-dict JSON "
|
||||
|
|
@ -1021,8 +1038,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
|
||||
async def _post_to_response_moderation_endpoint(
|
||||
self,
|
||||
response_data: Mapping[str, Any],
|
||||
request_data: Mapping[str, Any],
|
||||
response_data: Mapping[str, object],
|
||||
request_data: Mapping[str, object],
|
||||
) -> Mapping[str, Any]:
|
||||
"""Post the ``{request, response}`` envelope to the after_completion
|
||||
webhook and return its (possibly rewritten) response.
|
||||
|
|
@ -1039,7 +1056,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
"Response moderation service",
|
||||
)
|
||||
|
||||
async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, Any]) -> Mapping[str, Any]:
|
||||
async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> Mapping[str, Any]:
|
||||
"""Post a bare OpenAI request to the before_prompt webhook.
|
||||
|
||||
Returns ``{}`` (passthrough) or a synthetic chat.completion (block).
|
||||
|
|
@ -1054,7 +1071,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
chat.completion whose ``choices[0].message.content`` is the refusal
|
||||
explanation.
|
||||
"""
|
||||
choices: Final = service_response.get("choices")
|
||||
choices: Final[Sequence[_ServiceChoice] | None] = service_response.get("choices")
|
||||
if not choices:
|
||||
return None
|
||||
message: Final = choices[0].get("message") or _EMPTY_MAPPING
|
||||
|
|
@ -1086,7 +1103,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger):
|
|||
Expects service_response in OpenAI chat completion format:
|
||||
{"choices": [{"message": {"tool_calls": [...], "content": "..."}}]}
|
||||
"""
|
||||
choices: Final = service_response.get("choices") or ()
|
||||
choices: Final[Sequence[_ServiceChoice]] = service_response.get("choices") or ()
|
||||
if not choices:
|
||||
raise _MalformedToolBlockingResponseError("Response moderation service returned empty response")
|
||||
|
||||
|
|
|
|||
563
litellm/integrations/shadow_eval_logger.py
Normal file
563
litellm/integrations/shadow_eval_logger.py
Normal file
|
|
@ -0,0 +1,563 @@
|
|||
"""Shadow Eval Logger: samples a shadowed key's successful chat requests, duplicates each
|
||||
through the auto-router in a detached task, blind-judges real vs shadow, and appends one
|
||||
``LiteLLM_ShadowEvalAttempt`` row (verdict or error) as the feature's only hot-path write.
|
||||
Counts, status, and spend derive from those rows at read time, so nothing can disagree
|
||||
across pods or stop races; the hook reads active jobs through a short-TTL cache."""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import random
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs
|
||||
from litellm.litellm_core_utils.internal_call_metadata import sanitized_forwardable_call_metadata
|
||||
from litellm.litellm_core_utils.llm_judge import (
|
||||
default_router_provider,
|
||||
extract_text_from_content,
|
||||
judge_acompletion,
|
||||
parse_json_verdict,
|
||||
)
|
||||
from litellm.litellm_core_utils.redact_messages import should_redact_message_logging
|
||||
from litellm.types.utils import SHADOW_EVAL_JUDGE_CALL_ORIGIN, SHADOW_EVAL_ROUTER_CALL_ORIGIN
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.router import Router
|
||||
from litellm.types.utils import StandardLoggingPayload
|
||||
|
||||
# A job starting, stopping, or hitting its turn budget propagates to sampling within one
|
||||
# TTL; the turn budget can overshoot by at most one TTL of in-flight samples per pod.
|
||||
_JOBS_CACHE_TTL_SECONDS: Final = 10
|
||||
|
||||
# Concurrent shadow+judge pipelines per pod: a traffic spike turns into skipped samples
|
||||
# rather than an unbounded task pileup.
|
||||
_MAX_CONCURRENT_SHADOW_TASKS: Final = 16
|
||||
|
||||
# Total character budget for the judge's user prompt, however long the conversation and
|
||||
# the two responses are, so the prompt can never overflow a judge model's context window.
|
||||
_MAX_JUDGE_RESPONSE_CHARS: Final = 8_000
|
||||
_MAX_JUDGE_PROMPT_CHARS: Final = 24_000
|
||||
|
||||
# The judge answers with a small JSON object; a tighter budget truncates the JSON
|
||||
# mid-object and the attempt is lost to an error row.
|
||||
JUDGE_MAX_OUTPUT_TOKENS: Final = 500
|
||||
|
||||
_MAX_ERROR_CHARS: Final = 500
|
||||
|
||||
_EMPTY_METADATA: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
|
||||
_SAMPLED_CALL_TYPES: Final = frozenset({"completion", "acompletion"})
|
||||
|
||||
PAIRWISE_JUDGE_SYSTEM_PROMPT: Final = """You are an impartial quality judge comparing two responses to the same conversation.
|
||||
|
||||
The responses are labeled A and B in random order. You do not know which system produced which.
|
||||
|
||||
Criteria: correctness, completeness, clarity, conciseness.
|
||||
|
||||
Return ONLY valid JSON in this exact format, no other text:
|
||||
{
|
||||
"preference": "A" | "B" | "tie",
|
||||
"confidence": <0.0 to 1.0>,
|
||||
"reasoning": "<one sentence>"
|
||||
}"""
|
||||
|
||||
|
||||
class PairwiseVerdict(BaseModel):
|
||||
"""The judge's blind A/B verdict, validated at the parse boundary."""
|
||||
|
||||
preference: str = "tie"
|
||||
confidence: float = 0.0
|
||||
|
||||
|
||||
def _sample_hits(request_id: str, job_id: str, percentage: float) -> bool:
|
||||
"""Deterministically decide whether a request falls in the shadowed slice: hash-based
|
||||
rather than random so retries sample the same way and pods agree without coordination."""
|
||||
digest: Final = hashlib.sha256(f"{job_id}:{request_id}".encode()).digest()
|
||||
bucket: Final = int.from_bytes(digest[:8], "big") / float(2**64)
|
||||
return bucket * 100.0 < percentage
|
||||
|
||||
|
||||
def _judge_call_cost(response: object) -> float:
|
||||
"""Price a judge call, treating an unmapped judge model as free rather than fatal."""
|
||||
import litellm
|
||||
|
||||
try:
|
||||
return litellm.completion_cost(completion_response=response) or 0.0
|
||||
except Exception: # noqa: BLE001 # unmapped judge model: the verdict still counts, cost stays 0
|
||||
return 0.0
|
||||
|
||||
|
||||
def _unmask_preference(raw_preference: str, real_is_a: bool) -> str:
|
||||
"""Map the judge's blind A/B/tie verdict back to real/shadow/tie."""
|
||||
normalized: Final = raw_preference.strip().lower()
|
||||
if normalized == "a":
|
||||
return "real" if real_is_a else "shadow"
|
||||
if normalized == "b":
|
||||
return "shadow" if real_is_a else "real"
|
||||
return "tie"
|
||||
|
||||
|
||||
def _judge_user_prompt(conversation: str, response_a: str, response_b: str) -> str:
|
||||
"""The judge prompt under one total character budget: each response is capped, and
|
||||
the conversation tail gets whatever budget the responses left over."""
|
||||
a: Final = response_a[:_MAX_JUDGE_RESPONSE_CHARS]
|
||||
b: Final = response_b[:_MAX_JUDGE_RESPONSE_CHARS]
|
||||
conversation_budget: Final = _MAX_JUDGE_PROMPT_CHARS - len(a) - len(b)
|
||||
return (
|
||||
f"Conversation:\n{conversation[-conversation_budget:]}\n\n"
|
||||
f"Response A:\n{a}\n\n"
|
||||
f"Response B:\n{b}\n\n"
|
||||
"Which response is better?"
|
||||
)
|
||||
|
||||
|
||||
async def _key_or_team_is_over_budget(metadata: Mapping[str, object]) -> bool:
|
||||
"""Whether the shadowed key or its team is over budget, decided by the same owners
|
||||
the request path uses, so counter keys and thresholds can never drift from auth's.
|
||||
|
||||
Advisory and fail-open: real traffic on an over-budget key is already rejected at
|
||||
auth (so nothing reaches the success hook), and this gate only closes the race
|
||||
where the key crosses its budget while a request is in flight.
|
||||
"""
|
||||
try:
|
||||
from litellm.exceptions import BudgetExceededError
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
_team_max_budget_check,
|
||||
_virtual_key_max_budget_check,
|
||||
get_team_object,
|
||||
)
|
||||
from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache
|
||||
except ImportError:
|
||||
return False
|
||||
|
||||
auth: Final = metadata.get("user_api_key_auth")
|
||||
if not isinstance(auth, UserAPIKeyAuth):
|
||||
return False
|
||||
try:
|
||||
await _virtual_key_max_budget_check(valid_token=auth, proxy_logging_obj=proxy_logging_obj)
|
||||
if auth.team_id:
|
||||
team: Final = await get_team_object(
|
||||
team_id=auth.team_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
check_cache_only=True,
|
||||
)
|
||||
await _team_max_budget_check(team_object=team, valid_token=auth, proxy_logging_obj=proxy_logging_obj)
|
||||
except BudgetExceededError:
|
||||
return True
|
||||
except Exception as e: # noqa: BLE001 # advisory gate: a failed read must not block sampling
|
||||
verbose_logger.debug("shadow_eval: budget read failed: %s", e)
|
||||
return False
|
||||
|
||||
|
||||
def _request_was_routed_by(request_metadata: Mapping[str, object], router_name: str) -> bool:
|
||||
"""Duplicating a request the shadowed router already served compares the router to
|
||||
itself: guaranteed ties, judge spend for zero information."""
|
||||
decision: Final = request_metadata.get("routing_decision")
|
||||
if not isinstance(decision, Mapping):
|
||||
return False
|
||||
return decision.get("router_model_name") == router_name
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _CallFailure:
|
||||
"""A shadow or judge call that produced no usable response. cost carries any judge
|
||||
spend the failed attempt still billed, so job-level judge_spend never undercounts."""
|
||||
|
||||
error: str
|
||||
cost: float = 0.0
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ShadowResponse:
|
||||
"""A successful shadow call, with what the attempt row records."""
|
||||
|
||||
text: str
|
||||
model: str
|
||||
tier: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _JudgeVerdict:
|
||||
"""A parsed judge verdict, unmasked back to real/shadow/tie."""
|
||||
|
||||
preference: str
|
||||
confidence: float
|
||||
cost: float
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ActiveShadowEvalJob:
|
||||
"""One active job as the sampling path needs it: immutable config plus the attempt
|
||||
count as of the cache fill (the turn budget's staleness is bounded by the cache TTL)."""
|
||||
|
||||
id: str
|
||||
router_name: str
|
||||
shadow_percentage: float
|
||||
judge_model: str
|
||||
max_turns: int
|
||||
ends_at: datetime
|
||||
attempts: int
|
||||
|
||||
|
||||
def _as_utc(value: datetime) -> datetime:
|
||||
return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value
|
||||
|
||||
|
||||
_jobs_cache: Final = InMemoryCache(max_size_in_memory=4, default_ttl=_JOBS_CACHE_TTL_SECONDS)
|
||||
_JOBS_CACHE_KEY: Final = "shadow_eval:active_jobs"
|
||||
|
||||
|
||||
class ShadowEvalLogger(CustomLogger):
|
||||
"""Fires blind pairwise shadow evaluations for keys with an active shadow-eval job."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
router_provider: Callable[[], "Router | None"] | None = None,
|
||||
prisma_provider: Callable[[], "PrismaClient | None"] | None = None,
|
||||
jobs_cache: InMemoryCache | None = None,
|
||||
) -> None:
|
||||
"""Providers are callables so the proxy's lazily-initialized globals are resolved
|
||||
at call time, not at logger construction."""
|
||||
self._router_provider = router_provider or default_router_provider
|
||||
self._prisma_provider = prisma_provider or _default_prisma_provider
|
||||
self._jobs_cache = jobs_cache or _jobs_cache
|
||||
self._inflight_shadow_tasks: int = 0
|
||||
# Starts per job since the last cache fill, never decremented within a
|
||||
# generation; the refill absorbs written rows and resets.
|
||||
self._job_starts: dict[str, int] = {} # mutable-ok: per-generation counter
|
||||
|
||||
async def _active_jobs(self) -> Mapping[str, ActiveShadowEvalJob]:
|
||||
"""Active jobs by api_key_id, cache-first. A DB fault returns empty without
|
||||
caching, so sampling pauses for that request and the next one retries."""
|
||||
cached: Final = await self._jobs_cache.async_get_cache(_JOBS_CACHE_KEY)
|
||||
if cached is not None:
|
||||
return cached # pyright: ignore[reportReturnType] # cache stores exactly this mapping shape
|
||||
prisma: Final = self._prisma_provider()
|
||||
if prisma is None:
|
||||
return _EMPTY_JOBS
|
||||
try:
|
||||
records: Final = await prisma.db.litellm_shadowevaljob.find_many(
|
||||
where={ # mutable-ok: Prisma filter
|
||||
"stopped_at": None,
|
||||
"ends_at": {"gt": datetime.now(timezone.utc)}, # mutable-ok: Prisma filter
|
||||
},
|
||||
)
|
||||
grouped: Final = (
|
||||
await prisma.db.litellm_shadowevalattempt.group_by(
|
||||
by=["job_id"],
|
||||
count=True,
|
||||
where={"job_id": {"in": [str(record.id) for record in records]}}, # mutable-ok: Prisma filter
|
||||
)
|
||||
if records
|
||||
else ()
|
||||
)
|
||||
attempt_counts: Final = {str(row["job_id"]): int(row["_count"]["_all"]) for row in grouped or []}
|
||||
jobs: Final = {
|
||||
str(record.api_key_id): ActiveShadowEvalJob(
|
||||
id=str(record.id),
|
||||
router_name=str(record.router_name),
|
||||
shadow_percentage=float(record.shadow_percentage),
|
||||
judge_model=str(record.judge_model),
|
||||
max_turns=int(record.max_turns),
|
||||
ends_at=_as_utc(record.ends_at),
|
||||
attempts=attempt_counts.get(str(record.id), 0),
|
||||
)
|
||||
for record in records or []
|
||||
}
|
||||
await self._jobs_cache.async_set_cache(_JOBS_CACHE_KEY, jobs)
|
||||
self._job_starts = {} # rebind-ok: new generation, counts absorbed into the fill
|
||||
return jobs
|
||||
except Exception as e: # noqa: BLE001 # a DB blip must never break request logging
|
||||
verbose_logger.debug("shadow_eval: active-job read failed: %s", e)
|
||||
return _EMPTY_JOBS
|
||||
|
||||
#### hook ####
|
||||
|
||||
async def async_log_success_event(
|
||||
self,
|
||||
kwargs: Mapping[str, object],
|
||||
response_obj: object,
|
||||
start_time: object,
|
||||
end_time: object,
|
||||
) -> None:
|
||||
try:
|
||||
payload: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") # pyright: ignore[reportAssignmentType] # untyped callback kwargs
|
||||
if payload is None:
|
||||
return
|
||||
raw_meta: Final = get_litellm_metadata_from_kwargs(dict(kwargs)) # mutable-ok: helper needs dict
|
||||
request_metadata: Final = raw_meta if isinstance(raw_meta, Mapping) else _EMPTY_METADATA
|
||||
if request_metadata.get(INTERNAL_CALL_ORIGIN_METADATA_KEY):
|
||||
return # internal sub-call (our own shadow/judge, a classifier), not user traffic
|
||||
# redaction rewrites logged content before callbacks run, so this hook
|
||||
# only ever sees placeholders for a redacted request
|
||||
if should_redact_message_logging(dict(kwargs)): # mutable-ok: predicate takes a plain dict
|
||||
return
|
||||
metadata: Final = payload.get("metadata") or _EMPTY_METADATA
|
||||
api_key_hash: Final = metadata.get("user_api_key_hash")
|
||||
if not api_key_hash:
|
||||
return
|
||||
job: Final = (await self._active_jobs()).get(str(api_key_hash))
|
||||
if job is None:
|
||||
return
|
||||
if datetime.now(timezone.utc) >= job.ends_at:
|
||||
return
|
||||
if job.attempts + self._job_starts.get(job.id, 0) >= job.max_turns:
|
||||
return
|
||||
request_id: Final = payload.get("id") or ""
|
||||
if not request_id:
|
||||
return
|
||||
if not _sample_hits(request_id, job.id, job.shadow_percentage):
|
||||
return
|
||||
if payload.get("call_type") not in _SAMPLED_CALL_TYPES:
|
||||
return # only known chat-shaped traffic is comparable; unknown or missing types fail closed
|
||||
if _request_was_routed_by(request_metadata, job.router_name):
|
||||
return
|
||||
if self._inflight_shadow_tasks >= _MAX_CONCURRENT_SHADOW_TASKS:
|
||||
return
|
||||
raw_messages: Final = kwargs.get("messages")
|
||||
self._job_starts[job.id] = self._job_starts.get(job.id, 0) + 1
|
||||
self._inflight_shadow_tasks += 1
|
||||
task: Final = asyncio.create_task(
|
||||
self._run_shadow_eval(
|
||||
job=job,
|
||||
request_id=request_id,
|
||||
messages=tuple(m for m in raw_messages if isinstance(m, Mapping))
|
||||
if isinstance(raw_messages, Sequence)
|
||||
else (),
|
||||
response_obj=response_obj,
|
||||
real_model=payload.get("model") or "",
|
||||
model_parameters=MappingProxyType(
|
||||
dict(payload.get("model_parameters") or {}) # mutable-ok: frozen snapshot
|
||||
),
|
||||
parent_metadata=MappingProxyType(dict(request_metadata)), # mutable-ok: frozen snapshot
|
||||
)
|
||||
)
|
||||
task.add_done_callback(self._release_shadow_slot)
|
||||
except Exception as e: # noqa: BLE001 # logging hooks must never fail the request
|
||||
verbose_logger.debug("shadow_eval: failed to schedule task: %s", e)
|
||||
|
||||
def _release_shadow_slot(self, _task: "asyncio.Task[None]") -> None:
|
||||
self._inflight_shadow_tasks -= 1
|
||||
|
||||
#### the detached pipeline: one attempt row per sampled request, verdict or error ####
|
||||
|
||||
async def _run_shadow_eval(
|
||||
self,
|
||||
job: ActiveShadowEvalJob,
|
||||
request_id: str,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
response_obj: object,
|
||||
real_model: str,
|
||||
model_parameters: Mapping[str, object],
|
||||
parent_metadata: Mapping[str, object],
|
||||
) -> None:
|
||||
"""Budget gate -> shadow call -> blind judge -> one attempt row. The prisma gate
|
||||
sits above the dispatch so no provider spend happens without a place to record
|
||||
the outcome, and the budget read lives here rather than in the success hook."""
|
||||
prisma: Final = self._prisma_provider()
|
||||
try:
|
||||
if prisma is None:
|
||||
return
|
||||
real_text: Final = self._extract_response_text(response_obj)
|
||||
if not real_text or not messages:
|
||||
return
|
||||
if await _key_or_team_is_over_budget(parent_metadata):
|
||||
return
|
||||
|
||||
shadow: Final = await self._call_router_shadow(job.router_name, messages, model_parameters, parent_metadata)
|
||||
if isinstance(shadow, _CallFailure):
|
||||
await self._record_attempt(prisma, job, request_id, outcome="error", error=shadow.error)
|
||||
return
|
||||
|
||||
verdict: Final = await self._call_judge(
|
||||
judge_model=job.judge_model,
|
||||
messages=messages,
|
||||
real_text=real_text,
|
||||
shadow_text=shadow.text,
|
||||
parent_metadata=parent_metadata,
|
||||
)
|
||||
if isinstance(verdict, _CallFailure):
|
||||
await self._record_attempt(
|
||||
prisma,
|
||||
job,
|
||||
request_id,
|
||||
outcome="error",
|
||||
error=verdict.error,
|
||||
shadow=shadow,
|
||||
judge_cost=verdict.cost,
|
||||
)
|
||||
return
|
||||
await self._record_attempt(
|
||||
prisma,
|
||||
job,
|
||||
request_id,
|
||||
outcome=verdict.preference,
|
||||
shadow=shadow,
|
||||
real_model=real_model,
|
||||
confidence=verdict.confidence,
|
||||
judge_cost=verdict.cost,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # detached task: record what happened, never raise
|
||||
verbose_logger.debug("shadow_eval: pipeline failed for %s: %s", request_id, e)
|
||||
await self._record_attempt(prisma, job, request_id, outcome="error", error=f"pipeline error: {e}")
|
||||
|
||||
@staticmethod
|
||||
async def _record_attempt(
|
||||
prisma: "PrismaClient | None",
|
||||
job: ActiveShadowEvalJob,
|
||||
request_id: str,
|
||||
*,
|
||||
outcome: str,
|
||||
shadow: _ShadowResponse | None = None,
|
||||
real_model: str = "",
|
||||
confidence: float | None = None,
|
||||
judge_cost: float = 0.0,
|
||||
error: str | None = None,
|
||||
) -> None:
|
||||
if prisma is None:
|
||||
return
|
||||
try:
|
||||
await prisma.db.litellm_shadowevalattempt.create(
|
||||
data={ # mutable-ok: Prisma payload
|
||||
"job_id": job.id,
|
||||
"request_id": request_id,
|
||||
"outcome": outcome,
|
||||
"tier": shadow.tier if shadow else None,
|
||||
"real_model": real_model or None,
|
||||
"shadow_model": shadow.model if shadow else None,
|
||||
"confidence": confidence,
|
||||
"judge_cost": judge_cost,
|
||||
"error": error[:_MAX_ERROR_CHARS] if error else None,
|
||||
}
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # a lost row degrades sample size, nothing can disagree with it
|
||||
verbose_logger.debug("shadow_eval: attempt write failed for %s: %s", request_id, e)
|
||||
|
||||
async def _call_router_shadow(
|
||||
self,
|
||||
router_name: str,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
model_parameters: Mapping[str, object],
|
||||
parent_metadata: Mapping[str, object],
|
||||
) -> "_ShadowResponse | _CallFailure":
|
||||
"""Send the prompt through the auto-router being evaluated. The metadata carries
|
||||
the shadowed key's identity (spend attribution) and receives the router's routing
|
||||
decision write-back, read back for tier attribution."""
|
||||
router: Final = self._router_provider()
|
||||
if router is None:
|
||||
return _CallFailure("no router configured on this pod")
|
||||
shadow_metadata: Final[dict[str, object]] = ( # mutable-ok: router writes its routing decision back
|
||||
sanitized_forwardable_call_metadata(parent_metadata, SHADOW_EVAL_ROUTER_CALL_ORIGIN)
|
||||
)
|
||||
shadow_params: Final = { # mutable-ok: splatted as kwargs
|
||||
k: v for k, v in model_parameters.items() if k not in ("stream", "metadata")
|
||||
}
|
||||
try:
|
||||
response: Final = await router.acompletion(
|
||||
model=router_name,
|
||||
messages=messages, # pyright: ignore[reportArgumentType] # snapshot of the SDK's own message dicts
|
||||
metadata=shadow_metadata,
|
||||
num_retries=0,
|
||||
fallbacks=[], # mutable-ok: SDK kwarg; a failed shadow is a recorded error, never a spend multiplier
|
||||
**shadow_params,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # provider errors become error rows, not crashes
|
||||
verbose_logger.debug("shadow_eval: router call failed: %s", e)
|
||||
return _CallFailure(f"shadow router call failed: {e}")
|
||||
text: Final = self._extract_response_text(response)
|
||||
if not text:
|
||||
return _CallFailure("shadow router returned an empty response")
|
||||
raw_decision: Final = shadow_metadata.get("routing_decision")
|
||||
routing_decision: Final = raw_decision if isinstance(raw_decision, Mapping) else _EMPTY_METADATA
|
||||
raw_tier: Final = routing_decision.get("tier_label") or routing_decision.get("tier")
|
||||
return _ShadowResponse(
|
||||
text=text,
|
||||
model=str(getattr(response, "model", None) or routing_decision.get("routed_model") or ""),
|
||||
tier=str(raw_tier) if raw_tier is not None else None,
|
||||
)
|
||||
|
||||
async def _call_judge(
|
||||
self,
|
||||
judge_model: str,
|
||||
messages: Sequence[Mapping[str, object]],
|
||||
real_text: str,
|
||||
shadow_text: str,
|
||||
parent_metadata: Mapping[str, object],
|
||||
) -> "_JudgeVerdict | _CallFailure":
|
||||
"""Blind pairwise judge with A/B labels randomized to cancel position bias."""
|
||||
real_is_a: Final = random.random() < 0.5
|
||||
response_a: Final = real_text if real_is_a else shadow_text
|
||||
response_b: Final = shadow_text if real_is_a else real_text
|
||||
|
||||
conversation: Final = "\n".join(
|
||||
f"{str(m.get('role', 'user')).upper()}: {extract_text_from_content(m.get('content'))}"
|
||||
for m in messages
|
||||
if m.get("content") is not None
|
||||
)
|
||||
judge_metadata: Final = sanitized_forwardable_call_metadata(parent_metadata, SHADOW_EVAL_JUDGE_CALL_ORIGIN)
|
||||
judge_messages: Final = [ # mutable-ok: SDK takes a list
|
||||
{"role": "system", "content": PAIRWISE_JUDGE_SYSTEM_PROMPT}, # mutable-ok: SDK message
|
||||
{
|
||||
"role": "user",
|
||||
"content": _judge_user_prompt(conversation, response_a, response_b),
|
||||
}, # mutable-ok: SDK message
|
||||
]
|
||||
try:
|
||||
response: Final = await judge_acompletion(
|
||||
self._router_provider(),
|
||||
judge_model,
|
||||
judge_messages, # pyright: ignore[reportArgumentType] # plain SDK message dicts
|
||||
temperature=0,
|
||||
max_tokens=JUDGE_MAX_OUTPUT_TOKENS,
|
||||
metadata=judge_metadata,
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # judge outages become error rows, not crashes
|
||||
verbose_logger.debug("shadow_eval: judge call failed: %s", e)
|
||||
return _CallFailure(f"judge call failed: {e}")
|
||||
try:
|
||||
raw: Final = response["choices"][0]["message"]["content"] or ""
|
||||
verdict: Final = PairwiseVerdict.model_validate(parse_json_verdict(raw))
|
||||
except Exception as e: # noqa: BLE001 # malformed verdicts become error rows
|
||||
verbose_logger.debug("shadow_eval: unparseable judge verdict: %s", e)
|
||||
return _CallFailure(f"unparseable judge verdict: {e}", cost=_judge_call_cost(response))
|
||||
return _JudgeVerdict(
|
||||
preference=_unmask_preference(verdict.preference, real_is_a),
|
||||
confidence=max(0.0, min(1.0, verdict.confidence)),
|
||||
cost=_judge_call_cost(response),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _extract_response_text(response_obj: object) -> str:
|
||||
"""Extract the assistant's text from a ModelResponse-shaped object or dict."""
|
||||
try:
|
||||
content: Final = (
|
||||
response_obj["choices"][0]["message"]["content"]
|
||||
if isinstance(response_obj, Mapping)
|
||||
else response_obj.choices[0].message.content # pyright: ignore[reportAttributeAccessIssue] # duck-typed ModelResponse
|
||||
)
|
||||
except (AttributeError, KeyError, IndexError, TypeError):
|
||||
return ""
|
||||
return extract_text_from_content(content)
|
||||
|
||||
|
||||
_EMPTY_JOBS: Final[Mapping[str, ActiveShadowEvalJob]] = MappingProxyType({})
|
||||
|
||||
|
||||
def _default_prisma_provider() -> "PrismaClient | None":
|
||||
try:
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
except ImportError:
|
||||
return None
|
||||
return prisma_client
|
||||
|
|
@ -9,8 +9,8 @@ server-side using litellm router's search tools.
|
|||
import asyncio
|
||||
import math
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
from collections.abc import AsyncIterator, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, TypedDict, cast
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -37,10 +37,17 @@ from litellm.types.integrations.custom_logger import (
|
|||
AgenticLoopRequestPatch,
|
||||
)
|
||||
from litellm.types.integrations.websearch_interception import (
|
||||
AnthropicSearchQuery,
|
||||
AnthropicServerToolUseBlock,
|
||||
WebSearchInterceptionConfig,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import CallTypes, LlmProviders
|
||||
from litellm.types.utils import (
|
||||
AgenticLoopParams,
|
||||
CallTypes,
|
||||
LlmProviders,
|
||||
StandardLoggingUserAPIKeyMetadata,
|
||||
)
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -66,6 +73,23 @@ WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: Final = "_websearch_interception_emit_native_b
|
|||
WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks"
|
||||
|
||||
|
||||
class _PlanMetadataView(TypedDict):
|
||||
websearch_native_blocks: Sequence[Mapping[str, object]] | None
|
||||
|
||||
|
||||
class _AgenticLoopParamsView(TypedDict):
|
||||
agentic_loop_params: AgenticLoopParams
|
||||
|
||||
|
||||
class _WebSearchSettingsView(TypedDict):
|
||||
websearch_interception_params: WebSearchInterceptionConfig
|
||||
|
||||
|
||||
class _SearchToolConfig(TypedDict, total=False):
|
||||
search_tool_name: str
|
||||
litellm_params: Mapping[str, object] | None
|
||||
|
||||
|
||||
class WebSearchInterceptionLogger(CustomLogger):
|
||||
"""
|
||||
CustomLogger that intercepts WebSearch tool calls for models that don't
|
||||
|
|
@ -263,7 +287,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
return None
|
||||
|
||||
# Check if request has tools with native web_search
|
||||
tools: Final = kwargs.get("tools")
|
||||
tools: Final[Sequence[dict[str, object]] | None] = kwargs.get("tools")
|
||||
if not tools:
|
||||
return None
|
||||
|
||||
|
|
@ -312,7 +336,9 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
|
||||
return kwargs
|
||||
|
||||
def _convert_responses_tools(self, kwargs: Mapping[str, object], tools: list[dict[str, object]]) -> dict | None:
|
||||
def _convert_responses_tools(
|
||||
self, kwargs: Mapping[str, object], tools: Sequence[dict[str, object]]
|
||||
) -> dict[str, object] | None:
|
||||
"""Convert Responses API web search tools to the LiteLLM standard function tool."""
|
||||
if not any(is_web_search_tool_responses(tool) for tool in tools):
|
||||
return None
|
||||
|
|
@ -377,7 +403,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _tool_name(tool: dict[str, Any]) -> str | None:
|
||||
def _tool_name(tool: Mapping[str, object]) -> object:
|
||||
"""Effective tool name, handling OpenAI ``function`` wrapper shape."""
|
||||
fn: Final = tool.get("function")
|
||||
if tool.get("type") == "function" and isinstance(fn, dict):
|
||||
|
|
@ -385,7 +411,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
return tool.get("name")
|
||||
|
||||
@classmethod
|
||||
def _sync_forced_tool_choice(cls, tool_choice: Any, converted_tools: list[dict[str, object]]) -> object:
|
||||
def _sync_forced_tool_choice(cls, tool_choice: object, converted_tools: Sequence[Mapping[str, object]]) -> object:
|
||||
"""Repoint a forced ``tool_choice`` at ``litellm_web_search`` when it
|
||||
names a web-search tool that was just converted away.
|
||||
|
||||
|
|
@ -453,7 +479,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
kwargs[WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY] = True
|
||||
|
||||
# Convert native web search tools to LiteLLM standard
|
||||
converted_tools: Final = []
|
||||
converted_tools: Final[list[dict[str, object]]] = []
|
||||
for tool in tools:
|
||||
if is_web_search_tool(tool):
|
||||
standard_tool = get_litellm_web_search_tool()
|
||||
|
|
@ -824,7 +850,10 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
Anthropic-native clients (Claude Desktop, the Anthropic SDK) can
|
||||
render citations / sources alongside the model's textual reply.
|
||||
"""
|
||||
native_blocks: Final = plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY)
|
||||
metadata_view: Final[_PlanMetadataView] = {
|
||||
"websearch_native_blocks": plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY)
|
||||
}
|
||||
native_blocks: Final = metadata_view["websearch_native_blocks"]
|
||||
if not native_blocks:
|
||||
return response
|
||||
return self._inject_native_blocks(response, native_blocks)
|
||||
|
|
@ -833,22 +862,48 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
def _build_native_result_blocks(
|
||||
tool_calls: list[dict],
|
||||
structured_results: list[SearchResponse | None],
|
||||
) -> list[dict[str, object]]:
|
||||
"""Build one ``web_search_tool_result`` block per tool_call."""
|
||||
blocks: Final[list[dict[str, object]]] = []
|
||||
for i, tool_call in enumerate(tool_calls):
|
||||
tool_use_id = tool_call.get("id") or ""
|
||||
structured = structured_results[i] if i < len(structured_results) else None
|
||||
blocks.append(
|
||||
WebSearchTransformation.build_web_search_tool_result_block(
|
||||
tool_use_id=tool_use_id,
|
||||
search_response=structured,
|
||||
)
|
||||
) -> tuple[Mapping[str, object], ...]:
|
||||
"""
|
||||
Build a ``server_tool_use`` + ``web_search_tool_result`` pair per tool_call.
|
||||
|
||||
The pair is what Anthropic's spec requires: a bare result block, or one
|
||||
keyed by the model's ``toolu_...`` id instead of a ``srvtoolu_...`` one,
|
||||
is rejected on replay ("String should match pattern '^srvtoolu_'") and
|
||||
leaves native clients without a search to attach the sources to.
|
||||
"""
|
||||
return tuple(
|
||||
block
|
||||
for i, tool_call in enumerate(tool_calls)
|
||||
for block in WebSearchInterceptionLogger._native_result_pair(
|
||||
query=WebSearchInterceptionLogger._tool_call_query(tool_call),
|
||||
search_response=structured_results[i] if i < len(structured_results) else None,
|
||||
)
|
||||
return blocks
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _inject_native_blocks(response: Any, native_blocks: list[dict[str, object]]) -> Any:
|
||||
def _tool_call_query(tool_call: Mapping[str, object]) -> str:
|
||||
tool_input: Final = tool_call.get("input")
|
||||
if not isinstance(tool_input, Mapping):
|
||||
return ""
|
||||
query: Final = tool_input.get("query")
|
||||
return query if isinstance(query, str) else ""
|
||||
|
||||
@staticmethod
|
||||
def _native_result_pair(
|
||||
query: str,
|
||||
search_response: SearchResponse | None,
|
||||
) -> tuple[Mapping[str, object], Mapping[str, object]]:
|
||||
tool_use_id: Final = f"srvtoolu_{uuid.uuid4().hex}"
|
||||
return (
|
||||
AnthropicServerToolUseBlock(id=tool_use_id, input=AnthropicSearchQuery(query=query)).model_dump(),
|
||||
WebSearchTransformation.build_web_search_tool_result_block(
|
||||
tool_use_id=tool_use_id,
|
||||
search_response=search_response,
|
||||
),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any:
|
||||
"""Prepend native blocks to response content, dict or object form."""
|
||||
if not native_blocks:
|
||||
return response
|
||||
|
|
@ -1243,8 +1298,10 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
kwargs_for_followup: Final = self._prepare_followup_kwargs(kwargs)
|
||||
|
||||
if logging_obj is not None:
|
||||
agentic_params: Final = logging_obj.model_call_details.get("agentic_loop_params", {})
|
||||
full_model_name = agentic_params.get("model", model)
|
||||
agentic_view: Final[_AgenticLoopParamsView] = {
|
||||
"agentic_loop_params": logging_obj.model_call_details.get("agentic_loop_params", {})
|
||||
}
|
||||
full_model_name = agentic_view["agentic_loop_params"].get("model", model)
|
||||
verbose_logger.debug(
|
||||
"WebSearchInterception: Built anthropic request patch [call_id=%s model=%s messages=%d searches=%d]",
|
||||
_call_id,
|
||||
|
|
@ -1288,6 +1345,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router)
|
||||
search_provider: str | None = None
|
||||
search_litellm_params: dict[str, Any] = {}
|
||||
search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool)
|
||||
if search_tool is not None:
|
||||
await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs)
|
||||
search_litellm_params = dict(search_tool.get("litellm_params", {}) or {})
|
||||
|
|
@ -1304,12 +1362,30 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
verbose_logger.debug(
|
||||
"WebSearchInterception: Executing search for '%s' using provider '%s'", query, search_provider
|
||||
)
|
||||
user_api_key_auth: Final = self._get_user_api_key_auth_from_kwargs(kwargs)
|
||||
search_metadata: Final = (
|
||||
None
|
||||
if user_api_key_auth is None
|
||||
else self._build_search_request_metadata(
|
||||
user_api_key_auth=user_api_key_auth,
|
||||
search_tool_name=search_tool_name,
|
||||
)
|
||||
)
|
||||
search_kwargs: Final = {
|
||||
key: value
|
||||
for key, value in search_litellm_params.items()
|
||||
if key != "search_provider" and value is not None
|
||||
}
|
||||
result: Final = await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs)
|
||||
result: Final = (
|
||||
await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs)
|
||||
if search_metadata is None
|
||||
else await litellm.asearch(
|
||||
query=query,
|
||||
search_provider=search_provider,
|
||||
litellm_metadata=search_metadata,
|
||||
**search_kwargs,
|
||||
)
|
||||
)
|
||||
|
||||
# Format using transformation function
|
||||
search_result_text: Final = WebSearchTransformation.format_search_response(result)
|
||||
|
|
@ -1366,6 +1442,35 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
team_object=team_object,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _build_search_request_metadata(
|
||||
user_api_key_auth: "UserAPIKeyAuth",
|
||||
search_tool_name: str | None,
|
||||
) -> Mapping[str, object]:
|
||||
"""
|
||||
Spend-tracking metadata for the intercepted search, so its provider cost is logged
|
||||
and billed against the key/user/team that made the originating LLM request instead
|
||||
of being dropped by the proxy's spend hook for lack of an owner.
|
||||
"""
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
|
||||
user_api_key_metadata: Final[StandardLoggingUserAPIKeyMetadata] = (
|
||||
LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_auth)
|
||||
)
|
||||
return { # mutable-ok: litellm's metadata channel is a plain dict its logging path reads and enriches
|
||||
**user_api_key_metadata,
|
||||
"model_group": search_tool_name,
|
||||
"user_api_key": user_api_key_auth.api_key,
|
||||
"user_api_key_auth": user_api_key_auth,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _selected_search_tool_name(search_tool: Mapping[str, object] | None) -> str | None:
|
||||
if search_tool is None:
|
||||
return None
|
||||
search_tool_name: Final = search_tool.get("search_tool_name")
|
||||
return search_tool_name if isinstance(search_tool_name, str) and search_tool_name else None
|
||||
|
||||
@staticmethod
|
||||
def _get_user_api_key_auth_from_kwargs(kwargs: Mapping[str, object] | None) -> "UserAPIKeyAuth | None":
|
||||
if not kwargs:
|
||||
|
|
@ -1387,7 +1492,7 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
|
||||
return None
|
||||
|
||||
def _select_search_tool_from_router(self, llm_router: object) -> dict[str, Any] | None:
|
||||
def _select_search_tool_from_router(self, llm_router: object) -> "_SearchToolConfig | None":
|
||||
if llm_router is None or not hasattr(llm_router, "search_tools"):
|
||||
return None
|
||||
search_tools: Final = list(getattr(llm_router, "search_tools") or [])
|
||||
|
|
@ -1395,9 +1500,9 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
|
||||
def _select_search_tool_from_list(
|
||||
self,
|
||||
search_tools: list[dict[str, Any]],
|
||||
search_tools: list[_SearchToolConfig],
|
||||
source: str,
|
||||
) -> dict[str, Any] | None:
|
||||
) -> "_SearchToolConfig | None":
|
||||
if self.search_tool_name:
|
||||
matching_tools = [tool for tool in search_tools if tool.get("search_tool_name") == self.search_tool_name]
|
||||
if matching_tools:
|
||||
|
|
@ -1592,8 +1697,8 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
|
||||
@staticmethod
|
||||
def initialize_from_proxy_config(
|
||||
litellm_settings: dict[str, Any],
|
||||
callback_specific_params: dict[str, Any],
|
||||
litellm_settings: Mapping[str, WebSearchInterceptionConfig],
|
||||
callback_specific_params: Mapping[str, object],
|
||||
) -> "WebSearchInterceptionLogger":
|
||||
"""
|
||||
Static method to initialize WebSearchInterceptionLogger from proxy config.
|
||||
|
|
@ -1617,7 +1722,10 @@ class WebSearchInterceptionLogger(CustomLogger):
|
|||
# Get websearch_interception_params from litellm_settings or callback_specific_params
|
||||
websearch_params: WebSearchInterceptionConfig = {}
|
||||
if "websearch_interception_params" in litellm_settings:
|
||||
websearch_params = litellm_settings["websearch_interception_params"]
|
||||
settings_view: Final[_WebSearchSettingsView] = {
|
||||
"websearch_interception_params": litellm_settings["websearch_interception_params"]
|
||||
}
|
||||
websearch_params = settings_view["websearch_interception_params"]
|
||||
elif "websearch_interception" in callback_specific_params and isinstance(
|
||||
callback_specific_params["websearch_interception"], dict
|
||||
):
|
||||
|
|
|
|||
|
|
@ -412,6 +412,15 @@ class WebSearchTransformation:
|
|||
block that should accompany the model's text reply when the original
|
||||
request used a native ``web_search_*`` tool.
|
||||
|
||||
The spec'd shape carries page text only in ``encrypted_content``, an
|
||||
opaque server-issued blob that we cannot mint. Emitting the four spec
|
||||
fields alone would drop the snippet entirely, leaving the client (and
|
||||
the model, on any replayed follow-up turn) with URLs and titles but no
|
||||
evidence to answer from, forcing a fetch per result. So the snippet is
|
||||
carried in an additive ``snippet`` key alongside the spec fields.
|
||||
``encrypted_content`` stays empty rather than holding plaintext, which
|
||||
would assert encryption semantics that do not hold.
|
||||
|
||||
Spec reference:
|
||||
https://docs.anthropic.com/en/api/web-search-tool
|
||||
|
||||
|
|
@ -438,6 +447,7 @@ class WebSearchTransformation:
|
|||
"title": title,
|
||||
"page_age": page_age,
|
||||
"encrypted_content": "",
|
||||
"snippet": getattr(r, "snippet", "") or "",
|
||||
}
|
||||
)
|
||||
return {
|
||||
|
|
|
|||
|
|
@ -2,8 +2,8 @@
|
|||
Handler for transforming interactions API requests to litellm.responses requests.
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncIterator, Coroutine, Iterator
|
||||
from typing import Any, Final, cast
|
||||
from collections.abc import AsyncIterator, Callable, Coroutine, Iterator
|
||||
from typing import Any, Final
|
||||
|
||||
import litellm
|
||||
from litellm.interactions.litellm_responses_transformation.streaming_iterator import (
|
||||
|
|
@ -37,7 +37,7 @@ class LiteLLMResponsesInteractionsHandler:
|
|||
) -> (
|
||||
InteractionsAPIResponse
|
||||
| Iterator[InteractionsAPIStreamingResponse]
|
||||
| Coroutine[Any, Any, InteractionsAPIResponse | AsyncIterator[InteractionsAPIStreamingResponse]]
|
||||
| Coroutine[object, object, InteractionsAPIResponse | AsyncIterator[InteractionsAPIStreamingResponse]]
|
||||
):
|
||||
"""
|
||||
Handle Interactions API request by calling litellm.responses().
|
||||
|
|
@ -55,13 +55,15 @@ class LiteLLMResponsesInteractionsHandler:
|
|||
InteractionsAPIResponse or streaming iterator
|
||||
"""
|
||||
# Transform interactions request to responses request
|
||||
responses_request = LiteLLMResponsesInteractionsConfig.transform_interactions_request_to_responses_request(
|
||||
model=model,
|
||||
input=input,
|
||||
optional_params=optional_params,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
stream=stream,
|
||||
**kwargs,
|
||||
responses_request: Final = (
|
||||
LiteLLMResponsesInteractionsConfig.transform_interactions_request_to_responses_request(
|
||||
model=model,
|
||||
input=input,
|
||||
optional_params=optional_params,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
stream=stream,
|
||||
**kwargs,
|
||||
)
|
||||
)
|
||||
|
||||
if _is_async:
|
||||
|
|
@ -76,7 +78,10 @@ class LiteLLMResponsesInteractionsHandler:
|
|||
# Call litellm.responses()
|
||||
# Note: litellm.responses() returns Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]
|
||||
# but the type checker may see it as a coroutine in some contexts
|
||||
responses_response: Final = litellm.responses(
|
||||
responses_fn: Final[Callable[..., ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]] = vars(litellm)[
|
||||
"responses"
|
||||
]
|
||||
responses_response: Final = responses_fn(
|
||||
**responses_request,
|
||||
)
|
||||
|
||||
|
|
@ -92,8 +97,7 @@ class LiteLLMResponsesInteractionsHandler:
|
|||
)
|
||||
|
||||
# At this point, responses_response must be ResponsesAPIResponse (not streaming)
|
||||
# Cast to satisfy type checker since we've already checked it's not a streaming iterator
|
||||
responses_api_response: Final = cast(ResponsesAPIResponse, responses_response)
|
||||
responses_api_response: Final = responses_response
|
||||
|
||||
# Transform responses response to interactions response
|
||||
return LiteLLMResponsesInteractionsConfig.transform_responses_response_to_interactions_response(
|
||||
|
|
@ -112,7 +116,10 @@ class LiteLLMResponsesInteractionsHandler:
|
|||
"""Async handler for interactions API requests."""
|
||||
# Call litellm.aresponses()
|
||||
# Note: litellm.aresponses() returns Union[ResponsesAPIResponse, BaseResponsesAPIStreamingIterator]
|
||||
responses_response: Final = await litellm.aresponses(
|
||||
aresponses_fn: Final[
|
||||
Callable[..., Coroutine[object, object, ResponsesAPIResponse | BaseResponsesAPIStreamingIterator]]
|
||||
] = vars(litellm)["aresponses"]
|
||||
responses_response: Final = await aresponses_fn(
|
||||
**responses_request,
|
||||
)
|
||||
|
||||
|
|
@ -128,8 +135,7 @@ class LiteLLMResponsesInteractionsHandler:
|
|||
)
|
||||
|
||||
# At this point, responses_response must be ResponsesAPIResponse (not streaming)
|
||||
# Cast to satisfy type checker since we've already checked it's not a streaming iterator
|
||||
responses_api_response: Final = cast(ResponsesAPIResponse, responses_response)
|
||||
responses_api_response: Final = responses_response
|
||||
|
||||
# Transform responses response to interactions response
|
||||
return LiteLLMResponsesInteractionsConfig.transform_responses_response_to_interactions_response(
|
||||
|
|
|
|||
|
|
@ -2,12 +2,16 @@
|
|||
Transformation utilities for bridging Interactions API to Responses API.
|
||||
|
||||
This module handles transforming between:
|
||||
- Interactions API format (Google's format with Turn[], system_instruction, etc.)
|
||||
- Interactions API format (Google's format with Step[]/Turn[], system_instruction, etc.)
|
||||
- Responses API format (OpenAI's format with input[], instructions, etc.)
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, cast
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from litellm.types.interactions import (
|
||||
InteractionInput,
|
||||
InteractionsAPIOptionalRequestParams,
|
||||
|
|
@ -19,6 +23,8 @@ from litellm.types.llms.openai import (
|
|||
ResponsesAPIResponse,
|
||||
)
|
||||
|
||||
_STEP_TYPE_ROLES: Final = MappingProxyType({"user_input": "user", "model_output": "assistant"})
|
||||
|
||||
|
||||
class LiteLLMResponsesInteractionsConfig:
|
||||
"""Configuration class for transforming between Interactions API and Responses API."""
|
||||
|
|
@ -91,112 +97,94 @@ class LiteLLMResponsesInteractionsConfig:
|
|||
|
||||
Interactions API input can be:
|
||||
- string: "Hello"
|
||||
- Turn[]: [{"role": "user", "content": [...]}]
|
||||
- Content object
|
||||
- Step[]: [{"type": "user_input", "content": [...]}, {"type": "model_output", "content": [...]}]
|
||||
- Turn[] (legacy): [{"role": "user", "content": [...]}]
|
||||
- Content | Content[]: one user message worth of content parts
|
||||
|
||||
Responses API input is:
|
||||
- string: "Hello"
|
||||
- Message[]: [{"role": "user", "content": [...]}]
|
||||
- Message[]: [{"role": "user", "content": [{"type": "input_text", ...}]}]
|
||||
"""
|
||||
if isinstance(input, str):
|
||||
# ResponseInputParam accepts str
|
||||
return cast(ResponseInputParam, input)
|
||||
|
||||
if isinstance(input, list):
|
||||
# Turn[] format - convert to Responses API Message[] format
|
||||
messages: Final = []
|
||||
for turn in input:
|
||||
if isinstance(turn, dict):
|
||||
role = turn.get("role", "user")
|
||||
content = turn.get("content", [])
|
||||
transformed: Final = (
|
||||
[
|
||||
LiteLLMResponsesInteractionsConfig._transform_history_item(item)
|
||||
for item in input
|
||||
if LiteLLMResponsesInteractionsConfig._is_history_item(item)
|
||||
]
|
||||
if any(LiteLLMResponsesInteractionsConfig._is_history_item(item) for item in input)
|
||||
else [
|
||||
{
|
||||
"role": "user",
|
||||
"content": LiteLLMResponsesInteractionsConfig._transform_content_array(input, "user"),
|
||||
}
|
||||
]
|
||||
)
|
||||
return cast(ResponseInputParam, transformed)
|
||||
|
||||
# Transform content array
|
||||
transformed_content = LiteLLMResponsesInteractionsConfig._transform_content_array(content)
|
||||
|
||||
messages.append(
|
||||
{
|
||||
"role": role,
|
||||
"content": transformed_content,
|
||||
}
|
||||
)
|
||||
elif isinstance(turn, Turn):
|
||||
# Pydantic model
|
||||
role = turn.role if hasattr(turn, "role") else "user"
|
||||
content = turn.content if hasattr(turn, "content") else []
|
||||
|
||||
# Ensure content is a list for _transform_content_array
|
||||
# Cast to List[Any] to handle various content types
|
||||
if isinstance(content, list):
|
||||
content_list: list[Any] = list(content)
|
||||
elif content is not None:
|
||||
content_list = [content]
|
||||
else:
|
||||
content_list = []
|
||||
|
||||
transformed_content = LiteLLMResponsesInteractionsConfig._transform_content_array(content_list)
|
||||
|
||||
messages.append(
|
||||
{
|
||||
"role": role,
|
||||
"content": transformed_content,
|
||||
}
|
||||
)
|
||||
|
||||
return cast(ResponseInputParam, messages)
|
||||
|
||||
# Single content object - wrap in message
|
||||
if isinstance(input, dict):
|
||||
raw_content: Final = input.get("content")
|
||||
content_items: Final = raw_content if isinstance(raw_content, list) else [input]
|
||||
return cast(
|
||||
ResponseInputParam,
|
||||
[
|
||||
{
|
||||
"role": "user",
|
||||
"content": LiteLLMResponsesInteractionsConfig._transform_content_array(
|
||||
input.get("content", []) if isinstance(input.get("content"), list) else [input]
|
||||
),
|
||||
"content": LiteLLMResponsesInteractionsConfig._transform_content_array(content_items, "user"),
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
# Fallback: convert to string
|
||||
return cast(ResponseInputParam, str(input))
|
||||
|
||||
@staticmethod
|
||||
def _transform_content_array(content: list[Any]) -> list[dict[str, Any]]:
|
||||
"""Transform Interactions API content array to Responses API format."""
|
||||
if not isinstance(content, list):
|
||||
# Single content item - wrap in array
|
||||
content = [content]
|
||||
def _is_history_item(item: object) -> bool:
|
||||
if isinstance(item, Turn):
|
||||
return True
|
||||
return isinstance(item, dict) and ("role" in item or item.get("type") in _STEP_TYPE_ROLES)
|
||||
|
||||
transformed: Final[list[dict[str, Any]]] = []
|
||||
for item in content:
|
||||
if isinstance(item, dict):
|
||||
# Already in dict format, pass through
|
||||
transformed.append(item)
|
||||
elif isinstance(item, str):
|
||||
# Plain string - wrap in text format
|
||||
transformed.append({"type": "text", "text": item})
|
||||
else:
|
||||
# Pydantic model or other - convert to dict
|
||||
if hasattr(item, "model_dump"):
|
||||
dumped = item.model_dump()
|
||||
if isinstance(dumped, dict):
|
||||
transformed.append(dumped)
|
||||
else:
|
||||
# Fallback: wrap in text format
|
||||
transformed.append({"type": "text", "text": str(dumped)})
|
||||
elif hasattr(item, "dict"):
|
||||
dumped = item.dict()
|
||||
if isinstance(dumped, dict):
|
||||
transformed.append(dumped)
|
||||
else:
|
||||
# Fallback: wrap in text format
|
||||
transformed.append({"type": "text", "text": str(dumped)})
|
||||
else:
|
||||
# Fallback: wrap in text format
|
||||
transformed.append({"type": "text", "text": str(item)})
|
||||
@staticmethod
|
||||
def _transform_history_item(item: object) -> Mapping[str, object]:
|
||||
raw: Final = item.model_dump(exclude_none=True) if isinstance(item, Turn) else item
|
||||
fields: Final = raw if isinstance(raw, Mapping) else {}
|
||||
role: Final = LiteLLMResponsesInteractionsConfig._responses_role(fields)
|
||||
raw_content: Final = fields.get("content")
|
||||
content_items: Final = (
|
||||
raw_content if isinstance(raw_content, list) else [] if raw_content is None else [raw_content]
|
||||
)
|
||||
return {
|
||||
"role": role,
|
||||
"content": LiteLLMResponsesInteractionsConfig._transform_content_array(content_items, role),
|
||||
}
|
||||
|
||||
return transformed
|
||||
@staticmethod
|
||||
def _responses_role(item: Mapping[str, object]) -> str:
|
||||
step_role: Final = _STEP_TYPE_ROLES.get(str(item.get("type", "")))
|
||||
if step_role is not None:
|
||||
return step_role
|
||||
raw_role: Final = str(item.get("role") or "user")
|
||||
return "assistant" if raw_role == "model" else raw_role
|
||||
|
||||
@staticmethod
|
||||
def _transform_content_array(content: Sequence[object], role: str) -> Sequence[Mapping[str, object]]:
|
||||
"""Transform Interactions API content parts to Responses API parts for the given role."""
|
||||
return [LiteLLMResponsesInteractionsConfig._transform_content_item(item, role) for item in content]
|
||||
|
||||
@staticmethod
|
||||
def _transform_content_item(item: object, role: str) -> Mapping[str, object]:
|
||||
text_type: Final = "output_text" if role == "assistant" else "input_text"
|
||||
if isinstance(item, str):
|
||||
return {"type": text_type, "text": item}
|
||||
if isinstance(item, Mapping):
|
||||
if item.get("type") == "text":
|
||||
return {"type": text_type, "text": str(item.get("text", ""))}
|
||||
return item
|
||||
if isinstance(item, BaseModel):
|
||||
return LiteLLMResponsesInteractionsConfig._transform_content_item(item.model_dump(exclude_none=True), role)
|
||||
return {"type": text_type, "text": str(item)}
|
||||
|
||||
@staticmethod
|
||||
def transform_responses_response_to_interactions_response(
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ This module has no dependencies on proxy code and can be safely imported at the
|
|||
import json
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
|
|
@ -71,7 +72,7 @@ def get_litellm_gateway_api_key(
|
|||
return token_data["key"]
|
||||
|
||||
|
||||
def is_cli_token_fresh(token_data: dict, buffer_hours: float = 0.1) -> bool:
|
||||
def is_cli_token_fresh(token_data: Mapping[str, object], buffer_hours: float = 0.1) -> bool:
|
||||
"""Check whether a cached CLI token (as stored in token.json) is still
|
||||
within its expiration window. Used by `lite auth print-token` to fail
|
||||
fast, without a network round trip, once the cached token is past
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
# What is this?
|
||||
## Helper utilities
|
||||
import copy
|
||||
from collections.abc import Iterable
|
||||
from collections.abc import Iterable, Mapping
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal
|
||||
|
||||
import httpx
|
||||
|
|
@ -181,7 +181,7 @@ def add_missing_spend_metadata_to_litellm_metadata(litellm_metadata: dict, metad
|
|||
|
||||
|
||||
def get_metadata_variable_name_from_kwargs(
|
||||
kwargs: dict,
|
||||
kwargs: Mapping[str, object],
|
||||
) -> Literal["metadata", "litellm_metadata"]:
|
||||
"""
|
||||
Helper to return what the "metadata" field should be called in the request data
|
||||
|
|
|
|||
94
litellm/litellm_core_utils/internal_call_metadata.py
Normal file
94
litellm/litellm_core_utils/internal_call_metadata.py
Normal file
|
|
@ -0,0 +1,94 @@
|
|||
"""Metadata a request forwards to the internal LLM sub-calls it triggers.
|
||||
|
||||
Internal features (the auto-router's classifier and embeddings, shadow eval's shadow and
|
||||
judge calls) bill real provider spend that nobody typed a prompt for. That spend must land
|
||||
on the same key/team/org/user as the request that caused it, so the sub-call carries the
|
||||
caller's identity metadata, minus two things that must never be forwarded as-is:
|
||||
|
||||
* ``user_api_key_budget_reservation`` (and the reservation nested inside
|
||||
``user_api_key_auth``) belongs to the parent completion. If a sub-call's cost callback
|
||||
sees it, that callback finalizes the reservation and the parent's own callback then
|
||||
skips incrementing the key/team budget counters, losing the parent's spend.
|
||||
``user_api_key_auth`` itself is kept, sanitized, because model access-group filtering
|
||||
needs it.
|
||||
* The sub-call is stamped with ``INTERNAL_CALL_ORIGIN_METADATA_KEY`` so its spend log row
|
||||
records that it is not traffic the caller sent.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY
|
||||
from litellm.types.utils import InternalCallOrigin
|
||||
|
||||
BUDGET_RESERVATION_METADATA_KEYS: Final = frozenset({"user_api_key_budget_reservation"})
|
||||
|
||||
_USER_API_KEY_AUTH_KEY: Final = "user_api_key_auth"
|
||||
|
||||
FORWARDABLE_IDENTITY_METADATA_KEYS: Final = frozenset(
|
||||
{
|
||||
"user_api_key",
|
||||
"user_api_key_hash",
|
||||
"user_api_key_alias",
|
||||
"user_api_key_team_id",
|
||||
"user_api_key_org_id",
|
||||
"user_api_key_user_id",
|
||||
"user_api_key_end_user_id",
|
||||
_USER_API_KEY_AUTH_KEY,
|
||||
}
|
||||
)
|
||||
"""The caller-identity subset a detached sub-call needs to be attributed and
|
||||
budget-checked like the request that spawned it. Everything else on the parent's metadata
|
||||
(routing decision, guardrail state, logging payload) describes the parent call and would
|
||||
be a lie on a sub-call that runs after it returned."""
|
||||
|
||||
|
||||
def sanitize_user_api_key_auth(auth: object) -> object:
|
||||
"""Copy of the auth object with its budget reservation removed; the cost callback
|
||||
falls back to reading the reservation from inside the auth object."""
|
||||
if isinstance(auth, dict):
|
||||
return {k: v for k, v in auth.items() if k != "budget_reservation"} # mutable-ok: SDK metadata value
|
||||
reservation: Final[object] = getattr(auth, "budget_reservation", None)
|
||||
model_copy: Final[object] = getattr(auth, "model_copy", None)
|
||||
if reservation is not None and callable(model_copy):
|
||||
return model_copy(update={"budget_reservation": None}) # mutable-ok: pydantic update payload
|
||||
return auth
|
||||
|
||||
|
||||
def _sanitized(parent_metadata: Mapping[str, object]) -> dict[str, object]: # mutable-ok: SDK metadata kwarg
|
||||
return { # mutable-ok: SDK metadata kwarg
|
||||
k: sanitize_user_api_key_auth(v) if k == _USER_API_KEY_AUTH_KEY else v
|
||||
for k, v in parent_metadata.items()
|
||||
if k not in BUDGET_RESERVATION_METADATA_KEYS
|
||||
}
|
||||
|
||||
|
||||
def forwarded_internal_call_metadata(
|
||||
parent_metadata: Mapping[str, object] | None,
|
||||
call_origin: InternalCallOrigin,
|
||||
) -> dict[str, object]: # mutable-ok: SDK metadata kwarg
|
||||
"""Parent metadata, minus its budget reservation, stamped with the sub-call's origin.
|
||||
|
||||
For sub-calls made inside the parent request (classifier, embeddings), where the
|
||||
parent's full context still describes the call being made.
|
||||
"""
|
||||
if not parent_metadata:
|
||||
return {} # mutable-ok: SDK metadata kwarg
|
||||
return _sanitized(parent_metadata) | { # mutable-ok: SDK metadata kwarg
|
||||
INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin
|
||||
}
|
||||
|
||||
|
||||
def sanitized_forwardable_call_metadata(
|
||||
parent_metadata: Mapping[str, object],
|
||||
call_origin: InternalCallOrigin,
|
||||
) -> dict[str, object]: # mutable-ok: SDK metadata kwarg
|
||||
"""Just the caller's identity, stamped with the sub-call's origin.
|
||||
|
||||
For sub-calls detached from the parent request (shadow eval), which outlive it and
|
||||
must not inherit per-request state such as its routing decision or logging payload.
|
||||
"""
|
||||
identity: Final = {k: v for k, v in parent_metadata.items() if k in FORWARDABLE_IDENTITY_METADATA_KEYS}
|
||||
return _sanitized(identity) | {INTERNAL_CALL_ORIGIN_METADATA_KEY: call_origin} # mutable-ok: SDK metadata kwarg
|
||||
|
|
@ -10,7 +10,7 @@ import subprocess
|
|||
import sys
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from datetime import datetime as dt_object
|
||||
from functools import lru_cache
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast
|
||||
|
|
@ -25,7 +25,15 @@ from litellm import (
|
|||
log_raw_request_response,
|
||||
turn_off_message_logging,
|
||||
)
|
||||
from litellm._logging import _is_debugging_on, _redact_string, verbose_logger
|
||||
from litellm._logging import (
|
||||
_is_debugging_on,
|
||||
_redact_string,
|
||||
session_id_var,
|
||||
set_session_id,
|
||||
set_trace_id,
|
||||
trace_id_var,
|
||||
verbose_logger,
|
||||
)
|
||||
from litellm._uuid import uuid
|
||||
from litellm.batches.batch_utils import _handle_completed_batch
|
||||
from litellm.caching.caching import DualCache, InMemoryCache
|
||||
|
|
@ -168,6 +176,9 @@ from .initialize_dynamic_callback_params import (
|
|||
from .specialty_caches.dynamic_logging_cache import DynamicLoggingCache
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mcp.types import EmbeddedResource, ImageContent, TextContent
|
||||
|
||||
from litellm.integrations.otel.logger import OpenTelemetryV2
|
||||
from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig
|
||||
try:
|
||||
from litellm_enterprise.enterprise_callbacks.callback_controls import (
|
||||
|
|
@ -203,14 +214,30 @@ except Exception as e:
|
|||
PagerDutyAlerting = CustomLogger
|
||||
EnterpriseCallbackControls = None
|
||||
EnterpriseStandardLoggingPayloadSetupVAR = None
|
||||
_in_memory_loggers: Final[list[Any]] = []
|
||||
if TYPE_CHECKING:
|
||||
from litellm.integrations.generic_api.generic_api_callback import (
|
||||
GenericAPILogger as _GenericAPILoggerCls,
|
||||
)
|
||||
|
||||
_STANDARD_LOGGING_METADATA_KEYS: Final[frozenset] = frozenset(StandardLoggingMetadata.__annotations__.keys())
|
||||
_GENERIC_API_LOGGER_CLS: Final = _GenericAPILoggerCls
|
||||
_RESEND_EMAIL_LOGGER_FACTORY: Final = CustomLogger
|
||||
_SENDGRID_EMAIL_LOGGER_FACTORY: Final = CustomLogger
|
||||
_SMTP_EMAIL_LOGGER_FACTORY: Final = CustomLogger
|
||||
_PAGERDUTY_ALERTING_FACTORY: Final = CustomLogger
|
||||
else:
|
||||
_GENERIC_API_LOGGER_CLS: Final = GenericAPILogger
|
||||
_RESEND_EMAIL_LOGGER_FACTORY: Final = ResendEmailLogger
|
||||
_SENDGRID_EMAIL_LOGGER_FACTORY: Final = SendGridEmailLogger
|
||||
_SMTP_EMAIL_LOGGER_FACTORY: Final = SMTPEmailLogger
|
||||
_PAGERDUTY_ALERTING_FACTORY: Final = PagerDutyAlerting
|
||||
_in_memory_loggers: Final[list[CustomLogger]] = []
|
||||
|
||||
_STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = frozenset(StandardLoggingMetadata.__annotations__.keys())
|
||||
|
||||
### GLOBAL VARIABLES ###
|
||||
|
||||
# Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys
|
||||
_CUSTOM_PRICING_KEYS: Final[frozenset] = frozenset(CustomPricingLiteLLMParams.model_fields.keys())
|
||||
_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys())
|
||||
|
||||
sentry_sdk_instance = None
|
||||
capture_exception = None
|
||||
|
|
@ -313,6 +340,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
applied_guardrails: list[str] | None = None,
|
||||
kwargs: dict | None = None,
|
||||
log_raw_request_response: bool = False,
|
||||
supports_correlation_logging: bool = True,
|
||||
):
|
||||
_input: Final[str | None] = messages # save original value of messages
|
||||
if messages is not None:
|
||||
|
|
@ -338,6 +366,36 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self.call_type = call_type
|
||||
self.litellm_call_id = litellm_call_id
|
||||
self.litellm_trace_id: str = litellm_trace_id if litellm_trace_id else str(uuid.uuid4())
|
||||
|
||||
# Capture the pre-call *value* (not a contextvars.Token) so restoration works
|
||||
# even if this attempt's own logging ends up dispatched onto a different
|
||||
# asyncio Task/context (e.g. via asyncio.create_task or the logging worker) -
|
||||
# a Token can only be reset in the exact Context where it was created.
|
||||
self._pre_call_trace_id: str = trace_id_var.get()
|
||||
self._pre_call_session_id: str = session_id_var.get()
|
||||
_sid: Final = kwargs.get("litellm_session_id") if kwargs else None
|
||||
self.litellm_session_id: str = str(_sid) if _sid else ""
|
||||
# supports_correlation_logging is False for calls originating from the
|
||||
# sync client entry point (wrapper() in utils.py): a plain OS thread
|
||||
# has no per-call context isolation the way an asyncio Task does, and
|
||||
# a thread pool's worker threads are recycled across unrelated
|
||||
# requests, so stamping trace_id/session_id there risks one request's
|
||||
# ids leaking into a different, later request on the same thread. Sync
|
||||
# support is deferred to a follow-up PR with its own safe-restore
|
||||
# mechanism; async calls (the proxy's only call path) are unaffected.
|
||||
if supports_correlation_logging:
|
||||
set_trace_id(self.litellm_trace_id)
|
||||
set_session_id(self.litellm_session_id)
|
||||
# set_trace_id()/set_session_id() sanitize (strip control chars, bound
|
||||
# length) before storing, so the contextvar's actual value can differ
|
||||
# from self.litellm_trace_id/litellm_session_id. Capture what was
|
||||
# really stored - _restore_correlation_context_if_unclaimed() must
|
||||
# compare against this, not the raw ids, or a caller-supplied id
|
||||
# containing control characters/oversized input would never match
|
||||
# and cleanup would be skipped forever.
|
||||
self._own_trace_id: str = trace_id_var.get()
|
||||
self._own_session_id: str = session_id_var.get()
|
||||
|
||||
self.function_id = function_id
|
||||
self.streaming_chunks: list[Any] = [] # for generating complete stream response
|
||||
self.sync_streaming_chunks: list[Any] = [] # for generating complete stream response
|
||||
|
|
@ -1246,7 +1304,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
verbose_logger.exception("LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e)
|
||||
return response_obj
|
||||
|
||||
def _parse_post_mcp_call_hook_response(self, response: MCPPostCallResponseObject | None) -> Any:
|
||||
def _parse_post_mcp_call_hook_response(
|
||||
self, response: MCPPostCallResponseObject | None
|
||||
) -> "Sequence[TextContent | ImageContent | EmbeddedResource] | None":
|
||||
"""
|
||||
Parse the response from the post_mcp_tool_call_hook
|
||||
|
||||
|
|
@ -1690,7 +1750,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self.completion_start_time = completion_start_time
|
||||
self.model_call_details["completion_start_time"] = self.completion_start_time
|
||||
|
||||
def normalize_logging_result(self, result: Any) -> Any:
|
||||
def normalize_logging_result(self, result: Any) -> object:
|
||||
"""
|
||||
Some endpoints return a different type of result than what is expected by the logging system.
|
||||
This function is used to normalize the result to the expected type.
|
||||
|
|
@ -1726,7 +1786,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
)
|
||||
return logging_result
|
||||
|
||||
def _merge_hidden_params_from_response_into_metadata(self, logging_result: Any) -> None:
|
||||
def _merge_hidden_params_from_response_into_metadata(self, logging_result: object) -> None:
|
||||
"""
|
||||
Copy response._hidden_params into litellm_params.metadata['hidden_params'].
|
||||
|
||||
|
|
@ -1787,7 +1847,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None:
|
||||
emit_standard_logging_payload(standard_logging_payload)
|
||||
|
||||
def _build_standard_logging_payload(self, init_response_obj: Any, start_time: Any, end_time: Any) -> Any:
|
||||
def _build_standard_logging_payload(
|
||||
self, init_response_obj: object, start_time: Any, end_time: Any
|
||||
) -> StandardLoggingPayload | None:
|
||||
"""Build StandardLoggingPayload and accumulate its construction time."""
|
||||
_start: Final = time.time()
|
||||
payload: Final = get_standard_logging_object_payload(
|
||||
|
|
@ -1908,7 +1970,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
def _is_recognized_call_type_for_logging(
|
||||
self,
|
||||
logging_result: Any,
|
||||
logging_result: object,
|
||||
):
|
||||
"""
|
||||
Returns True if the call type is recognized for logging (eg. ModelResponse, ModelResponseStream, etc.)
|
||||
|
|
@ -1992,7 +2054,67 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
if complete_streaming_response is not None:
|
||||
await self.async_success_handler(result=complete_streaming_response)
|
||||
|
||||
def success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs):
|
||||
def _restore_correlation_context(self) -> None:
|
||||
"""Restore trace_id/session_id contextvars to their pre-call value.
|
||||
|
||||
Without this, a nested LiteLLM call sharing the same asyncio Task as an
|
||||
outer request (e.g. a guardrail's own LLM-as-judge call, an MCP sampling
|
||||
call) would leave the outer request's subsequent log lines stamped with
|
||||
the nested call's trace_id/session_id instead of its own.
|
||||
|
||||
Uses a plain set() of the captured pre-call value rather than
|
||||
contextvars.Token-based reset(), since this can end up called from a
|
||||
different asyncio Task/context than __init__ ran in (e.g. the request
|
||||
task's own wrapper() finally block, plus async_success_handler
|
||||
dispatched separately via asyncio.create_task/the logging worker) -
|
||||
reset() only works in the exact Context a Token was created in and
|
||||
raises otherwise. Deliberately NOT idempotent/guarded: each distinct
|
||||
Task that calls this needs its own restore to actually take effect in
|
||||
that Task's view of the contextvars, so calling it multiple times
|
||||
(once per Task involved in this attempt) is required, not just safe.
|
||||
"""
|
||||
set_trace_id(self._pre_call_trace_id)
|
||||
set_session_id(self._pre_call_session_id)
|
||||
|
||||
def _restore_correlation_context_if_unclaimed(self) -> None:
|
||||
"""Guarded variant for __del__-triggered cleanup only.
|
||||
|
||||
__del__ can fire arbitrarily late (delayed by cyclic GC, possibly
|
||||
after the consuming Task/thread has already moved on to a different,
|
||||
still-active call). Unconditionally restoring in that case would
|
||||
stomp the active call's trace_id/session_id with this abandoned
|
||||
stream's stale pre-call snapshot. Only restore if the contextvars
|
||||
still hold the ids *this* call set - i.e. nothing has claimed them
|
||||
since - so an unrelated active call is never overwritten.
|
||||
"""
|
||||
if trace_id_var.get() == self._own_trace_id and session_id_var.get() == self._own_session_id:
|
||||
self._restore_correlation_context()
|
||||
|
||||
def success_handler(
|
||||
self,
|
||||
result: Any = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml)
|
||||
start_time: datetime.datetime | None = None,
|
||||
end_time: datetime.datetime | None = None,
|
||||
cache_hit: bool | None = None,
|
||||
**kwargs: Any, # kwargs-ok: forwarded to _success_handler_body
|
||||
) -> None:
|
||||
"""Restores trace_id/session_id contextvars once this attempt's own success
|
||||
logging (including any nested calls its callbacks trigger) is fully done."""
|
||||
try:
|
||||
return self._success_handler_body(
|
||||
result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs
|
||||
)
|
||||
finally:
|
||||
self._restore_correlation_context()
|
||||
|
||||
def _success_handler_body(
|
||||
self,
|
||||
result: Any = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml)
|
||||
start_time: datetime.datetime | None = None,
|
||||
end_time: datetime.datetime | None = None,
|
||||
cache_hit: bool | None = None,
|
||||
**kwargs: Any, # kwargs-ok: forwarded from success_handler
|
||||
) -> None:
|
||||
verbose_logger.debug("Logging Details LiteLLM-Success Call: Cache_hit=%s", cache_hit)
|
||||
if not self.should_run_logging(event_type="sync_success"): # prevent double logging
|
||||
return
|
||||
|
|
@ -2399,7 +2521,31 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
e,
|
||||
)
|
||||
|
||||
async def async_success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs):
|
||||
async def async_success_handler(
|
||||
self,
|
||||
result: Any = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml)
|
||||
start_time: datetime.datetime | None = None,
|
||||
end_time: datetime.datetime | None = None,
|
||||
cache_hit: bool | None = None,
|
||||
**kwargs: Any, # kwargs-ok: forwarded to _async_success_handler_body
|
||||
) -> None:
|
||||
"""Restores trace_id/session_id contextvars once this attempt's own success
|
||||
logging (including any nested calls its callbacks trigger) is fully done."""
|
||||
try:
|
||||
return await self._async_success_handler_body(
|
||||
result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs
|
||||
)
|
||||
finally:
|
||||
self._restore_correlation_context()
|
||||
|
||||
async def _async_success_handler_body(
|
||||
self,
|
||||
result: Any = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml)
|
||||
start_time: datetime.datetime | None = None,
|
||||
end_time: datetime.datetime | None = None,
|
||||
cache_hit: bool | None = None,
|
||||
**kwargs: Any, # kwargs-ok: forwarded from async_success_handler
|
||||
) -> None:
|
||||
"""
|
||||
Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions.
|
||||
"""
|
||||
|
|
@ -2791,7 +2937,32 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
kwargs=self.model_call_details,
|
||||
)
|
||||
|
||||
def failure_handler(self, exception, traceback_exception, start_time=None, end_time=None):
|
||||
def failure_handler(
|
||||
self,
|
||||
exception: Exception,
|
||||
traceback_exception: str,
|
||||
start_time: datetime.datetime | None = None,
|
||||
end_time: datetime.datetime | None = None,
|
||||
) -> None:
|
||||
"""Restores trace_id/session_id contextvars once this attempt's own failure
|
||||
logging (including any nested calls its callbacks trigger) is fully done."""
|
||||
try:
|
||||
return self._failure_handler_body(
|
||||
exception=exception,
|
||||
traceback_exception=traceback_exception,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
finally:
|
||||
self._restore_correlation_context()
|
||||
|
||||
def _failure_handler_body(
|
||||
self,
|
||||
exception: Exception,
|
||||
traceback_exception: str,
|
||||
start_time: datetime.datetime | None = None,
|
||||
end_time: datetime.datetime | None = None,
|
||||
) -> None:
|
||||
verbose_logger.debug("Logging Details LiteLLM-Failure Call: %s", litellm.failure_callback)
|
||||
if not self.should_run_logging(event_type="sync_failure"): # prevent double logging
|
||||
return
|
||||
|
|
@ -2960,7 +3131,32 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while failure logging %s", e
|
||||
)
|
||||
|
||||
async def async_failure_handler(self, exception, traceback_exception, start_time=None, end_time=None):
|
||||
async def async_failure_handler(
|
||||
self,
|
||||
exception: Exception,
|
||||
traceback_exception: str,
|
||||
start_time: datetime.datetime | None = None,
|
||||
end_time: datetime.datetime | None = None,
|
||||
) -> None:
|
||||
"""Restores trace_id/session_id contextvars once this attempt's own failure
|
||||
logging (including any nested calls its callbacks trigger) is fully done."""
|
||||
try:
|
||||
return await self._async_failure_handler_body(
|
||||
exception=exception,
|
||||
traceback_exception=traceback_exception,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
finally:
|
||||
self._restore_correlation_context()
|
||||
|
||||
async def _async_failure_handler_body(
|
||||
self,
|
||||
exception: Exception,
|
||||
traceback_exception: str,
|
||||
start_time: datetime.datetime | None = None,
|
||||
end_time: datetime.datetime | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions.
|
||||
"""
|
||||
|
|
@ -3266,7 +3462,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
model=self.model,
|
||||
messages=[],
|
||||
logging_obj=self,
|
||||
optional_params={},
|
||||
optional_params=self.optional_params or {},
|
||||
api_key="",
|
||||
request_data={},
|
||||
encoding=litellm.encoding,
|
||||
|
|
@ -3287,6 +3483,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
),
|
||||
model_response=litellm.ModelResponse(),
|
||||
json_mode=None,
|
||||
speed=self.optional_params.get("speed") if self.optional_params else None,
|
||||
)
|
||||
return result
|
||||
|
||||
|
|
@ -4043,7 +4240,7 @@ def _init_custom_logger_compatible_class(
|
|||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, PagerDutyAlerting):
|
||||
return callback
|
||||
pagerduty_logger: Final = PagerDutyAlerting(**custom_logger_init_args)
|
||||
pagerduty_logger: Final = _PAGERDUTY_ALERTING_FACTORY(**custom_logger_init_args)
|
||||
_in_memory_loggers.append(pagerduty_logger)
|
||||
return pagerduty_logger
|
||||
elif logging_integration == "anthropic_cache_control_hook":
|
||||
|
|
@ -4073,7 +4270,7 @@ def _init_custom_logger_compatible_class(
|
|||
return _gcs_pubsub_logger
|
||||
elif logging_integration == "generic_api":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, GenericAPILogger):
|
||||
if isinstance(callback, _GENERIC_API_LOGGER_CLS):
|
||||
return callback
|
||||
generic_api_logger: Final = GenericAPILogger()
|
||||
_in_memory_loggers.append(generic_api_logger)
|
||||
|
|
@ -4082,21 +4279,21 @@ def _init_custom_logger_compatible_class(
|
|||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, ResendEmailLogger):
|
||||
return callback
|
||||
resend_email_logger: Final = ResendEmailLogger()
|
||||
resend_email_logger: Final = _RESEND_EMAIL_LOGGER_FACTORY()
|
||||
_in_memory_loggers.append(resend_email_logger)
|
||||
return resend_email_logger
|
||||
elif logging_integration == "sendgrid_email":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, SendGridEmailLogger):
|
||||
return callback
|
||||
sendgrid_email_logger: Final = SendGridEmailLogger()
|
||||
sendgrid_email_logger: Final = _SENDGRID_EMAIL_LOGGER_FACTORY()
|
||||
_in_memory_loggers.append(sendgrid_email_logger)
|
||||
return sendgrid_email_logger
|
||||
elif logging_integration == "smtp_email":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, SMTPEmailLogger):
|
||||
return callback
|
||||
smtp_email_logger: Final = SMTPEmailLogger()
|
||||
smtp_email_logger: Final = _SMTP_EMAIL_LOGGER_FACTORY()
|
||||
_in_memory_loggers.append(smtp_email_logger)
|
||||
return smtp_email_logger
|
||||
elif logging_integration == "humanloop":
|
||||
|
|
@ -4163,7 +4360,7 @@ def _init_custom_logger_compatible_class(
|
|||
return None
|
||||
|
||||
|
||||
def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list) -> Any | None:
|
||||
def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[CustomLogger]) -> "OpenTelemetryV2 | None":
|
||||
"""If ``LITELLM_OTEL_V2`` is on, build (or reuse) a single ``OpenTelemetryV2``
|
||||
instance configured via the preset for ``callback_name``.
|
||||
|
||||
|
|
@ -4194,7 +4391,7 @@ def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list) -> An
|
|||
return v2_logger
|
||||
|
||||
|
||||
def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list) -> None:
|
||||
def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list[CustomLogger]) -> None:
|
||||
"""
|
||||
Auto-initialize ArizePhoenixLogger when Phoenix env vars are detected.
|
||||
|
||||
|
|
@ -4421,7 +4618,7 @@ def get_custom_logger_compatible_class(
|
|||
return callback
|
||||
elif logging_integration == "generic_api":
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, GenericAPILogger):
|
||||
if isinstance(callback, _GENERIC_API_LOGGER_CLS):
|
||||
return callback
|
||||
elif logging_integration == "resend_email":
|
||||
for callback in _in_memory_loggers:
|
||||
|
|
@ -5061,33 +5258,61 @@ class StandardLoggingPayloadSetup:
|
|||
return end_time_float - start_time_float
|
||||
|
||||
@staticmethod
|
||||
def _get_standard_logging_payload_trace_id(
|
||||
def get_standard_logging_payload_trace_id(
|
||||
logging_obj: Logging,
|
||||
litellm_params: dict,
|
||||
litellm_params: Mapping[str, Any],
|
||||
) -> str:
|
||||
"""
|
||||
Returns the `litellm_trace_id` for this request
|
||||
|
||||
This helps link sessions when multiple requests are made in a single session
|
||||
|
||||
Gated behind `litellm.request_correlation_in_logs`:
|
||||
- Off (default): legacy behavior, preserved for backward compatibility -
|
||||
`litellm_session_id` takes priority over `litellm_trace_id` since historically
|
||||
this field doubled as the session-grouping field.
|
||||
- On: `litellm_trace_id` takes priority - trace_id and session_id are independent,
|
||||
see `get_standard_logging_payload_session_id` for session tracking.
|
||||
"""
|
||||
dynamic_litellm_session_id: Final = litellm_params.get("litellm_session_id")
|
||||
dynamic_litellm_trace_id: Final = litellm_params.get("litellm_trace_id")
|
||||
metadata: Final = litellm_params.get("metadata")
|
||||
metadata_session_id: Final = metadata.get("session_id") if metadata else None
|
||||
metadata_trace_id: Final = metadata.get("trace_id") if metadata else None
|
||||
|
||||
# Note: we recommend using `litellm_session_id` for session tracking
|
||||
# `litellm_trace_id` is an internal litellm param
|
||||
ordered_candidates: Final[tuple[Any, Any, Any, Any]] = (
|
||||
(dynamic_litellm_trace_id, dynamic_litellm_session_id, metadata_trace_id, metadata_session_id)
|
||||
if litellm.request_correlation_in_logs
|
||||
else (dynamic_litellm_session_id, dynamic_litellm_trace_id, metadata_session_id, metadata_trace_id)
|
||||
)
|
||||
for candidate in ordered_candidates:
|
||||
if candidate:
|
||||
return str(candidate)
|
||||
return logging_obj.litellm_trace_id
|
||||
|
||||
@staticmethod
|
||||
def get_standard_logging_payload_session_id(
|
||||
logging_obj: Logging,
|
||||
litellm_params: Mapping[str, Any],
|
||||
) -> str:
|
||||
"""
|
||||
Returns the end-user/conversation `litellm_session_id` for this request, independent of trace_id.
|
||||
|
||||
Only populated when `litellm.request_correlation_in_logs` is enabled - off by default
|
||||
to avoid changing existing StandardLoggingPayload shape for callers who haven't opted in.
|
||||
Unlike `get_standard_logging_payload_trace_id`, this never falls back to a generated
|
||||
per-call trace id: it's empty when the caller never supplied a session id.
|
||||
"""
|
||||
if not litellm.request_correlation_in_logs:
|
||||
return ""
|
||||
dynamic_litellm_session_id: Final = litellm_params.get("litellm_session_id")
|
||||
if dynamic_litellm_session_id:
|
||||
return str(dynamic_litellm_session_id)
|
||||
elif dynamic_litellm_trace_id:
|
||||
return str(dynamic_litellm_trace_id)
|
||||
# Fallback: use metadata.session_id or metadata.trace_id for call chaining
|
||||
metadata: Final = litellm_params.get("metadata") or {}
|
||||
metadata_session_id: Final = metadata.get("session_id")
|
||||
metadata_trace_id: Final = metadata.get("trace_id")
|
||||
metadata: Final = litellm_params.get("metadata")
|
||||
metadata_session_id: Final = metadata.get("session_id") if metadata else None
|
||||
if metadata_session_id:
|
||||
return str(metadata_session_id)
|
||||
if metadata_trace_id:
|
||||
return str(metadata_trace_id)
|
||||
return logging_obj.litellm_trace_id
|
||||
return logging_obj.litellm_session_id
|
||||
|
||||
@staticmethod
|
||||
def _get_user_agent_tags(proxy_server_request: dict) -> list[str] | None:
|
||||
|
|
@ -5392,7 +5617,11 @@ def get_standard_logging_object_payload(
|
|||
payload: Final[StandardLoggingPayload] = StandardLoggingPayload(
|
||||
id=str(id),
|
||||
litellm_call_id=kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"),
|
||||
trace_id=StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id(
|
||||
trace_id=StandardLoggingPayloadSetup.get_standard_logging_payload_trace_id(
|
||||
logging_obj=logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
),
|
||||
session_id=StandardLoggingPayloadSetup.get_standard_logging_payload_session_id(
|
||||
logging_obj=logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
),
|
||||
|
|
|
|||
|
|
@ -2,6 +2,7 @@
|
|||
Helper utilities for tracking the cost of built-in tools.
|
||||
"""
|
||||
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
import litellm
|
||||
|
|
@ -23,6 +24,14 @@ from litellm.types.utils import (
|
|||
)
|
||||
|
||||
|
||||
def _usage_reports_server_side_web_search_calls(usage: Usage) -> bool:
|
||||
details: Final = getattr(usage, "server_side_tool_usage_details", None)
|
||||
if not isinstance(details, Mapping):
|
||||
return False
|
||||
calls: Final = details.get("web_search_calls")
|
||||
return isinstance(calls, int) and calls > 0
|
||||
|
||||
|
||||
class StandardBuiltInToolCostTracking:
|
||||
"""
|
||||
Helper class for tracking the cost of built-in tools
|
||||
|
|
@ -351,6 +360,10 @@ class StandardBuiltInToolCostTracking:
|
|||
# and _handle_web_search_cost() is never called.
|
||||
if hasattr(usage, "server_tool_use") and _get_web_search_requests(usage.server_tool_use) is not None:
|
||||
return True
|
||||
# xAI reports usage.server_side_tool_usage_details.web_search_calls; a searched
|
||||
# answer with no url_citation annotations has no other chat-path signal
|
||||
if _usage_reports_server_side_web_search_calls(usage):
|
||||
return True
|
||||
return False
|
||||
elif isinstance(response_object, ResponsesAPIResponse):
|
||||
# response api explicitly includes web_search_call in the output
|
||||
|
|
@ -370,6 +383,8 @@ class StandardBuiltInToolCostTracking:
|
|||
)
|
||||
):
|
||||
return True
|
||||
if _usage_reports_server_side_web_search_calls(usage):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
|
|
@ -432,7 +447,9 @@ class StandardBuiltInToolCostTracking:
|
|||
"""
|
||||
output: Final = response_object.output
|
||||
for output_item in output:
|
||||
_output_type: str | None = getattr(output_item, "type", None)
|
||||
_output_type: str | None = (
|
||||
output_item.get("type") if isinstance(output_item, dict) else getattr(output_item, "type", None)
|
||||
)
|
||||
if _output_type == output_type:
|
||||
return True
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -49,6 +49,12 @@ _SERVICE_TIER_TO_COST_KEY_SUFFIX: Final[Mapping[str, str]] = MappingProxyType(
|
|||
}
|
||||
)
|
||||
|
||||
_INCLUSIVE_THRESHOLD_PROVIDERS: Final = frozenset({"xai"})
|
||||
|
||||
|
||||
def _uses_inclusive_token_thresholds(custom_llm_provider: str | None) -> bool:
|
||||
return custom_llm_provider in _INCLUSIVE_THRESHOLD_PROVIDERS
|
||||
|
||||
|
||||
def _get_token_detail_value(details: object, key: str) -> int | None:
|
||||
if isinstance(details, dict):
|
||||
|
|
@ -202,7 +208,11 @@ def _parse_above_token_threshold(key: str) -> float:
|
|||
|
||||
|
||||
def _get_token_base_cost(
|
||||
model_info: ModelInfo, usage: Usage, service_tier: str | None = None
|
||||
model_info: ModelInfo,
|
||||
usage: Usage,
|
||||
service_tier: str | None = None,
|
||||
*,
|
||||
threshold_is_inclusive: bool = False,
|
||||
) -> tuple[float, float, float, float, float]:
|
||||
"""
|
||||
Return prompt cost, completion cost, and cache costs for a given model and usage.
|
||||
|
|
@ -210,6 +220,9 @@ def _get_token_base_cost(
|
|||
If input_tokens > threshold and `input_cost_per_token_above_[x]k_tokens` or `input_cost_per_token_above_[x]_tokens` is set,
|
||||
then we use the corresponding threshold cost for all token types.
|
||||
|
||||
`threshold_is_inclusive` switches that comparison to >=, for providers such as xAI
|
||||
that bill the higher tier once the prompt reaches the threshold.
|
||||
|
||||
Returns:
|
||||
Tuple[float, float, float, float] - (prompt_cost, completion_cost, cache_creation_cost, cache_read_cost)
|
||||
"""
|
||||
|
|
@ -262,7 +275,7 @@ def _get_token_base_cost(
|
|||
# Handle both formats: _above_128k_tokens and _above_128_tokens
|
||||
threshold_str = key.split("_above_")[1].split("_tokens")[0]
|
||||
threshold = _parse_above_token_threshold(key)
|
||||
if usage.prompt_tokens > threshold:
|
||||
if usage.prompt_tokens > threshold or (threshold_is_inclusive and usage.prompt_tokens == threshold):
|
||||
# Prefer a service_tier-specific above-threshold key when available,
|
||||
# e.g. input_cost_per_token_priority_above_200k_tokens for Gemini
|
||||
# ON_DEMAND_PRIORITY. Falls back to the standard key automatically
|
||||
|
|
@ -777,7 +790,12 @@ def generic_cost_per_token(
|
|||
cache_creation_cost,
|
||||
cache_creation_cost_above_1hr,
|
||||
cache_read_cost,
|
||||
) = _get_token_base_cost(model_info=model_info, usage=usage, service_tier=service_tier)
|
||||
) = _get_token_base_cost(
|
||||
model_info=model_info,
|
||||
usage=usage,
|
||||
service_tier=service_tier,
|
||||
threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider),
|
||||
)
|
||||
|
||||
prompt_cost = _calculate_input_cost(
|
||||
prompt_tokens_details=prompt_tokens_details,
|
||||
|
|
@ -909,7 +927,12 @@ def get_token_type_cost_breakdown(
|
|||
cache_creation_cost_rate,
|
||||
cache_creation_cost_above_1hr_rate,
|
||||
cache_read_cost_rate,
|
||||
) = _get_token_base_cost(model_info=model_info, usage=usage, service_tier=service_tier)
|
||||
) = _get_token_base_cost(
|
||||
model_info=model_info,
|
||||
usage=usage,
|
||||
service_tier=service_tier,
|
||||
threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider),
|
||||
)
|
||||
|
||||
reasoning_tokens = (
|
||||
_parse_completion_tokens_details(usage)["reasoning_tokens"]
|
||||
|
|
@ -996,9 +1019,13 @@ def calculate_image_response_cost_from_usage(
|
|||
input_tokens_details: Final = getattr(usage, "input_tokens_details", None)
|
||||
prompt_tokens_details: PromptTokensDetailsWrapper | None = None
|
||||
if input_tokens_details is not None:
|
||||
# input_tokens_details may be a dict (e.g. OpenAI image edit responses)
|
||||
# or an object; read it tolerantly like the output side below, so image
|
||||
# input tokens are priced at input_cost_per_image_token instead of
|
||||
# silently falling back to the text rate.
|
||||
prompt_tokens_details = PromptTokensDetailsWrapper(
|
||||
text_tokens=getattr(input_tokens_details, "text_tokens", None),
|
||||
image_tokens=getattr(input_tokens_details, "image_tokens", None),
|
||||
text_tokens=_get_token_detail_value(input_tokens_details, "text_tokens"),
|
||||
image_tokens=_get_token_detail_value(input_tokens_details, "image_tokens"),
|
||||
cached_tokens=0,
|
||||
)
|
||||
|
||||
|
|
|
|||
87
litellm/litellm_core_utils/llm_judge.py
Normal file
87
litellm/litellm_core_utils/llm_judge.py
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
"""Shared primitives for LLM-judge features (llm_as_a_judge guardrail, shadow eval)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import litellm
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm import Router
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
JSON_FENCE_RE: Final = re.compile(r"```(?:json)?\s*(.*?)\s*```", re.DOTALL | re.IGNORECASE)
|
||||
|
||||
|
||||
def default_router_provider() -> Router | None:
|
||||
try:
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
except ImportError:
|
||||
return None
|
||||
|
||||
return llm_router
|
||||
|
||||
|
||||
def parse_json_verdict(raw: str) -> dict[str, object]: # mutable-ok: plain parsed-JSON payload
|
||||
"""Parse a judge's JSON verdict, tolerating markdown fences and surrounding prose."""
|
||||
text = raw.strip() # rebind-ok: progressively narrowed to the JSON payload
|
||||
fenced: Final = JSON_FENCE_RE.search(text)
|
||||
if fenced is not None:
|
||||
text = fenced.group(1).strip() # rebind-ok: progressively narrowed to the JSON payload
|
||||
parsed: object
|
||||
try:
|
||||
parsed = json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
start: Final = text.find("{")
|
||||
end: Final = text.rfind("}")
|
||||
if start == -1 or end <= start:
|
||||
raise
|
||||
parsed = json.loads(text[start : end + 1])
|
||||
if not isinstance(parsed, dict):
|
||||
raise ValueError("judge response is not a JSON object")
|
||||
return {str(k): v for k, v in parsed.items()} # mutable-ok: plain parsed-JSON payload
|
||||
|
||||
|
||||
def extract_text_from_content(content: object) -> str:
|
||||
"""Return plain text from a message content field (str or multimodal list)."""
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
return " ".join(
|
||||
str(part.get("text", "")) for part in content if isinstance(part, dict) and part.get("type") == "text"
|
||||
)
|
||||
return ""
|
||||
|
||||
|
||||
def router_resolves_model(router: Router | None, model: str) -> bool:
|
||||
"""Whether the model name resolves through the proxy's router (configured deployment
|
||||
or model-group alias), the same check the judge dispatch itself makes, so start-time
|
||||
validation cannot accept a name the call path then fails on."""
|
||||
return router is not None and bool(model in router.model_group_alias or router.get_model_list(model_name=model))
|
||||
|
||||
|
||||
async def judge_acompletion(
|
||||
router: Router | None,
|
||||
judge_model: str,
|
||||
messages: list[AllMessageValues], # mutable-ok: the SDK acompletion signature takes a list
|
||||
**params: object,
|
||||
) -> ModelResponse:
|
||||
"""Dispatch a judge call through the proxy's router when the judge model is a
|
||||
configured deployment (DB-stored credentials work), through the SDK for
|
||||
provider-qualified public names. The router path never retries or falls back:
|
||||
a failed judge call is the caller's counted failure, not a spend multiplier.
|
||||
Sampling preferences are advisory: models that removed sampling params (e.g.
|
||||
claude-sonnet-5) drop them instead of rejecting the judge call."""
|
||||
if router_resolves_model(router, judge_model):
|
||||
return await router.acompletion( # pyright: ignore[reportOptionalMemberAccess] # router_resolves_model implies router is not None
|
||||
model=judge_model,
|
||||
messages=messages,
|
||||
num_retries=0,
|
||||
fallbacks=[],
|
||||
drop_params=True,
|
||||
**params,
|
||||
)
|
||||
return await litellm.acompletion(model=judge_model, messages=messages, num_retries=0, drop_params=True, **params)
|
||||
|
|
@ -6,7 +6,7 @@ import io
|
|||
import json
|
||||
import mimetypes
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from os import PathLike
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, cast
|
||||
|
|
@ -1775,3 +1775,24 @@ def split_concatenated_json_objects(raw: str) -> list[dict[str, Any]]:
|
|||
idx = end_idx
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def text_completion_prompt_to_messages(prompt: object) -> tuple[AllMessageValues, ...]:
|
||||
"""
|
||||
Wrap an OpenAI ``/v1/completions`` ``prompt`` into Chat Completion messages.
|
||||
|
||||
Mirrors what ``litellm.text_completion`` does on the real-time path: a
|
||||
string becomes a single user message, and a list of strings becomes one
|
||||
user message per element. Pre-tokenized prompts (``list[int]`` /
|
||||
``list[list[int]]``) are only meaningful for the OpenAI-family text
|
||||
endpoints, so they are rejected here rather than silently forwarded, as is
|
||||
an empty prompt, which every chat-shaped provider rejects downstream.
|
||||
"""
|
||||
prompt_type_name: Final = type(prompt).__name__
|
||||
if isinstance(prompt, str) and prompt:
|
||||
return (ChatCompletionUserMessage(role="user", content=prompt),)
|
||||
entries: Final = cast("Sequence[object]", prompt) if isinstance(prompt, Sequence) else ()
|
||||
string_entries: Final = tuple(entry for entry in entries if isinstance(entry, str) and entry)
|
||||
if string_entries and len(string_entries) == len(entries):
|
||||
return tuple(ChatCompletionUserMessage(role="user", content=entry) for entry in string_entries)
|
||||
raise ValueError(f"`prompt` must be a non-empty string or a non-empty list of strings. Got: {prompt_type_name}.")
|
||||
|
|
|
|||
|
|
@ -549,7 +549,7 @@ def _fetch_and_extract_template(
|
|||
return chat_template, bos_token, eos_token
|
||||
|
||||
|
||||
async def ahf_chat_template(model: str, messages: list, chat_template: Any | None = None):
|
||||
async def ahf_chat_template(model: str, messages: list, chat_template: str | None = None):
|
||||
"""HuggingFace chat template (async version)"""
|
||||
from litellm.litellm_core_utils.prompt_templates.huggingface_template_handler import (
|
||||
_aget_chat_template_file,
|
||||
|
|
@ -576,7 +576,7 @@ async def ahf_chat_template(model: str, messages: list, chat_template: Any | Non
|
|||
)
|
||||
|
||||
|
||||
def hf_chat_template(model: str, messages: list, chat_template: Any | None = None):
|
||||
def hf_chat_template(model: str, messages: list, chat_template: str | None = None):
|
||||
"""HuggingFace chat template (sync version)"""
|
||||
from litellm.litellm_core_utils.prompt_templates.huggingface_template_handler import (
|
||||
_get_chat_template_file,
|
||||
|
|
@ -1130,7 +1130,7 @@ def convert_to_azure_openai_messages(
|
|||
|
||||
|
||||
def infer_protocol_value(
|
||||
value: Any,
|
||||
value: object,
|
||||
) -> Literal[
|
||||
"string_value",
|
||||
"number_value",
|
||||
|
|
@ -1702,7 +1702,9 @@ def convert_function_to_anthropic_tool_invoke(
|
|||
_name: Final = get_attribute_or_key(function_call, "name") or ""
|
||||
_arguments: Final = get_attribute_or_key(function_call, "arguments")
|
||||
|
||||
tool_input = parse_tool_call_arguments(_arguments, tool_name=_name, context="Anthropic function to tool invoke")
|
||||
tool_input: Final = parse_tool_call_arguments(
|
||||
_arguments, tool_name=_name, context="Anthropic function to tool invoke"
|
||||
)
|
||||
|
||||
anthropic_tool_invoke: Final = [
|
||||
AnthropicMessagesToolUseParam(
|
||||
|
|
@ -1764,7 +1766,7 @@ def convert_to_anthropic_tool_invoke(
|
|||
|
||||
Fixes: https://github.com/BerriAI/litellm/issues/17737
|
||||
"""
|
||||
anthropic_tool_invoke: Final[list[AnthropicMessagesToolUseParam | dict[str, Any]]] = []
|
||||
anthropic_tool_invoke: Final[list[AnthropicMessagesToolUseParam | dict[str, object]]] = []
|
||||
|
||||
for tool in tool_calls:
|
||||
if not get_attribute_or_key(tool, "type") == "function":
|
||||
|
|
@ -1785,7 +1787,7 @@ def convert_to_anthropic_tool_invoke(
|
|||
# Server tool IDs start with "srvtoolu_"
|
||||
if tool_id.startswith("srvtoolu_"):
|
||||
# Create server_tool_use block instead of tool_use
|
||||
_anthropic_server_tool_use: dict[str, Any] = {
|
||||
_anthropic_server_tool_use: dict[str, object] = {
|
||||
"type": "server_tool_use",
|
||||
"id": tool_id,
|
||||
"name": tool_name,
|
||||
|
|
@ -2177,7 +2179,7 @@ def _is_orphaned_tool_result(
|
|||
return False
|
||||
|
||||
|
||||
def _declared_tool_call_ids(message: Mapping[str, Any]) -> frozenset[str]:
|
||||
def _declared_tool_call_ids(message: Mapping[str, object]) -> frozenset[str]:
|
||||
tool_calls: Final = message.get("tool_calls")
|
||||
if not isinstance(tool_calls, list):
|
||||
return frozenset()
|
||||
|
|
@ -2186,7 +2188,7 @@ def _declared_tool_call_ids(message: Mapping[str, Any]) -> frozenset[str]:
|
|||
)
|
||||
|
||||
|
||||
def group_tool_exchanges(messages: Sequence[Mapping[str, Any]]) -> tuple[tuple[int, ...], ...]:
|
||||
def group_tool_exchanges(messages: Sequence[Mapping[str, object]]) -> tuple[tuple[int, ...], ...]:
|
||||
"""Group message indices into tool exchanges: an assistant row that made
|
||||
tool calls, together with the tool rows answering the ids it declared.
|
||||
|
||||
|
|
@ -2204,7 +2206,7 @@ def group_tool_exchanges(messages: Sequence[Mapping[str, Any]]) -> tuple[tuple[i
|
|||
return tuple(_iter_tool_exchange_groups(messages))
|
||||
|
||||
|
||||
def _iter_tool_exchange_groups(messages: Sequence[Mapping[str, Any]]) -> Iterator[tuple[int, ...]]:
|
||||
def _iter_tool_exchange_groups(messages: Sequence[Mapping[str, object]]) -> Iterator[tuple[int, ...]]:
|
||||
index = 0
|
||||
while index < len(messages):
|
||||
declared = _declared_tool_call_ids(messages[index])
|
||||
|
|
@ -2409,7 +2411,7 @@ def anthropic_messages_pt(
|
|||
# Convert ChatCompletionImageUrlObject to dict if needed
|
||||
image_url_value = m["image_url"]
|
||||
if isinstance(image_url_value, str):
|
||||
image_url_input: str | dict[str, Any] = image_url_value
|
||||
image_url_input: str | dict[str, object] = image_url_value
|
||||
else:
|
||||
# ChatCompletionImageUrlObject or dict case - convert to dict
|
||||
image_url_input = {
|
||||
|
|
@ -3179,7 +3181,7 @@ def _load_image_from_url(image_url):
|
|||
try:
|
||||
# Send a GET request to the image URL
|
||||
client: Final = HTTPHandler(concurrent_limit=1)
|
||||
response: Final = safe_get(client, image_url)
|
||||
response: Final[httpx.Response] = safe_get(client, image_url)
|
||||
response.raise_for_status() # Raise an exception for HTTP errors
|
||||
|
||||
# Check the response's content type to ensure it is an image
|
||||
|
|
@ -3382,7 +3384,7 @@ class BedrockImageProcessor:
|
|||
@staticmethod
|
||||
def _post_call_image_processing(response: httpx.Response, image_url: str = "") -> tuple[str, str]:
|
||||
# Check the response's content type to ensure it is an image
|
||||
content_type = response.headers.get("content-type")
|
||||
content_type: str | None = response.headers.get("content-type")
|
||||
|
||||
# Use helper function to infer content type with fallback logic
|
||||
content_type = infer_content_type_from_url_and_content(
|
||||
|
|
@ -3406,7 +3408,7 @@ class BedrockImageProcessor:
|
|||
params={"concurrent_limit": 1},
|
||||
)
|
||||
# Send a GET request to the image URL
|
||||
response: Final = await async_safe_get(client, image_url)
|
||||
response: Final[httpx.Response] = await async_safe_get(client, image_url)
|
||||
response.raise_for_status() # Raise an exception for HTTP errors
|
||||
|
||||
return BedrockImageProcessor._post_call_image_processing(response, image_url)
|
||||
|
|
@ -3419,7 +3421,7 @@ class BedrockImageProcessor:
|
|||
try:
|
||||
client: Final = HTTPHandler(concurrent_limit=1)
|
||||
# Send a GET request to the image URL
|
||||
response: Final = safe_get(client, image_url)
|
||||
response: Final[httpx.Response] = safe_get(client, image_url)
|
||||
response.raise_for_status() # Raise an exception for HTTP errors
|
||||
|
||||
return BedrockImageProcessor._post_call_image_processing(response, image_url)
|
||||
|
|
@ -3967,6 +3969,36 @@ def _rename_duplicate_bedrock_document_names(
|
|||
return contents
|
||||
|
||||
|
||||
BEDROCK_DOCUMENT_PLACEHOLDER_TEXT: Final = "."
|
||||
|
||||
|
||||
def _with_text_when_document_only(message: BedrockMessageBlock) -> BedrockMessageBlock:
|
||||
blocks: Final = message["content"]
|
||||
needs_text: Final = (
|
||||
message["role"] == "user"
|
||||
and any("document" in block for block in blocks)
|
||||
and all("text" not in block for block in blocks)
|
||||
)
|
||||
if not needs_text:
|
||||
return message
|
||||
placeholder: Final = BedrockContentBlock(text=BEDROCK_DOCUMENT_PLACEHOLDER_TEXT)
|
||||
cut: Final = len(blocks) - 1 if "cachePoint" in blocks[-1] else len(blocks)
|
||||
return BedrockMessageBlock(role="user", content=[*blocks[:cut], placeholder, *blocks[cut:]])
|
||||
|
||||
|
||||
def _ensure_document_messages_have_text(
|
||||
contents: list[BedrockMessageBlock],
|
||||
) -> list[BedrockMessageBlock]:
|
||||
"""
|
||||
Bedrock Converse rejects any user message that carries a document block
|
||||
without a sibling text block ("A text block must be included when using
|
||||
documents"), e.g. Claude Code sends the PDF as a document-only user turn.
|
||||
Inject a placeholder text block, kept ahead of a trailing cachePoint so
|
||||
the caller's cache boundary stays the final block.
|
||||
"""
|
||||
return [_with_text_when_document_only(message) for message in contents]
|
||||
|
||||
|
||||
def _sort_bedrock_assistant_content_blocks(
|
||||
blocks: list[BedrockContentBlock],
|
||||
) -> list[BedrockContentBlock]:
|
||||
|
|
@ -4535,7 +4567,7 @@ class BedrockConverseMessagesProcessor:
|
|||
llm_provider=llm_provider,
|
||||
)
|
||||
|
||||
return _rename_duplicate_bedrock_document_names(contents)
|
||||
return _ensure_document_messages_have_text(_rename_duplicate_bedrock_document_names(contents))
|
||||
|
||||
@staticmethod
|
||||
def translate_thinking_blocks_to_reasoning_content_blocks(
|
||||
|
|
@ -4911,7 +4943,7 @@ def _bedrock_converse_messages_pt(
|
|||
llm_provider=llm_provider,
|
||||
)
|
||||
|
||||
return _rename_duplicate_bedrock_document_names(contents)
|
||||
return _ensure_document_messages_have_text(_rename_duplicate_bedrock_document_names(contents))
|
||||
|
||||
|
||||
def make_valid_bedrock_tool_name(input_tool_name: str) -> str:
|
||||
|
|
@ -5328,10 +5360,10 @@ def get_attribute_or_key(tool_or_function, attribute, default=None):
|
|||
class NormalizedToolCall(TypedDict):
|
||||
id: str | None
|
||||
name: str | None
|
||||
arguments: dict[str, Any]
|
||||
arguments: dict[str, object]
|
||||
|
||||
|
||||
def _parse_tool_call_arguments(raw: Any, tool_name: str | None, context: str) -> dict[str, Any]:
|
||||
def _parse_tool_call_arguments(raw: Any, tool_name: str | None, context: str) -> dict[str, object]:
|
||||
# Anthropic's tool_use blocks already carry a parsed dict in "input";
|
||||
# chat completions and the Responses API carry a JSON string that may be
|
||||
# truncated by the model, so route those through the repair-aware parser.
|
||||
|
|
@ -5352,12 +5384,12 @@ def _parse_tool_call_arguments(raw: Any, tool_name: str | None, context: str) ->
|
|||
|
||||
|
||||
def _tool_calls_from_chat_completion_response(
|
||||
response: Any, include_all_choices: bool = False
|
||||
response: object, include_all_choices: bool = False
|
||||
) -> list[NormalizedToolCall]:
|
||||
choices: Final = get_attribute_or_key(response, "choices", None)
|
||||
if not (isinstance(choices, list) and choices):
|
||||
return []
|
||||
tool_calls: Final[list[Any]] = []
|
||||
tool_calls: Final[list[object]] = []
|
||||
for choice in choices if include_all_choices else choices[:1]:
|
||||
message = get_attribute_or_key(choice, "message", None)
|
||||
choice_tool_calls = get_attribute_or_key(message, "tool_calls", None) if message else None
|
||||
|
|
@ -5383,7 +5415,7 @@ def _tool_calls_from_chat_completion_response(
|
|||
return result
|
||||
|
||||
|
||||
def _tool_calls_from_responses_api_response(response: Any) -> list[NormalizedToolCall]:
|
||||
def _tool_calls_from_responses_api_response(response: object) -> list[NormalizedToolCall]:
|
||||
output: Final = get_attribute_or_key(response, "output", None)
|
||||
if not isinstance(output, list):
|
||||
return []
|
||||
|
|
@ -5406,7 +5438,7 @@ def _tool_calls_from_responses_api_response(response: Any) -> list[NormalizedToo
|
|||
return result
|
||||
|
||||
|
||||
def _tool_calls_from_anthropic_messages_response(response: Any) -> list[NormalizedToolCall]:
|
||||
def _tool_calls_from_anthropic_messages_response(response: object) -> list[NormalizedToolCall]:
|
||||
content: Final = get_attribute_or_key(response, "content", None)
|
||||
if not isinstance(content, list):
|
||||
return []
|
||||
|
|
@ -5425,7 +5457,7 @@ def _tool_calls_from_anthropic_messages_response(response: Any) -> list[Normaliz
|
|||
return result
|
||||
|
||||
|
||||
def get_tool_calls_from_response(response: Any, include_all_choices: bool = False) -> list[NormalizedToolCall]:
|
||||
def get_tool_calls_from_response(response: object, include_all_choices: bool = False) -> list[NormalizedToolCall]:
|
||||
"""
|
||||
Extract tool/function calls from a response object into a normalized
|
||||
``{"id", "name", "arguments"}`` shape, regardless of which API surface
|
||||
|
|
@ -5456,7 +5488,7 @@ def get_tool_calls_from_response(response: Any, include_all_choices: bool = Fals
|
|||
return []
|
||||
|
||||
|
||||
def has_tool_with_name(tools: Any, tool_name: str) -> bool:
|
||||
def has_tool_with_name(tools: object, tool_name: str) -> bool:
|
||||
"""
|
||||
Check whether a tools list (as sent to an LLM) includes a tool with the
|
||||
given name, regardless of shape: OpenAI-style function tools
|
||||
|
|
@ -5482,9 +5514,9 @@ def has_tool_with_name(tools: Any, tool_name: str) -> bool:
|
|||
|
||||
|
||||
def resolve_structured_messages(
|
||||
messages: list[dict[str, Any]] | None,
|
||||
messages: list[dict[str, object]] | None,
|
||||
request_kwargs: dict[str, Any],
|
||||
) -> list[dict[str, Any]] | None:
|
||||
) -> list[dict[str, object]] | None:
|
||||
"""
|
||||
Normalize a request's messages to OpenAI-spec chat-completions shape,
|
||||
regardless of which API surface produced them (chat completions,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, cast
|
||||
|
||||
import litellm
|
||||
|
|
@ -20,11 +21,24 @@ from .litellm_logging import Logging as LiteLLMLogging
|
|||
if TYPE_CHECKING:
|
||||
from websockets.asyncio.client import ClientConnection
|
||||
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
CLIENT_CONNECTION_CLASS = ClientConnection
|
||||
else:
|
||||
CLIENT_CONNECTION_CLASS = Any
|
||||
|
||||
|
||||
class _ClientWebSocketExceptions(Protocol):
|
||||
ConnectionClosed: type[Exception]
|
||||
|
||||
|
||||
class _ClientWebSocket(Protocol):
|
||||
exceptions: _ClientWebSocketExceptions
|
||||
|
||||
async def send_text(self, data: str) -> None: ...
|
||||
async def receive_text(self) -> str: ...
|
||||
|
||||
|
||||
class RealtimeEventNormalizer(Protocol):
|
||||
def should_drop(self, event: object) -> bool: ...
|
||||
def normalize(self, event: dict) -> dict: ...
|
||||
|
|
@ -48,13 +62,13 @@ class RealTimeStreaming:
|
|||
logging_obj: LiteLLMLogging,
|
||||
provider_config: BaseRealtimeConfig | None = None,
|
||||
model: str = "",
|
||||
user_api_key_dict: Any | None = None,
|
||||
user_api_key_dict: object | None = None,
|
||||
request_data: dict | None = None,
|
||||
backend_uses_beta_protocol: bool | None = None,
|
||||
force_transcription_model: str | None = None,
|
||||
event_normalizer: RealtimeEventNormalizer | None = None,
|
||||
):
|
||||
self.websocket = websocket
|
||||
self.websocket: _ClientWebSocket = websocket
|
||||
self.backend_ws = backend_ws
|
||||
self.logging_obj = logging_obj
|
||||
self.messages: list[OpenAIRealtimeEvents] = []
|
||||
|
|
@ -127,7 +141,7 @@ class RealTimeStreaming:
|
|||
]
|
||||
)
|
||||
_CLIENT_AUDIO_BUFFER_COMMIT_TYPES = frozenset(["input_audio_buffer.commit", "input_audio_buffer.end"])
|
||||
_AUDIO_FORMAT_MAP: dict[str, dict[str, Any]] = {
|
||||
_AUDIO_FORMAT_MAP: dict[str, dict[str, str | int]] = {
|
||||
"pcm16": {"type": "audio/pcm", "rate": 24000},
|
||||
"g711_ulaw": {"type": "audio/G711-ulaw", "rate": 8000},
|
||||
"g711_alaw": {"type": "audio/G711-alaw", "rate": 8000},
|
||||
|
|
@ -281,6 +295,7 @@ class RealTimeStreaming:
|
|||
if event_obj.get("type") != "response.done":
|
||||
return
|
||||
response: Final = cast(dict[str, Any], event_obj.get("response", {}))
|
||||
item: Mapping[str, object]
|
||||
for item in response.get("output", []):
|
||||
if item.get("type") == "function_call":
|
||||
self.tool_calls.append(
|
||||
|
|
@ -384,7 +399,7 @@ class RealTimeStreaming:
|
|||
return message
|
||||
|
||||
try:
|
||||
message_obj: Final = json.loads(message)
|
||||
message_obj: Final[Mapping[str, object]] = json.loads(message)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return message
|
||||
|
||||
|
|
@ -487,7 +502,7 @@ class RealTimeStreaming:
|
|||
if self._backend_setup_complete and not self._flushing_pending_messages_until_setup:
|
||||
return False
|
||||
try:
|
||||
msg_obj: Final = json.loads(message)
|
||||
msg_obj: Final[Mapping[str, object]] = json.loads(message)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return False
|
||||
return msg_obj.get("type") in RealTimeStreaming._CLIENT_AUDIO_BUFFER_TYPES
|
||||
|
|
@ -555,7 +570,7 @@ class RealTimeStreaming:
|
|||
def _event_to_client_json(self, event: dict) -> str:
|
||||
return json.dumps(self._normalize_event_for_ga_client(event))
|
||||
|
||||
async def _send_event_to_client(self, event: Any, event_str: str) -> bool:
|
||||
async def _send_event_to_client(self, event: object, event_str: str) -> bool:
|
||||
if self._should_drop_event_from_client(event):
|
||||
return False
|
||||
if isinstance(event, dict):
|
||||
|
|
@ -595,12 +610,12 @@ class RealTimeStreaming:
|
|||
|
||||
def _make_disable_auto_response_message(self) -> str:
|
||||
"""Return a session.update that disables VAD auto-response."""
|
||||
turn_detection: Final[dict[str, Any]] = {
|
||||
turn_detection: Final[dict[str, str | bool]] = {
|
||||
"type": "server_vad",
|
||||
"create_response": False,
|
||||
}
|
||||
if self._backend_uses_beta_protocol:
|
||||
session: dict[str, Any] = {"turn_detection": turn_detection}
|
||||
session: dict[str, object] = {"turn_detection": turn_detection}
|
||||
else:
|
||||
session = {
|
||||
"type": "realtime",
|
||||
|
|
@ -654,7 +669,7 @@ class RealTimeStreaming:
|
|||
|
||||
def _has_realtime_guardrails_for_event_hooks(
|
||||
self,
|
||||
event_hooks: list[Any],
|
||||
event_hooks: Sequence["GuardrailEventHooks"],
|
||||
) -> bool:
|
||||
"""Return True if any callback would run for one of ``event_hooks``."""
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
|
|
@ -699,7 +714,7 @@ class RealTimeStreaming:
|
|||
transcript: str,
|
||||
item_id: str | None = None,
|
||||
pre_block_backend_message: str | None = None,
|
||||
event_hooks: list[Any] | None = None,
|
||||
event_hooks: Sequence["GuardrailEventHooks"] | None = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Run registered guardrails on realtime text (transcript, user message, tool output).
|
||||
|
|
@ -753,7 +768,7 @@ class RealTimeStreaming:
|
|||
raise
|
||||
# Extract the human-readable error from the detail dict (HTTPException)
|
||||
# or fall back to str(e) for plain ValueError.
|
||||
detail = getattr(e, "detail", None)
|
||||
detail: object | None = getattr(e, "detail", None)
|
||||
if isinstance(detail, dict):
|
||||
safe_msg = detail.get("error") or str(e)
|
||||
elif detail is not None:
|
||||
|
|
@ -826,7 +841,7 @@ class RealTimeStreaming:
|
|||
return True
|
||||
return False
|
||||
|
||||
async def _handle_provider_config_message(self, raw_response) -> None:
|
||||
async def _handle_provider_config_message(self, raw_response: str) -> None:
|
||||
"""Process a backend message when a provider_config is set (transformed path)."""
|
||||
returned_object: Final = self.provider_config.transform_realtime_response(
|
||||
raw_response,
|
||||
|
|
@ -910,7 +925,7 @@ class RealTimeStreaming:
|
|||
await self._send_event_to_client(event, event_str)
|
||||
|
||||
@staticmethod
|
||||
def _parse_backend_event(raw_response: str) -> dict | None:
|
||||
def _parse_backend_event(raw_response: str) -> dict[str, object] | None:
|
||||
"""Parse a backend frame once. Returns None for non-JSON or non-object frames."""
|
||||
try:
|
||||
event: Final = json.loads(raw_response)
|
||||
|
|
@ -1020,7 +1035,7 @@ class RealTimeStreaming:
|
|||
objects and any test doubles that expose a .scope dict.
|
||||
"""
|
||||
try:
|
||||
headers: Final = websocket.scope.get("headers", [])
|
||||
headers: Final[Sequence[tuple[bytes | str, bytes | str]]] = websocket.scope.get("headers", [])
|
||||
for name, value in headers:
|
||||
if isinstance(name, bytes):
|
||||
name = name.decode("latin-1")
|
||||
|
|
@ -1071,9 +1086,9 @@ class RealTimeStreaming:
|
|||
session["output_modalities"] = ["text"]
|
||||
|
||||
# 3-7. Lift flat audio fields into the nested audio object
|
||||
audio: Final[dict[str, Any]] = {}
|
||||
inp: Final[dict[str, Any]] = {}
|
||||
out: Final[dict[str, Any]] = {}
|
||||
audio: Final[dict[str, object]] = {}
|
||||
inp: Final[dict[str, object]] = {}
|
||||
out: Final[dict[str, object]] = {}
|
||||
|
||||
# voice → audio.output.voice
|
||||
if "voice" in session:
|
||||
|
|
@ -1190,7 +1205,7 @@ class RealTimeStreaming:
|
|||
# model; check them with the same guardrail used for
|
||||
# user text so an attacker cannot smuggle blocked
|
||||
# content into a function_call_output.
|
||||
output = item.get("output", "")
|
||||
output: object = item.get("output", "")
|
||||
output_text = output if isinstance(output, str) else json.dumps(output)
|
||||
if output_text:
|
||||
# Build the sanitized function_call_output up
|
||||
|
|
@ -1241,7 +1256,7 @@ class RealTimeStreaming:
|
|||
# interaction turn.
|
||||
continue
|
||||
elif item.get("role") == "user":
|
||||
content_list = item.get("content", [])
|
||||
content_list: Sequence[object] = item.get("content", [])
|
||||
texts = [
|
||||
c.get("text", "")
|
||||
for c in content_list
|
||||
|
|
@ -1280,7 +1295,7 @@ class RealTimeStreaming:
|
|||
and not self._guardrail_turn_detection_update_sent
|
||||
and self._has_audio_transcription_guardrails()
|
||||
):
|
||||
session = msg_obj.setdefault("session", {})
|
||||
session: object = msg_obj.setdefault("session", {})
|
||||
if isinstance(session, dict):
|
||||
existing_td = session.get("turn_detection")
|
||||
if not isinstance(existing_td, dict):
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import time
|
|||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from itertools import groupby
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Union, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, TypedDict, Union, cast
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.types.llms.openai import (
|
||||
|
|
@ -30,6 +30,7 @@ from litellm.types.utils import (
|
|||
from litellm.utils import print_verbose, token_counter
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import (
|
||||
UsagePerChunk,
|
||||
)
|
||||
|
|
@ -39,6 +40,60 @@ if TYPE_CHECKING:
|
|||
)
|
||||
|
||||
|
||||
class _ThinkingBlockFragment(TypedDict, total=False):
|
||||
type: str | None
|
||||
data: str | None
|
||||
thinking: str | None
|
||||
signature: str | None
|
||||
|
||||
|
||||
class _ThinkingDelta(TypedDict, total=False):
|
||||
thinking_blocks: Sequence[_ThinkingBlockFragment]
|
||||
|
||||
|
||||
class _ThinkingChoice(TypedDict, total=False):
|
||||
delta: _ThinkingDelta
|
||||
|
||||
|
||||
class _ThinkingChunk(TypedDict):
|
||||
choices: Sequence[_ThinkingChoice]
|
||||
|
||||
|
||||
class _ContentChoice(TypedDict, total=False):
|
||||
delta: Mapping[str, str | None]
|
||||
|
||||
|
||||
class _ContentChunk(TypedDict):
|
||||
choices: Sequence[_ContentChoice]
|
||||
|
||||
|
||||
class _AudioDelta(TypedDict, total=False):
|
||||
audio: ChatCompletionAudioDelta | None
|
||||
|
||||
|
||||
class _AudioChoice(TypedDict, total=False):
|
||||
delta: _AudioDelta
|
||||
|
||||
|
||||
class _AudioChunk(TypedDict):
|
||||
choices: Sequence[_AudioChoice]
|
||||
|
||||
|
||||
class _UsageBearingChunk(TypedDict, total=False):
|
||||
usage: Usage | None
|
||||
_hidden_params: Mapping[str, str]
|
||||
|
||||
|
||||
class _UsageSummary(TypedDict):
|
||||
prompt_tokens: int | None
|
||||
completion_tokens: int | None
|
||||
cache_creation_input_tokens: int | None
|
||||
cache_read_input_tokens: int | None
|
||||
completion_tokens_details: CompletionTokensDetails | None
|
||||
prompt_tokens_details: PromptTokensDetailsWrapper | None
|
||||
cost: float | None
|
||||
|
||||
|
||||
def capture_cache_creation_token_details(
|
||||
prompt_tokens_details: PromptTokensDetailsWrapper | None,
|
||||
current: CacheCreationTokenDetails | None,
|
||||
|
|
@ -78,7 +133,7 @@ class ChunkProcessor:
|
|||
return []
|
||||
|
||||
first_chunk: Final = chunks[0]
|
||||
first_hidden_params: dict[str, Any] = {}
|
||||
first_hidden_params: dict[str, object] = {}
|
||||
if isinstance(first_chunk, dict):
|
||||
candidate = first_chunk.get("_hidden_params", {})
|
||||
if isinstance(candidate, dict):
|
||||
|
|
@ -90,7 +145,7 @@ class ChunkProcessor:
|
|||
|
||||
if first_hidden_params.get("created_at"):
|
||||
|
||||
def _created_at(chunk: Any) -> int | float:
|
||||
def _created_at(chunk: object) -> int | float:
|
||||
if isinstance(chunk, dict):
|
||||
params = chunk.get("_hidden_params", {})
|
||||
else:
|
||||
|
|
@ -103,7 +158,7 @@ class ChunkProcessor:
|
|||
return chunks
|
||||
|
||||
def update_model_response_with_hidden_params(
|
||||
self, model_response: ModelResponse, chunk: dict[str, Any] | None = None
|
||||
self, model_response: ModelResponse, chunk: Mapping[str, dict[str, object]] | None = None
|
||||
) -> ModelResponse:
|
||||
if chunk is None:
|
||||
return model_response
|
||||
|
|
@ -115,13 +170,13 @@ class ChunkProcessor:
|
|||
@staticmethod
|
||||
def apply_provider_assembled_streaming_metadata(
|
||||
response: ModelResponse,
|
||||
chunks: list[Any],
|
||||
logging_obj: Any | None = None,
|
||||
chunks: list[object],
|
||||
logging_obj: "Logging | None" = None,
|
||||
) -> None:
|
||||
if not chunks:
|
||||
return
|
||||
|
||||
model: Final = getattr(response, "model", None)
|
||||
model: Final[str | None] = getattr(response, "model", None)
|
||||
if not model:
|
||||
return
|
||||
|
||||
|
|
@ -159,7 +214,7 @@ class ChunkProcessor:
|
|||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_chunk_id(chunks: list[dict[str, Any]]) -> str:
|
||||
def _get_chunk_id(chunks: Sequence[Mapping[str, str]]) -> str:
|
||||
"""
|
||||
Chunks:
|
||||
[{"id": ""}, {"id": "1"}, {"id": "1"}]
|
||||
|
|
@ -170,7 +225,7 @@ class ChunkProcessor:
|
|||
return ""
|
||||
|
||||
@staticmethod
|
||||
def _get_model_from_chunks(chunks: list[dict[str, Any]], first_chunk_model: str) -> str:
|
||||
def _get_model_from_chunks(chunks: Sequence[Mapping[str, str]], first_chunk_model: str) -> str:
|
||||
"""
|
||||
Get the actual model from chunks, preferring a model that differs from the first chunk.
|
||||
|
||||
|
|
@ -456,7 +511,7 @@ class ChunkProcessor:
|
|||
)
|
||||
|
||||
def get_combined_content(
|
||||
self, chunks: list[dict[str, Any]], delta_key: str = "content"
|
||||
self, chunks: Sequence["_ContentChunk"], delta_key: str = "content"
|
||||
) -> ChatCompletionAssistantContentValue:
|
||||
content_list: Final[list[str]] = []
|
||||
for chunk in chunks:
|
||||
|
|
@ -475,7 +530,7 @@ class ChunkProcessor:
|
|||
return combined_content
|
||||
|
||||
def get_combined_thinking_content(
|
||||
self, chunks: list[dict[str, Any]]
|
||||
self, chunks: Sequence["_ThinkingChunk"]
|
||||
) -> list[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] | None:
|
||||
from litellm.types.llms.openai import (
|
||||
ChatCompletionRedactedThinkingBlock,
|
||||
|
|
@ -532,10 +587,10 @@ class ChunkProcessor:
|
|||
return thinking_blocks
|
||||
return None
|
||||
|
||||
def get_combined_reasoning_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAssistantContentValue:
|
||||
def get_combined_reasoning_content(self, chunks: Sequence["_ContentChunk"]) -> ChatCompletionAssistantContentValue:
|
||||
return self.get_combined_content(chunks, delta_key="reasoning_content")
|
||||
|
||||
def get_combined_audio_content(self, chunks: list[dict[str, Any]]) -> ChatCompletionAudioResponse:
|
||||
def get_combined_audio_content(self, chunks: Sequence["_AudioChunk"]) -> ChatCompletionAudioResponse:
|
||||
base64_data_list: Final[list[str]] = []
|
||||
transcript_list: Final[list[str]] = []
|
||||
expires_at: int | None = None
|
||||
|
|
@ -544,7 +599,7 @@ class ChunkProcessor:
|
|||
for chunk in chunks:
|
||||
choices = chunk["choices"]
|
||||
for choice in choices:
|
||||
delta = choice.get("delta") or {}
|
||||
delta: _AudioDelta = choice.get("delta") or {}
|
||||
audio: ChatCompletionAudioDelta | None = delta.get("audio")
|
||||
if audio is not None:
|
||||
for k, v in audio.items():
|
||||
|
|
@ -565,7 +620,7 @@ class ChunkProcessor:
|
|||
id=id,
|
||||
)
|
||||
|
||||
def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> dict:
|
||||
def _usage_chunk_calculation_helper(self, usage_chunk: Usage) -> "_UsageSummary":
|
||||
prompt_tokens = 0
|
||||
completion_tokens = 0
|
||||
## anthropic prompt caching information ##
|
||||
|
|
@ -623,8 +678,8 @@ class ChunkProcessor:
|
|||
return reasoning_tokens
|
||||
|
||||
@staticmethod
|
||||
def _extract_usage_chunk(chunk: dict[str, Any] | ModelResponse | ModelResponseStream) -> Usage | None:
|
||||
usage_chunk: Usage | dict[str, Any] | None = None
|
||||
def _extract_usage_chunk(chunk: "_UsageBearingChunk | ModelResponse | ModelResponseStream") -> Usage | None:
|
||||
usage_chunk: Usage | None = None
|
||||
if hasattr(chunk, "usage") and chunk.usage is not None:
|
||||
usage_chunk = chunk.usage
|
||||
elif "usage" in chunk:
|
||||
|
|
@ -640,7 +695,7 @@ class ChunkProcessor:
|
|||
|
||||
def _calculate_usage_per_chunk(
|
||||
self,
|
||||
chunks: list[dict[str, Any] | ModelResponse],
|
||||
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
|
||||
) -> "UsagePerChunk":
|
||||
from litellm.types.litellm_core_utils.streaming_chunk_builder_utils import (
|
||||
UsagePerChunk,
|
||||
|
|
@ -721,13 +776,7 @@ class ChunkProcessor:
|
|||
"web_search_requests",
|
||||
)
|
||||
|
||||
prompt_tokens_details = (
|
||||
cast(
|
||||
PromptTokensDetailsWrapper | None,
|
||||
usage_chunk_dict["prompt_tokens_details"],
|
||||
)
|
||||
or prompt_tokens_details
|
||||
)
|
||||
prompt_tokens_details = usage_chunk_dict["prompt_tokens_details"] or prompt_tokens_details
|
||||
|
||||
cache_creation_token_details = capture_cache_creation_token_details(
|
||||
prompt_tokens_details, cache_creation_token_details
|
||||
|
|
@ -754,11 +803,31 @@ class ChunkProcessor:
|
|||
completion_tokens_details=completion_tokens_details,
|
||||
prompt_tokens_details=prompt_tokens_details,
|
||||
cost=cost,
|
||||
inference_geo=self._last_provider_pricing_field(chunks, "inference_geo"),
|
||||
speed=self._last_provider_pricing_field(chunks, "speed"),
|
||||
)
|
||||
|
||||
def _last_provider_pricing_field(
|
||||
self,
|
||||
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
|
||||
field: str,
|
||||
) -> str | None:
|
||||
"""
|
||||
Last value of a provider-specific usage field that changes pricing but is not a
|
||||
declared ``Usage`` field, e.g. Anthropic's ``speed`` (fast mode multiplies
|
||||
non-cache token cost) and ``inference_geo``.
|
||||
"""
|
||||
values: Final = [
|
||||
value
|
||||
for chunk in chunks
|
||||
if (usage_chunk := self._extract_usage_chunk(chunk)) is not None
|
||||
and isinstance(value := getattr(usage_chunk, field, None), str)
|
||||
]
|
||||
return values[-1] if values else None
|
||||
|
||||
@staticmethod
|
||||
def _reset_anthropic_cursor_completion_tokens(
|
||||
chunks: list[dict[str, Any] | ModelResponse],
|
||||
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
|
||||
completion_tokens: int,
|
||||
completion_usage_updates: int,
|
||||
) -> int:
|
||||
|
|
@ -797,7 +866,7 @@ class ChunkProcessor:
|
|||
|
||||
def calculate_usage(
|
||||
self,
|
||||
chunks: list[dict[str, Any] | ModelResponse],
|
||||
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
|
||||
model: str,
|
||||
completion_output: str,
|
||||
messages: list | None = None,
|
||||
|
|
@ -851,8 +920,8 @@ class ChunkProcessor:
|
|||
setattr(returned_usage, "cache_read_input_tokens", cache_read_input_tokens) # for anthropic
|
||||
if completion_tokens_details is not None:
|
||||
if isinstance(completion_tokens_details, CompletionTokensDetails):
|
||||
returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper(
|
||||
**completion_tokens_details.model_dump()
|
||||
returned_usage.completion_tokens_details = CompletionTokensDetailsWrapper.model_validate(
|
||||
completion_tokens_details.model_dump()
|
||||
)
|
||||
else:
|
||||
returned_usage.completion_tokens_details = completion_tokens_details
|
||||
|
|
@ -885,7 +954,16 @@ class ChunkProcessor:
|
|||
|
||||
# Return a new usage object with the new values
|
||||
|
||||
returned_usage = Usage(**returned_usage.model_dump())
|
||||
provider_pricing_fields: Final = {
|
||||
field: value
|
||||
for field, value in (
|
||||
("inference_geo", calculated_usage_per_chunk["inference_geo"]),
|
||||
("speed", calculated_usage_per_chunk["speed"]),
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
|
||||
returned_usage = Usage(**returned_usage.model_dump(), **provider_pricing_fields)
|
||||
|
||||
return returned_usage
|
||||
|
||||
|
|
|
|||
|
|
@ -6,13 +6,14 @@ import logging
|
|||
import threading
|
||||
import time
|
||||
import traceback
|
||||
from collections.abc import AsyncIterator, Callable, Iterator
|
||||
from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Final, NoReturn, TypeVar, cast
|
||||
from typing import Any, Final, NoReturn, Protocol, TypeVar, cast
|
||||
|
||||
import anyio
|
||||
import httpx
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import NotRequired, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm import verbose_logger
|
||||
|
|
@ -54,7 +55,7 @@ _SYNC_ITER_EXHAUSTED: Final = object()
|
|||
_GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__)
|
||||
|
||||
|
||||
def _next_sync_or_exhausted(it: Any) -> Any:
|
||||
def _next_sync_or_exhausted(it: Any) -> object:
|
||||
"""
|
||||
Call next(it) from a thread and return _SYNC_ITER_EXHAUSTED on StopIteration.
|
||||
|
||||
|
|
@ -68,7 +69,7 @@ def _next_sync_or_exhausted(it: Any) -> Any:
|
|||
return _SYNC_ITER_EXHAUSTED
|
||||
|
||||
|
||||
def is_async_iterable(obj: Any) -> bool:
|
||||
def is_async_iterable(obj: object) -> bool:
|
||||
"""
|
||||
Check if an object is an async iterable (can be used with 'async for').
|
||||
|
||||
|
|
@ -81,7 +82,7 @@ def is_async_iterable(obj: Any) -> bool:
|
|||
return isinstance(obj, collections.abc.AsyncIterable)
|
||||
|
||||
|
||||
def print_verbose(print_statement):
|
||||
def print_verbose(print_statement: object):
|
||||
try:
|
||||
if litellm.set_verbose:
|
||||
print(print_statement) # noqa: T201
|
||||
|
|
@ -96,18 +97,70 @@ class _ProviderChunkParsed:
|
|||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ProviderChunkEarlyReturn:
|
||||
value: Any
|
||||
value: "ModelResponseStream | None"
|
||||
|
||||
|
||||
_ProviderChunkResult = _ProviderChunkParsed | _ProviderChunkEarlyReturn
|
||||
|
||||
|
||||
class _PredibaseStreamData(TypedDict):
|
||||
token: NotRequired[Mapping[str, str]]
|
||||
details: Mapping[str, str]
|
||||
generated_text: str | None
|
||||
error: str | None
|
||||
|
||||
|
||||
class _Ai21StreamData(TypedDict):
|
||||
completions: Sequence[Mapping[str, Mapping[str, str]]]
|
||||
|
||||
|
||||
class _MaritalkStreamData(TypedDict):
|
||||
answer: str
|
||||
|
||||
|
||||
class _NlpCloudStreamData(TypedDict):
|
||||
generated_text: str
|
||||
|
||||
|
||||
class _AlephAlphaStreamData(TypedDict):
|
||||
completions: Sequence[Mapping[str, str]]
|
||||
|
||||
|
||||
class _AzureStreamChoice(TypedDict):
|
||||
delta: Mapping[str, str] | None
|
||||
finish_reason: str | None
|
||||
|
||||
|
||||
class _AzureStreamData(TypedDict):
|
||||
choices: Sequence[_AzureStreamChoice]
|
||||
|
||||
|
||||
class _BasetenModelOutput(TypedDict):
|
||||
data: NotRequired[Sequence[str]]
|
||||
|
||||
|
||||
class _BasetenStreamData(TypedDict):
|
||||
token: NotRequired[Mapping[str, str]]
|
||||
model_output: NotRequired["_BasetenModelOutput | str"]
|
||||
completion: NotRequired[object]
|
||||
|
||||
|
||||
class _DeltaDumpDict(TypedDict):
|
||||
role: NotRequired[str | None]
|
||||
tool_calls: NotRequired[Sequence[Mapping[str, object]]]
|
||||
|
||||
|
||||
class _TextCompletionChoiceLike(Protocol):
|
||||
text: str
|
||||
finish_reason: str | None
|
||||
|
||||
|
||||
class CustomStreamWrapper:
|
||||
def __init__(
|
||||
self,
|
||||
completion_stream,
|
||||
model,
|
||||
logging_obj: Any,
|
||||
logging_obj: LiteLLMLoggingObject,
|
||||
custom_llm_provider: str | None = None,
|
||||
stream_options=None,
|
||||
make_call: Callable | None = None,
|
||||
|
|
@ -186,7 +239,7 @@ class CustomStreamWrapper:
|
|||
# Snapshot assumes self._hidden_params is populated from litellm_params
|
||||
# at init and never mutated during the stream. If that ever changes,
|
||||
# this cache must be removed.
|
||||
self._base_hidden_params: dict[str, Any] = {
|
||||
self._base_hidden_params: dict[str, object] = {
|
||||
**self._hidden_params,
|
||||
"response_cost": None,
|
||||
}
|
||||
|
|
@ -213,7 +266,75 @@ class CustomStreamWrapper:
|
|||
def __aiter__(self) -> AsyncIterator["ModelResponseStream"]:
|
||||
return self
|
||||
|
||||
def _restore_consumer_correlation_context(self, *, guarded: bool = False) -> None:
|
||||
"""Restore trace_id/session_id in the *consuming* thread/task/context.
|
||||
|
||||
wrapper_async() deliberately skips restoring correlation context when
|
||||
it returns a stream, so log lines emitted while the caller iterates it
|
||||
still carry this call's ids (see request_correlation_in_logs).
|
||||
wrapper() (the sync path) never stamps anything in the first place -
|
||||
see Logging.__init__'s supports_correlation_logging - so this method
|
||||
is an inert no-op for sync-created streams, harmless to call anyway
|
||||
since the class is shared between __next__ and __anext__.
|
||||
But the terminal success/failure handlers this stream dispatches to
|
||||
finish the job run on a *different* Task/thread (asyncio.create_task,
|
||||
threading.Thread, or the shared executor) - restoring there fixes up
|
||||
that detached context, not the one actually running the caller's
|
||||
`for`/`async for` loop. Call this at every point control genuinely
|
||||
returns to that consuming context: natural exhaustion (StopIteration/
|
||||
StopAsyncIteration), a raised failure, or explicit aclose(). Never let
|
||||
this raise - it must not break the caller's actual stream handling.
|
||||
|
||||
guarded=True (only __del__ uses this) skips the restore unless the
|
||||
contextvars still hold the ids this stream's own call set, so a
|
||||
delayed finalizer never overwrites a different, still-active call
|
||||
that has since taken over the same Task/thread's context.
|
||||
"""
|
||||
try:
|
||||
logging_obj: Final = getattr(self, "logging_obj", None)
|
||||
if logging_obj is None:
|
||||
return
|
||||
method_name: Final = (
|
||||
"_restore_correlation_context_if_unclaimed" if guarded else "_restore_correlation_context"
|
||||
)
|
||||
restore: Final = getattr(logging_obj, method_name, None)
|
||||
if restore is not None:
|
||||
restore()
|
||||
except Exception as restore_error: # noqa: BLE001 # best-effort cleanup; must not raise into the caller
|
||||
verbose_logger.debug("could not restore correlation context: %s", restore_error)
|
||||
|
||||
def __del__(self) -> None:
|
||||
"""Best-effort correlation-context cleanup for an abandoned async stream.
|
||||
|
||||
Only meaningfully applies to streams created by wrapper_async(): it
|
||||
leaves contextvars "open" across the caller's iteration, so if the
|
||||
caller never fully consumes the stream - stops early, drops the
|
||||
reference, cancels it - none of the exit points
|
||||
_restore_consumer_correlation_context() is called from ever run. For a
|
||||
sync stream (wrapper()), this is a no-op in practice: wrapper() never
|
||||
stamps trace_id/session_id for sync calls in the first place (see
|
||||
Logging.__init__'s supports_correlation_logging), so there is nothing
|
||||
for this to clean up.
|
||||
|
||||
This is a best-effort fallback, not a guarantee: __del__ timing is
|
||||
unpredictable (delayed by cyclic GC, not guaranteed at interpreter
|
||||
shutdown, and may run on a different thread), so this can only reduce
|
||||
how long the leak persists, not eliminate it. That's an acceptable
|
||||
trade specifically because its blast radius is bounded to the one
|
||||
asyncio Task this stream's own call ran in - each async call has its
|
||||
own copy of the contextvars, and Tasks (unlike a thread pool's worker
|
||||
threads) are never recycled across requests, so a delayed or missed
|
||||
cleanup here can never misattribute a *different* request's logs.
|
||||
guarded=True additionally ensures it never clobbers a different,
|
||||
still-active call's context within that same Task if this fires late.
|
||||
"""
|
||||
self._restore_consumer_correlation_context(guarded=True)
|
||||
|
||||
async def aclose(self):
|
||||
# Restore the consumer's outer context only after the underlying
|
||||
# provider stream's own close (and its diagnostic logging below, if
|
||||
# closing fails) completes - not before - so those log lines still
|
||||
# carry this closing stream's own trace_id/session_id.
|
||||
if self.completion_stream is not None:
|
||||
stream_to_close: Final = self.completion_stream
|
||||
self.completion_stream = None
|
||||
|
|
@ -233,6 +354,7 @@ class CustomStreamWrapper:
|
|||
"CustomStreamWrapper.aclose: error closing completion_stream: %s",
|
||||
e,
|
||||
)
|
||||
self._restore_consumer_correlation_context()
|
||||
|
||||
def check_send_stream_usage(self, stream_options: dict | None):
|
||||
return stream_options is not None and stream_options.get("include_usage", False) is True
|
||||
|
|
@ -347,7 +469,7 @@ class CustomStreamWrapper:
|
|||
finish_reason = ""
|
||||
print_verbose(f"chunk: {chunk}")
|
||||
if chunk.startswith("data:"):
|
||||
data_json: Final = json.loads(chunk[5:])
|
||||
data_json: Final[_PredibaseStreamData] = json.loads(chunk[5:])
|
||||
print_verbose(f"data json: {data_json}")
|
||||
if "token" in data_json and "text" in data_json["token"]:
|
||||
text = data_json["token"]["text"]
|
||||
|
|
@ -377,7 +499,7 @@ class CustomStreamWrapper:
|
|||
|
||||
def handle_ai21_chunk(self, chunk): # fake streaming
|
||||
chunk = chunk.decode("utf-8")
|
||||
data_json: Final = json.loads(chunk)
|
||||
data_json: Final[_Ai21StreamData] = json.loads(chunk)
|
||||
try:
|
||||
text: Final = data_json["completions"][0]["data"]["text"]
|
||||
is_finished: Final = True
|
||||
|
|
@ -392,7 +514,7 @@ class CustomStreamWrapper:
|
|||
|
||||
def handle_maritalk_chunk(self, chunk): # fake streaming
|
||||
chunk = chunk.decode("utf-8")
|
||||
data_json: Final = json.loads(chunk)
|
||||
data_json: Final[_MaritalkStreamData] = json.loads(chunk)
|
||||
try:
|
||||
text: Final = data_json["answer"]
|
||||
is_finished: Final = True
|
||||
|
|
@ -413,7 +535,7 @@ class CustomStreamWrapper:
|
|||
if self.model and "dolphin" in self.model:
|
||||
chunk = self.process_chunk(chunk=chunk)
|
||||
else:
|
||||
data_json: Final = json.loads(chunk)
|
||||
data_json: Final[_NlpCloudStreamData] = json.loads(chunk)
|
||||
chunk = data_json["generated_text"]
|
||||
text = chunk
|
||||
if "[DONE]" in text:
|
||||
|
|
@ -430,7 +552,7 @@ class CustomStreamWrapper:
|
|||
|
||||
def handle_aleph_alpha_chunk(self, chunk):
|
||||
chunk = chunk.decode("utf-8")
|
||||
data_json: Final = json.loads(chunk)
|
||||
data_json: Final[_AlephAlphaStreamData] = json.loads(chunk)
|
||||
try:
|
||||
text: Final = data_json["completions"][0]["completion"]
|
||||
is_finished: Final = True
|
||||
|
|
@ -458,7 +580,7 @@ class CustomStreamWrapper:
|
|||
"finish_reason": finish_reason,
|
||||
}
|
||||
elif chunk.startswith("data:"):
|
||||
data_json: Final = json.loads(chunk[5:]) # chunk.startswith("data:"):
|
||||
data_json: Final[_AzureStreamData] = json.loads(chunk[5:]) # chunk.startswith("data:"):
|
||||
try:
|
||||
if len(data_json["choices"]) > 0:
|
||||
delta: Final = data_json["choices"][0]["delta"]
|
||||
|
|
@ -547,7 +669,7 @@ class CustomStreamWrapper:
|
|||
text = ""
|
||||
is_finished = False
|
||||
finish_reason = None
|
||||
choices: Final = getattr(chunk, "choices", [])
|
||||
choices: Final[Sequence[_TextCompletionChoiceLike]] = getattr(chunk, "choices", [])
|
||||
if len(choices) > 0:
|
||||
text = choices[0].text
|
||||
if choices[0].finish_reason is not None:
|
||||
|
|
@ -568,7 +690,7 @@ class CustomStreamWrapper:
|
|||
is_finished = False
|
||||
finish_reason = None
|
||||
usage = None
|
||||
choices: Final = getattr(chunk, "choices", [])
|
||||
choices: Final[Sequence[_TextCompletionChoiceLike]] = getattr(chunk, "choices", [])
|
||||
if len(choices) > 0:
|
||||
text = choices[0].text
|
||||
if choices[0].finish_reason is not None:
|
||||
|
|
@ -585,12 +707,12 @@ class CustomStreamWrapper:
|
|||
except Exception as e:
|
||||
raise e
|
||||
|
||||
def handle_baseten_chunk(self, chunk):
|
||||
def handle_baseten_chunk(self, chunk) -> str:
|
||||
try:
|
||||
chunk = chunk.decode("utf-8")
|
||||
if len(chunk) > 0:
|
||||
if chunk.startswith("data:"):
|
||||
data_json = json.loads(chunk[5:])
|
||||
data_json: _BasetenStreamData = json.loads(chunk[5:])
|
||||
if "token" in data_json and "text" in data_json["token"]:
|
||||
return data_json["token"]["text"]
|
||||
else:
|
||||
|
|
@ -1256,13 +1378,14 @@ class CustomStreamWrapper:
|
|||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
if "usage" in response_obj is not None:
|
||||
_codestral_usage: Final[Usage] = response_obj["usage"]
|
||||
setattr(
|
||||
model_response,
|
||||
"usage",
|
||||
litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].prompt_tokens,
|
||||
completion_tokens=response_obj["usage"].completion_tokens,
|
||||
total_tokens=response_obj["usage"].total_tokens,
|
||||
prompt_tokens=_codestral_usage.prompt_tokens,
|
||||
completion_tokens=_codestral_usage.completion_tokens,
|
||||
total_tokens=_codestral_usage.total_tokens,
|
||||
),
|
||||
)
|
||||
elif self.custom_llm_provider == "azure_text":
|
||||
|
|
@ -1405,7 +1528,7 @@ class CustomStreamWrapper:
|
|||
is None
|
||||
):
|
||||
t.function.arguments = ""
|
||||
_json_delta: Final = delta.model_dump()
|
||||
_json_delta: Final[_DeltaDumpDict] = delta.model_dump()
|
||||
if "role" not in _json_delta or _json_delta["role"] is None:
|
||||
_json_delta["role"] = "assistant" # mistral's api returns role as None
|
||||
if "tool_calls" in _json_delta and isinstance(_json_delta["tool_calls"], list):
|
||||
|
|
@ -1675,7 +1798,7 @@ class CustomStreamWrapper:
|
|||
usage.cost, copy it into _hidden_params so litellm's cost
|
||||
calculator uses it instead of a token-based estimate.
|
||||
"""
|
||||
_usage: Final = getattr(response, "usage", None)
|
||||
_usage: Final[Usage | None] = getattr(response, "usage", None)
|
||||
if _usage is not None and hasattr(_usage, "cost") and _usage.cost is not None:
|
||||
if "additional_headers" not in response._hidden_params:
|
||||
response._hidden_params["additional_headers"] = {}
|
||||
|
|
@ -1839,6 +1962,7 @@ class CustomStreamWrapper:
|
|||
if self.sent_stream_usage is False and self.send_stream_usage is True:
|
||||
self.sent_stream_usage = True
|
||||
return response
|
||||
self._restore_consumer_correlation_context()
|
||||
raise # Re-raise StopIteration
|
||||
else:
|
||||
self.sent_last_chunk = True
|
||||
|
|
@ -1852,6 +1976,19 @@ class CustomStreamWrapper:
|
|||
processed_chunk,
|
||||
cache_hit,
|
||||
) # log response
|
||||
# Deliberately do NOT restore context here even though
|
||||
# completion_stream is already exhausted: this chunk is still
|
||||
# real data belonging to this call, and the caller's own
|
||||
# (application-level) log statements processing it run
|
||||
# immediately after this return, in this same synchronous
|
||||
# frame - restoring first would make those lines carry the
|
||||
# wrong ids, which is exactly what leaving context open during
|
||||
# iteration is meant to prevent (see
|
||||
# _restore_consumer_correlation_context's docstring). A caller
|
||||
# that keeps iterating gets cleaned up on its next __next__()
|
||||
# call (immediate StopIteration, handled above); one that
|
||||
# stops right here relies on aclose() or the best-effort
|
||||
# __del__ guard instead.
|
||||
return processed_chunk
|
||||
except Exception as e:
|
||||
traceback_exception: Final = traceback.format_exc()
|
||||
|
|
@ -1879,8 +2016,12 @@ class CustomStreamWrapper:
|
|||
cache_hit = False
|
||||
if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response":
|
||||
cache_hit = True
|
||||
self._check_max_streaming_duration()
|
||||
try:
|
||||
# Inside the try (not before it) so a raised litellm.Timeout flows
|
||||
# through the same except Exception -> _handle_stream_fallback_error
|
||||
# path as every other failure, restoring the consumer's correlation
|
||||
# context - a check before the try would bypass that entirely.
|
||||
self._check_max_streaming_duration()
|
||||
if self.completion_stream is None:
|
||||
await self.fetch_stream()
|
||||
|
||||
|
|
@ -2083,10 +2224,17 @@ class CustomStreamWrapper:
|
|||
)
|
||||
)
|
||||
|
||||
self._restore_consumer_correlation_context()
|
||||
raise StopAsyncIteration # Re-raise StopIteration
|
||||
else:
|
||||
self.sent_last_chunk = True
|
||||
processed_chunk: Final = self.finish_reason_handler()
|
||||
# see sync __next__'s sibling branch: deliberately do NOT restore
|
||||
# here - this chunk is still this call's own data, and restoring
|
||||
# before returning it would corrupt the caller's own log
|
||||
# statements processing it. A caller that keeps iterating gets
|
||||
# cleaned up on the next __anext__() call; one that stops here
|
||||
# relies on aclose() or the best-effort __del__ guard.
|
||||
return processed_chunk
|
||||
|
||||
def _log_stream_failure_and_raise(self, e: Exception) -> NoReturn:
|
||||
|
|
@ -2138,7 +2286,12 @@ class CustomStreamWrapper:
|
|||
"""
|
||||
from litellm.exceptions import MidStreamFallbackError
|
||||
|
||||
# Map to OpenAI exception format
|
||||
# Map to OpenAI exception format. Some providers' mappers (e.g.
|
||||
# _map_anthropic_exception, _map_aleph_alpha_exception) synchronously
|
||||
# log a debug diagnostic (the raw status code) as part of mapping -
|
||||
# restore the consumer's outer context only after this completes, so
|
||||
# that diagnostic log line still carries the failing stream's own
|
||||
# trace_id/session_id instead of the consumer's (or an empty one).
|
||||
if isinstance(e, OpenAIError):
|
||||
mapped_exception: Exception = e
|
||||
else:
|
||||
|
|
@ -2152,6 +2305,7 @@ class CustomStreamWrapper:
|
|||
)
|
||||
except Exception as mapping_error:
|
||||
mapped_exception = mapping_error
|
||||
self._restore_consumer_correlation_context()
|
||||
|
||||
def _normalize_status_code(exc: Exception) -> int | None:
|
||||
"""Best-effort status_code extraction."""
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ Pattern Overview:
|
|||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
|
|
@ -110,14 +111,10 @@ EMPTY_EXTRACTED_INPUT: Final = ExtractedInput(scanned=(), images=())
|
|||
|
||||
|
||||
class AnthropicMessagesHandler(BaseTranslation):
|
||||
"""
|
||||
Handler for processing Anthropic messages with guardrails.
|
||||
"""Process Anthropic messages with guardrails.
|
||||
|
||||
This class provides methods to:
|
||||
1. Process input messages (pre-call hook)
|
||||
2. Process output responses (post-call hook)
|
||||
|
||||
Methods can be overridden to customize behavior for different message formats.
|
||||
In-sequence system entries are untrusted client input. This handler scans and preserves
|
||||
them through guardrail rewrites; downstream provider handling is out of scope.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
|
|
@ -331,16 +328,30 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
skip_tool: Final = effective_skip_tool_message_for_guardrail(guardrail_to_apply)
|
||||
scan_only_tool_results: Final = effective_scan_only_tool_results_for_guardrail(guardrail_to_apply)
|
||||
|
||||
chat_completion_compatible_request: Final = self._translate_to_openai(data)
|
||||
# Exclude only the trusted top-level prompt. In-sequence system entries are untrusted
|
||||
# and must stay aligned with texts_to_check for positional masking. When the top-level
|
||||
# prompt is included, the pre-existing count mismatch disables positional masking.
|
||||
translation_source: Final = { # mutable-ok: API message payload
|
||||
key: value for key, value in data.items() if key != "system"
|
||||
}
|
||||
chat_completion_compatible_request: Final = self._translate_to_openai(translation_source)
|
||||
|
||||
full_structured_messages: Final = cast(
|
||||
list[AllMessageValues],
|
||||
chat_completion_compatible_request.get("messages", []),
|
||||
)
|
||||
has_midturn_system_message: Final = any(
|
||||
str(message.get("role") or "").lower() == "system" for message in full_structured_messages
|
||||
)
|
||||
hoisted_system_message: Final = None if skip_system else self._hoisted_top_level_system_message(data)
|
||||
if hoisted_system_message is not None:
|
||||
full_structured_messages.insert(0, hoisted_system_message)
|
||||
# skip_system already excluded the trusted top-level prompt (it is simply not hoisted);
|
||||
# in-sequence system entries are untrusted and always stay in scope.
|
||||
scoped_message_indices: Final = scoped_structured_message_indices(
|
||||
full_structured_messages,
|
||||
scan_only_tool_results=scan_only_tool_results,
|
||||
skip_system=skip_system,
|
||||
skip_system=False,
|
||||
skip_tool=skip_tool,
|
||||
)
|
||||
structured_messages: Final = [full_structured_messages[index] for index in scoped_message_indices]
|
||||
|
|
@ -422,6 +433,8 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
scoped_indices=scoped_message_indices,
|
||||
guardrailed_scoped=guardrailed_structured_messages,
|
||||
),
|
||||
hoisted_system_message=hoisted_system_message,
|
||||
preserve_system_messages=has_midturn_system_message,
|
||||
)
|
||||
else:
|
||||
# Step 3: Map guardrail responses back to original message structure
|
||||
|
|
@ -435,36 +448,150 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def _write_back_structured_messages(data: dict, structured_messages: list) -> None:
|
||||
"""Convert compressed structured_messages back to Anthropic format and write to data.
|
||||
def _hoisted_top_level_system_message(
|
||||
self, data: dict
|
||||
) -> AllMessageValues | None: # mutable-ok: API message payload
|
||||
"""Return the system message produced by translating the top-level prompt."""
|
||||
system: Final = data.get("system")
|
||||
if not system:
|
||||
return None
|
||||
probe: Final = self._translate_to_openai(
|
||||
{ # mutable-ok: API message payload
|
||||
"model": data.get("model") or "",
|
||||
"messages": [], # mutable-ok: API message payload
|
||||
"system": system,
|
||||
}
|
||||
)
|
||||
hoisted: Final = probe.get("messages") or [] # mutable-ok: API message payload
|
||||
return hoisted[0] if hoisted else None
|
||||
|
||||
``anthropic_messages_pt`` merges every run of consecutive user/tool rows
|
||||
into a single message, so a turn carrying only tool results and the user
|
||||
turn that follows it come back fused, and the request the model sees no
|
||||
longer has the boundaries the client sent. Converting a row at a time
|
||||
would keep them apart but breaks tool pairing: an assistant row whose
|
||||
tool results sit outside its own call reads as an orphaned tool call,
|
||||
and under ``modify_params`` the sanitizer answers it with a synthetic
|
||||
"tool execution skipped" result and drops the real one. Converting each
|
||||
assistant row together with the tool rows that answer it, and every
|
||||
other row on its own, satisfies both.
|
||||
"""
|
||||
@staticmethod
|
||||
def _openai_system_message_to_anthropic(
|
||||
message: dict[str, Any],
|
||||
) -> dict[str, Any] | None: # mutable-ok: API message payload
|
||||
"""Convert an OpenAI system message to the client's Anthropic-shaped entry."""
|
||||
content: Final = message.get("content")
|
||||
if isinstance(content, str):
|
||||
return (
|
||||
{"role": "system", "content": content} if content else None # mutable-ok: API message payload
|
||||
) # mutable-ok: API message payload
|
||||
if not isinstance(content, list):
|
||||
return None
|
||||
blocks: Final[list[dict[str, Any]]] = [] # mutable-ok: API message payload
|
||||
for block in content:
|
||||
if not isinstance(block, dict) or block.get("type") != "text":
|
||||
continue
|
||||
text = block.get("text")
|
||||
if not isinstance(text, str) or not text:
|
||||
continue
|
||||
anthropic_block: dict[str, Any] = { # mutable-ok: API message payload
|
||||
"type": "text",
|
||||
"text": text,
|
||||
} # mutable-ok: API message payload
|
||||
cache_control = block.get("cache_control")
|
||||
if cache_control:
|
||||
anthropic_block["cache_control"] = deepcopy(cache_control)
|
||||
blocks.append(anthropic_block)
|
||||
return (
|
||||
{"role": "system", "content": blocks} if blocks else None # mutable-ok: API message payload
|
||||
) # mutable-ok: API message payload
|
||||
|
||||
@staticmethod
|
||||
def _is_hoisted_top_level_system(message: object, hoisted_system_message: object) -> bool:
|
||||
"""Match the hoisted prompt by identity, or by value after serialization."""
|
||||
if hoisted_system_message is None:
|
||||
return False
|
||||
if message is hoisted_system_message:
|
||||
return True
|
||||
return (
|
||||
isinstance(message, dict) and isinstance(hoisted_system_message, dict) and message == hoisted_system_message
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_system(message: object) -> bool:
|
||||
"""Whether the row is an in-sequence system message."""
|
||||
return isinstance(message, dict) and str(message.get("role") or "").lower() == "system"
|
||||
|
||||
@staticmethod
|
||||
def _defer_systems_inside_tool_exchanges(
|
||||
structured_messages: list, # mutable-ok: API message payload
|
||||
) -> list:
|
||||
"""Hold a system row until the tool exchange around it completes so the call/result pair converts together."""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import group_tool_exchanges
|
||||
|
||||
non_system_positions: Final[list[int]] = [
|
||||
index
|
||||
for index, message in enumerate(structured_messages)
|
||||
if not AnthropicMessagesHandler._is_system(message)
|
||||
]
|
||||
exchange_end_for_start: Final[dict[int, int]] = {
|
||||
non_system_positions[group[0]]: non_system_positions[group[-1]]
|
||||
for group in group_tool_exchanges([structured_messages[index] for index in non_system_positions])
|
||||
if len(group) > 1
|
||||
}
|
||||
ordered: Final[list] = [] # mutable-ok: API message payload
|
||||
deferred_systems: Final[list] = [] # mutable-ok: API message payload
|
||||
open_exchange_end = -1 # rebind-ok: advances to the enclosing exchange's last index
|
||||
for index, message in enumerate(structured_messages):
|
||||
if AnthropicMessagesHandler._is_system(message) and index < open_exchange_end:
|
||||
deferred_systems.append(message)
|
||||
continue
|
||||
open_exchange_end = exchange_end_for_start.get(index, open_exchange_end)
|
||||
ordered.append(message)
|
||||
if index >= open_exchange_end and deferred_systems:
|
||||
ordered.extend(deferred_systems)
|
||||
deferred_systems.clear()
|
||||
ordered.extend(deferred_systems)
|
||||
return ordered
|
||||
|
||||
@staticmethod
|
||||
def _write_back_structured_messages(
|
||||
data: dict, # mutable-ok: API message payload
|
||||
structured_messages: list, # mutable-ok: API message payload
|
||||
hoisted_system_message: object = None,
|
||||
preserve_system_messages: bool = False,
|
||||
) -> None:
|
||||
"""Write a guardrail's structured-message rewrite back without losing corrections."""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
anthropic_messages_pt,
|
||||
group_tool_exchanges,
|
||||
)
|
||||
|
||||
_is_system: Final = AnthropicMessagesHandler._is_system
|
||||
model: Final = str(data.get("model") or "")
|
||||
non_system: Final = [m for m in structured_messages if m.get("role") != "system"]
|
||||
groups: Final = tuple([non_system[index] for index in group] for group in group_tool_exchanges(non_system)) or (
|
||||
non_system,
|
||||
)
|
||||
converted: Final = [
|
||||
message
|
||||
for group in groups
|
||||
for message in anthropic_messages_pt(messages=group, model=model, llm_provider="anthropic")
|
||||
]
|
||||
converted: Final[list] = [] # mutable-ok: API message payload
|
||||
|
||||
def _convert_run(run: list) -> None: # mutable-ok: API message payload
|
||||
for group in group_tool_exchanges(run):
|
||||
converted.extend(
|
||||
anthropic_messages_pt(
|
||||
messages=[run[index] for index in group], # mutable-ok: API message payload
|
||||
model=model,
|
||||
llm_provider="anthropic",
|
||||
)
|
||||
)
|
||||
|
||||
ordered: Final = AnthropicMessagesHandler._defer_systems_inside_tool_exchanges(structured_messages)
|
||||
run: Final[list] = [] # mutable-ok: API message payload
|
||||
hoisted_dropped = False # rebind-ok: flips once the hoisted prompt is dropped
|
||||
for message in ordered:
|
||||
if not _is_system(message):
|
||||
run.append(message)
|
||||
continue
|
||||
_convert_run(run)
|
||||
run.clear()
|
||||
if not hoisted_dropped and AnthropicMessagesHandler._is_hoisted_top_level_system(
|
||||
message, hoisted_system_message
|
||||
):
|
||||
hoisted_dropped = True
|
||||
continue
|
||||
if preserve_system_messages:
|
||||
anthropic_system = AnthropicMessagesHandler._openai_system_message_to_anthropic(message)
|
||||
if anthropic_system is not None:
|
||||
converted.append(anthropic_system)
|
||||
_convert_run(run)
|
||||
if not any(not _is_system(message) for message in converted):
|
||||
converted.extend(anthropic_messages_pt(messages=[], model=model, llm_provider="anthropic"))
|
||||
for msg in converted:
|
||||
content = msg.get("content")
|
||||
if isinstance(content, list):
|
||||
|
|
@ -473,6 +600,31 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
block.pop("cache_control", None)
|
||||
data["messages"] = converted
|
||||
|
||||
@staticmethod
|
||||
def _extract_midturn_system_text(
|
||||
message: dict[str, Any], # mutable-ok: API message payload
|
||||
msg_idx: int,
|
||||
) -> ExtractedInput:
|
||||
"""Match the adapter's filtering so positional guardrail write-back stays aligned."""
|
||||
content: Final = message.get("content")
|
||||
if isinstance(content, str):
|
||||
if not content:
|
||||
return EMPTY_EXTRACTED_INPUT
|
||||
return ExtractedInput(scanned=(ScannedText(content, MessageContentTarget(msg_idx)),), images=())
|
||||
if not isinstance(content, list):
|
||||
return EMPTY_EXTRACTED_INPUT
|
||||
return ExtractedInput(
|
||||
scanned=tuple(
|
||||
ScannedText(text_str, ContentBlockTextTarget(msg_idx, content_idx))
|
||||
for content_idx, content_item in enumerate(content)
|
||||
if isinstance(content_item, dict)
|
||||
and content_item.get("type") == "text"
|
||||
and isinstance(text_str := content_item.get("text"), str)
|
||||
and text_str
|
||||
),
|
||||
images=(),
|
||||
)
|
||||
|
||||
def extract_request_tool_names(self, data: dict) -> list[str]:
|
||||
"""Extract tool names from Anthropic messages request (tools[].name)."""
|
||||
names: Final[list[str]] = []
|
||||
|
|
@ -490,11 +642,17 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
skip_tool_message: bool = False,
|
||||
scan_only_tool_results: bool = False,
|
||||
) -> ExtractedInput:
|
||||
"""Extract text content and images from a message.
|
||||
|
||||
In-sequence system entries are scanned even when ``skip_system_message`` is set:
|
||||
that flag covers only the trusted top-level prompt, which never appears here.
|
||||
"""
|
||||
Extract text content and images from a message.
|
||||
"""
|
||||
role: Final = str(message.get("role") or "").lower()
|
||||
if (skip_system_message and role == "system") or (skip_tool_message and role == "tool"):
|
||||
role: Final = str(message.get("role") or "")
|
||||
if role == "system":
|
||||
if scan_only_tool_results:
|
||||
return EMPTY_EXTRACTED_INPUT
|
||||
return cls._extract_midturn_system_text(message=message, msg_idx=msg_idx)
|
||||
if skip_tool_message and role.lower() == "tool":
|
||||
return EMPTY_EXTRACTED_INPUT
|
||||
|
||||
content: Final = message.get("content", None)
|
||||
|
|
|
|||
|
|
@ -58,6 +58,7 @@ from ..common_utils import AnthropicError, process_anthropic_headers
|
|||
from .transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY, AnthropicConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig
|
||||
|
||||
|
|
@ -206,7 +207,7 @@ class AnthropicChatCompletion(BaseLLM):
|
|||
client: AsyncHTTPHandler | None,
|
||||
encoding,
|
||||
api_key,
|
||||
logging_obj,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
stream,
|
||||
_is_function_call,
|
||||
data: dict,
|
||||
|
|
@ -324,7 +325,7 @@ class AnthropicChatCompletion(BaseLLM):
|
|||
print_verbose: Callable,
|
||||
encoding,
|
||||
api_key,
|
||||
logging_obj,
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
optional_params: dict,
|
||||
timeout: float | httpx.Timeout,
|
||||
litellm_params: dict,
|
||||
|
|
|
|||
|
|
@ -4,9 +4,12 @@ This file contains common utils for anthropic calls.
|
|||
|
||||
import copy
|
||||
import re
|
||||
from typing import Any, Final
|
||||
from collections.abc import Mapping, Sequence
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, Literal
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
||||
|
|
@ -1057,6 +1060,152 @@ def sanitize_tool_use_ids_in_anthropic_messages(messages: list[Any]) -> list[Any
|
|||
return out
|
||||
|
||||
|
||||
class _ReplayedSearchQuery(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
query: str = ""
|
||||
|
||||
|
||||
class _ReplayedWebSearchResult(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
type: Literal["web_search_result"]
|
||||
url: str = ""
|
||||
title: str = ""
|
||||
snippet: str = ""
|
||||
encrypted_content: str = ""
|
||||
|
||||
|
||||
class _ReplayedWebSearchToolResult(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
type: Literal["web_search_tool_result"]
|
||||
tool_use_id: str
|
||||
content: tuple[_ReplayedWebSearchResult, ...]
|
||||
|
||||
|
||||
class _ReplayedServerToolUse(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
type: Literal["server_tool_use"]
|
||||
id: str
|
||||
input: _ReplayedSearchQuery = _ReplayedSearchQuery()
|
||||
|
||||
|
||||
class _TextBlock(BaseModel):
|
||||
type: Literal["text"] = "text"
|
||||
text: str
|
||||
|
||||
|
||||
_WEB_SEARCH_TOOL_RESULT_ADAPTER: Final = TypeAdapter(_ReplayedWebSearchToolResult)
|
||||
_SERVER_TOOL_USE_ADAPTER: Final = TypeAdapter(_ReplayedServerToolUse)
|
||||
|
||||
|
||||
def _flattenable_web_search_tool_result(block: object) -> _ReplayedWebSearchToolResult | None:
|
||||
"""
|
||||
The parsed block when it is a ``web_search_tool_result`` carrying no
|
||||
``encrypted_content``, else None for anything Anthropic itself issued.
|
||||
|
||||
An empty ``content`` list is flattenable too. It is what the interceptor emits
|
||||
when a search legitimately returns nothing and when a search raises, and it
|
||||
carries neither evidence to preserve nor an ``encrypted_content`` to respect,
|
||||
so leaving it in place only buys the 400 this whole function exists to avoid.
|
||||
"""
|
||||
try:
|
||||
parsed: Final = _WEB_SEARCH_TOOL_RESULT_ADAPTER.validate_python(block)
|
||||
except ValidationError:
|
||||
return None
|
||||
if any(result.encrypted_content for result in parsed.content):
|
||||
return None
|
||||
return parsed
|
||||
|
||||
|
||||
def _replayed_server_tool_use(block: object) -> _ReplayedServerToolUse | None:
|
||||
try:
|
||||
return _SERVER_TOOL_USE_ADAPTER.validate_python(block)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _render_web_search_results(query: str, results: tuple[_ReplayedWebSearchResult, ...]) -> str:
|
||||
header: Final = f"Web search results for '{query}':" if query else "Web search results:"
|
||||
if not results:
|
||||
return f"{header}\n\nNo results were returned."
|
||||
body: Final = "\n\n".join(
|
||||
"\n".join(
|
||||
line
|
||||
for line in (
|
||||
f"Title: {result.title}" if result.title else "",
|
||||
f"URL: {result.url}" if result.url else "",
|
||||
f"Snippet: {result.snippet}" if result.snippet else "",
|
||||
)
|
||||
if line
|
||||
)
|
||||
for result in results
|
||||
)
|
||||
return f"{header}\n\n{body}" if body else header
|
||||
|
||||
|
||||
def _rewrite_replayed_web_search_block(
|
||||
block: object,
|
||||
flattenable: Mapping[str, _ReplayedWebSearchToolResult],
|
||||
queries: Mapping[str, str],
|
||||
) -> object | None:
|
||||
parsed_result: Final = _flattenable_web_search_tool_result(block)
|
||||
if parsed_result is not None:
|
||||
return _TextBlock(
|
||||
text=_render_web_search_results(queries.get(parsed_result.tool_use_id, ""), parsed_result.content)
|
||||
).model_dump()
|
||||
parsed_use: Final = _replayed_server_tool_use(block)
|
||||
if parsed_use is not None and parsed_use.id in flattenable:
|
||||
return None
|
||||
return block
|
||||
|
||||
|
||||
def _flatten_web_search_results_in_message(message: object) -> object:
|
||||
if not isinstance(message, Mapping) or not isinstance(message.get("content"), Sequence):
|
||||
return message
|
||||
content: Final = message["content"]
|
||||
if isinstance(content, str):
|
||||
return message
|
||||
flattenable: Final = MappingProxyType(
|
||||
{
|
||||
parsed.tool_use_id: parsed
|
||||
for parsed in (_flattenable_web_search_tool_result(block) for block in content)
|
||||
if parsed is not None
|
||||
}
|
||||
)
|
||||
if not flattenable:
|
||||
return message
|
||||
queries: Final = MappingProxyType(
|
||||
{
|
||||
parsed.id: parsed.input.query
|
||||
for parsed in (_replayed_server_tool_use(block) for block in content)
|
||||
if parsed is not None
|
||||
}
|
||||
)
|
||||
rewritten: Final = tuple(_rewrite_replayed_web_search_block(block, flattenable, queries) for block in content)
|
||||
return {**message, "content": [b for b in rewritten if b is not None]} # mutable-ok: JSON wire format
|
||||
|
||||
|
||||
def flatten_unencrypted_web_search_results_in_anthropic_messages( # mutable-ok: as sibling sanitizers
|
||||
messages: list[Any],
|
||||
) -> list[Any]:
|
||||
"""
|
||||
Return a new message list with replayed ``web_search_tool_result`` blocks that
|
||||
carry no ``encrypted_content`` rewritten into plain ``text`` blocks holding the
|
||||
same title / url / snippet evidence.
|
||||
|
||||
``encrypted_content`` is an opaque blob only Anthropic's own search backend can
|
||||
mint, so blocks synthesized by LiteLLM (websearch interception against a search
|
||||
provider) are rejected with ``Invalid encrypted_content in search_result block``
|
||||
when a native client loops them back as history. Flattening them keeps the
|
||||
evidence in the conversation instead of 400ing the follow-up turn, and leaves
|
||||
genuine Anthropic-issued blocks untouched.
|
||||
"""
|
||||
return [_flatten_web_search_results_in_message(m) for m in messages] # mutable-ok: JSON wire format
|
||||
|
||||
|
||||
def process_anthropic_headers(headers: httpx.Headers | dict) -> dict:
|
||||
openai_headers: Final = {}
|
||||
if "anthropic-ratelimit-requests-limit" in headers:
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue