From ecd5ad49f709e549e2f28e9bc71d4f49e3efb1d0 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 19:24:46 +0000 Subject: [PATCH 01/10] ci: always run the UI API types sync check so it can be required Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/check-ui-api-types.yml | 42 +++++++++++++++++++----- 1 file changed, 34 insertions(+), 8 deletions(-) diff --git a/.github/workflows/check-ui-api-types.yml b/.github/workflows/check-ui-api-types.yml index 597daebd720..587d6595bcc 100644 --- a/.github/workflows/check-ui-api-types.yml +++ b/.github/workflows/check-ui-api-types.yml @@ -2,17 +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 + pull-requests: read + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true jobs: check-sync: @@ -25,17 +27,35 @@ jobs: with: persist-credentials: false + - name: Detect changes that can affect the generated types + id: changes + env: + GH_TOKEN: ${{ github.token }} + REPO: ${{ github.repository }} + PR_NUMBER: ${{ github.event.pull_request.number }} + run: | + set -euo pipefail + files="$(gh api --paginate "repos/$REPO/pulls/$PR_NUMBER/files" --jq '.[].filename')" + if printf '%s\n' "$files" | 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$)'; then + echo "relevant=true" >> "$GITHUB_OUTPUT" + else + 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 +66,17 @@ 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: Generate Prisma client + if: steps.changes.outputs.relevant == 'true' env: PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache 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 +84,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." From bc71564d1d4132005f9c7fa62279b88b8dd43ca1 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 19:58:22 +0000 Subject: [PATCH 02/10] build(deps): defer the second pypdf advisory until the 6.15.0 bump GHSA-fp3f-mc75-235c was published at 19:29 UTC today, hours after #36212 landed, and covers a different pypdf issue (large memory use on big /ToUnicode streams) that lands on the same 6.15.0 fix. It reds osv-scan on every PR into staging again. pypdf 6.15.0 is still inside the repo's P3D exclude-newer window until 08-09, so uv cannot lock it yet and widening exclude-newer to pull it in early would weaken the freshness guard for every package. This gets the same 08-12 ignoreUntil as the first pypdf entry so both drop together in the bump. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- osv-scanner.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/osv-scanner.toml b/osv-scanner.toml index efcbe6c8c16..4ef612e3a70 100644 --- a/osv-scanner.toml +++ b/osv-scanner.toml @@ -3,6 +3,11 @@ id = "GHSA-fwg2-594c-jp42" ignoreUntil = 2026-08-12 reason = "pypdf 6.15.0 (the fix) published 2026-08-06 and is still inside the P3D exclude-newer window, so uv cannot lock it yet; bump and drop this entry from 2026-08-09" +[[IgnoredVulns]] +id = "GHSA-fp3f-mc75-235c" +ignoreUntil = 2026-08-12 +reason = "second pypdf advisory with the same 6.15.0 fix, published 2026-08-07 after the first; drop alongside GHSA-fwg2-594c-jp42 in the same bump" + [[IgnoredVulns]] id = "GHSA-w8v5-vhqr-4h9v" ignoreUntil = 2026-09-09 From 4306c4ad70e4f50f8415d5e47e13884e0152596a Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 21:07:27 +0000 Subject: [PATCH 03/10] ci: detect relevant changes from git instead of the paginated files API Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/check-ui-api-types.yml | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/.github/workflows/check-ui-api-types.yml b/.github/workflows/check-ui-api-types.yml index 587d6595bcc..02543d67a82 100644 --- a/.github/workflows/check-ui-api-types.yml +++ b/.github/workflows/check-ui-api-types.yml @@ -10,7 +10,6 @@ on: permissions: contents: read - pull-requests: read concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} @@ -26,19 +25,22 @@ 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 - env: - GH_TOKEN: ${{ github.token }} - REPO: ${{ github.repository }} - PR_NUMBER: ${{ github.event.pull_request.number }} run: | set -euo pipefail - files="$(gh api --paginate "repos/$REPO/pulls/$PR_NUMBER/files" --jq '.[].filename')" - if printf '%s\n' "$files" | 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$)'; then + 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 From 860e37597fd20e2fe62c952d674c8b5c0ebddae9 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 19:56:29 +0000 Subject: [PATCH 04/10] fix(a2a): align agent list annotation and test with the tuple return type PR #36020 changed AgentRegistry.get_agent_list to return tuple[AgentResponse, ...], and PR #35163 added a test asserting the result equals []. Both were green on their own branches and only collided once they were both on litellm_internal_staging, so proxy-server has been failing on every PR since with 'assert () == []'. Nothing user-facing was wrong: get_agents only iterates the result and rebuilds it with comprehensions, and FastAPI serializes a tuple to the same JSON array. The test expectation was simply stale, so it now compares against (). The get_agents local was still annotated list[AgentResponse] while two branches assign the registry tuple straight into it, so it widens to Sequence[AgentResponse]. That covers both the tuple and the list branches without pretending the value is mutable. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/agent_endpoints/endpoints.py | 2 +- tests/test_litellm/proxy/proxy_server/test_proxy_config.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 5bc5b5566f8..d348bc01153 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -253,7 +253,7 @@ async def get_agents( ) try: - returned_agents: list[AgentResponse] = [] + returned_agents: Sequence[AgentResponse] = () # Admin users get all agents if ( diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 30156849628..91a7e1bc2c2 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2638,4 +2638,4 @@ async def test_ProxyConfig__init_non_llm_configs_empty_agents_key_clears_remembe assert clean_agent_registry.config_agents == () clean_agent_registry.load_agents_from_db_and_config(db_agents=None) - assert clean_agent_registry.get_agent_list() == [] + assert clean_agent_registry.get_agent_list() == () From 2f70f4edbaf7aa042bac1b097bae5199774ad56b Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 7 Aug 2026 21:30:09 +0000 Subject: [PATCH 05/10] build(deps): bump nanoid to 3.3.17 in the dashboard lockfile GHSA-2v37-7h3g-55p8 (CVE-2026-67213) rates 8.2 against nanoid 3.3.16 and reds osv-scan on every PR into staging. Custom generators can loop indefinitely when size is zero, so a caller that passes through a zero size hangs the process. nanoid is transitive through the dashboard's toolchain and 3.3.17 is a patch release that published 08-03, so it already clears the .npmrc min-release-age=3 guard. The lock diff is the version, resolved url, and integrity hash for that one package. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/package-lock.json | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 34dea2a39f5..515a992bc85 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -10319,9 +10319,9 @@ "license": "MIT" }, "node_modules/nanoid": { - "version": "3.3.16", - "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.16.tgz", - "integrity": "sha512-bzlKTyNJ7+LdGIIwy8ijFpIqEQIvafahV7eYykJ8Cvh42EdJeODoJ6gUJXpQJvej1BddH8OqTXZNE/KfbWAu8Q==", + "version": "3.3.17", + "resolved": "https://registry.npmjs.org/nanoid/-/nanoid-3.3.17.tgz", + "integrity": "sha512-xQLf0A3HOMlgHq0n247/LRuAOYmB7dXJ/DvAxGvsSBij45XtBSmQycu+F8ODbHwns/XyFZagyL1+J0Offw1E0g==", "funding": [ { "type": "github", From 8db2fbaad0c064b61e8f33d0c21c06d36e63d05d Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 7 Aug 2026 15:53:59 -0700 Subject: [PATCH 06/10] feat(ui): show user email or alias in usage data export (#36232) * feat(ui): resolve user email/alias in usage export instead of raw user id * test(ui): cover email/alias resolution in usage export data builders * chore(ui): drop explanatory comment per repo comment policy --- .../EntityUsageExport/utils.test.ts | 172 ++++++++++++++++++ .../src/components/EntityUsageExport/utils.ts | 26 ++- 2 files changed, 189 insertions(+), 9 deletions(-) diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts index e91b7b73a1e..2f551176091 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.test.ts @@ -2197,4 +2197,176 @@ describe("EntityUsageExport utils", () => { }); }); }); + + describe("display name resolution from entity metadata", () => { + const entityMetrics = { + spend: 12.25, + api_requests: 40, + successful_requests: 39, + failed_requests: 1, + total_tokens: 900, + prompt_tokens: 500, + completion_tokens: 400, + cache_read_input_tokens: 20, + cache_creation_input_tokens: 10, + }; + + const makeSpendData = (entity: string, metadata?: Record): EntitySpendData => ({ + results: [ + { + date: "2025-04-01", + breakdown: { + entities: { + [entity]: { + metrics: entityMetrics, + metadata, + api_key_breakdown: { + key1: { + metrics: entityMetrics, + metadata: { key_alias: "prod-key" }, + }, + }, + }, + }, + }, + }, + ], + metadata: mockSpendData.metadata, + }); + + it("should export the user email as the entity label and keep the raw user id in the id column", () => { + const result = generateDailyData( + makeSpendData("user-123", { user_email: "ada@example.com", user_alias: "Ada" }), + "User", + ); + + expect(result).toHaveLength(1); + expect(result[0]["User"]).toBe("ada@example.com"); + expect(result[0]["User ID"]).toBe("user-123"); + }); + + it("should fall back to the user alias when the user has no email", () => { + const nullEmail = generateDailyData( + makeSpendData("user-123", { user_email: null, user_alias: "Ada Lovelace" }), + "User", + ); + const missingEmail = generateDailyData(makeSpendData("user-123", { user_alias: "Ada Lovelace" }), "User"); + + expect(nullEmail[0]["User"]).toBe("Ada Lovelace"); + expect(missingEmail[0]["User"]).toBe("Ada Lovelace"); + }); + + it("should fall back to the raw entity key when the entity carries no metadata", () => { + const noMetadata = generateDailyData(makeSpendData("my-tag"), "Tag"); + const emptyMetadata = generateDailyData(makeSpendData("customer-9", {}), "Customer"); + const blankNames = generateDailyData(makeSpendData("user-123", { user_email: null, user_alias: null }), "User"); + + expect(noMetadata[0]["Tag"]).toBe("my-tag"); + expect(emptyMetadata[0]["Customer"]).toBe("customer-9"); + expect(blankNames[0]["User"]).toBe("user-123"); + }); + + it("should prefer the team alias map over any alias in entity metadata", () => { + const result = generateDailyData( + makeSpendData("team-1", { team_alias: "Stale Alias", user_email: "ada@example.com" }), + "Team", + mockTeamAliasMap, + ); + + expect(result[0]["Team"]).toBe("Team One"); + }); + + it("should use the team alias from entity metadata when the alias map has no entry for the team", () => { + const result = generateDailyData( + makeSpendData("team-9", { team_alias: "Team Nine", user_email: "ada@example.com" }), + "Team", + mockTeamAliasMap, + ); + + expect(result[0]["Team"]).toBe("Team Nine"); + }); + + it("should resolve metadata.alias to the user email in getEntityBreakdown", () => { + const withEmail = getEntityBreakdown( + makeSpendData("user-123", { user_email: "ada@example.com", user_alias: "Ada" }), + ); + const withoutEmail = getEntityBreakdown(makeSpendData("user-123", { user_alias: "Ada" })); + + expect(withEmail[0].metadata.alias).toBe("ada@example.com"); + expect(withEmail[0].metadata.id).toBe("user-123"); + expect(withoutEmail[0].metadata.alias).toBe("Ada"); + }); + + it("should resolve the user email on every key row of the keys scope", () => { + const spendData: EntitySpendData = { + results: [ + { + date: "2025-04-01", + breakdown: { + entities: { + "user-123": { + metrics: entityMetrics, + metadata: { user_email: "ada@example.com", user_alias: "Ada" }, + api_key_breakdown: { + key1: { metrics: entityMetrics, metadata: { key_alias: "prod-key" } }, + key2: { metrics: entityMetrics, metadata: { key_alias: "dev-key" } }, + }, + }, + }, + }, + }, + ], + metadata: mockSpendData.metadata, + }; + + const result = generateDailyWithKeysData(spendData, "User"); + + expect(result).toHaveLength(2); + expect(result.map((r) => r["User"])).toEqual(["ada@example.com", "ada@example.com"]); + expect(result.map((r) => r["User ID"])).toEqual(["user-123", "user-123"]); + expect(result.find((r) => r["Key ID"] === "key1")?.["Key Alias"]).toBe("prod-key"); + expect(result.find((r) => r["Key ID"] === "key2")?.["Key Alias"]).toBe("dev-key"); + }); + + it("should resolve each entity's own email in the models scope", () => { + const spendData: EntitySpendData = { + results: [ + { + date: "2025-04-01", + breakdown: { + entities: { + "user-a": { + metrics: entityMetrics, + metadata: { user_email: "ada@example.com", user_alias: "Ada" }, + api_key_breakdown: { key1: { metrics: entityMetrics, metadata: {} } }, + }, + "user-b": { + metrics: entityMetrics, + metadata: { user_email: null, user_alias: "Grace" }, + api_key_breakdown: { key2: { metrics: entityMetrics, metadata: {} } }, + }, + }, + models: { + "claude-sonnet-4-5": { + metrics: entityMetrics, + api_key_breakdown: { + key1: { metrics: entityMetrics, metadata: {} }, + key2: { metrics: entityMetrics, metadata: {} }, + }, + }, + }, + }, + }, + ], + metadata: mockSpendData.metadata, + }; + + const result = generateDailyWithModelsData(spendData, "User"); + + expect(result).toHaveLength(2); + expect(result.every((r) => r.Model === "claude-sonnet-4-5")).toBe(true); + expect(result.find((r) => r["User ID"] === "user-a")?.["User"]).toBe("ada@example.com"); + expect(result.find((r) => r["User ID"] === "user-b")?.["User"]).toBe("Grace"); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts index 8ff389ded5e..53d100040d7 100644 --- a/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts +++ b/ui/litellm-dashboard/src/components/EntityUsageExport/utils.ts @@ -3,12 +3,18 @@ import type { DateRangePickerValue } from "@tremor/react"; import Papa from "papaparse"; import type { EntityBreakdown, EntitySpendData, EntityType, ExportMetadata, ExportScope } from "./types"; -// Resolve display name for an entity. For teams the teamAliasMap provides -// a human-readable alias; for every other entity type the entity key itself -// (tag name, org id, customer id, …) is already the correct label. -const resolveEntityDisplay = (entity: string, teamAliasMap: Record): { id: string; alias: string } => ({ +const resolveEntityDisplay = ( + entity: string, + teamAliasMap: Record, + entityMetadata?: Record, +): { id: string; alias: string } => ({ id: entity, - alias: teamAliasMap[entity] || entity, + alias: + teamAliasMap[entity] || + entityMetadata?.team_alias || + entityMetadata?.user_email || + entityMetadata?.user_alias || + entity, }); // Mirrors backend SpendMetrics fields (litellm/types/activity_tracking.py). @@ -68,7 +74,7 @@ export const getEntityBreakdown = ( spendData.results.forEach((day) => { Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { - const { id, alias } = resolveEntityDisplay(entity, teamAliasMap); + const { id, alias } = resolveEntityDisplay(entity, teamAliasMap, data.metadata); if (!entitySpend[entity]) { entitySpend[entity] = { @@ -113,7 +119,7 @@ export const generateDailyData = ( spendData.results.forEach((day) => { Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { - const { id, alias } = resolveEntityDisplay(entity, teamAliasMap); + const { id, alias } = resolveEntityDisplay(entity, teamAliasMap, data.metadata); dailyBreakdown.push({ Date: day.date, @@ -164,7 +170,7 @@ export const generateDailyWithKeysData = ( spendData.results.forEach((day) => { Object.entries(resolveEntities(day.breakdown)).forEach(([entity, data]: [string, any]) => { - const { id: entityId, alias: entityAlias } = resolveEntityDisplay(entity, teamAliasMap); + const { id: entityId, alias: entityAlias } = resolveEntityDisplay(entity, teamAliasMap, data.metadata); const apiKeyBreakdown = data.api_key_breakdown || {}; // Iterate through each API key in the breakdown @@ -241,11 +247,13 @@ export const generateDailyWithModelsData = ( spendData.results.forEach((day) => { const dailyEntityModels: { [key: string]: { [key: string]: any } } = {}; + const dailyEntityMetadata: { [key: string]: Record | undefined } = {}; Object.entries(resolveEntities(day.breakdown)).forEach(([entity, entityData]: [string, any]) => { if (!dailyEntityModels[entity]) { dailyEntityModels[entity] = {}; } + dailyEntityMetadata[entity] = entityData.metadata; Object.entries(day.breakdown.models || {}).forEach(([model, modelData]: [string, any]) => { const entityApiKeys = entityData.api_key_breakdown || {}; @@ -282,7 +290,7 @@ export const generateDailyWithModelsData = ( }); Object.entries(dailyEntityModels).forEach(([entity, models]) => { - const { id, alias } = resolveEntityDisplay(entity, teamAliasMap); + const { id, alias } = resolveEntityDisplay(entity, teamAliasMap, dailyEntityMetadata[entity]); Object.entries(models).forEach(([model, metrics]: [string, any]) => { dailyModelBreakdown.push({ From 3238ce840601c820f31b01f1731b704a6d3f8f2e Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 7 Aug 2026 17:03:34 -0700 Subject: [PATCH 07/10] feat(auto-router): track turns per complexity tier (LIT-5302) (#36209) * feat(auto-router): track turns per complexity tier (LIT-5302) Stamps complexity tier at decision time (rollup never re-derives from routed model, since tier->model mapping is mutable config). Records per-tier turn counts in LiteLLM_AutoRouterSession.tier_turns (jsonb), rolls up per router in benchmarks SQL via jsonb_object_agg, returns on AutoRouterBenchmarkGroup for dashboard turns/share metrics. Addresses Greptile/Bugbot findings: - Missing _SessionAggRow.tier_turns field: added with field_validator to parse jsonb text cast and handle NULL. Would 500 every benchmarks read. - Missing ::text cast on tier parameter: Postgres fails type inference on parameterized CASE/IS NULL without explicit cast. Added to all usages. - Docstring false claim (only complexity routers produce tiers): quality router stamps numeric tier '1'/'2'/'3'. Per-type grouping in SQL prevents cross-contamination. Rewrote docstring to clarify isolation. - Comment convention violations: stripped per CLAUDE.md rule. - Test gaps: 8 unit tests for extraction/validation/aggregation, 7 behavior tests for SQL semantics against real Postgres. 12 mutations killed. Fixed fragile complexity_router test that broke on nested function calls. No API change; extends existing GET /auto_router/benchmarks response only. Co-Authored-By: Claude * fix(auto-router): address review findings on tier turns tracking - Guard router_type update so a mid-session reconfigure can't pool foreign tier names into tier_turns - Keep pinned turns attributed to the tier that actually serves them - Drop stray -- AlterTable comment from hand-written migration - Drop the now-unnecessary ::text/json.loads round-trip; prisma already returns tier_turns as a parsed dict Co-Authored-By: Claude * fix(auto-router): satisfy type-discipline lint gate - tier_turns fields: dict[str, int] -> Mapping[str, int] (LIT001, mutable collection in annotation); these are read-only after construction - _summed_agg_row: {} -> MappingProxyType({}) (LIT002, mutable dict literal) - default-fallback branch: replace the reassigned-without-Final fallback_tier with a Final default_model_first flag and a single ternary assignment (LIT010) Verified locally: type_discipline_gate.py, ruff_strict_gate.py, and type_check_gate.py all pass against the litellm_internal_staging merge-base; full test_complexity_router.py (374), auto_router management-endpoint tests (26), db-layer rollup tests (31), and the live-Postgres proxy_behavior rollup suite (17) all pass. Co-Authored-By: Claude --------- Co-authored-by: Claude --- .../migration.sql | 1 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/proxy/db/autorouter_session_rollup.py | 13 ++- .../auto_router_endpoints.py | 30 +++++- litellm/proxy/schema.prisma | 1 + .../complexity_router/complexity_router.py | 6 +- .../auto_router_endpoints.py | 11 ++ schema.prisma | 1 + .../spend/test_autorouter_session_rollup.py | 102 ++++++++++++++++++ .../db/test_autorouter_session_rollup.py | 22 +++- .../test_auto_router_endpoints.py | 39 +++++++ .../router_strategy/test_complexity_router.py | 56 ++++++++-- 12 files changed, 268 insertions(+), 15 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql new file mode 100644 index 00000000000..81b1cbc7ec3 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260807000000_add_autorouter_session_tier_turns/migration.sql @@ -0,0 +1 @@ +ALTER TABLE "LiteLLM_AutoRouterSession" ADD COLUMN IF NOT EXISTS "tier_turns" JSONB NOT NULL DEFAULT '{}'; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index b6557e3006d..9c871b65f40 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -1439,6 +1439,7 @@ 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") diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index da1652cdb61..b1f074c26b7 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -49,6 +49,7 @@ class AutoRouterTurnTransaction: cache_hit: bool cache_ttl_seconds: int | None cache_touched: bool + tier: str | None = None class TurnCacheFacts(NamedTuple): @@ -152,11 +153,13 @@ def build_autorouter_turn_transaction( return None usage_object_raw: Final = metadata.get("usage_object") cache: Final = turn_cache_facts(usage_object_raw if isinstance(usage_object_raw, Mapping) else None) + tier_raw: Final = routing_decision.get("tier") return AutoRouterTurnTransaction( api_key=api_key, session_id=_bounded_session_id(session_id), router_name=router_name, router_type=str(routing_decision.get("router_type") or "unknown"), + tier=tier_raw if isinstance(tier_raw, str) and tier_raw else None, model=model, turn_at=turn_at, total_tokens=int(payload.get("prompt_tokens") or 0) + int(payload.get("completion_tokens") or 0), @@ -184,6 +187,8 @@ _COVERED: Final = _p("covered") _CACHE_HIT: Final = _p("cache_hit") _CACHE_TTL: Final = _p("cache_ttl_seconds") _TOUCHED: Final = _p("cache_touched") +_TIER: Final = f"{_p('tier')}::text" +_TIER_DELTA: Final = f"(CASE WHEN {_TIER} IS NULL THEN '{{}}'::jsonb ELSE jsonb_build_object({_TIER}, 1) END)" _IN_ORDER: Final = f"{_TURN_AT}::timestamp >= t.last_turn_at" _SAME: Final = f"{_IN_ORDER} AND t.last_model = {_MODEL}" @@ -201,7 +206,7 @@ INSERT INTO "LiteLLM_AutoRouterSession" AS t ( last_model, models, turns, unordered_turns, covered_turns, cache_hits, same_model_turns, same_model_hits, first_visit_turns, first_visit_hits, return_turns, return_hits, return_expired_misses, return_within_ttl_misses, - ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend + ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, tier_turns ) VALUES ( {_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp, @@ -211,7 +216,8 @@ VALUES ( 0, 0, 0, 0, (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_5M_SECONDS} THEN 1 ELSE 0 END), (CASE WHEN {_CACHE_TTL}::int = {CACHE_TTL_1H_SECONDS} THEN 1 ELSE 0 END), - {_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8 + {_p("total_tokens")}::bigint, {_p("spend")}::float8, {_p("saved_spend")}::float8, + {_TIER_DELTA} ) ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET turns = t.turns + 1, @@ -242,6 +248,9 @@ ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET ELSE COALESCE((t.models -> {_MODEL} ->> 'ttl')::int, {_CACHE_TTL}::int) END) )), last_model = (CASE WHEN {_IN_ORDER} THEN {_MODEL} ELSE t.last_model END), + tier_turns = (CASE WHEN {_TIER} IS NOT NULL AND t.router_type = {_p("router_type")} + THEN t.tier_turns || jsonb_build_object({_TIER}, COALESCE((t.tier_turns ->> {_TIER})::int, 0) + 1) + ELSE t.tier_turns END), first_turn_at = LEAST(t.first_turn_at, EXCLUDED.first_turn_at), last_turn_at = GREATEST(t.last_turn_at, EXCLUDED.last_turn_at) """ diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index 141094f4d4c..d4221845b0c 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -4,8 +4,9 @@ AUTO ROUTER MANAGEMENT ENDPOINTS POST /auto_router/test_routing - Route one prompt through an unsaved complexity-router config """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final from pydantic import BaseModel, TypeAdapter @@ -260,6 +261,7 @@ async def preview_auto_router_routing( class _SessionAggRow(BaseModel): router_name: str router_type: str + tier_turns: Mapping[str, int] sessions: int turns: int unordered_turns: int @@ -284,6 +286,23 @@ class _SessionAggRow(BaseModel): _SESSION_AGG_ROWS: Final = TypeAdapter(list[_SessionAggRow]) _BENCHMARKS_SQL: Final = """ +WITH windowed AS ( + SELECT * FROM "LiteLLM_AutoRouterSession" + WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp +), +tier_maps AS ( + SELECT router_name, router_type, jsonb_object_agg(tier, tier_turns) AS tier_turns + FROM ( + SELECT router_name, router_type, kv.key AS tier, SUM((kv.value)::int)::int AS tier_turns + FROM windowed, LATERAL jsonb_each_text(tier_turns) AS kv + GROUP BY router_name, router_type, kv.key + ) per_tier + GROUP BY router_name, router_type +) +SELECT + agg.*, + COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns +FROM ( SELECT router_name, router_type, @@ -306,10 +325,11 @@ SELECT COALESCE(SUM(spend), 0)::float8 AS spend, COALESCE(SUM(saved_spend), 0)::float8 AS saved_spend, COALESCE(SUM(EXTRACT(EPOCH FROM (last_turn_at - first_turn_at))), 0)::float8 AS session_seconds -FROM "LiteLLM_AutoRouterSession" -WHERE last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp +FROM windowed GROUP BY router_name, router_type -ORDER BY SUM(spend) DESC +) agg +LEFT JOIN tier_maps USING (router_name, router_type) +ORDER BY agg.spend DESC """ @@ -366,6 +386,7 @@ def _summed_agg_row(rows: Sequence[_SessionAggRow]) -> _SessionAggRow: return _SessionAggRow( router_name="", router_type="", + tier_turns=MappingProxyType({}), sessions=sum(row.sessions for row in rows), turns=sum(row.turns for row in rows), unordered_turns=sum(row.unordered_turns for row in rows), @@ -443,6 +464,7 @@ async def get_auto_router_benchmarks( AutoRouterBenchmarkGroup( router_name=row.router_name, router_type=row.router_type, + tier_turns=row.tier_turns, **_benchmark_totals(row).model_dump(), ) for row in rows diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index b6557e3006d..9c871b65f40 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1439,6 +1439,7 @@ 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") diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index a69509fc37a..f6ced0bb9d2 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -1730,6 +1730,7 @@ class ComplexityRouter(CustomLogger): routing_decision=self._build_routing_decision( routed_model=routed_model, cause=cause, + tier=self._tier_for_model(routed_model), escalation_keyword=pin_escalation_keyword, escalated=escalated, conversation_continuing=conversation_continuing, @@ -1797,7 +1798,8 @@ class ComplexityRouter(CustomLogger): if user_message is None: verbose_router_logger.debug("ComplexityRouter: No user message found, routing to default model") - if not self.config.plugins and self.config.default_model: + default_model_first: Final = not self.config.plugins and self.config.default_model + if default_model_first: # No plugins configured: preserve the pre-existing default_model-first # priority exactly (changing it would be a silent behavior change for # every non-plugin user, not just a security fix). @@ -1809,12 +1811,14 @@ class ComplexityRouter(CustomLogger): routed_model = await self._pick_model_for_tier( ComplexityTier.MEDIUM, messages, resolved_messages, request_kwargs ) + fallback_tier: Final = None if default_model_first else ComplexityTier.MEDIUM return PreRoutingHookResponse( model=routed_model, messages=messages if has_original_messages else None, routing_decision=self._build_routing_decision( routed_model=routed_model, cause="default_fallback", + tier=fallback_tier, conversation_continuing=conversation_continuing, ), ) diff --git a/litellm/types/management_endpoints/auto_router_endpoints.py b/litellm/types/management_endpoints/auto_router_endpoints.py index 6c8fb96a729..6626dea6849 100644 --- a/litellm/types/management_endpoints/auto_router_endpoints.py +++ b/litellm/types/management_endpoints/auto_router_endpoints.py @@ -2,6 +2,7 @@ Types for auto-router management endpoints """ +from collections.abc import Mapping from typing import Final from pydantic import BaseModel, Field, field_validator @@ -120,6 +121,16 @@ class AutoRouterBenchmarkGroup(AutoRouterBenchmarkTotals): router_name: str = Field(description="The auto-router alias requests were sent to") router_type: str = Field(description="complexity, adaptive or quality") + tier_turns: Mapping[str, int] = Field( + default_factory=dict, + description="Turns per tier, keyed by the tier name the routing decision recorded at " + "request time (never re-derived at read time, since the tier-to-model mapping is " + "mutable config). Tier names are scoped to this group's router_type and are not " + "comparable across types: a complexity router reports 'simple'/'medium'/'complex'/" + "'reasoning', a quality router reports its numeric quality tier, and an adaptive router " + "records no tier at all. Turns no tier served (the classifier fell back to default_model) " + "are absent rather than pooled under a sentinel key, so the values may sum to less than turns", + ) class AutoRouterBenchmarksResponse(BaseModel): diff --git a/schema.prisma b/schema.prisma index b6557e3006d..9c871b65f40 100644 --- a/schema.prisma +++ b/schema.prisma @@ -1439,6 +1439,7 @@ 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") diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py index aa734ee22cc..65b70f13a3b 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py @@ -38,11 +38,13 @@ async def _turn( tokens: int = 100, spend: float = 0.01, saved: float = 0.02, + tier: "str | None" = None, ) -> None: touched: Final = 1 if (hit or ttl is not None or not covered) else 0 await db.execute_raw( UPSERT_AUTOROUTER_SESSION_SQL, key, session_id, router, router_type, model, at.isoformat(), tokens, spend, saved, covered, hit, ttl, touched, + tier, ) @@ -195,6 +197,106 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db assert [(row["router_type"], row["sessions"]) for row in matching] == [("complexity", 1), ("quality", 1)] +async def test_tier_turns_count_each_tier_that_served_a_turn(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0, tier="simple") + await _turn(db, key, "B", T0 + timedelta(seconds=10), tier="complex") + await _turn(db, key, "A", T0 + timedelta(seconds=20), tier="simple") + + assert (await _row(db, key))["tier_turns"] == {"simple": 2, "complex": 1} + + +async def test_an_untiered_turn_increments_no_tier_counter(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0, tier=None) + assert (await _row(db, key))["tier_turns"] == {} + + await _turn(db, key, "A", T0 + timedelta(seconds=10), tier="medium") + await _turn(db, key, "A", T0 + timedelta(seconds=20), tier=None) + row = await _row(db, key) + assert row["tier_turns"] == {"medium": 1} + assert row["turns"] == 3 + + +async def test_a_mid_session_router_type_change_keeps_foreign_tier_names_out_of_the_map(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0, router_type="complexity", tier="medium") + await _turn(db, key, "A", T0 + timedelta(seconds=10), router_type="quality", tier="2") + await _turn(db, key, "A", T0 + timedelta(seconds=20), router_type="complexity", tier="medium") + + row = await _row(db, key) + assert row["tier_turns"] == {"medium": 2} + assert row["turns"] == 3 + + +async def test_an_out_of_order_turn_still_counts_toward_its_tier(db): + key = f"k-{uuid.uuid4()}" + await _turn(db, key, "A", T0 + timedelta(seconds=60), tier="simple") + await _turn(db, key, "A", T0, tier="simple") + + row = await _row(db, key) + assert row["tier_turns"] == {"simple": 2} + assert row["unordered_turns"] == 1 + + +async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db): + key = f"k-{uuid.uuid4()}" + router = f"r-{uuid.uuid4()}" + await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier="simple") + await _turn(db, key, "A", T0 + timedelta(seconds=10), session_id=f"s-{uuid.uuid4()}", router=router, tier="simple") + await _turn(db, key, "B", T0 + timedelta(seconds=20), session_id=f"s-{uuid.uuid4()}", router=router, tier="complex") + await _turn(db, key, "C", T0 + timedelta(seconds=30), session_id=f"s-{uuid.uuid4()}", router=router, tier=None) + + rows = await db.query_raw( + _BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + ) + grouped = next(row for row in rows if row["router_name"] == router) + assert grouped["tier_turns"] == {"simple": 2, "complex": 1} + assert grouped["turns"] == 4 + + +async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(db): + key = f"k-{uuid.uuid4()}" + router = f"r-{uuid.uuid4()}" + await _turn( + db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, router_type="complexity", tier="medium" + ) + await _turn( + db, + key, + "A", + T0 + timedelta(seconds=10), + session_id=f"s-{uuid.uuid4()}", + router=router, + router_type="quality", + tier="2", + ) + + rows = await db.query_raw( + _BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + ) + by_type = {row["router_type"]: row["tier_turns"] for row in rows if row["router_name"] == router} + assert by_type == {"complexity": {"medium": 1}, "quality": {"2": 1}} + + +async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db): + key = f"k-{uuid.uuid4()}" + router = f"r-{uuid.uuid4()}" + await _turn(db, key, "A", T0, session_id=f"s-{uuid.uuid4()}", router=router, tier=None) + + rows = await db.query_raw( + _BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + ) + grouped = next(row for row in rows if row["router_name"] == router) + assert grouped["tier_turns"] == {} + + async def test_a_miss_that_touched_no_cache_does_not_advance_the_ttl_clock(db): key = f"k-{uuid.uuid4()}" await _turn(db, key, "A", T0, ttl=300) diff --git a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py index aa7d01bc880..0df11f224a2 100644 --- a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py +++ b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py @@ -93,6 +93,19 @@ class TestBuildTransaction: def test_requests_without_a_routing_decision_are_skipped(self, metadata: dict): assert _build(metadata=metadata) is None + def test_the_tier_the_decision_recorded_is_carried_onto_the_transaction(self): + transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "tier": "reasoning"})) + assert transaction is not None and transaction.tier == "reasoning" + + @pytest.mark.parametrize("tier", [None, "", 3, {"tier": "medium"}]) + def test_a_decision_without_a_usable_tier_records_no_tier(self, tier: object): + transaction = _build(metadata=_metadata(routing_decision={**ROUTING_DECISION, "tier": tier})) + assert transaction is not None and transaction.tier is None + + def test_a_decision_that_never_mentions_tier_records_no_tier(self): + transaction = _build() + assert transaction is not None and transaction.tier is None + def test_router_name_falls_back_to_the_payload_model_group(self): transaction = _build(metadata=_metadata(routing_decision={"router_type": "complexity"})) assert transaction is not None and transaction.router_name == "live-auto" @@ -167,7 +180,11 @@ class _FakeClient: self.db = _FakeDB(failures, poison_session) -def _transaction(session_id: str = "s1", at: datetime = datetime(2026, 8, 1, 12, 0, 0)) -> AutoRouterTurnTransaction: +def _transaction( + session_id: str = "s1", + at: datetime = datetime(2026, 8, 1, 12, 0, 0), + tier: str | None = "medium", +) -> AutoRouterTurnTransaction: return AutoRouterTurnTransaction( api_key="k1", session_id=session_id, @@ -182,6 +199,7 @@ def _transaction(session_id: str = "s1", at: datetime = datetime(2026, 8, 1, 12, cache_hit=False, cache_ttl_seconds=None, cache_touched=False, + tier=tier, ) @@ -201,7 +219,7 @@ class TestFlush: assert sql == UPSERT_AUTOROUTER_SESSION_SQL assert params == ( "k1", "s1", "live-auto", "complexity", "bedrock/haiku", - "2026-08-01T12:00:00", 100, 0.01, 0.02, 1, 0, None, 0, + "2026-08-01T12:00:00", 100, 0.01, 0.02, 1, 0, None, 0, "medium", ) def test_a_connect_error_retries_the_same_statement(self): diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 888db031515..3a995e27697 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -292,6 +292,7 @@ class TestAutoRouterBenchmarks: ROW = _SessionAggRow( router_name="live-auto", router_type="complexity", + tier_turns={}, sessions=4, turns=40, unordered_turns=1, @@ -377,6 +378,21 @@ class TestAutoRouterBenchmarks: assert totals.avg_turns_per_session == 10.0 assert totals.spend == 10.0 + def test_tier_names_stay_scoped_to_the_router_type_that_recorded_them(self): + quality = self.ROW.model_copy( + update={"router_name": "quality-auto", "router_type": "quality", "tier_turns": {"2": 7}} + ) + complexity = self.ROW.model_copy(update={"tier_turns": {"medium": 7}}) + assert complexity.tier_turns == {"medium": 7} + assert quality.tier_turns == {"2": 7} + + def test_summed_totals_carry_no_tier_map_because_names_are_router_scoped(self): + from litellm.proxy.management_endpoints.auto_router_endpoints import _summed_agg_row + + quality = self.ROW.model_copy(update={"router_type": "quality", "tier_turns": {"2": 7}}) + complexity = self.ROW.model_copy(update={"tier_turns": {"medium": 7}}) + assert _summed_agg_row([complexity, quality]).tier_turns == {} + @pytest.mark.asyncio async def test_non_admin_roles_cannot_read_benchmarks(self): from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks @@ -427,3 +443,26 @@ class TestAutoRouterBenchmarks: assert response.routers_in_scope == 1 assert response.groups[0].router_name == "live-auto" assert response.groups[0].saved_pct == response.totals.saved_pct == 75.0 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "wire_value, expected", [({"simple": 24, "complex": 16}, {"simple": 24, "complex": 16}), ({}, {})] + ) + async def test_the_tier_map_reaches_the_response_as_the_jsonb_column_returns_it( + self, wire_value: dict, expected: dict, monkeypatch: pytest.MonkeyPatch + ): + from litellm.proxy import proxy_server + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks + + class _DB: + async def query_raw(self, sql: str, *params: object): + return [{**TestAutoRouterBenchmarks.ROW.model_dump(), "tier_turns": wire_value}] + + monkeypatch.setattr(proxy_server, "prisma_client", type("P", (), {"db": _DB()})()) + + response = await get_auto_router_benchmarks( + user_api_key_dict=ADMIN, + start_date="2026-07-01", + end_date="2026-08-01", + ) + assert response.groups[0].tier_turns == expected diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 8e9e32f5898..356556f3563 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -3477,6 +3477,23 @@ class TestSessionAffinity: # Pinned to the first turn's model, not re-classified down to SIMPLE. assert second.model == "o1-preview" + @pytest.mark.asyncio + async def test_a_pinned_turn_reports_the_tier_that_serves_it(self, mock_router_instance, session_affinity_config): + mock_router_instance.cache = DualCache() + router = ComplexityRouter( + model_name="test-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=session_affinity_config, + ) + request_kwargs = self._request_kwargs("session-1") + await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE + ) + pinned = await router.async_pre_routing_hook( + model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE + ) + assert pinned.routing_decision["tier"] == "REASONING" + @pytest.mark.asyncio async def test_different_sessions_classify_independently(self, mock_router_instance, session_affinity_config): mock_router_instance.cache = DualCache() @@ -4362,7 +4379,24 @@ class TestRoutingDecisionContents: assert decision is not None assert decision["cause"] == "default_fallback" assert decision["routed_model"] == response.model - assert "tier" not in decision + assert decision.get("tier") == "MEDIUM" + + @pytest.mark.asyncio + async def test_a_default_model_fallback_claims_no_tier(self, mock_router_instance, basic_config): + router = ComplexityRouter( + model_name="test-complexity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={**basic_config, "default_model": "gpt-4o"}, + ) + response = await router.async_pre_routing_hook( + model="test-complexity-router", + request_kwargs={}, + messages=[{"role": "system", "content": "be nice"}], + ) + assert response is not None + assert response.routing_decision is not None + assert response.routing_decision["cause"] == "default_fallback" + assert "tier" not in response.routing_decision @pytest.mark.asyncio async def test_session_pin_decision(self, mock_router_instance, basic_config): @@ -5907,11 +5941,21 @@ class TestConversationShapeDiscriminator: ) builds = source.split("self._build_routing_decision(")[1:] assert builds - missing = [ - i - for i, block in enumerate(builds) - if "conversation_continuing=conversation_continuing" not in block.split("),")[0] - ] + missing = [] + for i, block in enumerate(builds): + depth = 0 + end = 0 + for j, char in enumerate(block): + if char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth < 0: + end = j + break + extracted = block[:end] + if "conversation_continuing=conversation_continuing" not in extracted: + missing.append(i) assert not missing, f"routing decisions {missing} do not carry the conversation shape" From e50a42051c531f8d95a0dbf905917c991f8ed8f5 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 7 Aug 2026 17:28:55 -0700 Subject: [PATCH 08/10] fix(websearch): restore snippet text in native web_search_tool_result blocks (LIT-5315) (#36228) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(websearch): restore snippet text in native web_search_tool_result blocks (LIT-5315) The build_web_search_tool_result_block method copied url/title/page_age but hardcoded encrypted_content to empty string, never reading SearchResult.snippet. This left every native block content-free, forcing clients to web_fetch each result to recover evidence—the reported symptom. The Anthropic spec carries page text only in encrypted_content (an opaque server-issued blob we cannot mint), so snippet is emitted as an additive key alongside the spec fields. encrypted_content stays empty rather than holding plaintext, which would assert encryption semantics that don't hold. The anthropic SDK's BaseModel sets extra='allow', so the additive snippet key survives SDK parsing. litellm has no typed model for web_search_result at all, so nothing drops it internally. Turn-2 replay behavior is unaffected: the empty encrypted_content already exists today. Tests: - Updated test_shape_with_results to assert snippet present - Added test_snippet_carried_for_every_result to cover multi-result ordering - Added test_missing_snippet_degrades_to_empty_string for edge case - Mutation check: reverting source-only yields 3 test failures, restored to 117 passed Fixes: LIT-5315 Co-Authored-By: Claude * fix(websearch): make synthesized web_search blocks replayable by native clients Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(websearch): flatten a resultless replayed search block so Bedrock accepts the next turn The flatten added for LIT-5315 bails when the replayed web_search_tool_result carries an empty content list, but that is exactly what the interceptor emits when a search legitimately returns nothing and when a search raises. The block survived into the outbound body, Bedrock rejected the tag, and the conversation died on the following turn just as it did before the flatten existed. An empty content list has no encrypted_content to respect and no evidence to preserve, so it flattens safely, and its paired server_tool_use goes with it. The rendered text now says so explicitly rather than emitting a bare header. Adds the multi-turn replay coverage that existed nowhere: the outbound Bedrock invoke body is asserted free of both block types, parametrized over the results-present and resultless cases, and built from the interceptor's own builder so the fixture cannot drift from what it emits. Resolves LIT-5320 * test(websearch): pin flatten idempotency for the agentic-loop re-entry The agentic loop re-enters the same /v1/messages entry point for its follow-up call and hands it the original client history, so the flatten runs again over already-flattened messages once per iteration. Bedrock always takes that path, since its config reports web search as natively handled and the short-circuit is skipped. A pass that appended the rendered text instead of replacing the block would duplicate the evidence on every iteration and re-ship the unsupported tag, and no existing single-pass test sees it. Mutation checked: keeping the original block alongside the rendered text fails this test on its own. --------- Co-authored-by: Claude Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Yassin Kortam --- .../websearch_interception/handler.py | 56 ++++-- .../websearch_interception/transformation.py | 10 + litellm/llms/anthropic/common_utils.py | 151 +++++++++++++- .../messages/handler.py | 3 + .../integrations/websearch_interception.py | 23 ++- .../test_websearch_native_blocks.py | 47 ++++- ...erimental_pass_through_messages_handler.py | 55 ++++++ .../anthropic/test_anthropic_common_utils.py | 186 ++++++++++++++++++ .../test_anthropic_claude3_transformation.py | 65 ++++++ 9 files changed, 572 insertions(+), 24 deletions(-) diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 71388134e98..9748db2dcd2 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -9,7 +9,7 @@ server-side using litellm router's search tools. import asyncio import math import uuid -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast import litellm @@ -37,6 +37,8 @@ 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 @@ -833,22 +835,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 diff --git a/litellm/integrations/websearch_interception/transformation.py b/litellm/integrations/websearch_interception/transformation.py index 795810a7c40..199ab020559 100644 --- a/litellm/integrations/websearch_interception/transformation.py +++ b/litellm/integrations/websearch_interception/transformation.py @@ -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 { diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 314cfef6d84..9aa5a4f465f 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -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: diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 3ef298aa336..c4b5cc628e2 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -14,6 +14,7 @@ from typing import Any, Final, cast import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, sanitize_tool_use_ids_in_anthropic_messages, strip_empty_text_blocks_from_anthropic_messages, ) @@ -222,6 +223,7 @@ async def anthropic_messages( # Replay of cross-provider tool history (e.g. kimi -> Anthropic) may carry # ids like ``functions.Bash:0`` that violate Anthropic's id pattern. messages = sanitize_tool_use_ids_in_anthropic_messages(messages) + messages = flatten_unencrypted_web_search_results_in_anthropic_messages(messages) from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, @@ -413,6 +415,7 @@ def anthropic_messages_handler( if not kwargs.pop("_litellm_messages_presanitized", False): messages = strip_empty_text_blocks_from_anthropic_messages(messages) messages = sanitize_tool_use_ids_in_anthropic_messages(messages) + messages = flatten_unencrypted_web_search_results_in_anthropic_messages(messages) from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, diff --git a/litellm/types/integrations/websearch_interception.py b/litellm/types/integrations/websearch_interception.py index 05537da67d7..90713b270be 100644 --- a/litellm/types/integrations/websearch_interception.py +++ b/litellm/types/integrations/websearch_interception.py @@ -2,7 +2,28 @@ Type definitions for WebSearch Interception integration. """ -from typing import TypedDict +from typing import Literal, TypedDict + +from pydantic import BaseModel + + +class AnthropicSearchQuery(BaseModel): + """``input`` of an Anthropic ``server_tool_use`` block for a web search.""" + + query: str + + +class AnthropicServerToolUseBlock(BaseModel): + """ + The ``server_tool_use`` block that must accompany a ``web_search_tool_result``. + + Anthropic requires the pair, with a ``srvtoolu_``-prefixed id shared by both. + """ + + type: Literal["server_tool_use"] = "server_tool_use" + id: str + name: Literal["web_search"] = "web_search" + input: AnthropicSearchQuery class WebSearchInterceptionConfig(TypedDict, total=False): diff --git a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py b/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py index 544abab8dcf..c859f9b2f55 100644 --- a/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py +++ b/tests/test_litellm/integrations/websearch_interception/test_websearch_native_blocks.py @@ -134,6 +134,30 @@ class TestBuildWebSearchToolResultBlock: assert first["title"] == "LiteLLM Docs" assert first["page_age"] == "2025-01-15" assert first["encrypted_content"] == "" + assert first["snippet"] == "Unified interface for LLMs." + + def test_snippet_carried_for_every_result(self): + # The snippet is the only field carrying page text. Losing it leaves the + # client and the model with nothing to answer from, forcing a fetch per + # result. + block = WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="toolu_abc", + search_response=_make_search_response(), + ) + assert [r["snippet"] for r in block["content"]] == [ + "Unified interface for LLMs.", + "Pay-per-use pricing model.", + ] + + def test_missing_snippet_degrades_to_empty_string(self): + response = SearchResponse( + results=[SearchResult(title="T", url="https://x/", snippet="")] + ) + block = WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="toolu_abc", + search_response=response, + ) + assert block["content"][0]["snippet"] == "" def test_handles_none_search_response(self): block = WebSearchTransformation.build_web_search_tool_result_block( @@ -223,11 +247,15 @@ class TestBuildPlanAttachesBlocks: ) blocks = plan.metadata.get(WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY) - assert isinstance(blocks, list) - assert len(blocks) == 1 - assert blocks[0]["type"] == "web_search_tool_result" - assert blocks[0]["tool_use_id"] == "toolu_one" - assert blocks[0]["content"][0]["url"] == "https://docs.litellm.ai/" + assert isinstance(blocks, tuple) + assert [b["type"] for b in blocks] == [ + "server_tool_use", + "web_search_tool_result", + ] + assert blocks[0]["id"].startswith("srvtoolu_") + assert blocks[0]["input"] == {"query": "what is litellm"} + assert blocks[1]["tool_use_id"] == blocks[0]["id"] + assert blocks[1]["content"][0]["url"] == "https://docs.litellm.ai/" @pytest.mark.asyncio async def test_metadata_does_not_carry_blocks_when_flag_absent(self): @@ -479,6 +507,9 @@ class TestLegacyPathMatchesNewPath: kwargs={WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: True}, ) - assert out["content"][0]["type"] == "web_search_tool_result" - assert out["content"][0]["tool_use_id"] == "toolu_legacy" - assert out["content"][1]["type"] == "text" + assert [b["type"] for b in out["content"]] == [ + "server_tool_use", + "web_search_tool_result", + "text", + ] + assert out["content"][1]["tool_use_id"] == out["content"][0]["id"] diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index df3db3d2c57..f11324ca376 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -732,6 +732,61 @@ def test_handler_skips_strip_when_presanitized(): assert result is not None +def test_handler_flattens_replayed_unencrypted_web_search_results(): + """Synthesized search blocks replayed as history must reach the provider as text.""" + from litellm.llms.anthropic.experimental_pass_through.messages import handler + + captured = {} + + def fake_base_handler(*args, **kwargs): + captured.update(kwargs) + return "stub" + + with patch.object( + handler.base_llm_http_handler, + "anthropic_messages_handler", + side_effect=fake_base_handler, + ): + handler.anthropic_messages_handler( + max_tokens=10, + messages=[ + {"role": "user", "content": "latest litellm version?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "latest litellm version"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [ + { + "type": "web_search_result", + "url": "https://github.com/BerriAI/litellm/releases", + "title": "Releases", + "page_age": None, + "encrypted_content": "", + "snippet": "Latest release v1.95.0", + } + ], + }, + ], + }, + {"role": "user", "content": "which version?"}, + ], + model="anthropic/claude-3-5-sonnet-20241022", + custom_llm_provider="anthropic", + ) + + replayed = captured["messages"][1]["content"] + assert [b["type"] for b in replayed] == ["text"] + assert "Snippet: Latest release v1.95.0" in replayed[0]["text"] + + def test_presanitized_flag_not_leaked_to_provider_params(): """The private sentinel must be popped, never forwarded as a request param.""" from litellm.llms.anthropic.experimental_pass_through.messages import handler diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 6ab0f2c08ab..9df72108332 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -10,8 +10,10 @@ Verifies that: - ANTHROPIC_API_KEY / ANTHROPIC_API_BASE take precedence over their aliases. """ +import json import os import sys +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -1457,6 +1459,190 @@ class TestAnthropicThinkingSignatureSelfHeal: out = strip_empty_text_blocks_from_anthropic_messages(msgs) assert [b["type"] for b in out[0]["content"]] == ["tool_result"] + def test_flatten_unencrypted_web_search_results_keeps_snippet_evidence(self): + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + msgs = [ + {"role": "user", "content": "latest litellm version?"}, + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "latest litellm version"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [ + { + "type": "web_search_result", + "url": "https://github.com/BerriAI/litellm/releases", + "title": "Releases", + "page_age": None, + "encrypted_content": "", + "snippet": "Latest release v1.95.0", + } + ], + }, + {"type": "text", "text": "v1.95.0"}, + ], + }, + ] + + out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + + assert out[0] is msgs[0] + assert [b["type"] for b in out[1]["content"]] == ["text", "text"] + flattened = out[1]["content"][0]["text"] + assert "Web search results for 'latest litellm version':" in flattened + assert "URL: https://github.com/BerriAI/litellm/releases" in flattened + assert "Snippet: Latest release v1.95.0" in flattened + assert msgs[1]["content"][0]["type"] == "server_tool_use" + + @pytest.mark.parametrize("results", [[], None], ids=["empty_list", "search_raised"]) + def test_flatten_unencrypted_web_search_results_flattens_a_resultless_search(self, results): + """A search that found nothing, or that raised, still has to be flattened. + + Both cases reach the client as ``content: []``, and leaving that block in + place ships an unsupported tag to Bedrock on the next turn just as surely + as a populated one does. + """ + from litellm.integrations.websearch_interception.transformation import ( + WebSearchTransformation, + ) + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + block = WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="srvtoolu_1", + search_response=None if results is None else SimpleNamespace(results=results), + ) + assert block["content"] == [], "fixture drifted from what the interceptor emits" + + msgs = [ + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "who won"}, + }, + block, + {"type": "text", "text": "I could not find that."}, + ], + } + ] + + out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + + assert [b["type"] for b in out[0]["content"]] == ["text", "text"] + assert out[0]["content"][0]["text"] == ("Web search results for 'who won':\n\nNo results were returned.") + + @pytest.mark.parametrize("results", [[SimpleNamespace(title="Rome", url="u", snippet="s", date=None)], []]) + def test_flatten_unencrypted_web_search_results_is_idempotent(self, results): + """Flattening twice must equal flattening once. + + The agentic loop re-enters the same entry point for its follow-up call and + hands it the original history, so this runs again on already-flattened + messages once per iteration. A pass that appended instead of replacing + would duplicate the evidence on every loop. + """ + from litellm.integrations.websearch_interception.transformation import ( + WebSearchTransformation, + ) + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + msgs = [ + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_1", "name": "web_search", "input": {"query": "when"}}, + WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="srvtoolu_1", + search_response=SimpleNamespace(results=results), + ), + {"type": "text", "text": "753 BC."}, + ], + } + ] + + once = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + twice = flatten_unencrypted_web_search_results_in_anthropic_messages(once) + + assert [b["type"] for b in once[0]["content"]] == ["text", "text"] + assert json.dumps(twice) == json.dumps(once) + + def test_flatten_unencrypted_web_search_results_preserves_real_anthropic_blocks(self): + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + msgs = [ + { + "role": "assistant", + "content": [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "q"}, + }, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": [ + { + "type": "web_search_result", + "url": "https://example.com", + "title": "Example", + "page_age": None, + "encrypted_content": "EqgfCioIARgBIiQ4", + } + ], + }, + ], + } + ] + + out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + + assert out[0] is msgs[0] + + def test_flatten_unencrypted_web_search_results_leaves_error_blocks_alone(self): + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + + msgs = [ + { + "role": "assistant", + "content": [ + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_1", + "content": { + "type": "web_search_tool_result_error", + "error_code": "max_uses_exceeded", + }, + } + ], + } + ] + + out = flatten_unencrypted_web_search_results_in_anthropic_messages(msgs) + + assert out[0] is msgs[0] + def test_sanitize_tool_use_ids_in_anthropic_messages(self): from litellm.llms.anthropic.common_utils import ( sanitize_tool_use_ids_in_anthropic_messages, diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 3b8b4af78d9..76bb11cc26d 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -4,6 +4,7 @@ import json import os import sys from datetime import datetime +from types import SimpleNamespace from unittest.mock import Mock import pytest @@ -2515,3 +2516,67 @@ def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag( assert thinking.get("type") == "enabled" assert isinstance(thinking.get("budget_tokens"), int) assert "output_config" not in flipped + + +@pytest.mark.parametrize( + "search_results, expected_evidence", + [ + pytest.param( + [SimpleNamespace(title="Rome", url="https://ex.com/rome", snippet="Founded 753 BC.", date=None)], + "Snippet: Founded 753 BC.", + id="search_returned_results", + ), + pytest.param([], "No results were returned.", id="search_returned_nothing"), + ], +) +def test_replayed_intercepted_search_turn_leaves_no_unsupported_block_for_bedrock(search_results, expected_evidence): + """A native client replaying an intercepted search turn must not 400 on Bedrock. + + ``websearch_interception`` hands Claude Desktop an Anthropic-native + ``server_tool_use`` + ``web_search_tool_result`` pair, and Anthropic's protocol + obliges the client to replay that assistant turn verbatim on every later turn. + Bedrock's Anthropic schema defines neither tag, so both have to be gone from the + outbound body by the time it is signed, with the search evidence carried forward + as text instead. Built from the real builder rather than a hand-written fixture + so the two cannot drift apart. + """ + from litellm.integrations.websearch_interception.transformation import ( + WebSearchTransformation, + ) + from litellm.llms.anthropic.common_utils import ( + flatten_unencrypted_web_search_results_in_anthropic_messages, + ) + from litellm.types.router import GenericLiteLLMParams + + replayed_turn = [ + { + "type": "server_tool_use", + "id": "srvtoolu_1", + "name": "web_search", + "input": {"query": "when was Rome founded"}, + }, + WebSearchTransformation.build_web_search_tool_result_block( + tool_use_id="srvtoolu_1", + search_response=SimpleNamespace(results=search_results), + ), + {"type": "text", "text": "Rome was founded in 753 BC."}, + ] + messages = [ + {"role": "user", "content": [{"type": "text", "text": "When was Rome founded?"}]}, + {"role": "assistant", "content": replayed_turn}, + {"role": "user", "content": [{"type": "text", "text": "Repeat the year."}]}, + ] + + body = AmazonAnthropicClaudeMessagesConfig().transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=flatten_unencrypted_web_search_results_in_anthropic_messages(messages), + anthropic_messages_optional_request_params={"max_tokens": 64}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + serialized = json.dumps(body) + assert "web_search_tool_result" not in serialized + assert "server_tool_use" not in serialized + assert expected_evidence in serialized + assert "Rome was founded in 753 BC." in serialized From 1a45bf9afebfb26e656b09c16df54a850fd035e3 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 17:45:30 -0700 Subject: [PATCH 09/10] fix(proxy): resolve entity access groups in the model listing endpoints (#36230) * fix(proxy): resolve entity access groups in the model listing endpoints Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): reuse the fetched team object when listing models Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): cover key-level access group resolution in model listing Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: shivam Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_checks.py | 5 +- litellm/proxy/auth/model_checks.py | 7 +- litellm/proxy/utils.py | 136 +++++++++++++++--- .../proxy/auth/test_model_checks.py | 25 ++++ .../proxy/utils/helpers/test_model_access.py | 117 +++++++++++++++ 5 files changed, 262 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 3da899a5610..d07ac0c5586 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -13,6 +13,7 @@ import asyncio import math import re import time +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast from fastapi import HTTPException, Request, status @@ -2834,7 +2835,7 @@ async def get_org_object( async def _get_resources_from_access_groups( - access_group_ids: list[str], + access_group_ids: Sequence[str], resource_field: Literal["access_model_names", "access_mcp_server_ids", "access_agent_ids"], prisma_client: PrismaClient | None = None, user_api_key_cache: UserApiKeyCache | None = None, @@ -2893,7 +2894,7 @@ async def _get_resources_from_access_groups( async def _get_models_from_access_groups( - access_group_ids: list[str], + access_group_ids: Sequence[str], prisma_client: PrismaClient | None = None, user_api_key_cache: UserApiKeyCache | None = None, proxy_logging_obj: ProxyLogging | None = None, diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index a6e5eb2a0a0..ff9211742f3 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -1,6 +1,7 @@ # What is this? ## Common checks for /v1/models and `/model/info` import copy +from collections.abc import Sequence from typing import Any, Final import litellm @@ -178,8 +179,8 @@ def get_team_models( def get_complete_model_list( - key_models: list[str], - team_models: list[str], + key_models: Sequence[str], + team_models: Sequence[str], proxy_model_list: list[str], user_model: str | None, infer_model_from_keys: bool | None, @@ -203,7 +204,7 @@ def get_complete_model_list( def append_unique(models): for model in models: - if model not in unique_models: + if model not in unique_models and model != SpecialModelNames.no_default_models.value: unique_models.append(model) if key_models: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 605455a2f73..e59c6adaf22 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -165,6 +165,7 @@ if TYPE_CHECKING: from prisma.client import TransactionManager from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction @@ -6419,6 +6420,74 @@ def construct_database_url_from_env_vars() -> str | None: return None +async def _get_validated_team_object( + user_api_key_dict: "UserAPIKeyAuth", + team_id: str, + prisma_client: "PrismaClient", + user_api_key_cache: "UserApiKeyCache", + proxy_logging_obj: "ProxyLogging", +) -> "LiteLLM_TeamTableCachedObj": + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.management_endpoints.team_endpoints import validate_membership + + team_object: Final = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object) + return team_object + + +async def _get_team_object_for_access_groups( + team_id: str | None, + prisma_client: Optional["PrismaClient"], + user_api_key_cache: Optional["UserApiKeyCache"], + proxy_logging_obj: Optional["ProxyLogging"], +) -> Optional["LiteLLM_TeamTableCachedObj"]: + from litellm.proxy.auth.auth_checks import get_team_object + + if team_id is None or prisma_client is None or user_api_key_cache is None or proxy_logging_obj is None: + return None + try: + return await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException: + verbose_proxy_logger.debug("Could not fetch team %s while listing models", team_id) + return None + + +async def _get_access_group_models( + user_api_key_dict: "UserAPIKeyAuth", + team_object: Optional["LiteLLM_TeamTableCachedObj"], + prisma_client: Optional["PrismaClient"], + user_api_key_cache: Optional["UserApiKeyCache"], + proxy_logging_obj: Optional["ProxyLogging"], +) -> tuple[str, ...]: + from litellm.proxy.auth.auth_checks import ( + _get_models_from_access_groups, + get_authorized_resources_from_key_access_groups, + ) + + team_group_models: Final = await _get_models_from_access_groups( + access_group_ids=(team_object.access_group_ids or ()) if team_object is not None else (), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + key_group_models: Final = await get_authorized_resources_from_key_access_groups( + valid_token=user_api_key_dict, + team_object=team_object, + resource_field="access_model_names", + ) + return tuple(dict.fromkeys((*team_group_models, *key_group_models))) + + async def get_available_models_for_user( user_api_key_dict: "UserAPIKeyAuth", llm_router: Optional["Router"], @@ -6450,13 +6519,11 @@ async def get_available_models_for_user( Returns: List of model names available to the user """ - from litellm.proxy.auth.auth_checks import get_team_object from litellm.proxy.auth.model_checks import ( get_complete_model_list, get_key_models, get_team_models, ) - from litellm.proxy.management_endpoints.team_endpoints import validate_membership # Get proxy model list and access groups if llm_router is None: @@ -6466,31 +6533,33 @@ async def get_available_models_for_user( proxy_model_list = llm_router.get_model_names() model_access_groups = llm_router.get_model_access_groups() - # Get key models - key_models = get_key_models( - user_api_key_dict=user_api_key_dict, - proxy_model_list=proxy_model_list, - model_access_groups=model_access_groups, - include_model_access_groups=include_model_access_groups, - ) - - # Get team models - team_models: list[str] = user_api_key_dict.team_models - - # If specific team_id is provided, validate and get team models - if team_id and prisma_client and proxy_logging_obj and user_api_key_cache: - key_models = [] - team_object: Final = await get_team_object( + requested_team_object: Final = ( + await _get_validated_team_object( + user_api_key_dict=user_api_key_dict, team_id=team_id, prisma_client=prisma_client, user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, ) - await validate_membership(user_api_key_dict=user_api_key_dict, team_table=team_object) - team_models = team_object.models + if team_id and prisma_client and proxy_logging_obj and user_api_key_cache + else None + ) - team_models = get_team_models( - team_models=team_models, + key_models: Final[Sequence[str]] = ( + () + if requested_team_object is not None + else get_key_models( + user_api_key_dict=user_api_key_dict, + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups, + ) + ) + + team_models: Final = get_team_models( + team_models=( + requested_team_object.models if requested_team_object is not None else user_api_key_dict.team_models + ), proxy_model_list=proxy_model_list, model_access_groups=model_access_groups, include_model_access_groups=include_model_access_groups, @@ -6498,10 +6567,31 @@ async def get_available_models_for_user( effective_team_id: Final = team_id or user_api_key_dict.team_id + access_group_models: Final = ( + await _get_access_group_models( + user_api_key_dict=user_api_key_dict, + team_object=requested_team_object + or await _get_team_object_for_access_groups( + team_id=effective_team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if key_models or team_models + else () + ) + + granted_key_models: Final = (*key_models, *access_group_models) if key_models else key_models + granted_team_models: Final = (*team_models, *access_group_models) if team_models else team_models + # Get complete model list all_models: Final = get_complete_model_list( - key_models=key_models, - team_models=team_models, + key_models=granted_key_models, + team_models=granted_team_models, proxy_model_list=proxy_model_list, user_model=user_model, infer_model_from_keys=general_settings.get("infer_model_from_keys", False), diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index f56b8e113a9..e6c0eaee3c4 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -735,3 +735,28 @@ def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion(): litellm.vertex_language_models.discard(fake_model) litellm.add_known_models(model_cost_map={}) assert fake_model not in litellm.models_by_provider["vertex_ai"] + +def test_get_complete_model_list_drops_no_default_models_sentinel(): + from litellm.proxy.auth.model_checks import get_complete_model_list + + result = get_complete_model_list( + key_models=["no-default-models", "model-a"], + team_models=[], + proxy_model_list=["model-a", "model-b"], + user_model=None, + infer_model_from_keys=False, + ) + assert result == ["model-a"] + + +def test_get_complete_model_list_sentinel_only_grants_nothing(): + from litellm.proxy.auth.model_checks import get_complete_model_list + + result = get_complete_model_list( + key_models=["no-default-models"], + team_models=["no-default-models"], + proxy_model_list=["model-a", "model-b"], + user_model=None, + infer_model_from_keys=False, + ) + assert result == [] diff --git a/tests/test_litellm/proxy/utils/helpers/test_model_access.py b/tests/test_litellm/proxy/utils/helpers/test_model_access.py index 59268e1427b..5fb4392eec6 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_model_access.py +++ b/tests/test_litellm/proxy/utils/helpers/test_model_access.py @@ -9,6 +9,7 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ( create_model_info_response, get_available_models_for_user, + hash_token, is_known_model, is_known_vector_store_index, model_dump_with_preserved_fields, @@ -404,3 +405,119 @@ async def test_get_available_models_for_user_error_path_complete_list_raises( general_settings={}, user_model=None, ) + + +@pytest.mark.asyncio +async def test_get_available_models_for_user_resolves_team_access_group_models( + monkeypatch, +): + from litellm.models.access_group import LiteLLM_AccessGroupTable + from litellm.models.team import LiteLLM_TeamTableCachedObj + + team = LiteLLM_TeamTableCachedObj( + team_id="team-1", + models=["no-default-models"], + access_group_ids=["ag-1"], + ) + access_group = LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="repro-group", + access_model_names=["model-a", "model-b"], + assigned_team_ids=["team-1"], + ) + + async def _get_team_object(**_kwargs): + return team + + async def _get_access_object(**_kwargs): + return access_group + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object) + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_access_object", _get_access_object) + + result = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", + user_id="user-1", + team_id="team-1", + models=["all-team-models"], + team_models=["no-default-models"], + ), + llm_router=_router_with_models(["model-a", "model-b", "model-c"]), + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert sorted(result) == ["model-a", "model-b"] + + +@pytest.mark.asyncio +async def test_get_available_models_for_user_without_access_groups_grants_nothing( + monkeypatch, +): + from litellm.models.team import LiteLLM_TeamTableCachedObj + + async def _get_team_object(**_kwargs): + return LiteLLM_TeamTableCachedObj(team_id="team-1", models=["no-default-models"]) + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object) + + result = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", + user_id="user-1", + team_id="team-1", + models=["all-team-models"], + team_models=["no-default-models"], + ), + llm_router=_router_with_models(["model-a", "model-b"]), + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert result == [] + +@pytest.mark.asyncio +async def test_get_available_models_for_user_resolves_key_access_group_models( + monkeypatch, +): + from litellm.models.access_group import LiteLLM_AccessGroupTable + from litellm.models.team import LiteLLM_TeamTableCachedObj + + async def _get_team_object(**_kwargs): + return LiteLLM_TeamTableCachedObj(team_id="team-1", models=["no-default-models"]) + + async def _get_access_object(**_kwargs): + return LiteLLM_AccessGroupTable( + access_group_id="ag-1", + access_group_name="key-group", + access_model_names=["model-b"], + assigned_key_ids=[hash_token("sk-test-key")], + ) + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_object", _get_team_object) + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_access_object", _get_access_object) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()) + + result = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", + user_id="user-1", + team_id="team-1", + models=["no-default-models"], + team_models=["no-default-models"], + access_group_ids=["ag-1"], + ), + llm_router=_router_with_models(["model-a", "model-b"]), + general_settings={}, + user_model=None, + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + user_api_key_cache=MagicMock(), + ) + assert result == ["model-b"] From 2a9aac70045282b82c67b87b52881c8af0521db1 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 7 Aug 2026 17:45:50 -0700 Subject: [PATCH 10/10] fix(ui): let access groups be a team's only model source, with hover provenance (#36234) * feat(proxy): return per-group model provenance on /team/info /team/info now carries access_group_details, one entry per resolved access group with its id, name, and model list, so the UI can attribute each inherited model to the group granting it. The batch resolver returns the access group rows keyed by id instead of a stringly dict of lists, and the team member budget helper returns a copy instead of mutating its parameter. Type discipline and basedpyright budgets ratchet down accordingly. * feat(ui): allow group-only teams and show model provenance on hover Team create and edit no longer require a model selection: an empty selection is saved as the no-default-models sentinel, never as a bare empty list, since an empty team model list means unrestricted access. The team info Models card now renders every badge with a hover tooltip naming how the team got that model: directly, via named access groups, or both, and group-granted badges stay visible when the direct list is empty or a sentinel. * refactor(proxy): dedupe access group ids and return copies instead of mutating Duplicate access_group_ids no longer amplify the /team/info response: ids collapse order-preserving before provenance is built, pinned by a regression test. The resolver returns a model_copy rather than mutating its parameter, and the team create call sends a new object instead of reassigning formValues.models. Budgets ratchet down further with the mutation removal. --- basedpyright-code-budget.json | 6 +- litellm/proxy/_types.py | 7 ++ .../management_endpoints/team_endpoints.py | 88 ++++++++++--------- .../test_team_endpoints.py | 79 +++++++++++++++-- type-discipline-budget.json | 6 +- .../src/components/Teams.test.tsx | 28 ++++++ ui/litellm-dashboard/src/components/Teams.tsx | 11 +-- .../src/components/team/TeamInfo.tsx | 41 +++++---- .../components/team/teamModelAccess.test.ts | 86 ++++++++++++++++++ .../src/components/team/teamModelAccess.ts | 82 +++++++++++++++++ 10 files changed, 357 insertions(+), 77 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/team/teamModelAccess.test.ts create mode 100644 ui/litellm-dashboard/src/components/team/teamModelAccess.ts diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index c8765eb0bd0..32b8eb3d4d0 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 45110 + "limit": 45098 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39838 + "limit": 39826 }, "reportUnknownParameterType": { "limit": 20237 }, "reportUnknownVariableType": { - "limit": 31383 + "limit": 31371 }, "reportUnnecessaryCast": { "limit": 122 diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5b7dc3a7c73..1fc05ac4653 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3894,12 +3894,19 @@ class OrganizationMemberUpdateResponse(MemberUpdateResponse): ########################################## +class TeamAccessGroupModelGrant(LiteLLMPydanticObjectBase): + access_group_id: str + access_group_name: str + models: tuple[str, ...] + + class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): team_member_budget_table: LiteLLM_BudgetTableFull | None = None # Resources inherited from access groups (separate from direct assignments) access_group_models: list[str] | None = None access_group_mcp_server_ids: list[str] | None = None access_group_agent_ids: list[str] | None = None + access_group_details: tuple[TeamAccessGroupModelGrant, ...] | None = None class TeamInfoResponseObject(TypedDict): diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index fe5a0e06d2e..0b99879f9fe 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -57,6 +57,7 @@ from litellm.proxy._types import ( SpecialManagementEndpointEnums, SpecialModelNames, SpecialProxyStrings, + TeamAccessGroupModelGrant, TeamAddMemberResponse, TeamInfoResponseObject, TeamInfoResponseObjectTeamTable, @@ -3829,7 +3830,7 @@ async def _add_team_member_budget_table( ) -> TeamInfoResponseObjectTeamTable: try: team_budget: Final = await _budget_db(prisma_client).find_unique(where={"budget_id": team_member_budget_id}) - team_info_response_object.team_member_budget_table = team_budget + return team_info_response_object.model_copy(update={"team_member_budget_table": team_budget}) except Exception: verbose_proxy_logger.info( "Team member budget table not found, passed team_member_budget_id=%s", team_member_budget_id @@ -3838,21 +3839,34 @@ async def _add_team_member_budget_table( return team_info_response_object -async def _resolve_team_access_group_resources(_team_info: TeamInfoResponseObjectTeamTable) -> None: - """Populate access_group_models / mcp_server_ids / agent_ids on the team - info response by resolving inherited resources from its access groups.""" +async def _resolve_team_access_group_resources( + _team_info: TeamInfoResponseObjectTeamTable, +) -> TeamInfoResponseObjectTeamTable: + """Return a copy of the team info with access_group_models / mcp_server_ids / + agent_ids / details resolved from its access groups.""" if not _team_info.access_group_ids: - return + return _team_info ag_lookup: Final = await _batch_resolve_access_group_resources(_team_info.access_group_ids) - models, mcp_ids, agent_ids = set(), set(), set() - for ag_id in _team_info.access_group_ids: - if ag_id in ag_lookup: - models.update(ag_lookup[ag_id]["models"]) - mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"]) - agent_ids.update(ag_lookup[ag_id]["agent_ids"]) - _team_info.access_group_models = list(models) - _team_info.access_group_mcp_server_ids = list(mcp_ids) - _team_info.access_group_agent_ids = list(agent_ids) + resolved_groups: Final = tuple( + ag_lookup[ag_id] for ag_id in dict.fromkeys(_team_info.access_group_ids) if ag_id in ag_lookup + ) + return _team_info.model_copy( + update={ + "access_group_models": list({m for group in resolved_groups for m in (group.access_model_names or [])}), + "access_group_mcp_server_ids": list( + {s for group in resolved_groups for s in (group.access_mcp_server_ids or [])} + ), + "access_group_agent_ids": list({a for group in resolved_groups for a in (group.access_agent_ids or [])}), + "access_group_details": tuple( + TeamAccessGroupModelGrant( + access_group_id=group.access_group_id, + access_group_name=group.access_group_name, + models=tuple(group.access_model_names or ()), + ) + for group in resolved_groups + ), + } + ) @router.get("/team/info", tags=["team management"], dependencies=[Depends(user_api_key_auth)]) @@ -3958,11 +3972,11 @@ async def team_info( ) # Resolve resources inherited from access groups - await _resolve_team_access_group_resources(_team_info) + resolved_team_info: Final = await _resolve_team_access_group_resources(_team_info) response_object: Final = TeamInfoResponseObject( team_id=team_id, - team_info=_team_info, + team_info=resolved_team_info, keys=keys, team_memberships=returned_tm, ) @@ -4391,32 +4405,21 @@ async def _build_team_list_where_conditions( async def _batch_resolve_access_group_resources( all_access_group_ids: list[str], -) -> dict[str, dict[str, list[str]]]: +) -> dict[str, LiteLLM_AccessGroupTable]: """ - Batch-fetch access groups in a single DB query and return a per-group - resource map. - - Returns {ag_id: {"models": [...], "mcp_server_ids": [...], "agent_ids": [...]}}. - Missing/invalid groups are silently omitted. + Batch-fetch access groups in a single DB query and return them keyed by + access_group_id. Missing/invalid groups are silently omitted. """ from litellm.proxy.proxy_server import prisma_client as _prisma_client if not all_access_group_ids or _prisma_client is None: return {} - unique_ids: Final = list(set(all_access_group_ids)) + unique_ids: Final = tuple(frozenset(all_access_group_ids)) rows: Final = await _access_group_db(_prisma_client).find_many( where={"access_group_id": {"in": unique_ids}}, ) - - result: Final[dict[str, dict[str, list[str]]]] = {} - for row in rows: - result[row.access_group_id] = { - "models": list(row.access_model_names or []), - "mcp_server_ids": list(row.access_mcp_server_ids or []), - "agent_ids": list(row.access_agent_ids or []), - } - return result + return {row.access_group_id: row for row in rows} def _convert_teams_to_response_models( @@ -4710,15 +4713,18 @@ async def list_team_v2( all_ag_ids: Final = [ag_id for t in team_items_with_ag for ag_id in (t.access_group_ids or [])] ag_lookup: Final = await _batch_resolve_access_group_resources(all_ag_ids) for team_item in team_items_with_ag: - models, mcp_ids, agent_ids = set(), set(), set() - for ag_id in team_item.access_group_ids or []: - if ag_id in ag_lookup: - models.update(ag_lookup[ag_id]["models"]) - mcp_ids.update(ag_lookup[ag_id]["mcp_server_ids"]) - agent_ids.update(ag_lookup[ag_id]["agent_ids"]) - team_item.access_group_models = list(models) - team_item.access_group_mcp_server_ids = list(mcp_ids) - team_item.access_group_agent_ids = list(agent_ids) + team_groups = tuple( + ag_lookup[ag_id] for ag_id in (team_item.access_group_ids or []) if ag_id in ag_lookup + ) + team_item.access_group_models = list( + {m for group in team_groups for m in (group.access_model_names or [])} + ) + team_item.access_group_mcp_server_ids = list( + {s for group in team_groups for s in (group.access_mcp_server_ids or [])} + ) + team_item.access_group_agent_ids = list( + {a for group in team_groups for a in (group.access_agent_ids or [])} + ) return { "teams": team_list, diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index eedffa1ea5f..a1cbc77b7a5 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -8637,9 +8637,9 @@ class TestBatchResolveAccessGroupResources: with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): result = await _batch_resolve_access_group_resources(["ag-1"]) - assert sorted(result["ag-1"]["models"]) == ["claude-3", "gpt-4"] - assert result["ag-1"]["mcp_server_ids"] == ["mcp-1"] - assert sorted(result["ag-1"]["agent_ids"]) == ["agent-1", "agent-2"] + assert sorted(result["ag-1"].access_model_names) == ["claude-3", "gpt-4"] + assert result["ag-1"].access_mcp_server_ids == ["mcp-1"] + assert sorted(result["ag-1"].access_agent_ids) == ["agent-1", "agent-2"] @pytest.mark.asyncio async def test_multiple_access_groups(self): @@ -8668,8 +8668,8 @@ class TestBatchResolveAccessGroupResources: with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): result = await _batch_resolve_access_group_resources(["ag-1", "ag-2"]) - assert result["ag-1"]["models"] == ["gpt-4"] - assert result["ag-2"]["models"] == ["gemini"] + assert result["ag-1"].access_model_names == ["gpt-4"] + assert result["ag-2"].access_model_names == ["gemini"] @pytest.mark.asyncio async def test_missing_access_group_omitted(self): @@ -8735,6 +8735,75 @@ class TestBatchResolveAccessGroupResources: assert "ag-1" in result +class TestResolveTeamAccessGroupResources: + """Tests for the per-team access group resolution on /team/info.""" + + @pytest.mark.asyncio + async def test_populates_flat_lists_and_per_group_details(self): + """access_group_details must attribute each model to the group granting it, + so the UI can show provenance on hover; flat lists stay for back-compat. + Duplicated ids must collapse to one entry (response amplification), and the + input object must stay untouched (resolution returns a copy).""" + from litellm.proxy._types import TeamInfoResponseObjectTeamTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_team_access_group_resources, + ) + + row1 = MagicMock() + row1.access_group_id = "ag-1" + row1.access_group_name = "shared-models" + row1.access_model_names = ["gpt-4", "claude-3"] + row1.access_mcp_server_ids = ["mcp-1"] + row1.access_agent_ids = [] + + row2 = MagicMock() + row2.access_group_id = "ag-2" + row2.access_group_name = "extra-models" + row2.access_model_names = ["claude-3", "gemini"] + row2.access_mcp_server_ids = [] + row2.access_agent_ids = ["agent-1"] + + fake_prisma = MagicMock() + fake_prisma.db.litellm_accessgrouptable.find_many = AsyncMock( + return_value=[row1, row2] + ) + + team_info = TeamInfoResponseObjectTeamTable( + team_id="team-1", access_group_ids=["ag-1", "ag-2", "ag-1", "ag-missing"] + ) + with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma): + resolved = await _resolve_team_access_group_resources(team_info) + + assert team_info.access_group_details is None + assert sorted(resolved.access_group_models or []) == [ + "claude-3", + "gemini", + "gpt-4", + ] + assert resolved.access_group_mcp_server_ids == ["mcp-1"] + assert resolved.access_group_agent_ids == ["agent-1"] + assert [ + (d.access_group_id, d.access_group_name, d.models) + for d in (resolved.access_group_details or []) + ] == [ + ("ag-1", "shared-models", ("gpt-4", "claude-3")), + ("ag-2", "extra-models", ("claude-3", "gemini")), + ] + + @pytest.mark.asyncio + async def test_no_access_groups_leaves_details_unset(self): + from litellm.proxy._types import TeamInfoResponseObjectTeamTable + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_team_access_group_resources, + ) + + team_info = TeamInfoResponseObjectTeamTable(team_id="team-1", access_group_ids=[]) + resolved = await _resolve_team_access_group_resources(team_info) + + assert resolved.access_group_details is None + assert resolved.access_group_models is None + + @pytest.mark.asyncio async def test_verify_team_access_denies_unauthorized_user(): """ diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 991f8eaa934..0a0cfe9a617 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23245 + "limit": 23235 }, "LIT002": { - "limit": 27179 + "limit": 27176 }, "LIT003": { "limit": 269 @@ -30,6 +30,6 @@ "limit": 16769 }, "LIT011": { - "limit": 5602 + "limit": 5598 } } diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index 2bda0f72cec..e8331294972 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -613,6 +613,34 @@ describe("Teams - access_group_ids in team create", () => { ); }); }); + + it("creates a team with no models selected, sending the no-default-models sentinel instead of an empty list", async () => { + renderWithQueryClient(); + + const createButton = screen.getAllByRole("button", { name: /create team/i })[0]; + act(() => { + fireEvent.click(createButton); + }); + + await waitFor(() => { + expect(screen.getByLabelText(/team name/i)).toBeInTheDocument(); + }); + + fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Group Only Team" } }); + + const createTeamSubmitButtons = screen.getAllByRole("button", { name: /create team/i }); + fireEvent.click(createTeamSubmitButtons[createTeamSubmitButtons.length - 1]); + + await waitFor(() => { + expect(teamCreateCall).toHaveBeenCalledWith( + "test-token", + expect.objectContaining({ + team_alias: "Group Only Team", + models: ["no-default-models"], + }), + ); + }); + }); }); describe("Teams - metadata key-value pairs in team create", () => { diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 95642b93019..42ff8bcc4c2 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -42,6 +42,7 @@ interface TeamProps { import DeleteResourceModal from "./common_components/DeleteResourceModal"; import { teamCreateCall } from "./networking"; +import { normalizeTeamModelSelection } from "./team/teamModelAccess"; import { ModelSelect } from "./ModelSelect/ModelSelect"; const canCreateOrManageTeams = ( @@ -351,7 +352,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser } } - await teamCreateCall(accessToken, formValues); + await teamCreateCall(accessToken, { ...formValues, models: normalizeTeamModelSelection(formValues.models) }); NotificationsManager.success("Team created"); await refreshTeams(); form.resetFields(); @@ -618,17 +619,11 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser label={ Models{" "} - + } - rules={[ - { - required: true, - message: "Please select at least one model", - }, - ]} name="models" > = new Set([ "disable_global_guardrails", ]); +const TEAM_MODEL_BADGE_COLORS: Record = { + "all-proxy": "red", + "no-default": "gray", + direct: "blue", + "access-group": "green", +}; + export interface TeamMembership { user_id: string; team_id: string; @@ -132,6 +145,7 @@ export interface TeamData { access_group_models?: string[]; access_group_mcp_server_ids?: string[]; access_group_agent_ids?: string[]; + access_group_details?: TeamAccessGroupModelGrant[]; router_settings?: Record; guardrails?: string[]; policies?: string[]; @@ -483,7 +497,7 @@ const TeamInfoView: React.FC = ({ const updateData: any = { team_id: teamId, team_alias: values.team_alias, - models: values.models, + models: normalizeTeamModelSelection(values.models), tpm_limit: sanitizeNumeric(values.tpm_limit), rpm_limit: sanitizeNumeric(values.rpm_limit), model_tpm_limit: modelTpmLimit, @@ -764,21 +778,14 @@ const TeamInfoView: React.FC = ({ Models
- {info.models.length === 0 || info.models.includes("all-proxy-models") ? ( - All proxy models - ) : ( - <> - {info.models.map((model: string, index: number) => ( - - {model} - - ))} - {(info.access_group_models || []).map((model: string, index: number) => ( - - {model} - - ))} - + {computeTeamModelBadges(info.models, info.access_group_models || [], info.access_group_details).map( + (badge, index) => ( + + + {badge.label} + + + ), )}
@@ -982,7 +989,7 @@ const TeamInfoView: React.FC = ({ { + it("substitutes the no-default-models sentinel for an empty selection", () => { + expect(normalizeTeamModelSelection([])).toEqual(["no-default-models"]); + expect(normalizeTeamModelSelection(undefined)).toEqual(["no-default-models"]); + }); + + it("passes a non-empty selection through untouched", () => { + expect(normalizeTeamModelSelection(["gpt-4o-mini"])).toEqual(["gpt-4o-mini"]); + expect(normalizeTeamModelSelection(["all-proxy-models"])).toEqual(["all-proxy-models"]); + }); +}); + +describe("computeTeamModelBadges", () => { + it("attributes group-only models to the groups granting them", () => { + const badges = computeTeamModelBadges(["sonnet-direct"], [], GRANTS); + expect(badges).toEqual([ + { + label: "sonnet-direct", + kind: "direct", + tooltip: "Granted directly in the team's model list", + }, + { label: "haiku", kind: "access-group", tooltip: "Granted via access groups shared, extra" }, + { label: "gpt-4o-mini", kind: "access-group", tooltip: "Granted via access group shared" }, + { label: "sonnet", kind: "access-group", tooltip: "Granted via access group extra" }, + ]); + }); + + it("marks a model both direct and group-granted on the direct badge, without a duplicate badge", () => { + const badges = computeTeamModelBadges(["haiku"], [], GRANTS); + expect(badges).toEqual([ + { + label: "haiku", + kind: "direct", + tooltip: "Granted directly in the team's model list, and also via access groups shared, extra", + }, + { label: "gpt-4o-mini", kind: "access-group", tooltip: "Granted via access group shared" }, + { label: "sonnet", kind: "access-group", tooltip: "Granted via access group extra" }, + ]); + }); + + it("shows the no-default-models sentinel as its own badge and keeps group badges visible", () => { + const badges = computeTeamModelBadges(["no-default-models"], [], [GRANTS[0]]); + expect(badges.map((b) => [b.label, b.kind])).toEqual([ + ["No default models", "no-default"], + ["haiku", "access-group"], + ["gpt-4o-mini", "access-group"], + ]); + }); + + it("still shows group badges when the empty model list grants everything", () => { + const badges = computeTeamModelBadges([], [], [GRANTS[0]]); + expect(badges[0]).toEqual({ + label: "All proxy models", + kind: "all-proxy", + tooltip: "The team's model list is empty, so it can access every model on the proxy", + }); + expect(badges.slice(1).map((b) => b.label)).toEqual(["haiku", "gpt-4o-mini"]); + }); + + it("distinguishes the all-proxy-models sentinel from an empty list in the tooltip", () => { + const badges = computeTeamModelBadges(["all-proxy-models"], [], []); + expect(badges).toEqual([ + { + label: "All proxy models", + kind: "all-proxy", + tooltip: "Granted by the All Proxy Models entry in the team's model list", + }, + ]); + }); + + it("falls back to the flat access_group_models list when per-group details are absent", () => { + const badges = computeTeamModelBadges(["direct-model"], ["haiku"], undefined); + expect(badges).toEqual([ + { label: "direct-model", kind: "direct", tooltip: "Granted directly in the team's model list" }, + { label: "haiku", kind: "access-group", tooltip: "Granted via an access group" }, + ]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/team/teamModelAccess.ts b/ui/litellm-dashboard/src/components/team/teamModelAccess.ts new file mode 100644 index 00000000000..91ddf0f4045 --- /dev/null +++ b/ui/litellm-dashboard/src/components/team/teamModelAccess.ts @@ -0,0 +1,82 @@ +export const ALL_PROXY_MODELS = "all-proxy-models"; +export const NO_DEFAULT_MODELS = "no-default-models"; + +export interface TeamAccessGroupModelGrant { + access_group_id: string; + access_group_name: string; + models: string[]; +} + +export type TeamModelBadgeKind = "all-proxy" | "no-default" | "direct" | "access-group"; + +export interface TeamModelBadge { + label: string; + kind: TeamModelBadgeKind; + tooltip: string; +} + +export function normalizeTeamModelSelection(models: string[] | undefined): string[] { + return models && models.length > 0 ? models : [NO_DEFAULT_MODELS]; +} + +const describeGroups = (names: string[]): string => + names.length > 1 ? `access groups ${names.join(", ")}` : `access group ${names[0]}`; + +export function computeTeamModelBadges( + models: string[], + accessGroupModels: string[], + accessGroupDetails: TeamAccessGroupModelGrant[] | undefined, +): TeamModelBadge[] { + const grants = accessGroupDetails ?? []; + const groupNamesFor = (model: string): string[] => + grants.filter((g) => g.models.includes(model)).map((g) => g.access_group_name); + const viaGroups = (model: string): string => { + const names = groupNamesFor(model); + return names.length > 0 ? describeGroups(names) : "an access group"; + }; + + const allProxy = models.length === 0 || models.includes(ALL_PROXY_MODELS); + const directModels = allProxy ? [] : models.filter((m) => m !== NO_DEFAULT_MODELS); + const groupModels = [...new Set(grants.length > 0 ? grants.flatMap((g) => g.models) : accessGroupModels)].filter( + (m) => !directModels.includes(m), + ); + + const allProxyBadge: TeamModelBadge = { + label: "All proxy models", + kind: "all-proxy", + tooltip: models.includes(ALL_PROXY_MODELS) + ? "Granted by the All Proxy Models entry in the team's model list" + : "The team's model list is empty, so it can access every model on the proxy", + }; + const noDefaultBadge: TeamModelBadge = { + label: "No default models", + kind: "no-default", + tooltip: "No models are granted directly. Access comes only from access groups", + }; + const headBadge = (): TeamModelBadge[] => { + if (allProxy) return [allProxyBadge]; + if (models.includes(NO_DEFAULT_MODELS)) return [noDefaultBadge]; + return []; + }; + + return [ + ...headBadge(), + ...directModels.map( + (m): TeamModelBadge => ({ + label: m, + kind: "direct", + tooltip: + groupNamesFor(m).length > 0 + ? `Granted directly in the team's model list, and also via ${viaGroups(m)}` + : "Granted directly in the team's model list", + }), + ), + ...groupModels.map( + (m): TeamModelBadge => ({ + label: m, + kind: "access-group", + tooltip: `Granted via ${viaGroups(m)}`, + }), + ), + ]; +}