diff --git a/.github/workflows/sync-together-ai-models.yml b/.github/workflows/sync-together-ai-models.yml index f1a8a841d0f..1daaadeabe2 100644 --- a/.github/workflows/sync-together-ai-models.yml +++ b/.github/workflows/sync-together-ai-models.yml @@ -32,7 +32,7 @@ jobs: echo "An open sync PR already exists on branch $open_pr; skipping this run." fi env: - GH_TOKEN: ${{ secrets.GH_TOKEN }} + GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }} - name: Run the sync if: steps.existing.outputs.open_pr == '' run: | @@ -65,4 +65,4 @@ jobs: --head "$branch" \ --base litellm_internal_staging env: - GH_TOKEN: ${{ secrets.GH_TOKEN }} + GH_TOKEN: ${{ secrets.GH_TOKEN || github.token }} diff --git a/.github/workflows/test-terraform-provider.yml b/.github/workflows/test-terraform-provider.yml index e46432e0e31..eb7b299fd1f 100644 --- a/.github/workflows/test-terraform-provider.yml +++ b/.github/workflows/test-terraform-provider.yml @@ -114,4 +114,4 @@ jobs: - name: Audit provider endpoints against the schema working-directory: terraform/provider - run: go run ./tools/endpointaudit -provider-dir ./litellm -spec "${RUNNER_TEMP}/openapi.json" + run: go run ./tools/endpointaudit -provider-dir ./litellm -spec "${RUNNER_TEMP}/openapi.json" -coverage-allowlist ./tools/endpointaudit/coverage_allowlist.txt diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 122bd82c657..962d1266fd7 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 18483 + "limit": 17271 }, "reportArgumentType": { - "limit": 2557 + "limit": 2539 }, "reportAssignmentType": { "limit": 319 @@ -12,25 +12,25 @@ "limit": 480 }, "reportCallIssue": { - "limit": 113 + "limit": 112 }, "reportConstantRedefinition": { "limit": 40 }, "reportDeprecated": { - "limit": 213 + "limit": 212 }, "reportDuplicateImport": { "limit": 19 }, "reportExplicitAny": { - "limit": 5960 + "limit": 5486 }, "reportFunctionMemberAccess": { "limit": 7 }, "reportGeneralTypeIssues": { - "limit": 105 + "limit": 101 }, "reportIncompatibleMethodOverride": { "limit": 56 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5659 + "limit": 5658 }, "reportMissingTypeArgument": { - "limit": 15482 + "limit": 15425 }, "reportMissingTypeStubs": { "limit": 40 @@ -72,7 +72,7 @@ "limit": 0 }, "reportOptionalMemberAccess": { - "limit": 1058 + "limit": 1055 }, "reportOptionalOperand": { "limit": 0 @@ -93,7 +93,7 @@ "limit": 213 }, "reportTypedDictNotRequiredAccess": { - "limit": 26 + "limit": 25 }, "reportUndefinedVariable": { "limit": 0 @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38779 + "limit": 38721 }, "reportUnknownParameterType": { - "limit": 19827 + "limit": 19778 }, "reportUnknownVariableType": { - "limit": 30348 + "limit": 30290 }, "reportUnnecessaryCast": { "limit": 117 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 831 + "limit": 829 }, "reportUntypedBaseClass": { "limit": 0 @@ -138,9 +138,9 @@ "limit": 138 }, "reportUnusedImport": { - "limit": 544 + "limit": 543 }, "reportUnusedVariable": { - "limit": 145 + "limit": 139 } } diff --git a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py index 18ac29b9781..8f4e999bb9d 100644 --- a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py @@ -7,7 +7,7 @@ GET - /audit/{id} - Get audit log by id GET - /audit - Get all audit logs """ -from typing import Any, Dict, List, Optional +from typing import TYPE_CHECKING, Final, Optional #### AUDIT LOGGING #### from fastapi import APIRouter, Depends, HTTPException, Query @@ -18,11 +18,16 @@ from litellm_enterprise.types.proxy.audit_logging_endpoints import ( from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.table_repositories import AuditLogRepository + +if TYPE_CHECKING: + from prisma import models as prisma_models router = APIRouter() -def _build_json_field_or_condition(json_key: str, value: str) -> Dict[str, Any]: +def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, object]: """ Build an OR condition that matches a value inside a JSON column at the given key, checking both before_value and updated_values. @@ -101,46 +106,37 @@ async def get_audit_logs( detail={"message": CommonProxyErrors.db_not_connected_error.value}, ) - # Build filter conditions - where_conditions: Dict[str, Any] = {} - if changed_by: - where_conditions["changed_by"] = changed_by - if changed_by_api_key: - where_conditions["changed_by_api_key"] = changed_by_api_key - if action: - where_conditions["action"] = action - if table_name: - where_conditions["table_name"] = table_name - if object_id: - where_conditions["object_id"] = object_id - if start_date or end_date: - date_filter: Dict[str, Any] = {} - if start_date: - date_filter["gte"] = start_date - if end_date: - date_filter["lte"] = end_date - where_conditions["updated_at"] = date_filter + date_filter: Final[dict[str, str]] = { + **({"gte": start_date} if start_date else {}), + **({"lte": end_date} if end_date else {}), + } # JSON field filters (PostgreSQL only) — each filter is AND'd with the # others, but checks both before_value and updated_values internally (OR). - if object_team_id: - where_conditions["AND"] = where_conditions.get("AND", []) + [ - _build_json_field_or_condition("team_id", object_team_id) - ] - if object_key_hash: - where_conditions["AND"] = where_conditions.get("AND", []) + [ - _build_json_field_or_condition("token", object_key_hash) - ] + json_field_conditions: Final[list[dict[str, object]]] = [ + *([_build_json_field_or_condition("team_id", object_team_id)] if object_team_id else []), + *([_build_json_field_or_condition("token", object_key_hash)] if object_key_hash else []), + ] - # Build sort conditions - order_by: Dict[str, Any] = {} - if sort_by and isinstance(sort_by, str): - order_by[sort_by] = sort_order - else: - order_by["updated_at"] = sort_order # Default sort by updated_at + # Build filter conditions + where_conditions: Final[dict[str, object]] = { + **({"changed_by": changed_by} if changed_by else {}), + **({"changed_by_api_key": changed_by_api_key} if changed_by_api_key else {}), + **({"action": action} if action else {}), + **({"table_name": table_name} if table_name else {}), + **({"object_id": object_id} if object_id else {}), + **({"updated_at": date_filter} if start_date or end_date else {}), + **({"AND": json_field_conditions} if json_field_conditions else {}), + } + + order_by: Final[dict[str, str]] = ( + {sort_by: sort_order} if sort_by and isinstance(sort_by, str) else {"updated_at": sort_order} + ) + + audit_log_table: Final[TableActions["prisma_models.LiteLLM_AuditLog"]] = AuditLogRepository(prisma_client).table # Get paginated results - audit_logs = await prisma_client.db.litellm_auditlog.find_many( + audit_logs: Final = await audit_log_table.find_many( where=where_conditions, order=order_by, skip=(page - 1) * page_size, @@ -148,13 +144,14 @@ async def get_audit_logs( ) # Get total count for pagination - total_count = await prisma_client.db.litellm_auditlog.count(where=where_conditions) - total_pages = -(-total_count // page_size) # Ceiling division + total_count: Final = await audit_log_table.count(where=where_conditions) + total_pages: Final = -(-total_count // page_size) # Ceiling division # Return paginated response return PaginatedAuditLogResponse( audit_logs=[ - AuditLogResponse(**audit_log.model_dump()) for audit_log in audit_logs + AuditLogResponse.model_validate(audit_log.model_dump()) + for audit_log in audit_logs ] if audit_logs else [], @@ -198,8 +195,10 @@ async def get_audit_log_by_id( detail={"message": CommonProxyErrors.db_not_connected_error.value}, ) + audit_log_table: Final[TableActions["prisma_models.LiteLLM_AuditLog"]] = AuditLogRepository(prisma_client).table + # Get the audit log by ID - audit_log = await prisma_client.db.litellm_auditlog.find_unique(where={"id": id}) + audit_log: Final = await audit_log_table.find_unique(where={"id": id}) if audit_log is None: raise HTTPException( @@ -207,4 +206,4 @@ async def get_audit_log_by_id( ) # Convert to response model - return AuditLogResponse(**audit_log.model_dump()) + return AuditLogResponse.model_validate(audit_log.model_dump()) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 39f8de0b0cc..cf2cee9b6ef 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -473,19 +473,56 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) page_size: Final = min(limit or 20, 100) - cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": after}, "skip": 1} if after else {} - - batches = await _managed_object_table(self.prisma_client).find_many( - where=where_clause, - take=page_size + 1, - order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], - **cursor_args, + matches: Final = await self._collect_listed_batches( + where_clause=where_clause, + after=after, + wanted=page_size + 1, + user_api_key_dict=user_api_key_dict, ) + return build_list_page(list(matches[:page_size]), has_more=len(matches) > page_size) - has_more = len(batches) > page_size + async def _collect_listed_batches( + self, + where_clause: Mapping[str, object], + after: Optional[str], + wanted: int, + user_api_key_dict: UserAPIKeyAuth, + ) -> tuple[LiteLLMBatch, ...]: + """Read chunks newest-first until ``wanted`` batches survive parsing and + file-id resolution or the caller's rows run out, so a run of rows that will + not parse refills the page instead of emptying it. The first chunk is + ``wanted`` rows, so a healthy page still costs one query; a scan that has to + continue widens to ``FILE_LIST_CONTINUATION_CHUNK_SIZE`` like ``afile_list``, + and every chunk advances the keyset cursor, so the walk ends once the + caller's rows are exhausted.""" + matches: tuple[LiteLLMBatch, ...] = () # rebind-ok: accumulates survivors across chunks + cursor_id: Optional[str] = after # rebind-ok: keyset cursor advances to each chunk's last row + chunk_size: int = wanted # rebind-ok: widens once a scan has to continue past the first chunk + while len(matches) < wanted: + cursor_args: _CursorPageArgs = {"cursor": {"unified_object_id": cursor_id}, "skip": 1} if cursor_id else {} + chunk = await _managed_object_table(self.prisma_client).find_many( + where=where_clause, + take=chunk_size, + order=[{"created_at": "desc"}, {"unified_object_id": "desc"}], + **cursor_args, + ) + matches = matches + await self._resolve_listed_rows( + rows=chunk, wanted=wanted - len(matches), user_api_key_dict=user_api_key_dict + ) + if len(chunk) < chunk_size: + break + cursor_id = chunk[-1].unified_object_id + chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE) + return matches + async def _resolve_listed_rows( + self, + rows: "Sequence[PrismaManagedObjectRow]", + wanted: int, + user_api_key_dict: UserAPIKeyAuth, + ) -> tuple[LiteLLMBatch, ...]: parsed_rows: Final = tuple( - (row, batch_obj) for row in batches[:page_size] if (batch_obj := _parse_managed_batch_row(row)) is not None + (row, batch_obj) for row in rows if (batch_obj := _parse_managed_batch_row(row)) is not None ) unified_id_by_raw_id: Final = await map_raw_file_ids_to_unified( raw_file_ids=frozenset( @@ -496,19 +533,19 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ), prisma_client=self.prisma_client, ) - resolved_batches: Final = [ - await self._resolve_listed_batch( + resolved: Final[list[LiteLLMBatch]] = [] # mutable-ok: resolution stops as soon as the page is full + for row, batch_obj in parsed_rows: + if len(resolved) == wanted: + break + resolved_batch = await self._resolve_listed_batch( row=row, batch_obj=batch_obj, unified_id_by_raw_id=unified_id_by_raw_id, user_api_key_dict=user_api_key_dict, ) - for row, batch_obj in parsed_rows - ] - return build_list_page( - [batch_obj for batch_obj in resolved_batches if batch_obj is not None], - has_more=has_more, - ) + if resolved_batch is not None: + resolved.append(resolved_batch) + return tuple(resolved) async def _resolve_listed_batch( self, diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index 5ec482c385a..b2eda76f9ae 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -15,6 +15,7 @@ from collections.abc import Sequence from typing import TYPE_CHECKING, Final from fastapi import APIRouter, Depends, HTTPException, Request +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -26,39 +27,50 @@ from litellm.proxy.management_helpers.utils import ( management_endpoint_wrapper, ) from litellm.proxy.utils import PrismaClient, handle_exception_on_proxy +from litellm.repositories.budget_repository import BudgetRepository +from litellm.repositories.object_permission_repository import ObjectPermissionRepository +from litellm.repositories.prisma_protocols import TableActions +from litellm.repositories.project_repository import ProjectRepository +from litellm.repositories.team_repository import TeamRepository +from litellm.repositories.user_repository import UserRepository +from litellm.repositories.verification_token_repository import VerificationTokenRepository if TYPE_CHECKING: from prisma import models as prisma_models - from prisma.actions import ( - LiteLLM_ProjectTableActions, - LiteLLM_TeamTableActions, - LiteLLM_VerificationTokenActions, - ) from litellm import Router router = APIRouter() - -def _team_table(prisma_client: PrismaClient) -> "LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]": - team_table: LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable] = prisma_client.db.litellm_teamtable - return team_table +_OBJECT_PERMISSION_PAYLOAD: Final = TypeAdapter(dict[str, object]) -def _project_table(prisma_client: PrismaClient) -> "LiteLLM_ProjectTableActions[prisma_models.LiteLLM_ProjectTable]": - project_table: LiteLLM_ProjectTableActions[prisma_models.LiteLLM_ProjectTable] = ( - prisma_client.db.litellm_projecttable - ) - return project_table +def _team_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_TeamTable"]: + return TeamRepository(prisma_client).table + + +def _project_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_ProjectTable"]: + return ProjectRepository(prisma_client).table def _verification_token_table( prisma_client: PrismaClient, -) -> "LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]": - verification_token_table: LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken] = ( - prisma_client.db.litellm_verificationtoken - ) - return verification_token_table +) -> TableActions["prisma_models.LiteLLM_VerificationToken"]: + return VerificationTokenRepository(prisma_client).table + + +def _budget_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_BudgetTable"]: + return BudgetRepository(prisma_client).table + + +def _object_permission_table( + prisma_client: PrismaClient, +) -> TableActions["prisma_models.LiteLLM_ObjectPermissionTable"]: + return ObjectPermissionRepository(prisma_client).table + + +def _user_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_UserTable"]: + return UserRepository(prisma_client).table def _jsonified(prisma_client: PrismaClient, payload: dict[str, object]) -> dict[str, object]: @@ -329,7 +341,7 @@ async def _create_budget_for_project( new_budget = _jsonified(prisma_client, budget_row.model_dump(exclude_none=True)) - _budget: prisma_models.LiteLLM_BudgetTable = await prisma_client.db.litellm_budgettable.create( + _budget: Final = await _budget_table(prisma_client).create( data={ **new_budget, "created_by": user_id or litellm_proxy_admin_name, @@ -352,10 +364,8 @@ async def _set_project_object_permission( return None if data.object_permission is not None: - created_object_permission: prisma_models.LiteLLM_ObjectPermissionTable = ( - await prisma_client.db.litellm_objectpermissiontable.create( - data=data.object_permission.model_dump(exclude_none=True), - ) + created_object_permission: Final = await _object_permission_table(prisma_client).create( + data=data.object_permission.model_dump(exclude_none=True), ) del data.object_permission return created_object_permission.object_permission_id @@ -586,10 +596,8 @@ async def new_project( new_project_row = _remove_budget_fields_from_project_data(new_project_row) verbose_proxy_logger.info(f"new_project_row: {json.dumps(new_project_row, indent=2)}") - response: prisma_models.LiteLLM_ProjectTable = await prisma_client.db.litellm_projecttable.create( - data={ - **new_project_row, # type: ignore - }, + response: Final = await _project_table(prisma_client).create( + data={**new_project_row}, include={"litellm_budget_table": True}, ) @@ -776,7 +784,7 @@ async def update_project( if budget_updates and existing_project.budget_id: # Update existing budget - await prisma_client.db.litellm_budgettable.update( + await _budget_table(prisma_client).update( where={"budget_id": existing_project.budget_id}, data={ **budget_updates, @@ -791,18 +799,17 @@ async def update_project( if "object_permission" in update_data: object_permission_data = update_data.pop("object_permission") if object_permission_data: + object_permission_payload: Final = _OBJECT_PERMISSION_PAYLOAD.validate_python(object_permission_data) if existing_project.object_permission_id: # Update existing permission - await prisma_client.db.litellm_objectpermissiontable.update( + await _object_permission_table(prisma_client).update( where={"object_permission_id": existing_project.object_permission_id}, - data=object_permission_data, + data=object_permission_payload, ) else: # Create new permission - created_permission: prisma_models.LiteLLM_ObjectPermissionTable = ( - await prisma_client.db.litellm_objectpermissiontable.create( - data=object_permission_data, - ) + created_permission: Final = await _object_permission_table(prisma_client).create( + data=object_permission_payload, ) update_data["object_permission_id"] = created_permission.object_permission_id @@ -818,7 +825,7 @@ async def update_project( update_data = _remove_budget_fields_from_project_data(update_data) # Update project - updated_project: prisma_models.LiteLLM_ProjectTable | None = await prisma_client.db.litellm_projecttable.update( + updated_project: Final = await _project_table(prisma_client).update( where={"project_id": data.project_id}, data=update_data, include={"litellm_budget_table": True, "object_permission": True}, @@ -1058,7 +1065,7 @@ async def list_projects( # Look up the user's team memberships via the reverse-index on # LiteLLM_UserTable.teams (maintained by team_member_add alongside # members_with_roles). This avoids a full scan of all team rows. - user_record: prisma_models.LiteLLM_UserTable | None = await prisma_client.db.litellm_usertable.find_unique( + user_record: Final = await _user_table(prisma_client).find_unique( where={"user_id": user_api_key_dict.user_id}, ) user_team_ids: list[str] = user_record.teams if user_record is not None and user_record.teams else [] diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ce28f737334..4388e561026 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2,6 +2,36 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "aho-corasick" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c982642fa9e8606056828ee9a8505737230110bb1099153c79efe865c59d12ba" +dependencies = [ + "memchr", +] + +[[package]] +name = "alloca" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5a7d05ea6aea7e9e64d25b9156ba2fee3fdd659e34e41063cd2fc7cd020d7f4" +dependencies = [ + "cc", +] + +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + [[package]] name = "arc-swap" version = "1.9.2" @@ -506,6 +536,12 @@ dependencies = [ "either", ] +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + [[package]] name = "cc" version = "1.3.0" @@ -541,6 +577,58 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + +[[package]] +name = "clap" +version = "4.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889" +dependencies = [ + "anstyle", + "clap_lex", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + [[package]] name = "cmake" version = "0.1.58" @@ -596,6 +684,72 @@ dependencies = [ "libc", ] +[[package]] +name = "criterion" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "950046b2aa2492f9a536f5f4f9a3de7b9e2476e575e05bd6c333371add4d98f3" +dependencies = [ + "alloca", + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "itertools", + "num-traits", + "oorandom", + "page_size", + "plotters", + "rayon", + "regex", + "serde", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8d80a2f4f5b554395e47b5d8305bc3d27813bacb73493eb1001e8f76dae29ea" +dependencies = [ + "cast", + "itertools", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "crypto-common" version = "0.1.7" @@ -856,6 +1010,17 @@ dependencies = [ "tracing", ] +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.17.1" @@ -1179,6 +1344,15 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" +[[package]] +name = "itertools" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "413ee7dfc52ee1a4949ceeb7dbc8a33f2d6c088194d9f922fb8318faf1f01186" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -1255,10 +1429,13 @@ dependencies = [ name = "litellm-python-bridge" version = "0.1.0" dependencies = [ + "criterion", "litellm-ai-gateway", "litellm-core", "pyo3", "pyo3-async-runtimes", + "pythonize", + "serde", "serde_json", "tokio", ] @@ -1340,6 +1517,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + [[package]] name = "openssl-probe" version = "0.2.1" @@ -1352,6 +1535,16 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" +[[package]] +name = "page_size" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d5b2194ed13191c1999ae0704b7839fb18384fa22e49b57eeaa97d79ce40da" +dependencies = [ + "libc", + "winapi", +] + [[package]] name = "percent-encoding" version = "2.3.2" @@ -1376,6 +1569,34 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" +[[package]] +name = "plotters" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" +dependencies = [ + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "plotters-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + [[package]] name = "portable-atomic" version = "1.14.0" @@ -1486,6 +1707,16 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "pythonize" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ec376e1216e0c929a74964ce2020012a1a39f32d80e78aa688721219ea7fb89" +dependencies = [ + "pyo3", + "serde", +] + [[package]] name = "quinn" version = "0.11.11" @@ -1613,12 +1844,61 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad8553b9b26413251cbf30e620595c7a41b3887f03da04579c0e6b0d6a06b4b2" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + [[package]] name = "regex-lite" version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + [[package]] name = "reqwest" version = "0.12.28" @@ -1774,6 +2054,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "schannel" version = "0.1.29" @@ -2099,6 +2388,16 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "tinyvec" version = "1.12.0" @@ -2363,6 +2662,16 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "want" version = "0.3.1" @@ -2475,6 +2784,37 @@ dependencies = [ "rustls-pki-types", ] +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + [[package]] name = "windows-link" version = "0.2.1" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 6d63be05d00..481ea3f8f66 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -18,6 +18,7 @@ litellm-ai-gateway = { path = "crates/ai-gateway", default-features = false } axum = "0.7" pyo3 = "0.29.0" pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } +pythonize = "0.29.0" rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls", "http2", "stream"] } serde = { version = "1.0", features = ["derive"] } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 20a9ba789ce..0c4a753f762 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -9,10 +9,23 @@ repository.workspace = true name = "_native" crate-type = ["cdylib"] +[features] +default = ["extension-module"] +extension-module = ["pyo3/extension-module"] + [dependencies] litellm-core = { workspace = true, features = ["bedrock-auth"] } litellm-ai-gateway = { workspace = true, default-features = false } -pyo3 = { workspace = true, features = ["extension-module"] } +pyo3.workspace = true pyo3-async-runtimes.workspace = true +pythonize.workspace = true +serde.workspace = true serde_json.workspace = true tokio.workspace = true + +[dev-dependencies] +criterion = "0.8.2" + +[[bench]] +name = "serialization" +harness = false diff --git a/litellm-rust/crates/python-bridge/benches/serialization.rs b/litellm-rust/crates/python-bridge/benches/serialization.rs new file mode 100644 index 00000000000..8a90cf667d0 --- /dev/null +++ b/litellm-rust/crates/python-bridge/benches/serialization.rs @@ -0,0 +1,103 @@ +use std::hint::black_box; +use std::time::Duration; + +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use pyo3::prelude::*; +use pyo3::types::PyDict; +use serde_json::{Value, json}; + +const PAYLOAD_SIZES: &[(&str, usize)] = &[ + ("1_KiB", 1024), + ("64_KiB", 64 * 1024), + ("1_MiB", 1024 * 1024), + ("4_MiB", 4 * 1024 * 1024), + ("16_MiB", 16 * 1024 * 1024), +]; + +fn former_json_roundtrip_from_py(py: Python<'_>, value: &Bound<'_, PyAny>) -> Value { + let json = py.import("json").expect("Python json module should import"); + let encoded: String = json + .call_method1("dumps", (value,)) + .expect("payload should serialize") + .extract() + .expect("json.dumps should return a string"); + serde_json::from_str(&encoded).expect("serialized JSON should parse") +} + +fn pythonize_from_py(value: &Bound<'_, PyAny>) -> Value { + pythonize::depythonize(value).expect("payload should depythonize") +} + +fn former_json_roundtrip_to_py(py: Python<'_>, value: &Value) -> Py { + let json = py.import("json").expect("Python json module should import"); + let encoded = serde_json::to_string(value).expect("response should serialize"); + json.call_method1("loads", (encoded,)) + .expect("serialized response should parse in Python") + .unbind() +} + +fn pythonize_to_py(py: Python<'_>, value: &Value) -> Py { + pythonize::pythonize(py, value) + .expect("response should pythonize") + .unbind() +} + +fn serialization(c: &mut Criterion) { + Python::initialize(); + Python::attach(|py| { + for &(label, payload_bytes) in PAYLOAD_SIZES { + let data_uri = format!("data:image/png;base64,{}", "A".repeat(payload_bytes)); + let document = PyDict::new(py); + document + .set_item("type", "image_url") + .expect("document type should be set"); + document + .set_item("image_url", &data_uri) + .expect("document URL should be set"); + let response = json!({ + "pages": [{ + "index": 0, + "markdown": "OCR text", + "images": [{"image_base64": data_uri}], + }], + "model": "mistral-ocr-latest", + "document_annotation": null, + "usage_info": {"pages_processed": 1}, + "object": "ocr", + }); + + c.bench_with_input( + BenchmarkId::new("python_to_rust_json", label), + &document, + |b, document| { + b.iter(|| former_json_roundtrip_from_py(py, black_box(document.as_any()))) + }, + ); + c.bench_with_input( + BenchmarkId::new("python_to_rust_pythonize", label), + &document, + |b, document| b.iter(|| pythonize_from_py(black_box(document.as_any()))), + ); + c.bench_with_input( + BenchmarkId::new("rust_to_python_json", label), + &response, + |b, response| b.iter(|| former_json_roundtrip_to_py(py, black_box(response))), + ); + c.bench_with_input( + BenchmarkId::new("rust_to_python_pythonize", label), + &response, + |b, response| b.iter(|| pythonize_to_py(py, black_box(response))), + ); + } + }); +} + +criterion_group! { + name = benches; + config = Criterion::default() + .sample_size(20) + .warm_up_time(Duration::from_secs(1)) + .measurement_time(Duration::from_secs(4)); + targets = serialization +} +criterion_main!(benches); diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index c6f81cf6916..f9e75f45f75 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -19,6 +19,9 @@ use pyo3::types::{PyAny, PyDict}; use serde_json::{Map, Value}; mod gil; +mod marshal; + +use marshal::{from_py, to_py}; pyo3::create_exception!( _native, @@ -41,35 +44,18 @@ type MarshaledOcrInputs = ( Option, ); -fn py_to_json(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult { - let json = py.import("json")?; - let encoded: String = json.call_method1("dumps", (value,))?.extract()?; - serde_json::from_str(&encoded).map_err(|err| PyValueError::new_err(err.to_string())) -} - -fn json_to_py(py: Python<'_>, value: Value) -> PyResult> { - let json = py.import("json")?; - let encoded = - serde_json::to_string(&value).map_err(|err| PyValueError::new_err(err.to_string()))?; - Ok(json.call_method1("loads", (encoded,))?.unbind()) -} - fn messages_response_to_py( py: Python<'_>, response: AnthropicMessagesResponse, ) -> PyResult> { - let value = - serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?; - json_to_py(py, value) + to_py(py, &response) } fn chat_completions_response_to_py( py: Python<'_>, response: ChatCompletionsResponse, ) -> PyResult> { - let value = - serde_json::to_value(response).map_err(|err| PyValueError::new_err(err.to_string()))?; - json_to_py(py, value) + to_py(py, &response) } fn core_error_to_pyerr(err: CoreError) -> PyErr { @@ -116,7 +102,7 @@ fn optional_object_to_map( value: Option>, ) -> PyResult> { match value { - Some(value) => match py_to_json(py, value.bind(py))? { + Some(value) => match from_py(value.bind(py))? { Value::Object(map) => Ok(map), _ => Err(PyValueError::new_err(format!("{name} must be a dict"))), }, @@ -139,7 +125,7 @@ fn marshal_headers( headers: Option>, ) -> PyResult> { let value = match headers { - Some(headers) => py_to_json(py, headers.bind(py))?, + Some(headers) => from_py(headers.bind(py))?, None => Value::Object(Map::new()), }; let Value::Object(headers) = value else { @@ -211,7 +197,7 @@ fn marshal_inputs( optional_params: Option>, timeout_seconds: Option, ) -> PyResult { - let document = py_to_json(py, document.bind(py))?; + let document = from_py(document.bind(py))?; let extra_headers = match extra_headers { Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), None => None, @@ -262,7 +248,7 @@ fn ocr( }); match result { - Ok(value) => json_to_py(py, value), + Ok(value) => to_py(py, &value), Err(err) => Err(core_error_to_pyerr(err)), } } @@ -307,7 +293,7 @@ fn aocr( .await .map_err(core_error_to_pyerr)?; - Python::attach(|py| json_to_py(py, value)) + Python::attach(|py| to_py(py, &value)) }) } @@ -325,7 +311,7 @@ fn transcription( optional_params: Option>, timeout_seconds: Option, ) -> PyResult> { - let audio = py_to_json(py, audio.bind(py))?; + let audio = from_py(audio.bind(py))?; let extra_headers = match extra_headers { Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), None => None, @@ -351,7 +337,7 @@ fn transcription( )) }); match result { - Ok(value) => json_to_py(py, value), + Ok(value) => to_py(py, &value), Err(err) => Err(core_error_to_pyerr(err)), } } @@ -370,7 +356,7 @@ fn atranscription( optional_params: Option>, timeout_seconds: Option, ) -> PyResult> { - let audio = py_to_json(py, audio.bind(py))?; + let audio = from_py(audio.bind(py))?; let extra_headers = match extra_headers { Some(headers) => Some(optional_object_to_map(py, "extra_headers", Some(headers))?), None => None, @@ -394,7 +380,7 @@ fn atranscription( }) .await .map_err(core_error_to_pyerr)?; - Python::attach(|py| json_to_py(py, value)) + Python::attach(|py| to_py(py, &value)) }) } @@ -406,7 +392,7 @@ fn marshal_messages_inputs( extra_headers: Option>, timeout_seconds: Option, ) -> PyResult { - let body = py_to_json(py, body.bind(py))?; + let body: Value = from_py(body.bind(py))?; if !body.is_object() { return Err(PyValueError::new_err("body must be a dict")); } @@ -498,7 +484,7 @@ fn marshal_chat_completions_inputs( extra_headers: Option>, timeout_seconds: Option, ) -> PyResult { - let messages = py_to_json(py, messages.bind(py))?; + let messages: Value = from_py(messages.bind(py))?; if !messages.is_array() { return Err(PyValueError::new_err("messages must be a list")); } @@ -527,7 +513,7 @@ fn chat_completions_decline( optional_params: Option>, custom_llm_provider: Option, ) -> PyResult> { - let messages = py_to_json(py, messages.bind(py))?; + let messages = from_py(messages.bind(py))?; let optional_params = optional_object_to_map(py, "optional_params", optional_params)?; Ok(chat_completions_decline_reason( &model, diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs new file mode 100644 index 00000000000..c3d0638427c --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -0,0 +1,20 @@ +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use serde::Serialize; +use serde::de::DeserializeOwned; + +pub fn from_py(value: &Bound<'_, PyAny>) -> PyResult +where + T: DeserializeOwned, +{ + pythonize::depythonize(value).map_err(|error| PyValueError::new_err(error.to_string())) +} + +pub fn to_py(py: Python<'_>, value: &T) -> PyResult> +where + T: Serialize + ?Sized, +{ + pythonize::pythonize(py, value) + .map(Bound::unbind) + .map_err(|error| PyValueError::new_err(error.to_string())) +} diff --git a/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs new file mode 100644 index 00000000000..6a6ede22e85 --- /dev/null +++ b/litellm-rust/crates/python-bridge/tests/marshal_boundary.rs @@ -0,0 +1,52 @@ +use std::fs; +use std::path::{Path, PathBuf}; + +const DISALLOWED_OUTSIDE_MARSHAL: &[&str] = &[ + "py.import(\"json\")", + "pythonize::", + "serde_json::to_string", + "serde_json::from_str", +]; + +fn source_root() -> PathBuf { + Path::new(env!("CARGO_MANIFEST_DIR")).join("src") +} + +fn rust_sources(directory: &Path) -> Vec { + fs::read_dir(directory) + .expect("bridge source directory should be readable") + .map(|entry| { + entry + .expect("bridge source entry should be readable") + .path() + }) + .flat_map(|path| { + if path.is_dir() { + rust_sources(&path) + } else if path.extension().is_some_and(|extension| extension == "rs") { + vec![path] + } else { + Vec::new() + } + }) + .collect() +} + +#[test] +fn serialization_is_centralized_in_marshal_module() { + let root = source_root(); + + for path in rust_sources(&root) { + if path == root.join("marshal.rs") { + continue; + } + let source = fs::read_to_string(&path).expect("bridge source should be readable"); + for disallowed in DISALLOWED_OUTSIDE_MARSHAL { + assert!( + !source.contains(disallowed), + "{} bypasses the typed marshal module with `{disallowed}`", + path.display() + ); + } + } +} diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 933464d3f23..d7e00a81b38 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -18,7 +18,10 @@ until they're actually needed. import importlib import sys from collections.abc import Callable -from typing import Any, Final, cast +from types import ModuleType +from typing import TYPE_CHECKING, Any, Final, cast + +from typing_extensions import ReadOnly, TypedDict # Import all the data structures that define what can be lazy-loaded # These are just lists of names and maps of where to find them @@ -53,6 +56,9 @@ from ._lazy_imports_registry import ( UTILS_NAMES, ) +if TYPE_CHECKING: + from tiktoken import Encoding + def get_litellm_globals() -> dict: """ @@ -78,10 +84,10 @@ def _get_utils_globals() -> dict: # They're separate from the main lazy import system because they have specific use cases # Lazy loader for default encoding - avoids importing heavy tiktoken library at startup -_default_encoding: Any | None = None +_default_encoding: "Encoding | None" = None -def _get_default_encoding() -> Any: +def _get_default_encoding() -> "Encoding": """ Lazily load and cache the default OpenAI encoding. @@ -100,10 +106,10 @@ def _get_default_encoding() -> Any: # Lazy loader for get_modified_max_tokens to avoid importing token_counter at module import time -_get_modified_max_tokens_func: Any | None = None +_get_modified_max_tokens_func: "Callable[..., int | None] | None" = None -def _get_modified_max_tokens() -> Any: +def _get_modified_max_tokens() -> "Callable[..., int | None]": """ Lazily load and cache the get_modified_max_tokens function. @@ -124,10 +130,10 @@ def _get_modified_max_tokens() -> Any: # Lazy loader for token_counter to avoid importing token_counter module at module import time -_token_counter_new_func: Any | None = None +_token_counter_new_func: "Callable[..., int] | None" = None -def _get_token_counter_new() -> Any: +def _get_token_counter_new() -> "Callable[..., int]": """ Lazily load and cache the token_counter function (aliased as token_counter_new). @@ -154,10 +160,10 @@ def _get_token_counter_new() -> Any: # This registry maps attribute names (like "ModelResponse") to handler functions # It's built once the first time someone accesses a lazy-loaded attribute # Example: {"ModelResponse": _lazy_import_utils, "Cache": _lazy_import_caching, ...} -_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], Any]] | None = None +_LAZY_IMPORT_REGISTRY: dict[str, Callable[[str], object]] | None = None -def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]: +def _get_lazy_import_registry() -> dict[str, Callable[[str], object]]: """ Build the registry that maps attribute names to their handler functions. @@ -206,7 +212,18 @@ def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]: return _LAZY_IMPORT_REGISTRY -def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> Any: +class _AttributeView(TypedDict): + """Holds one module attribute so the lazily fetched value is read back as ``object``.""" + + value: ReadOnly[object] + + +def _module_attribute(module: ModuleType, attr_name: str) -> object: + attribute: Final[_AttributeView] = {"value": getattr(module, attr_name)} + return attribute["value"] + + +def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> object: """ Generic function that handles lazy importing for most attributes. @@ -255,7 +272,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate # Step 6: Get the actual attribute from the module # Example: getattr(utils_module, "ModelResponse") returns the ModelResponse class - value: Final = getattr(module, attr_name) + value: Final = _module_attribute(module, attr_name) # Step 7: Cache it so we don't have to import again next time _globals[name] = value @@ -272,62 +289,62 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate # The registry (above) maps attribute names to these handler functions. -def _lazy_import_utils(name: str) -> Any: +def _lazy_import_utils(name: str) -> object: """Handler for utils module attributes (ModelResponse, token_counter, etc.)""" return _generic_lazy_import(name, _UTILS_IMPORT_MAP, "Utils") -def _lazy_import_cost_calculator(name: str) -> Any: +def _lazy_import_cost_calculator(name: str) -> object: """Handler for cost calculator functions (completion_cost, cost_per_token, etc.)""" return _generic_lazy_import(name, _COST_CALCULATOR_IMPORT_MAP, "Cost calculator") -def _lazy_import_token_counter(name: str) -> Any: +def _lazy_import_token_counter(name: str) -> object: """Handler for token counter utilities""" return _generic_lazy_import(name, _TOKEN_COUNTER_IMPORT_MAP, "Token counter") -def _lazy_import_bedrock_types(name: str) -> Any: +def _lazy_import_bedrock_types(name: str) -> object: """Handler for Bedrock type aliases""" return _generic_lazy_import(name, _BEDROCK_TYPES_IMPORT_MAP, "Bedrock types") -def _lazy_import_types_utils(name: str) -> Any: +def _lazy_import_types_utils(name: str) -> object: """Handler for types from litellm.types.utils (BudgetConfig, ImageObject, etc.)""" return _generic_lazy_import(name, _TYPES_UTILS_IMPORT_MAP, "Types utils") -def _lazy_import_caching(name: str) -> Any: +def _lazy_import_caching(name: str) -> object: """Handler for caching classes (Cache, DualCache, RedisCache, etc.)""" return _generic_lazy_import(name, _CACHING_IMPORT_MAP, "Caching") -def _lazy_import_dotprompt(name: str) -> Any: +def _lazy_import_dotprompt(name: str) -> object: """Handler for dotprompt integration globals""" return _generic_lazy_import(name, _DOTPROMPT_IMPORT_MAP, "Dotprompt") -def _lazy_import_types(name: str) -> Any: +def _lazy_import_types(name: str) -> object: """Handler for type classes (GuardrailItem, etc.)""" return _generic_lazy_import(name, _TYPES_IMPORT_MAP, "Types") -def _lazy_import_llm_configs(name: str) -> Any: +def _lazy_import_llm_configs(name: str) -> object: """Handler for LLM config classes (AnthropicConfig, OpenAILikeChatConfig, etc.)""" return _generic_lazy_import(name, _LLM_CONFIGS_IMPORT_MAP, "LLM config") -def _lazy_import_litellm_logging(name: str) -> Any: +def _lazy_import_litellm_logging(name: str) -> object: """Handler for litellm_logging module (Logging, modify_integration)""" return _generic_lazy_import(name, _LITELLM_LOGGING_IMPORT_MAP, "Litellm logging") -def _lazy_import_llm_provider_logic(name: str) -> Any: +def _lazy_import_llm_provider_logic(name: str) -> object: """Handler for LLM provider logic functions (get_llm_provider, etc.)""" return _generic_lazy_import(name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic") -def _lazy_import_utils_module(name: str) -> Any: +def _lazy_import_utils_module(name: str) -> object: """ Handler for utils module lazy imports. @@ -355,7 +372,7 @@ def _lazy_import_utils_module(name: str) -> Any: module = importlib.import_module(module_path) # Get the actual attribute from the module - value: Final = getattr(module, attr_name) + value: Final = _module_attribute(module, attr_name) # Cache it so we don't have to import again next time _globals[name] = value @@ -370,7 +387,7 @@ def _lazy_import_utils_module(name: str) -> Any: # These handlers have custom logic that doesn't fit the generic pattern -def _lazy_import_llm_client_cache(name: str) -> Any: +def _lazy_import_llm_client_cache(name: str) -> object: """ Handler for LLM client cache - has special logic for singleton instance. @@ -386,8 +403,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any: return _globals[name] # Import the class - module: Final = importlib.import_module("litellm.caching.llm_caching_handler") - LLMClientCache: Final = getattr(module, "LLMClientCache") + from litellm.caching.llm_caching_handler import LLMClientCache # If they want the class itself, return it if name == "LLMClientCache": @@ -403,7 +419,7 @@ def _lazy_import_llm_client_cache(name: str) -> Any: raise AttributeError(f"LLM client cache lazy import: unknown attribute {name!r}") -def _lazy_import_http_handlers(name: str) -> Any: +def _lazy_import_http_handlers(name: str) -> object: """ Handler for HTTP clients - has special logic for creating client instances. diff --git a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py index 15cf77708f9..838c0fd8373 100644 --- a/litellm/a2a_protocol/litellm_completion_bridge/transformation.py +++ b/litellm/a2a_protocol/litellm_completion_bridge/transformation.py @@ -17,11 +17,27 @@ A2A Streaming Events: - Artifact update (kind: "artifact-update") - Content/artifact delivery """ +from collections.abc import Mapping, MutableMapping, Sequence from datetime import datetime, timezone -from typing import Any, Final +from typing import TYPE_CHECKING, Final from uuid import uuid4 +from pydantic import JsonValue, TypeAdapter, ValidationError + from litellm._logging import verbose_logger +from litellm.types.utils import ModelResponse + +if TYPE_CHECKING: + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + +_STR_KEY_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + + +def _as_object_mapping(value: object) -> Mapping[str, object]: + try: + return _STR_KEY_MAPPING_ADAPTER.validate_python(value) + except ValidationError: + return {} class A2AStreamingContext: @@ -30,7 +46,7 @@ class A2AStreamingContext: Tracks task_id, context_id, and message accumulation. """ - def __init__(self, request_id: str, input_message: dict[str, Any]): + def __init__(self, request_id: str, input_message: Mapping[str, JsonValue]): self.request_id = request_id self.task_id = str(uuid4()) self.context_id = str(uuid4()) @@ -46,44 +62,46 @@ class A2ACompletionBridgeTransformation: """ @staticmethod - def _extract_text_from_a2a_parts(parts: list[dict[str, Any]]) -> str: + def _text_from_a2a_part(part: JsonValue) -> str | None: + if not isinstance(part, dict): + return None + text: Final = part.get("text") + if text is None: + return None + if part.get("kind") not in (None, "", "text"): + return None + return str(text) + + @staticmethod + def _extract_text_from_a2a_parts(parts: Sequence[JsonValue]) -> str: """Extract text from A2A parts (with or without explicit ``kind``).""" - content_parts: Final[list[str]] = [] - for part in parts: - if not isinstance(part, dict): - continue - kind = part.get("kind") - text = part.get("text") - if text is None: - continue - if kind in (None, "", "text"): - content_parts.append(str(text)) - return "\n".join(content_parts) + extracted: Final = (A2ACompletionBridgeTransformation._text_from_a2a_part(part) for part in parts) + return "\n".join(text for text in extracted if text is not None) @staticmethod def get_forward_metadata( - a2a_message: dict[str, Any], - params: dict[str, Any] | None = None, - ) -> dict[str, Any] | None: + a2a_message: Mapping[str, JsonValue], + params: Mapping[str, JsonValue] | None = None, + ) -> Mapping[str, JsonValue] | None: """ Merge A2A metadata from MessageSendParams and the message for downstream providers. Forwarded once on the LangGraph run payload (``metadata``), not duplicated on each input message — see ``apply_forward_metadata_to_completion_params``. """ - merged: Final[dict[str, Any]] = {} - if params and isinstance(params.get("metadata"), dict): - merged.update(params["metadata"]) + params_metadata: Final = params.get("metadata") if params else None message_metadata: Final = a2a_message.get("metadata") - if isinstance(message_metadata, dict): - merged.update(message_metadata) + merged: Final[dict[str, JsonValue]] = { + **(params_metadata if isinstance(params_metadata, dict) else {}), + **(message_metadata if isinstance(message_metadata, dict) else {}), + } return merged or None @staticmethod def apply_forward_metadata_to_completion_params( - completion_params: dict[str, Any], - a2a_message: dict[str, Any], - params: dict[str, Any] | None = None, + completion_params: MutableMapping[str, object], + a2a_message: Mapping[str, JsonValue], + params: Mapping[str, JsonValue] | None = None, ) -> None: """ Attach A2A metadata to completion kwargs for provider bridges (e.g. LangGraph). @@ -97,24 +115,20 @@ class A2ACompletionBridgeTransformation: if not forward_metadata: return - extra_body = completion_params.get("extra_body") - if not isinstance(extra_body, dict): - extra_body = {} + extra_body: Final = _as_object_mapping(completion_params.get("extra_body")) # Layer client-supplied A2A metadata under any agent-owner-configured # ``extra_body.metadata`` so the configured keys remain authoritative # and an A2A caller cannot overwrite server-set run metadata. - existing_metadata: Final = extra_body.get("metadata") - existing_dict: Final[dict[str, Any]] = existing_metadata if isinstance(existing_metadata, dict) else {} - merged_metadata: Final[dict[str, Any]] = {**forward_metadata, **existing_dict} - extra_body = {**extra_body, "metadata": merged_metadata} - completion_params["extra_body"] = extra_body + existing_dict: Final = _as_object_mapping(extra_body.get("metadata")) + merged_metadata: Final[dict[str, object]] = {**forward_metadata, **existing_dict} + completion_params["extra_body"] = {**extra_body, "metadata": merged_metadata} verbose_logger.debug("A2A -> completion forward metadata keys=%s", list(forward_metadata.keys())) @staticmethod def a2a_message_to_openai_messages( - a2a_message: dict[str, Any], - ) -> list[dict[str, Any]]: + a2a_message: Mapping[str, JsonValue], + ) -> list[dict[str, object]]: """ Transform an A2A message to OpenAI message format. @@ -125,25 +139,19 @@ class A2ACompletionBridgeTransformation: List of OpenAI-format messages """ role: Final = a2a_message.get("role", "user") - parts = a2a_message.get("parts", []) + raw_parts: Final = a2a_message.get("parts", []) # Map A2A roles to OpenAI roles - openai_role = role - if role == "user": - openai_role = "user" - elif role == "assistant": - openai_role = "assistant" - elif role == "system": - openai_role = "system" - - if not isinstance(parts, list): - parts = [] + openai_role: Final = ( + "user" if role == "user" else "assistant" if role == "assistant" else "system" if role == "system" else role + ) + parts: Final = raw_parts if isinstance(raw_parts, list) else [] content: Final = A2ACompletionBridgeTransformation._extract_text_from_a2a_parts(parts) # Do not attach A2A message.metadata here — the completion bridge forwards it # once at run level via extra_body.metadata (LangGraph POST /runs/wait shape). - openai_message: Final[dict[str, Any]] = {"role": openai_role, "content": content} + openai_message: Final[dict[str, object]] = {"role": openai_role, "content": content} verbose_logger.debug( "A2A -> OpenAI transform: role=%s -> %s, content_length=%s", role, openai_role, len(content) @@ -151,11 +159,20 @@ class A2ACompletionBridgeTransformation: return [openai_message] + @staticmethod + def _extract_response_content(response: "ModelResponse | CustomStreamWrapper") -> str: + if not isinstance(response, ModelResponse) or not response.choices: + return "" + choice: Final = response.choices[0] + if not choice.message: + return "" + return choice.message.content or "" + @staticmethod def openai_response_to_a2a_response( - response: Any, + response: "ModelResponse | CustomStreamWrapper", request_id: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Transform a LiteLLM ModelResponse to A2A SendMessageResponse format. @@ -166,12 +183,7 @@ class A2ACompletionBridgeTransformation: Returns: A2A SendMessageResponse dict """ - # Extract content from response - content = "" - if hasattr(response, "choices") and response.choices: - choice: Final = response.choices[0] - if hasattr(choice, "message") and choice.message: - content = choice.message.content or "" + content: Final = A2ACompletionBridgeTransformation._extract_response_content(response) # Build A2A message a2a_message: Final = { @@ -182,7 +194,7 @@ class A2ACompletionBridgeTransformation: } # Build A2A response - a2a_response: Final = { + a2a_response: Final[dict[str, object]] = { "jsonrpc": "2.0", "id": request_id, "result": a2a_message, @@ -200,7 +212,7 @@ class A2ACompletionBridgeTransformation: @staticmethod def create_task_event( ctx: A2AStreamingContext, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Create the initial task event with status 'submitted'. @@ -235,7 +247,7 @@ class A2ACompletionBridgeTransformation: state: str, final: bool = False, message_text: str | None = None, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Create a status update event. @@ -245,7 +257,7 @@ class A2ACompletionBridgeTransformation: final: Whether this is the final event message_text: Optional message text for 'working' status """ - status: Final[dict[str, Any]] = { + status: Final[dict[str, object]] = { "state": state, "timestamp": A2ACompletionBridgeTransformation._get_timestamp(), } @@ -277,7 +289,7 @@ class A2ACompletionBridgeTransformation: def create_artifact_update_event( ctx: A2AStreamingContext, text: str, - ) -> dict[str, Any]: + ) -> dict[str, object]: """ Create an artifact update event with content. diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 56b8089b0af..0e8b8136c19 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -86,7 +86,7 @@ A2ACardResolver: Final = LiteLLMA2ACardResolver def _set_usage_on_logging_obj( - kwargs: dict[str, Any], + kwargs: Mapping[str, object], prompt_tokens: int, completion_tokens: int, ) -> None: @@ -99,7 +99,7 @@ def _set_usage_on_logging_obj( completion_tokens: Number of output tokens """ litellm_logging_obj: Final = kwargs.get("litellm_logging_obj") - if litellm_logging_obj is not None: + if isinstance(litellm_logging_obj, Logging): usage: Final = litellm.Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, @@ -109,7 +109,7 @@ def _set_usage_on_logging_obj( def _set_agent_id_on_logging_obj( - kwargs: dict[str, Any], + kwargs: Mapping[str, object], agent_id: str | None, ) -> None: """ @@ -123,7 +123,7 @@ def _set_agent_id_on_logging_obj( return litellm_logging_obj: Final = kwargs.get("litellm_logging_obj") - if litellm_logging_obj is not None: + if isinstance(litellm_logging_obj, Logging): # Set agent_id directly on model_call_details (same pattern as custom_llm_provider) litellm_logging_obj.model_call_details["agent_id"] = agent_id @@ -132,7 +132,7 @@ _A2A_COST_PARAM_KEYS: Final = ("cost_per_query", "input_cost_per_token", "output def _set_litellm_params_on_logging_obj( - kwargs: dict[str, Any], + kwargs: Mapping[str, object], litellm_params: Mapping[str, object], ) -> None: """ @@ -144,18 +144,22 @@ def _set_litellm_params_on_logging_obj( context, so merge the pricing keys in rather than replacing the dict. """ logging_obj: Final = kwargs.get("litellm_logging_obj") - if logging_obj is None: + if not isinstance(logging_obj, Logging): return - cost_params = {key: litellm_params[key] for key in _A2A_COST_PARAM_KEYS if litellm_params.get(key) is not None} + cost_params: Final = { + key: litellm_params[key] for key in _A2A_COST_PARAM_KEYS if litellm_params.get(key) is not None + } if not cost_params: return - existing: Final = logging_obj.model_call_details.get("litellm_params") or {} - logging_obj.model_call_details["litellm_params"] = {**existing, **cost_params} + logging_obj.model_call_details["litellm_params"] = { + **(logging_obj.model_call_details.get("litellm_params") or {}), + **cost_params, + } -def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: dict[str, Any]) -> str: +def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: Mapping[str, object]) -> str: """ Extract agent info and set model/custom_llm_provider for cost tracking. @@ -175,7 +179,7 @@ def _get_a2a_model_info(a2a_client: "A2AClientType", kwargs: dict[str, Any]) -> # Set on litellm_logging_obj if available (for standard logging payload) litellm_logging_obj: Final = kwargs.get("litellm_logging_obj") - if litellm_logging_obj is not None: + if isinstance(litellm_logging_obj, Logging): litellm_logging_obj.model = model litellm_logging_obj.custom_llm_provider = custom_llm_provider litellm_logging_obj.model_call_details["model"] = model @@ -498,7 +502,7 @@ async def asend_message( response: Final = LiteLLMSendMessageResponse.from_a2a_response(a2a_response, request_id=str(request.id)) # Calculate token usage from request and response - response_dict: Final[dict[str, object]] = a2a_response.model_dump(mode="json", exclude_none=True) + response_dict: Final[dict[str, object]] = a2a_response.root.model_dump(mode="json", exclude_none=True) ( prompt_tokens, completion_tokens, diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 2aa7b527c57..c8360a81c7a 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -390,7 +390,7 @@ def _handle_retrieve_batch_providers_without_provider_config( custom_llm_provider: Literal[ "openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "litellm_proxy", "anthropic" ] = "openai", - logging_obj: Any | None = None, + logging_obj: LiteLLMLoggingObj | None = None, ): api_base: str | None = None if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: diff --git a/litellm/integrations/bitbucket/bitbucket_client.py b/litellm/integrations/bitbucket/bitbucket_client.py index e06e5ab358f..9c964d8c10c 100644 --- a/litellm/integrations/bitbucket/bitbucket_client.py +++ b/litellm/integrations/bitbucket/bitbucket_client.py @@ -4,11 +4,38 @@ BitBucket API client for fetching .prompt files from BitBucket repositories. import base64 import urllib.parse -from typing import Any, Final +from collections.abc import Mapping +from typing import Final, TypedDict + +from typing_extensions import NotRequired, ReadOnly from litellm.llms.custom_httpx.http_handler import HTTPHandler +class BitBucketSrcEntry(TypedDict): + path: ReadOnly[NotRequired[str]] + type: ReadOnly[NotRequired[str]] + + +class BitBucketSrcListing(TypedDict): + values: ReadOnly[NotRequired[list[BitBucketSrcEntry]]] + + +class BitBucketBranch(TypedDict): + name: ReadOnly[NotRequired[str]] + type: ReadOnly[NotRequired[str]] + + +class BitBucketBranchListing(TypedDict): + values: ReadOnly[NotRequired[list[BitBucketBranch]]] + + +class BitBucketFileMetadata(TypedDict): + content_type: ReadOnly[str | None] + content_length: ReadOnly[str | None] + last_modified: ReadOnly[str | None] + + def _sanitize_file_path(file_path: str) -> str: """Reject path traversal and URL-encode each path segment.""" if "#" in file_path or "?" in file_path: @@ -31,7 +58,7 @@ class BitBucketClient: - Branch-specific file fetching """ - def __init__(self, config: dict[str, Any]): + def __init__(self, config: Mapping[str, object]): """ Initialize the BitBucket client. @@ -135,8 +162,8 @@ class BitBucketClient: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data: Final = response.json() - files: Final = [] + data: Final[BitBucketSrcListing] = response.json() + files: Final[list[str]] = [] for item in data.get("values", []): if item.get("type") == "commit_file": @@ -162,7 +189,7 @@ class BitBucketClient: else: raise Exception(f"Error listing files in '{directory_path}': {e}") - def get_repository_info(self) -> dict[str, Any]: + def get_repository_info(self) -> Mapping[str, object]: """ Get information about the repository. @@ -191,7 +218,7 @@ class BitBucketClient: except Exception: return False - def get_branches(self) -> list[dict[str, Any]]: + def get_branches(self) -> list[BitBucketBranch]: """ Get list of branches in the repository. @@ -204,12 +231,12 @@ class BitBucketClient: response: Final = self.http_handler.get(url, headers=self.headers) response.raise_for_status() - data: Final = response.json() + data: Final[BitBucketBranchListing] = response.json() return data.get("values", []) except Exception as e: raise Exception(f"Failed to get branches: {e}") - def get_file_metadata(self, file_path: str) -> dict[str, Any] | None: + def get_file_metadata(self, file_path: str) -> BitBucketFileMetadata | None: """ Get metadata about a file (size, last modified, etc.). diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index 76720682101..321c7896d63 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -7,7 +7,7 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan. import time import uuid -from typing import Any, ClassVar, Final, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, cast from litellm._logging import verbose_logger from litellm.compression import compress @@ -22,6 +22,9 @@ from litellm.types.integrations.custom_logger import ( ) from litellm.types.utils import CallTypes +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + LITELLM_CONTENT_RETRIEVE_TOOL_NAME: Final = "litellm_content_retrieve" _CACHE_TTL_SECONDS: Final = 15 * 60 @@ -222,7 +225,7 @@ class CompressionInterceptionLogger(CustomLogger): response: Any, anthropic_messages_provider_config: Any, anthropic_messages_optional_request_params: dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj | None", stream: bool, kwargs: dict, ) -> AgenticLoopPlan: diff --git a/litellm/integrations/focus/destinations/mavvrik_destination.py b/litellm/integrations/focus/destinations/mavvrik_destination.py index dad5526eb18..4e7765b9e5d 100644 --- a/litellm/integrations/focus/destinations/mavvrik_destination.py +++ b/litellm/integrations/focus/destinations/mavvrik_destination.py @@ -9,9 +9,12 @@ Flow: from __future__ import annotations import gzip -from typing import Any, Final +from collections.abc import Mapping +from typing import Final, Protocol from urllib.parse import urlparse +from typing_extensions import NotRequired, ReadOnly, TypedDict + from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -28,6 +31,34 @@ _MAVVRIK_ALLOWED_SUFFIXES: Final = (".mavvrik.dev", ".mavvrik.ai", ".mavvrik.app _GCS_CHUNK_SIZE: Final = 8 * 1024 * 1024 # 8 MB +class MavvrikRegisterBody(TypedDict): + metricsMarker: ReadOnly[NotRequired[int | str]] + + +class MavvrikUploadUrlBody(TypedDict): + url: ReadOnly[NotRequired[str]] + + +class _RegisterResponse(Protocol): + def json(self) -> MavvrikRegisterBody: ... + + +class _UploadUrlResponse(Protocol): + def json(self) -> MavvrikUploadUrlBody: ... + + +def _register_body(response: _RegisterResponse) -> MavvrikRegisterBody: + return response.json() + + +def _upload_url_body(response: _UploadUrlResponse) -> MavvrikUploadUrlBody: + return response.json() + + +def _header_value(headers: Mapping[str, str], name: str) -> str | None: + return headers.get(name) + + def _validate_api_endpoint(api_endpoint: str) -> None: if not api_endpoint.startswith("https://"): raise ValueError("MAVVRIK_API_ENDPOINT must be an HTTPS URL") @@ -56,12 +87,12 @@ class FocusMavvrikDestination(FocusDestination): self, *, prefix: str, - config: dict[str, Any] | None = None, + config: Mapping[str, str] | None = None, ) -> None: - config = config or {} - api_key: Final = config.get("api_key") - api_endpoint: Final = config.get("api_endpoint") - connection_id: Final = config.get("connection_id") + resolved_config: Final[Mapping[str, str]] = config or {} + api_key: Final = resolved_config.get("api_key") + api_endpoint: Final = resolved_config.get("api_endpoint") + connection_id: Final = resolved_config.get("connection_id") if not api_key: raise ValueError( @@ -100,7 +131,7 @@ class FocusMavvrikDestination(FocusDestination): def _auth_headers(self) -> dict[str, str]: return {"Content-Type": "application/json", "x-api-key": self.api_key} - async def _ensure_registered(self) -> int | None: + async def _ensure_registered(self) -> int | str | None: """POST agent endpoint to register/initialize the connector (once per instance). Returns metricsMarker from the Mavvrik response — the last date index @@ -127,7 +158,7 @@ class FocusMavvrikDestination(FocusDestination): if resp.status_code >= 400: raise RuntimeError(f"Mavvrik FOCUS destination: register failed ({resp.status_code}): {resp.text[:200]}") self._registered = True - metrics_marker: Final = resp.json().get("metricsMarker", 0) + metrics_marker: Final = _register_body(resp).get("metricsMarker", 0) verbose_logger.debug( "Mavvrik FOCUS destination: connector registered (metricsMarker=%s)", metrics_marker, @@ -148,7 +179,7 @@ class FocusMavvrikDestination(FocusDestination): raise RuntimeError( f"Mavvrik FOCUS destination: failed to get signed URL ({resp.status_code}): {resp.text[:200]}" ) - signed_url: Final = resp.json().get("url") + signed_url: Final = _upload_url_body(resp).get("url") if not signed_url: raise RuntimeError(f"Mavvrik FOCUS destination: response missing 'url' field: {resp.json()}") _validate_gcs_url(signed_url, "signed URL") @@ -190,7 +221,7 @@ class FocusMavvrikDestination(FocusDestination): f"Mavvrik FOCUS destination: GCS session init failed ({init_resp.status_code}): {init_resp.text[:400]}" ) - session_uri: Final = init_resp.headers.get("Location") + session_uri: Final = _header_value(init_resp.headers, "Location") if not session_uri: raise RuntimeError("Mavvrik FOCUS destination: GCS session init missing Location header") _validate_gcs_url(session_uri, "session URI") @@ -264,7 +295,7 @@ class FocusMavvrikDestination(FocusDestination): ) verbose_logger.debug("Mavvrik FOCUS destination: metricsMarker advanced to %s", date_epoch) - async def get_metrics_marker(self) -> int | None: + async def get_metrics_marker(self) -> int | str | None: """Register with Mavvrik and return the current metricsMarker. Always calls the Mavvrik register API — unlike deliver() which skips @@ -287,7 +318,7 @@ class FocusMavvrikDestination(FocusDestination): if resp.status_code >= 400: raise RuntimeError(f"Mavvrik FOCUS destination: register failed ({resp.status_code}): {resp.text[:200]}") self._registered = True - metrics_marker: Final = resp.json().get("metricsMarker", 0) + metrics_marker: Final = _register_body(resp).get("metricsMarker", 0) verbose_logger.debug("Mavvrik FOCUS destination: got metricsMarker=%s", metrics_marker) return metrics_marker diff --git a/litellm/integrations/prometheus_helpers/prometheus_api.py b/litellm/integrations/prometheus_helpers/prometheus_api.py index 9f77f87a670..e111474bd4d 100644 --- a/litellm/integrations/prometheus_helpers/prometheus_api.py +++ b/litellm/integrations/prometheus_helpers/prometheus_api.py @@ -7,6 +7,9 @@ import time from datetime import datetime, timedelta from typing import Final +from pydantic import BaseModel, TypeAdapter +from typing_extensions import ReadOnly, TypedDict + from litellm import get_secret from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( @@ -18,10 +21,32 @@ PROMETHEUS_URL: Final[str | None] = get_secret("PROMETHEUS_URL") PROMETHEUS_SELECTED_INSTANCE: Final[str | None] = get_secret("PROMETHEUS_SELECTED_INSTANCE") async_http_handler: Final = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) +_RAW_JSON_PAYLOAD: Final = TypeAdapter(object) + + +class PrometheusRangeSample(BaseModel): + """One ``matrix`` series of the Prometheus HTTP query API.""" + + metric: dict[str, object] + values: list[tuple[float, str]] + + +class PrometheusQueryData(BaseModel): + result: list[PrometheusRangeSample] + + +class PrometheusQueryResponse(BaseModel): + data: PrometheusQueryData + + +class PrometheusDailySpend(TypedDict): + date: ReadOnly[str] + spend: ReadOnly[float] + async def get_metric_from_prometheus( metric_name: str, -): +) -> list[PrometheusRangeSample]: # Get the start of the current day in Unix timestamp if PROMETHEUS_URL is None: raise ValueError("PROMETHEUS_URL not set please set 'PROMETHEUS_URL=<>' in .env") @@ -31,13 +56,13 @@ async def get_metric_from_prometheus( response: Final = await async_http_handler.get( f"{PROMETHEUS_URL}/api/v1/query", params={"query": query, "time": now} ) # End of the day - _json_response: Final = response.json() + _json_response: Final = _RAW_JSON_PAYLOAD.validate_python(response.json()) verbose_logger.debug("json response from prometheus /query api %s", _json_response) - results: Final = response.json()["data"]["result"] + results: Final = PrometheusQueryResponse.model_validate(_json_response).data.result return results -async def get_fallback_metric_from_prometheus(): +async def get_fallback_metric_from_prometheus() -> str: """ Gets fallback metrics from prometheus for the last 24 hours """ @@ -55,17 +80,17 @@ async def get_fallback_metric_from_prometheus(): verbose_logger.debug("response json %s", response_json) for result in response_json: verbose_logger.debug("result= %s", result) - metric = result["metric"] - metric_values = result["values"] + metric_labels = result.metric + metric_values = result.values most_recent_value = metric_values[0] if PROMETHEUS_SELECTED_INSTANCE is not None: - if metric.get("instance") != PROMETHEUS_SELECTED_INSTANCE: + if metric_labels.get("instance") != PROMETHEUS_SELECTED_INSTANCE: continue value = int(float(most_recent_value[1])) # Convert value to integer - primary_model = metric.get("primary_model", "Unknown") - fallback_model = metric.get("fallback_model", "Unknown") + primary_model = metric_labels.get("primary_model", "Unknown") + fallback_model = metric_labels.get("fallback_model", "Unknown") response_message += f"`{value} successful fallback requests` with primary model=`{primary_model}` -> fallback model=`{fallback_model}`" response_message += "\n" verbose_logger.debug("response message %s", response_message) @@ -96,7 +121,7 @@ def _quote_promql_string_literal(value: str) -> str: return json.dumps(value, ensure_ascii=False) -async def get_daily_spend_from_prometheus(api_key: str | None): +async def get_daily_spend_from_prometheus(api_key: str | None) -> list[PrometheusDailySpend]: """ Expected Response Format: [ @@ -133,17 +158,16 @@ async def get_daily_spend_from_prometheus(api_key: str | None): } response: Final = await async_http_handler.get(url, params=params) - _json_response: Final = response.json() + _json_response: Final = _RAW_JSON_PAYLOAD.validate_python(response.json()) verbose_logger.debug("json response from prometheus /query api %s", _json_response) - results: Final = response.json()["data"]["result"] - formatted_results: Final = [] - - for result in results: - metric_data = result["values"] - for timestamp, value in metric_data: - # Convert timestamp to ISO 8601 string with UTC offset - date = datetime.fromtimestamp(float(timestamp)).isoformat() + "+00:00" - spend = float(value) - formatted_results.append({"date": date, "spend": spend}) + results: Final = PrometheusQueryResponse.model_validate(_json_response).data.result + formatted_results: Final[list[PrometheusDailySpend]] = [ + { + "date": datetime.fromtimestamp(float(timestamp)).isoformat() + "+00:00", + "spend": float(value), + } + for result in results + for timestamp, value in result.values + ] return formatted_results diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 10b4c0dd433..9f6ae72fb3a 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -17,6 +17,7 @@ import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import DEFAULT_S3_BATCH_SIZE, DEFAULT_S3_FLUSH_INTERVAL_SECONDS from litellm.integrations.s3 import get_s3_object_key, resolve_sse_params +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM @@ -222,7 +223,10 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): protocol: Final = "https://" if self.s3_endpoint_url.startswith("https://") else "http://" return f"{protocol}{self.s3_bucket_name}.{endpoint_host}/{encoded_key}" return f"{self.s3_endpoint_url}/{self.s3_bucket_name}/{encoded_key}" - return f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}.amazonaws.com/{encoded_key}" + return ( + f"https://{self.s3_bucket_name}.s3.{self.s3_region_name}." + f"{get_aws_dns_suffix(self.s3_region_name)}/{encoded_key}" + ) def _sse_headers(self) -> Mapping[str, str]: candidates: Final = { diff --git a/litellm/litellm_core_utils/aws_partition.py b/litellm/litellm_core_utils/aws_partition.py new file mode 100644 index 00000000000..f8ca3aa4473 --- /dev/null +++ b/litellm/litellm_core_utils/aws_partition.py @@ -0,0 +1,55 @@ +import re +from types import MappingProxyType +from typing import Final, NamedTuple + + +class AwsPartition(NamedTuple): + partition: str + dns_suffix: str + + +_COMMERCIAL_PARTITION: Final = AwsPartition(partition="aws", dns_suffix="amazonaws.com") + +_PARTITIONS_BY_REGION_PREFIX: Final = MappingProxyType( + { + "cn-": AwsPartition(partition="aws-cn", dns_suffix="amazonaws.com.cn"), + "us-gov-": AwsPartition(partition="aws-us-gov", dns_suffix="amazonaws.com"), + "us-isob-": AwsPartition(partition="aws-iso-b", dns_suffix="sc2s.sgov.gov"), + "us-isof-": AwsPartition(partition="aws-iso-f", dns_suffix="csp.hci.ic.gov"), + "us-iso-": AwsPartition(partition="aws-iso", dns_suffix="c2s.ic.gov"), + "eu-isoe-": AwsPartition(partition="aws-iso-e", dns_suffix="cloud.adc-e.uk"), + } +) + +_BEDROCK_ARN_PATTERN: Final = re.compile(r"arn:aws(?:-[a-z0-9-]+)?:bedrock") +_BEDROCK_ARN_PREFIX_PATTERN: Final = re.compile(r"\Aarn:aws(?:-[a-z0-9-]+)?:bedrock:") +_AWS_ARN_PATTERN: Final = re.compile(r"arn:aws(?:-[a-z0-9-]+)?:") + + +def get_aws_partition(aws_region_name: str | None) -> AwsPartition: + if not aws_region_name: + return _COMMERCIAL_PARTITION + return next( + (partition for prefix, partition in _PARTITIONS_BY_REGION_PREFIX.items() if aws_region_name.startswith(prefix)), + _COMMERCIAL_PARTITION, + ) + + +def get_aws_dns_suffix(aws_region_name: str | None) -> str: + return get_aws_partition(aws_region_name).dns_suffix + + +def get_aws_arn_prefix(aws_region_name: str | None) -> str: + return f"arn:{get_aws_partition(aws_region_name).partition}:" + + +def contains_bedrock_arn(value: str) -> bool: + return _BEDROCK_ARN_PATTERN.search(value) is not None + + +def is_bedrock_arn(value: str) -> bool: + return _BEDROCK_ARN_PREFIX_PATTERN.match(value) is not None + + +def contains_aws_arn(value: str) -> bool: + return _AWS_ARN_PATTERN.search(value) is not None diff --git a/litellm/litellm_core_utils/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index c833d57b6a9..04824a5bf39 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -2,6 +2,7 @@ from collections.abc import Mapping from typing import Final import litellm +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH def _form_field_value(value: object) -> str: @@ -13,18 +14,31 @@ def _form_field_value(value: object) -> str: def _flatten_form_field(key: str, value: object) -> tuple[tuple[str, str], ...]: - if isinstance(value, Mapping): - return tuple( - item for subkey, subvalue in value.items() for item in _flatten_form_field(f"{key}[{subkey}]", subvalue) - ) - if isinstance(value, (list, tuple)): - return tuple(item for entry in value for item in _flatten_form_field(f"{key}[]", entry)) - if value is None: - return () - serialized: Final = _form_field_value(value) - if not serialized: - return () - return ((key, serialized),) + pending_fields: Final[ # mutable-ok: depth-capped stack walks nested JSON into multipart names + list[tuple[str, object, int]] + ] = [ # mutable-ok: depth-capped stack walks nested JSON into multipart names + (key, value, 0) + ] + flat_fields: Final[list[tuple[str, str]]] = [] # mutable-ok: local accumulator + while pending_fields: + current_key, current_value, depth = pending_fields.pop() + if depth > DEFAULT_MAX_RECURSE_DEPTH: + raise ValueError("form field nesting exceeds max depth") + if isinstance(current_value, Mapping): + pending_fields.extend( + (f"{current_key}[{subkey}]", subvalue, depth + 1) + for subkey, subvalue in reversed(tuple(current_value.items())) + ) + continue + if isinstance(current_value, (list, tuple)): + pending_fields.extend((f"{current_key}[]", entry, depth + 1) for entry in reversed(tuple(current_value))) + continue + if current_value is None: + continue + serialized = _form_field_value(current_value) + if serialized: + flat_fields.append((current_key, serialized)) + return tuple(flat_fields) def _is_form_scalar(value: object) -> bool: @@ -32,23 +46,36 @@ def _is_form_scalar(value: object) -> bool: def _flatten_form_data_field(key: str, value: object) -> tuple[tuple[str, str | tuple[str, ...]], ...]: - if isinstance(value, Mapping): - return tuple( - item - for subkey, subvalue in value.items() - for item in _flatten_form_data_field(f"{key}[{subkey}]", subvalue) - ) - if isinstance(value, (list, tuple)): - if all(_is_form_scalar(entry) for entry in value): - serialized_fields: Final = tuple(field for entry in value if (field := _form_field_value(entry))) - return ((key, serialized_fields),) if serialized_fields else () - return tuple(item for entry in value for item in _flatten_form_data_field(f"{key}[]", entry)) - if value is None: - return () - serialized: Final = _form_field_value(value) - if not serialized: - return () - return ((key, serialized),) + pending_fields: Final[ # mutable-ok: depth-capped stack walks nested JSON into multipart names + list[tuple[str, object, int]] + ] = [ # mutable-ok: depth-capped stack walks nested JSON into multipart names + (key, value, 0) + ] + flat_fields: Final[list[tuple[str, str | tuple[str, ...]]]] = [] # mutable-ok: local accumulator + while pending_fields: + current_key, current_value, depth = pending_fields.pop() + if depth > DEFAULT_MAX_RECURSE_DEPTH: + raise ValueError("form field nesting exceeds max depth") + if isinstance(current_value, Mapping): + pending_fields.extend( + (f"{current_key}[{subkey}]", subvalue, depth + 1) + for subkey, subvalue in reversed(tuple(current_value.items())) + ) + continue + if isinstance(current_value, (list, tuple)): + if all(_is_form_scalar(entry) for entry in current_value): + serialized_fields = tuple(field for entry in current_value if (field := _form_field_value(entry))) + if serialized_fields: + flat_fields.append((current_key, serialized_fields)) + continue + pending_fields.extend((f"{current_key}[]", entry, depth + 1) for entry in reversed(tuple(current_value))) + continue + if current_value is None: + continue + serialized = _form_field_value(current_value) + if serialized: + flat_fields.append((current_key, serialized)) + return tuple(flat_fields) def flatten_form_field_values(*sources: Mapping[str, object] | None) -> tuple[tuple[str, str | tuple[str, ...]], ...]: diff --git a/litellm/litellm_core_utils/model_response_utils.py b/litellm/litellm_core_utils/model_response_utils.py index ea4be1c856f..7bf667164ae 100644 --- a/litellm/litellm_core_utils/model_response_utils.py +++ b/litellm/litellm_core_utils/model_response_utils.py @@ -2,9 +2,21 @@ Utility functions for ModelResponse and ModelResponseStream objects. """ -from typing import Any, Final +from collections.abc import Mapping +from typing import Final -from litellm.types.utils import Delta, ModelResponseBase, ModelResponseStream +from typing_extensions import ReadOnly, TypedDict + +from litellm.types.utils import Delta, ModelResponseBase, ModelResponseStream, StreamingChoices + + +class _AttributeView(TypedDict): + value: ReadOnly[object] + + +def _attribute_of(source: object, name: str) -> object: + attribute: Final[_AttributeView] = {"value": getattr(source, name)} + return attribute["value"] def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: @@ -40,10 +52,10 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: return False # Check model_extra for dynamically added fields (this is where Pydantic stores them) - if hasattr(model_response, "model_extra") and model_response.model_extra: - for extra_field_name, extra_field_value in model_response.model_extra.items(): - if _has_meaningful_content(extra_field_value): - return False + stream_extra_fields: Final[Mapping[str, object]] = model_response.model_extra or {} + for extra_field_value in stream_extra_fields.values(): + if _has_meaningful_content(extra_field_value): + return False # Check for any non-base fields that are set # Access model_fields on the class, not the instance, to avoid Pydantic 2.11+ deprecation warnings @@ -57,7 +69,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: continue # Check if any other field has meaningful content - model_response_value = getattr(model_response, model_response_field, None) + model_response_value: object = getattr(model_response, model_response_field, None) if _has_meaningful_content(model_response_value): return False @@ -71,7 +83,7 @@ def is_model_response_stream_empty(model_response: ModelResponseStream) -> bool: return True -def _has_meaningful_content(value: Any) -> bool: +def _has_meaningful_content(value: object) -> bool: """ Check if a value contains meaningful content. @@ -102,7 +114,7 @@ def _has_meaningful_content(value: Any) -> bool: return True -def _is_choice_non_empty(choice: Any) -> bool: +def _is_choice_non_empty(choice: StreamingChoices) -> bool: """ Deep check if a choice contains any meaningful content. @@ -113,41 +125,41 @@ def _is_choice_non_empty(choice: Any) -> bool: bool: True if the choice has meaningful content, False otherwise """ # Check finish_reason - if hasattr(choice, "finish_reason") and choice.finish_reason is not None: + if getattr(choice, "finish_reason", None) is not None: return True # Check logprobs - if hasattr(choice, "logprobs") and choice.logprobs is not None: + if getattr(choice, "logprobs", None) is not None: return True # Check enhancements (if present) - if hasattr(choice, "enhancements") and choice.enhancements is not None: + if getattr(choice, "enhancements", None) is not None: return True # Deep check delta object - if hasattr(choice, "delta") and choice.delta is not None: - if _is_delta_non_empty(choice.delta): - return True + choice_delta: Final[Delta | None] = getattr(choice, "delta", None) + if choice_delta is not None and _is_delta_non_empty(choice_delta): + return True # Check model_extra for dynamically added fields on the choice - if hasattr(choice, "model_extra") and choice.model_extra: - for extra_field_name, extra_field_value in choice.model_extra.items(): - # Skip certain structural fields that are just default/None placeholders - if extra_field_name == "index" and extra_field_value == 0: - continue - if extra_field_name in {"finish_reason", "logprobs"} and extra_field_value is None: - continue - if extra_field_name == "delta": - continue - if _has_meaningful_content(extra_field_value): - return True + choice_extra_fields: Final[Mapping[str, object]] = choice.model_extra or {} + for extra_field_name, extra_field_value in choice_extra_fields.items(): + # Skip certain structural fields that are just default/None placeholders + if extra_field_name == "index" and extra_field_value == 0: + continue + if extra_field_name in {"finish_reason", "logprobs"} and extra_field_value is None: + continue + if extra_field_name == "delta": + continue + if _has_meaningful_content(extra_field_value): + return True # Check for any other non-standard fields on the choice for attr_name in dir(choice): # Skip private attributes, methods, and known empty fields if ( attr_name.startswith("_") - or callable(getattr(choice, attr_name)) + or callable(_attribute_of(choice, attr_name)) or attr_name.startswith("model_") or attr_name in { @@ -160,8 +172,8 @@ def _is_choice_non_empty(choice: Any) -> bool: ): continue - attr_value = getattr(choice, attr_name, None) - if _has_meaningful_content(attr_value): + choice_attr_value: object = getattr(choice, attr_name, None) + if _has_meaningful_content(choice_attr_value): return True return False @@ -178,20 +190,20 @@ def _is_delta_non_empty(delta: Delta) -> bool: bool: True if the delta has meaningful content, False otherwise """ # Check model_extra for dynamically added fields (this is where Pydantic stores them) - if hasattr(delta, "model_extra") and delta.model_extra: - for extra_field_name, extra_field_value in delta.model_extra.items(): - # Even structural fields are meaningful if they have actual content - if _has_meaningful_content(extra_field_value): - return True + delta_extra_fields: Final[Mapping[str, object]] = delta.model_extra or {} + for extra_field_value in delta_extra_fields.values(): + # Even structural fields are meaningful if they have actual content + if _has_meaningful_content(extra_field_value): + return True # Check all regular attributes of the delta object for attr_name in dir(delta): # Skip private attributes, methods, and Pydantic-specific fields - if attr_name.startswith("_") or callable(getattr(delta, attr_name)) or attr_name.startswith("model_"): + if attr_name.startswith("_") or callable(_attribute_of(delta, attr_name)) or attr_name.startswith("model_"): continue - attr_value = getattr(delta, attr_name, None) - if _has_meaningful_content(attr_value): + delta_attr_value: object = getattr(delta, attr_name, None) + if _has_meaningful_content(delta_attr_value): return True return False diff --git a/litellm/litellm_core_utils/url_utils.py b/litellm/litellm_core_utils/url_utils.py index 0a59eaa75d3..125baa4743a 100644 --- a/litellm/litellm_core_utils/url_utils.py +++ b/litellm/litellm_core_utils/url_utils.py @@ -21,13 +21,61 @@ Admins can opt out via two ``litellm`` globals (wired from proxy config): import socket from ipaddress import ip_address, ip_network -from typing import Any, Final +from typing import Any, Final, Protocol from urllib.parse import quote, urlparse, urlunparse import httpx +from typing_extensions import ReadOnly, TypedDict import litellm +_SockAddr = tuple[str, int] | tuple[str, int, int, int] | tuple[int, bytes] + + +class _LocationHeaderView(TypedDict): + location: ReadOnly[object] + + +class _ResponseView(TypedDict): + response: ReadOnly[httpx.Response] + + +class _UrlFetcher(Protocol): + """The slice of ``httpx.Client`` / ``HTTPHandler`` that ``safe_get`` drives.""" + + def get( + self, + url: str, + *, + headers: dict[str, str] | None = None, + follow_redirects: bool = False, + ) -> httpx.Response: ... + + +class _AsyncUrlFetcher(Protocol): + """The slice of ``httpx.AsyncClient`` / ``AsyncHTTPHandler`` that ``async_safe_get`` drives.""" + + async def get( + self, + url: str, + *, + headers: dict[str, str] | None = None, + follow_redirects: bool = False, + ) -> httpx.Response: ... + + +class _FetcherView(TypedDict): + fetcher: ReadOnly[_UrlFetcher] + + +class _AsyncFetcherView(TypedDict): + fetcher: ReadOnly[_AsyncUrlFetcher] + + +class _CallerHeadersView(TypedDict): + headers: ReadOnly[dict[str, str]] + + # Globally-routable IPs that are cloud-internal. Everything else # non-public is caught by ``not ip.is_global`` (RFC 6890, as implemented by # Python's ``ipaddress`` module). This list only holds IPs that are @@ -44,7 +92,7 @@ class SSRFError(ValueError): """Raised when a URL targets a blocked network.""" -def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") -> str: +def encode_url_path_segment(value: object, *, field_name: str = "path parameter") -> str: """Percent-encode one user-controlled URL path segment. ``urllib.parse.quote(..., safe="")`` intentionally leaves RFC 3986 @@ -64,7 +112,7 @@ def encode_url_path_segment(value: Any, *, field_name: str = "path parameter") - return quote(value_str, safe="") -def encode_url_path_segments(value: Any, *, field_name: str = "path") -> str: +def encode_url_path_segments(value: object, *, field_name: str = "path") -> str: """Percent-encode a user-controlled URL path made of multiple segments. Empty segments are rejected, so leading, trailing, or consecutive slashes @@ -77,11 +125,7 @@ def encode_url_path_segments(value: Any, *, field_name: str = "path") -> str: if value_str == "": raise ValueError(f"{field_name} is required") - encoded_segments: Final = [] - for segment in value_str.split("/"): - encoded_segments.append(encode_url_path_segment(segment, field_name=field_name)) - - return "/".join(encoded_segments) + return "/".join(encode_url_path_segment(segment, field_name=field_name) for segment in value_str.split("/")) def _is_blocked_ip(addr: str) -> bool: @@ -202,7 +246,7 @@ def _format_host_header(hostname: str, port: int, default_port: int) -> str: return f"{bracketed}:{port}" -def _sockaddr_host(sockaddr: Any) -> str: +def _sockaddr_host(sockaddr: _SockAddr) -> str: """Return the host element of a ``getaddrinfo`` sockaddr as ``str``. ``getaddrinfo`` with ``IPPROTO_TCP`` returns AF_INET / AF_INET6 sockaddrs @@ -285,8 +329,8 @@ def validate_url(url: str) -> tuple[str, str]: raise SSRFError(f"No addresses found for '{hostname}'") if not is_allowlisted: - for family, type_, proto, canonname, sockaddr in addrinfo: - resolved_ip = _sockaddr_host(sockaddr) + for addrinfo_entry in addrinfo: + resolved_ip = _sockaddr_host(addrinfo_entry[4]) if _is_blocked_ip(resolved_ip): raise SSRFError( f"URL targets a blocked address ({resolved_ip}). " @@ -363,9 +407,10 @@ def assert_same_origin(candidate_url: str, expected_url: str) -> None: _MAX_REDIRECTS: Final = 10 -def _extract_redirect_url(response: Any, request_url: str) -> str: +def _extract_redirect_url(response: httpx.Response, request_url: str) -> str: """Extract and resolve the redirect target from a response's Location header.""" - location: Final = response.headers.get("location") + header_view: Final[_LocationHeaderView] = {"location": response.headers.get("location")} + location: Final = header_view["location"] if not isinstance(location, str) or not location: raise SSRFError("Redirect response has no Location header") # Resolve relative URLs against the request URL @@ -393,14 +438,17 @@ def safe_get(client: Any, url: str, **kwargs: Any) -> Any: """ if not getattr(litellm, "user_url_validation", True): kwargs.setdefault("follow_redirects", True) - return client.get(url, **kwargs) + unvalidated: Final[_ResponseView] = {"response": client.get(url, **kwargs)} + return unvalidated["response"] + fetcher_view: Final[_FetcherView] = {"fetcher": client} + fetcher: Final = fetcher_view["fetcher"] kwargs.pop("follow_redirects", None) - caller_headers: Final = kwargs.pop("headers", {}) + headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})} for _ in range(_MAX_REDIRECTS): validated_url, original_host = validate_url(url) - response = client.get( + response = fetcher.get( validated_url, - headers={**caller_headers, "Host": original_host}, + headers={**headers_view["headers"], "Host": original_host}, follow_redirects=False, **kwargs, ) @@ -416,14 +464,17 @@ async def async_safe_get(client: Any, url: str, **kwargs: Any) -> Any: """Async version of safe_get.""" if not getattr(litellm, "user_url_validation", True): kwargs.setdefault("follow_redirects", True) - return await client.get(url, **kwargs) + unvalidated: Final[_ResponseView] = {"response": await client.get(url, **kwargs)} + return unvalidated["response"] + fetcher_view: Final[_AsyncFetcherView] = {"fetcher": client} + fetcher: Final = fetcher_view["fetcher"] kwargs.pop("follow_redirects", None) - caller_headers: Final = kwargs.pop("headers", {}) + headers_view: Final[_CallerHeadersView] = {"headers": kwargs.pop("headers", {})} for _ in range(_MAX_REDIRECTS): validated_url, original_host = validate_url(url) - response = await client.get( + response = await fetcher.get( validated_url, - headers={**caller_headers, "Host": original_host}, + headers={**headers_view["headers"], "Host": original_host}, follow_redirects=False, **kwargs, ) diff --git a/litellm/llms/a2a/chat/transformation.py b/litellm/llms/a2a/chat/transformation.py index 8ebf8958416..f6cb14c0836 100644 --- a/litellm/llms/a2a/chat/transformation.py +++ b/litellm/llms/a2a/chat/transformation.py @@ -4,7 +4,7 @@ A2A Protocol Transformation for LiteLLM import uuid from collections.abc import Iterator -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final import httpx @@ -20,6 +20,11 @@ from ..common_utils import ( ) from .streaming_iterator import A2AModelResponseIterator +if TYPE_CHECKING: + import tiktoken + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + class A2AConfig(BaseConfig): """ @@ -246,12 +251,12 @@ class A2AConfig(BaseConfig): model: str, raw_response: httpx.Response, model_response: ModelResponse, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", request_data: dict, messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/aiml/image_generation/transformation.py b/litellm/llms/aiml/image_generation/transformation.py index ba641c0a752..4f4cd074165 100644 --- a/litellm/llms/aiml/image_generation/transformation.py +++ b/litellm/llms/aiml/image_generation/transformation.py @@ -14,6 +14,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -169,7 +171,7 @@ class AimlImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/aiohttp_openai/chat/transformation.py b/litellm/llms/aiohttp_openai/chat/transformation.py index 21adab2d5b1..530896bf9b0 100644 --- a/litellm/llms/aiohttp_openai/chat/transformation.py +++ b/litellm/llms/aiohttp_openai/chat/transformation.py @@ -16,6 +16,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, ModelResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -66,7 +68,7 @@ class AiohttpOpenAIChatConfig(OpenAILikeChatConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/amazon_nova/chat/transformation.py b/litellm/llms/amazon_nova/chat/transformation.py index c26182643df..7551fb28c21 100644 --- a/litellm/llms/amazon_nova/chat/transformation.py +++ b/litellm/llms/amazon_nova/chat/transformation.py @@ -2,7 +2,7 @@ Translate from OpenAI's `/v1/chat/completions` to Amazon Nova's `/v1/chat/completions` """ -from typing import Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -16,6 +16,9 @@ from litellm.types.utils import ModelResponse from ...openai_like.chat.transformation import OpenAILikeChatConfig +if TYPE_CHECKING: + import tiktoken + class AmazonNovaChatConfig(OpenAILikeChatConfig): max_completion_tokens: int | None = None @@ -83,7 +86,7 @@ class AmazonNovaChatConfig(OpenAILikeChatConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/anthropic/batches/transformation.py b/litellm/llms/anthropic/batches/transformation.py index 3f8fd2c27f4..6b39adc511e 100644 --- a/litellm/llms/anthropic/batches/transformation.py +++ b/litellm/llms/anthropic/batches/transformation.py @@ -12,6 +12,8 @@ from litellm.types.llms.openai import AllMessageValues, CreateBatchRequest from litellm.types.utils import LiteLLMBatch, LlmProviders, ModelResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj LoggingClass = LiteLLMLoggingObj @@ -261,7 +263,7 @@ class AnthropicBatchesConfig(BaseBatchesConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 15fc482b34e..47116a8f8fb 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -92,6 +92,8 @@ from ..common_utils import ( ) if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj LoggingClass = LiteLLMLoggingObj @@ -2575,7 +2577,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/anthropic/completion/transformation.py b/litellm/llms/anthropic/completion/transformation.py index d4e2b3db166..b15b0159bd9 100644 --- a/litellm/llms/anthropic/completion/transformation.py +++ b/litellm/llms/anthropic/completion/transformation.py @@ -7,7 +7,7 @@ Litellm provider slug: `anthropic_text/` import json import time from collections.abc import AsyncIterator, Iterator -from typing import Final +from typing import TYPE_CHECKING, Final import httpx @@ -32,6 +32,9 @@ from litellm.types.utils import ( Usage, ) +if TYPE_CHECKING: + import tiktoken + class AnthropicTextError(BaseLLMException): def __init__(self, status_code, message): @@ -182,7 +185,7 @@ class AnthropicTextConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: str, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: @@ -202,9 +205,10 @@ class AnthropicTextConfig(BaseConfig): model_response.choices[0].finish_reason = completion_response["stop_reason"] ## CALCULATING USAGE - prompt_tokens: Final = len(encoding.encode(prompt)) ##[TODO] use the anthropic tokenizer here + tokenizer: Final = encoding if encoding is not None else litellm.encoding + prompt_tokens: Final = len(tokenizer.encode(prompt)) ##[TODO] use the anthropic tokenizer here completion_tokens: Final = len( - encoding.encode(model_response["choices"][0]["message"].get("content", "")) + tokenizer.encode(model_response["choices"][0]["message"].get("content", "")) ) ##[TODO] use the anthropic tokenizer here model_response.created = int(time.time()) diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py index 41795fa0f32..902808647c0 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py @@ -2,7 +2,7 @@ import inspect from collections.abc import Awaitable, Callable -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Final, TypeAlias from litellm._logging import verbose_logger from litellm.types.llms.anthropic import AppliedEdit @@ -11,7 +11,13 @@ from .constants import CLEAR_TOOL_USES_EDIT_TYPE, COMPACT_EDIT_TYPE from .editors import apply_clear_tool_uses_20250919, apply_compact_20260112 from .result import PolyfillResult -EditorFn = Callable[..., Any] +if TYPE_CHECKING: + from litellm.proxy._types import UserAPIKeyAuth + from litellm.router import Router + +EditorResult: TypeAlias = "PolyfillResult | tuple[list[dict[str, object]], AppliedEdit | None]" + +EditorFn: TypeAlias = "Callable[..., EditorResult | Awaitable[EditorResult]]" _EDITOR_REGISTRY: Final[dict[str, EditorFn]] = { CLEAR_TOOL_USES_EDIT_TYPE: apply_clear_tool_uses_20250919, @@ -19,23 +25,31 @@ _EDITOR_REGISTRY: Final[dict[str, EditorFn]] = { } -def _normalize_spec( - spec: dict[str, Any] | list[dict[str, Any]] | None, -) -> list[dict[str, Any]] | None: - """Accept Anthropic-native dict form or OpenAI list form; return edits list.""" - if isinstance(spec, list): - # Local import to avoid an import cycle at module load. - from litellm.llms.anthropic.chat.transformation import AnthropicConfig - - spec = AnthropicConfig.map_openai_context_management_to_anthropic(spec) - - edits: Final = spec.get("edits") if isinstance(spec, dict) else None +def _edits_from(normalized: dict[str, object] | None) -> list[dict[str, object]] | None: + edits: Final = normalized.get("edits") if isinstance(normalized, dict) else None if not edits or not isinstance(edits, list): return None return [edit for edit in edits if isinstance(edit, dict)] -def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult: +def _normalize_spec( + spec: dict[str, object] | list[dict[str, object]] | None, +) -> list[dict[str, object]] | None: + """Accept Anthropic-native dict form or OpenAI list form; return edits list.""" + if isinstance(spec, list): + # Local import to avoid an import cycle at module load. + from litellm.llms.anthropic.chat.transformation import AnthropicConfig + + return _edits_from(AnthropicConfig.map_openai_context_management_to_anthropic(spec)) + + return _edits_from(spec) + + +def _wrap_editor_return( + raw: EditorResult, + *, + fallback_system: str | list[dict[str, object]] | None, +) -> PolyfillResult: """Coerce an editor's native return shape into a ``PolyfillResult``. v0 sync editors (e.g. ``clear_tool_uses_20250919``) return a 2-tuple @@ -46,7 +60,7 @@ def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult: return raw # Legacy 2-tuple return — sync editors don't mutate ``system``, so # carry the caller's value forward. - messages, applied = cast(tuple[list[dict[str, Any]], Any], raw) + messages, applied = raw return PolyfillResult( messages=messages, system=fallback_system, @@ -57,13 +71,13 @@ def _wrap_editor_return(raw: Any, *, fallback_system: Any) -> PolyfillResult: async def apply_context_management( *, model: str, - messages: list[dict[str, Any]], - tools: list[dict[str, Any]] | None, - system: Any, - context_management_spec: dict[str, Any] | list[dict[str, Any]] | None, - litellm_metadata: dict[str, Any] | None = None, - llm_router: Any = None, - user_api_key_auth: Any = None, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + system: str | list[dict[str, object]] | None, + context_management_spec: dict[str, object] | list[dict[str, object]] | None, + litellm_metadata: dict[str, object] | None = None, + llm_router: "Router | None" = None, + user_api_key_auth: "UserAPIKeyAuth | None" = None, ) -> PolyfillResult: """Run edits in order; return a single ``PolyfillResult``. @@ -92,22 +106,30 @@ async def apply_context_management( ) continue - kwargs: dict[str, Any] = { - "model": model, - "messages": current_messages, - "tools": tools, - "system": current_system, - "edit_spec": edit_spec, - } # Only async editors accept these — passing them to sync v0 editors # would break their signature. - if inspect.iscoroutinefunction(editor): - kwargs["litellm_metadata"] = litellm_metadata - kwargs["llm_router"] = llm_router - kwargs["user_api_key_auth"] = user_api_key_auth - raw_result = await cast(Callable[..., Awaitable[Any]], editor)(**kwargs) - else: - raw_result = editor(**kwargs) + editor_is_async = inspect.iscoroutinefunction(editor) + editor_return = ( + editor( + model=model, + messages=current_messages, + tools=tools, + system=current_system, + edit_spec=edit_spec, + litellm_metadata=litellm_metadata, + llm_router=llm_router, + user_api_key_auth=user_api_key_auth, + ) + if editor_is_async + else editor( + model=model, + messages=current_messages, + tools=tools, + system=current_system, + edit_spec=edit_spec, + ) + ) + raw_result = editor_return if isinstance(editor_return, (PolyfillResult, tuple)) else await editor_return result = _wrap_editor_return(raw_result, fallback_system=current_system) diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/clear_tool_uses.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/clear_tool_uses.py index 393c0507d2b..00ecb315bf1 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/clear_tool_uses.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/clear_tool_uses.py @@ -2,6 +2,8 @@ from typing import Any, Final, cast +from typing_extensions import ReadOnly, TypedDict + import litellm from litellm._logging import verbose_logger from litellm.types.llms.anthropic import AppliedEdit @@ -14,7 +16,18 @@ from ..constants import ( from ..placeholders import build_cleared_tool_result_content -def _count_tool_uses(messages: list[dict[str, Any]]) -> int: +class ClearToolUsesEditSpec(TypedDict, total=False): + """The ``clear_tool_uses_20250919`` entry of a ``context_management`` spec.""" + + type: ReadOnly[str] + trigger: ReadOnly[dict[str, object]] + keep: ReadOnly[dict[str, object]] + clear_at_least: ReadOnly[object] + exclude_tools: ReadOnly[object] + clear_tool_inputs: ReadOnly[object] + + +def _count_tool_uses(messages: list[dict[str, object]]) -> int: """Return the number of tool_use content blocks across all messages. Only counts blocks with a string ``id`` to stay consistent with @@ -32,7 +45,7 @@ def _count_tool_uses(messages: list[dict[str, Any]]) -> int: return count -def _collect_tool_use_ids_in_order(messages: list[dict[str, Any]]) -> list[str]: +def _collect_tool_use_ids_in_order(messages: list[dict[str, object]]) -> list[str]: """Return tool_use ids in the chronological order they appear in messages.""" ids: Final[list[str]] = [] for msg in messages: @@ -47,10 +60,10 @@ def _collect_tool_use_ids_in_order(messages: list[dict[str, Any]]) -> list[str]: def _trigger_met( - trigger: dict[str, Any], + trigger: dict[str, object], model: str, - messages: list[dict[str, Any]], - tools: list[dict[str, Any]] | None, + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, ) -> tuple[bool, int | None]: """Return (trigger_met, input_tokens if counted for reuse).""" trigger_type: Final = trigger.get("type", "input_tokens") @@ -73,7 +86,7 @@ def _trigger_met( return current_tokens > threshold, current_tokens -def _resolve_keep_count(keep: dict[str, Any]) -> int: +def _resolve_keep_count(keep: dict[str, object]) -> int: keep_type: Final = keep.get("type", "tool_uses") if keep_type != "tool_uses": return DEFAULT_KEEP_TOOL_USES @@ -84,7 +97,7 @@ def _resolve_keep_count(keep: dict[str, Any]) -> int: def _last_completed_tool_use_id( - messages: list[dict[str, Any]], + messages: list[dict[str, object]], ) -> str | None: """Latest completed tool_result id; never cleared.""" last_id: str | None = None @@ -99,17 +112,19 @@ def _last_completed_tool_use_id( return last_id -def _clear_tool_results(messages: list[dict[str, Any]], ids_to_clear: set) -> tuple[list[dict[str, Any]], int]: +def _clear_tool_results( + messages: list[dict[str, object]], ids_to_clear: set[str] +) -> tuple[list[dict[str, object]], int]: """Clear matching tool_result content; return (messages, cleared_count).""" cleared = 0 - new_messages: Final[list[dict[str, Any]]] = [] + new_messages: Final[list[dict[str, object]]] = [] for msg in messages: content = msg.get("content") if not isinstance(content, list): new_messages.append(msg) continue - new_blocks: list[Any] = [] + new_blocks: list[object] = [] mutated = False for block in content: if ( @@ -138,11 +153,11 @@ def _clear_tool_results(messages: list[dict[str, Any]], ids_to_clear: set) -> tu def apply_clear_tool_uses_20250919( *, model: str, - messages: list[dict[str, Any]], - tools: list[dict[str, Any]] | None, - system: Any, - edit_spec: dict[str, Any], -) -> tuple[list[dict[str, Any]], AppliedEdit | None]: + messages: list[dict[str, object]], + tools: list[dict[str, object]] | None, + system: str | list[dict[str, object]] | None, + edit_spec: ClearToolUsesEditSpec, +) -> tuple[list[dict[str, object]], AppliedEdit | None]: """Apply clear_tool_uses; return (messages, AppliedEdit or None).""" ignored_knobs = [knob for knob in ("clear_at_least", "exclude_tools", "clear_tool_inputs") if knob in edit_spec] for ignored_knob in ignored_knobs: @@ -153,11 +168,11 @@ def apply_clear_tool_uses_20250919( CLEAR_TOOL_USES_EDIT_TYPE, ) - trigger: Final = edit_spec.get("trigger") or { + trigger: Final[dict[str, object]] = edit_spec.get("trigger") or { "type": "input_tokens", "value": DEFAULT_INPUT_TOKENS_TRIGGER, } - keep: Final = edit_spec.get("keep") or { + keep: Final[dict[str, object]] = edit_spec.get("keep") or { "type": "tool_uses", "value": DEFAULT_KEEP_TOOL_USES, } diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py index 55a85c011d0..171f5156594 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/agentic_streaming_iterator.py @@ -18,11 +18,14 @@ import asyncio import contextlib import json from collections.abc import AsyncIterator -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, cast from litellm._logging import verbose_logger from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + HOLD_BACK_PING_INTERVAL_SECONDS: Final = 15.0 SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES: Final = ( b"event: error\n" @@ -181,7 +184,7 @@ class AgenticAnthropicStreamingIterator: messages: list[dict], anthropic_messages_provider_config: Any, anthropic_messages_optional_request_params: dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", custom_llm_provider: str, kwargs: dict, hold_back: bool = False, diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py index 3e314f76a3e..69985bcdaa3 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/handler.py @@ -571,7 +571,34 @@ def anthropic_messages_handler( anthropic_messages_provider_config = OpenAILikeAnthropicMessagesConfig() if anthropic_messages_provider_config is None: # Route to Responses API for OpenAI / Azure, chat/completions for everything else. - _shared_kwargs: Final = dict( + if _should_route_to_responses_api(custom_llm_provider, original_model, model): + return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler( + max_tokens=max_tokens, + messages=messages, + model=original_model, + metadata=metadata, + stop_sequences=stop_sequences, + stream=stream, + system=system, + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + _is_async=is_async, + api_key=api_key, + api_base=api_base, + client=client, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) + + # The in-gateway context_management polyfill runs inside + # ``async_anthropic_messages_handler`` so it can ``await`` the + # summarization model for ``compact_20260112``. ``context_management`` + # is passed through as a regular kwarg. + return LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( max_tokens=max_tokens, messages=messages, model=original_model, @@ -592,16 +619,6 @@ def anthropic_messages_handler( custom_llm_provider=custom_llm_provider, **kwargs, ) - if _should_route_to_responses_api(custom_llm_provider, original_model, model): - return LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler(**_shared_kwargs) - - # The in-gateway context_management polyfill runs inside - # ``async_anthropic_messages_handler`` so it can ``await`` the - # summarization model for ``compact_20260112``. ``context_management`` - # is passed through as a regular kwarg. - return LiteLLMMessagesToCompletionTransformationHandler.anthropic_messages_handler( - **_shared_kwargs, - ) if custom_llm_provider is None: raise ValueError( diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 8387dd8310d..a282d5f4d4f 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -342,7 +342,7 @@ class BaseAnthropicMessagesStreamingIterator: self.start_time = datetime.now() self.completion_start_time: datetime | None = None - async def _handle_streaming_logging(self, collected_chunks: list[bytes]): + async def _handle_streaming_logging(self, collected_chunks: list[bytes], *, stream_teardown: bool = False): """Handle the logging after all chunks have been collected.""" from litellm.proxy.pass_through_endpoints.streaming_handler import ( PassThroughStreamingHandler, @@ -354,21 +354,26 @@ class BaseAnthropicMessagesStreamingIterator: if self.completion_start_time is not None: self.litellm_logging_obj.completion_start_time = self.completion_start_time self.litellm_logging_obj.model_call_details["completion_start_time"] = self.completion_start_time + logging_coroutine: Final = PassThroughStreamingHandler._route_streaming_logging_to_handler( + litellm_logging_obj=self.litellm_logging_obj, + passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, + url_route="/v1/messages", + request_body=self.request_body or {}, + endpoint_type=EndpointType.ANTHROPIC, + start_time=self.start_time, + raw_bytes=collected_chunks, + end_time=end_time, + ) + deferred_dispatch_armed: Final = ( + getattr(self.litellm_logging_obj, "_on_deferred_stream_complete", None) is not None + ) + if deferred_dispatch_armed and not stream_teardown: + self.litellm_logging_obj._deferred_stream_complete_args = (logging_coroutine,) + return # Enqueue on the rooted logging worker rather than asyncio.create_task: # this also runs during generator teardown after a client disconnect, # where an unrooted task could be garbage-collected before it bills. - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=PassThroughStreamingHandler._route_streaming_logging_to_handler( - litellm_logging_obj=self.litellm_logging_obj, - passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, - url_route="/v1/messages", - request_body=self.request_body or {}, - endpoint_type=EndpointType.ANTHROPIC, - start_time=self.start_time, - raw_bytes=collected_chunks, - end_time=end_time, - ) - ) + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) def get_async_streaming_response_iterator( self, @@ -433,7 +438,7 @@ class BaseAnthropicMessagesStreamingIterator: # post-loop logging below never runs and the tokens already streamed # (and billed by the provider) would never reach spend tracking. See LIT-5839. if collected_chunks: - await self._handle_streaming_logging(collected_chunks) + await self._handle_streaming_logging(collected_chunks, stream_teardown=True) raise if not saw_terminal_event: diff --git a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py index c1ea39fd72c..b6ec9520e79 100644 --- a/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/experimental_pass_through/responses_adapters/handler.py @@ -5,10 +5,11 @@ Used when the target model is an OpenAI or Azure model. """ from collections.abc import AsyncIterator, Coroutine, Mapping -from typing import Any, Final +from typing import Any, Final, TypeAlias import litellm from litellm.types.llms.anthropic import ( + AllAnthropicMessageValues, AllAnthropicToolsValues, AnthropicMessagesRequest, AnthropicOutputConfig, @@ -23,6 +24,8 @@ from ..utils import local_model_name from .streaming_iterator import AnthropicResponsesStreamWrapper from .transformation import LiteLLMAnthropicToResponsesAPIAdapter +AnthropicRequestMessages: TypeAlias = list[AllAnthropicMessageValues] | list[dict[str, object]] + _ADAPTER: Final = LiteLLMAnthropicToResponsesAPIAdapter() @@ -34,22 +37,22 @@ def _forwarded_kwargs(extra_kwargs: Mapping[str, object] | None) -> Mapping[str, def _build_responses_kwargs( *, max_tokens: int, - messages: list[dict], + messages: AnthropicRequestMessages, model: str, - context_management: dict | None = None, - metadata: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[AllAnthropicToolsValues | dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: AnthropicOutputSchema | None = None, - extra_kwargs: dict[str, Any] | None = None, + extra_kwargs: Mapping[str, object] | None = None, ) -> dict[str, Any]: """ Build the kwargs dict to pass directly to litellm.responses() / litellm.aresponses(). @@ -83,30 +86,32 @@ def _build_responses_kwargs( anthropic_request: Final = AnthropicMessagesRequest(**request_data) responses_kwargs: Final = _ADAPTER.translate_request(anthropic_request) + forwarded_kwargs: Final = _forwarded_kwargs(extra_kwargs) # Normalize reasoning effort based on model capabilities # (e.g. "max" → "xhigh"/"high", "minimal" → "low" if unsupported) reasoning: Final = responses_kwargs.get("reasoning") - if isinstance(reasoning, dict) and "effort" in reasoning: - from litellm.llms.anthropic.experimental_pass_through.utils import ( - normalize_reasoning_effort_value, - ) + if isinstance(reasoning, dict): + effort: Final[object] = reasoning.get("effort") + if isinstance(effort, str): + from litellm.llms.anthropic.experimental_pass_through.utils import ( + normalize_reasoning_effort_value, + ) - effort: Final = reasoning["effort"] - normalized: Final = normalize_reasoning_effort_value( - effort, - model=model, - custom_llm_provider=(extra_kwargs or {}).get("custom_llm_provider"), - ) - if normalized != effort: - responses_kwargs["reasoning"] = {**reasoning, "effort": normalized} + provider_hint: Final = forwarded_kwargs.get("custom_llm_provider") + normalized: Final = normalize_reasoning_effort_value( + effort, + model=model, + custom_llm_provider=provider_hint if isinstance(provider_hint, str) else None, + ) + if normalized != effort: + responses_kwargs["reasoning"] = {**reasoning, "effort": normalized} if stream: responses_kwargs["stream"] = True # Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.) excluded: Final = {"anthropic_messages"} - forwarded_kwargs: Final = _forwarded_kwargs(extra_kwargs) for key, value in forwarded_kwargs.items(): if key == "litellm_logging_obj" and value is not None: from litellm.litellm_core_utils.litellm_logging import ( @@ -140,18 +145,18 @@ class LiteLLMMessagesToResponsesAPIHandler: @staticmethod async def async_anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: AnthropicRequestMessages, model: str, - context_management: dict | None = None, - metadata: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[AllAnthropicToolsValues | dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: AnthropicOutputSchema | None = None, @@ -193,18 +198,18 @@ class LiteLLMMessagesToResponsesAPIHandler: @staticmethod def anthropic_messages_handler( max_tokens: int, - messages: list[dict], + messages: AnthropicRequestMessages, model: str, - context_management: dict | None = None, - metadata: dict | None = None, + context_management: dict[str, object] | None = None, + metadata: dict[str, object] | None = None, output_config: AnthropicOutputConfig | None = None, stop_sequences: list[str] | None = None, stream: bool | None = False, system: str | None = None, temperature: float | None = None, - thinking: dict | None = None, - tool_choice: dict | None = None, - tools: list[AllAnthropicToolsValues | dict] | None = None, + thinking: dict[str, object] | None = None, + tool_choice: dict[str, object] | None = None, + tools: list[AllAnthropicToolsValues | dict[str, object]] | None = None, top_k: int | None = None, top_p: float | None = None, output_format: AnthropicOutputSchema | None = None, diff --git a/litellm/llms/anthropic/skills/transformation.py b/litellm/llms/anthropic/skills/transformation.py index 566322bbdd6..448e2dc2584 100644 --- a/litellm/llms/anthropic/skills/transformation.py +++ b/litellm/llms/anthropic/skills/transformation.py @@ -2,9 +2,10 @@ Anthropic Skills API configuration and transformations """ -from typing import Any, Final +from typing import Final import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger from litellm.litellm_core_utils.url_utils import encode_url_path_segment @@ -22,6 +23,8 @@ from litellm.types.llms.anthropic_skills import ( from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders +_RAW_JSON_PAYLOAD: Final = TypeAdapter(object) + class AnthropicSkillsConfig(BaseSkillsAPIConfig): """Anthropic-specific Skills API configuration""" @@ -104,10 +107,10 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> Skill: """Transform Anthropic response to Skill object""" - response_json: Final = raw_response.json() + response_json: Final = _RAW_JSON_PAYLOAD.validate_python(raw_response.json()) verbose_logger.debug("Transforming create skill response: %s", response_json) - return Skill(**response_json) + return Skill.model_validate(response_json) def transform_list_skills_request( self, @@ -122,13 +125,12 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): url: Final = self.get_complete_url(api_base=api_base, endpoint="skills") # Build query parameters - query_params: Final[dict[str, Any]] = {} - if "limit" in list_params and list_params["limit"]: - query_params["limit"] = list_params["limit"] - if "page" in list_params and list_params["page"]: - query_params["page"] = list_params["page"] - if "source" in list_params and list_params["source"]: - query_params["source"] = list_params["source"] + limit: Final = list_params.get("limit") + page: Final = list_params.get("page") + source: Final = list_params.get("source") + query_params: Final[dict[str, int | str]] = { + key: value for key, value in (("limit", limit), ("page", page), ("source", source)) if value + } verbose_logger.debug( "List skills request made to Anthropic Skills endpoint with params: %s", @@ -143,10 +145,10 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> ListSkillsResponse: """Transform Anthropic response to ListSkillsResponse""" - response_json: Final = raw_response.json() + response_json: Final = _RAW_JSON_PAYLOAD.validate_python(raw_response.json()) verbose_logger.debug("Transforming list skills response: %s", response_json) - return ListSkillsResponse(**response_json) + return ListSkillsResponse.model_validate(response_json) def transform_get_skill_request( self, @@ -168,10 +170,10 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> Skill: """Transform Anthropic response to Skill object""" - response_json: Final = raw_response.json() + response_json: Final = _RAW_JSON_PAYLOAD.validate_python(raw_response.json()) verbose_logger.debug("Transforming get skill response: %s", response_json) - return Skill(**response_json) + return Skill.model_validate(response_json) def transform_delete_skill_request( self, @@ -193,7 +195,7 @@ class AnthropicSkillsConfig(BaseSkillsAPIConfig): logging_obj: LiteLLMLoggingObj, ) -> DeleteSkillResponse: """Transform Anthropic response to DeleteSkillResponse""" - response_json: Final = raw_response.json() + response_json: Final = _RAW_JSON_PAYLOAD.validate_python(raw_response.json()) verbose_logger.debug("Transforming delete skill response: %s", response_json) - return DeleteSkillResponse(**response_json) + return DeleteSkillResponse.model_validate(response_json) diff --git a/litellm/llms/aws_polly/text_to_speech/transformation.py b/litellm/llms/aws_polly/text_to_speech/transformation.py index 68630335ca7..8f96f80d15e 100644 --- a/litellm/llms/aws_polly/text_to_speech/transformation.py +++ b/litellm/llms/aws_polly/text_to_speech/transformation.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, Final, Union import httpx +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.base_llm.text_to_speech.transformation import ( BaseTextToSpeechConfig, TextToSpeechRequestData, @@ -238,7 +239,7 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): return api_base.rstrip("/") + "/v1/speech" aws_region_name: Final = litellm_params.get("aws_region_name", self.DEFAULT_REGION) - return f"https://polly.{aws_region_name}.amazonaws.com/v1/speech" + return f"https://polly.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/v1/speech" def is_ssml_input(self, input: str) -> bool: """ diff --git a/litellm/llms/azure/audio_transcriptions.py b/litellm/llms/azure/audio_transcriptions.py index 3ab0bd18b45..4a5ed2ccb0c 100644 --- a/litellm/llms/azure/audio_transcriptions.py +++ b/litellm/llms/azure/audio_transcriptions.py @@ -1,5 +1,5 @@ from collections.abc import Coroutine -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final from openai import AsyncAzureOpenAI, AzureOpenAI from pydantic import BaseModel @@ -16,6 +16,9 @@ from litellm.utils import ( from .azure import AzureChatCompletion from .common_utils import AzureOpenAIError +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + class AzureAudioTranscription(AzureChatCompletion): def audio_transcriptions( @@ -23,7 +26,7 @@ class AzureAudioTranscription(AzureChatCompletion): model: str, audio_file: FileTypes, optional_params: dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", model_response: TranscriptionResponse, timeout: float, max_retries: int, @@ -112,7 +115,7 @@ class AzureAudioTranscription(AzureChatCompletion): data: dict, model_response: TranscriptionResponse, timeout: float, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", api_version: str | None = None, api_key: str | None = None, api_base: str | None = None, diff --git a/litellm/llms/azure/chat/gpt_transformation.py b/litellm/llms/azure/chat/gpt_transformation.py index 30fc3635d9d..2df4ab731ab 100644 --- a/litellm/llms/azure/chat/gpt_transformation.py +++ b/litellm/llms/azure/chat/gpt_transformation.py @@ -23,6 +23,8 @@ from ...base_llm.chat.transformation import BaseConfig from ..common_utils import AzureOpenAIError if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj LoggingClass = LiteLLMLoggingObj @@ -271,7 +273,7 @@ class AzureOpenAIConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/azure/completion/handler.py b/litellm/llms/azure/completion/handler.py index 728968e12e7..80934e994f6 100644 --- a/litellm/llms/azure/completion/handler.py +++ b/litellm/llms/azure/completion/handler.py @@ -193,7 +193,7 @@ class AzureTextCompletion(BaseAzureLLM): data: dict, timeout: Any, model_response: ModelResponse, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj, max_retries: int, azure_ad_token: str | None = None, client=None, # this is the AsyncAzureOpenAI diff --git a/litellm/llms/azure/files/handler.py b/litellm/llms/azure/files/handler.py index 4f93896699f..67bf47c2359 100644 --- a/litellm/llms/azure/files/handler.py +++ b/litellm/llms/azure/files/handler.py @@ -48,7 +48,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): verbose_logger.debug("create_file_data=%s", create_file_data) response = await openai_client.files.create(**self._prepare_create_file_data(create_file_data)) verbose_logger.debug("create_file_response=%s", response) - return OpenAIFileObject(**response.model_dump()) + return OpenAIFileObject.model_validate(response.model_dump()) def create_file( self, @@ -60,8 +60,8 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): timeout: float | httpx.Timeout, max_retries: int | None, client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, - litellm_params: dict | None = None, - ) -> OpenAIFileObject | Coroutine[Any, Any, OpenAIFileObject]: + litellm_params: dict[str, object] | None = None, + ) -> OpenAIFileObject | Coroutine[object, object, OpenAIFileObject]: openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( litellm_params=litellm_params or {}, api_key=api_key, @@ -84,7 +84,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): response: Final = cast(AzureOpenAI | OpenAI, openai_client).files.create( **self._prepare_create_file_data(create_file_data) ) - return OpenAIFileObject(**response.model_dump()) + return OpenAIFileObject.model_validate(response.model_dump()) async def afile_content( self, @@ -104,8 +104,8 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): max_retries: int | None, api_version: str | None = None, client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, - litellm_params: dict | None = None, - ) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]: + litellm_params: dict[str, object] | None = None, + ) -> HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]: openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( litellm_params=litellm_params or {}, api_key=api_key, @@ -150,7 +150,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): max_retries: int | None, api_version: str | None = None, client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, - litellm_params: dict | None = None, + litellm_params: dict[str, object] | None = None, ): openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( litellm_params=litellm_params or {}, @@ -200,7 +200,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): organization: str | None = None, api_version: str | None = None, client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, - litellm_params: dict | None = None, + litellm_params: dict[str, object] | None = None, ): openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( litellm_params=litellm_params or {}, @@ -252,7 +252,7 @@ class AzureOpenAIFilesAPI(BaseAzureLLM): purpose: str | None = None, api_version: str | None = None, client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = None, - litellm_params: dict | None = None, + litellm_params: dict[str, object] | None = None, ): openai_client: AzureOpenAI | AsyncAzureOpenAI | OpenAI | AsyncOpenAI | None = self.get_azure_openai_client( litellm_params=litellm_params or {}, diff --git a/litellm/llms/azure_ai/agents/transformation.py b/litellm/llms/azure_ai/agents/transformation.py index b81e6b0d62d..60ce81a23c7 100644 --- a/litellm/llms/azure_ai/agents/transformation.py +++ b/litellm/llms/azure_ai/agents/transformation.py @@ -34,6 +34,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -295,7 +297,7 @@ class AzureAIAgentsConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/azure_ai/azure_model_router/transformation.py b/litellm/llms/azure_ai/azure_model_router/transformation.py index 61cbc213b11..9e35e396e15 100644 --- a/litellm/llms/azure_ai/azure_model_router/transformation.py +++ b/litellm/llms/azure_ai/azure_model_router/transformation.py @@ -5,7 +5,7 @@ The Model Router is a special Azure AI deployment that automatically routes requ to the best available model. It has specific cost tracking requirements. """ -from typing import Any, Final +from typing import TYPE_CHECKING, Final from httpx import Response @@ -14,6 +14,9 @@ from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse +if TYPE_CHECKING: + import tiktoken + class AzureModelRouterConfig(AzureAIStudioConfig): """ @@ -56,7 +59,7 @@ class AzureModelRouterConfig(AzureAIStudioConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 9e7161120cc..7fe9d3dec52 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -1,7 +1,7 @@ import copy import enum import re -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Final, cast from urllib.parse import urlparse import httpx @@ -25,6 +25,9 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ModelResponse, ProviderField from litellm.utils import _add_path_to_api_base, supports_tool_choice +if TYPE_CHECKING: + import tiktoken + class AzureFoundryErrorStrings(str, enum.Enum): SET_EXTRA_PARAMETERS_TO_PASS_THROUGH = "Set extra-parameters to 'pass-through'" @@ -258,7 +261,7 @@ class AzureAIStudioConfig(OpenAIConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py index 02e62f27d02..64f81956ad7 100644 --- a/litellm/llms/azure_ai/image_generation/mai_transformation.py +++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py @@ -11,6 +11,7 @@ from litellm.types.utils import ImageResponse from litellm.utils import convert_to_model_response_object if TYPE_CHECKING: + import tiktoken from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj @@ -199,7 +200,7 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py index a0e427eab9a..f5126f81006 100644 --- a/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py +++ b/litellm/llms/azure_ai/ocr/document_intelligence/transformation.py @@ -12,7 +12,7 @@ import asyncio import re import time from collections.abc import Mapping -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final from urllib.parse import quote import httpx @@ -41,6 +41,9 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.secret_managers.main import get_secret_str +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + AZURE_DOCUMENT_INTELLIGENCE_API_KEY_ENV_VAR: Final = "AZURE_DOCUMENT_INTELLIGENCE_API_KEY" @@ -676,7 +679,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): self, model: str, raw_response: httpx.Response, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", **kwargs, ) -> OCRResponse: """ @@ -751,7 +754,7 @@ class AzureDocumentIntelligenceOCRConfig(BaseOCRConfig): self, model: str, raw_response: httpx.Response, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", **kwargs, ) -> OCRResponse: """ diff --git a/litellm/llms/base.py b/litellm/llms/base.py index 7dec5509c46..8f6f45f4d35 100644 --- a/litellm/llms/base.py +++ b/litellm/llms/base.py @@ -6,6 +6,7 @@ import httpx import litellm if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.utils import ModelResponse, TextCompletionResponse @@ -19,7 +20,7 @@ class BaseLLM: response: httpx.Response, model_response: "ModelResponse", stream: bool, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", optional_params: dict, api_key: str, data: dict | str, @@ -38,7 +39,7 @@ class BaseLLM: response: httpx.Response, model_response: "TextCompletionResponse", stream: bool, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", optional_params: dict, api_key: str, data: dict | str, diff --git a/litellm/llms/base_llm/audio_transcription/transformation.py b/litellm/llms/base_llm/audio_transcription/transformation.py index da1776d8dc7..b323c4812b5 100644 --- a/litellm/llms/base_llm/audio_transcription/transformation.py +++ b/litellm/llms/base_llm/audio_transcription/transformation.py @@ -12,6 +12,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import FileTypes, ModelResponse, TranscriptionResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -110,7 +112,7 @@ class BaseAudioTranscriptionConfig(BaseConfig, ABC): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/base_llm/bridges/completion_transformation.py b/litellm/llms/base_llm/bridges/completion_transformation.py index 2d5879dc8e3..87b55152d09 100644 --- a/litellm/llms/base_llm/bridges/completion_transformation.py +++ b/litellm/llms/base_llm/bridges/completion_transformation.py @@ -4,9 +4,10 @@ Bridge for transforming API requests to another API requests from abc import ABC, abstractmethod from collections.abc import AsyncIterator, Iterator -from typing import TYPE_CHECKING, Any, Union +from typing import TYPE_CHECKING, Union if TYPE_CHECKING: + import tiktoken from pydantic import BaseModel from litellm import LiteLLMLoggingObj, ModelResponse @@ -38,7 +39,7 @@ class CompletionTransformationBridge(ABC): messages: list["AllMessageValues"], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> "ModelResponse": diff --git a/litellm/llms/base_llm/chat/transformation.py b/litellm/llms/base_llm/chat/transformation.py index 11c763ceb9a..bbe1cc85df1 100644 --- a/litellm/llms/base_llm/chat/transformation.py +++ b/litellm/llms/base_llm/chat/transformation.py @@ -21,6 +21,8 @@ from litellm.types.llms.openai import ( ) if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.utils import ModelResponse @@ -342,7 +344,7 @@ class BaseConfig(ABC): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> "ModelResponse": diff --git a/litellm/llms/base_llm/completion/transformation.py b/litellm/llms/base_llm/completion/transformation.py index c38199b0966..fb472dfa63b 100644 --- a/litellm/llms/base_llm/completion/transformation.py +++ b/litellm/llms/base_llm/completion/transformation.py @@ -8,6 +8,8 @@ from litellm.types.llms.openai import AllMessageValues, OpenAITextCompletionUser from litellm.types.utils import ModelResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -66,7 +68,7 @@ class BaseTextCompletionConfig(BaseConfig, ABC): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/base_llm/embedding/transformation.py b/litellm/llms/base_llm/embedding/transformation.py index 0330c0118bd..da87dcc7f98 100644 --- a/litellm/llms/base_llm/embedding/transformation.py +++ b/litellm/llms/base_llm/embedding/transformation.py @@ -8,6 +8,8 @@ from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues from litellm.types.utils import EmbeddingResponse, ModelResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -78,7 +80,7 @@ class BaseEmbeddingConfig(BaseConfig, ABC): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/base_llm/files/transformation.py b/litellm/llms/base_llm/files/transformation.py index b20fe0f1560..7a7088c2fb5 100644 --- a/litellm/llms/base_llm/files/transformation.py +++ b/litellm/llms/base_llm/files/transformation.py @@ -20,6 +20,8 @@ from litellm.types.utils import LlmProviders, ModelResponse from ..chat.transformation import BaseConfig if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from litellm.router import Router as _Router from litellm.types.llms.openai import HttpxBinaryResponseContent @@ -207,7 +209,7 @@ class BaseFilesConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/base_llm/image_generation/transformation.py b/litellm/llms/base_llm/image_generation/transformation.py index 4ce4add0432..4616441133e 100644 --- a/litellm/llms/base_llm/image_generation/transformation.py +++ b/litellm/llms/base_llm/image_generation/transformation.py @@ -11,6 +11,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -91,7 +93,7 @@ class BaseImageGenerationConfig(ABC): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/base_llm/image_variations/transformation.py b/litellm/llms/base_llm/image_variations/transformation.py index beae828c301..d3e02139e0e 100644 --- a/litellm/llms/base_llm/image_variations/transformation.py +++ b/litellm/llms/base_llm/image_variations/transformation.py @@ -17,6 +17,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -80,7 +82,7 @@ class BaseImageVariationConfig(BaseConfig, ABC): image: FileTypes, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, ) -> ImageResponse: pass @@ -96,7 +98,7 @@ class BaseImageVariationConfig(BaseConfig, ABC): image: FileTypes, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, ) -> ImageResponse: pass @@ -123,7 +125,7 @@ class BaseImageVariationConfig(BaseConfig, ABC): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 4332848e545..852cfaa24f2 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -23,6 +23,7 @@ from litellm.constants import ( BEDROCK_MAX_POLICY_SIZE, STS_CREDENTIAL_EXPIRY_SAFETY_MARGIN_SECONDS, ) +from litellm.litellm_core_utils.aws_partition import contains_bedrock_arn, get_aws_dns_suffix from litellm.litellm_core_utils.dd_tracing import tracer from litellm.secret_managers.main import get_secret, get_secret_str @@ -348,7 +349,7 @@ class BaseAWSLLM: def _get_aws_region_from_model_arn(self, model: str | None) -> str | None: try: # First check if the string contains the expected prefix - if not isinstance(model, str) or "arn:aws:bedrock" not in model: + if not isinstance(model, str) or not contains_bedrock_arn(model): return None # Split the ARN and check if we have enough parts @@ -625,24 +626,29 @@ class BaseAWSLLM: return match.group(1) if match else None @staticmethod - def _resolve_sts_region(aws_sts_endpoint: str | None = None) -> str | None: - """STS signing region: parsed from aws_sts_endpoint else AWS_REGION / AWS_DEFAULT_REGION.""" + def _resolve_sts_region( + aws_sts_endpoint: str | None = None, + aws_region_name: str | None = None, + ) -> str | None: + """STS signing region: parsed from aws_sts_endpoint, else AWS_REGION / AWS_DEFAULT_REGION, else the configured aws_region_name.""" return ( BaseAWSLLM._parse_sts_region_from_endpoint(aws_sts_endpoint) or os.getenv("AWS_REGION") or os.getenv("AWS_DEFAULT_REGION") + or aws_region_name ) def _build_sts_client_kwargs( self, aws_sts_endpoint: str | None = None, ssl_verify: bool | str | None = None, + aws_region_name: str | None = None, ) -> dict: """STS client kwargs with aligned endpoint_url and region_name (SigV4).""" kwargs: Final[dict] = {"verify": self._get_ssl_verify(ssl_verify)} if aws_sts_endpoint is not None: kwargs["endpoint_url"] = aws_sts_endpoint - sts_region: Final = self._resolve_sts_region(aws_sts_endpoint) + sts_region: Final = self._resolve_sts_region(aws_sts_endpoint, aws_region_name) if sts_region is not None: kwargs["region_name"] = sts_region return kwargs @@ -837,6 +843,7 @@ class BaseAWSLLM: sts_client_kwargs: Final = self._build_sts_client_kwargs( aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, + aws_region_name=aws_region_name, ) with tracer.trace("boto3.client(sts)"): @@ -948,6 +955,7 @@ class BaseAWSLLM: aws_external_id: str | None = None, aws_sts_endpoint: str | None = None, ssl_verify: bool | str | None = None, + aws_region_name: str | None = None, ) -> dict: """Handle cross-account role assumption for IRSA.""" import boto3 @@ -961,6 +969,7 @@ class BaseAWSLLM: irsa_sts_kwargs: Final = self._build_sts_client_kwargs( aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, + aws_region_name=aws_region_name, ) # Create an STS client without credentials @@ -1017,6 +1026,7 @@ class BaseAWSLLM: aws_external_id: str | None = None, aws_sts_endpoint: str | None = None, ssl_verify: bool | str | None = None, + aws_region_name: str | None = None, ) -> dict: """Handle same-account role assumption for IRSA.""" import boto3 @@ -1024,6 +1034,7 @@ class BaseAWSLLM: irsa_sts_kwargs: Final = self._build_sts_client_kwargs( aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, + aws_region_name=aws_region_name, ) verbose_logger.debug("Same account role assumption, using automatic IRSA") @@ -1153,6 +1164,7 @@ class BaseAWSLLM: aws_external_id, aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, + aws_region_name=aws_region_name, ) else: sts_response = self._handle_irsa_same_account( @@ -1161,6 +1173,7 @@ class BaseAWSLLM: aws_external_id, aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, + aws_region_name=aws_region_name, ) return self._extract_credentials_and_ttl(sts_response) @@ -1182,6 +1195,7 @@ class BaseAWSLLM: sts_client_kwargs: Final = self._build_sts_client_kwargs( aws_sts_endpoint=aws_sts_endpoint, ssl_verify=ssl_verify, + aws_region_name=aws_region_name, ) if aws_access_key_id is None and aws_secret_access_key is None: with tracer.trace("boto3.client(sts)"): @@ -1363,14 +1377,15 @@ class BaseAWSLLM: """ Select the default endpoint url based on the endpoint type - Default endpoint url is https://bedrock-runtime.{aws_region_name}.amazonaws.com + Default endpoint url is https://bedrock-runtime.{aws_region_name}.{partition dns suffix} """ + dns_suffix: Final = get_aws_dns_suffix(aws_region_name) if endpoint_type == "agent": - return f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com" + return f"https://bedrock-agent-runtime.{aws_region_name}.{dns_suffix}" elif endpoint_type == "agentcore": - return f"https://bedrock-agentcore.{aws_region_name}.amazonaws.com" + return f"https://bedrock-agentcore.{aws_region_name}.{dns_suffix}" else: - return f"https://bedrock-runtime.{aws_region_name}.amazonaws.com" + return f"https://bedrock-runtime.{aws_region_name}.{dns_suffix}" def _get_boto_credentials_from_optional_params( self, optional_params: dict, model: str | None = None diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py index 6efdd17f98d..4b500897642 100644 --- a/litellm/llms/bedrock/batches/handler.py +++ b/litellm/llms/bedrock/batches/handler.py @@ -1,9 +1,11 @@ +from collections.abc import Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Final, cast from openai.types.batch import BatchRequestCounts from openai.types.batch import Metadata as OpenAIBatchMetadata +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.types.utils import LiteLLMBatch if TYPE_CHECKING: @@ -68,6 +70,19 @@ def _predict_output_file_uri(output_prefix: str, input_uri: str, job_id: str | N return f"{output_prefix}{job_id}/{input_basename}.out" +def _record_counts_from_response(response: Mapping[str, object]) -> BatchRequestCounts | None: + total_records: Final = response.get("totalRecordCount") + success_records: Final = response.get("successRecordCount") + if not isinstance(total_records, int) or not isinstance(success_records, int): + return None + error_records: Final = response.get("errorRecordCount") + return BatchRequestCounts( + total=total_records, + completed=success_records, + failed=error_records if isinstance(error_records, int) else 0, + ) + + def _to_epoch(value: Any) -> int | None: if value is None: return None @@ -271,11 +286,11 @@ class BedrockBatchesHandler: ``aws_external_id``). Unknown keys are ignored. Returns: - ``LiteLLMBatch`` shaped like an OpenAI Batch resource. Note that - ``request_counts`` is always ``(0, 0, 0)`` because - ``GetModelInvocationJob`` does not surface per-record counts; - callers that need accurate counts should parse - ``manifest.json.out`` from the output S3 prefix. + ``LiteLLMBatch`` shaped like an OpenAI Batch resource. + ``request_counts`` maps ``GetModelInvocationJob``'s + ``totalRecordCount`` / ``successRecordCount`` / ``errorRecordCount`` + when the provider reports them, and is ``None`` when it does not + (older botocore, or a status that omits counts). """ try: import boto3 @@ -323,7 +338,9 @@ class BedrockBatchesHandler: api_key="", additional_args={ "complete_input_dict": {"jobIdentifier": batch_id}, - "api_base": (f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{url_path_id}"), + "api_base": ( + f"https://bedrock.{region}.{get_aws_dns_suffix(region)}/model-invocation-job/{url_path_id}" + ), }, ) @@ -386,7 +403,7 @@ class BedrockBatchesHandler: failed_at=completed_at if openai_status == "failed" else None, cancelled_at=completed_at if openai_status == "cancelled" else None, expired_at=completed_at if openai_status == "expired" else None, - request_counts=BatchRequestCounts(total=0, completed=0, failed=0), + request_counts=_record_counts_from_response(response), metadata=openai_batch_metadata, completion_window="24h", endpoint="/v1/chat/completions", diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 04f395f2bf1..7729cdfdb0d 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -1,11 +1,12 @@ import os import re import time -from typing import Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast from httpx import Headers, Response from pydantic import TypeAdapter, ValidationError +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix, is_bedrock_arn from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_BATCH_PREFIX, ) @@ -34,6 +35,9 @@ from ..common_utils import ( resolve_s3_encryption_key_id, ) +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + # Bedrock batch input files are uploaded as # s3://bucket/litellm-bedrock-files-{model, ":" -> "-"}-{uuid4}.jsonl (see # BedrockFilesTransformation._get_s3_object_name). A uuid4 is always 36 hex/dash @@ -138,8 +142,10 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): aws_region_name: Final = self._get_aws_region_name(request_params, model) # Bedrock model invocation job endpoint - # Format: https://bedrock.{region}.amazonaws.com/model-invocation-job - bedrock_endpoint: Final = f"https://bedrock.{aws_region_name}.amazonaws.com/model-invocation-job" + # Format: https://bedrock.{region}.{partition dns suffix}/model-invocation-job + bedrock_endpoint: Final = ( + f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/model-invocation-job" + ) return bedrock_endpoint @@ -238,8 +244,9 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): # For Bedrock, we need to return a pre-signed request with AWS auth headers # Use common utility for AWS signing request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) + aws_region_name: Final = self._get_aws_region_name(request_params, model) endpoint_url: Final = ( - f"https://bedrock.{self._get_aws_region_name(request_params, model)}.amazonaws.com/model-invocation-job" + f"https://bedrock.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/model-invocation-job" ) signed_headers, signed_data = self.common_utils.sign_aws_request( service_name="bedrock", @@ -261,7 +268,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): self, model: str | None, raw_response: Response, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", litellm_params: dict, ) -> LiteLLMBatch: """ @@ -371,7 +378,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): """ # For Bedrock, batch_id should be the full job ARN # The GetModelInvocationJob API expects the full ARN as the identifier - if not batch_id.startswith("arn:aws:bedrock:"): + if not is_bedrock_arn(batch_id): raise ValueError(f"Invalid batch_id format. Expected ARN, got: {batch_id}") # Extract the job identifier from the ARN - use the full ARN path part @@ -390,7 +397,9 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): import urllib.parse as _ul encoded_arn: Final = _ul.quote(batch_id, safe="") - endpoint_url: Final = f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}" + endpoint_url: Final = ( + f"https://bedrock.{region}.{get_aws_dns_suffix(region)}/model-invocation-job/{encoded_arn}" + ) # Use common utility for AWS signing request_params: Final = merge_bedrock_aws_request_params(litellm_params, optional_params) @@ -527,7 +536,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): self, model: str | None, raw_response: Response, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", litellm_params: dict, ) -> LiteLLMBatch: """ diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py index 4a2db621421..690040dd93b 100644 --- a/litellm/llms/bedrock/chat/agentcore/transformation.py +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -13,6 +13,7 @@ import httpx from litellm._logging import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.prompt_templates.common_utils import ( convert_content_list_to_str, ) @@ -38,6 +39,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -97,7 +100,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): if aws_bedrock_runtime_endpoint: base_url = aws_bedrock_runtime_endpoint else: - base_url = f"https://bedrock-agentcore.{region}.amazonaws.com" + base_url = f"https://bedrock-agentcore.{region}.{get_aws_dns_suffix(region)}" # Based on boto3 client.invoke_agent_runtime, the path is: # /runtimes/{URL-ENCODED-ARN}/invocations?qualifier= @@ -974,7 +977,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index d0d97b41692..d22b225b0bd 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -7,7 +7,7 @@ import json import time import types from collections.abc import Mapping -from typing import Final, Literal, cast, overload +from typing import TYPE_CHECKING, Final, Literal, cast, overload import httpx @@ -94,6 +94,9 @@ from ..common_utils import ( normalize_bedrock_opus_output_config_effort, ) +if TYPE_CHECKING: + import tiktoken + # Computer use tool prefixes supported by Bedrock BEDROCK_COMPUTER_USE_TOOLS: Final = [ "computer_use_preview", @@ -1770,7 +1773,7 @@ class AmazonConverseConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_agent/transformation.py b/litellm/llms/bedrock/chat/invoke_agent/transformation.py index 2198e19cd7e..e30ec731d8c 100644 --- a/litellm/llms/bedrock/chat/invoke_agent/transformation.py +++ b/litellm/llms/bedrock/chat/invoke_agent/transformation.py @@ -37,6 +37,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -436,7 +438,7 @@ class AmazonInvokeAgentConfig(BaseConfig, BaseAWSLLM): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py index d86c756ca99..5a3f4f17b8b 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_deepseek_transformation.py @@ -1,4 +1,4 @@ -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, cast from httpx import Response @@ -24,6 +24,9 @@ from litellm.types.utils import ( from .amazon_llama_transformation import AmazonLlamaConfig +if TYPE_CHECKING: + import tiktoken + class AmazonDeepSeekR1Config(AmazonLlamaConfig): def transform_response( @@ -36,7 +39,7 @@ class AmazonDeepSeekR1Config(AmazonLlamaConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py index a8275f1d35f..91c3a363c31 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_moonshot_transformation.py @@ -21,6 +21,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from litellm.types.utils import ModelResponse @@ -200,7 +202,7 @@ class AmazonMoonshotConfig(AmazonInvokeConfig, MoonshotChatConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> "ModelResponse": diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py index 361f53d6ace..5f8ab94b00c 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py @@ -6,7 +6,7 @@ Inherits from `AmazonConverseConfig` Nova + Invoke API Tutorial: https://docs.aws.amazon.com/nova/latest/userguide/using-invoke-api.html """ -from typing import Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -18,6 +18,9 @@ from litellm.types.utils import ModelResponse from ..converse_transformation import AmazonConverseConfig from .base_invoke_transformation import AmazonInvokeConfig +if TYPE_CHECKING: + import tiktoken + class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig): """ @@ -70,7 +73,7 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py index a775db2ebc7..c78375c37bb 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py @@ -7,7 +7,7 @@ The main difference is in the response format: Qwen2 uses "text" field while Qwe Qwen2 + Invoke API Tutorial: https://docs.aws.amazon.com/bedrock/latest/userguide/invoke-imported-model.html """ -from typing import Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -20,6 +20,9 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse, Usage +if TYPE_CHECKING: + import tiktoken + class AmazonQwen2Config(AmazonQwen3Config): """ @@ -41,7 +44,7 @@ class AmazonQwen2Config(AmazonQwen3Config): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py index 7db8d77ff84..e251fb15725 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py @@ -6,7 +6,7 @@ Inherits from `AmazonInvokeConfig` Qwen3 + Invoke API Tutorial: https://docs.aws.amazon.com/bedrock/latest/userguide/invoke-imported-model.html """ -from typing import Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -18,6 +18,9 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse, Usage +if TYPE_CHECKING: + import tiktoken + class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): """ @@ -167,7 +170,7 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py index 591de36dc18..cd8066cda4d 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_twelvelabs_pegasus_transformation.py @@ -25,6 +25,8 @@ from litellm.types.utils import ModelResponse, Usage from litellm.utils import get_base64_str if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -188,7 +190,7 @@ class AmazonTwelveLabsPegasusConfig(AmazonInvokeConfig, BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index b8b07af59c6..40b90014f3b 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -29,6 +29,8 @@ from litellm.types.utils import ModelResponse from litellm.utils import _supports_factory if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -397,7 +399,7 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 333326a766b..37121d2ece7 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -34,6 +34,8 @@ from litellm.types.utils import ModelResponse, Usage from litellm.utils import CustomStreamWrapper if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -286,7 +288,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 4ad20772ed0..72e3cc1b326 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -21,6 +21,7 @@ import httpx import litellm from litellm import verbose_logger +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -434,15 +435,15 @@ def init_bedrock_client( ssl_verify: Final = _get_bedrock_client_ssl_verify() ### SET REGION NAME - if region_name: - pass - elif aws_region_name: - region_name = aws_region_name - elif litellm_aws_region_name: - region_name = litellm_aws_region_name - elif standard_aws_region_name: - region_name = standard_aws_region_name - else: + resolved_region_name: Final = next( + ( + candidate + for candidate in (region_name, aws_region_name, litellm_aws_region_name, standard_aws_region_name) + if isinstance(candidate, str) and candidate + ), + None, + ) + if resolved_region_name is None: raise BedrockError( message="AWS region not set: set AWS_REGION_NAME or AWS_REGION env variable or in .env file", status_code=401, @@ -455,7 +456,7 @@ def init_bedrock_client( elif env_aws_bedrock_runtime_endpoint: endpoint_url = env_aws_bedrock_runtime_endpoint else: - endpoint_url = f"https://bedrock-runtime.{region_name}.amazonaws.com" + endpoint_url = f"https://bedrock-runtime.{resolved_region_name}.{get_aws_dns_suffix(resolved_region_name)}" import boto3 @@ -492,7 +493,7 @@ def init_bedrock_client( aws_access_key_id=sts_response["Credentials"]["AccessKeyId"], aws_secret_access_key=sts_response["Credentials"]["SecretAccessKey"], aws_session_token=sts_response["Credentials"]["SessionToken"], - region_name=region_name, + region_name=resolved_region_name, endpoint_url=endpoint_url, config=config, verify=ssl_verify, @@ -513,7 +514,7 @@ def init_bedrock_client( aws_access_key_id=sts_response["Credentials"]["AccessKeyId"], aws_secret_access_key=sts_response["Credentials"]["SecretAccessKey"], aws_session_token=sts_response["Credentials"]["SessionToken"], - region_name=region_name, + region_name=resolved_region_name, endpoint_url=endpoint_url, config=config, verify=ssl_verify, @@ -526,7 +527,7 @@ def init_bedrock_client( service_name="bedrock-runtime", aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, - region_name=region_name, + region_name=resolved_region_name, endpoint_url=endpoint_url, config=config, verify=ssl_verify, @@ -536,7 +537,7 @@ def init_bedrock_client( client = boto3.Session(profile_name=aws_profile_name).client( service_name="bedrock-runtime", - region_name=region_name, + region_name=resolved_region_name, endpoint_url=endpoint_url, config=config, verify=ssl_verify, @@ -547,7 +548,7 @@ def init_bedrock_client( client = boto3.client( service_name="bedrock-runtime", - region_name=region_name, + region_name=resolved_region_name, endpoint_url=endpoint_url, config=config, verify=ssl_verify, diff --git a/litellm/llms/bedrock/count_tokens/transformation.py b/litellm/llms/bedrock/count_tokens/transformation.py index f87a3bc3452..48fc41ed12b 100644 --- a/litellm/llms/bedrock/count_tokens/transformation.py +++ b/litellm/llms/bedrock/count_tokens/transformation.py @@ -6,7 +6,10 @@ to AWS Bedrock's CountTokens API format and vice versa. """ import re -from typing import Any, Final +from collections.abc import Mapping +from typing import Final, Literal + +from pydantic import JsonValue from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.bedrock.common_utils import get_bedrock_base_model @@ -17,6 +20,48 @@ from litellm.llms.bedrock.common_utils import get_bedrock_base_model DEFAULT_ANTHROPIC_INVOKE_MODEL_MAX_TOKENS: Final = 1024 +def _json_dict(value: JsonValue) -> dict[str, JsonValue]: + return value if isinstance(value, dict) else {} + + +def _json_list(value: JsonValue) -> list[JsonValue]: + return value if isinstance(value, list) else [] + + +def _to_converse_content(content: JsonValue) -> list[JsonValue]: + if isinstance(content, str): + return [{"text": content}] + if isinstance(content, list): + return content + return [] + + +def _to_converse_message(message: JsonValue) -> dict[str, JsonValue]: + fields: Final = _json_dict(message) + return { + "role": fields.get("role"), + "content": _to_converse_content(fields.get("content", "")), + } + + +def _sanitized_bedrock_tool_name(raw_name: JsonValue) -> str: + name: Final = re.sub(r"[^a-zA-Z0-9_]", "_", raw_name if isinstance(raw_name, str) else "") + prefixed: Final = name if not name or name[0].isalpha() else f"t_{name}" + return prefixed[:64] + + +def _to_bedrock_tool_spec(tool: JsonValue) -> dict[str, JsonValue]: + fields: Final = _json_dict(tool) + name: Final = _sanitized_bedrock_tool_name(fields.get("name", "")) + return { + "toolSpec": { + "name": name, + "description": fields.get("description") or name, + "inputSchema": {"json": fields.get("input_schema", {"type": "object", "properties": {}})}, + } + } + + class BedrockCountTokensConfig(BaseAWSLLM): """ Configuration and transformation logic for AWS Bedrock CountTokens API. @@ -27,7 +72,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): - Response: {"inputTokens": } """ - def _detect_input_type(self, request_data: dict[str, Any]) -> str: + def _detect_input_type(self, request_data: Mapping[str, JsonValue]) -> Literal["converse", "invokeModel"]: """ Detect whether to use 'converse' or 'invokeModel' input format. @@ -57,8 +102,8 @@ class BedrockCountTokensConfig(BaseAWSLLM): def transform_anthropic_to_bedrock_count_tokens( self, - request_data: dict[str, Any], - ) -> dict[str, Any]: + request_data: Mapping[str, JsonValue], + ) -> dict[str, JsonValue]: """ Transform request to Bedrock CountTokens format. Supports both Converse and InvokeModel input types. @@ -95,27 +140,16 @@ class BedrockCountTokensConfig(BaseAWSLLM): else: return self._transform_to_invoke_model_format(request_data) - def _transform_to_converse_format(self, request_data: dict[str, Any]) -> dict[str, Any]: + def _transform_to_converse_format(self, request_data: Mapping[str, JsonValue]) -> dict[str, JsonValue]: """Transform to Converse input format, including system and tools.""" - messages: Final = request_data.get("messages", []) + messages: Final = _json_list(request_data.get("messages")) system: Final = request_data.get("system") tools: Final = request_data.get("tools") # Transform messages - user_messages: Final = [] - for message in messages: - transformed_message: dict[str, Any] = { - "role": message.get("role"), - "content": [], - } - content = message.get("content", "") - if isinstance(content, str): - transformed_message["content"].append({"text": content}) - elif isinstance(content, list): - transformed_message["content"] = content - user_messages.append(transformed_message) + user_messages: Final[list[JsonValue]] = [_to_converse_message(message) for message in messages] - converse_input: Final[dict[str, Any]] = {"messages": user_messages} + converse_input: Final[dict[str, JsonValue]] = {"messages": user_messages} # Transform system prompt (string or list of blocks → Bedrock format) system_blocks: Final = self._transform_system(system) @@ -129,7 +163,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): return {"input": {"converse": converse_input}} - def _transform_system(self, system: Any | None) -> list[dict[str, Any]]: + def _transform_system(self, system: JsonValue) -> list[JsonValue]: """Transform Anthropic system prompt to Bedrock system blocks.""" if system is None: return [] @@ -140,36 +174,16 @@ class BedrockCountTokensConfig(BaseAWSLLM): return [{"text": block.get("text", "")} for block in system if isinstance(block, dict)] return [] - def _transform_tools(self, tools: list[dict[str, Any]] | None) -> dict[str, Any] | None: + def _transform_tools(self, tools: JsonValue) -> dict[str, JsonValue] | None: """Transform Anthropic tools to Bedrock toolConfig format.""" if not tools: return None - bedrock_tools: Final = [] - for tool in tools: - name = tool.get("name", "") - # Bedrock tool names must match [a-zA-Z][a-zA-Z0-9_]* and max 64 chars - name = re.sub(r"[^a-zA-Z0-9_]", "_", name) - if name and not name[0].isalpha(): - name = "t_" + name - name = name[:64] - - description = tool.get("description") or name - input_schema = tool.get("input_schema", {"type": "object", "properties": {}}) - - bedrock_tools.append( - { - "toolSpec": { - "name": name, - "description": description, - "inputSchema": {"json": input_schema}, - } - } - ) + bedrock_tools: Final[list[JsonValue]] = [_to_bedrock_tool_spec(tool) for tool in _json_list(tools)] return {"tools": bedrock_tools} - def _transform_to_invoke_model_format(self, request_data: dict[str, Any]) -> dict[str, Any]: + def _transform_to_invoke_model_format(self, request_data: Mapping[str, JsonValue]) -> dict[str, JsonValue]: """Transform to InvokeModel input format.""" import base64 import json @@ -223,7 +237,9 @@ class BedrockCountTokensConfig(BaseAWSLLM): return endpoint - def transform_bedrock_response_to_anthropic(self, bedrock_response: dict[str, Any]) -> dict[str, Any]: + def transform_bedrock_response_to_anthropic( + self, bedrock_response: Mapping[str, JsonValue] + ) -> dict[str, JsonValue]: """ Transform Bedrock CountTokens response to Anthropic format. @@ -241,7 +257,7 @@ class BedrockCountTokensConfig(BaseAWSLLM): return {"input_tokens": input_tokens} - def validate_count_tokens_request(self, request_data: dict[str, Any]) -> None: + def validate_count_tokens_request(self, request_data: Mapping[str, JsonValue]) -> None: """ Validate the incoming count tokens request. Supports both Converse and InvokeModel input formats. diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index 082bf7ee2d9..c34ca7750e2 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -6,7 +6,7 @@ import copy import json import urllib.parse from collections.abc import Callable -from typing import Any, Final, get_args +from typing import TYPE_CHECKING, Any, Final, get_args import httpx @@ -37,6 +37,9 @@ from .amazon_titan_v2_transformation import AmazonTitanV2Config from .cohere_transformation import BedrockCohereEmbeddingConfig from .twelvelabs_marengo_transformation import TwelveLabsMarengoEmbeddingConfig +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + class BedrockEmbedding(BaseAWSLLM): def _load_credentials( @@ -58,6 +61,7 @@ class BedrockEmbedding(BaseAWSLLM): aws_profile_name: Final = optional_params.pop("aws_profile_name", None) aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None) + aws_external_id: Final = optional_params.pop("aws_external_id", None) ### SET REGION NAME ### if aws_region_name is None: @@ -84,6 +88,7 @@ class BedrockEmbedding(BaseAWSLLM): aws_role_name=aws_role_name, aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, ) return credentials, aws_region_name @@ -233,7 +238,7 @@ class BedrockEmbedding(BaseAWSLLM): endpoint_url: str, aws_region_name: str, model: str, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL, api_key: str | None = None, is_async_invoke: bool | None = False, @@ -301,7 +306,7 @@ class BedrockEmbedding(BaseAWSLLM): endpoint_url: str, aws_region_name: str, model: str, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL, api_key: str | None = None, is_async_invoke: bool | None = False, diff --git a/litellm/llms/bedrock/files/transformation.py b/litellm/llms/bedrock/files/transformation.py index b034696594a..f442608a288 100644 --- a/litellm/llms/bedrock/files/transformation.py +++ b/litellm/llms/bedrock/files/transformation.py @@ -20,6 +20,7 @@ from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.constants import BEDROCK_INVOKE_PROVIDERS_LITERAL from litellm.files.utils import FilesAPIUtils +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.cloud_storage_security import ( BEDROCK_MANAGED_S3_BATCH_PREFIX, BEDROCK_MANAGED_S3_PREFIXES, @@ -413,7 +414,8 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): # S3 endpoint URL format s3_endpoint_url: Final = ( - request_params.get("s3_endpoint_url") or f"https://s3.{aws_region_name}.amazonaws.com" + request_params.get("s3_endpoint_url") + or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" ).rstrip("/") return f"{s3_endpoint_url}/{bucket_name}/{encoded_object_name}" @@ -1249,7 +1251,9 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig): region_params: Final[dict[str, str | None]] = {"aws_region_name": region_preference} aws_region_name: Final = self._get_aws_region_name(optional_params=region_params, model="") - s3_endpoint_url = (request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.amazonaws.com").rstrip("/") + s3_endpoint_url = ( + request_params.s3_endpoint_url or f"https://s3.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" + ).rstrip("/") url: Final = f"{s3_endpoint_url}/{bucket_name}/{encode_s3_object_key_for_url(object_key)}" litellm_params[S3_SIGNED_GET_HEADERS_PARAM] = self._sign_s3_get_request( diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index 3eeb3cb9fc6..96d7a79c6d8 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -7,18 +7,66 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic. import asyncio import contextlib import json -from typing import Any, Final +from typing import Final, Protocol -from pydantic import TypeAdapter +from pydantic import JsonValue, TypeAdapter from litellm._logging import _redact_string, verbose_proxy_logger +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.types.realtime import RealtimeResponseTransformInput from ..base_aws_llm import BaseAWSLLM from ..common_utils import BedrockError from .transformation import BedrockRealtimeConfig _CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) +_CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) + + +def _json_dict(value: JsonValue) -> dict[str, JsonValue]: + return value if isinstance(value, dict) else {} + + +def _json_str(value: JsonValue) -> str | None: + return value if isinstance(value, str) else None + + +class RealtimeClientWebSocket(Protocol): + """The client-facing websocket surface the realtime bridge talks to.""" + + async def receive_text(self) -> str: ... + + async def send_text(self, data: str) -> None: ... + + async def close(self, code: int = 1000, reason: str | None = None) -> None: ... + + +class BedrockInputStream(Protocol): + async def send(self, event: object) -> None: ... + + async def close(self) -> None: ... + + +class BedrockPayloadPart(Protocol): + @property + def bytes_(self) -> bytes | None: ... + + +class BedrockOutputChunk(Protocol): + @property + def value(self) -> BedrockPayloadPart | None: ... + + +class BedrockOutputStream(Protocol): + async def receive(self) -> BedrockOutputChunk | None: ... + + +class BedrockBidirectionalStream(Protocol): + @property + def input_stream(self) -> BedrockInputStream: ... + + async def await_output(self) -> tuple[object, BedrockOutputStream]: ... class BedrockRealtime(BaseAWSLLM): @@ -30,7 +78,7 @@ class BedrockRealtime(BaseAWSLLM): async def async_realtime( self, model: str, - websocket: Any, + websocket: RealtimeClientWebSocket, logging_obj: LiteLLMLogging, api_base: str | None = None, api_key: str | None = None, @@ -81,7 +129,7 @@ class BedrockRealtime(BaseAWSLLM): elif aws_bedrock_runtime_endpoint is not None: endpoint_uri = aws_bedrock_runtime_endpoint else: - endpoint_uri = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com" + endpoint_uri = f"https://bedrock-runtime.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" verbose_proxy_logger.debug("Bedrock Realtime: Connecting to %s with model %s", endpoint_uri, model) @@ -132,7 +180,7 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect") # Track state for transformation - session_state: Final = { + session_state: Final[RealtimeResponseTransformInput] = { "current_output_item_id": None, "current_response_id": None, "current_conversation_id": None, @@ -182,11 +230,11 @@ class BedrockRealtime(BaseAWSLLM): async def _forward_client_to_bedrock( self, - client_ws: Any, - bedrock_stream: Any, + client_ws: RealtimeClientWebSocket, + bedrock_stream: BedrockBidirectionalStream, transformation_config: BedrockRealtimeConfig, model: str, - session_state: dict, + session_state: RealtimeResponseTransformInput, logging_obj: LiteLLMLogging | None = None, ): """Forward messages from client WebSocket to Bedrock stream.""" @@ -223,11 +271,11 @@ class BedrockRealtime(BaseAWSLLM): client_message_type: str | None = None requested_modalities: list[str] | None = None with contextlib.suppress(Exception): - parsed_client_message = json.loads(message) - client_message_type = parsed_client_message.get("type") + parsed_client_message = _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message)) + client_message_type = _json_str(parsed_client_message.get("type")) if client_message_type == "session.update": requested_modalities = _CLIENT_MODALITIES_ADAPTER.validate_python( - parsed_client_message.get("session", {}).get("modalities") + _json_dict(parsed_client_message.get("session")).get("modalities") ) if client_message_type == "session.update": await client_ws.send_text( @@ -246,12 +294,12 @@ class BedrockRealtime(BaseAWSLLM): async def _forward_bedrock_to_client( self, - bedrock_stream: Any, - client_ws: Any, + bedrock_stream: BedrockBidirectionalStream, + client_ws: RealtimeClientWebSocket, transformation_config: BedrockRealtimeConfig, model: str, logging_obj: LiteLLMLogging, - session_state: dict, + session_state: RealtimeResponseTransformInput, ): """Forward messages from Bedrock stream to client WebSocket.""" try: @@ -264,13 +312,12 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended") break - if result.value and result.value.bytes_: - bedrock_response = result.value.bytes_.decode("utf-8") + payload_bytes = result.value.bytes_ if result.value else None + if payload_bytes: + bedrock_response = payload_bytes.decode("utf-8") verbose_proxy_logger.debug("Bedrock Realtime: Received from Bedrock: %s", bedrock_response[:200]) # Transform Bedrock format to OpenAI format - from litellm.types.realtime import RealtimeResponseTransformInput - realtime_response_transform_input: RealtimeResponseTransformInput = { "current_output_item_id": session_state.get("current_output_item_id"), "current_response_id": session_state.get("current_response_id"), diff --git a/litellm/llms/black_forest_labs/image_generation/transformation.py b/litellm/llms/black_forest_labs/image_generation/transformation.py index 5953ad1996b..119ffff1c34 100644 --- a/litellm/llms/black_forest_labs/image_generation/transformation.py +++ b/litellm/llms/black_forest_labs/image_generation/transformation.py @@ -29,6 +29,8 @@ from ..common_utils import ( ) if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -256,7 +258,7 @@ class BlackForestLabsImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/bytez/chat/transformation.py b/litellm/llms/bytez/chat/transformation.py index becd3f2d67e..d9a0c98b6db 100644 --- a/litellm/llms/bytez/chat/transformation.py +++ b/litellm/llms/bytez/chat/transformation.py @@ -23,6 +23,8 @@ from litellm.utils import CustomStreamWrapper, ModelResponse, Usage from ..common_utils import API_BASE, BytezError if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -185,7 +187,7 @@ class BytezChatConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/chatgpt/authenticator.py b/litellm/llms/chatgpt/authenticator.py index 6a3d278a74c..563826c2b93 100644 --- a/litellm/llms/chatgpt/authenticator.py +++ b/litellm/llms/chatgpt/authenticator.py @@ -2,9 +2,11 @@ import base64 import json import os import time -from typing import Any, Final +from collections.abc import Mapping +from typing import Final, TypeAlias import httpx +from pydantic import JsonValue, TypeAdapter, ValidationError from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import _get_httpx_client @@ -27,6 +29,16 @@ DEVICE_CODE_TIMEOUT_SECONDS: Final = 15 * 60 DEVICE_CODE_COOLDOWN_SECONDS: Final = 5 * 60 DEVICE_CODE_POLL_SLEEP_SECONDS: Final = 5 +OPENAI_AUTH_CLAIM_KEY: Final = "https://api.openai.com/auth" + +JsonObject: TypeAlias = Mapping[str, JsonValue] + +_JSON_OBJECT_ADAPTER: Final = TypeAdapter(JsonObject) + + +def _optional_str(value: JsonValue | None) -> str | None: + return value if isinstance(value, str) else None + class Authenticator: def __init__(self) -> None: @@ -43,10 +55,10 @@ class Authenticator: def get_access_token(self) -> str: auth_data: Final = self._read_auth_file() if auth_data: - access_token: Final = auth_data.get("access_token") + access_token: Final = _optional_str(auth_data.get("access_token")) if access_token and not self._is_token_expired(auth_data, access_token): return access_token - refresh_token: Final = auth_data.get("refresh_token") + refresh_token: Final = _optional_str(auth_data.get("refresh_token")) if refresh_token: try: refreshed: Final = self._refresh_tokens(refresh_token) @@ -67,48 +79,47 @@ class Authenticator: auth_data: Final = self._read_auth_file() if not auth_data: return None - account_id: Final = auth_data.get("account_id") + account_id: Final = _optional_str(auth_data.get("account_id")) if account_id: return account_id id_token: Final = auth_data.get("id_token") access_token: Final = auth_data.get("access_token") - derived: Final = self._extract_account_id(id_token or access_token) + derived: Final = self._extract_account_id(_optional_str(id_token or access_token)) if derived: - auth_data["account_id"] = derived - self._write_auth_file(auth_data) + self._write_auth_file({**auth_data, "account_id": derived}) return derived def _ensure_token_dir(self) -> None: if not os.path.exists(self.token_dir): os.makedirs(self.token_dir, exist_ok=True) - def _read_auth_file(self) -> dict[str, Any] | None: + def _read_auth_file(self) -> JsonObject | None: try: with open(self.auth_file, "r") as f: - return json.load(f) + return _JSON_OBJECT_ADAPTER.validate_python(json.load(f)) except OSError: return None - except json.JSONDecodeError as exc: + except (json.JSONDecodeError, ValidationError) as exc: verbose_logger.warning("Invalid ChatGPT auth file: %s", exc) return None - def _write_auth_file(self, data: dict[str, Any]) -> None: + def _write_auth_file(self, data: JsonObject) -> None: try: with open(self.auth_file, "w") as f: json.dump(data, f) except OSError as exc: verbose_logger.error("Failed to write ChatGPT auth file: %s", exc) - def _is_token_expired(self, auth_data: dict[str, Any], access_token: str) -> bool: - expires_at = auth_data.get("expires_at") - if expires_at is None: - expires_at = self._get_expires_at(access_token) - if expires_at: - auth_data["expires_at"] = expires_at - self._write_auth_file(auth_data) - if expires_at is None: + def _is_token_expired(self, auth_data: JsonObject, access_token: str) -> bool: + stored_expires_at: Final = auth_data.get("expires_at") + if isinstance(stored_expires_at, (int, float)): + return time.time() >= float(stored_expires_at) - TOKEN_EXPIRY_SKEW_SECONDS + derived_expires_at: Final = self._get_expires_at(access_token) + if derived_expires_at: + self._write_auth_file({**auth_data, "expires_at": derived_expires_at}) + if derived_expires_at is None: return True - return time.time() >= float(expires_at) - TOKEN_EXPIRY_SKEW_SECONDS + return time.time() >= float(derived_expires_at) - TOKEN_EXPIRY_SKEW_SECONDS def _get_expires_at(self, token: str) -> int | None: claims: Final = self._decode_jwt_claims(token) @@ -117,15 +128,14 @@ class Authenticator: return int(exp) return None - def _decode_jwt_claims(self, token: str) -> dict[str, Any]: + def _decode_jwt_claims(self, token: str) -> JsonObject: try: parts: Final = token.split(".") if len(parts) < 2: return {} - payload_b64 = parts[1] - payload_b64 += "=" * (-len(payload_b64) % 4) + payload_b64: Final = parts[1] + "=" * (-len(parts[1]) % 4) payload_bytes: Final = base64.urlsafe_b64decode(payload_b64) - return json.loads(payload_bytes.decode("utf-8")) + return _JSON_OBJECT_ADAPTER.validate_python(json.loads(payload_bytes.decode("utf-8"))) except Exception: return {} @@ -133,7 +143,7 @@ class Authenticator: if not token: return None claims: Final = self._decode_jwt_claims(token) - auth_claims: Final = claims.get("https://api.openai.com/auth") + auth_claims: Final = claims.get(OPENAI_AUTH_CLAIM_KEY) if isinstance(auth_claims, dict): account_id: Final = auth_claims.get("chatgpt_account_id") if isinstance(account_id, str) and account_id: @@ -170,7 +180,7 @@ class Authenticator: json={"client_id": CHATGPT_CLIENT_ID}, ) resp.raise_for_status() - data: Final = resp.json() + data: Final = _JSON_OBJECT_ADAPTER.validate_python(resp.json()) except httpx.HTTPStatusError as exc: raise GetDeviceCodeError( message=f"Failed to request device code: {exc}", @@ -182,8 +192,8 @@ class Authenticator: status_code=400, ) - device_auth_id: Final = data.get("device_auth_id") - user_code: Final = data.get("user_code") or data.get("usercode") + device_auth_id: Final = _optional_str(data.get("device_auth_id")) + user_code: Final = _optional_str(data.get("user_code") or data.get("usercode")) interval: Final = data.get("interval") if not device_auth_id or not user_code: raise GetDeviceCodeError( @@ -210,16 +220,16 @@ class Authenticator: }, ) if resp.status_code == 200: - data = resp.json() - if all( - key in data - for key in ( - "authorization_code", - "code_challenge", - "code_verifier", - ) - ): - return data + data = _JSON_OBJECT_ADAPTER.validate_python(resp.json()) + authorization_code = _optional_str(data.get("authorization_code")) + code_challenge = _optional_str(data.get("code_challenge")) + code_verifier = _optional_str(data.get("code_verifier")) + if authorization_code and code_challenge and code_verifier: + return { + "authorization_code": authorization_code, + "code_challenge": code_challenge, + "code_verifier": code_verifier, + } if resp.status_code in (403, 404): time.sleep(max(interval, DEVICE_CODE_POLL_SLEEP_SECONDS)) continue @@ -262,7 +272,7 @@ class Authenticator: content=body, ) resp.raise_for_status() - data: Final = resp.json() + data: Final = _JSON_OBJECT_ADAPTER.validate_python(resp.json()) except httpx.HTTPStatusError as exc: raise GetAccessTokenError( message=f"Token exchange failed: {exc}", @@ -274,15 +284,18 @@ class Authenticator: status_code=400, ) - if not all(key in data for key in ("access_token", "refresh_token", "id_token")): + access_token: Final = _optional_str(data.get("access_token")) + refresh_token: Final = _optional_str(data.get("refresh_token")) + id_token: Final = _optional_str(data.get("id_token")) + if not access_token or not refresh_token or not id_token: raise GetAccessTokenError( message=f"Token exchange response missing fields: {data}", status_code=400, ) return { - "access_token": data["access_token"], - "refresh_token": data["refresh_token"], - "id_token": data["id_token"], + "access_token": access_token, + "refresh_token": refresh_token, + "id_token": id_token, } def _refresh_tokens(self, refresh_token: str) -> dict[str, str]: @@ -298,7 +311,7 @@ class Authenticator: }, ) resp.raise_for_status() - data: Final = resp.json() + data: Final = _JSON_OBJECT_ADAPTER.validate_python(resp.json()) except httpx.HTTPStatusError as exc: raise RefreshAccessTokenError( message=f"Refresh token failed: {exc}", @@ -310,8 +323,8 @@ class Authenticator: status_code=400, ) - access_token: Final = data.get("access_token") - id_token: Final = data.get("id_token") + access_token: Final = _optional_str(data.get("access_token")) + id_token: Final = _optional_str(data.get("id_token")) if not access_token or not id_token: raise RefreshAccessTokenError( message=f"Refresh response missing fields: {data}", @@ -320,14 +333,14 @@ class Authenticator: refreshed: Final = { "access_token": access_token, - "refresh_token": data.get("refresh_token", refresh_token), + "refresh_token": _optional_str(data.get("refresh_token")) or refresh_token, "id_token": id_token, } auth_data: Final = self._build_auth_record(refreshed) self._write_auth_file(auth_data) return refreshed - def _build_auth_record(self, tokens: dict[str, str]) -> dict[str, Any]: + def _build_auth_record(self, tokens: dict[str, str]) -> JsonObject: access_token: Final = tokens.get("access_token") id_token: Final = tokens.get("id_token") expires_at: Final = self._get_expires_at(access_token) if access_token else None @@ -340,31 +353,30 @@ class Authenticator: "account_id": account_id, } - def _get_device_code_cooldown_remaining(self, auth_data: dict[str, Any] | None) -> float: + def _get_device_code_cooldown_remaining(self, auth_data: JsonObject | None) -> float: if not auth_data: return 0.0 - requested_at = auth_data.get("device_code_requested_at") + requested_at: Final = auth_data.get("device_code_requested_at") if not isinstance(requested_at, (int, float, str)): return 0.0 try: - requested_at = float(requested_at) + requested_seconds: Final = float(requested_at) except (TypeError, ValueError): return 0.0 - elapsed: Final = time.time() - requested_at + elapsed: Final = time.time() - requested_seconds remaining: Final = DEVICE_CODE_COOLDOWN_SECONDS - elapsed return max(0.0, remaining) def _record_device_code_request(self) -> None: auth_data: Final = self._read_auth_file() or {} - auth_data["device_code_requested_at"] = time.time() - self._write_auth_file(auth_data) + self._write_auth_file({**auth_data, "device_code_requested_at": time.time()}) def _wait_for_access_token(self, timeout_seconds: float) -> str | None: deadline: Final = time.time() + timeout_seconds while time.time() < deadline: auth_data = self._read_auth_file() if auth_data: - access_token = auth_data.get("access_token") + access_token = _optional_str(auth_data.get("access_token")) if access_token and not self._is_token_expired(auth_data, access_token): return access_token sleep_for = min(DEVICE_CODE_POLL_SLEEP_SECONDS, max(0.0, deadline - time.time())) diff --git a/litellm/llms/chatgpt/chat/streaming_utils.py b/litellm/llms/chatgpt/chat/streaming_utils.py index 57f679947f6..d4a168a6984 100644 --- a/litellm/llms/chatgpt/chat/streaming_utils.py +++ b/litellm/llms/chatgpt/chat/streaming_utils.py @@ -4,7 +4,27 @@ Streaming utilities for ChatGPT provider. Normalizes non-spec-compliant tool_call chunks from the ChatGPT backend API. """ -from typing import Any, Final +from collections.abc import Awaitable +from typing import Final, Protocol + +from litellm.types.utils import ( + ChatCompletionDeltaCustomToolCall, + ChatCompletionDeltaToolCall, + Delta, + ModelResponseStream, +) + + +class ChatGPTChunkStream(Protocol): + """A ChatGPT chunk source driven either synchronously or asynchronously.""" + + def __next__(self) -> ModelResponseStream: ... + + def __anext__(self) -> Awaitable[ModelResponseStream]: ... + + +def _first_choice_delta(chunk: ModelResponseStream) -> Delta | None: + return chunk.choices[0].delta class ChatGPTToolCallNormalizer: @@ -20,13 +40,13 @@ class ChatGPTToolCallNormalizer: chunks to the consumer. """ - def __init__(self, stream: Any): - self._stream = stream + def __init__(self, stream: ChatGPTChunkStream): + self._stream: Final = stream self._seen_ids: dict[str, int] = {} # tool_call_id -> assigned_index self._next_index: int = 0 self._last_id: str | None = None # tracks which tool call the next delta belongs to - def __getattr__(self, name: str) -> Any: + def __getattr__(self, name: str) -> object: return getattr(self._stream, name) def __iter__(self): @@ -35,30 +55,30 @@ class ChatGPTToolCallNormalizer: def __aiter__(self): return self - def __next__(self): + def __next__(self) -> ModelResponseStream: while True: chunk = next(self._stream) result = self._normalize(chunk) if result is not None: return result - async def __anext__(self): + async def __anext__(self) -> ModelResponseStream: while True: chunk = await self._stream.__anext__() result = self._normalize(chunk) if result is not None: return result - def _normalize(self, chunk: Any) -> Any: + def _normalize(self, chunk: ModelResponseStream) -> ModelResponseStream | None: """Fix tool_calls in the chunk. Returns None to skip duplicate chunks.""" if not chunk.choices: return chunk - delta: Final = chunk.choices[0].delta + delta: Final = _first_choice_delta(chunk) if delta is None or not delta.tool_calls: return chunk - normalized: Final = [] + normalized: Final[list[ChatCompletionDeltaToolCall | ChatCompletionDeltaCustomToolCall]] = [] for tc in delta.tool_calls: if tc.id and tc.id not in self._seen_ids: # New tool call — assign correct index diff --git a/litellm/llms/chatgpt/responses/transformation.py b/litellm/llms/chatgpt/responses/transformation.py index 8e4bbf1d3c9..b96e06be3d8 100644 --- a/litellm/llms/chatgpt/responses/transformation.py +++ b/litellm/llms/chatgpt/responses/transformation.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.core_helpers import process_response_headers @@ -28,6 +28,9 @@ from ..common_utils import ( get_chatgpt_default_instructions, ) +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): def __init__(self) -> None: @@ -107,7 +110,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig): self, model: str, raw_response: Any, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", ): body_text: Final = raw_response.text or "" if not self._should_parse_as_sse(raw_response=raw_response, body_text=body_text): diff --git a/litellm/llms/clarifai/chat/transformation.py b/litellm/llms/clarifai/chat/transformation.py index f5227966aef..76d35467497 100644 --- a/litellm/llms/clarifai/chat/transformation.py +++ b/litellm/llms/clarifai/chat/transformation.py @@ -13,6 +13,8 @@ from litellm.types.utils import ModelResponse from ...openai.chat.gpt_transformation import OpenAIGPTConfig if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -85,7 +87,7 @@ class ClarifaiConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/cohere/chat/transformation.py b/litellm/llms/cohere/chat/transformation.py index 3560683c49b..319603b0dad 100644 --- a/litellm/llms/cohere/chat/transformation.py +++ b/litellm/llms/cohere/chat/transformation.py @@ -15,6 +15,8 @@ from ..common_utils import ModelResponseIterator as CohereModelResponseIterator from ..common_utils import validate_environment as cohere_validate_environment if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -225,7 +227,7 @@ class CohereChatConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/cohere/chat/v2_transformation.py b/litellm/llms/cohere/chat/v2_transformation.py index a7db03924b6..4252e7d02e9 100644 --- a/litellm/llms/cohere/chat/v2_transformation.py +++ b/litellm/llms/cohere/chat/v2_transformation.py @@ -20,6 +20,8 @@ from ..common_utils import CohereError, CohereV2ModelResponseIterator from ..common_utils import validate_environment as cohere_validate_environment if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -189,7 +191,7 @@ class CohereV2ChatConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/cohere/embed/handler.py b/litellm/llms/cohere/embed/handler.py index 3384839da85..3cebf6b9a90 100644 --- a/litellm/llms/cohere/embed/handler.py +++ b/litellm/llms/cohere/embed/handler.py @@ -3,8 +3,7 @@ Legacy /v1/embedding handler for Bedrock Cohere. """ import json -from collections.abc import Callable -from typing import Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -20,6 +19,9 @@ from litellm.types.utils import EmbeddingResponse from .v1_transformation import CohereEmbeddingConfig +if TYPE_CHECKING: + import tiktoken + def validate_environment(api_key, headers: dict): # Create a lowercase key lookup to avoid duplicate headers with different cases @@ -58,7 +60,7 @@ async def async_embedding( api_base: str, api_key: str | None, headers: dict, - encoding: Callable, + encoding: "tiktoken.Encoding | None", client: AsyncHTTPHandler | None = None, ): ## LOGGING @@ -120,7 +122,7 @@ def embedding( logging_obj: LiteLLMLoggingObj, optional_params: dict, headers: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", data: dict | CohereEmbeddingRequest | None = None, complete_api_base: str | None = None, api_key: str | None = None, diff --git a/litellm/llms/cohere/rerank/guardrail_translation/handler.py b/litellm/llms/cohere/rerank/guardrail_translation/handler.py index b5e49bd922e..84cb551190a 100644 --- a/litellm/llms/cohere/rerank/guardrail_translation/handler.py +++ b/litellm/llms/cohere/rerank/guardrail_translation/handler.py @@ -13,6 +13,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.rerank import RerankResponse @@ -42,7 +43,7 @@ class CohereRerankHandler(BaseTranslation): self, data: dict, guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> Any: """ Process input text fields ('query' and 'instruction') by applying @@ -94,7 +95,7 @@ class CohereRerankHandler(BaseTranslation): self, response: "RerankResponse", guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: Any | None = None, request_data: dict | None = None, ) -> Any: diff --git a/litellm/llms/cometapi/image_generation/transformation.py b/litellm/llms/cometapi/image_generation/transformation.py index 3c643f5ce36..03c820de198 100644 --- a/litellm/llms/cometapi/image_generation/transformation.py +++ b/litellm/llms/cometapi/image_generation/transformation.py @@ -13,6 +13,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -130,7 +132,7 @@ class CometAPIImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/compactifai/chat/transformation.py b/litellm/llms/compactifai/chat/transformation.py index 44e1ab15801..63a5427d211 100644 --- a/litellm/llms/compactifai/chat/transformation.py +++ b/litellm/llms/compactifai/chat/transformation.py @@ -14,6 +14,8 @@ from litellm.types.utils import ModelResponse from ...openai.chat.gpt_transformation import OpenAIGPTConfig if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -49,7 +51,7 @@ class CompactifAIChatConfig(OpenAIGPTConfig): messages: list, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index e9140e63cb3..7035ce58ae1 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -26,6 +26,8 @@ from litellm.types.utils import HttpHandlerRequestFields, ImageResponse, LlmProv from litellm.utils import CustomStreamWrapper, ModelResponse, ProviderConfigManager if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -266,7 +268,7 @@ class BaseLLMAIOHTTPHandler: messages: list, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, client: ClientSession | None = None, ): diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index 344a53d87f6..b6586481fd3 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -5,20 +5,38 @@ import os import ssl import typing import urllib.request -from collections.abc import Callable -from typing import Any, ClassVar, Final +from collections.abc import Callable, Generator +from typing import ClassVar, Final import aiohttp import aiohttp.client_exceptions import aiohttp.http_exceptions import httpx from aiohttp.client import ClientResponse, ClientSession +from pydantic import BaseModel, TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.secret_managers.main import str_to_bool -AIOHTTP_EXC_MAP: Final[dict] = { + +class HttpxTimeoutExtension(BaseModel): + connect: float | None = None + read: float | None = None + write: float | None = None + pool: float | None = None + + +class AiohttpSslRequestOption(TypedDict, total=False): + ssl: ReadOnly[bool | ssl.SSLContext] + + +_TIMEOUT_EXTENSION: Final = TypeAdapter(HttpxTimeoutExtension) +_EMPTY_TIMEOUT: Final[HttpxTimeoutExtension] = HttpxTimeoutExtension() +_NO_SSL_OVERRIDE: Final[AiohttpSslRequestOption] = {} + +AIOHTTP_EXC_MAP: Final[dict[type[BaseException], type[Exception]]] = { # Order matters here, most specific exception first # Timeout related exceptions asyncio.TimeoutError: httpx.TimeoutException, @@ -58,11 +76,11 @@ except ImportError: @contextlib.contextmanager -def map_aiohttp_exceptions() -> typing.Iterator[None]: +def map_aiohttp_exceptions() -> Generator[None, None, None]: try: yield except Exception as exc: - mapped_exc = None + mapped_exc: type[Exception] | None = None for from_exc, to_exc in AIOHTTP_EXC_MAP.items(): if not isinstance(exc, from_exc): @@ -222,7 +240,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): if session.closed: return - session_loop: Final = getattr(session, "_loop", None) + session_loop: Final[asyncio.AbstractEventLoop | None] = getattr(session, "_loop", None) try: current_loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop() except RuntimeError: @@ -278,7 +296,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): # Check if the existing session is still valid for the current event loop try: - session_loop: Final = getattr(self.client, "_loop", None) + session_loop: Final[asyncio.AbstractEventLoop | None] = getattr(self.client, "_loop", None) current_loop: Final = asyncio.get_running_loop() # If session is from a different or closed loop, recreate it @@ -312,7 +330,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): self, client_session: ClientSession, request: httpx.Request, - timeout: dict, + timeout: HttpxTimeoutExtension, proxy: str | None, sni_hostname: str | None, ssl_verify: bool | ssl.SSLContext | None = None, @@ -323,7 +341,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): Args: client_session: The aiohttp ClientSession to use request: The httpx Request to send - timeout: Timeout settings dict with 'connect', 'read', 'pool' keys + timeout: Timeout settings with 'connect', 'read', 'pool' fields proxy: Optional proxy URL sni_hostname: Optional SNI hostname for SSL ssl_verify: Optional SSL verification setting (False to disable, SSLContext for custom) @@ -346,25 +364,24 @@ class LiteLLMAiohttpTransport(AiohttpTransport): # Only pass ssl kwarg when explicitly configured, to avoid # overriding the session/connector defaults with None (which is # not a valid value for aiohttp's ssl parameter). - request_kwargs: Final[dict[str, Any]] = { - "method": request.method, - "url": YarlURL(str(request.url), encoded=True), - "headers": request.headers, - "data": data, - "allow_redirects": False, - "auto_decompress": False, - "timeout": ClientTimeout( - sock_connect=timeout.get("connect"), - sock_read=timeout.get("read"), - connect=timeout.get("pool"), - ), - "proxy": proxy, - "server_hostname": sni_hostname, - } - if ssl_verify is not None: - request_kwargs["ssl"] = ssl_verify + ssl_option: Final[AiohttpSslRequestOption] = _NO_SSL_OVERRIDE if ssl_verify is None else {"ssl": ssl_verify} - response: Final = await client_session.request(**request_kwargs).__aenter__() + response: Final = await client_session.request( + method=request.method, + url=YarlURL(str(request.url), encoded=True), + headers=request.headers, + data=data, + allow_redirects=False, + auto_decompress=False, + timeout=ClientTimeout( + sock_connect=timeout.connect, + sock_read=timeout.read, + connect=timeout.pool, + ), + proxy=proxy, + server_hostname=sni_hostname, + **ssl_option, + ).__aenter__() return response @@ -372,8 +389,8 @@ class LiteLLMAiohttpTransport(AiohttpTransport): self, request: httpx.Request, ) -> httpx.Response: - timeout: Final = request.extensions.get("timeout", {}) - sni_hostname: Final = request.extensions.get("sni_hostname") + timeout: Final = _TIMEOUT_EXTENSION.validate_python(request.extensions.get("timeout", _EMPTY_TIMEOUT)) + sni_hostname: Final[str | None] = request.extensions.get("sni_hostname") # Use helper to ensure we have a valid session for the current event loop client_session = self._get_valid_client_session() diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 417ecd80be2..573ba85416f 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -165,6 +165,7 @@ def _rust_responses_websocket_enabled( from .http_handler import get_shared_realtime_ssl_context if TYPE_CHECKING: + import tiktoken from aiohttp import ClientSession from websockets.asyncio.client import ClientConnection @@ -405,7 +406,7 @@ class BaseLLMHTTPHandler: messages: list, optional_params: dict, litellm_params: dict, - encoding: object, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, client: AsyncHTTPHandler | None = None, json_mode: bool = False, @@ -471,7 +472,7 @@ class BaseLLMHTTPHandler: api_base: str | None, custom_llm_provider: str, model_response: ModelResponse, - encoding: object, + encoding: "tiktoken.Encoding | None", logging_obj: LiteLLMLoggingObj, optional_params: dict, timeout: float | httpx.Timeout, diff --git a/litellm/llms/custom_llm.py b/litellm/llms/custom_llm.py index fcd41d11499..c70b9b81b42 100644 --- a/litellm/llms/custom_llm.py +++ b/litellm/llms/custom_llm.py @@ -25,6 +25,7 @@ from .base import BaseLLM if TYPE_CHECKING: from litellm import CustomStreamWrapper + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj class CustomLLMError(Exception): # use this for all your exceptions @@ -134,7 +135,7 @@ class CustomLLM(BaseLLM): api_base: str | None, model_response: ImageResponse, optional_params: dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, ) -> ImageResponse: @@ -148,7 +149,7 @@ class CustomLLM(BaseLLM): api_key: str | None, # dynamically set api_key - https://docs.litellm.ai/docs/set_keys#api_key api_base: str | None, # dynamically set api_base - https://docs.litellm.ai/docs/set_keys#api_base optional_params: dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> ImageResponse: @@ -160,7 +161,7 @@ class CustomLLM(BaseLLM): input: list, model_response: EmbeddingResponse, print_verbose: Callable, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", optional_params: dict, api_key: str | None = None, api_base: str | None = None, @@ -175,7 +176,7 @@ class CustomLLM(BaseLLM): input: list, model_response: EmbeddingResponse, print_verbose: Callable, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", optional_params: dict, api_key: str | None = None, api_base: str | None = None, @@ -193,7 +194,7 @@ class CustomLLM(BaseLLM): api_key: str | None, api_base: str | None, optional_params: dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", timeout: float | httpx.Timeout | None = None, client: HTTPHandler | None = None, ) -> ImageResponse: @@ -208,7 +209,7 @@ class CustomLLM(BaseLLM): api_key: str | None, api_base: str | None, optional_params: dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", timeout: float | httpx.Timeout | None = None, client: AsyncHTTPHandler | None = None, ) -> ImageResponse: diff --git a/litellm/llms/dashscope/image_generation/transformation.py b/litellm/llms/dashscope/image_generation/transformation.py index e655e2ea87d..a7f0e98865f 100644 --- a/litellm/llms/dashscope/image_generation/transformation.py +++ b/litellm/llms/dashscope/image_generation/transformation.py @@ -38,6 +38,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -157,7 +159,7 @@ class DashScopeImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/databricks/chat/transformation.py b/litellm/llms/databricks/chat/transformation.py index 8a625569cfa..c587146005f 100644 --- a/litellm/llms/databricks/chat/transformation.py +++ b/litellm/llms/databricks/chat/transformation.py @@ -136,6 +136,8 @@ def _split_parallel_tool_calls(messages: list[AllMessageValues]) -> list[AllMess if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -603,7 +605,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/databricks/common_utils.py b/litellm/llms/databricks/common_utils.py index b8dc98f2582..7695b1cb35e 100644 --- a/litellm/llms/databricks/common_utils.py +++ b/litellm/llms/databricks/common_utils.py @@ -13,6 +13,7 @@ Authentication priority: import os import re from typing import Any, Final, Literal +from urllib.parse import urlsplit, urlunsplit from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -224,11 +225,8 @@ class DatabricksBase: """ import requests - # Extract workspace URL from api_base - workspace_url = api_base.rstrip("/") - if "/serving-endpoints" in workspace_url: - workspace_url = workspace_url.replace("/serving-endpoints", "") - + api_base_parts: Final = urlsplit(api_base) + workspace_url: Final = urlunsplit((api_base_parts.scheme, api_base_parts.netloc, "", "", "")) token_url: Final = f"{workspace_url}/oidc/v1/token" try: diff --git a/litellm/llms/fal_ai/image_generation/bria_transformation.py b/litellm/llms/fal_ai/image_generation/bria_transformation.py index 5cfe6a67523..c528550811a 100644 --- a/litellm/llms/fal_ai/image_generation/bria_transformation.py +++ b/litellm/llms/fal_ai/image_generation/bria_transformation.py @@ -8,6 +8,8 @@ from litellm.types.utils import ImageObject, ImageResponse from .transformation import FalAIBaseConfig if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -185,7 +187,7 @@ class FalAIBriaConfig(FalAIBaseConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py index 6e962978a43..228dd9257ce 100644 --- a/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py +++ b/litellm/llms/fal_ai/image_generation/flux_pro_v11_ultra_transformation.py @@ -8,6 +8,8 @@ from litellm.types.utils import ImageObject, ImageResponse from .transformation import FalAIBaseConfig if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -192,7 +194,7 @@ class FalAIFluxProV11UltraConfig(FalAIBaseConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py b/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py index 2c6716f1365..04b4f426878 100644 --- a/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py +++ b/litellm/llms/fal_ai/image_generation/ideogram_v3_transformation.py @@ -8,6 +8,8 @@ from litellm.types.utils import ImageObject, ImageResponse from .transformation import FalAIBaseConfig if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -148,7 +150,7 @@ class FalAIIdeogramV3Config(FalAIBaseConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fal_ai/image_generation/imagen4_transformation.py b/litellm/llms/fal_ai/image_generation/imagen4_transformation.py index 28332a1f867..8a6665b2585 100644 --- a/litellm/llms/fal_ai/image_generation/imagen4_transformation.py +++ b/litellm/llms/fal_ai/image_generation/imagen4_transformation.py @@ -8,6 +8,8 @@ from litellm.types.utils import ImageObject, ImageResponse from .transformation import FalAIBaseConfig if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -180,7 +182,7 @@ class FalAIImagen4Config(FalAIBaseConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py b/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py index a5f0c086379..4880dfec7e3 100644 --- a/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py +++ b/litellm/llms/fal_ai/image_generation/recraft_v3_transformation.py @@ -8,6 +8,8 @@ from litellm.types.utils import ImageObject, ImageResponse from .transformation import FalAIBaseConfig if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -170,7 +172,7 @@ class FalAIRecraftV3Config(FalAIBaseConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py b/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py index 500aa859fe8..bc3a4d07282 100644 --- a/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py +++ b/litellm/llms/fal_ai/image_generation/stable_diffusion_transformation.py @@ -8,6 +8,8 @@ from litellm.types.utils import ImageObject, ImageResponse from .transformation import FalAIBaseConfig if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -206,7 +208,7 @@ class FalAIStableDiffusionConfig(FalAIBaseConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fal_ai/image_generation/transformation.py b/litellm/llms/fal_ai/image_generation/transformation.py index b65f9585730..7a114677b2d 100644 --- a/litellm/llms/fal_ai/image_generation/transformation.py +++ b/litellm/llms/fal_ai/image_generation/transformation.py @@ -13,6 +13,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -76,7 +78,7 @@ class FalAIBaseConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index 4e9731ef485..b6a5ee40672 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -1,6 +1,6 @@ import json from collections.abc import AsyncIterator, Iterator, Mapping -from typing import Any, Final, Literal, cast +from typing import TYPE_CHECKING, Any, Final, Literal, cast import httpx @@ -45,6 +45,9 @@ from ..common_utils import ( resolve_fireworks_resource_name, ) +if TYPE_CHECKING: + import tiktoken + def _extract_fireworks_hidden_params(payload: dict) -> dict: """ @@ -691,7 +694,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/gemini/image_edit/transformation.py b/litellm/llms/gemini/image_edit/transformation.py index 67b1f97a3a2..e6c22dc60b4 100644 --- a/litellm/llms/gemini/image_edit/transformation.py +++ b/litellm/llms/gemini/image_edit/transformation.py @@ -120,7 +120,7 @@ class GeminiImageEditConfig(BaseImageEditConfig): self, model: str, raw_response: httpx.Response, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj, ) -> ImageResponse: model_response: Final = ImageResponse() try: diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 3943c0a7dae..d009fe4cd72 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -24,6 +24,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -171,7 +173,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/gigachat/chat/transformation.py b/litellm/llms/gigachat/chat/transformation.py index 6d75c311084..b859a843251 100644 --- a/litellm/llms/gigachat/chat/transformation.py +++ b/litellm/llms/gigachat/chat/transformation.py @@ -22,6 +22,8 @@ from ..authenticator import get_access_token from ..file_handler import upload_file_sync if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -391,7 +393,7 @@ class GigaChatConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/github_copilot/chat/transformation.py b/litellm/llms/github_copilot/chat/transformation.py index 27a0028ce4a..8634b374f1b 100644 --- a/litellm/llms/github_copilot/chat/transformation.py +++ b/litellm/llms/github_copilot/chat/transformation.py @@ -1,6 +1,6 @@ import json import os -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final import httpx @@ -17,6 +17,9 @@ from ..common_utils import ( get_copilot_default_headers, ) +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + class GithubCopilotConfig(OpenAIConfig): def __init__( @@ -272,7 +275,7 @@ class GithubCopilotConfig(OpenAIConfig): model: str, raw_response: httpx.Response, model_response: "ModelResponse", - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", request_data: dict, messages: list[AllMessageValues], optional_params: dict, diff --git a/litellm/llms/groq/chat/transformation.py b/litellm/llms/groq/chat/transformation.py index c5e6bc13153..41a2df17c6f 100644 --- a/litellm/llms/groq/chat/transformation.py +++ b/litellm/llms/groq/chat/transformation.py @@ -3,7 +3,7 @@ Translate from OpenAI's `/v1/chat/completions` to Groq's `/v1/chat/completions` """ from collections.abc import AsyncIterator, Coroutine, Iterator -from typing import Any, Final, Literal, cast, overload +from typing import TYPE_CHECKING, Any, Final, Literal, cast, overload import httpx from pydantic import BaseModel, TypeAdapter, ValidationError @@ -26,6 +26,9 @@ from litellm.types.utils import ModelResponse, ModelResponseStream, ServerToolUs from ...openai_like.chat.transformation import OpenAILikeChatConfig +if TYPE_CHECKING: + import tiktoken + GROQ_COMPOUND_MODELS: Final = frozenset({"compound", "compound-mini"}) @@ -283,7 +286,7 @@ class GroqChatConfig(OpenAILikeChatConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/hosted_vllm/videos/__init__.py b/litellm/llms/hosted_vllm/videos/__init__.py new file mode 100644 index 00000000000..89aa5ef2e8b --- /dev/null +++ b/litellm/llms/hosted_vllm/videos/__init__.py @@ -0,0 +1,9 @@ +from litellm.llms.base_llm.videos.transformation import BaseVideoConfig + +from .transformation import HostedVLLMVideoConfig + +__all__ = ("HostedVLLMVideoConfig",) + + +def get_hosted_vllm_video_config(model: str | None) -> BaseVideoConfig: + return HostedVLLMVideoConfig() diff --git a/litellm/llms/hosted_vllm/videos/transformation.py b/litellm/llms/hosted_vllm/videos/transformation.py new file mode 100644 index 00000000000..96cbfc3cf70 --- /dev/null +++ b/litellm/llms/hosted_vllm/videos/transformation.py @@ -0,0 +1,206 @@ +"""Video generation for Hosted VLLM (vLLM-Omni OpenAI-compatible /v1/videos).""" + +import json +from collections.abc import Mapping +from io import BufferedReader +from types import MappingProxyType +from typing import Final +from urllib.parse import urlparse + +from httpx._types import FileTypes, RequestFiles + +from litellm.images.utils import ImageEditRequestUtils +from litellm.litellm_core_utils.url_utils import SSRFError, validate_url +from litellm.llms.openai.videos.transformation import OpenAIVideoConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams +from litellm.types.videos.main import VideoCreateOptionalRequestParams + +_EXCLUDED_FORM_KEYS: Final = frozenset( + { + "model", + "prompt", + "extra_headers", + "extra_query", + "extra_body", + "timeout", + "custom_llm_provider", + "input_reference", + "characters", + } +) + +_VLLM_OMNI_VIDEO_PARAMS: Final = ( + "image_reference", + "video_reference", + "audio_reference", + "width", + "height", + "num_frames", + "fps", + "num_inference_steps", + "guidance_scale", + "guidance_scale_2", + "boundary_ratio", + "flow_shift", + "true_cfg_scale", + "seed", + "generate_sound", + "sound_duration", + "negative_prompt", + "enable_frame_interpolation", + "frame_interpolation_exp", + "frame_interpolation_scale", + "frame_interpolation_model_path", + "lora", + "extra_params", + "aspect_ratio", +) + +_REFERENCE_URL_KEYS: Final = MappingProxyType( + { + "image_reference": "image_url", + "video_reference": "video_url", + "audio_reference": "audio_url", + } +) + + +def _serialize_form_value(value: object) -> str: + if isinstance(value, str): + return value + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (Mapping, list, tuple)): + return json.dumps(value) + return str(value) + + +def _maybe_json(value: object) -> object: + if not isinstance(value, str): + return value + stripped: Final = value.strip() + if not stripped or stripped[0] not in "{[": + return value + return json.loads(stripped) + + +def _reject_unsafe_media_url(url: str) -> None: + scheme: Final = urlparse(url).scheme.lower() + if scheme in ("", "data"): + return + if scheme not in ("http", "https"): + raise SSRFError(f"URL scheme '{scheme}' is not allowed") + validate_url(url) + + +def _reject_unsafe_urls_in_item(url_key: str, item: object) -> None: + if not isinstance(item, Mapping): + return + url: Final = item.get(url_key) + if isinstance(url, str): + _reject_unsafe_media_url(url) + + +def _reject_unsafe_media_urls(field_name: str, value: object) -> None: + url_key: Final = _REFERENCE_URL_KEYS.get(field_name) + if url_key is None: + return + parsed: Final = _maybe_json(value) + if isinstance(parsed, list): + for item in parsed: + _reject_unsafe_urls_in_item(url_key, item) + return + if isinstance(parsed, Mapping): + _reject_unsafe_urls_in_item(url_key, parsed) + + +def _form_value(key: str, value: object) -> str: + _reject_unsafe_media_urls(key, value) + return _serialize_form_value(value) + + +def _input_reference_file(reference: object) -> tuple[str, FileTypes]: + content_type: Final = ImageEditRequestUtils.get_image_content_type(reference) + if isinstance(reference, BufferedReader): + return ("input_reference", (reference.name, reference, content_type)) + return ("input_reference", ("input_reference.png", reference, content_type)) + + +class HostedVLLMVideoConfig(OpenAIVideoConfig): + """ + vLLM-Omni videos API is OpenAI-compatible but requires multipart/form-data. + + https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/videos_api/ + """ + + def get_supported_openai_params(self, model: str) -> list: # mutable-ok: BaseVideoConfig contract + return [ # mutable-ok: BaseVideoConfig returns list + *super().get_supported_openai_params(model), + *_VLLM_OMNI_VIDEO_PARAMS, + ] + + def map_openai_params( + self, + video_create_optional_params: VideoCreateOptionalRequestParams, + model: str, + drop_params: bool, + ) -> dict: # mutable-ok: BaseVideoConfig contract; extra_body merge mutates this dict + return { # mutable-ok: VideoGenerationRequestUtils.update/pop extra_body onto this mapping + key: value for key, value in video_create_optional_params.items() if value is not None + } + + def validate_environment( + self, + headers: dict, # mutable-ok: BaseVideoConfig contract + model: str, + api_key: str | None = None, + litellm_params: GenericLiteLLMParams | None = None, + ) -> dict: # mutable-ok: BaseVideoConfig contract + resolved_key: Final = ( + (litellm_params.api_key if litellm_params is not None else None) + or api_key + or get_secret_str("HOSTED_VLLM_API_KEY") + or "fake-api-key" + ) + return {**headers, "Authorization": f"Bearer {resolved_key}"} # mutable-ok: httpx headers are a dict + + def get_complete_url( + self, + model: str, + api_base: str | None, + litellm_params: dict, # mutable-ok: BaseVideoConfig contract + ) -> str: + resolved_api_base: Final = api_base or get_secret_str("HOSTED_VLLM_API_BASE") + if resolved_api_base is None: + raise ValueError( + "api_base not set for Hosted VLLM videos API. " + "Set via api_base parameter or HOSTED_VLLM_API_BASE environment variable" + ) + trimmed: Final = resolved_api_base.rstrip("/") + if trimmed.endswith("/v1"): + return f"{trimmed}/videos" + return f"{trimmed}/v1/videos" + + def transform_video_create_request( + self, + model: str, + prompt: str, + api_base: str, + video_create_optional_request_params: dict, # mutable-ok: BaseVideoConfig contract + litellm_params: GenericLiteLLMParams, + headers: dict, # mutable-ok: BaseVideoConfig contract + ) -> tuple[dict, RequestFiles, str]: # mutable-ok: BaseVideoConfig contract + data: Final = { # mutable-ok: BaseVideoConfig contract returns a data dict + "model": model, + "prompt": prompt, + **{ # mutable-ok: spread remaining Omni form fields into that data dict + key: _form_value(key, value) + for key, value in video_create_optional_request_params.items() + if key not in _EXCLUDED_FORM_KEYS and value is not None + }, + } + input_reference: Final = video_create_optional_request_params.get("input_reference") + if input_reference is None: + return data, (), api_base + return data, (_input_reference_file(input_reference),), api_base diff --git a/litellm/llms/langflow/chat/transformation.py b/litellm/llms/langflow/chat/transformation.py index 6b837007f21..17ae7017cf6 100644 --- a/litellm/llms/langflow/chat/transformation.py +++ b/litellm/llms/langflow/chat/transformation.py @@ -14,6 +14,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.utils import CustomStreamWrapper @@ -223,7 +225,7 @@ class LangFlowConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/langgraph/chat/transformation.py b/litellm/llms/langgraph/chat/transformation.py index c72246114b8..84d79e6bd31 100644 --- a/litellm/llms/langgraph/chat/transformation.py +++ b/litellm/llms/langgraph/chat/transformation.py @@ -23,6 +23,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.utils import CustomStreamWrapper @@ -413,7 +415,7 @@ class LangGraphConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/lemonade/chat/transformation.py b/litellm/llms/lemonade/chat/transformation.py index 4ea96df0ac4..553478aec16 100644 --- a/litellm/llms/lemonade/chat/transformation.py +++ b/litellm/llms/lemonade/chat/transformation.py @@ -2,7 +2,7 @@ Translate from OpenAI's `/v1/chat/completions` to Lemonade's `/v1/chat/completions` """ -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final from urllib.parse import quote import httpx @@ -18,6 +18,9 @@ from litellm.types.utils import ModelResponse from ...openai_like.chat.transformation import OpenAILikeChatConfig +if TYPE_CHECKING: + import tiktoken + class LemonadeChatConfig(OpenAILikeChatConfig): _DEFAULT_API_KEY = "lemonade" @@ -228,7 +231,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/litellm_proxy/skills/transformation.py b/litellm/llms/litellm_proxy/skills/transformation.py index 33c26801617..c972dc349c9 100644 --- a/litellm/llms/litellm_proxy/skills/transformation.py +++ b/litellm/llms/litellm_proxy/skills/transformation.py @@ -7,8 +7,10 @@ API requests to database operations via LiteLLMSkillsHandler. Pattern follows litellm/llms/litellm_proxy/responses/transformation.py """ -from collections.abc import Coroutine -from typing import TYPE_CHECKING, Any, Final, Optional +from collections.abc import Coroutine, Sequence +from typing import TYPE_CHECKING, Final, Optional + +from pydantic import JsonValue from litellm.types.llms.anthropic_skills import ( DeleteSkillResponse, @@ -19,7 +21,7 @@ from litellm.types.utils import LlmProviders if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj - from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy._types import LiteLLM_SkillsTable, UserAPIKeyAuth class LiteLLMSkillsTransformationHandler: @@ -40,18 +42,18 @@ class LiteLLMSkillsTransformationHandler: display_title: str | None = None, description: str | None = None, instructions: str | None = None, - files: list[Any] | None = None, + files: Sequence[object] | None = None, file_content: bytes | None = None, file_name: str | None = None, file_type: str | None = None, - metadata: dict[str, Any] | None = None, + metadata: dict[str, JsonValue] | None = None, user_id: str | None = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, _is_async: bool = False, logging_obj: Optional["LiteLLMLoggingObj"] = None, litellm_call_id: str | None = None, **kwargs, - ) -> Skill | Coroutine[Any, Any, Skill]: + ) -> Skill | Coroutine[object, object, Skill]: """ Create a skill in LiteLLM database. @@ -127,7 +129,7 @@ class LiteLLMSkillsTransformationHandler: file_content: bytes | None = None, file_name: str | None = None, file_type: str | None = None, - metadata: dict[str, Any] | None = None, + metadata: dict[str, JsonValue] | None = None, user_id: str | None = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, ) -> Skill: @@ -163,7 +165,7 @@ class LiteLLMSkillsTransformationHandler: litellm_call_id: str | None = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, **kwargs, - ) -> ListSkillsResponse | Coroutine[Any, Any, ListSkillsResponse]: + ) -> ListSkillsResponse | Coroutine[object, object, ListSkillsResponse]: """ List skills from LiteLLM database. @@ -235,7 +237,7 @@ class LiteLLMSkillsTransformationHandler: litellm_call_id: str | None = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, **kwargs, - ) -> Skill | Coroutine[Any, Any, Skill]: + ) -> Skill | Coroutine[object, object, Skill]: """ Get a skill from LiteLLM database. @@ -296,7 +298,7 @@ class LiteLLMSkillsTransformationHandler: litellm_call_id: str | None = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, **kwargs, - ) -> DeleteSkillResponse | Coroutine[Any, Any, DeleteSkillResponse]: + ) -> DeleteSkillResponse | Coroutine[object, object, DeleteSkillResponse]: """ Delete a skill from LiteLLM database. @@ -352,7 +354,7 @@ class LiteLLMSkillsTransformationHandler: type=result.get("type", "skill_deleted"), ) - def _db_skill_to_response(self, db_skill: Any) -> Skill: + def _db_skill_to_response(self, db_skill: "LiteLLM_SkillsTable") -> Skill: """ Convert a database skill record to Anthropic-compatible Skill response. @@ -362,21 +364,8 @@ class LiteLLMSkillsTransformationHandler: Returns: Skill object """ - created_at = "" - updated_at = "" - - if hasattr(db_skill, "created_at") and db_skill.created_at: - created_at = ( - db_skill.created_at.isoformat() - if hasattr(db_skill.created_at, "isoformat") - else str(db_skill.created_at) - ) - if hasattr(db_skill, "updated_at") and db_skill.updated_at: - updated_at = ( - db_skill.updated_at.isoformat() - if hasattr(db_skill.updated_at, "isoformat") - else str(db_skill.updated_at) - ) + created_at: Final = db_skill.created_at.isoformat() if db_skill.created_at else "" + updated_at: Final = db_skill.updated_at.isoformat() if db_skill.updated_at else "" return Skill( id=db_skill.skill_id, diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 0c95fd4df07..a76a8a3e98c 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -7,7 +7,7 @@ Docs - https://docs.mistral.ai/api/ """ from collections.abc import AsyncIterator, Coroutine, Iterator -from typing import Any, Final, Literal, cast, get_type_hints, overload +from typing import TYPE_CHECKING, Any, Final, Literal, cast, get_type_hints, overload import httpx @@ -26,6 +26,9 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse, ModelResponseStream from litellm.utils import convert_to_model_response_object +if TYPE_CHECKING: + import tiktoken + class MistralConfig(OpenAIGPTConfig): """ @@ -550,7 +553,7 @@ class MistralConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/mistral/ocr/guardrail_translation/handler.py b/litellm/llms/mistral/ocr/guardrail_translation/handler.py index 303e212e888..2af8172c992 100644 --- a/litellm/llms/mistral/ocr/guardrail_translation/handler.py +++ b/litellm/llms/mistral/ocr/guardrail_translation/handler.py @@ -13,6 +13,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.ocr.transformation import OCRResponse @@ -33,7 +34,7 @@ class OCRHandler(BaseTranslation): self, data: dict, guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> Any: """ Process OCR input by applying guardrails to the document reference. @@ -87,7 +88,7 @@ class OCRHandler(BaseTranslation): self, response: "OCRResponse", guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: Any | None = None, request_data: dict | None = None, ) -> Any: diff --git a/litellm/llms/mistral/ocr/transformation.py b/litellm/llms/mistral/ocr/transformation.py index 354e41c61bf..78c8dd11171 100644 --- a/litellm/llms/mistral/ocr/transformation.py +++ b/litellm/llms/mistral/ocr/transformation.py @@ -2,7 +2,7 @@ Mistral OCR transformation implementation. """ -from typing import Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -15,6 +15,9 @@ from litellm.llms.base_llm.ocr.transformation import ( ) from litellm.secret_managers.main import get_secret_str +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + MISTRAL_OCR_API_KEY_ENV_VAR: Final = "MISTRAL_API_KEY" @@ -198,7 +201,7 @@ class MistralOCRConfig(BaseOCRConfig): self, model: str, raw_response: httpx.Response, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", **kwargs, ) -> OCRResponse: """ diff --git a/litellm/llms/nlp_cloud/chat/transformation.py b/litellm/llms/nlp_cloud/chat/transformation.py index a06786d2163..17c547618d3 100644 --- a/litellm/llms/nlp_cloud/chat/transformation.py +++ b/litellm/llms/nlp_cloud/chat/transformation.py @@ -14,6 +14,8 @@ from litellm.utils import ModelResponse, Usage from ..common_utils import NLPCloudError if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj LoggingClass = LiteLLMLoggingObj @@ -173,7 +175,7 @@ class NLPCloudConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/oci/chat/cohere.py b/litellm/llms/oci/chat/cohere.py index 7ae438fd4cd..384e7ec4cf8 100644 --- a/litellm/llms/oci/chat/cohere.py +++ b/litellm/llms/oci/chat/cohere.py @@ -8,10 +8,11 @@ response parsing, and streaming chunk parsing for models served with import datetime import json +from collections.abc import Iterable, Mapping, Sequence from typing import Any, Final import httpx -from pydantic import ValidationError +from pydantic import JsonValue, TypeAdapter, ValidationError from litellm.llms.oci.chat.generic import ( _normalize_oci_finish_reason, @@ -35,7 +36,7 @@ from litellm.types.llms.oci import ( CohereToolMessage, CohereToolResult, ) -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ChatCompletionAssistantToolCall from litellm.types.utils import ( Choices, Delta, @@ -46,19 +47,60 @@ from litellm.types.utils import ( ) -def _extract_text_content(content: Any) -> str: - """Return the plain-text representation of a message content value.""" +def _json_dict(value: JsonValue) -> dict[str, JsonValue]: + return value if isinstance(value, dict) else {} + + +def _json_list(value: JsonValue) -> list[JsonValue]: + return value if isinstance(value, list) else [] + + +def _json_str(value: JsonValue) -> str: + return value if isinstance(value, str) else "" + + +def _content_block_text(block: Mapping[str, object]) -> str: + if not isinstance(block, dict) or block.get("type") != "text": + return "" + text: Final = block.get("text", "") + return text if isinstance(text, str) else "" + + +def _content_text(content: str | Iterable[Mapping[str, object]] | None) -> str: if content is None: return "" if isinstance(content, str): return content if isinstance(content, list): - return "".join( - item.get("text", "") for item in content if isinstance(item, dict) and item.get("type") == "text" - ) + return "".join(_content_block_text(block) for block in content) return str(content) +def _extract_text_content(content: Any) -> str: + """Return the plain-text representation of a message content value.""" + return _content_text(content) + + +_TOOL_ARGUMENTS_ADAPTER: Final = TypeAdapter(dict[str, object]) + + +def _parsed_tool_arguments(raw_arguments: str | dict[str, object]) -> dict[str, object]: + if not isinstance(raw_arguments, str): + return raw_arguments + try: + return _TOOL_ARGUMENTS_ADAPTER.validate_json(raw_arguments) + except ValidationError: + return {} + + +def _to_cohere_tool_call(tool_call: ChatCompletionAssistantToolCall) -> CohereToolCall: + function_fields: Final = tool_call.get("function", {}) + return CohereToolCall( + name=str(function_fields.get("name", "")), + parameters=_parsed_tool_arguments(function_fields.get("arguments", "{}")), + ) + + def adapt_messages_to_cohere_standard( messages: list[AllMessageValues], ) -> list[CohereMessage]: @@ -78,21 +120,12 @@ def adapt_messages_to_cohere_standard( """ # First pass: build tool_call_id → CohereToolCall so tool-result messages can # reference the originating call by name and parameters. - tool_call_lookup: Final[dict[str, CohereToolCall]] = {} - for msg in messages: - if msg.get("role") == "assistant": - tool_calls_raw: Any = msg.get("tool_calls") or [] - for tc in tool_calls_raw: - tc_id = tc.get("id", "") - raw_args = tc.get("function", {}).get("arguments", "{}") - try: - params: dict[str, object] = json.loads(raw_args) if isinstance(raw_args, str) else raw_args - except json.JSONDecodeError: - params = {} - tool_call_lookup[tc_id] = CohereToolCall( - name=str(tc.get("function", {}).get("name", "")), - parameters=params, - ) + tool_call_lookup: Final = { + tool_call.get("id", ""): _to_cohere_tool_call(tool_call) + for msg in messages + if msg.get("role") == "assistant" and "tool_calls" in msg + for tool_call in msg["tool_calls"] or [] + } last_user_index: Final = next( (i for i in range(len(messages) - 1, -1, -1) if messages[i].get("role") == "user"), @@ -107,24 +140,11 @@ def adapt_messages_to_cohere_standard( role = msg.get("role") content = _extract_text_content(msg.get("content")) - tool_calls: list[CohereToolCall] | None = None - if role == "assistant" and msg.get("tool_calls"): - tool_calls = [] - for tc in msg["tool_calls"]: # pyright: ignore[reportOptionalIterable] # truthiness check above rules out None - raw_arguments = tc.get("function", {}).get("arguments", {}) - if isinstance(raw_arguments, str): - try: - arguments: dict[str, object] = json.loads(raw_arguments) - except json.JSONDecodeError: - arguments = {} - else: - arguments = raw_arguments - tool_calls.append( - CohereToolCall( - name=str(tc.get("function", {}).get("name", "")), - parameters=arguments, - ) - ) + tool_calls = ( + [_to_cohere_tool_call(tool_call) for tool_call in msg["tool_calls"]] + if role == "assistant" and "tool_calls" in msg and msg["tool_calls"] + else None + ) if role == "user": chat_history.append(CohereMessage(role="USER", message=content)) @@ -150,8 +170,41 @@ def adapt_messages_to_cohere_standard( return chat_history +def _resolved_oci_parameter_schema(raw_parameters: dict[str, JsonValue]) -> JsonValue: + return sanitize_oci_schema(resolve_oci_schema_anyof(resolve_oci_schema_refs(raw_parameters))) + + +def _cohere_parameter_definition(param_schema: dict[str, JsonValue], is_required: bool) -> CohereParameterDefinition: + json_type: Final = _json_str(param_schema.get("type")) or "string" + return CohereParameterDefinition( + description=enrich_cohere_param_description(_json_str(param_schema.get("description")), param_schema), + type=OCI_JSON_TO_PYTHON_TYPES.get(json_type, json_type), + isRequired=is_required, + ) + + +def _cohere_parameter_definitions(resolved_schema: JsonValue) -> dict[str, CohereParameterDefinition]: + schema_fields: Final = _json_dict(resolved_schema) + required: Final = _json_list(schema_fields.get("required")) + return { + param_name: _cohere_parameter_definition(_json_dict(param_schema), param_name in required) + for param_name, param_schema in _json_dict(schema_fields.get("properties")).items() + } + + +def _to_cohere_tool(tool: Mapping[str, JsonValue]) -> CohereTool: + function_def: Final = _json_dict(tool.get("function")) + return CohereTool( + name=_json_str(function_def.get("name")), + description=_json_str(function_def.get("description")), + parameterDefinitions=_cohere_parameter_definitions( + _resolved_oci_parameter_schema(_json_dict(function_def.get("parameters"))) + ), + ) + + def adapt_tool_definitions_to_cohere_standard( - tools: list[dict[str, Any]], + tools: Sequence[Mapping[str, JsonValue]], ) -> list[CohereTool]: """Adapt OpenAI-format tool definitions to the OCI Cohere format. @@ -160,45 +213,18 @@ def adapt_tool_definitions_to_cohere_standard( - Embeds unsupported constraints (enum, format, range, pattern) into the parameter description so the model can still see them. """ - cohere_tools: Final = [] - for tool in tools: - function_def = tool.get("function", {}) - raw_params = function_def.get("parameters", {}) - - resolved = sanitize_oci_schema(resolve_oci_schema_anyof(resolve_oci_schema_refs(raw_params))) - properties = resolved.get("properties", {}) - required = resolved.get("required", []) - - parameter_definitions = {} - for param_name, param_schema in properties.items(): - json_type = param_schema.get("type", "string") - python_type = OCI_JSON_TO_PYTHON_TYPES.get(json_type, json_type) - parameter_definitions[param_name] = CohereParameterDefinition( - description=enrich_cohere_param_description(param_schema.get("description", ""), param_schema), - type=python_type, - isRequired=param_name in required, - ) - - cohere_tools.append( - CohereTool( - name=function_def.get("name", ""), - description=function_def.get("description", ""), - parameterDefinitions=parameter_definitions, - ) - ) - - return cohere_tools + return [_to_cohere_tool(tool) for tool in tools] def handle_cohere_response( - json_response: dict, + json_response: Mapping[str, JsonValue], model: str, model_response: ModelResponse, raw_response: httpx.Response, ) -> ModelResponse: """Parse a non-streaming Cohere OCI response into a LiteLLM ModelResponse.""" try: - cohere_response: Final = CohereChatResult(**json_response) + cohere_response: Final = CohereChatResult.model_validate(json_response) except (TypeError, ValidationError) as e: raise OCIError( message=f"Response cannot be casted to CohereChatResult: {e}", @@ -258,7 +284,7 @@ def handle_cohere_response( def handle_cohere_stream_chunk( - dict_chunk: dict, + dict_chunk: Mapping[str, JsonValue], prior_tool_calls_emitted: bool = False, prior_text_emitted: bool = False, ) -> ModelResponseStream: @@ -279,7 +305,7 @@ def handle_cohere_stream_chunk( the text is passed through so the response content isn't silently lost. """ try: - typed_chunk: Final = CohereStreamChunk(**dict_chunk) + typed_chunk: Final = CohereStreamChunk.model_validate(dict_chunk) except (TypeError, ValidationError) as e: raise OCIError( status_code=500, diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 94494a87bba..98e23a59eea 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -65,6 +65,8 @@ from litellm.types.utils import ( from litellm.utils import supports_reasoning if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -601,7 +603,7 @@ class OCIChatConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index d6aa1f1743b..de626b468f0 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -31,6 +31,8 @@ from litellm.types.utils import ModelResponse, ModelResponseStream from ..common_utils import OllamaError if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -319,7 +321,7 @@ class OllamaChatConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: str, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py index 65edd5cb718..dccc83efed4 100644 --- a/litellm/llms/ollama/completion/transformation.py +++ b/litellm/llms/ollama/completion/transformation.py @@ -31,6 +31,8 @@ from litellm.types.utils import ( from ..common_utils import OllamaError, OllamaModelInfo, _convert_image if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -246,7 +248,7 @@ class OllamaConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: str, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: @@ -323,9 +325,10 @@ class OllamaConfig(BaseConfig): model_response.created = int(time.time()) model_response.model = "ollama/" + model _prompt: Final = request_data.get("prompt", "") + tokenizer: Final = encoding if encoding is not None else litellm.encoding prompt_tokens: Final = response_json.get( "prompt_eval_count", - len(encoding.encode(_prompt, disallowed_special=())), + len(tokenizer.encode(_prompt, disallowed_special=())), ) completion_tokens: Final = response_json.get( "eval_count", len(response_json.get("message", dict()).get("content", "")) diff --git a/litellm/llms/oobabooga/chat/oobabooga.py b/litellm/llms/oobabooga/chat/oobabooga.py index 8655d8c28c8..cd118a0af29 100644 --- a/litellm/llms/oobabooga/chat/oobabooga.py +++ b/litellm/llms/oobabooga/chat/oobabooga.py @@ -1,6 +1,6 @@ import json from collections.abc import Callable -from typing import Any, Final +from typing import TYPE_CHECKING, Final import litellm from litellm.llms.custom_httpx.http_handler import _get_httpx_client @@ -9,6 +9,9 @@ from litellm.utils import EmbeddingResponse, ModelResponse, Usage from ..common_utils import OobaboogaError from .transformation import OobaboogaConfig +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + oobabooga_config: Final = OobaboogaConfig() @@ -92,7 +95,7 @@ def embedding( model_response: EmbeddingResponse, api_key: str | None, api_base: str | None, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", optional_params: dict, encoding=None, ): diff --git a/litellm/llms/oobabooga/chat/transformation.py b/litellm/llms/oobabooga/chat/transformation.py index f695b2226e3..43d627102b6 100644 --- a/litellm/llms/oobabooga/chat/transformation.py +++ b/litellm/llms/oobabooga/chat/transformation.py @@ -11,6 +11,8 @@ from litellm.types.utils import ModelResponse, Usage from ..common_utils import OobaboogaError if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj LoggingClass = LiteLLMLoggingObj @@ -37,7 +39,7 @@ class OobaboogaConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index 9b7c5a3f857..5894658e5d2 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -54,6 +54,8 @@ from litellm.utils import convert_to_model_response_object from ..common_utils import OpenAIError if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from litellm.llms.base_llm.base_utils import BaseTokenCounter from litellm.types.llms.openai import ChatCompletionToolParam @@ -595,7 +597,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index e411dc497fc..54673c77f80 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -47,6 +47,7 @@ from litellm.types.utils import ( if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj class OpenAIChatCompletionsHandler(BaseTranslation): @@ -75,7 +76,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): self, data: dict, guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> Any: """ Process input messages by applying guardrails to text content. @@ -324,7 +325,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): self, response: "ModelResponse", guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: Any | None = None, request_data: dict | None = None, ) -> Any: @@ -435,7 +436,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): self, responses_so_far: list["ModelResponseStream"], guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: Any | None = None, request_data: dict | None = None, stream_transform_sink: StreamTransformSink | None = None, @@ -485,7 +486,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): *, responses_so_far: list["ModelResponseStream"], guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None, + litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: Any | None, request_data: dict | None, ) -> list["ModelResponseStream"]: @@ -621,7 +622,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): *, responses_so_far: list["ModelResponseStream"], guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None, + litellm_logging_obj: "LiteLLMLoggingObj | None", user_api_key_dict: Any | None, request_data: dict | None, sink: StreamTransformSink, diff --git a/litellm/llms/openai/completion/guardrail_translation/handler.py b/litellm/llms/openai/completion/guardrail_translation/handler.py index 2c8c61ebf4e..f3557d4017e 100644 --- a/litellm/llms/openai/completion/guardrail_translation/handler.py +++ b/litellm/llms/openai/completion/guardrail_translation/handler.py @@ -13,6 +13,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.utils import TextCompletionResponse @@ -33,7 +34,7 @@ class OpenAITextCompletionHandler(BaseTranslation): self, data: dict, guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> Any: """ Process input prompt by applying guardrails to text content. @@ -120,7 +121,7 @@ class OpenAITextCompletionHandler(BaseTranslation): self, response: "TextCompletionResponse", guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: Any | None = None, request_data: dict | None = None, ) -> Any: diff --git a/litellm/llms/openai/embeddings/guardrail_translation/handler.py b/litellm/llms/openai/embeddings/guardrail_translation/handler.py index 280b0783e52..ef464e8a849 100644 --- a/litellm/llms/openai/embeddings/guardrail_translation/handler.py +++ b/litellm/llms/openai/embeddings/guardrail_translation/handler.py @@ -13,6 +13,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.utils import EmbeddingResponse @@ -35,7 +36,7 @@ class OpenAIEmbeddingsHandler(BaseTranslation): self, data: dict, guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> Any: """ Process input text by applying guardrails to text content. @@ -70,7 +71,7 @@ class OpenAIEmbeddingsHandler(BaseTranslation): data: dict, input_data: str, guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None, + litellm_logging_obj: "LiteLLMLoggingObj | None", ) -> dict: """Process a single string input through the guardrail.""" inputs: Final = GenericGuardrailAPIInputs(texts=[input_data]) @@ -99,7 +100,7 @@ class OpenAIEmbeddingsHandler(BaseTranslation): data: dict, input_data: list[str | int | list[int]], guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None, + litellm_logging_obj: "LiteLLMLoggingObj | None", ) -> dict: """Process a list input through the guardrail (if it contains strings).""" if len(input_data) == 0: @@ -144,7 +145,7 @@ class OpenAIEmbeddingsHandler(BaseTranslation): self, response: "EmbeddingResponse", guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: Any | None = None, request_data: dict | None = None, ) -> Any: diff --git a/litellm/llms/openai/image_generation/dall_e_2_transformation.py b/litellm/llms/openai/image_generation/dall_e_2_transformation.py index accdbf29efa..74936cf1895 100644 --- a/litellm/llms/openai/image_generation/dall_e_2_transformation.py +++ b/litellm/llms/openai/image_generation/dall_e_2_transformation.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -10,6 +10,7 @@ from litellm.types.utils import ImageResponse from litellm.utils import convert_to_model_response_object if TYPE_CHECKING: + import tiktoken from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj @@ -51,7 +52,7 @@ class DallE2ImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/openai/image_generation/dall_e_3_transformation.py b/litellm/llms/openai/image_generation/dall_e_3_transformation.py index 02a287d375a..5c561d011a9 100644 --- a/litellm/llms/openai/image_generation/dall_e_3_transformation.py +++ b/litellm/llms/openai/image_generation/dall_e_3_transformation.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -10,6 +10,7 @@ from litellm.types.utils import ImageResponse from litellm.utils import convert_to_model_response_object if TYPE_CHECKING: + import tiktoken from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj @@ -51,7 +52,7 @@ class DallE3ImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/openai/image_generation/gpt_transformation.py b/litellm/llms/openai/image_generation/gpt_transformation.py index 28abb136557..05494c497ca 100644 --- a/litellm/llms/openai/image_generation/gpt_transformation.py +++ b/litellm/llms/openai/image_generation/gpt_transformation.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -10,6 +10,7 @@ from litellm.types.utils import ImageResponse from litellm.utils import convert_to_model_response_object if TYPE_CHECKING: + import tiktoken from litellm.litellm_core_utils.logging import Logging as LiteLLMLoggingObj @@ -60,7 +61,7 @@ class GPTImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/openai/image_generation/guardrail_translation/handler.py b/litellm/llms/openai/image_generation/guardrail_translation/handler.py index e6f1c7efc31..b1d64fb1c09 100644 --- a/litellm/llms/openai/image_generation/guardrail_translation/handler.py +++ b/litellm/llms/openai/image_generation/guardrail_translation/handler.py @@ -13,6 +13,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.utils import ImageResponse @@ -32,7 +33,7 @@ class OpenAIImageGenerationHandler(BaseTranslation): self, data: dict, guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> Any: """ Process input prompt by applying guardrails to text content. @@ -82,7 +83,7 @@ class OpenAIImageGenerationHandler(BaseTranslation): self, response: "ImageResponse", guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: Any | None = None, request_data: dict | None = None, ) -> Any: diff --git a/litellm/llms/openai/image_variations/transformation.py b/litellm/llms/openai/image_variations/transformation.py index be171bb3522..afd2909b697 100644 --- a/litellm/llms/openai/image_variations/transformation.py +++ b/litellm/llms/openai/image_variations/transformation.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import TYPE_CHECKING from aiohttp import ClientResponse from httpx import Headers, Response @@ -11,6 +11,9 @@ from litellm.types.utils import FileTypes, HttpHandlerRequestFields, ImageRespon from ...base_llm.image_variations.transformation import BaseImageVariationConfig from ..common_utils import OpenAIError +if TYPE_CHECKING: + import tiktoken + class OpenAIImageVariationConfig(BaseImageVariationConfig): def get_supported_openai_params(self, model: str) -> list[OpenAIImageVariationOptionalParams]: @@ -50,7 +53,7 @@ class OpenAIImageVariationConfig(BaseImageVariationConfig): image: FileTypes, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, ) -> ImageResponse: return model_response @@ -65,7 +68,7 @@ class OpenAIImageVariationConfig(BaseImageVariationConfig): image: FileTypes, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, ) -> ImageResponse: return model_response diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index ee0efb88a38..6e66c998acf 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -7,6 +7,7 @@ from urllib.parse import urlparse import httpx if TYPE_CHECKING: + import tiktoken from aiohttp import ClientSession import openai @@ -264,7 +265,7 @@ class OpenAIConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: object, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: @@ -1345,7 +1346,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): data: dict, model_response: ModelResponse, timeout: float, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj, api_key: str | None = None, api_base: str | None = None, client=None, @@ -1408,7 +1409,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): prompt: str, timeout: float, optional_params: dict, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj, api_key: str | None = None, api_base: str | None = None, model_response: ImageResponse | None = None, diff --git a/litellm/llms/openai/speech/guardrail_translation/handler.py b/litellm/llms/openai/speech/guardrail_translation/handler.py index ea3bd6e6c53..9e338e80632 100644 --- a/litellm/llms/openai/speech/guardrail_translation/handler.py +++ b/litellm/llms/openai/speech/guardrail_translation/handler.py @@ -13,6 +13,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.llms.openai import HttpxBinaryResponseContent @@ -31,7 +32,7 @@ class OpenAITextToSpeechHandler(BaseTranslation): self, data: dict, guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> Any: """ Process input text by applying guardrails. @@ -80,7 +81,7 @@ class OpenAITextToSpeechHandler(BaseTranslation): self, response: "HttpxBinaryResponseContent", guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: Any | None = None, request_data: dict | None = None, ) -> Any: diff --git a/litellm/llms/openai/transcriptions/guardrail_translation/handler.py b/litellm/llms/openai/transcriptions/guardrail_translation/handler.py index 0b8a88d64b0..97fd1038d35 100644 --- a/litellm/llms/openai/transcriptions/guardrail_translation/handler.py +++ b/litellm/llms/openai/transcriptions/guardrail_translation/handler.py @@ -13,6 +13,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.utils import TranscriptionResponse @@ -31,7 +32,7 @@ class OpenAIAudioTranscriptionHandler(BaseTranslation): self, data: dict, guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> Any: """ Process input - not applicable for audio transcription. @@ -55,7 +56,7 @@ class OpenAIAudioTranscriptionHandler(BaseTranslation): self, response: "TranscriptionResponse", guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: Any | None = None, request_data: dict | None = None, ) -> Any: diff --git a/litellm/llms/openai_like/chat/transformation.py b/litellm/llms/openai_like/chat/transformation.py index f0fd7db7f9f..030710c8b2d 100644 --- a/litellm/llms/openai_like/chat/transformation.py +++ b/litellm/llms/openai_like/chat/transformation.py @@ -13,6 +13,8 @@ from litellm.types.utils import ModelResponse from ...openai.chat.gpt_transformation import OpenAIGPTConfig if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -129,7 +131,7 @@ class OpenAILikeChatConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index 71c21f14351..77a902149d9 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -8,7 +8,7 @@ Docs: https://openrouter.ai/docs/parameters from collections.abc import AsyncIterator, Iterator from enum import Enum -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, cast import httpx @@ -22,6 +22,11 @@ from litellm.types.utils import ModelResponse, ModelResponseStream from ...openai.chat.gpt_transformation import OpenAIGPTConfig from ..common_utils import OpenRouterException +if TYPE_CHECKING: + import tiktoken + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + class CacheControlSupportedModels(str, Enum): """Models that support cache_control in content blocks.""" @@ -172,12 +177,12 @@ class OpenrouterConfig(OpenAIGPTConfig): model: str, raw_response: httpx.Response, model_response: ModelResponse, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", request_data: dict, messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/openrouter/image_generation/transformation.py b/litellm/llms/openrouter/image_generation/transformation.py index 3342a6e4c71..6bbda324336 100644 --- a/litellm/llms/openrouter/image_generation/transformation.py +++ b/litellm/llms/openrouter/image_generation/transformation.py @@ -50,6 +50,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj else: LiteLLMLoggingObj = Any @@ -317,7 +319,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/perplexity/chat/transformation.py b/litellm/llms/perplexity/chat/transformation.py index bf33103b480..354f7692fd5 100644 --- a/litellm/llms/perplexity/chat/transformation.py +++ b/litellm/llms/perplexity/chat/transformation.py @@ -2,7 +2,7 @@ Translate from OpenAI's `/v1/chat/completions` to Perplexity's `/v1/chat/completions` """ -from typing import Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -14,6 +14,9 @@ from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import AllMessageValues, ChatCompletionAnnotation, ChatCompletionAnnotationURLCitation from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage +if TYPE_CHECKING: + import tiktoken + class PerplexityChatConfig(OpenAIGPTConfig): @property @@ -72,7 +75,7 @@ class PerplexityChatConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/petals/completion/transformation.py b/litellm/llms/petals/completion/transformation.py index 97b021bb119..3e0de14a7b2 100644 --- a/litellm/llms/petals/completion/transformation.py +++ b/litellm/llms/petals/completion/transformation.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import TYPE_CHECKING, Final from httpx import Headers, Response @@ -13,6 +13,9 @@ from litellm.types.utils import ModelResponse from ..common_utils import PetalsError +if TYPE_CHECKING: + import tiktoken + class PetalsConfig(BaseConfig): """ @@ -109,7 +112,7 @@ class PetalsConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/recraft/image_generation/transformation.py b/litellm/llms/recraft/image_generation/transformation.py index 2b7b44c7233..3a04e0a62b4 100644 --- a/litellm/llms/recraft/image_generation/transformation.py +++ b/litellm/llms/recraft/image_generation/transformation.py @@ -14,6 +14,8 @@ from litellm.types.llms.recraft import RecraftImageGenerationRequestParams from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -120,7 +122,7 @@ class RecraftImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/reducto/ocr/transformation.py b/litellm/llms/reducto/ocr/transformation.py index 84d5164cf87..a7216e4ec40 100644 --- a/litellm/llms/reducto/ocr/transformation.py +++ b/litellm/llms/reducto/ocr/transformation.py @@ -1,4 +1,4 @@ -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final import httpx @@ -17,6 +17,9 @@ from litellm.llms.reducto.common import ( upload_bytes_sync, ) +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + class _BaseReductoOCRConfig(BaseOCRConfig): def map_ocr_params( @@ -127,7 +130,7 @@ class _BaseReductoOCRConfig(BaseOCRConfig): self, model: str, raw_response: httpx.Response, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", **kwargs, ) -> OCRResponse: response_json: Final = raw_response.json() diff --git a/litellm/llms/replicate/chat/transformation.py b/litellm/llms/replicate/chat/transformation.py index 4cee5489fe0..769160c6ced 100644 --- a/litellm/llms/replicate/chat/transformation.py +++ b/litellm/llms/replicate/chat/transformation.py @@ -19,6 +19,8 @@ from litellm.utils import token_counter from ..common_utils import ReplicateError if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj LoggingClass = LiteLLMLoggingObj @@ -235,7 +237,7 @@ class ReplicateConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/runwayml/image_generation/transformation.py b/litellm/llms/runwayml/image_generation/transformation.py index 344c8ae2d7c..cde65addb65 100644 --- a/litellm/llms/runwayml/image_generation/transformation.py +++ b/litellm/llms/runwayml/image_generation/transformation.py @@ -20,6 +20,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -294,7 +296,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: @@ -369,7 +371,7 @@ class RunwayMLImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/sagemaker/chat/handler.py b/litellm/llms/sagemaker/chat/handler.py index b3e9ed671fc..3f62b7276df 100644 --- a/litellm/llms/sagemaker/chat/handler.py +++ b/litellm/llms/sagemaker/chat/handler.py @@ -5,6 +5,7 @@ from typing import Final import httpx +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.utils import ModelResponse, get_secret @@ -34,6 +35,7 @@ class SagemakerChatHandler(BaseAWSLLM): optional_params.pop("aws_bedrock_runtime_endpoint", None) # https://bedrock-runtime.{region_name}.amazonaws.com aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None) + aws_external_id: Final = optional_params.pop("aws_external_id", None) ### SET REGION NAME ### if aws_region_name is None: @@ -60,6 +62,7 @@ class SagemakerChatHandler(BaseAWSLLM): aws_role_name=aws_role_name, aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, ) return credentials, aws_region_name @@ -79,10 +82,11 @@ class SagemakerChatHandler(BaseAWSLLM): raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") sigv4: Final = SigV4Auth(credentials, "sagemaker", aws_region_name) + dns_suffix: Final = get_aws_dns_suffix(aws_region_name) if optional_params.get("stream") is True: - api_base = f"https://runtime.sagemaker.{aws_region_name}.amazonaws.com/endpoints/{model}/invocations-response-stream" + api_base = f"https://runtime.sagemaker.{aws_region_name}.{dns_suffix}/endpoints/{model}/invocations-response-stream" else: - api_base = f"https://runtime.sagemaker.{aws_region_name}.amazonaws.com/endpoints/{model}/invocations" + api_base = f"https://runtime.sagemaker.{aws_region_name}.{dns_suffix}/endpoints/{model}/invocations" sagemaker_base_url: Final = optional_params.get("sagemaker_base_url", None) if sagemaker_base_url is not None: diff --git a/litellm/llms/sagemaker/chat/transformation.py b/litellm/llms/sagemaker/chat/transformation.py index 37ddd813d6f..04995f32d97 100644 --- a/litellm/llms/sagemaker/chat/transformation.py +++ b/litellm/llms/sagemaker/chat/transformation.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any, Final, cast import httpx from httpx._models import Headers +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.logging_utils import track_llm_api_timing from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -93,10 +94,11 @@ class SagemakerChatConfig(OpenAIGPTConfig, BaseAWSLLM): model=model, model_id=None, ) + dns_suffix: Final = get_aws_dns_suffix(aws_region_name) if stream is True: - api_base = f"https://runtime.sagemaker.{aws_region_name}.amazonaws.com/endpoints/{model}/invocations-response-stream" + api_base = f"https://runtime.sagemaker.{aws_region_name}.{dns_suffix}/endpoints/{model}/invocations-response-stream" else: - api_base = f"https://runtime.sagemaker.{aws_region_name}.amazonaws.com/endpoints/{model}/invocations" + api_base = f"https://runtime.sagemaker.{aws_region_name}.{dns_suffix}/endpoints/{model}/invocations" sagemaker_base_url: Final = cast(str | None, optional_params.get("sagemaker_base_url")) if sagemaker_base_url is not None: diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index 84cad56f0d4..fb8074d3682 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -1,13 +1,14 @@ import json from collections.abc import Callable from copy import deepcopy -from typing import Any, Final, cast +from typing import Final, cast import httpx import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import ( @@ -57,6 +58,7 @@ class SagemakerLLM(BaseAWSLLM): optional_params.pop("aws_bedrock_runtime_endpoint", None) # https://bedrock-runtime.{region_name}.amazonaws.com aws_web_identity_token: Final = optional_params.pop("aws_web_identity_token", None) aws_sts_endpoint: Final = optional_params.pop("aws_sts_endpoint", None) + aws_external_id: Final = optional_params.pop("aws_external_id", None) ### SET REGION NAME ### if aws_region_name is None: @@ -83,6 +85,7 @@ class SagemakerLLM(BaseAWSLLM): aws_role_name=aws_role_name, aws_web_identity_token=aws_web_identity_token, aws_sts_endpoint=aws_sts_endpoint, + aws_external_id=aws_external_id, ) return credentials, aws_region_name @@ -104,10 +107,11 @@ class SagemakerLLM(BaseAWSLLM): raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") sigv4: Final = SigV4Auth(credentials, "sagemaker", aws_region_name) + dns_suffix: Final = get_aws_dns_suffix(aws_region_name) if optional_params.get("stream") is True: - api_base = f"https://runtime.sagemaker.{aws_region_name}.amazonaws.com/endpoints/{model}/invocations-response-stream" + api_base = f"https://runtime.sagemaker.{aws_region_name}.{dns_suffix}/endpoints/{model}/invocations-response-stream" else: - api_base = f"https://runtime.sagemaker.{aws_region_name}.amazonaws.com/endpoints/{model}/invocations" + api_base = f"https://runtime.sagemaker.{aws_region_name}.{dns_suffix}/endpoints/{model}/invocations" sagemaker_base_url: Final = optional_params.get("sagemaker_base_url", None) if sagemaker_base_url is not None: @@ -404,7 +408,7 @@ class SagemakerLLM(BaseAWSLLM): encoding, model_response: ModelResponse, model_id: str | None, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj, litellm_params: dict, headers: dict, ): @@ -467,7 +471,7 @@ class SagemakerLLM(BaseAWSLLM): encoding, model_response: ModelResponse, optional_params: dict, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj, model_id: str | None, headers: dict, litellm_params: dict, diff --git a/litellm/llms/sagemaker/completion/transformation.py b/litellm/llms/sagemaker/completion/transformation.py index f0962a8eb66..576018f0046 100644 --- a/litellm/llms/sagemaker/completion/transformation.py +++ b/litellm/llms/sagemaker/completion/transformation.py @@ -24,6 +24,8 @@ from litellm.utils import token_counter from ..common_utils import SagemakerError if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -196,7 +198,7 @@ class SagemakerConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: str, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/sagemaker/embedding/cohere_transformation.py b/litellm/llms/sagemaker/embedding/cohere_transformation.py index b05e146a966..4687ff6b3f4 100644 --- a/litellm/llms/sagemaker/embedding/cohere_transformation.py +++ b/litellm/llms/sagemaker/embedding/cohere_transformation.py @@ -13,6 +13,7 @@ Reference: https://docs.cohere.com/v2/reference/embed from typing import TYPE_CHECKING, Any, cast if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.llms.openai import AllEmbeddingInputValues from httpx._models import Headers, Response @@ -90,7 +91,7 @@ class SagemakerCohereEmbeddingConfig(BaseEmbeddingConfig): model: str, raw_response: Response, model_response: "EmbeddingResponse", - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", api_key: str | None = None, request_data: dict = {}, optional_params: dict = {}, diff --git a/litellm/llms/sagemaker/embedding/transformation.py b/litellm/llms/sagemaker/embedding/transformation.py index 04bf040098e..97940929b09 100644 --- a/litellm/llms/sagemaker/embedding/transformation.py +++ b/litellm/llms/sagemaker/embedding/transformation.py @@ -7,6 +7,7 @@ In the Huggingface TGI format. from typing import TYPE_CHECKING, Any, Final if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.llms.openai import AllEmbeddingInputValues from httpx._models import Headers, Response @@ -84,7 +85,7 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig): model: str, raw_response: Response, model_response: "EmbeddingResponse", - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", api_key: str | None = None, request_data: dict = {}, optional_params: dict = {}, diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py index a376e9c60b3..d64d7a57281 100755 --- a/litellm/llms/sap/chat/transformation.py +++ b/litellm/llms/sap/chat/transformation.py @@ -15,6 +15,8 @@ from litellm.types.utils import ModelResponse from ...openai.chat.gpt_transformation import OpenAIGPTConfig if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -381,7 +383,7 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/soniox/common_utils.py b/litellm/llms/soniox/common_utils.py index e2b18736e0e..cb33f0b8996 100644 --- a/litellm/llms/soniox/common_utils.py +++ b/litellm/llms/soniox/common_utils.py @@ -2,7 +2,8 @@ Shared utilities for the Soniox provider (https://soniox.com). """ -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final, TypeAlias from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( SubtitleToken, @@ -11,6 +12,8 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( ) from litellm.llms.base_llm.chat.transformation import BaseLLMException +SonioxToken: TypeAlias = Mapping[str, object] + # Soniox API base URL. SONIOX_API_BASE: Final[str] = "https://api.soniox.com" @@ -68,7 +71,15 @@ def get_soniox_api_base(api_base: str | None = None) -> str: return base.rstrip("/") -def render_soniox_tokens(tokens: list[dict[str, Any]]) -> str: +def _token_text(value: object) -> str: + return value if isinstance(value, str) else "" + + +def _token_milliseconds(value: object) -> int | None: + return value if isinstance(value, int) else None + + +def render_soniox_tokens(tokens: Sequence[SonioxToken]) -> str: """ Render a list of Soniox tokens to a readable transcript string. @@ -85,11 +96,11 @@ def render_soniox_tokens(tokens: list[dict[str, Any]]) -> str: return "" text_parts: Final[list[str]] = [] - current_speaker: Any | None = None - current_language: Any | None = None + current_speaker: object = None + current_language: object = None for token in tokens: - text = token.get("text", "") + text = _token_text(token.get("text", "")) speaker = token.get("speaker") language = token.get("language") is_translation = token.get("translation_status") == "translation" @@ -107,23 +118,27 @@ def render_soniox_tokens(tokens: list[dict[str, Any]]) -> str: current_language = language prefix = "[Translation] " if is_translation else "" text_parts.append(f"\n{prefix}[{current_language}] ") - text = text.lstrip() if isinstance(text, str) else text + text = text.lstrip() text_parts.append(text) return "".join(text_parts) -def _soniox_token_to_subtitle_token(token: dict[str, Any]) -> SubtitleToken: +def _token_speaker(value: object) -> str | int | None: + return value if isinstance(value, str | int) else None + + +def _soniox_token_to_subtitle_token(token: SonioxToken) -> SubtitleToken: return SubtitleToken( - text=token.get("text", ""), - start_ms=token.get("start_ms"), - end_ms=token.get("end_ms"), - speaker=token.get("speaker"), + text=_token_text(token.get("text", "")), + start_ms=_token_milliseconds(token.get("start_ms")), + end_ms=_token_milliseconds(token.get("end_ms")), + speaker=_token_speaker(token.get("speaker")), ) -def render_soniox_tokens_as_srt(tokens: list[dict[str, Any]]) -> str: +def render_soniox_tokens_as_srt(tokens: Sequence[SonioxToken]) -> str: """ Render Soniox tokens as SRT (SubRip) subtitle format. @@ -132,7 +147,7 @@ def render_soniox_tokens_as_srt(tokens: list[dict[str, Any]]) -> str: return render_subtitle_tokens_as_srt(tuple(_soniox_token_to_subtitle_token(token) for token in tokens)) -def render_soniox_tokens_as_vtt(tokens: list[dict[str, Any]]) -> str: +def render_soniox_tokens_as_vtt(tokens: Sequence[SonioxToken]) -> str: """ Render Soniox tokens as WebVTT subtitle format. diff --git a/litellm/llms/stability/image_generation/transformation.py b/litellm/llms/stability/image_generation/transformation.py index 804613ea161..cf3576a9404 100644 --- a/litellm/llms/stability/image_generation/transformation.py +++ b/litellm/llms/stability/image_generation/transformation.py @@ -26,6 +26,8 @@ from litellm.types.llms.stability import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -205,7 +207,7 @@ class StabilityImageGenerationConfig(BaseImageGenerationConfig): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/topaz/image_variations/transformation.py b/litellm/llms/topaz/image_variations/transformation.py index 3c914eb6a4c..f4753c8ba17 100644 --- a/litellm/llms/topaz/image_variations/transformation.py +++ b/litellm/llms/topaz/image_variations/transformation.py @@ -2,7 +2,7 @@ import base64 import time from collections.abc import Mapping from io import BytesIO -from typing import Any, Final +from typing import TYPE_CHECKING, Final from aiohttp import ClientResponse from httpx import Headers, Response @@ -22,6 +22,9 @@ from litellm.types.utils import ( from ...base_llm.image_variations.transformation import BaseImageVariationConfig from ..common_utils import TopazException, TopazModelInfo +if TYPE_CHECKING: + import tiktoken + class TopazImageVariationConfig(TopazModelInfo, BaseImageVariationConfig): def get_supported_openai_params(self, model: str) -> list[OpenAIImageVariationOptionalParams]: @@ -136,7 +139,7 @@ class TopazImageVariationConfig(TopazModelInfo, BaseImageVariationConfig): image: FileTypes, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, ) -> ImageResponse: image_content: Final = await raw_response.read() @@ -155,7 +158,7 @@ class TopazImageVariationConfig(TopazModelInfo, BaseImageVariationConfig): image: FileTypes, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, ) -> ImageResponse: image_content: Final = raw_response.content diff --git a/litellm/llms/triton/completion/transformation.py b/litellm/llms/triton/completion/transformation.py index 5f1986c6124..98a68ba2c36 100644 --- a/litellm/llms/triton/completion/transformation.py +++ b/litellm/llms/triton/completion/transformation.py @@ -4,7 +4,7 @@ Translates from OpenAI's `/v1/chat/completions` endpoint to Triton's `/generate` import json from collections.abc import AsyncIterator, Iterator -from typing import Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal from httpx import Headers, Response @@ -28,6 +28,9 @@ from litellm.types.utils import ( from ..common_utils import TritonError +if TYPE_CHECKING: + import tiktoken + class TritonConfig(BaseConfig): """ @@ -92,7 +95,7 @@ class TritonConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: @@ -212,7 +215,7 @@ class TritonGenerateConfig(TritonConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: @@ -277,7 +280,7 @@ class TritonInferConfig(TritonConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/vertex_ai/agent_engine/transformation.py b/litellm/llms/vertex_ai/agent_engine/transformation.py index 76aaa4895e2..e430d9e2280 100644 --- a/litellm/llms/vertex_ai/agent_engine/transformation.py +++ b/litellm/llms/vertex_ai/agent_engine/transformation.py @@ -29,6 +29,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import Choices, Message, ModelResponse, Usage if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.utils import CustomStreamWrapper @@ -283,7 +285,7 @@ class VertexAgentEngineConfig(BaseConfig, VertexBase): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 6481b67fad7..377cd9f3437 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -1,8 +1,9 @@ import json from collections.abc import Coroutine -from typing import Any, Final +from typing import TYPE_CHECKING, Final, Protocol import httpx +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.litellm_core_utils.url_utils import ( @@ -20,11 +21,47 @@ from litellm.types.llms.openai import CreateBatchRequest from litellm.types.llms.vertex_ai import ( VERTEX_CREDENTIALS_TYPES, VertexAIBatchPredictionJob, + VertexBatchPredictionResponse, ) from litellm.types.utils import LiteLLMBatch from .transformation import VertexAIBatchTransformation +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + +class _VertexBatchJsonSource(Protocol): + """An HTTP response whose JSON body is a single Vertex AI batch prediction job.""" + + def json(self) -> VertexBatchPredictionResponse: ... + + +class _VertexBatchListJsonSource(Protocol): + """An HTTP response whose JSON body is a page of Vertex AI batch prediction jobs.""" + + def json(self) -> dict[str, object]: ... + + +class _VertexBatchPayloadView(TypedDict): + """Holds one decoded batch prediction job so the payload reads back typed.""" + + payload: ReadOnly[VertexBatchPredictionResponse] + + +class _FetchedResponseView(TypedDict): + """Holds one ``safe_get`` result so the response reads back as ``httpx.Response``.""" + + response: ReadOnly[httpx.Response] + + +def _vertex_batch_payload(response: _VertexBatchJsonSource) -> VertexBatchPredictionResponse: + return response.json() + + +def _vertex_batch_list_payload(response: _VertexBatchListJsonSource) -> dict[str, object]: + return response.json() + class VertexAIBatchPrediction(VertexLLM): def __init__(self, gcs_bucket_name: str, *args, **kwargs): @@ -41,7 +78,7 @@ class VertexAIBatchPrediction(VertexLLM): vertex_location: str | None, timeout: float | httpx.Timeout, max_retries: int | None, - ) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]: + ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: sync_handler: Final = _get_httpx_client() access_token, project_id = self._ensure_access_token( @@ -98,7 +135,8 @@ class VertexAIBatchPrediction(VertexLLM): data=json.dumps(vertex_batch_request), ) - _json_response: Final = response.json() + payload_view: Final[_VertexBatchPayloadView] = {"payload": response.json()} + _json_response: Final = payload_view["payload"] vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( response=_json_response ) @@ -128,7 +166,8 @@ class VertexAIBatchPrediction(VertexLLM): ) raise - _json_response: Final = response.json() + payload_view: Final[_VertexBatchPayloadView] = {"payload": response.json()} + _json_response: Final = payload_view["payload"] vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( response=_json_response ) @@ -154,8 +193,8 @@ class VertexAIBatchPrediction(VertexLLM): vertex_location: str | None, timeout: float | httpx.Timeout, max_retries: int | None, - logging_obj: Any | None = None, - ) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]: + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: sync_handler: Final = _get_httpx_client() access_token, project_id = self._ensure_access_token( @@ -231,20 +270,22 @@ class VertexAIBatchPrediction(VertexLLM): # rebind / private / cloud-metadata targets are rejected; the # proxy auth gate already blocks malicious clientside ``api_base`` # at the boundary — this is defense-in-depth for SDK callers. - response: Final = safe_get( - sync_handler, - api_base, - headers=headers, - ) + fetched: Final[_FetchedResponseView] = { + "response": safe_get( + sync_handler, + api_base, + headers=headers, + ) + } + response: Final = fetched["response"] if response.status_code != 200: raise VertexAIError( status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" ) - _json_response: Final = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( - response=_json_response + response=_vertex_batch_payload(response) ) return vertex_batch_response @@ -252,7 +293,7 @@ class VertexAIBatchPrediction(VertexLLM): self, api_base: str, headers: dict[str, str], - logging_obj: Any | None = None, + logging_obj: "LiteLLMLoggingObj | None" = None, ) -> LiteLLMBatch: client: Final = get_async_httpx_client( llm_provider=litellm.LlmProviders.VERTEX_AI, @@ -284,19 +325,21 @@ class VertexAIBatchPrediction(VertexLLM): # request kwargs, so wrap the fetch in ``async_safe_get`` to reject # DNS-rebind / private / cloud-metadata targets. Defense-in-depth # behind the proxy auth gate's clientside ``api_base`` check. - response: Final = await async_safe_get( - client, - api_base, - headers=headers, - ) + fetched: Final[_FetchedResponseView] = { + "response": await async_safe_get( + client, + api_base, + headers=headers, + ) + } + response: Final = fetched["response"] if response.status_code != 200: raise VertexAIError( status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" ) - _json_response: Final = response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( - response=_json_response + response=_vertex_batch_payload(response) ) return vertex_batch_response @@ -345,11 +388,9 @@ class VertexAIBatchPrediction(VertexLLM): "Authorization": f"Bearer {access_token}", } - params: Final[dict[str, Any]] = {} - if limit is not None: - params["pageSize"] = str(limit) - if after is not None: - params["pageToken"] = after + limit_params: Final[dict[str, str]] = {"pageSize": str(limit)} if limit is not None else {} + after_params: Final[dict[str, str]] = {"pageToken": after} if after is not None else {} + params: Final = {**limit_params, **after_params} if _is_async is True: return self._async_list_batches( @@ -369,7 +410,7 @@ class VertexAIBatchPrediction(VertexLLM): status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" ) - _json_response: Final = response.json() + _json_response: Final = _vertex_batch_list_payload(response) vertex_batch_response: Final = ( VertexAIBatchTransformation.transform_vertex_ai_batch_list_response_to_openai_list_response( response=_json_response @@ -381,7 +422,7 @@ class VertexAIBatchPrediction(VertexLLM): self, api_base: str, headers: dict[str, str], - params: dict[str, Any], + params: dict[str, str], ): client: Final = get_async_httpx_client( llm_provider=litellm.LlmProviders.VERTEX_AI, @@ -396,7 +437,7 @@ class VertexAIBatchPrediction(VertexLLM): status_code=response.status_code, message=f"Error: {response.status_code} {response.text}" ) - _json_response: Final = response.json() + _json_response: Final = _vertex_batch_list_payload(response) vertex_batch_response: Final = ( VertexAIBatchTransformation.transform_vertex_ai_batch_list_response_to_openai_list_response( response=_json_response @@ -414,7 +455,7 @@ class VertexAIBatchPrediction(VertexLLM): vertex_location: str | None, timeout: float | httpx.Timeout, max_retries: int | None, - ) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]: + ) -> LiteLLMBatch | Coroutine[object, object, LiteLLMBatch]: access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_project, @@ -494,9 +535,8 @@ class VertexAIBatchPrediction(VertexLLM): message=f"Error: {retrieve_response.status_code} {retrieve_response.text}", ) - _json_response: Final = retrieve_response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( - response=_json_response + response=_vertex_batch_payload(retrieve_response) ) return vertex_batch_response @@ -541,8 +581,7 @@ class VertexAIBatchPrediction(VertexLLM): message=f"Error: {retrieve_response.status_code} {retrieve_response.text}", ) - _json_response: Final = retrieve_response.json() vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_response_to_openai_batch_response( - response=_json_response + response=_vertex_batch_payload(retrieve_response) ) return vertex_batch_response diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index 13c1ba5a697..f81d4ca777e 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -3,7 +3,7 @@ Google AI Studio /batchEmbedContents Embeddings Endpoint """ import json -from typing import Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal import httpx @@ -29,6 +29,9 @@ from .batch_embed_content_transformation import ( transform_openai_input_gemini_embed_content, ) +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + class GoogleBatchEmbeddings(VertexLLM): @staticmethod @@ -125,7 +128,7 @@ class GoogleBatchEmbeddings(VertexLLM): model_response: EmbeddingResponse, custom_llm_provider: Literal["gemini", "vertex_ai"], optional_params: dict, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", api_key: str | None = None, api_base: str | None = None, encoding=None, @@ -290,7 +293,7 @@ class GoogleBatchEmbeddings(VertexLLM): use_embed_content: bool = False, api_key: str | None = None, optional_params: dict | None = None, - logging_obj: Any | None = None, + logging_obj: "LiteLLMLoggingObj | None" = None, ) -> EmbeddingResponse: if client is None: _params: Final = {} diff --git a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py index 67c6bff4381..5889a8eba06 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_gemini_transformation.py @@ -2,7 +2,7 @@ import base64 import json import os from io import BufferedReader, BytesIO -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast import httpx from httpx._types import RequestFiles @@ -14,6 +14,11 @@ from litellm.llms.vertex_ai.common_utils import get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM from litellm.secret_managers.main import get_secret_str from litellm.types.images.main import ImageEditOptionalRequestParams +from litellm.types.llms.vertex_ai import ( + GenerateContentResponseBody, + HttpxContentType, + HttpxPartType, +) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import FileTypes, ImageObject, ImageResponse, OpenAIImage @@ -25,6 +30,16 @@ else: LiteLLMLoggingObj = Any +class _GenerateContentSource(Protocol): + """An HTTP response whose JSON body is a Gemini ``generateContent`` result.""" + + def json(self) -> GenerateContentResponseBody: ... + + +def _generate_content_payload(response: _GenerateContentSource) -> GenerateContentResponseBody: + return response.json() + + class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): """ Vertex AI Gemini Image Edit Configuration @@ -46,16 +61,13 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): image_edit_optional_params: ImageEditOptionalRequestParams, model: str, drop_params: bool, - ) -> dict[str, Any]: + ) -> dict[str, str]: supported_params: Final = self.get_supported_openai_params(model) - filtered_params = {key: value for key, value in image_edit_optional_params.items() if key in supported_params} + if "size" not in supported_params or "size" not in image_edit_optional_params: + return {} - mapped_params: Final[dict[str, Any]] = {} - - if "size" in filtered_params: - mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(filtered_params["size"]) - - return mapped_params + size: Final = image_edit_optional_params.get("size") + return {"aspectRatio": self._map_size_to_aspect_ratio(size or "")} def _resolve_vertex_project(self) -> str | None: return ( @@ -86,12 +98,12 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): def validate_environment( self, - headers: dict, + headers: dict[str, str], model: str, api_key: str | None = None, - litellm_params: dict | None = None, + litellm_params: dict[str, object] | None = None, api_base: str | None = None, - ) -> dict: + ) -> dict[str, str]: headers = headers or {} litellm_params = litellm_params or {} @@ -116,7 +128,7 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): self, model: str, api_base: str | None, - litellm_params: dict, + litellm_params: dict[str, object], ) -> str: """ Get the complete URL for Vertex AI Gemini generateContent API @@ -148,38 +160,36 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): model: str, prompt: str | None, image: FileTypes | None, - image_edit_optional_request_params: dict[str, Any], + image_edit_optional_request_params: dict[str, object], litellm_params: GenericLiteLLMParams, - headers: dict, - ) -> tuple[dict[str, Any], RequestFiles | None]: + headers: dict[str, str], + ) -> tuple[dict[str, object], RequestFiles | None]: inline_parts: Final = self._prepare_inline_image_parts(image) if image else [] if not inline_parts: raise ValueError("Vertex AI Gemini image edit requires at least one image.") # Build parts list with image and prompt (if provided) - parts: Final = inline_parts.copy() - if prompt is not None and prompt != "": - parts.append({"text": prompt}) + text_parts: Final[list[HttpxPartType]] = [{"text": prompt}] if prompt is not None and prompt != "" else [] + parts: Final[list[HttpxPartType]] = [*inline_parts, *text_parts] # Correct format for Vertex AI Gemini image editing - contents: Final = {"role": "USER", "parts": parts} - - request_body: Final[dict[str, Any]] = {"contents": contents} - - # Generation config with proper structure for image editing - generation_config: Final[dict[str, Any]] = {"response_modalities": ["IMAGE"]} + contents: Final[dict[str, object]] = {"role": "USER", "parts": parts} # Add image-specific configuration - image_config: Final[dict[str, Any]] = {} - if "aspectRatio" in image_edit_optional_request_params: - image_config["aspect_ratio"] = image_edit_optional_request_params["aspectRatio"] + image_config: Final = ( + {"aspect_ratio": image_edit_optional_request_params["aspectRatio"]} + if "aspectRatio" in image_edit_optional_request_params + else None + ) - if image_config: - generation_config["image_config"] = image_config + # Generation config with proper structure for image editing + generation_config: Final[dict[str, object]] = { + key: value for key, value in (("response_modalities", ["IMAGE"]), ("image_config", image_config)) if value + } - request_body["generationConfig"] = generation_config + request_body: Final[dict[str, object]] = {"contents": contents, "generationConfig": generation_config} - payload: Final[Any] = json.dumps(request_body) + payload: Final = json.dumps(request_body) empty_files: Final = cast(RequestFiles, []) return cast(tuple[dict[str, Any], RequestFiles | None], (payload, empty_files)) @@ -187,11 +197,11 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): self, model: str, raw_response: httpx.Response, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj, ) -> ImageResponse: model_response: Final = ImageResponse() try: - response_json: Final = raw_response.json() + response_json: Final = _generate_content_payload(raw_response) except Exception as exc: raise self.get_error_class( error_message=f"Error transforming image edit response: {exc}", @@ -200,20 +210,15 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): ) candidates: Final = response_json.get("candidates", []) - data_list: Final[list[ImageObject]] = [] - - for candidate in candidates: - content = candidate.get("content", {}) - parts = content.get("parts", []) - for part in parts: - inline_data = part.get("inlineData") - if inline_data and inline_data.get("data"): - data_list.append( - ImageObject( - b64_json=inline_data["data"], - url=None, - ) - ) + contents: Final[list[HttpxContentType]] = [ + candidate["content"] for candidate in candidates if "content" in candidate + ] + parts: Final[list[HttpxPartType]] = [part for content in contents for part in content.get("parts", [])] + data_list: Final[list[ImageObject]] = [ + ImageObject(b64_json=b64_json, url=None) + for part in parts + if (inline_data := part.get("inlineData")) and (b64_json := inline_data.get("data")) + ] model_response.data = cast(list[OpenAIImage], data_list) return model_response @@ -229,30 +234,18 @@ class VertexAIGeminiImageEditConfig(BaseImageEditConfig, VertexLLM): } return aspect_ratio_map.get(size, "1:1") - def _prepare_inline_image_parts(self, image: FileTypes | list[FileTypes]) -> list[dict[str, Any]]: - images: list[FileTypes] - if isinstance(image, list): - images = image - else: - images = [image] - - inline_parts: Final[list[dict[str, Any]]] = [] - for img in images: - if img is None: - continue - - mime_type = ImageEditRequestUtils.get_image_content_type(img) - image_bytes = self._read_all_bytes(img) - inline_parts.append( - { - "inlineData": { - "mimeType": mime_type, - "data": base64.b64encode(image_bytes).decode("utf-8"), - } + def _prepare_inline_image_parts(self, image: FileTypes | list[FileTypes]) -> list[HttpxPartType]: + images: Final[list[FileTypes]] = image if isinstance(image, list) else [image] + return [ + { + "inlineData": { + "mimeType": ImageEditRequestUtils.get_image_content_type(img), + "data": base64.b64encode(self._read_all_bytes(img)).decode("utf-8"), } - ) - - return inline_parts + } + for img in images + if img is not None + ] def _read_all_bytes(self, image: FileTypes) -> bytes: if isinstance(image, bytes): diff --git a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py index 9c6e943dc04..c6ad5928b74 100644 --- a/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_edit/vertex_imagen_transformation.py @@ -195,7 +195,7 @@ class VertexAIImagenImageEditConfig(BaseImageEditConfig, VertexLLM): self, model: str, raw_response: httpx.Response, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj, ) -> ImageResponse: model_response: Final = ImageResponse() try: diff --git a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py index 2d7d78efa48..6a5bb484540 100644 --- a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py +++ b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py @@ -1,5 +1,5 @@ import json -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final import httpx from openai.types.image import Image @@ -14,6 +14,9 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import Ver from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES from litellm.types.utils import ImageResponse +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + class VertexImageGeneration(VertexLLM): def process_image_generation_response( @@ -74,7 +77,7 @@ class VertexImageGeneration(VertexLLM): vertex_location: str | None, vertex_credentials: VERTEX_CREDENTIALS_TYPES | None, model_response: ImageResponse, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", model: str = "imagegeneration", # vertex ai uses imagegeneration as the default model client: Any | None = None, optional_params: dict | None = None, @@ -173,7 +176,7 @@ class VertexImageGeneration(VertexLLM): vertex_location: str | None, vertex_credentials: VERTEX_CREDENTIALS_TYPES | None, model_response: ImageResponse, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", model: str = "imagegeneration", # vertex ai uses imagegeneration as the default model client: AsyncHTTPHandler | None = None, optional_params: dict | None = None, diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 799307f98c7..d7a2491c04a 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -24,6 +24,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -282,7 +284,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py index 64d5b55d3f4..8faf7b0d484 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_imagen_transformation.py @@ -20,6 +20,8 @@ from litellm.types.llms.openai import ( from litellm.types.utils import ImageObject, ImageResponse if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -212,7 +214,7 @@ class VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): request_data: dict, optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ImageResponse: diff --git a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py index 962dfe52c0a..c80a02c3683 100644 --- a/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/rag_api/transformation.py @@ -1,6 +1,8 @@ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, Protocol import httpx +from typing_extensions import ReadOnly, TypedDict from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig from litellm.llms.vertex_ai.common_utils import get_vertex_base_url @@ -25,6 +27,66 @@ else: LiteLLMLoggingObj = Any +class VertexRagPageSpan(TypedDict, total=False): + """Page range a retrieved chunk came from, as ``:retrieveContexts`` returns it.""" + + firstPage: ReadOnly[int] + lastPage: ReadOnly[int] + + +class VertexRagContext(TypedDict, total=False): + """One retrieved chunk in a Vertex AI RAG ``:retrieveContexts`` response.""" + + text: ReadOnly[str] + sourceUri: ReadOnly[str] + sourceDisplayName: ReadOnly[str] + score: ReadOnly[float] + pageSpan: ReadOnly[VertexRagPageSpan] + + +class VertexRagContextGroup(TypedDict, total=False): + contexts: ReadOnly[list[VertexRagContext]] + + +class VertexRagRetrieveContextsResponse(TypedDict, total=False): + contexts: ReadOnly[VertexRagContextGroup] + + +class VertexRagCorpusResponse(TypedDict, total=False): + """A RAG corpus resource, as ``POST /ragCorpora`` returns it.""" + + name: ReadOnly[str] + display_name: ReadOnly[str] + createTime: ReadOnly[object] + labels: ReadOnly[object] + + +class _SearchQueryView(TypedDict): + """Holds the logged search query so the model call detail reads back as ``str``.""" + + query: ReadOnly[str] + + +class _RetrieveContextsSource(Protocol): + """An HTTP response whose JSON body is a Vertex AI RAG ``:retrieveContexts`` result.""" + + def json(self) -> VertexRagRetrieveContextsResponse: ... + + +class _RagCorpusSource(Protocol): + """An HTTP response whose JSON body is a Vertex AI RAG corpus resource.""" + + def json(self) -> VertexRagCorpusResponse: ... + + +def _retrieve_contexts_payload(response: _RetrieveContextsSource) -> VertexRagRetrieveContextsResponse: + return response.json() + + +def _rag_corpus_payload(response: _RagCorpusSource) -> VertexRagCorpusResponse: + return response.json() + + class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): """ Configuration for Vertex AI Vector Store RAG API @@ -96,8 +158,8 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, - extra_body: dict[str, Any] | None = None, - ) -> tuple[str, dict[str, Any]]: + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: """ Transform search request for Vertex AI RAG API """ @@ -120,12 +182,6 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): # Just the corpus ID, construct full path full_rag_corpus = f"projects/{vertex_project}/locations/{vertex_location}/ragCorpora/{vector_store_id}" - # Build the request body for Vertex AI RAG API - request_body: Final[dict[str, Any]] = { - "vertex_rag_store": {"rag_resources": [{"rag_corpus": full_rag_corpus}]}, - "query": {"text": query}, - } - ######################################################### # Update logging object with details of the request ######################################################### @@ -133,22 +189,28 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): # Add optional parameters max_num_results: Final = vector_store_search_optional_params.get("max_num_results") - if max_num_results is not None: - request_body["query"]["rag_retrieval_config"] = {"top_k": max_num_results} - - # Add filters if provided filters: Final = vector_store_search_optional_params.get("filters") - if filters is not None: - if "rag_retrieval_config" not in request_body["query"]: - request_body["query"]["rag_retrieval_config"] = {} - request_body["query"]["rag_retrieval_config"]["filter"] = filters - - # Add ranking options if provided ranking_options: Final = vector_store_search_optional_params.get("ranking_options") - if ranking_options is not None: - if "rag_retrieval_config" not in request_body["query"]: - request_body["query"]["rag_retrieval_config"] = {} - request_body["query"]["rag_retrieval_config"]["ranking"] = ranking_options + rag_retrieval_config: Final[Mapping[str, object]] = { + key: value + for key, value in ( + ("top_k", max_num_results), + ("filter", filters), + ("ranking", ranking_options), + ) + if value is not None + } + + # Build the request body for Vertex AI RAG API + query_body: Final[Mapping[str, object]] = { + key: value + for key, value in (("text", query), ("rag_retrieval_config", rag_retrieval_config or None)) + if value is not None + } + request_body: Final[dict[str, object]] = { + "vertex_rag_store": {"rag_resources": [{"rag_corpus": full_rag_corpus}]}, + "query": query_body, + } return url, request_body @@ -159,12 +221,13 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): Transform Vertex AI RAG API response to standard vector store search response """ try: - response_json: Final = response.json() + response_json: Final = _retrieve_contexts_payload(response) # Extract contexts from Vertex AI response - handle nested structure - contexts: Final = response_json.get("contexts", {}).get("contexts", []) + context_group: Final[VertexRagContextGroup] = response_json.get("contexts", {}) + contexts: Final = context_group.get("contexts", []) # Transform contexts to standard format - search_results: Final = [] + search_results: Final[list[VectorStoreSearchResult]] = [] for context in contexts: content = [ VectorStoreResultContent( @@ -182,7 +245,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): filename = source_display_name if source_display_name else "Unknown Document" # Build attributes with available metadata - attributes = {} + attributes: dict[str, object] = {} if source_uri: attributes["sourceUri"] = source_uri if source_display_name: @@ -202,9 +265,10 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): ) search_results.append(result) + query_view: Final[_SearchQueryView] = {"query": litellm_logging_obj.model_call_details.get("query", "")} return VectorStoreSearchResponse( object="vector_store.search_results.page", - search_query=litellm_logging_obj.model_call_details.get("query", ""), + search_query=query_view["query"], data=search_results, ) @@ -219,22 +283,25 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): self, vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, api_base: str, - ) -> tuple[str, dict[str, Any]]: + ) -> tuple[str, dict[str, object]]: """ Transform create request for Vertex AI RAG Corpus """ url: Final = f"{api_base}/ragCorpora" # Base URL for creating RAG corpus - # Build the request body for Vertex AI RAG Corpus creation - request_body: Final[dict[str, Any]] = { - "display_name": vector_store_create_optional_params.get("name", "litellm-vector-store"), - "description": "Vector store created via LiteLLM", - } - # Add metadata if provided metadata: Final = vector_store_create_optional_params.get("metadata") - if metadata is not None: - request_body["labels"] = metadata + + # Build the request body for Vertex AI RAG Corpus creation + request_body: Final[dict[str, object]] = { + key: value + for key, value in ( + ("display_name", vector_store_create_optional_params.get("name", "litellm-vector-store")), + ("description", "Vector store created via LiteLLM"), + ("labels", metadata), + ) + if value is not None + } return url, request_body @@ -243,7 +310,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase): Transform Vertex AI RAG Corpus creation response to standard vector store response """ try: - response_json: Final = response.json() + response_json: Final = _rag_corpus_payload(response) # Extract the corpus ID from the response name corpus_name: Final = response_json.get("name", "") diff --git a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py index a0597769f7b..0bcf16ee06f 100644 --- a/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py +++ b/litellm/llms/vertex_ai/vector_stores/search_api/transformation.py @@ -1,6 +1,8 @@ -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final, Protocol import httpx +from typing_extensions import ReadOnly, TypedDict from litellm import get_model_info from litellm.exceptions import BadRequestError @@ -50,6 +52,52 @@ VERTEX_SEARCH_DATASTORE_EXTRA_BODY_FIELDS: Final = frozenset(VertexSearchDataSto VERTEX_SEARCH_ENGINE_EXTRA_BODY_FIELDS: Final = frozenset(VertexSearchEngineExtraBody.__annotations__) +class VertexSearchSnippet(TypedDict, total=False): + snippet: ReadOnly[str] + htmlSnippet: ReadOnly[str] + + +class VertexSearchDerivedStructData(TypedDict, total=False): + """The ``derivedStructData`` blob Discovery Engine attaches to each search hit.""" + + title: ReadOnly[str] + link: ReadOnly[str] + displayLink: ReadOnly[str] + formattedUrl: ReadOnly[str] + snippets: ReadOnly[list[VertexSearchSnippet]] + + +class VertexSearchDocument(TypedDict, total=False): + derivedStructData: ReadOnly[VertexSearchDerivedStructData] + + +class VertexSearchHit(TypedDict, total=False): + id: ReadOnly[str] + document: ReadOnly[VertexSearchDocument] + + +class VertexSearchApiResponse(TypedDict, total=False): + """Body of a Discovery Engine ``:search`` response.""" + + results: ReadOnly[list[VertexSearchHit]] + + +class _SearchQueryView(TypedDict): + """Holds the logged search query so the model call detail reads back as ``str``.""" + + query: ReadOnly[str] + + +class _VertexSearchApiSource(Protocol): + """An HTTP response whose JSON body is a Discovery Engine ``:search`` result.""" + + def json(self) -> VertexSearchApiResponse: ... + + +def _vertex_search_payload(response: _VertexSearchApiSource) -> VertexSearchApiResponse: + return response.json() + + class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): """ Configuration for Vertex AI Search API Vector Store @@ -61,7 +109,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): super().__init__() @staticmethod - def get_supported_extra_body_fields(is_engine: bool = False) -> frozenset: + def get_supported_extra_body_fields(is_engine: bool = False) -> frozenset[str]: """ Native SearchRequest fields callers may forward via ``extra_body``. @@ -75,7 +123,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): return VERTEX_SEARCH_DATASTORE_EXTRA_BODY_FIELDS @classmethod - def _filter_extra_body(cls, extra_body: dict[str, Any], is_engine: bool = False) -> dict[str, Any]: + def _filter_extra_body(cls, extra_body: Mapping[str, object], is_engine: bool = False) -> dict[str, object]: """ Validate ``extra_body`` against the supported-field allowlist for the active serving config (engine/app vs data store). @@ -196,8 +244,8 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): api_base: str, litellm_logging_obj: LiteLLMLoggingObj, litellm_params: dict, - extra_body: dict[str, Any] | None = None, - ) -> tuple[str, dict[str, Any]]: + extra_body: Mapping[str, object] | None = None, + ) -> tuple[str, dict[str, object]]: """ Transform a search request for the Vertex AI Search (Discovery Engine) API. @@ -222,7 +270,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): is_engine: Final = bool(litellm_params.get("vertex_engine_id")) - request_body: Final[dict[str, Any]] = {"query": query, "pageSize": 10} + request_body: Final[dict[str, object]] = {"query": query, "pageSize": 10} max_num_results: Final = vector_store_search_optional_params.get("max_num_results") if max_num_results is not None: request_body["pageSize"] = max_num_results @@ -256,7 +304,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): } """ try: - response_json: Final = response.json() + response_json: Final = _vertex_search_payload(response) # Extract results from Vertex AI Search API response results: Final = response_json.get("results", []) @@ -264,8 +312,8 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): # Transform results to standard format search_results: Final[list[VectorStoreSearchResult]] = [] for result in results: - document = result.get("document", {}) - derived_data = document.get("derivedStructData", {}) + document: VertexSearchDocument = result.get("document", {}) + derived_data: VertexSearchDerivedStructData = document.get("derivedStructData", {}) # Extract text content from snippets snippets = derived_data.get("snippets", []) @@ -329,9 +377,10 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase): ) search_results.append(result_obj) + query_view: Final[_SearchQueryView] = {"query": litellm_logging_obj.model_call_details.get("query", "")} return VectorStoreSearchResponse( object="vector_store.search_results.page", - search_query=litellm_logging_obj.model_call_details.get("query", ""), + search_query=query_view["query"], data=search_results, ) diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index a8430455323..d7ad69593c6 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -1,6 +1,6 @@ # What is this? ## Handler file for calling claude-3 on vertex ai -from typing import Any, Final +from typing import TYPE_CHECKING, Final import httpx @@ -12,6 +12,9 @@ from litellm.types.utils import ModelResponse from ....anthropic.chat.transformation import AnthropicConfig from .output_params_utils import sanitize_vertex_anthropic_output_params +if TYPE_CHECKING: + import tiktoken + class VertexAIError(Exception): def __init__(self, status_code, message): @@ -183,7 +186,7 @@ class VertexAIAnthropicConfig(AnthropicConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py index 7b0c26f5881..279035c455d 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/llama3/transformation.py @@ -1,6 +1,6 @@ import types from collections.abc import AsyncIterator, Iterator -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final import httpx @@ -20,6 +20,9 @@ from litellm.types.utils import ( from ...common_utils import VertexAIError +if TYPE_CHECKING: + import tiktoken + class VertexAILlama3Config(OpenAIGPTConfig): """ @@ -109,7 +112,7 @@ class VertexAILlama3Config(OpenAIGPTConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py index 6c955d9bab1..58cf7c7e702 100644 --- a/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py +++ b/litellm/llms/vertex_ai/vertex_gemma_models/transformation.py @@ -9,7 +9,7 @@ The actual message transformation reuses OpenAIGPTConfig since Gemma uses OpenAI """ from collections.abc import Callable -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, cast import httpx @@ -23,6 +23,11 @@ from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse +if TYPE_CHECKING: + import tiktoken + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + class VertexGemmaConfig(OpenAIGPTConfig): """ @@ -210,7 +215,7 @@ class VertexGemmaConfig(OpenAIGPTConfig): custom_prompt_dict: dict, model_response: ModelResponse, print_verbose: Callable, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", optional_params: dict, acompletion: bool, litellm_params: dict, @@ -265,12 +270,12 @@ class VertexGemmaConfig(OpenAIGPTConfig): api_key: str, model_response: ModelResponse, print_verbose: Callable, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", optional_params: dict, litellm_params: dict, client: HTTPHandler | httpx.Client | None = None, timeout: float | httpx.Timeout | None = None, - encoding: Any = None, + encoding: "tiktoken.Encoding | None" = None, ): """Synchronous completion request""" from litellm.utils import convert_to_model_response_object @@ -355,12 +360,12 @@ class VertexGemmaConfig(OpenAIGPTConfig): api_key: str, model_response: ModelResponse, print_verbose: Callable, - logging_obj: Any, + logging_obj: "LiteLLMLoggingObj", optional_params: dict, litellm_params: dict, client: AsyncHTTPHandler | httpx.AsyncClient | None = None, timeout: float | httpx.Timeout | None = None, - encoding: Any = None, + encoding: "tiktoken.Encoding | None" = None, ): """Asynchronous completion request""" from litellm.utils import convert_to_model_response_object diff --git a/litellm/llms/watsonx/completion/transformation.py b/litellm/llms/watsonx/completion/transformation.py index 2645d099ee4..0b4c9ae917a 100644 --- a/litellm/llms/watsonx/completion/transformation.py +++ b/litellm/llms/watsonx/completion/transformation.py @@ -20,6 +20,8 @@ from ..common_utils import ( ) if TYPE_CHECKING: + import tiktoken + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj LiteLLMLoggingObj = _LiteLLMLoggingObj @@ -278,7 +280,7 @@ class IBMWatsonXAIConfig(IBMWatsonXMixin, BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: str, + encoding: "tiktoken.Encoding | None", api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: diff --git a/litellm/llms/xai/oauth.py b/litellm/llms/xai/oauth.py index 37dae93a725..e8196ec6cb9 100644 --- a/litellm/llms/xai/oauth.py +++ b/litellm/llms/xai/oauth.py @@ -8,11 +8,13 @@ import threading import time import uuid import webbrowser +from collections.abc import Mapping from http.server import BaseHTTPRequestHandler, HTTPServer -from typing import Any, Final +from typing import Final from urllib.parse import parse_qs, urlencode, urlparse import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.constants import XAI_API_BASE @@ -31,6 +33,40 @@ XAI_OAUTH_CALLBACK_TIMEOUT_SECONDS: Final = 180 _XAI_OAUTH_REFRESH_LOCK: Final = threading.Lock() +class XAIOAuthRecord(TypedDict): + access_token: ReadOnly[str] + refresh_token: ReadOnly[str] + id_token: ReadOnly[str | None] + token_type: ReadOnly[str] + token_endpoint: ReadOnly[str] + expires_at: ReadOnly[float | None] + + +class _TokenPayload(TypedDict): + access_token: NotRequired[ReadOnly[str]] + refresh_token: NotRequired[ReadOnly[str]] + id_token: NotRequired[ReadOnly[str | None]] + token_type: NotRequired[ReadOnly[str]] + expires_in: NotRequired[ReadOnly[float]] + + +class _DiscoveryDocument(TypedDict): + authorization_endpoint: NotRequired[ReadOnly[str]] + token_endpoint: NotRequired[ReadOnly[str]] + + +class _AuthFileView(TypedDict): + record: ReadOnly[XAIOAuthRecord | None] + + +class _TokenPayloadView(TypedDict): + payload: ReadOnly[_TokenPayload | None] + + +class _DiscoveryView(TypedDict): + document: ReadOnly[_DiscoveryDocument] + + class XAIOAuthError(Exception): pass @@ -75,7 +111,7 @@ class _CallbackHandler(BaseHTTPRequestHandler): ) self.wfile.write(body) - def log_message(self, format: str, *args: Any) -> None: + def log_message(self, format: str, *args: object) -> None: return @@ -115,7 +151,7 @@ class XAIOAuthAuthenticator: refreshed: Final = self._refresh_tokens(locked_auth_data) return refreshed["access_token"] - def login(self, force: bool = False, no_browser: bool = False) -> dict[str, Any]: + def login(self, force: bool = False, no_browser: bool = False) -> XAIOAuthRecord: existing: Final = self._read_auth_file() if existing and not force and existing.get("access_token"): if not self._is_expired(existing): @@ -177,15 +213,16 @@ class XAIOAuthAuthenticator: except OSError: verbose_logger.debug("Could not chmod xAI OAuth token directory") - def _read_auth_file(self) -> dict[str, Any] | None: + def _read_auth_file(self) -> XAIOAuthRecord | None: try: with open(self.auth_file, "r") as f: - data: Final = json.load(f) + loaded: Final[_AuthFileView] = {"record": json.load(f)} + data: Final = loaded["record"] return data if isinstance(data, dict) else None except (OSError, json.JSONDecodeError): return None - def _write_auth_file(self, data: dict[str, Any]) -> None: + def _write_auth_file(self, data: XAIOAuthRecord) -> None: self._ensure_token_dir() tmp_file: Final = os.path.join( self.token_dir, @@ -216,7 +253,7 @@ class XAIOAuthAuthenticator: pass raise - def _is_expired(self, auth_data: dict[str, Any]) -> bool: + def _is_expired(self, auth_data: XAIOAuthRecord) -> bool: expires_at: Final = auth_data.get("expires_at") if expires_at is None: return True @@ -234,9 +271,10 @@ class XAIOAuthAuthenticator: f"xAI OAuth discovery request failed: {exc.response.status_code} {exc.response.text}" ) from exc try: - data: Final = response.json() + discovered: Final[_DiscoveryView] = {"document": response.json()} except ValueError as exc: raise XAIOAuthError("xAI OAuth discovery response was not valid JSON") from exc + data: Final = discovered["document"] authorization_endpoint: Final = data.get("authorization_endpoint") token_endpoint: Final = data.get("token_endpoint") if not authorization_endpoint or not token_endpoint: @@ -304,7 +342,7 @@ class XAIOAuthAuthenticator: server.server_close() raise XAIOAuthError("Timed out waiting for xAI OAuth callback") - def _exchange_token(self, token_endpoint: str, data: dict[str, str]) -> dict[str, Any]: + def _exchange_token(self, token_endpoint: str, data: dict[str, str]) -> _TokenPayload: try: response: Final = self._client().post( token_endpoint, @@ -320,19 +358,20 @@ class XAIOAuthAuthenticator: f"xAI OAuth token request failed: {exc.response.status_code} {exc.response.text}" ) from exc try: - body: Final = response.json() + exchanged: Final[_TokenPayloadView] = {"payload": response.json()} except ValueError as exc: raise XAIOAuthError("xAI OAuth token response was not valid JSON") from exc + body: Final = exchanged["payload"] if not isinstance(body, dict): raise XAIOAuthError("xAI OAuth token response was not an object") return body def _build_auth_record( self, - token_payload: dict[str, Any], + token_payload: _TokenPayload, token_endpoint: str, fallback_refresh_token: str | None = None, - ) -> dict[str, Any]: + ) -> XAIOAuthRecord: access_token: Final = token_payload.get("access_token") refresh_token: Final = token_payload.get("refresh_token") or fallback_refresh_token if not access_token: @@ -353,7 +392,7 @@ class XAIOAuthAuthenticator: "expires_at": expires_at, } - def _refresh_tokens(self, auth_data: dict[str, Any]) -> dict[str, Any]: + def _refresh_tokens(self, auth_data: XAIOAuthRecord) -> XAIOAuthRecord: token_endpoint = auth_data.get("token_endpoint") if not token_endpoint: token_endpoint = self._discover()["token_endpoint"] @@ -379,5 +418,5 @@ class XAIOAuthAuthenticator: return refreshed -def should_use_xai_oauth(litellm_params: dict[str, Any] | None) -> bool: +def should_use_xai_oauth(litellm_params: Mapping[str, object] | None) -> bool: return bool((litellm_params or {}).get("use_xai_oauth")) diff --git a/litellm/passthrough/utils.py b/litellm/passthrough/utils.py index df39b8fad48..7eb14fcc118 100644 --- a/litellm/passthrough/utils.py +++ b/litellm/passthrough/utils.py @@ -6,6 +6,7 @@ import httpx from litellm._logging import verbose_logger from litellm.constants import PASS_THROUGH_HEADER_PREFIX +from litellm.litellm_core_utils.aws_partition import contains_aws_arn # Headers that must not be overwritten via the x-pass- forwarding mechanism. # Includes standard credential/auth headers and protocol-level headers that @@ -126,7 +127,7 @@ class CommonUtils: import re # Early exit: if no ARN detected, return unchanged - if "arn:aws:" not in endpoint: + if not contains_aws_arn(endpoint): return endpoint # Handle all patterns in one go - more efficient and cleaner diff --git a/litellm/provider_endpoints_support_backup.json b/litellm/provider_endpoints_support_backup.json index 86c14fb4cd8..ead26ab65c5 100644 --- a/litellm/provider_endpoints_support_backup.json +++ b/litellm/provider_endpoints_support_backup.json @@ -1180,7 +1180,8 @@ "files": true, "rerank": true, "a2a": true, - "interactions": true + "interactions": true, + "video_generations": true } }, "huggingface": { diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index d5461f01ed8..425f82794e6 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -903,8 +903,9 @@ class MCPRequestHandler: NotSessionBearer, SessionBearerAdmitted, SessionBearerInvalid, + SessionSigningConfigError, + active_session_signing_keys, resolve_session_bearer, - session_keys_from_master_key, ) from litellm.proxy.proxy_server import master_key @@ -913,7 +914,10 @@ class MCPRequestHandler: await MCPRequestHandler._run_pre_db_read_auth_checks(request=request, route=route) - keys: Final = session_keys_from_master_key(master_key) + keys: Final = active_session_signing_keys(master_key) + if isinstance(keys, SessionSigningConfigError): + verbose_logger.error("mcp gateway session admission rejected: %s", keys.detail) + raise HTTPException(status_code=500, detail="Server misconfigured: mcp_session_token_signing is invalid") result: Final = resolve_session_bearer(authorization_value, keys, datetime.now(timezone.utc)) match result: case SessionBearerAdmitted(): diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 314c80adbc4..853a07972c1 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -65,15 +65,16 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credentials import ( SessionRefreshOpened, + SessionSigningConfigError, + active_session_signing_keys, open_session_refresh_bearer, - session_keys_from_master_key, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( SESSION_REFRESH_TTL_SECONDS, MintedSessionToken, SessionAudience, - SessionKeys, SessionPrincipal, + SessionSigningKeys, mint_session_refresh_token, mint_session_token, ) @@ -885,7 +886,7 @@ class _SingleUseGuard: return "first" if count == 1 else "replayed" -def _session_token_pair(principal: SessionPrincipal, keys: SessionKeys, now: datetime) -> Response: +def _session_token_pair(principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime) -> Response: access: Final = mint_session_token(principal, keys, now) refresh: Final = mint_session_refresh_token(principal, keys, now) if not isinstance(access, MintedSessionToken) or not isinstance(refresh, MintedSessionToken): @@ -912,7 +913,7 @@ class _ProxyCredentialTokenResponse(TypedDict): def _proxy_credential_response( - minted: MintedProxyCredential, principal: SessionPrincipal, keys: SessionKeys, now: datetime + minted: MintedProxyCredential, principal: SessionPrincipal, keys: SessionSigningKeys, now: datetime ) -> Response: """The proxy-API token response: the access token is the very credential ``lite login`` stores (accepted on every proxy route with user and team attribution), and @@ -998,7 +999,10 @@ async def aggregate_token( if master_key is None: verbose_logger.error("mcp_gateway_dcr token grant rejected: no master_key configured") return _oauth_error(500, "server_error", "the gateway has no master key configured") - keys: Final = session_keys_from_master_key(master_key) + keys: Final = active_session_signing_keys(master_key) + if isinstance(keys, SessionSigningConfigError): + verbose_logger.error("mcp_gateway_dcr token grant rejected: %s", keys.detail) + return _oauth_error(500, "server_error", "the gateway session signing configuration is invalid") now: Final = datetime.now(timezone.utc) issue: Final = _GrantIssuer( request=request, @@ -1043,7 +1047,7 @@ class _GrantIssuer: self, request: Request, resource: str | None, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, reload_user: ReloadUser, mint_proxy_credential: MintProxyCredential, @@ -1146,7 +1150,7 @@ async def _refresh_token_grant( refresh_token: str | None, client_id: str, resource: str | None, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, issue: _GrantIssuer, ) -> Response: @@ -1182,7 +1186,10 @@ async def revoke_refresh_token(token: str, client_id: str, master_key: str | Non if master_key is None: verbose_logger.error("mcp_gateway_dcr revoke rejected: no master_key configured") return _oauth_error(500, "server_error", "the gateway has no master key configured") - keys: Final = session_keys_from_master_key(master_key) + keys: Final = active_session_signing_keys(master_key) + if isinstance(keys, SessionSigningConfigError): + verbose_logger.error("mcp_gateway_dcr revoke rejected: %s", keys.detail) + return _oauth_error(500, "server_error", "the gateway session signing configuration is invalid") now: Final = datetime.now(timezone.utc) opened: Final = open_session_refresh_bearer(token, keys, now, expected_client_id=client_id) if isinstance(opened, SessionRefreshOpened): diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py index e836e2bd363..4918229c2b8 100644 --- a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py +++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py @@ -39,6 +39,7 @@ if TYPE_CHECKING: from mcp.types import CallToolResult from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj class MCPGuardrailTranslationHandler(BaseTranslation): @@ -48,7 +49,7 @@ class MCPGuardrailTranslationHandler(BaseTranslation): self, data: dict[str, Any], guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> dict[str, Any]: mcp_tool_name: Final = data.get("mcp_tool_name") or data.get("name") mcp_arguments = data.get("mcp_arguments") or data.get("arguments") @@ -99,7 +100,7 @@ class MCPGuardrailTranslationHandler(BaseTranslation): self, response: "CallToolResult", guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, user_api_key_dict: Any | None = None, request_data: dict | None = None, ) -> Any: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py index 70a04ac290a..df2bbdba345 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_credentials.py @@ -20,13 +20,16 @@ from datetime import datetime from functools import lru_cache from typing import Final, Literal, TypeAlias -from pydantic import BaseModel, ConfigDict, SecretStr +from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( + AsymmetricSessionKeys, OpenedSessionToken, SessionExpired, SessionKeys, SessionPrincipal, + SessionRotatedPublicKey, + SessionSigningKeys, is_session_refresh_token, is_session_token, open_session_refresh_token, @@ -68,6 +71,99 @@ def session_keys_from_master_key(master_key: str) -> SessionKeys: return SessionKeys(signing_key=SecretStr(signing)) +class SessionSigningPreviousKey(BaseModel): + """One retired key in ``mcp_session_token_signing.previous_public_keys``: its ``kid`` + and the PEM public half (inline or an ``os.environ/`` reference).""" + + model_config = ConfigDict(frozen=True, extra="forbid") + kid: str = Field(min_length=1) + public_key: str = Field(min_length=1) + + +class MCPSessionTokenSigningSettings(BaseModel): + """The ``general_settings.mcp_session_token_signing`` block: opt-in asymmetric signing + for the gateway session tokens. Absent, the gateway keeps the backward-compatible + HS256 key derived from ``master_key``. ``private_key`` and each ``public_key`` accept + a PEM string inline or an ``os.environ/`` (or secret manager) reference.""" + + model_config = ConfigDict(frozen=True, extra="forbid") + algorithm: Literal["RS256"] + kid: str = Field(min_length=1) + private_key: str = Field(min_length=1) + previous_public_keys: tuple[SessionSigningPreviousKey, ...] = () + + +class SessionSigningConfigError(BaseModel): + """``mcp_session_token_signing`` is present but unusable (bad shape, unresolvable + secret reference, or a key that is not a loadable RSA PEM); the caller fails closed + with a server error instead of silently falling back to HS256.""" + + model_config = ConfigDict(frozen=True) + tag: Literal["session_signing_config_error"] = "session_signing_config_error" + detail: str + + +def _resolve_key_material(value: str) -> str | None: + if not value.startswith("os.environ/"): + return value + from litellm.secret_managers.main import get_secret_str # noqa: PLC0415 # heavy import kept off the pure path + + return get_secret_str(value) + + +def resolve_session_signing_keys( + master_key: str, + raw_settings: object | None, +) -> SessionSigningKeys | SessionSigningConfigError: + """Turn the operator's ``mcp_session_token_signing`` setting into signing key material. + + ``None`` (the setting absent) keeps the backward-compatible HS256 key derived from + ``master_key``. A present setting must fully validate into RS256 material; any defect + is a ``SessionSigningConfigError`` value so token issuance and admission fail closed + rather than minting under a key the operator did not intend. + """ + if raw_settings is None: + return session_keys_from_master_key(master_key) + try: + settings: Final = MCPSessionTokenSigningSettings.model_validate(raw_settings) + except ValidationError as exc: + return SessionSigningConfigError(detail=f"mcp_session_token_signing is malformed: {exc}") + private_pem: Final = _resolve_key_material(settings.private_key) + if private_pem is None: + return SessionSigningConfigError(detail="mcp_session_token_signing.private_key reference did not resolve") + resolved_previous: Final = tuple( + (previous.kid, _resolve_key_material(previous.public_key)) for previous in settings.previous_public_keys + ) + unresolved: Final = tuple(kid for kid, pem in resolved_previous if pem is None) + if unresolved: + return SessionSigningConfigError( + detail=f"mcp_session_token_signing.previous_public_keys reference did not resolve for kid(s): {', '.join(unresolved)}" + ) + try: + return AsymmetricSessionKeys( + private_key_pem=SecretStr(private_pem), + kid=settings.kid, + previous_public_keys=tuple( + SessionRotatedPublicKey(kid=kid, public_key_pem=pem) + for kid, pem in resolved_previous + if pem is not None + ), + ) + except ValidationError as exc: + return SessionSigningConfigError( + detail=f"mcp_session_token_signing keys are not usable RSA PEM material: {exc}" + ) + + +def active_session_signing_keys(master_key: str) -> SessionSigningKeys | SessionSigningConfigError: + """Wiring helper for the token endpoint and the admission edge: resolve the signing + keys from the live ``general_settings.mcp_session_token_signing`` block, or derive the + default HS256 key from ``master_key`` when the block is absent.""" + from litellm.proxy.proxy_server import general_settings # noqa: PLC0415 # circular import at module load + + return resolve_session_signing_keys(master_key, general_settings.get("mcp_session_token_signing")) + + class NotSessionBearer(BaseModel): """The bearer is not session-shaped; admission continues on its normal path.""" @@ -116,7 +212,7 @@ def is_session_bearer_shaped(authorization_value: str) -> bool: def resolve_session_bearer( authorization_value: str, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, ) -> SessionBearerResult: """Classify an ``Authorization`` value presented at the aggregate MCP edge. @@ -166,7 +262,7 @@ SessionRefreshResult: TypeAlias = SessionRefreshOpened | SessionRefreshInvalid def open_session_refresh_bearer( refresh_value: str, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, expected_client_id: str, ) -> SessionRefreshResult: diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py index 2c7b970ca0e..6824f96f927 100644 --- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py +++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/session_token.py @@ -8,8 +8,11 @@ is therefore a stable REFERENCE, not an authorization: admission reloads the liv record and policy on every request, so deactivating the user (or their team) kills outstanding sessions immediately without a revocation store. -Wire shape: ``llm_session_`` (access) / ``llm_srefresh_`` (refresh) + an HS256 JWT, -the same signing approach as :mod:`.envelope`. Claims are ``iss``/``iat``/``exp`` +Wire shape: ``llm_session_`` (access) / ``llm_srefresh_`` (refresh) + a JWT signed with +the injected key material: HS256 under the default master-key-derived secret (the same +signing approach as :mod:`.envelope`), or RS256 under an operator-provided RSA private +key (:class:`AsymmetricSessionKeys`) so downstream validators hold only the public half. +Claims are ``iss``/``iat``/``exp`` plus ``jti`` (per-mint uniqueness, so two tokens minted in the same second never collide and a future revocation list has a stable handle), ``kind``, ``user_id``, and ``client_id``; ``client_id`` binds the refresh token @@ -31,11 +34,16 @@ injected ``now``); the strict pydantic claims model is the sole, total type gate from __future__ import annotations import secrets +from collections import Counter from datetime import datetime, timedelta +from functools import lru_cache from typing import Final, Literal, TypeAlias import jwt -from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError +from cryptography.exceptions import UnsupportedAlgorithm +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from pydantic import BaseModel, ConfigDict, Field, SecretStr, ValidationError, field_validator, model_validator SESSION_TOKEN_PREFIX: Final = "llm_session_" """Marker prefix on every serialized session ACCESS token so the admission edge can cheaply @@ -71,6 +79,11 @@ limits while bounding hostile input before JWT parsing.""" _SESSION_JWT_ALGORITHM: Final = "HS256" +_SESSION_RSA_ALGORITHM: Final = "RS256" + +_MIN_RSA_KEY_BITS: Final = 2048 +"""RFC 7518 section 3.3: RS256 requires a key of at least 2048 bits.""" + SessionTokenKind = Literal["session", "session_refresh"] """Which credential a session token is. Stamped into the signed claims and required to match on open, so a signature-valid token of one kind cannot be replayed as the other even if its @@ -120,6 +133,85 @@ class SessionKeys(BaseModel): signing_key: SecretStr = Field(min_length=32) +class SessionRotatedPublicKey(BaseModel): + """The public half of a retired signing key, kept verifiable under its ``kid`` during a + rotation window so tokens minted before the rotation stay valid until they expire.""" + + model_config = ConfigDict(frozen=True) + kid: str = Field(min_length=1) + public_key_pem: str = Field(min_length=1) + + @field_validator("public_key_pem") + @classmethod + def _pem_is_an_rsa_public_key(cls, value: str) -> str: + try: + loaded: Final = serialization.load_pem_public_key(value.encode()) + except (ValueError, TypeError, UnsupportedAlgorithm) as exc: + raise ValueError(f"public_key_pem is not a loadable PEM public key: {exc}") from exc + if not isinstance(loaded, rsa.RSAPublicKey): + raise ValueError("public_key_pem must be an RSA public key in PEM format") # noqa: TRY004 # pydantic validators must raise ValueError + if loaded.key_size < _MIN_RSA_KEY_BITS: + raise ValueError(f"public_key_pem must be an RSA key of at least {_MIN_RSA_KEY_BITS} bits") + return value + + +class AsymmetricSessionKeys(BaseModel): + """Injected RS256 key material: the issuer-held RSA private key and the stable ``kid`` + stamped into every minted token's JOSE header, plus the public halves of previously + rotated keys that verification still accepts while their tokens age out. Downstream + validators never need the private key: :func:`session_public_key_pem` yields the + public half to distribute.""" + + model_config = ConfigDict(frozen=True) + private_key_pem: SecretStr + kid: str = Field(min_length=1) + previous_public_keys: tuple[SessionRotatedPublicKey, ...] = () + + @field_validator("private_key_pem") + @classmethod + def _pem_is_a_strong_rsa_private_key(cls, value: SecretStr) -> SecretStr: + try: + loaded: Final = serialization.load_pem_private_key(value.get_secret_value().encode(), password=None) + except (ValueError, TypeError, UnsupportedAlgorithm) as exc: + raise ValueError(f"private_key_pem is not a loadable unencrypted PEM private key: {exc}") from exc + if not isinstance(loaded, rsa.RSAPrivateKey): + raise ValueError("private_key_pem must be an unencrypted RSA private key in PEM format") # noqa: TRY004 # pydantic validators must raise ValueError + if loaded.key_size < _MIN_RSA_KEY_BITS: + raise ValueError(f"private_key_pem must be an RSA key of at least {_MIN_RSA_KEY_BITS} bits") + return value + + @model_validator(mode="after") + def _kids_are_unique(self) -> AsymmetricSessionKeys: + kids: Final = (self.kid, *(previous.kid for previous in self.previous_public_keys)) + duplicates: Final = tuple(kid for kid, count in Counter(kids).items() if count > 1) + if duplicates: + raise ValueError( + f"every kid must be unique across the current and previous keys; duplicated: {', '.join(duplicates)}" + ) + return self + + +SessionSigningKeys: TypeAlias = SessionKeys | AsymmetricSessionKeys +"""Every key material shape the mints and openers accept: the default master-key-derived +HS256 secret, or operator-configured RS256 RSA keys.""" + + +@lru_cache(maxsize=8) +def _public_key_pem_from_private(private_key_pem: str) -> str: + loaded: Final = serialization.load_pem_private_key(private_key_pem.encode(), password=None) + return ( + loaded.public_key() + .public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo) + .decode() + ) + + +def session_public_key_pem(keys: AsymmetricSessionKeys) -> str: + """The PEM public half of the current RS256 signing key: the only material a downstream + validator (an external gateway verifying ``kid``-matched tokens) ever needs.""" + return _public_key_pem_from_private(keys.private_key_pem.get_secret_value()) + + class MintedSessionToken(BaseModel): """A minted session token: the client-held bearer value and when it expires.""" @@ -221,7 +313,7 @@ def is_session_refresh_token(candidate: str) -> bool: def mint_session_token( principal: SessionPrincipal, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, ) -> MintedSessionToken | SessionTokenMintError: """Mint the short-lived session ACCESS token for ``principal``. @@ -241,7 +333,7 @@ def mint_session_token( def mint_session_refresh_token( principal: SessionPrincipal, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, ) -> MintedSessionToken | SessionTokenMintError: """Mint the long-lived session REFRESH token for ``principal``. @@ -262,7 +354,7 @@ def mint_session_refresh_token( def open_session_token( candidate: str, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, ) -> OpenedSessionToken | SessionTokenOpenError: """Validate a session ACCESS ``candidate`` and recover the principal. @@ -275,7 +367,7 @@ def open_session_token( def open_session_refresh_token( candidate: str, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, ) -> OpenedSessionToken | SessionTokenOpenError: """Validate a session REFRESH ``candidate`` and recover the principal. @@ -292,7 +384,7 @@ def _mint( prefix: str, principal: SessionPrincipal, expires_at: datetime, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, ) -> MintedSessionToken | SessionTokenTooLarge: """Sign the claims for either token kind and enforce the size cap. Shared by both mints @@ -309,20 +401,33 @@ def _mint( audience=principal.audience, team_id=principal.team_id, ) - token: Final = prefix + jwt.encode( - claims.model_dump(exclude_none=True), keys.signing_key.get_secret_value(), algorithm=_SESSION_JWT_ALGORITHM - ) + token: Final = prefix + _sign_claims(claims, keys) size_bytes: Final = len(token.encode("utf-8")) if size_bytes > MAX_SESSION_TOKEN_BYTES: return SessionTokenTooLarge(size_bytes=size_bytes, max_bytes=MAX_SESSION_TOKEN_BYTES) return MintedSessionToken(token=SecretStr(token), expires_at=expires_at) +def _sign_claims(claims: _SessionClaims, keys: SessionSigningKeys) -> str: + """Sign the claim set under whichever key material was injected: RS256 with the ``kid`` + in the JOSE header (so a validator can pick the right public key), or the default + HS256 secret with no header extras (byte-compatible with every pre-RS256 token).""" + payload: Final = claims.model_dump(exclude_none=True) + if isinstance(keys, AsymmetricSessionKeys): + return jwt.encode( + payload, + keys.private_key_pem.get_secret_value(), + algorithm=_SESSION_RSA_ALGORITHM, + headers={"kid": keys.kid}, + ) + return jwt.encode(payload, keys.signing_key.get_secret_value(), algorithm=_SESSION_JWT_ALGORITHM) + + def _open( candidate: str, prefix: str, expected_kind: SessionTokenKind, - keys: SessionKeys, + keys: SessionSigningKeys, now: datetime, ) -> OpenedSessionToken | SessionTokenOpenError: """Prefix-route, size-bound, signature-verify, kind-check, and expiry-check an @@ -337,7 +442,7 @@ def _open( return SessionMalformed() if len(candidate.encode("utf-8", "surrogatepass")) > MAX_SESSION_TOKEN_BYTES: return SessionMalformed() - claims: Final = _decode_claims(candidate.removeprefix(prefix), keys.signing_key) + claims: Final = _decode_claims(candidate.removeprefix(prefix), keys) if not isinstance(claims, _SessionClaims): return claims if claims.kind != expected_kind: @@ -356,14 +461,51 @@ def _open( ) +class _VerificationMaterial(BaseModel): + model_config = ConfigDict(frozen=True) + key: SecretStr + algorithm: Literal["HS256", "RS256"] + + +def _verification_material( + compact: str, + keys: SessionSigningKeys, +) -> _VerificationMaterial | SessionBadSignature | SessionMalformed: + """Pick the single key and algorithm the candidate is allowed to verify under. + + HS256 mode has exactly one secret. RS256 mode routes by the JOSE header ``kid``: the + current key's derived public half, or a retired key's stored public half during a + rotation window. An unknown or missing ``kid`` is ``SessionBadSignature`` (a foreign + key), and an undecodable header is ``SessionMalformed``. The algorithm is pinned per + key shape, never read from the header, so an HS256 token can never be verified + against a public key or vice versa. + """ + if isinstance(keys, SessionKeys): + return _VerificationMaterial(key=keys.signing_key, algorithm=_SESSION_JWT_ALGORITHM) + try: + header: Final = jwt.get_unverified_header(compact) + except jwt.InvalidTokenError: + return SessionMalformed() + kid: Final = header.get("kid") + if kid == keys.kid: + return _VerificationMaterial(key=SecretStr(session_public_key_pem(keys)), algorithm=_SESSION_RSA_ALGORITHM) + for previous in keys.previous_public_keys: + if previous.kid == kid: + return _VerificationMaterial(key=SecretStr(previous.public_key_pem), algorithm=_SESSION_RSA_ALGORITHM) + return SessionBadSignature() + + def _decode_claims( compact: str, - signing_key: SecretStr, + keys: SessionSigningKeys, ) -> _SessionClaims | SessionBadSignature | SessionMalformed: - """Verify the HS256 signature and shape of an attacker-controlled compact JWT. + """Verify the signature and shape of an attacker-controlled compact JWT. ``compact`` is fully hostile and bounded to ``MAX_SESSION_TOKEN_BYTES`` by the caller. - PyJWT's ``iat``/``nbf``/``exp`` validators are disabled: they raise on hostile claim + The accepted algorithm is pinned by :func:`_verification_material` from the injected + key shape, so ``alg`` confusion (``none``, or HS256 signed with a public key as the + secret) fails before or at signature verification. PyJWT's ``iat``/``nbf``/``exp`` + validators are disabled: they raise on hostile claim types and, for ``iat``/``nbf``, compare against the wall clock rather than the injected ``now`` (``exp`` is checked by the caller against ``now``). Apart from a signature mismatch, every decode failure is ``SessionMalformed``: a non-UTF-8 candidate surfaces @@ -371,11 +513,14 @@ def _decode_claims( ``TypeError`` from PyJWT's claim validators, and a wrong issuer or structurally invalid token as an ``InvalidTokenError``. ``_SessionClaims`` is the total type gate. """ + material: Final = _verification_material(compact, keys) + if not isinstance(material, _VerificationMaterial): + return material try: payload: Final = jwt.decode( compact, - signing_key.get_secret_value(), - algorithms=[_SESSION_JWT_ALGORITHM], + material.key.get_secret_value(), + algorithms=[material.algorithm], issuer=SESSION_ISSUER, options={ "verify_exp": False, diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index fa33a307438..6d9a907324d 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -1,10 +1,10 @@ import asyncio import hashlib import json -from collections.abc import Iterator, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from datetime import datetime, timezone from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypedDict +from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, TypedDict import litellm from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -114,6 +114,13 @@ def object_permission_table( return table +def _dump_agent_params(raw: Mapping[str, object]) -> dict[str, object]: + model_dump: Final[Callable[[], dict[str, object]] | None] = getattr(raw, "model_dump", None) + if model_dump is not None: + return model_dump() + return dict(raw) if raw else {} + + class GrantMigrationResult(NamedTuple): rewritten: int missed: int @@ -205,7 +212,7 @@ class AgentRegistry: def load_agents_from_db_and_config( self, agent_config: Sequence[AgentConfig] | None = None, - db_agents: list[dict[str, Any]] | None = None, + db_agents: Sequence[Mapping[str, object]] | None = None, ): """ Rebuild the registry from the DB rows plus the agents declared in config.yaml. @@ -227,7 +234,7 @@ class AgentRegistry: if not isinstance(db_agent, dict): raise ValueError("db_agents must be a list of dictionaries") - self.register_agent(agent_config=AgentResponse(**db_agent)) + self.register_agent(agent_config=AgentResponse.model_validate(db_agent)) self.load_agents_from_config(agent_config if agent_config is not None else self.config_agents) return self.agent_list @@ -295,19 +302,13 @@ class AgentRegistry: agent_name: Final = agent.get("agent_name") # Serialize litellm_params - litellm_params_obj: Final[Any] = agent.get("litellm_params", {}) - if hasattr(litellm_params_obj, "model_dump"): - litellm_params_dict = litellm_params_obj.model_dump() - else: - litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {} + litellm_params_obj: Final = agent.get("litellm_params", {}) + litellm_params_dict: Final[dict[str, object]] = _dump_agent_params(litellm_params_obj) litellm_params: Final[str] = safe_dumps(litellm_params_dict) # Serialize agent_card_params - agent_card_params_obj: Final[Any] = agent.get("agent_card_params", {}) - if hasattr(agent_card_params_obj, "model_dump"): - agent_card_params_dict = agent_card_params_obj.model_dump() - else: - agent_card_params_dict = dict(agent_card_params_obj) if agent_card_params_obj else {} + agent_card_params_obj: Final = agent.get("agent_card_params", {}) + agent_card_params_dict: Final[dict[str, object]] = _dump_agent_params(agent_card_params_obj) agent_card_params: Final[str] = safe_dumps(agent_card_params_dict) # Handle object_permission (MCP tool access for agent) @@ -408,7 +409,7 @@ class AgentRegistry: existing_agent: Final = dict(existing_row) augment_agent: Final = {**existing_agent, **agent} - update_data: Final[dict[str, Any]] = {} + update_data: Final[dict[str, object]] = {} if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") if augment_agent.get("litellm_params"): @@ -476,19 +477,13 @@ class AgentRegistry: agent_name: Final = agent.get("agent_name") # Serialize litellm_params - litellm_params_obj: Final[Any] = agent.get("litellm_params", {}) - if hasattr(litellm_params_obj, "model_dump"): - litellm_params_dict = litellm_params_obj.model_dump() - else: - litellm_params_dict = dict(litellm_params_obj) if litellm_params_obj else {} + litellm_params_obj: Final = agent.get("litellm_params", {}) + litellm_params_dict: Final[dict[str, object]] = _dump_agent_params(litellm_params_obj) litellm_params: Final[str] = safe_dumps(litellm_params_dict) # Serialize agent_card_params - agent_card_params_obj: Final[Any] = agent.get("agent_card_params", {}) - if hasattr(agent_card_params_obj, "model_dump"): - agent_card_params_dict = agent_card_params_obj.model_dump() - else: - agent_card_params_dict = dict(agent_card_params_obj) if agent_card_params_obj else {} + agent_card_params_obj: Final = agent.get("agent_card_params", {}) + agent_card_params_dict: Final[dict[str, object]] = _dump_agent_params(agent_card_params_obj) agent_card_params: Final[str] = safe_dumps(agent_card_params_dict) # Serialize static_headers for update diff --git a/litellm/proxy/client/cli/commands/_cli_context.py b/litellm/proxy/client/cli/commands/_cli_context.py new file mode 100644 index 00000000000..74c29653d16 --- /dev/null +++ b/litellm/proxy/client/cli/commands/_cli_context.py @@ -0,0 +1,19 @@ +from typing import Final + +import click +from typing_extensions import ReadOnly, TypedDict + + +class CliContextValues(TypedDict): + """Values the top-level CLI group stores on the click context.""" + + base_url: ReadOnly[str] + api_key: ReadOnly[str | None] + + +_UNSET_CLI_CONTEXT: Final[CliContextValues] = {"base_url": "", "api_key": None} + + +def cli_context_values(ctx: click.Context) -> CliContextValues: + values: Final[CliContextValues] = getattr(ctx, "obj", _UNSET_CLI_CONTEXT) + return values diff --git a/litellm/proxy/client/cli/commands/chat.py b/litellm/proxy/client/cli/commands/chat.py index c88d89dab2d..780695a37bb 100644 --- a/litellm/proxy/client/cli/commands/chat.py +++ b/litellm/proxy/client/cli/commands/chat.py @@ -1,6 +1,7 @@ import json import sys -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Final import click import requests @@ -8,15 +9,42 @@ from rich.console import Console from rich.panel import Panel from rich.prompt import Prompt from rich.table import Table +from typing_extensions import NotRequired, ReadOnly, TypedDict from ... import Client from ...chat import ChatClient +from ._cli_context import cli_context_values -def _get_available_models(ctx: click.Context) -> list[dict[str, Any]]: +class _MessagesView(TypedDict): + messages: ReadOnly[list[dict[str, str]]] + + +class _StreamDelta(TypedDict): + content: ReadOnly[NotRequired[str]] + + +class _StreamChoice(TypedDict): + delta: ReadOnly[NotRequired[_StreamDelta]] + + +class _StreamChunkView(TypedDict): + choices: ReadOnly[Sequence[_StreamChoice]] + + +class _StreamErrorBody(TypedDict): + error: ReadOnly[NotRequired[Mapping[str, object]]] + + +class _ErrorBodyView(TypedDict): + body: ReadOnly[_StreamErrorBody] + + +def _get_available_models(ctx: click.Context) -> Sequence[Mapping[str, object]]: """Get list of available models from the proxy server""" try: - client: Final = Client(base_url=ctx.obj["base_url"], api_key=ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + client: Final = Client(base_url=context["base_url"], api_key=context["api_key"]) models_list: Final = client.models.list() # Ensure we return a list of dictionaries if isinstance(models_list, list): @@ -28,7 +56,7 @@ def _get_available_models(ctx: click.Context) -> list[dict[str, Any]]: return [] -def _select_model(console: Console, available_models: list[dict[str, Any]]) -> str | None: +def _select_model(console: Console, available_models: Sequence[Mapping[str, object]]) -> str | None: """Interactive model selection""" if not available_models: console.print("[yellow]No models available or could not fetch models list.[/yellow]") @@ -42,7 +70,7 @@ def _select_model(console: Console, available_models: list[dict[str, Any]]) -> s table.add_column("Owned By", style="yellow") MAX_MODELS_TO_DISPLAY: Final = 200 - models_to_display: Final[list[dict[str, Any]]] = available_models[:MAX_MODELS_TO_DISPLAY] + models_to_display: Final = available_models[:MAX_MODELS_TO_DISPLAY] for i, model in enumerate(models_to_display): # Limit to first 200 models table.add_row(str(i + 1), str(model.get("id", "")), str(model.get("owned_by", ""))) @@ -62,7 +90,7 @@ def _select_model(console: Console, available_models: list[dict[str, Any]]) -> s try: index = int(choice) - 1 if 0 <= index < len(available_models): - return available_models[index]["id"] + return str(available_models[index]["id"]) else: console.print( f"[red]Invalid index. Please enter a number between 1 and {len(available_models)}[/red]" @@ -132,10 +160,11 @@ def chat( console.print("[red]No model selected. Exiting.[/red]") return - client: Final = ChatClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + client: Final = ChatClient(context["base_url"], context["api_key"]) # Initialize conversation history - messages: list[dict[str, Any]] = [] + messages: list[dict[str, str]] = [] # Add system message if provided if system: @@ -238,7 +267,7 @@ def _show_help(console: Console): console.print(Panel(help_text, title="Help")) -def _show_history(console: Console, messages: list[dict[str, Any]]): +def _show_history(console: Console, messages: list[dict[str, str]]): """Show conversation history""" if not messages: console.print("[yellow]No conversation history.[/yellow]") @@ -260,7 +289,7 @@ def _show_history(console: Console, messages: list[dict[str, Any]]): ) -def _save_conversation(console: Console, messages: list[dict[str, Any]], command: str): +def _save_conversation(console: Console, messages: list[dict[str, str]], command: str): """Save conversation to a file""" parts: Final = command.split() if len(parts) < 2: @@ -279,7 +308,7 @@ def _save_conversation(console: Console, messages: list[dict[str, Any]], command console.print(f"[red]Error saving conversation: {e}[/red]") -def _load_conversation(console: Console, command: str, system: str | None) -> list[dict[str, Any]]: +def _load_conversation(console: Console, command: str, system: str | None) -> list[dict[str, str]]: """Load conversation from a file""" parts: Final = command.split() if len(parts) < 2: @@ -292,9 +321,9 @@ def _load_conversation(console: Console, command: str, system: str | None) -> li try: with open(filename, "r") as f: - messages: Final = json.load(f) + loaded: Final[_MessagesView] = {"messages": json.load(f)} console.print(f"[green]Conversation loaded from {filename}[/green]") - return messages + return loaded["messages"] except FileNotFoundError: console.print(f"[red]File not found: {filename}[/red]") except Exception as e: @@ -309,10 +338,10 @@ def _load_conversation(console: Console, command: str, system: str | None) -> li def _handle_special_commands( console: Console, user_input: str, - messages: list[dict[str, Any]], + messages: list[dict[str, str]], system: str | None, ctx: click.Context, -) -> tuple[bool, list[dict[str, Any]], str | None]: +) -> tuple[bool, list[dict[str, str]], str | None]: """Handle special chat commands. Returns (should_exit, updated_messages, updated_model)""" if user_input.lower() in ["/quit", "/exit", "/q"]: console.print("[yellow]Chat session ended.[/yellow]") @@ -321,11 +350,9 @@ def _handle_special_commands( _show_help(console) return False, messages, None elif user_input.lower() == "/clear": - new_messages = [] - if system: - new_messages.append({"role": "system", "content": system}) + cleared_messages: Final[list[dict[str, str]]] = [{"role": "system", "content": system}] if system else [] console.print("[green]Conversation history cleared.[/green]") - return False, new_messages, None + return False, cleared_messages, None elif user_input.lower() == "/history": _show_history(console, messages) return False, messages, None @@ -353,7 +380,7 @@ def _stream_response( console: Console, client: ChatClient, model: str, - messages: list[dict[str, Any]], + messages: list[dict[str, str]], temperature: float, max_tokens: int | None, ) -> str | None: @@ -366,8 +393,9 @@ def _stream_response( temperature=temperature, max_tokens=max_tokens, ): - if "choices" in chunk and len(chunk["choices"]) > 0: - delta = chunk["choices"][0].get("delta", {}) + streamed: _StreamChunkView = {"choices": chunk.get("choices", ())} + if len(streamed["choices"]) > 0: + delta = streamed["choices"][0].get("delta", {}) content = delta.get("content", "") if content: assistant_content += content @@ -380,8 +408,8 @@ def _stream_response( except requests.exceptions.HTTPError as e: console.print(f"\n[red]Error: HTTP {e.response.status_code}[/red]") try: - error_body: Final = e.response.json() - console.print(f"[red]{error_body.get('error', {}).get('message', 'Unknown error')}[/red]") + error_body: Final[_ErrorBodyView] = {"body": e.response.json()} + console.print(f"[red]{error_body['body'].get('error', {}).get('message', 'Unknown error')}[/red]") except json.JSONDecodeError: console.print(f"[red]{e.response.text}[/red]") return None diff --git a/litellm/proxy/client/cli/commands/credentials.py b/litellm/proxy/client/cli/commands/credentials.py index c550b39d33f..2c4080dbeb2 100644 --- a/litellm/proxy/client/cli/commands/credentials.py +++ b/litellm/proxy/client/cli/commands/credentials.py @@ -1,12 +1,36 @@ import json +from collections.abc import Sequence from typing import Final, Literal import click import requests import rich from rich.table import Table +from typing_extensions import NotRequired, ReadOnly, TypedDict from ...credentials import CredentialsManagementClient +from ._cli_context import cli_context_values + + +class _CredentialInfo(TypedDict): + custom_llm_provider: ReadOnly[NotRequired[str]] + + +class _CredentialItem(TypedDict): + credential_name: ReadOnly[NotRequired[str]] + credential_info: ReadOnly[NotRequired[_CredentialInfo]] + + +class _CredentialsListView(TypedDict): + credentials: ReadOnly[Sequence[_CredentialItem]] + + +class _JsonObjectView(TypedDict): + value: ReadOnly[dict[str, object]] + + +class _JsonBodyView(TypedDict): + body: ReadOnly[object] @click.group() @@ -25,7 +49,8 @@ def credentials(): @click.pass_context def list(ctx: click.Context, output_format: Literal["table", "json"]): """List all credentials""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + client: Final = CredentialsManagementClient(context["base_url"], context["api_key"]) response: Final = client.list() assert isinstance(response, dict) @@ -39,7 +64,8 @@ def list(ctx: click.Context, output_format: Literal["table", "json"]): table.add_column("Custom LLM Provider", style="green") # Add rows - for cred in response.get("credentials", []): + listed: Final[_CredentialsListView] = {"credentials": response.get("credentials", [])} + for cred in listed["credentials"]: info = cred.get("credential_info", {}) table.add_row( str(cred.get("credential_name", "")), @@ -66,21 +92,22 @@ def list(ctx: click.Context, output_format: Literal["table", "json"]): @click.pass_context def create(ctx: click.Context, credential_name: str, info: str, values: str): """Create a new credential""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + client: Final = CredentialsManagementClient(context["base_url"], context["api_key"]) try: - credential_info: Final = json.loads(info) - credential_values: Final = json.loads(values) + credential_info: Final[_JsonObjectView] = {"value": json.loads(info)} + credential_values: Final[_JsonObjectView] = {"value": json.loads(values)} except json.JSONDecodeError as e: raise click.BadParameter(f"Invalid JSON: {e}") try: - response: Final = client.create(credential_name, credential_info, credential_values) + response: Final = client.create(credential_name, credential_info["value"], credential_values["value"]) rich.print_json(data=response) except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) try: - error_body: Final = e.response.json() - rich.print_json(data=error_body) + error_body: Final[_JsonBodyView] = {"body": e.response.json()} + rich.print_json(data=error_body["body"]) except json.JSONDecodeError: click.echo(e.response.text, err=True) raise click.Abort() @@ -91,15 +118,16 @@ def create(ctx: click.Context, credential_name: str, info: str, values: str): @click.pass_context def delete(ctx: click.Context, credential_name: str): """Delete a credential by name""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + client: Final = CredentialsManagementClient(context["base_url"], context["api_key"]) try: response: Final = client.delete(credential_name) rich.print_json(data=response) except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) try: - error_body: Final = e.response.json() - rich.print_json(data=error_body) + error_body: Final[_JsonBodyView] = {"body": e.response.json()} + rich.print_json(data=error_body["body"]) except json.JSONDecodeError: click.echo(e.response.text, err=True) raise click.Abort() @@ -110,6 +138,7 @@ def delete(ctx: click.Context, credential_name: str): @click.pass_context def get(ctx: click.Context, credential_name: str): """Get a credential by name""" - client: Final = CredentialsManagementClient(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + client: Final = CredentialsManagementClient(context["base_url"], context["api_key"]) response: Final = client.get(credential_name) rich.print_json(data=response) diff --git a/litellm/proxy/client/cli/commands/teams.py b/litellm/proxy/client/cli/commands/teams.py index e814ac84ebb..1f91d5559d8 100644 --- a/litellm/proxy/client/cli/commands/teams.py +++ b/litellm/proxy/client/cli/commands/teams.py @@ -1,14 +1,33 @@ """Team management commands for LiteLLM CLI.""" +from collections.abc import Mapping, Sequence from typing import Any, Final import click import requests from rich.console import Console from rich.table import Table +from typing_extensions import ReadOnly, TypedDict from litellm.proxy.client import Client +from ._cli_context import cli_context_values + + +class _TeamRow(TypedDict): + team_alias: ReadOnly[str | None] + team_id: ReadOnly[str | None] + models: ReadOnly[Sequence[str]] + max_budget: ReadOnly[object] + + +class _TeamModelsView(TypedDict): + models: ReadOnly[Sequence[str]] + + +class _ErrorBodyView(TypedDict): + body: ReadOnly[Mapping[str, object]] + @click.group() def teams(): @@ -32,10 +51,14 @@ def display_teams_table(teams: list[dict[str, Any]]) -> None: table.add_column("Role", style="red") for i, team in enumerate(teams): - team_alias = team.get("team_alias") or "N/A" - team_id = team.get("team_id", "N/A") - models = team.get("models", []) - max_budget = team.get("max_budget") + row: _TeamRow = { + "team_alias": team.get("team_alias") or "N/A", + "team_id": team.get("team_id", "N/A"), + "models": team.get("models", []), + "max_budget": team.get("max_budget"), + } + models = row["models"] + max_budget = row["max_budget"] # Format models list if models: @@ -55,7 +78,7 @@ def display_teams_table(teams: list[dict[str, Any]]) -> None: # This would need to be implemented based on actual API response structure pass - table.add_row(str(i + 1), team_alias, team_id, models_str, budget_str, role) + table.add_row(str(i + 1), row["team_alias"], row["team_id"], models_str, budget_str, role) console.print(table) @@ -64,7 +87,8 @@ def display_teams_table(teams: list[dict[str, Any]]) -> None: @click.pass_context def list(ctx: click.Context): """List teams that you belong to""" - client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + client: Final = Client(context["base_url"], context["api_key"]) try: # Use list() for simpler response structure (returns array directly) @@ -72,8 +96,8 @@ def list(ctx: click.Context): display_teams_table(teams) except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) - error_body: Final = e.response.json() - click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) + error_body: Final[_ErrorBodyView] = {"body": e.response.json()} + click.echo(f"Details: {error_body['body'].get('detail', 'Unknown error')}", err=True) raise click.Abort() except Exception as e: click.echo(f"Error: {e}", err=True) @@ -84,7 +108,8 @@ def list(ctx: click.Context): @click.pass_context def available(ctx: click.Context): """List teams that are available to join""" - client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"]) + context: Final = cli_context_values(ctx) + client: Final = Client(context["base_url"], context["api_key"]) try: teams: Final = client.teams.get_available() @@ -96,8 +121,8 @@ def available(ctx: click.Context): click.echo("No available teams to join.") except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) - error_body: Final = e.response.json() - click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) + error_body: Final[_ErrorBodyView] = {"body": e.response.json()} + click.echo(f"Details: {error_body['body'].get('detail', 'Unknown error')}", err=True) except Exception as e: click.echo(f"Error: {e}", err=True) raise click.Abort() @@ -108,8 +133,9 @@ def available(ctx: click.Context): @click.pass_context def assign_key(ctx: click.Context, team_id: str | None): """Assign your current CLI key to a team""" - client: Final = Client(ctx.obj["base_url"], ctx.obj["api_key"]) - api_key: Final = ctx.obj["api_key"] + context: Final = cli_context_values(ctx) + client: Final = Client(context["base_url"], context["api_key"]) + api_key: Final = context["api_key"] if not api_key: click.echo("No API key found. Please login first using 'litellm login'") @@ -145,17 +171,17 @@ def assign_key(ctx: click.Context, team_id: str | None): teams = client.teams.list() for team in teams: if team.get("team_id") == team_id: - models = team.get("models", []) - if models: - click.echo(f"You can now access models: {', '.join(models)}") + team_models: _TeamModelsView = {"models": team.get("models", [])} + if team_models["models"]: + click.echo(f"You can now access models: {', '.join(team_models['models'])}") else: click.echo("You can now access all available models") break except requests.exceptions.HTTPError as e: click.echo(f"Error: HTTP {e.response.status_code}", err=True) - error_body: Final = e.response.json() - click.echo(f"Details: {error_body.get('detail', 'Unknown error')}", err=True) + error_body: Final[_ErrorBodyView] = {"body": e.response.json()} + click.echo(f"Details: {error_body['body'].get('detail', 'Unknown error')}", err=True) raise click.Abort() except Exception as e: click.echo(f"Error: {e}", err=True) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index cbd50da9c3e..ff6c8d1b1f8 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -4,7 +4,7 @@ import json import logging import math import traceback -from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType @@ -2372,6 +2372,21 @@ class ProxyBaseLLMRequestProcessing: ) logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete + elif ( + _post_call_guardrails_active + and route_type == "anthropic_messages" + and self._is_streaming_response(response) + ): + from litellm.litellm_core_utils.logging_worker import ( + GLOBAL_LOGGING_WORKER, + ) + + async def _on_deferred_native_stream_complete( + logging_coroutine: Coroutine[object, object, object], + ) -> None: + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) + + logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete if route_type == "allm_passthrough_route": # Check if response is an async generator diff --git a/litellm/proxy/common_utils/get_routes.py b/litellm/proxy/common_utils/get_routes.py index 2118a6610b4..28e58808c8a 100644 --- a/litellm/proxy/common_utils/get_routes.py +++ b/litellm/proxy/common_utils/get_routes.py @@ -2,71 +2,87 @@ Utility class for getting routes from a FastAPI app. """ -from typing import Any, Final +from collections.abc import Sequence +from typing import Final, Protocol from starlette.routing import BaseRoute +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_logger +class NamedEndpoint(Protocol): + __name__: str + + +class RouteInfo(TypedDict): + path: ReadOnly[str | None] + methods: ReadOnly[Sequence[str] | None] + name: ReadOnly[str | None] + endpoint: ReadOnly[str | None] + mounted_app: NotRequired[ReadOnly[bool]] + + class GetRoutes: @staticmethod def get_app_routes( route: BaseRoute, - endpoint_route: Any, - ) -> list[dict[str, Any]]: + endpoint_route: NamedEndpoint, + ) -> list[RouteInfo]: """ Get routes for a regular route. """ - routes: Final[list[dict[str, Any]]] = [] - route_info: Final = { + route_info: Final[RouteInfo] = { "path": getattr(route, "path", None), "methods": getattr(route, "methods", None), "name": getattr(route, "name", None), "endpoint": (endpoint_route.__name__ if getattr(route, "endpoint", None) else None), } - routes.append(route_info) - return routes + return [route_info] @staticmethod def get_routes_for_mounted_app( route: BaseRoute, - ) -> list[dict[str, Any]]: + ) -> list[RouteInfo]: """ Get routes for a mounted sub-application. """ - routes: Final[list[dict[str, Any]]] = [] - mount_path: Final = getattr(route, "path", "") - sub_app: Final = getattr(route, "app", None) - if sub_app and hasattr(sub_app, "routes"): - for sub_route in sub_app.routes: - # Get endpoint - either from endpoint attribute or app attribute - endpoint_func = getattr(sub_route, "endpoint", None) or getattr(sub_route, "app", None) - - if endpoint_func is not None: - sub_route_path = getattr(sub_route, "path", "") - full_path = mount_path.rstrip("/") + sub_route_path - - route_info = { - "path": full_path, - "methods": getattr(sub_route, "methods", ["GET", "POST"]), - "name": getattr(sub_route, "name", None), - "endpoint": GetRoutes._safe_get_endpoint_name(endpoint_func), - "mounted_app": True, - } - routes.append(route_info) - return routes + mount_path: Final[str] = getattr(route, "path", "") + sub_app: Final[object] = getattr(route, "app", None) + if not sub_app or not hasattr(sub_app, "routes"): + return [] + sub_routes: Final[Sequence[object]] = getattr(sub_app, "routes", ()) + return [ + sub_route_info + for sub_route in sub_routes + if (sub_route_info := GetRoutes._mounted_sub_route_info(mount_path, sub_route)) is not None + ] @staticmethod - def _safe_get_endpoint_name(endpoint_function: Any) -> str | None: + def _mounted_sub_route_info(mount_path: str, sub_route: object) -> RouteInfo | None: + endpoint_func: Final[object] = getattr(sub_route, "endpoint", None) or getattr(sub_route, "app", None) + if endpoint_func is None: + return None + sub_route_path: Final[str] = getattr(sub_route, "path", "") + return { + "path": mount_path.rstrip("/") + sub_route_path, + "methods": getattr(sub_route, "methods", ["GET", "POST"]), + "name": getattr(sub_route, "name", None), + "endpoint": GetRoutes._safe_get_endpoint_name(endpoint_func), + "mounted_app": True, + } + + @staticmethod + def _safe_get_endpoint_name(endpoint_function: object) -> str | None: """ Safely get the name of the endpoint function. """ try: if hasattr(endpoint_function, "__name__"): - return getattr(endpoint_function, "__name__") + endpoint_name: Final[str] = getattr(endpoint_function, "__name__", "") + return endpoint_name elif hasattr(endpoint_function, "__class__") and hasattr(endpoint_function.__class__, "__name__"): - return getattr(endpoint_function.__class__, "__name__") + return endpoint_function.__class__.__name__ else: return None except Exception: diff --git a/litellm/proxy/db/create_views.py b/litellm/proxy/db/create_views.py index 5ea9cba8018..10daeee4e7b 100644 --- a/litellm/proxy/db/create_views.py +++ b/litellm/proxy/db/create_views.py @@ -1,21 +1,21 @@ -from typing import Any, Final, Protocol +from collections.abc import Mapping, Sequence +from typing import Final, Protocol from litellm import verbose_logger -_db = Any - class SupportsExecuteRaw(Protocol): - """The one database operation create_view_tolerating_race needs. - - Narrower than the `_db = Any` the rest of this module still uses, so the - helper's contract is checkable at its call sites without retyping every - function here. - """ + """The one database operation create_view_tolerating_race needs.""" async def execute_raw(self, query: str, *args: object) -> int: ... +class SupportsRawQueries(SupportsExecuteRaw, Protocol): + """The database operations the view bootstrap needs: probe a relation, then create it.""" + + async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ... + + # Markers that indicate a view/relation does not yet exist in the database. # Keeping these in one place avoids repeating the check across all view blocks # and prevents overly broad matches (e.g. bare 'undefined' would also match @@ -46,7 +46,7 @@ async def create_view_tolerating_race(db: SupportsExecuteRaw, view_name: str, dd verbose_logger.debug("%s already created by a concurrent replica", view_name) -async def create_missing_views(db: _db): +async def create_missing_views(db: SupportsRawQueries) -> None: """ -------------------------------------------------- NOTE: Copy of `litellm/db_scripts/create_views.py`. @@ -246,7 +246,7 @@ async def create_missing_views(db: _db): await create_view_tolerating_race(db, "Last30dTopEndUsersSpend", sql_query) -async def should_create_missing_views(db: _db) -> bool: +async def should_create_missing_views(db: SupportsRawQueries) -> bool: """ Run only on first time startup. diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 853c033c37e..ad92902221a 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -6,8 +6,9 @@ This is to prevent deadlocks and improve reliability import asyncio import json -from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, cast +from collections.abc import Mapping, Sequence +from functools import reduce +from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast from redis.exceptions import RedisError @@ -53,6 +54,52 @@ if TYPE_CHECKING: else: PrismaClient = Any +BufferedSpendTransactions: TypeAlias = DBSpendUpdateTransactions | Mapping[str, BaseDailySpendTransaction] + +_SpendTransactionField: TypeAlias = Literal[ + "user_list_transactions", + "end_user_list_transactions", + "key_list_transactions", + "team_list_transactions", + "team_member_list_transactions", + "org_list_transactions", + "tag_list_transactions", + "agent_list_transactions", +] + +_SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = ( + "user_list_transactions", + "end_user_list_transactions", + "key_list_transactions", + "team_list_transactions", + "team_member_list_transactions", + "org_list_transactions", + "tag_list_transactions", + "agent_list_transactions", +) + +_ValueT = TypeVar("_ValueT") + + +def _accumulated_spend(totals: Mapping[str, float], entities: Mapping[str, float]) -> dict[str, float]: + return {**totals, **{entity_id: totals.get(entity_id, 0) + amount for entity_id, amount in entities.items()}} + + +def _entity_transactions(transaction: DBSpendUpdateTransactions, field: _SpendTransactionField) -> dict[str, float]: + entities: Final[dict[str, float] | None] = transaction.get(field) + return entities if isinstance(entities, dict) else {} + + +def _merged_entity_transactions( + list_of_transactions: Sequence[DBSpendUpdateTransactions], + field: _SpendTransactionField, +) -> dict[str, float]: + return reduce( + _accumulated_spend, + (_entity_transactions(transaction, field) for transaction in list_of_transactions), + {}, + ) + class RedisUpdateBuffer: """ @@ -86,7 +133,7 @@ class RedisUpdateBuffer: async def _store_transactions_in_redis( self, - transactions: Any, + transactions: Mapping[str, BaseDailySpendTransaction] | None, redis_key: str, service_type: ServiceTypes, ) -> None: @@ -183,7 +230,9 @@ class RedisUpdateBuffer: return # Get all transactions - db_spend_update_transactions = await spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + db_spend_update_transactions: Final = ( + await spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions() + ) daily_spend_update_transactions: Final = ( await daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) @@ -204,7 +253,7 @@ class RedisUpdateBuffer: verbose_proxy_logger.debug("ALL DAILY SPEND UPDATE TRANSACTIONS: %s", daily_spend_update_transactions) # Build a list of rpush operations, skipping empty/None transaction sets - _queue_configs: Final[list[tuple[Any, str, ServiceTypes]]] = [ + _queue_configs: Final[list[tuple[BufferedSpendTransactions | None, str, ServiceTypes]]] = [ ( db_spend_update_transactions, REDIS_UPDATE_BUFFER_KEY, @@ -435,14 +484,12 @@ class RedisUpdateBuffer: """ Gets the number of transactions to store in Redis """ - num_transactions = 0 - for v in db_spend_update_transactions.values(): - if isinstance(v, dict): - num_transactions += len(v) - return num_transactions + return sum( + len(_entity_transactions(db_spend_update_transactions, field)) for field in _SPEND_TRANSACTION_FIELDS + ) @staticmethod - def _remove_prefix_from_keys(data: dict[str, Any], prefix: str) -> dict[str, Any]: + def _remove_prefix_from_keys(data: Mapping[str, _ValueT], prefix: str) -> dict[str, _ValueT]: """ Removes the specified prefix from the keys of a dictionary. """ @@ -489,7 +536,7 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_transactions: Final = await self.redis_cache.async_lpop( + list_of_transactions: Final[str | list[str] | None] = await self.redis_cache.async_lpop( key=REDIS_UPDATE_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, ) @@ -577,7 +624,7 @@ class RedisUpdateBuffer: db_spend = self._combine_list_of_transactions(parsed) # Slots 1-5: daily spend categories - daily_results: Final[list[dict[str, Any] | None]] = [] + daily_results: Final[list[dict[str, BaseDailySpendTransaction] | None]] = [] for slot in range(1, 6): slot_result = raw_results[slot] if slot_result is None: @@ -612,6 +659,23 @@ class RedisUpdateBuffer: service_type=ServiceTypes.REDIS_DAILY_TAG_SPEND_UPDATE_QUEUE, ) + async def _lpop_daily_spend_transactions( + self, + redis_key: str, + ) -> list[dict[str, BaseDailySpendTransaction]] | None: + """ + Drains a daily spend buffer key and parses each popped item as JSON. + """ + if self.redis_cache is None: + return None + list_of_transactions: Final[list[str] | None] = await self.redis_cache.async_lpop( + key=redis_key, + count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + ) + if list_of_transactions is None: + return None + return [json.loads(transaction) for transaction in list_of_transactions] + async def get_all_daily_spend_update_transactions_from_redis_buffer( self, ) -> dict[str, DailyUserSpendTransaction] | None: @@ -620,13 +684,11 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_transactions: Final = await self.redis_cache.async_lpop( - key=REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY ) - if list_of_transactions is None: + if list_of_daily_spend_update_transactions is None: return None - list_of_daily_spend_update_transactions = [json.loads(transaction) for transaction in list_of_transactions] return cast( dict[str, DailyUserSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( @@ -642,13 +704,11 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_transactions: Final = await self.redis_cache.async_lpop( - key=REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( + REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY ) - if list_of_transactions is None: + if list_of_daily_spend_update_transactions is None: return None - list_of_daily_spend_update_transactions = [json.loads(transaction) for transaction in list_of_transactions] return cast( dict[str, DailyTeamSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( @@ -664,13 +724,11 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_transactions: Final = await self.redis_cache.async_lpop( - key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( + REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY ) - if list_of_transactions is None: + if list_of_daily_spend_update_transactions is None: return None - list_of_daily_spend_update_transactions = [json.loads(transaction) for transaction in list_of_transactions] return cast( dict[str, DailyOrganizationSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( @@ -686,13 +744,11 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_transactions: Final = await self.redis_cache.async_lpop( - key=REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( + REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY ) - if list_of_transactions is None: + if list_of_daily_spend_update_transactions is None: return None - list_of_daily_spend_update_transactions = [json.loads(transaction) for transaction in list_of_transactions] return cast( dict[str, DailyEndUserSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( @@ -708,13 +764,11 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_transactions: Final = await self.redis_cache.async_lpop( - key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( + REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY ) - if list_of_transactions is None: + if list_of_daily_spend_update_transactions is None: return None - list_of_daily_spend_update_transactions = [json.loads(transaction) for transaction in list_of_transactions] return cast( dict[str, DailyAgentSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( @@ -730,13 +784,11 @@ class RedisUpdateBuffer: """ if self.redis_cache is None: return None - list_of_transactions: Final = await self.redis_cache.async_lpop( - key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, - count=MAX_REDIS_BUFFER_DEQUEUE_COUNT, + list_of_daily_spend_update_transactions: Final = await self._lpop_daily_spend_transactions( + REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY ) - if list_of_transactions is None: + if list_of_daily_spend_update_transactions is None: return None - list_of_daily_spend_update_transactions = [json.loads(transaction) for transaction in list_of_transactions] return cast( dict[str, DailyTagSpendTransaction], DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions( @@ -746,7 +798,7 @@ class RedisUpdateBuffer: @staticmethod def _parse_list_of_transactions( - list_of_transactions: Any | list[Any], + list_of_transactions: str | list[str], ) -> list[DBSpendUpdateTransactions]: """ Parses the list of transactions from Redis @@ -763,40 +815,19 @@ class RedisUpdateBuffer: """ Combines the list of transactions into a single DBSpendUpdateTransactions object """ - # Initialize a new combined transaction object with empty dictionaries - combined_transaction: Final = DBSpendUpdateTransactions( - user_list_transactions={}, - end_user_list_transactions={}, - key_list_transactions={}, - team_list_transactions={}, - team_member_list_transactions={}, - org_list_transactions={}, - tag_list_transactions={}, - agent_list_transactions={}, + return DBSpendUpdateTransactions( + user_list_transactions=_merged_entity_transactions(list_of_transactions, "user_list_transactions"), + end_user_list_transactions=_merged_entity_transactions(list_of_transactions, "end_user_list_transactions"), + key_list_transactions=_merged_entity_transactions(list_of_transactions, "key_list_transactions"), + team_list_transactions=_merged_entity_transactions(list_of_transactions, "team_list_transactions"), + team_member_list_transactions=_merged_entity_transactions( + list_of_transactions, "team_member_list_transactions" + ), + org_list_transactions=_merged_entity_transactions(list_of_transactions, "org_list_transactions"), + tag_list_transactions=_merged_entity_transactions(list_of_transactions, "tag_list_transactions"), + agent_list_transactions=_merged_entity_transactions(list_of_transactions, "agent_list_transactions"), ) - # Define the transaction fields to process - transaction_fields: Final = [ - "user_list_transactions", - "end_user_list_transactions", - "key_list_transactions", - "team_list_transactions", - "team_member_list_transactions", - "org_list_transactions", - "tag_list_transactions", - "agent_list_transactions", - ] - - # Loop through each transaction and combine the values - for transaction in list_of_transactions: - # Process each field type - for field in transaction_fields: - if transaction.get(field): - for entity_id, amount in transaction[field].items(): - combined_transaction[field][entity_id] = combined_transaction[field].get(entity_id, 0) + amount - - return combined_transaction - async def _emit_new_item_added_to_redis_buffer_event( self, service: ServiceTypes, diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 22fc32a898a..1929e7d3fc8 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -6,11 +6,15 @@ otherwise PrismaClient uses the writer-only PrismaWrapper directly. import os from collections.abc import Callable -from typing import Any, Final +from datetime import timedelta +from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger from litellm.proxy.db.prisma_client import PrismaWrapper +if TYPE_CHECKING: + from prisma.types import HttpConfig + # Per-model action methods that read from the database. These are routed to # the read replica when one is configured. _MODEL_READ_METHODS: Final = frozenset( @@ -43,15 +47,15 @@ class _RoutedActions: def __init__( self, - writer_actions: Any, - reader_actions: Any, + writer_actions: object, + reader_actions: object, should_use_reader: Callable[[], bool], ): self._writer_actions = writer_actions self._reader_actions = reader_actions self._should_use_reader = should_use_reader - def __getattr__(self, name: str) -> Any: + def __getattr__(self, name: str) -> object: if name in _MODEL_READ_METHODS and self._should_use_reader(): return getattr(self._reader_actions, name) return getattr(self._writer_actions, name) @@ -135,21 +139,21 @@ class RoutingPrismaWrapper: return not self._reader_unavailable @staticmethod - async def _try_connect(client: PrismaWrapper, *args: Any, **kwargs: Any) -> Exception | None: + async def _try_connect(client: PrismaWrapper, timeout: int | timedelta | None = None) -> Exception | None: if client.is_connected() is True: return None try: - await client.connect(*args, **kwargs) + await client.connect(timeout) return None except Exception as e: return e - async def connect(self, *args: Any, **kwargs: Any) -> None: - writer_error: Final = await self._try_connect(self._writer, *args, **kwargs) + async def connect(self, timeout: int | timedelta | None = None) -> None: + writer_error: Final = await self._try_connect(self._writer, timeout) if writer_error is None: self._writer_unavailable = False verbose_proxy_logger.info("[writer] DB connected") - reader_error: Final = await self._try_connect(self._reader, *args, **kwargs) + reader_error: Final = await self._try_connect(self._reader, timeout) if reader_error is None: self._reader_unavailable = False verbose_proxy_logger.info("[reader] DB connected") @@ -176,11 +180,11 @@ class RoutingPrismaWrapper: writer_error, ) - async def disconnect(self, *args: Any, **kwargs: Any) -> None: + async def disconnect(self, timeout: float | timedelta | None = None) -> None: first_error: BaseException | None = None for client in (self._writer, self._reader): try: - await client.disconnect(*args, **kwargs) + await client.disconnect(timeout) except Exception as e: if first_error is None: first_error = e @@ -206,7 +210,7 @@ class RoutingPrismaWrapper: async def recreate_prisma_client( self, new_db_url: str, - http_client: Any | None = None, + http_client: "HttpConfig | None" = None, *, expected_generation: int | None = None, ) -> bool: @@ -245,7 +249,7 @@ class RoutingPrismaWrapper: ) return True - async def _recreate_reader(self, http_client: Any | None = None) -> None: + async def _recreate_reader(self, http_client: "HttpConfig | None" = None) -> None: """Resolve the reader URL and recreate its Prisma client. Token-authenticated readers regenerate their token (host/port/user came @@ -266,13 +270,13 @@ class RoutingPrismaWrapper: def __getattr__(self, name: str) -> Any: if name in _TOP_LEVEL_READ_METHODS: return getattr(self.read_target, name) - writer_attr: Final = getattr(self._writer, name) + writer_attr: Final[object] = getattr(self._writer, name) # Per-model action accessors are non-callable instances that expose # both `find_many` and `create`. Methods like execute_raw / batch_ / # tx are callables and stay on the writer untouched. if not callable(writer_attr) and hasattr(writer_attr, "find_many") and hasattr(writer_attr, "create"): try: - reader_attr: Final = getattr(self._reader, name) + reader_attr: Final[object] = getattr(self._reader, name) except AttributeError: return writer_attr return _RoutedActions(writer_attr, reader_attr, self._should_use_reader) diff --git a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py index 1ca4652b9f9..bf2aa1f76e0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py +++ b/litellm/proxy/guardrails/guardrail_hooks/block_code_execution/block_code_execution.py @@ -41,7 +41,9 @@ LANGUAGE_ALIASES: Final[dict[str, str]] = { } # Tags that indicate non-executable / plain text (lower confidence when block-all) -NON_EXECUTABLE_TAGS: Final[frozenset] = frozenset({"text", "plaintext", "plain", "markdown", "md", "output", "result"}) +NON_EXECUTABLE_TAGS: Final[frozenset[str]] = frozenset( + {"text", "plaintext", "plain", "markdown", "md", "output", "result"} +) # Regex: fenced code block with optional language tag. Handles ```lang\n...\n``` # Content between fences; does not handle nested ``` inside body (documented edge case). @@ -486,7 +488,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): new_text: Final = "".join(parts) return new_text, should_raise - def _raise_block_error(self, language: str, is_output: bool, request_data: dict) -> None: + def _raise_block_error(self, language: str, is_output: bool, request_data: dict[str, object]) -> None: if language == "execution_request": msg = "Content blocked: execution request detected" else: @@ -510,7 +512,7 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: @@ -551,15 +553,16 @@ class BlockCodeExecutionGuardrail(CustomGuardrail): exception_str = str(e) raise finally: - guardrail_response: list[dict] | str = [dict(d) for d in detections] - if status != "success" and not detections: - guardrail_response = exception_str + detection_dicts: Final[list[dict[str, object]]] = [dict(d) for d in detections] + guardrail_response: Final[list[dict[str, object]] | str] = ( + exception_str if status != "success" and not detections else detection_dicts + ) max_confidence: float | None = None for d in detections: c = d.get("confidence") if c is not None and (max_confidence is None or c > max_confidence): max_confidence = c - tracing_kw: Final[dict[str, Any]] = { + tracing_kw: Final[GuardrailTracingDetail] = { "guardrail_id": self.guardrail_name, "detection_method": "fenced_code_block", "match_details": guardrail_response, diff --git a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py index 37e4c72bf96..5c14d03f50e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/compresr/compresr.py @@ -1071,7 +1071,7 @@ class CompresrGuardrail(CustomGuardrail): response: Any, anthropic_messages_provider_config: Any, anthropic_messages_optional_request_params: dict, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj | None, stream: bool, kwargs: dict, ) -> AgenticLoopPlan: diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py index 1eb6d2d1bb7..830dec8d80d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py @@ -36,6 +36,7 @@ Example: block when response rejects the user (input_type response only): import asyncio import threading +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast from fastapi import HTTPException @@ -59,9 +60,9 @@ if TYPE_CHECKING: class CustomCodeGuardrailError(Exception): """Raised when custom code guardrail execution fails.""" - def __init__(self, message: str, details: dict[str, Any] | None = None) -> None: + def __init__(self, message: str, details: Mapping[str, object] | None = None) -> None: super().__init__(message) - self.details = details or {} + self.details: Mapping[str, object] = details or {} class CustomCodeCompilationError(CustomCodeGuardrailError): @@ -116,8 +117,8 @@ class CustomCodeGuardrail(CustomGuardrail): guardrail_name: Name of this guardrail instance **kwargs: Additional arguments passed to CustomGuardrail """ - self.custom_code = custom_code - self._compiled_function: Any | None = None + self.custom_code: str = custom_code + self._compiled_function: Callable[..., object] | None = None self._compile_lock = threading.Lock() self._compile_error: str | None = None @@ -191,7 +192,7 @@ class CustomCodeGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: @@ -233,15 +234,14 @@ class CustomCodeGuardrail(CustomGuardrail): safe_request_data: Final = self._prepare_safe_request_data(request_data) # Execute the custom function - handle both sync and async functions - result = self._compiled_function(inputs, safe_request_data, input_type) + raw_result: Final = self._compiled_function(inputs, safe_request_data, input_type) # If the function is async (returns a coroutine), await it - if asyncio.iscoroutine(result): - result = await result + resolved_result: Final[object] = await raw_result if asyncio.iscoroutine(raw_result) else raw_result # Process the result return self._process_result( - result=result, + result=resolved_result, inputs=inputs, request_data=request_data, input_type=input_type, @@ -263,7 +263,7 @@ class CustomCodeGuardrail(CustomGuardrail): }, ) from e - def _prepare_safe_request_data(self, request_data: dict) -> dict[str, Any]: + def _prepare_safe_request_data(self, request_data: Mapping[str, object]) -> dict[str, object]: """ Prepare a safe subset of request_data for code execution. @@ -286,9 +286,9 @@ class CustomCodeGuardrail(CustomGuardrail): def _process_result( self, - result: Any, + result: object, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], ) -> GenericGuardrailAPIInputs: """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py index da12222f233..35f1e6e6515 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/sandbox.py @@ -14,8 +14,11 @@ We subclass it to permit those specific nodes, while keeping every other restriction intact. """ +import ast import operator -from typing import Any, Final +from collections.abc import Callable, Mapping +from types import CodeType +from typing import Final from RestrictedPython import ( RestrictingNodeTransformer, @@ -45,20 +48,20 @@ class AsyncAwareTransformer(RestrictingNodeTransformer): ``node_contents_visit`` so their children still get transformed. """ - def visit_AsyncFunctionDef(self, node: Any) -> Any: + def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> ast.AST: return self.visit_FunctionDef(node) - def visit_AsyncFor(self, node: Any) -> Any: + def visit_AsyncFor(self, node: ast.AsyncFor) -> ast.AST: return self.node_contents_visit(node) - def visit_AsyncWith(self, node: Any) -> Any: + def visit_AsyncWith(self, node: ast.AsyncWith) -> ast.AST: return self.node_contents_visit(node) - def visit_Await(self, node: Any) -> Any: + def visit_Await(self, node: ast.Await) -> ast.AST: return self.node_contents_visit(node) -_INPLACE_OPS: Final[dict[str, Any]] = { +_INPLACE_OPS: Final[Mapping[str, Callable[[object, object], object]]] = { "+=": operator.iadd, "-=": operator.isub, "*=": operator.imul, @@ -75,7 +78,7 @@ _INPLACE_OPS: Final[dict[str, Any]] = { } -def _inplacevar_(op: str, x: Any, y: Any) -> Any: +def _inplacevar_(op: str, x: object, y: object) -> object: # RestrictedPython rewrites ``x += 1`` on a simple name into # ``x = _inplacevar_("+=", x, 1)``. The package deliberately ships no # default, so we dispatch through ``operator``'s in-place helpers, which @@ -86,7 +89,7 @@ def _inplacevar_(op: str, x: Any, y: Any) -> Any: return fn(x, y) -def _build_sandbox_builtins() -> dict[str, Any]: +def _build_sandbox_builtins() -> dict[str, object]: # ``limited_builtins`` overrides ``list``/``tuple``/``range`` from # ``safe_builtins`` with bounds-checking variants (e.g. ``limited_range`` # rejects ``range(10**18)``). ``utility_builtins`` adds ``set``, @@ -98,25 +101,26 @@ def _build_sandbox_builtins() -> dict[str, Any]: } -def build_sandbox_globals() -> dict[str, Any]: +def build_sandbox_globals() -> dict[str, object]: """Assemble the globals dict for executing guardrail code. Includes the LiteLLM-provided primitives (``regex_match``, ``http_get``, ``allow``/``block``/``modify``, etc.) plus the RestrictedPython guards that the compiled bytecode expects to find by name. """ - sandbox: Final[dict[str, Any]] = get_custom_code_primitives().copy() - sandbox["__builtins__"] = _build_sandbox_builtins() - sandbox["_getattr_"] = safer_getattr - sandbox["_getitem_"] = default_guarded_getitem - sandbox["_getiter_"] = default_guarded_getiter - sandbox["_iter_unpack_sequence_"] = guarded_iter_unpack_sequence - sandbox["_write_"] = full_write_guard - sandbox["_inplacevar_"] = _inplacevar_ - return sandbox + return { + **get_custom_code_primitives(), + "__builtins__": _build_sandbox_builtins(), + "_getattr_": safer_getattr, + "_getitem_": default_guarded_getitem, + "_getiter_": default_guarded_getiter, + "_iter_unpack_sequence_": guarded_iter_unpack_sequence, + "_write_": full_write_guard, + "_inplacevar_": _inplacevar_, + } -def compile_sandboxed(source: str, filename: str = "") -> Any: +def compile_sandboxed(source: str, filename: str = "") -> CodeType: """Compile guardrail source with RestrictedPython's AST transformer. Raises ``SyntaxError`` on either a Python syntax error or a restricted diff --git a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py index c0f72af7576..214d4b486d4 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py +++ b/litellm/proxy/guardrails/guardrail_hooks/deepkeep/deepkeep.py @@ -7,9 +7,10 @@ import os from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol import httpx +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version @@ -27,6 +28,12 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionToolCallChunk, + ChatCompletionToolParam, + ) + from litellm.types.utils import ChatCompletionMessageToolCall GUARDRAIL_NAME: Final = "deepkeep" @@ -34,6 +41,39 @@ GUARDRAIL_NAME: Final = "deepkeep" _DEEPKEEP_GUARDRAIL_ENDPOINT: Final = "/v3/openai/beta/litellm_basic_guardrail_api" +class DeepKeepFirewallResponse(TypedDict): + """Body returned by the DeepKeep firewall endpoint.""" + + action: ReadOnly[NotRequired[str]] + blocked_reason: ReadOnly[NotRequired[str]] + texts: ReadOnly[NotRequired["list[str]"]] + images: ReadOnly[NotRequired["list[str]"]] + tools: ReadOnly[NotRequired["list[ChatCompletionToolParam]"]] + tool_calls: ReadOnly[NotRequired["list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall]"]] + structured_messages: ReadOnly[NotRequired["list[AllMessageValues]"]] + + +class _DeepKeepInitKwargsView(TypedDict): + """Typed read of the guardrail name carried in the untyped base-guardrail kwargs.""" + + guardrail_name: ReadOnly[str] + + +class _DeepKeepMetadataSource(TypedDict, total=False): + """Typed read of the two untyped metadata mappings this guardrail merges.""" + + litellm_metadata: ReadOnly[Mapping[str, object]] + metadata: ReadOnly[Mapping[str, object]] + + +class _FirewallResponseBody(Protocol): + def json(self) -> DeepKeepFirewallResponse: ... + + +def _firewall_response_body(response: _FirewallResponseBody) -> DeepKeepFirewallResponse: + return response.json() + + class DeepKeepGuardrailMissingSecrets(Exception): """Exception raised when DeepKeep API key or firewall_id is missing.""" @@ -125,14 +165,16 @@ class DeepKeepGuardrail(CustomGuardrail): super().__init__(**kwargs) + init_view: Final[_DeepKeepInitKwargsView] = {"guardrail_name": kwargs.get("guardrail_name", "unknown")} + verbose_proxy_logger.debug( "DeepKeep guardrail initialized: guardrail_name=%s, api_base=%s, firewall_id=%s", - kwargs.get("guardrail_name", "unknown"), + init_view["guardrail_name"], self.api_base, self.firewall_id, ) - def _extract_user_api_key_metadata(self, request_data: dict) -> dict[str, Any]: + def _extract_user_api_key_metadata(self, request_data: _DeepKeepMetadataSource) -> dict[str, object]: """ Extract user API key metadata from request_data for the DeepKeep API. @@ -142,11 +184,11 @@ class DeepKeepGuardrail(CustomGuardrail): Returns: Dictionary with user API key metadata fields. """ - result_metadata: Final[dict[str, Any]] = {} + result_metadata: Final[dict[str, object]] = {} litellm_metadata: Final = request_data.get("litellm_metadata", {}) top_level_metadata: Final = request_data.get("metadata", {}) - metadata_dict: Final = {**top_level_metadata, **litellm_metadata} + metadata_dict: Final[Mapping[str, object]] = {**top_level_metadata, **litellm_metadata} if not metadata_dict: return result_metadata @@ -219,7 +261,7 @@ class DeepKeepGuardrail(CustomGuardrail): ) -> GenericGuardrailAPIInputs: """Handle errors from the DeepKeep API with fail-open/fail-closed logic.""" if is_unreachable and self.unreachable_fallback == "fail_open": - http_status_code: Final = getattr(getattr(error, "response", None), "status_code", None) + http_status_code: Final[int | None] = getattr(getattr(error, "response", None), "status_code", None) return self._fail_open_passthrough( inputs=inputs, input_type=input_type, @@ -233,12 +275,12 @@ class DeepKeepGuardrail(CustomGuardrail): @staticmethod def _build_return_inputs( *, - response_json: dict[str, Any], - texts: list, - images: Any | None, - tools: Any | None, - tool_calls: Any | None, - structured_messages: Any | None, + response_json: DeepKeepFirewallResponse, + texts: list[str], + images: "list[str] | None", + tools: "list[ChatCompletionToolParam] | None", + tool_calls: "list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] | None", + structured_messages: "list[AllMessageValues] | None", ) -> GenericGuardrailAPIInputs: """Merge original inputs with any guardrail-modified values from the API response. @@ -248,22 +290,27 @@ class DeepKeepGuardrail(CustomGuardrail): silently discarded in favour of the original content. """ return_inputs: Final = GenericGuardrailAPIInputs(texts=texts) - if response_json.get("texts") is not None: - return_inputs["texts"] = response_json["texts"] - if response_json.get("images") is not None: - return_inputs["images"] = response_json["images"] + texts_override: Final = response_json.get("texts") + if texts_override is not None: + return_inputs["texts"] = texts_override + images_override: Final = response_json.get("images") + if images_override is not None: + return_inputs["images"] = images_override elif images is not None: return_inputs["images"] = images - if response_json.get("tools") is not None: - return_inputs["tools"] = response_json["tools"] + tools_override: Final = response_json.get("tools") + if tools_override is not None: + return_inputs["tools"] = tools_override elif tools is not None: return_inputs["tools"] = tools - if response_json.get("tool_calls") is not None: - return_inputs["tool_calls"] = response_json["tool_calls"] + tool_calls_override: Final = response_json.get("tool_calls") + if tool_calls_override is not None: + return_inputs["tool_calls"] = tool_calls_override elif tool_calls is not None: return_inputs["tool_calls"] = tool_calls - if response_json.get("structured_messages") is not None: - return_inputs["structured_messages"] = response_json["structured_messages"] + structured_messages_override: Final = response_json.get("structured_messages") + if structured_messages_override is not None: + return_inputs["structured_messages"] = structured_messages_override elif structured_messages is not None: return_inputs["structured_messages"] = structured_messages return return_inputs @@ -309,7 +356,7 @@ class DeepKeepGuardrail(CustomGuardrail): request_body: Final = request_data.get("body") or {} # Merge additional provider-specific params from config and dynamic params - additional_params: Final[dict[str, Any]] = {"firewall_id": self.firewall_id} + additional_params: Final[dict[str, object]] = {"firewall_id": self.firewall_id} dynamic_params: Final = self.get_guardrail_dynamic_request_body_params(request_body) if dynamic_params: additional_params.update({k: v for k, v in dynamic_params.items() if k != "firewall_id"}) @@ -318,7 +365,7 @@ class DeepKeepGuardrail(CustomGuardrail): user_metadata: Final = self._extract_user_api_key_metadata(request_data) # Build request payload - guardrail_request: Final[dict[str, Any]] = { + guardrail_request: Final[dict[str, object]] = { "litellm_call_id": (logging_obj.litellm_call_id if logging_obj else None), "litellm_trace_id": (logging_obj.litellm_trace_id if logging_obj else None), "texts": texts, @@ -343,7 +390,7 @@ class DeepKeepGuardrail(CustomGuardrail): ) response.raise_for_status() - response_json: Final = response.json() + response_json: Final = _firewall_response_body(response) verbose_proxy_logger.debug("DeepKeep guardrail response: %s", response_json) diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 84c6b220e62..b9cc90107e2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -740,7 +740,7 @@ class HeadroomGuardrail(CustomGuardrail): response: Any, anthropic_messages_provider_config: Any, anthropic_messages_optional_request_params: dict, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj | None, stream: bool, kwargs: dict, ) -> AgenticLoopPlan: diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py index 507dd645953..c595952bec2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py @@ -24,11 +24,12 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import ChatCompletionToolCallChunk from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import ( HiddenlayerAction, HiddenlayerMessages, ) -from litellm.types.utils import GenericGuardrailAPIInputs +from litellm.types.utils import ChatCompletionMessageToolCall, GenericGuardrailAPIInputs if TYPE_CHECKING: from pydantic import BaseModel @@ -46,8 +47,12 @@ class _HiddenlayerAnalysisEntry(TypedDict, total=False): detected: bool +class _HiddenlayerModifiedMessage(TypedDict): + content: ReadOnly[str | list[Mapping[str, str]]] + + class _HiddenlayerModifiedSide(TypedDict): - messages: Any + messages: ReadOnly[list[_HiddenlayerModifiedMessage]] class _HiddenlayerResponse(TypedDict, total=False): @@ -56,8 +61,16 @@ class _HiddenlayerResponse(TypedDict, total=False): modified_data: Mapping[str, _HiddenlayerModifiedSide] +class _ProxyServerRequest(TypedDict, total=False): + headers: ReadOnly[dict[str, str]] + + +class _HiddenlayerRequestData(TypedDict, total=False): + proxy_server_request: ReadOnly[_ProxyServerRequest] + + class _LoggedCallMetadata(TypedDict, total=False): - headers: ReadOnly[Mapping[str, str]] + headers: ReadOnly[dict[str, str]] class _LoggedCallLitellmParams(TypedDict, total=False): @@ -65,7 +78,7 @@ class _LoggedCallLitellmParams(TypedDict, total=False): class _HiddenlayerOutputMessage(TypedDict, total=False): - content: ReadOnly[str | Sequence[Mapping[str, str]]] + content: ReadOnly[str | list[Mapping[str, str]]] class _HiddenlayerChoiceMessage(TypedDict, total=False): @@ -81,6 +94,15 @@ class _HiddenlayerV2Output(TypedDict, total=False): choices: ReadOnly[Sequence[_HiddenlayerChoice]] +class _HiddenlayerV2OutputView(TypedDict): + """Typed read of the untyped JSON body returned by the HiddenLayer detection endpoints.""" + + evaluation: ReadOnly[_HiddenlayerV2Output] + + +_HiddenlayerV2Payload = Mapping[str, object] | list[ChatCompletionToolCallChunk] | list[ChatCompletionMessageToolCall] + + class _LoggedCallDetails(Protocol): """Logging object view that exposes its untyped call details with the shape this guardrail reads.""" @@ -94,7 +116,25 @@ class _TokenPayloadSource(Protocol): def json(self) -> Mapping[str, str]: ... -def _logged_request_headers(logging_obj: _LoggedCallDetails) -> Mapping[str, str]: +class _InteractionPayloadSource(Protocol): + """Response view that decodes the HiddenLayer v1 interaction body with the shape this guardrail reads.""" + + def json(self) -> _HiddenlayerResponse: ... + + +def _interaction_body(response: _InteractionPayloadSource) -> _HiddenlayerResponse: + return response.json() + + +def _proxy_server_request(request_data: _HiddenlayerRequestData) -> _ProxyServerRequest | None: + return request_data.get("proxy_server_request") + + +def _proxy_request_headers(request_data: _HiddenlayerRequestData) -> dict[str, str]: + return request_data.get("proxy_server_request", {}).get("headers", {}) + + +def _logged_request_headers(logging_obj: _LoggedCallDetails) -> dict[str, str]: return logging_obj.model_call_details.get("litellm_params", {}).get("metadata", {}).get("headers", {}) @@ -204,7 +244,7 @@ class HiddenlayerGuardrail(CustomGuardrail): # from the logging object. It ends up working out that on the request, we parse the # hiddenlayer params from the raw request and then retrieve those same headers # from the logger object on the response from the model. - headers = request_data.get("proxy_server_request", {}).get("headers", {}) + headers = _proxy_request_headers(request_data) if not headers and logging_obj and logging_obj.model_call_details: headers = _logged_request_headers(logging_obj) @@ -309,7 +349,7 @@ class HiddenlayerGuardrail(CustomGuardrail): headers=headers, ) response.raise_for_status() - result: _HiddenlayerResponse = response.json() + result: _HiddenlayerResponse = _interaction_body(response) verbose_proxy_logger.debug("Hiddenlayer reponse: %s", result) @@ -333,7 +373,7 @@ class HiddenlayerGuardrail(CustomGuardrail): raise e response.raise_for_status() - result = response.json() + result = _interaction_body(response) verbose_proxy_logger.debug("Hiddenlayer reponse: %s", result) return result @@ -401,13 +441,13 @@ class HiddenlayerGuardrailV2(CustomGuardrail): # from the logging object. It ends up working out that on the request, we parse the # hiddenlayer params from the raw request and then retrieve those same headers # from the logger object on the response from the model. - headers = request_data.get("proxy_server_request", {}).get("headers", {}) + headers = _proxy_request_headers(request_data) if not headers and logging_obj and logging_obj.model_call_details: - headers = logging_obj.model_call_details.get("litellm_params", {}).get("metadata", {}).get("headers", {}) + headers = _logged_request_headers(logging_obj) # put our roundtrip id in the header to the model so we get it on the way back from the model if "hl-roundtrip-id" not in headers: - proxy_req: Final = request_data.get("proxy_server_request") + proxy_req: Final = _proxy_server_request(request_data) if proxy_req is not None and "headers" in proxy_req: proxy_req["headers"]["hl-roundtrip-id"] = str(uuid4()) headers["hl-roundtrip-id"] = proxy_req["headers"]["hl-roundtrip-id"] @@ -417,7 +457,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): if "hl-requester-id" not in hl_headers: hl_headers["hl-requester-id"] = "LiteLLM" - payload: object + payload: _HiddenlayerV2Payload if input_type == "request": payload = { "messages": inputs.get("structured_messages"), @@ -445,7 +485,8 @@ class HiddenlayerGuardrailV2(CustomGuardrail): response: Final = await self._call_hiddenlayer(payload, input_type, hl_headers) output: Final = response.json() - evaluated_output: Final[_HiddenlayerV2Output] = output + output_view: Final[_HiddenlayerV2OutputView] = {"evaluation": output} + evaluated_output: Final = output_view["evaluation"] if _header_value(response.headers, "hl-runtime-action", "").lower() == "block": raise HTTPException( @@ -456,7 +497,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): }, ) - new_texts: Final = [] + new_texts: Final[list[str]] = [] if input_type == "request": inputs["structured_messages"] = output @@ -484,7 +525,7 @@ class HiddenlayerGuardrailV2(CustomGuardrail): async def _call_hiddenlayer( self, - payload: Any, + payload: _HiddenlayerV2Payload, input_type: Literal["request", "response"], hl_headers: dict[str, str], ) -> httpx.Response: diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py index e3f67f0024b..172b1440ca3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_as_a_judge/__init__.py @@ -1,10 +1,11 @@ """LLM-as-a-Judge guardrail: uses an LLM to score responses against weighted criteria.""" -from collections.abc import Callable +from collections.abc import Callable, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Any, Final, Generic, Literal, Optional, TypeVar from fastapi import HTTPException +from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger @@ -22,6 +23,7 @@ if TYPE_CHECKING: from litellm import Router from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.guardrails import Guardrail, LitellmParams + from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import StandardLoggingEvalInformation JUDGE_SYSTEM_PROMPT = """You are a quality judge. Evaluate the assistant's response against the criteria provided. @@ -40,29 +42,53 @@ _default_router_provider: Final = default_router_provider _parse_judge_verdict: Final = parse_json_verdict _extract_text_from_content: Final = extract_text_from_content +_ParamT = TypeVar("_ParamT") + + +class _LitellmParamView(TypedDict, Generic[_ParamT]): + """Typed read of a single entry in an untyped ``litellm_params`` mapping.""" + + value: ReadOnly[_ParamT] + + +class JudgeCriterion(TypedDict): + """A single weighted criterion the judge scores the response against.""" + + name: ReadOnly[NotRequired[str]] + description: ReadOnly[NotRequired[str]] + weight: ReadOnly[NotRequired[float]] + + +class JudgeMessage(TypedDict): + """The parts of a conversation message the judge prompt renders.""" + + role: ReadOnly[NotRequired[str]] + content: ReadOnly[NotRequired[object]] + def _get_litellm_param( litellm_params: "LitellmParams", guardrail: "Guardrail", key: str, - default: Any = None, -) -> Any: - val: Final = getattr(litellm_params, key, None) + default: _ParamT, +) -> _ParamT: + val: Final[_ParamT | None] = getattr(litellm_params, key, None) if val is not None: return val raw: Final = guardrail.get("litellm_params") if isinstance(raw, dict) and key in raw: - return raw[key] + entry: Final[_LitellmParamView[_ParamT]] = {"value": raw[key]} + return entry["value"] if raw is not None and not isinstance(raw, dict): - attr: Final = getattr(raw, key, None) + attr: Final[_ParamT | None] = getattr(raw, key, None) if attr is not None: return attr return default def _build_judge_prompt( - criteria: list[dict[str, Any]], - messages: list[dict[str, Any]], + criteria: Sequence[JudgeCriterion], + messages: Sequence[JudgeMessage], response_text: str, ) -> str: criteria_block: Final = "\n".join( @@ -87,7 +113,7 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): self, guardrail_name: str, judge_model: str, - criteria: list[dict[str, Any]], + criteria: Sequence[JudgeCriterion], overall_threshold: float = 80.0, on_failure: Literal["block", "log"] = "block", event_hook: GuardrailEventHooks | list[GuardrailEventHooks] | None = None, @@ -121,10 +147,10 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): async def _run_judge( self, - messages: list[dict[str, Any]], + messages: Sequence[JudgeMessage], response_text: str, - ) -> dict[str, Any]: - judge_messages: Final = [ + ) -> dict[str, object]: + judge_messages: Final[list[AllMessageValues]] = [ {"role": "system", "content": JUDGE_SYSTEM_PROMPT}, { "role": "user", @@ -159,10 +185,10 @@ class LLMAsAJudgeGuardrail(CustomGuardrail): start_time: Final = datetime.now() status: GuardrailStatus = "success" - judge_result: dict[str, Any] = {} + judge_result: dict[str, object] = {} try: - messages: Final[list[dict[str, Any]]] = request_data.get("messages") or [] + messages: Final[Sequence[JudgeMessage]] = request_data.get("messages") or [] try: judge_result = await self._run_judge(messages, response_text) @@ -238,11 +264,11 @@ def initialize_guardrail( if not guardrail_name: raise ValueError("llm_as_a_judge guardrail requires a guardrail_name") - judge_model: Final = _get_litellm_param(litellm_params, guardrail, "judge_model") + judge_model: Final[str] = _get_litellm_param(litellm_params, guardrail, "judge_model", "") if not judge_model: raise ValueError("llm_as_a_judge guardrail requires judge_model in litellm_params") - criteria: Final = _get_litellm_param(litellm_params, guardrail, "criteria") or [] + criteria: Final[Sequence[JudgeCriterion]] = _get_litellm_param(litellm_params, guardrail, "criteria", ()) or () if not criteria: raise ValueError("llm_as_a_judge guardrail requires at least one criterion") @@ -250,13 +276,13 @@ def initialize_guardrail( if abs(weight_total - 100) > 0.5: raise ValueError(f"llm_as_a_judge criterion weights must sum to 100 (got {weight_total})") - on_failure: Final = _get_litellm_param(litellm_params, guardrail, "on_failure", "block") + on_failure: Final[Literal["block", "log"]] = _get_litellm_param(litellm_params, guardrail, "on_failure", "block") if on_failure not in _VALID_ON_FAILURE: raise ValueError(f"llm_as_a_judge on_failure must be 'block' or 'log', got '{on_failure}'") overall_threshold: Final = float(_get_litellm_param(litellm_params, guardrail, "overall_threshold", 80.0)) - mode: Final = _get_litellm_param(litellm_params, guardrail, "mode") + mode: Final[str | None] = _get_litellm_param(litellm_params, guardrail, "mode", None) event_hook: GuardrailEventHooks | None = None if isinstance(mode, str) and mode in {e.value for e in GuardrailEventHooks}: event_hook = GuardrailEventHooks(mode) diff --git a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/mcp_end_user_permission.py b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/mcp_end_user_permission.py index 4e8eec6a14e..1e246922fca 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/mcp_end_user_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/mcp_end_user_permission/mcp_end_user_permission.py @@ -22,6 +22,7 @@ from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel GUARDRAIL_NAME: Final = "mcp_end_user_permission" @@ -54,7 +55,7 @@ class MCPEndUserPermissionGuardrail(CustomGuardrail): inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"] = "request", - logging_obj: Any | None = None, + logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: """ Filters MCP tools the end user cannot access based on their diff --git a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py index 6644a3d3902..b31ed4b0f4a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ovalix/ovalix.py @@ -25,6 +25,7 @@ from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -196,7 +197,7 @@ class OvalixGuardrail(CustomGuardrail): inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"], - logging_obj: Any | None = None, + logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: """ Apply Ovalix guardrail to the given inputs (request or response text). diff --git a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py index c25f704567e..f780f4dd67d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/promptguard/promptguard.py @@ -7,7 +7,9 @@ before and after LLM calls. """ import os -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict + +from typing_extensions import ReadOnly from litellm._logging import verbose_proxy_logger from litellm.exceptions import GuardrailRaisedException @@ -20,6 +22,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -34,6 +37,22 @@ _DEFAULT_API_BASE: Final = "https://api.promptguard.co" _GUARD_ENDPOINT: Final = "/api/v1/guard" +class PromptGuardGuardAPIResponse(TypedDict, total=False): + """Body returned by the PromptGuard ``/api/v1/guard`` endpoint.""" + + decision: ReadOnly[str] + threat_type: ReadOnly[str] + event_id: ReadOnly[str] + confidence: ReadOnly[float] + redacted_messages: ReadOnly[list[AllMessageValues]] + + +class PromptGuardHTTPView(TypedDict): + """Typed read of the untyped JSON body returned by the httpx client.""" + + guard_response: ReadOnly[PromptGuardGuardAPIResponse] + + class PromptGuardMissingCredentials(Exception): pass @@ -96,7 +115,7 @@ class PromptGuardGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: @@ -114,7 +133,7 @@ class PromptGuardGuardrail(CustomGuardrail): direction: Final = "input" if input_type == "request" else "output" - payload: Final[dict[str, Any]] = { + payload: Final[dict[str, object]] = { "messages": messages, "direction": direction, } @@ -144,7 +163,8 @@ class PromptGuardGuardrail(CustomGuardrail): timeout=10.0, ) response.raise_for_status() - result: Final = response.json() + view: Final[PromptGuardHTTPView] = {"guard_response": response.json()} + result: Final = view["guard_response"] except Exception as exc: verbose_proxy_logger.error("PromptGuard API error: %s", str(exc)) if self.block_on_error: @@ -187,7 +207,7 @@ class PromptGuardGuardrail(CustomGuardrail): return inputs @staticmethod - def _extract_texts_from_messages(messages: list) -> list[str]: + def _extract_texts_from_messages(messages: list[AllMessageValues]) -> list[str]: """Extract text content from user-role messages only. Only user messages are extracted to avoid injecting system or diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 88aa55fd4a9..8b57bdca2fe 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -1372,8 +1372,25 @@ class DBHealthCache(TypedDict): db_health_cache: DBHealthCache = {"status": "unknown", "last_updated": datetime.now()} +# Bounds each DB round-trip on the probe path so a hung connection during a +# failover cannot make the probe fail by timeout (k8s default timeoutSeconds: 5). +DB_READINESS_CHECK_TIMEOUT_SECONDS: Final = 2.0 +# One deadline for the whole probe-path DB check (initial check + reconnect + +# re-check, including reconnect lock waits), kept under timeoutSeconds: 5. +DB_READINESS_PROBE_DEADLINE_SECONDS: Final = 4.0 -async def _db_health_readiness_check(): + +async def _db_health_readiness_check() -> DBHealthCache: + try: + return await asyncio.wait_for( + _db_health_readiness_check_unbounded(), + timeout=DB_READINESS_PROBE_DEADLINE_SECONDS, + ) + except asyncio.TimeoutError: + return {"status": "disconnected", "last_updated": db_health_cache["last_updated"]} + + +async def _db_health_readiness_check_unbounded() -> DBHealthCache: from litellm.proxy.proxy_server import prisma_client global db_health_cache @@ -1387,7 +1404,7 @@ async def _db_health_readiness_check(): db_health_cache = {"status": "disconnected", "last_updated": datetime.now()} return db_health_cache - await prisma_client.health_check() + await asyncio.wait_for(prisma_client.health_check(), timeout=DB_READINESS_CHECK_TIMEOUT_SECONDS) db_health_cache = {"status": "connected", "last_updated": datetime.now()} return db_health_cache except Exception as e: @@ -1395,8 +1412,15 @@ async def _db_health_readiness_check(): if PrismaDBExceptionHandler.is_database_transport_error(e): try: verbose_proxy_logger.warning("_db_health_readiness_check: health_check failed, attempting reconnect") - await prisma_client.attempt_db_reconnect(reason="health_readiness_check") - await prisma_client.health_check() + await prisma_client.attempt_db_reconnect( + reason="health_readiness_check", + timeout_seconds=DB_READINESS_CHECK_TIMEOUT_SECONDS, + lock_timeout_seconds=DB_READINESS_CHECK_TIMEOUT_SECONDS, + ) + await asyncio.wait_for( + prisma_client.health_check(), + timeout=DB_READINESS_CHECK_TIMEOUT_SECONDS, + ) verbose_proxy_logger.info("_db_health_readiness_check: reconnect succeeded") db_health_cache = { "status": "connected", @@ -1580,7 +1604,14 @@ async def _get_health_readiness_details( # serve requests that depend on persisted state (keys, budgets, # spend logs). Return 503 so orchestrators take this pod out of # rotation; "Not connected" (no DB configured at all) stays 200. - if response is not None and db_health_status["status"] != "connected": + # With allow_requests_on_db_unavailable the proxy keeps serving + # during a DB outage, so the pod must stay in rotation (200) and + # report the DB state through the body instead. + if ( + response is not None + and db_health_status["status"] != "connected" + and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() + ): response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE return { "status": "healthy", @@ -1671,7 +1702,10 @@ async def _resolve_public_readiness_db(response: Response) -> str: return "Not connected" db_health_status: Final = await _db_health_readiness_check() - if db_health_status["status"] != "connected": + if ( + db_health_status["status"] != "connected" + and not PrismaDBExceptionHandler.should_allow_request_on_db_unavailable() + ): response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE return db_health_status["status"] diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 6abfca1d3a0..f593e94b36f 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -1,7 +1,7 @@ import asyncio import traceback from datetime import datetime -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import verbose_proxy_logger @@ -36,6 +36,9 @@ from litellm.types.utils import ( ) from litellm.utils import get_end_user_id_for_cost_tracking +if TYPE_CHECKING: + from litellm.proxy.utils import ProxyLogging + _UNATTRIBUTED_TRACKABLE_CALL_TYPES: Final[frozenset[str]] = frozenset( { CallTypes.pass_through.value, @@ -555,7 +558,7 @@ def _get_request_tags_for_cost_tracking( async def _update_database_and_spend_counters( - proxy_logging_obj: Any, + proxy_logging_obj: "ProxyLogging", increment_spend_counters: Any, user_api_key: str | None, user_id: str | None, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 194ef5d756c..b7688790e54 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1948,6 +1948,13 @@ async def add_litellm_data_to_request( for key, value in data["litellm_metadata"].items(): if key not in data[_metadata_variable_name]: data[_metadata_variable_name][key] = value + if _metadata_variable_name == "metadata": + data["metadata"]["tags"] = LiteLLMProxyRequestSetup._merge_tags( # pyright: ignore[reportPrivateUsage] # same-module helper, budget blocks the unsuppressed idiom sibling call sites use + request_tags=data["metadata"].get("tags"), + tags_to_add=data["litellm_metadata"].get("tags"), + ) + if _metadata_variable_name == "metadata": + data.pop("litellm_metadata", None) data = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( data=data, diff --git a/litellm/proxy/management_endpoints/cache_settings_endpoints.py b/litellm/proxy/management_endpoints/cache_settings_endpoints.py index 77ff77c9a88..40124bd19a4 100644 --- a/litellm/proxy/management_endpoints/cache_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/cache_settings_endpoints.py @@ -15,7 +15,7 @@ from datetime import datetime, timezone from typing import TYPE_CHECKING, Any, Final, Protocol from fastapi import APIRouter, Depends, Header, HTTPException -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm._redis import _redis_kwargs_from_environment @@ -41,6 +41,8 @@ if TYPE_CHECKING: router: Final = APIRouter() +_STORED_CACHE_SETTINGS_ADAPTER: Final = TypeAdapter(dict[str, object]) + class _CacheConfigRow(Protocol): @property @@ -61,13 +63,13 @@ def _cache_config_table(prisma_client: "PrismaClient") -> _CacheConfigTable: # Sentinel passwords never leave the server in a GET response. `url` is here # because a Redis/Valkey URL can embed a password inline # (e.g. redis://:secret@host:6379/1). -_CACHE_SENSITIVE_FIELDS: Final[set] = {"password", "sentinel_password", "url"} +_CACHE_SENSITIVE_FIELDS: Final[set[str]] = {"password", "sentinel_password", "url"} # The env fallback resolves the full set of redis.Redis kwargs, which includes # credential-bearing params (azure_client_secret, ssl_password, ...) that are # not cache UI fields. Only overlay fields the settings page actually renders, # so the read never surfaces a credential the UI does not manage. -_CACHE_SETTINGS_FIELD_NAMES: Final[frozenset] = frozenset(field.field_name for field in CACHE_SETTINGS_FIELDS) +_CACHE_SETTINGS_FIELD_NAMES: Final[frozenset[str]] = frozenset(field.field_name for field in CACHE_SETTINGS_FIELDS) # Classifier used, alongside _CACHE_SENSITIVE_FIELDS, to redact any # credential-bearing key before it leaves the server (`url` is kept in the @@ -78,7 +80,7 @@ _CREDENTIAL_CLASSIFIER: Final = SensitiveDataMasker() _REDACTED_VALUE: Final = "***REDACTED***" -_URL_OVERRIDDEN_CONNECTION_FIELDS: Final[frozenset] = frozenset({"host", "port", "db", "password", "username"}) +_URL_OVERRIDDEN_CONNECTION_FIELDS: Final[frozenset[str]] = frozenset({"host", "port", "db", "password", "username"}) def _resolve_cache_url_precedence(settings: Mapping[str, object]) -> dict[str, Any]: @@ -160,7 +162,7 @@ def _has_connection_target(value: object) -> bool: # Every field that identifies which Redis a credential belongs to, across node # (host/port/url), cluster (redis_startup_nodes), and sentinel # (sentinel_nodes/service_name) modes. A stored secret is bound to these. -_CONNECTION_TARGET_FIELDS: Final[tuple] = ( +_CONNECTION_TARGET_FIELDS: Final[tuple[str, ...]] = ( "host", "port", "url", @@ -363,8 +365,6 @@ class CacheSettingsManager: Initialize cache settings from database into the router on startup. Only reinitializes if cache params have changed. """ - import json - try: cache_config: Final = await call_with_db_reconnect_retry( prisma_client, @@ -374,10 +374,11 @@ class CacheSettingsManager: if cache_config is not None and cache_config.cache_settings: # Parse cache settings JSON cache_settings_json: Final = cache_config.cache_settings - if isinstance(cache_settings_json, str): - cache_settings_dict = json.loads(cache_settings_json) - else: - cache_settings_dict = cache_settings_json + cache_settings_dict: Final[dict[str, object]] = ( + _STORED_CACHE_SETTINGS_ADAPTER.validate_json(cache_settings_json) + if isinstance(cache_settings_json, str) + else dict(cache_settings_json) + ) # Decrypt cache settings decrypted_settings: Final = proxy_config._decrypt_db_variables(variables_dict=cache_settings_dict) diff --git a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py index 4bc53678c23..1096954536a 100644 --- a/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py +++ b/litellm/proxy/management_endpoints/policy_endpoints/ai_policy_suggester.py @@ -9,6 +9,7 @@ from typing import Final import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL +from litellm.proxy._types import ProxyErrorTypes, ProxyException SUGGEST_TOOL: Final = { "type": "function", @@ -60,6 +61,18 @@ class AiPolicySuggester: system_prompt: Final = self._build_system_prompt(templates) user_prompt: Final = self._build_user_prompt(attack_examples, description) model = model or DEFAULT_COMPETITOR_DISCOVERY_MODEL + custom_llm_provider: Final = model.split("/", 1)[0] if "/" in model else None + supported_params: Final = litellm.get_supported_openai_params( + model=model, + custom_llm_provider=custom_llm_provider, + ) + if supported_params is not None and "tools" not in supported_params: + raise ProxyException( + message=(f"AI policy suggestion requires tool calling; model '{model}' does not support it"), + type=ProxyErrorTypes.validation_error.value, + param="model", + code=400, + ) try: response: Final = await litellm.acompletion( @@ -74,6 +87,7 @@ class AiPolicySuggester: "function": {"name": "select_policy_templates"}, }, temperature=0.2, + drop_params=True, ) tool_calls: Final = response.choices[0].message.tool_calls diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index ddfdb56ac2c..992ed0d814d 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1351,15 +1351,16 @@ def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool: provider response briefly lags before the output id populates). Retiring in that window loses the spend record forever. Retire only once we can prove there is nothing left to recover: the output file has actually arrived, or the provider - reports no successful request lines. When counts are unknown, stay eligible so - the next poller pass revisits it. (#37713) + reported a positive total with zero successful request lines, proving it + enumerated the batch and none succeeded. A zero or unknown total means counts + are unreported, so stay eligible and let the next poller pass revisit it. (#37713) """ if response.output_file_id is not None: return True request_counts = response.request_counts if request_counts is None: return False - return request_counts.completed == 0 + return request_counts.total > 0 and request_counts.completed == 0 async def update_batch_in_database( diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 92bbd58ed90..9bc90260de1 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -8,7 +8,7 @@ import asyncio import traceback from collections.abc import Mapping -from typing import Any, BinaryIO, Final, cast, get_args +from typing import Any, BinaryIO, Final, TypedDict, cast, get_args import httpx from fastapi import ( @@ -23,6 +23,7 @@ from fastapi import ( status, ) from pydantic import TypeAdapter +from typing_extensions import ReadOnly import litellm from litellm import CreateFileRequest, get_secret_str @@ -83,6 +84,13 @@ router: Final = APIRouter() _MAX_BATCH_FILE_SIZE_MB_ADAPTER: Final = TypeAdapter(int | None) + +class UploadedFileInfo(TypedDict): + filename: ReadOnly[str | None] + content_type: ReadOnly[str | None] + size: ReadOnly[int | None] + + files_config = None @@ -526,6 +534,22 @@ async def create_file( proxy_config=proxy_config, ) + uploaded_file_info: Final[UploadedFileInfo] = { + "filename": file.filename, + "content_type": file.content_type, + "size": file.size, + } + data["purpose"] = purpose + data["file"] = uploaded_file_info + hooked_data: Final = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=user_api_key_dict, + data=data, + call_type="acreate_file", + ) + data = hooked_data if hooked_data is not None else data + data.pop("purpose", None) + data.pop("file", None) + # /v1/files stores its proxy metadata under litellm_metadata, not metadata request_metadata: Final = data.get("metadata") or data.get("litellm_metadata") or EMPTY_MAPPING scan_result: Final = await _scan_batch_upload( diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 9a3bc82c6fa..525e7099b89 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -26,6 +26,7 @@ from litellm.constants import ( ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS, BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES, ) +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._types import * @@ -1057,7 +1058,7 @@ async def bedrock_proxy_route( except ImportError: raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.") - aws_region_name: Final = litellm.utils.get_secret(secret_name="AWS_REGION_NAME") + aws_region_name: Final = get_secret_str(secret_name="AWS_REGION_NAME") if not _is_bedrock_agent_runtime_route(endpoint=endpoint): return await bedrock_llm_proxy_route( endpoint=endpoint, @@ -1072,7 +1073,7 @@ async def bedrock_proxy_route( detail="bedrock-agent-runtime pass-through is disabled on this proxy.", ) - base_target_url: Final = f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com" + base_target_url: Final = f"https://bedrock-agent-runtime.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}" encoded_endpoint = httpx.URL(endpoint).path # Ensure endpoint starts with '/' for proper URL construction @@ -1205,7 +1206,7 @@ async def comprehend_medical_proxy_route( "X-Amz-Target": f"{COMPREHEND_MEDICAL_TARGET_PREFIX}.{operation}", } ) - target_url: Final = f"https://comprehendmedical.{aws_region_name}.amazonaws.com/" + target_url: Final = f"https://comprehendmedical.{aws_region_name}.{get_aws_dns_suffix(aws_region_name)}/" _request: Final = AWSRequest(method="POST", url=target_url, data=json.dumps(data), headers=headers) sigv4.add_auth(_request) prepped: Final = _request.prepare() diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0f97473312c..80e991640a9 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1527,7 +1527,7 @@ class ProxyLogging: async def _process_prompt_template( self, data: dict, - litellm_logging_obj: Any, + litellm_logging_obj: "LiteLLMLoggingObj", prompt_id: str, prompt_version: int | None, call_type: CallTypesLiteral, @@ -3262,8 +3262,14 @@ class ProxyLogging: # through each of them adds N pass-through trampolines per chunk for # zero behavior change. Skip the chain entirely and stream through. if not caps.iterator_overrides: - async for chunk in response: - yield chunk + try: + async for chunk in response: + yield chunk + except (GeneratorExit, asyncio.CancelledError): + raise + except Exception: + ProxyLogging._fire_deferred_stream_logging(request_data) + raise ProxyLogging._fire_deferred_stream_logging(request_data) return @@ -3316,9 +3322,14 @@ class ProxyLogging: ), ) - # Actually iterate through the chained async generator and yield chunks - async for chunk in current_response: - yield chunk + try: + async for chunk in current_response: + yield chunk + except (GeneratorExit, asyncio.CancelledError): + raise + except Exception: + ProxyLogging._fire_deferred_stream_logging(request_data) + raise # Fire deferred logging AFTER all guardrail end-of-stream blocks # completed. unified_guardrail writes guardrail_information during @@ -5664,12 +5675,8 @@ class PrismaClient: return True acquire_task: Final = asyncio.create_task(_acquire_reconnect_lock()) - done, _pending = await asyncio.wait( - {acquire_task}, - timeout=lock_timeout_seconds, - return_when=asyncio.FIRST_COMPLETED, - ) - if acquire_task not in done: + + async def _abandon_acquire_task() -> None: acquire_task.cancel() try: await acquire_task @@ -5684,6 +5691,18 @@ class PrismaClient: self._db_reconnect_lock.release() except RuntimeError: pass + + try: + done, _pending = await asyncio.wait( + {acquire_task}, + timeout=lock_timeout_seconds, + return_when=asyncio.FIRST_COMPLETED, + ) + except asyncio.CancelledError: + await asyncio.shield(_abandon_acquire_task()) + raise + if acquire_task not in done: + await _abandon_acquire_task() verbose_proxy_logger.debug( "Skipping DB reconnect attempt due to lock acquisition timeout. reason=%s timeout=%ss", reason, diff --git a/litellm/rag/ingestion/bedrock_ingestion.py b/litellm/rag/ingestion/bedrock_ingestion.py index f721c204318..3d7056f8176 100644 --- a/litellm/rag/ingestion/bedrock_ingestion.py +++ b/litellm/rag/ingestion/bedrock_ingestion.py @@ -17,6 +17,7 @@ import uuid from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_logger +from litellm.litellm_core_utils.aws_partition import get_aws_arn_prefix from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.rag.ingestion.base_ingestion import BaseRAGIngestion @@ -52,11 +53,12 @@ def _normalize_principal_arn(caller_arn: str, account_id: str) -> str: """ if ":assumed-role/" in caller_arn: # Extract role name from assumed-role ARN - # Format: arn:aws:sts::ACCOUNT:assumed-role/ROLE-NAME/SESSION-NAME + # Format: arn:PARTITION:sts::ACCOUNT:assumed-role/ROLE-NAME/SESSION-NAME + partition: Final = caller_arn.split(":")[1] parts: Final = caller_arn.split("/") if len(parts) >= 2: role_name: Final = parts[1] - return f"arn:aws:iam::{account_id}:role/{role_name}" + return f"arn:{partition}:iam::{account_id}:role/{role_name}" return caller_arn @@ -294,7 +296,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): normalized_caller_arn: Final = _normalize_principal_arn(caller_arn, account_id) verbose_logger.debug("Caller ARN: %s, Normalized: %s", caller_arn, normalized_caller_arn) - principals = [f"arn:aws:iam::{account_id}:root", normalized_caller_arn] + principals = [f"{get_aws_arn_prefix(self.aws_region_name)}iam::{account_id}:root", normalized_caller_arn] # Deduplicate in case caller is root principals = list(set(principals)) @@ -454,7 +456,10 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): "Condition": { "StringEquals": {"aws:SourceAccount": account_id}, "ArnLike": { - "aws:SourceArn": f"arn:aws:bedrock:{self.aws_region_name}:{account_id}:knowledge-base/*" + "aws:SourceArn": ( + f"{get_aws_arn_prefix(self.aws_region_name)}bedrock:" + f"{self.aws_region_name}:{account_id}:knowledge-base/*" + ) }, }, } @@ -475,7 +480,10 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): { "Effect": "Allow", "Action": ["bedrock:InvokeModel"], - "Resource": [f"arn:aws:bedrock:{self.aws_region_name}::foundation-model/{self.embedding_model}"], + "Resource": [ + f"{get_aws_arn_prefix(self.aws_region_name)}bedrock:" + f"{self.aws_region_name}::foundation-model/{self.embedding_model}" + ], }, { "Effect": "Allow", @@ -486,8 +494,8 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): "Effect": "Allow", "Action": ["s3:GetObject", "s3:ListBucket"], "Resource": [ - f"arn:aws:s3:::{self.s3_bucket}", - f"arn:aws:s3:::{self.s3_bucket}/*", + f"{get_aws_arn_prefix(self.aws_region_name)}s3:::{self.s3_bucket}", + f"{get_aws_arn_prefix(self.aws_region_name)}s3:::{self.s3_bucket}/*", ], }, ], @@ -517,7 +525,10 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): knowledgeBaseConfiguration={ "type": "VECTOR", "vectorKnowledgeBaseConfiguration": { - "embeddingModelArn": f"arn:aws:bedrock:{self.aws_region_name}::foundation-model/{self.embedding_model}", + "embeddingModelArn": ( + f"{get_aws_arn_prefix(self.aws_region_name)}bedrock:" + f"{self.aws_region_name}::foundation-model/{self.embedding_model}" + ), }, }, storageConfiguration={ @@ -562,7 +573,7 @@ class BedrockRAGIngestion(BaseRAGIngestion, BaseAWSLLM): dataSourceConfiguration={ "type": "S3", "s3Configuration": { - "bucketArn": f"arn:aws:s3:::{self.s3_bucket}", + "bucketArn": f"{get_aws_arn_prefix(self.aws_region_name)}s3:::{self.s3_bucket}", "inclusionPrefixes": [self.s3_prefix], }, }, diff --git a/litellm/rag/ingestion/vertex_ai_ingestion.py b/litellm/rag/ingestion/vertex_ai_ingestion.py index 748c2b0d2b2..07f9f346d08 100644 --- a/litellm/rag/ingestion/vertex_ai_ingestion.py +++ b/litellm/rag/ingestion/vertex_ai_ingestion.py @@ -10,7 +10,10 @@ Based on: https://docs.cloud.google.com/vertex-ai/generative-ai/docs/model-refer from __future__ import annotations import json -from typing import TYPE_CHECKING, Any, Final +from collections.abc import Mapping +from typing import TYPE_CHECKING, Final + +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.llms.custom_httpx.http_handler import ( @@ -26,6 +29,42 @@ if TYPE_CHECKING: from litellm.types.rag import RAGIngestOptions +def _present_fields(fields: tuple[tuple[str, object], ...]) -> Mapping[str, object]: + return {name: value for name, value in fields if value} + + +class VertexRagResourceName(TypedDict, total=False): + name: ReadOnly[str] + + +class VertexRagOperation(TypedDict, total=False): + """A Vertex AI long-running operation resource, as the RAG Engine API returns it.""" + + done: ReadOnly[bool] + name: ReadOnly[str] + error: ReadOnly[object] + response: ReadOnly[VertexRagResourceName] + + +class VertexRagFileUpload(TypedDict, total=False): + """Body of a ``ragFiles:upload`` response.""" + + name: ReadOnly[str] + ragFile: ReadOnly[VertexRagResourceName] + + +class _RagOperationView(TypedDict): + """Holds one decoded long-running operation so the JSON body reads back typed.""" + + operation: ReadOnly[VertexRagOperation] + + +class _RagFileUploadView(TypedDict): + """Holds one decoded ``ragFiles:upload`` body so the JSON body reads back typed.""" + + upload: ReadOnly[VertexRagFileUpload] + + class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): """ Vertex AI RAG Engine ingestion implementation. @@ -148,26 +187,20 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): url: Final = f"{base_url}/v1beta1/projects/{self.project_id}/locations/{self.location}/ragCorpora" # Build request body with camelCase keys (Vertex AI API format) - request_body: Final[dict[str, Any]] = { - "displayName": display_name, - } - - if description: - request_body["description"] = description - - # Add vector database config if specified vector_db_config: Final = self.vector_store_config.get("vector_db_config") - if vector_db_config: - request_body["vectorDbConfig"] = vector_db_config - - # Add embedding model config if specified embedding_model: Final = self.vector_store_config.get("embedding_model") - if embedding_model: - if "vectorDbConfig" not in request_body: - request_body["vectorDbConfig"] = {} - request_body["vectorDbConfig"]["ragEmbeddingModelConfig"] = { - "vertexPredictionEndpoint": {"endpoint": embedding_model} - } + embedding_model_config: Final = ( + {"ragEmbeddingModelConfig": {"vertexPredictionEndpoint": {"endpoint": embedding_model}}} + if embedding_model + else None + ) + vector_db_section: Final = ( + {**(vector_db_config or {}), **embedding_model_config} if embedding_model_config else vector_db_config + ) + request_body: Final = { + "displayName": display_name, + **_present_fields((("description", description), ("vectorDbConfig", vector_db_section))), + } verbose_logger.debug("Creating RAG corpus: %s", url) verbose_logger.debug("Request body: %s", json.dumps(request_body, indent=2)) @@ -190,7 +223,8 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): verbose_logger.error(error_msg) raise Exception(error_msg) - response_data: Final = response.json() + operation_view: Final[_RagOperationView] = {"operation": response.json()} + response_data: Final = operation_view["operation"] verbose_logger.debug("Create corpus response: %s", json.dumps(response_data, indent=2)) # The response is a long-running operation @@ -257,12 +291,13 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): verbose_logger.error(error_msg) raise Exception(error_msg) - operation_data = response.json() + operation_view: _RagOperationView = {"operation": response.json()} + operation_data: VertexRagOperation = operation_view["operation"] if operation_data.get("done"): # Check for errors if "error" in operation_data: - error = operation_data["error"] + error = operation_data.get("error") raise Exception(f"Operation failed: {error}") # Extract corpus name from response @@ -308,39 +343,30 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): url: Final = f"{base_url}/upload/v1beta1/{rag_corpus_id}/ragFiles:upload" # Build metadata for the file with snake_case keys (as per upload API docs) - metadata: Final[dict[str, Any]] = { - "rag_file": { - "display_name": filename, - } + description: Final = self.vector_store_config.get("file_description") + rag_file: Final = { + "display_name": filename, + **_present_fields((("description", description),)), } - # Add description if provided - description: Final = self.vector_store_config.get("file_description") - if description: - metadata["rag_file"]["description"] = description - # Add chunking configuration if provided - chunking_strategy: Final = self.chunking_strategy - if chunking_strategy and isinstance(chunking_strategy, dict): - chunk_size: Final = chunking_strategy.get("chunk_size") - chunk_overlap: Final = chunking_strategy.get("chunk_overlap") - - if chunk_size or chunk_overlap: - if "upload_rag_file_config" not in metadata: - metadata["upload_rag_file_config"] = {} - - metadata["upload_rag_file_config"]["rag_file_transformation_config"] = { - "rag_file_chunking_config": {"fixed_length_chunking": {}} + chunking_strategy: Final[Mapping[str, object]] = self.chunking_strategy + chunk_size: Final = chunking_strategy.get("chunk_size") + chunk_overlap: Final = chunking_strategy.get("chunk_overlap") + fixed_length_chunking: Final = _present_fields((("chunk_size", chunk_size), ("chunk_overlap", chunk_overlap))) + upload_rag_file_config: Final = ( + { + "rag_file_transformation_config": { + "rag_file_chunking_config": {"fixed_length_chunking": fixed_length_chunking} } - - chunking_config: Final = metadata["upload_rag_file_config"]["rag_file_transformation_config"][ - "rag_file_chunking_config" - ]["fixed_length_chunking"] - - if chunk_size: - chunking_config["chunk_size"] = chunk_size - if chunk_overlap: - chunking_config["chunk_overlap"] = chunk_overlap + } + if fixed_length_chunking + else None + ) + metadata: Final = { + "rag_file": rag_file, + **_present_fields((("upload_rag_file_config", upload_rag_file_config),)), + } verbose_logger.debug("Uploading file to RAG corpus: %s", url) verbose_logger.debug("Metadata: %s", json.dumps(metadata, indent=2)) @@ -375,11 +401,11 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): # Parse response to get file ID try: - response_data: Final = response.json() + upload_view: Final[_RagFileUploadView] = {"upload": response.json()} + response_data: Final = upload_view["upload"] # The response should contain the rag_file resource name - file_id = response_data.get("ragFile", {}).get("name", "") - if not file_id: - file_id = response_data.get("name", "") + rag_file_name: Final = response_data.get("ragFile", {}).get("name", "") + file_id: Final = rag_file_name or response_data.get("name", "") verbose_logger.debug("Upload complete. File ID: %s", file_id) return file_id @@ -413,25 +439,30 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): base_url: Final = get_vertex_base_url(self.location) url: Final = f"{base_url}/v1beta1/{rag_corpus_id}/ragFiles:import" - # Build request body with camelCase keys (Vertex AI API format) - request_body: Final[dict[str, Any]] = {"importRagFilesConfig": {"gcsSource": {"uris": gcs_uris}}} - # Add chunking configuration if provided - chunking_strategy: Final = self.chunking_strategy - if chunking_strategy and isinstance(chunking_strategy, dict): - chunk_size: Final = chunking_strategy.get("chunk_size") - chunk_overlap: Final = chunking_strategy.get("chunk_overlap") - - if chunk_size or chunk_overlap: - request_body["importRagFilesConfig"]["ragFileChunkingConfig"] = { - "chunkSize": chunk_size or 1024, - "chunkOverlap": chunk_overlap or 200, - } + chunking_strategy: Final[Mapping[str, object]] = self.chunking_strategy + chunk_size: Final = chunking_strategy.get("chunk_size") + chunk_overlap: Final = chunking_strategy.get("chunk_overlap") # Add max embedding requests per minute if specified max_embedding_qpm: Final = self.vector_store_config.get("max_embedding_requests_per_min") - if max_embedding_qpm: - request_body["importRagFilesConfig"]["maxEmbeddingRequestsPerMin"] = max_embedding_qpm + + # Build request body with camelCase keys (Vertex AI API format) + chunking_config: Final = ( + {"chunkSize": chunk_size or 1024, "chunkOverlap": chunk_overlap or 200} + if chunk_size or chunk_overlap + else None + ) + import_config: Final = { + "gcsSource": {"uris": gcs_uris}, + **_present_fields( + ( + ("ragFileChunkingConfig", chunking_config), + ("maxEmbeddingRequestsPerMin", max_embedding_qpm), + ) + ), + } + request_body: Final = {"importRagFilesConfig": import_config} verbose_logger.debug("Importing files from GCS: %s", url) verbose_logger.debug("Request body: %s", json.dumps(request_body, indent=2)) @@ -455,7 +486,8 @@ class VertexAIRAGIngestion(BaseRAGIngestion, VertexBase): verbose_logger.error(error_msg) raise Exception(error_msg) - response_data: Final = response.json() + operation_view: Final[_RagOperationView] = {"operation": response.json()} + response_data: Final = operation_view["operation"] operation_name: Final = response_data.get("name", "") verbose_logger.debug("Import operation started: %s", operation_name) diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 3184c1bb0c7..924574537f3 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -252,7 +252,18 @@ def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[li return fallback_model_group, generic_fallback_idx -PROVIDER_SCOPED_RESOURCE_KEYS: Final = ("input_file_id", "training_file") +PROVIDER_SCOPED_RESOURCE_KEYS: Final = ("input_file_id", "training_file", "batch_id", "file_id", "fine_tuning_job_id") +PROVIDER_SCOPED_RESOURCE_FUNCTION_NAMES: Final = frozenset( + { + "_acreate_batch", + "_acancel_batch", + "acreate_fine_tuning_job", + "acancel_fine_tuning_job", + "aretrieve_fine_tuning_job", + "afile_content", + "afile_delete", + } +) PROVIDER_SCOPED_CREATION_FUNCTION_NAMES: Final = frozenset({"_acreate_file"}) @@ -284,13 +295,23 @@ async def _is_fallback_target_authorized( def references_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool: """ - True when the request names a file that only exists under one provider's credentials. + True when a file, batch, or fine-tuning job operation names an id that only exists + under one provider's credentials. - Batch and fine-tuning jobs are created from a file the caller already uploaded, and - that file lives in the account of the deployment that stored it. Handing the id to a - different model group can only fail, and the second provider's error replaces the - error the caller actually needs to see. + Each of those ids lives in the account of the deployment that issued it. Handing it to + a different model group asks a provider about an id it never issued, which costs an + extra round trip that can only answer not-found. Generic calls dispatched through + `Router._ageneric_api_call_with_fallbacks` carry the real handler in + `original_generic_function`, so both slots are checked. Gating on the handler name + keeps completion-style requests eligible for cross-group fallback even when a caller + passes a stray extra body field that happens to share one of these key names. """ + handler_names: Final = tuple( + getattr(kwargs.get(function_key), "__name__", None) + for function_key in ("original_function", "original_generic_function") + ) + if all(name not in PROVIDER_SCOPED_RESOURCE_FUNCTION_NAMES for name in handler_names): + return False return any(kwargs.get(key) for key in PROVIDER_SCOPED_RESOURCE_KEYS) @@ -371,7 +392,7 @@ async def run_async_fallback( continue if same_model_group_only and _get_fallback_target_model_group(mg) != original_model_group: verbose_router_logger.info( - "Skipping fallback to model_group = %s: request is pinned to model_group = %s by its uploaded file", + "Skipping fallback to model_group = %s: request names a resource owned by model_group = %s", mask_sensitive_structure(mg), original_model_group, ) diff --git a/litellm/rust_bridge/responses_websocket.py b/litellm/rust_bridge/responses_websocket.py index 20fc2634a8a..0634867af1c 100644 --- a/litellm/rust_bridge/responses_websocket.py +++ b/litellm/rust_bridge/responses_websocket.py @@ -3,7 +3,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any, Final, Protocol +from typing import Final, Protocol import httpx from websockets.exceptions import ConnectionClosedOK @@ -12,15 +12,22 @@ from litellm.rust_bridge.loader import get_native_bridge from litellm.rust_bridge.timeouts import timeout_to_seconds +class RustResponsesWebSocket(Protocol): + async def send_text(self, text: str) -> None: ... + + async def recv_text(self) -> str | None: ... + + async def close(self) -> None: ... + + class RustResponsesWebSocketConnection(Protocol): @classmethod - def connect( + async def connect( cls, url: str, headers: dict[str, str], timeout_seconds: float | None, - ) -> Any: - raise NotImplementedError + ) -> RustResponsesWebSocket: ... class _Unset: @@ -32,7 +39,7 @@ _UNSET: Final[_Unset] = _Unset() @dataclass(slots=True) class _RustResponsesWebSocketState: - connection: Any = None + connection: RustResponsesWebSocketConnection | None = None _STATE: Final[_RustResponsesWebSocketState] = _RustResponsesWebSocketState() @@ -40,27 +47,27 @@ _STATE: Final[_RustResponsesWebSocketState] = _RustResponsesWebSocketState() def set_rust_responses_websocket( *, - connection: Any = _UNSET, + connection: RustResponsesWebSocketConnection | None | _Unset = _UNSET, ) -> None: if not isinstance(connection, _Unset): _STATE.connection = connection -def load_rust_responses_websocket() -> Any: +def load_rust_responses_websocket() -> RustResponsesWebSocketConnection | None: if _STATE.connection is not None: return _STATE.connection native_bridge: Final = get_native_bridge() if native_bridge is None: return None - try: - return native_bridge.ResponsesWebSocketConnection - except AttributeError: - return None + connection_type: Final[RustResponsesWebSocketConnection | None] = getattr( + native_bridge, "ResponsesWebSocketConnection", None + ) + return connection_type class _ConnectionAdapter: - def __init__(self, connection: Any): - self._connection = connection + def __init__(self, connection: RustResponsesWebSocket): + self._connection: Final[RustResponsesWebSocket] = connection async def send(self, text: str) -> None: await self._connection.send_text(text) diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index 38a2ddd0bfc..2c7f1f8389d 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -22,12 +22,14 @@ import httpx import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, ) from litellm.proxy._types import KeyManagementSystem +from litellm.secret_managers.main import get_secret_str from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.secret_managers.main import KeyManagementSettings @@ -556,13 +558,15 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager): boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(optional_params) - # Get endpoint - _, endpoint_url = self.get_runtime_endpoint( - api_base=None, - aws_bedrock_runtime_endpoint=boto3_credentials_info.aws_bedrock_runtime_endpoint, - aws_region_name=boto3_credentials_info.aws_region_name, + region_name: Final = boto3_credentials_info.aws_region_name + explicit_runtime_endpoint: Final = boto3_credentials_info.aws_bedrock_runtime_endpoint or get_secret_str( + "AWS_BEDROCK_RUNTIME_ENDPOINT" + ) + endpoint_url: Final = ( + explicit_runtime_endpoint.replace("bedrock-runtime", "secretsmanager") + if explicit_runtime_endpoint + else f"https://secretsmanager.{region_name}.{get_aws_dns_suffix(region_name)}" ) - endpoint_url = endpoint_url.replace("bedrock-runtime", "secretsmanager") # Use provided request_data if available, otherwise build default data if request_data: diff --git a/litellm/secret_managers/secret_manager_handler.py b/litellm/secret_managers/secret_manager_handler.py index c77d2505d0e..f1f38c384cc 100644 --- a/litellm/secret_managers/secret_manager_handler.py +++ b/litellm/secret_managers/secret_manager_handler.py @@ -6,14 +6,64 @@ Handles retrieving secrets from different secret management systems. import base64 import os -from typing import Any, Final +from collections.abc import Mapping +from typing import Any, Final, Generic, Protocol, TypeVar + +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import print_verbose -from litellm.types.secret_managers.main import KeyManagementSystem +from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagementSystem + +_ClientT = TypeVar("_ClientT") -def _is_base64(s): +class _SecretManagerClientView(TypedDict, Generic[_ClientT]): + """Typed read of the untyped secret manager handle configured for this key manager.""" + + client: ReadOnly[_ClientT] + + +class _AzureKeyVaultSecret(Protocol): + @property + def value(self) -> str | None: ... + + +class _AzureKeyVaultClient(Protocol): + def get_secret(self, name: str) -> _AzureKeyVaultSecret: ... + + +class _GoogleKmsDecryptResponse(Protocol): + @property + def plaintext(self) -> bytes: ... + + +class _GoogleKmsClient(Protocol): + def decrypt(self, request: Mapping[str, object]) -> _GoogleKmsDecryptResponse: ... + + +class _AwsKmsClient(Protocol): + def decrypt(self, CiphertextBlob: bytes) -> Mapping[str, bytes]: ... + + +class _GoogleSecretManagerClient(Protocol): + def get_secret_from_google_secret_manager(self, secret_name: str) -> str | None: ... + + +class _SyncSecretReader(Protocol): + def sync_read_secret(self, secret_name: str) -> str | None: ... + + +class _InfisicalSecret(Protocol): + @property + def secret_value(self) -> str | None: ... + + +class _InfisicalClient(Protocol): + def get_secret(self, secret_name: str) -> _InfisicalSecret: ... + + +def _is_base64(s: str) -> bool: """Check if a string is valid base64.""" import binascii @@ -27,7 +77,7 @@ def get_secret_from_manager( client: Any, key_manager: str, secret_name: str, - key_management_settings: Any | None = None, + key_management_settings: KeyManagementSettings | None = None, ) -> str | None: """ Get a secret from the configured secret manager. @@ -46,34 +96,41 @@ def get_secret_from_manager( Exception: For other errors during secret retrieval """ secret = None + raw_view: Final[_SecretManagerClientView[object]] = {"client": client} + client_object: Final = raw_view["client"] if ( key_manager == KeyManagementSystem.AZURE_KEY_VAULT.value - or type(client).__module__ + "." + type(client).__name__ == "azure.keyvault.secrets._client.SecretClient" + or type(client_object).__module__ + "." + type(client_object).__name__ + == "azure.keyvault.secrets._client.SecretClient" ): # support Azure Secret Client - from azure.keyvault.secrets import SecretClient - secret = client.get_secret(secret_name).value + azure_view: Final[_SecretManagerClientView[_AzureKeyVaultClient]] = {"client": client} + azure_client: Final = azure_view["client"] + secret = azure_client.get_secret(secret_name).value elif ( - key_manager == KeyManagementSystem.GOOGLE_KMS.value or client.__class__.__name__ == "KeyManagementServiceClient" + key_manager == KeyManagementSystem.GOOGLE_KMS.value + or client_object.__class__.__name__ == "KeyManagementServiceClient" ): - encrypted_secret: Any = os.getenv(secret_name) + encrypted_secret: Final = os.getenv(secret_name) if encrypted_secret is None: raise ValueError("Google KMS requires the encrypted secret to be in the environment!") b64_flag: Final = _is_base64(encrypted_secret) if b64_flag is True: # if passed in as encoded b64 string - encrypted_secret = base64.b64decode(encrypted_secret) - ciphertext: Final = encrypted_secret + ciphertext: Final = base64.b64decode(encrypted_secret) else: raise ValueError( "Google KMS requires the encrypted secret to be encoded in base64" ) # fix for this vulnerability https://huntr.com/bounties/ae623c2f-b64b-4245-9ed4-f13a0a5824ce - response = client.decrypt( + google_kms_view: Final[_SecretManagerClientView[_GoogleKmsClient]] = {"client": client} + google_kms_client: Final = google_kms_view["client"] + google_kms_response: Final = google_kms_client.decrypt( request={ "name": litellm._google_kms_resource_name, "ciphertext": ciphertext, } ) - secret = response.plaintext.decode("utf-8") # assumes the original value was encoded with utf-8 + secret = google_kms_response.plaintext.decode("utf-8") # assumes the original value was encoded with utf-8 elif key_manager == KeyManagementSystem.AWS_KMS.value: """ @@ -85,13 +142,13 @@ def get_secret_from_manager( # Decode the base64 encoded ciphertext ciphertext_blob: Final = base64.b64decode(encrypted_value) - # Set up the parameters for the decrypt call - params: Final = {"CiphertextBlob": ciphertext_blob} # Perform the decryption - response = client.decrypt(**params) + aws_kms_view: Final[_SecretManagerClientView[_AwsKmsClient]] = {"client": client} + aws_kms_client: Final = aws_kms_view["client"] + aws_kms_response: Final = aws_kms_client.decrypt(CiphertextBlob=ciphertext_blob) # Extract and decode the plaintext - plaintext: Final = response["Plaintext"] + plaintext: Final = aws_kms_response["Plaintext"] secret = plaintext.decode("utf-8") if isinstance(secret, str): secret = secret.strip() @@ -114,7 +171,9 @@ def get_secret_from_manager( elif key_manager == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value: try: - secret = client.get_secret_from_google_secret_manager(secret_name) + google_secret_manager_view: Final[_SecretManagerClientView[_GoogleSecretManagerClient]] = {"client": client} + google_secret_manager_client: Final = google_secret_manager_view["client"] + secret = google_secret_manager_client.get_secret_from_google_secret_manager(secret_name) print_verbose(f"secret from google secret manager: [set={secret is not None}]") if secret is None: raise ValueError(f"No secret found in Google Secret Manager for {secret_name}") @@ -124,7 +183,9 @@ def get_secret_from_manager( elif key_manager == KeyManagementSystem.HASHICORP_VAULT.value: try: - secret = client.sync_read_secret(secret_name=secret_name) + hashicorp_view: Final[_SecretManagerClientView[_SyncSecretReader]] = {"client": client} + hashicorp_client: Final = hashicorp_view["client"] + secret = hashicorp_client.sync_read_secret(secret_name=secret_name) if secret is None: raise ValueError(f"No secret found in Hashicorp Secret Manager for {secret_name}") except Exception as e: @@ -133,7 +194,9 @@ def get_secret_from_manager( elif key_manager == KeyManagementSystem.CYBERARK.value: try: - secret = client.sync_read_secret(secret_name=secret_name) + cyberark_view: Final[_SecretManagerClientView[_SyncSecretReader]] = {"client": client} + cyberark_client: Final = cyberark_view["client"] + secret = cyberark_client.sync_read_secret(secret_name=secret_name) if secret is None: raise ValueError(f"No secret found in CyberArk Secret Manager for {secret_name}") except Exception as e: @@ -153,13 +216,16 @@ def get_secret_from_manager( raise ValueError(f"No secret found in Custom Secret Manager for {secret_name}") else: raise ValueError( - f"Custom secret manager client must be an instance of CustomSecretManager, got {type(client).__name__}" + "Custom secret manager client must be an instance of CustomSecretManager, " + f"got {type(client_object).__name__}" ) elif key_manager == "local": secret = os.getenv(secret_name) else: # assume the default is infisicial client - secret = client.get_secret(secret_name).secret_value + infisical_view: Final[_SecretManagerClientView[_InfisicalClient]] = {"client": client} + infisical_client: Final = infisical_view["client"] + secret = infisical_client.get_secret(secret_name).secret_value return secret diff --git a/litellm/types/agents.py b/litellm/types/agents.py index c85507b77c1..2cb42ce3fac 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal @@ -25,7 +26,7 @@ class AgentExtension(TypedDict, total=False): uri: str # required description: str | None required: bool | None - params: dict[str, Any] | None + params: dict[str, object] | None # AgentCapabilities @@ -70,10 +71,10 @@ class MutualTLSSecurityScheme(SecuritySchemeBase, total=False): class OAuthFlows(TypedDict, total=False): """Defines the configuration for the supported OAuth 2.0 flows.""" - authorizationCode: dict[str, Any] | None - clientCredentials: dict[str, Any] | None - implicit: dict[str, Any] | None - password: dict[str, Any] | None + authorizationCode: dict[str, object] | None + clientCredentials: dict[str, object] | None + implicit: dict[str, object] | None + password: dict[str, object] | None class OAuth2SecurityScheme(SecuritySchemeBase, total=False): @@ -129,7 +130,7 @@ class AgentCardSignature(TypedDict, total=False): protected: str # required signature: str # required - header: dict[str, Any] | None + header: dict[str, object] | None # AgentCard @@ -179,7 +180,7 @@ class AgentObjectPermission(TypedDict, total=False): class AgentConfig(TypedDict, total=False): agent_name: Required[str] agent_card_params: Required[AgentCard] - litellm_params: dict[str, Any] # allow for any future litellm params + litellm_params: dict[str, object] # allow for any future litellm params object_permission: AgentObjectPermission tpm_limit: int | None rpm_limit: int | None @@ -192,7 +193,7 @@ class AgentConfig(TypedDict, total=False): class PatchAgentRequest(TypedDict, total=False): agent_name: str agent_card_params: AgentCard - litellm_params: dict[str, Any] + litellm_params: dict[str, object] object_permission: AgentObjectPermission tpm_limit: int | None rpm_limit: int | None @@ -214,9 +215,9 @@ class AgentKeySummary(BaseModel): class AgentResponse(BaseModel): agent_id: str agent_name: str - litellm_params: dict[str, Any] | None = None + litellm_params: dict[str, object] | None = None agent_card_params: dict[str, Any] - object_permission: dict[str, Any] | None = None + object_permission: dict[str, object] | None = None spend: float | None = None tpm_limit: int | None = None rpm_limit: int | None = None @@ -251,7 +252,7 @@ class AgentCreateResponse(LiteLLMPydanticObjectBase): name: str | None = None model_config = {"extra": "allow"} - _hidden_params: dict = PrivateAttr(default_factory=dict) + _hidden_params: dict[str, object] = PrivateAttr(default_factory=dict) class AgentDeleteResult(LiteLLMPydanticObjectBase): @@ -265,7 +266,7 @@ class AgentDeleteResult(LiteLLMPydanticObjectBase): deleted: bool = True model_config = {"extra": "allow"} - _hidden_params: dict = PrivateAttr(default_factory=dict) + _hidden_params: dict[str, object] = PrivateAttr(default_factory=dict) class AgentListResponse(LiteLLMPydanticObjectBase): @@ -275,11 +276,11 @@ class AgentListResponse(LiteLLMPydanticObjectBase): a plain dict so no fields are silently dropped. """ - agents: list[dict[str, Any]] = [] + agents: list[dict[str, object]] = [] next_page_token: str | None = None model_config = {"extra": "allow"} - _hidden_params: dict = PrivateAttr(default_factory=dict) + _hidden_params: dict[str, object] = PrivateAttr(default_factory=dict) class AgentVersionsResponse(LiteLLMPydanticObjectBase): @@ -289,11 +290,11 @@ class AgentVersionsResponse(LiteLLMPydanticObjectBase): field of the form ``agents/{agent_id}/versions/{uuid}``. """ - agent_versions: list[dict[str, Any]] = [] + agent_versions: list[dict[str, object]] = [] next_page_token: str | None = None model_config = {"extra": "allow"} - _hidden_params: dict = PrivateAttr(default_factory=dict) + _hidden_params: dict[str, object] = PrivateAttr(default_factory=dict) class AgentMakePublicResponse(BaseModel): @@ -307,9 +308,9 @@ class MakeAgentsPublicRequest(BaseModel): def _normalize_a2a_jsonrpc_response( - response_dict: dict[str, Any], - request_id: Any | None = None, -) -> dict[str, Any]: + response_dict: Mapping[str, object], + request_id: object | None = None, +) -> dict[str, object]: """ Ensure JSON-RPC responses include ``id`` when the caller supplied one. @@ -347,22 +348,22 @@ class LiteLLMSendMessageResponse(LiteLLMPydanticObjectBase): # A2A response fields id: str | StrictInt | None = None jsonrpc: str = "2.0" - result: dict[str, Any] | None = None - error: dict[str, Any] | None = None + result: dict[str, object] | None = None + error: dict[str, object] | None = None # LiteLLM usage tracking - usage: dict[str, Any] | None = None + usage: dict[str, object] | None = None model_config = {"extra": "allow"} # LiteLLM private attributes for logging/cost tracking - _hidden_params: dict = PrivateAttr(default_factory=dict) + _hidden_params: dict[str, object] = PrivateAttr(default_factory=dict) @classmethod def from_a2a_response( cls, response: "SendMessageResponse", - request_id: Any | None = None, + request_id: object | None = None, ) -> "LiteLLMSendMessageResponse": """ Create a LiteLLMSendMessageResponse from an a2a SDK SendMessageResponse. @@ -377,13 +378,13 @@ class LiteLLMSendMessageResponse(LiteLLMPydanticObjectBase): response_dict: Final = _normalize_a2a_jsonrpc_response( response.model_dump(mode="json", exclude_none=True), request_id=request_id ) - return cls(**response_dict) + return cls.model_validate(response_dict) @classmethod def from_dict( cls, - response_dict: dict[str, Any], - request_id: Any | None = None, + response_dict: Mapping[str, object], + request_id: object | None = None, ) -> "LiteLLMSendMessageResponse": """ Create a LiteLLMSendMessageResponse from a dict. @@ -395,4 +396,4 @@ class LiteLLMSendMessageResponse(LiteLLMPydanticObjectBase): Returns: LiteLLMSendMessageResponse with _hidden_params support """ - return cls(**_normalize_a2a_jsonrpc_response(response_dict, request_id=request_id)) + return cls.model_validate(_normalize_a2a_jsonrpc_response(response_dict, request_id=request_id)) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 737361e5413..f0319a7c664 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -544,6 +544,8 @@ CallTypesLiteral = Literal[ "_arealtime", "create_batch", "acreate_batch", + "create_file", + "acreate_file", "pass_through_endpoint", "allm_passthrough_route", "anthropic_messages", diff --git a/litellm/utils.py b/litellm/utils.py index 5cd9bfc5f32..fa2226dbf2c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9122,6 +9122,10 @@ class ProviderConfigManager: from litellm.llms.runwayml.videos.transformation import RunwayMLVideoConfig return RunwayMLVideoConfig() + elif LlmProviders.HOSTED_VLLM == provider: + from litellm.llms.hosted_vllm.videos import get_hosted_vllm_video_config + + return get_hosted_vllm_video_config(model) return None @staticmethod diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 1d8d374c2c4..7c7d508856f 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -1277,7 +1277,8 @@ "files": true, "rerank": true, "a2a": true, - "interactions": true + "interactions": true, + "video_generations": true } }, "huggingface": { diff --git a/pyproject.toml b/pyproject.toml index eba9e5afc98..a0db4d49467 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,7 +69,7 @@ proxy = [ "mcp>=1.28.1,<2.0", "litellm-proxy-extras==0.4.90", "litellm-enterprise==0.1.61", - "RestrictedPython>=8.1,<9.0", + "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", "polars>=1.38.1,<2.0", diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 0418eeaac8f..c60988eccc0 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3014 + "limit": 3012 }, "ANN002": { "limit": 71 @@ -9,13 +9,13 @@ "limit": 827 }, "ANN201": { - "limit": 2011 + "limit": 2003 }, "ANN202": { - "limit": 847 + "limit": 845 }, "ANN204": { - "limit": 706 + "limit": 702 }, "ANN205": { "limit": 112 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1153 + "limit": 655 }, "ASYNC230": { "limit": 11 @@ -39,7 +39,7 @@ "limit": 503 }, "B009": { - "limit": 58 + "limit": 52 }, "B010": { "limit": 190 @@ -78,7 +78,7 @@ "limit": 1 }, "C901": { - "limit": 312 + "limit": 311 }, "D419": { "limit": 6 @@ -117,7 +117,7 @@ "limit": 1 }, "PERF102": { - "limit": 27 + "limit": 23 }, "PERF401": { "limit": 12 @@ -177,7 +177,7 @@ "limit": 8 }, "RUF019": { - "limit": 38 + "limit": 32 }, "RUF046": { "limit": 4 @@ -198,7 +198,7 @@ "limit": 58 }, "SIM102": { - "limit": 317 + "limit": 315 }, "SIM103": { "limit": 119 @@ -231,7 +231,7 @@ "limit": 5 }, "TID251": { - "limit": 1201 + "limit": 1117 }, "TRY002": { "limit": 524 @@ -246,7 +246,7 @@ "limit": 113 }, "TRY300": { - "limit": 859 + "limit": 857 }, "UP028": { "limit": 2 diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 3bd3c4d1d6c..842bfb4bdb1 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -16,11 +16,31 @@ longer signal it. ### Added +- **jwt_key_mapping**: New `litellm_jwt_key_mapping` resource for the proxy's JWT to virtual key mappings, so JWT clients identified by a claim (`client_id`, `azp`, `sub`) map to virtual keys and inherit their models, budgets and rate limits. Supports `description` and `is_active`, rotating the mapped key in place, and forces replacement when the claim name or value changes - **team**: `soft_budget`, `tags`, and `soft_budget_alerting_emails` attributes on `litellm_team`, matching what `/team/new` and `/team/update` already accept; `soft_budget_alerting_emails` is sent under `metadata`, where the proxy reads it +- **user**: New `litellm_user` resource and `litellm_user` / `litellm_users` data sources for managing internal users +- **budget**: New `litellm_budget` resource and `litellm_budget` / `litellm_budgets` data sources for reusable budget objects +- **tag**: New `litellm_tag` resource and `litellm_tag` / `litellm_tags` data sources for spend and routing tags +- **project**: New `litellm_project` resource and `litellm_project` / `litellm_projects` data sources +- **guardrail**: New `litellm_guardrail` resource and `litellm_guardrail` / `litellm_guardrails` data sources; `litellm_params` is sensitive and never read back into state +- **prompt**: New `litellm_prompt` resource and `litellm_prompt` / `litellm_prompts` data sources for prompt templates +- **agent**: New `litellm_agent` resource and `litellm_agent` / `litellm_agents` data sources for A2A agents +- **search_tool**: New `litellm_search_tool` resource and `litellm_search_tool` / `litellm_search_tools` data sources +- **access groups**: New `litellm_access_group` and `litellm_unified_access_group` resources with matching singular and plural data sources +- **fallback**: New `litellm_fallback` resource and data source for per-model fallbacks (general, context window and content policy) +- **block resources**: New `litellm_key_block` and `litellm_team_block` resources to manage the blocked state of existing keys and teams +- **data sources for existing resources**: New `litellm_key` / `litellm_keys`, `litellm_team` / `litellm_teams`, `litellm_model` / `litellm_models`, `litellm_organization` / `litellm_organizations` and `litellm_mcp_server` / `litellm_mcp_servers` data sources +- **key**: New arguments `budget_id`, `enforced_params`, `allowed_routes`, `allowed_passthrough_routes`, `rpm_limit_type`, `tpm_limit_type`, `prompts`, `organization_id` and `project_id` +- **team**: New arguments `model_aliases`, `guardrails`, `prompts`, `team_member_budget`, `team_member_budget_duration`, `team_member_rpm_limit`, `team_member_tpm_limit`, `team_member_key_duration`, `model_rpm_limit`, `model_tpm_limit`, `allowed_passthrough_routes`, `rpm_limit_type` and `tpm_limit_type` +- **import**: `terraform import` support for `litellm_team`, `litellm_model`, `litellm_organization`, `litellm_mcp_server`, `litellm_vector_store` and every new resource ### Fixed - **team**: Read now decodes the `team_info` envelope `/team/info` actually returns, so team attributes refresh from the proxy instead of always falling back to the prior state +- **key**: Read now unwraps the `info` envelope `/key/info` actually returns; previously reads mapped nothing back into state, so drift on a key was never detected +- **key**: Updates no longer send an empty `budget_duration`, which the proxy rejects with a 400; any update to a key without a configured `budget_duration` previously failed outright +- **key**: A config-supplied `key` value (write-only) is now forwarded to `/key/generate`; previously it was silently dropped and the proxy generated a random key instead +- **security**: The `litellm_key` data source and `litellm_key_block` resource normalize raw `sk-` keys to their SHA-256 token hash before building request URLs and resource IDs, so plaintext keys no longer land in reverse-proxy access logs, Terraform plan output, or state IDs ### Changed diff --git a/terraform/provider/README.md b/terraform/provider/README.md index fe67d6aa430..0a6d15c7844 100644 --- a/terraform/provider/README.md +++ b/terraform/provider/README.md @@ -1,10 +1,10 @@ # LiteLLM Terraform Provider -This Terraform provider allows you to manage LiteLLM resources through Infrastructure as Code. It provides support for managing models, teams, team members, and API keys via the LiteLLM REST API. +This Terraform provider allows you to manage LiteLLM resources through Infrastructure as Code. It provides support for managing models, teams, team members, API keys, users, organizations, budgets, tags, projects, guardrails, prompts, agents, search tools, access groups, fallbacks, MCP servers, credentials and vector stores via the LiteLLM REST API, along with read-only data sources for each of them. ## Source of truth -This directory (`terraform/provider/` in [BerriAI/litellm](https://github.com/BerriAI/litellm)) is the source of truth for the provider. [BerriAI/terraform-provider-litellm](https://github.com/BerriAI/terraform-provider-litellm) is a thin release mirror that the public Terraform Registry ingests from; do not open PRs there. Changes land here, where CI builds the provider, runs its tests, and statically audits every endpoint the provider calls against the proxy's generated OpenAPI schema (`tools/endpointaudit/`), so the provider cannot drift from the LiteLLM API silently. Releases are published by mirroring this directory into the split repo and tagging it, which triggers the goreleaser workflow there (see `RELEASING.md`) +This directory (`terraform/provider/` in [BerriAI/litellm](https://github.com/BerriAI/litellm)) is the source of truth for the provider. [BerriAI/terraform-provider-litellm](https://github.com/BerriAI/terraform-provider-litellm) is a thin release mirror that the public Terraform Registry ingests from; do not open PRs there. Changes land here, where CI builds the provider, runs its tests, and statically audits every endpoint the provider calls against the proxy's generated OpenAPI schema (`tools/endpointaudit/`), so the provider cannot drift from the LiteLLM API silently. The same audit runs in reverse as a coverage gate: every management endpoint in the schema must be covered by a resource or data source, or carry a documented entry in `tools/endpointaudit/coverage_allowlist.txt`, and stale allowlist entries fail CI. Releases are published by mirroring this directory into the split repo and tagging it, which triggers the goreleaser workflow there (see `RELEASING.md`) ## Versioning @@ -151,6 +151,7 @@ For full details on the litellm_key resource, see the [key resource - litellm_mcp_server: Manage MCP (Model Context Protocol) servers. [Documentation](docs/resources/mcp_server.md) - litellm_credential: Manage credentials for secure authentication. [Documentation](docs/resources/credential.md) - litellm_vector_store: Manage vector stores for embeddings and RAG. [Documentation](docs/resources/vector_store.md) +- litellm_jwt_key_mapping: Map JWT claim values to virtual keys for per-client budgets and limits. [Documentation](docs/resources/jwt_key_mapping.md) ### Available Data Sources diff --git a/terraform/provider/docs/data-sources/access_group.md b/terraform/provider/docs/data-sources/access_group.md new file mode 100644 index 00000000000..a1a8db8bd25 --- /dev/null +++ b/terraform/provider/docs/data-sources/access_group.md @@ -0,0 +1,34 @@ +--- +page_title: "litellm_access_group Data Source - terraform-provider-litellm" +subcategory: "" +description: |- + Retrieves information about an existing LiteLLM model access group. +--- + +# litellm_access_group (Data Source) + +Retrieves information about an existing LiteLLM model access group by name. + +## Example Usage + +```terraform +data "litellm_access_group" "production" { + access_group = "production-models" +} + +output "production_models" { + value = data.litellm_access_group.production.model_names +} +``` + +## Argument Reference + +* `access_group` - (Required) Name of the access group to look up. + +## Attribute Reference + +* `id` - The access group name. + +* `model_names` - List of model names in the access group. + +* `deployment_count` - Number of deployments tagged with this access group. diff --git a/terraform/provider/docs/data-sources/access_groups.md b/terraform/provider/docs/data-sources/access_groups.md new file mode 100644 index 00000000000..a81ac5d0772 --- /dev/null +++ b/terraform/provider/docs/data-sources/access_groups.md @@ -0,0 +1,33 @@ +--- +page_title: "litellm_access_groups Data Source - terraform-provider-litellm" +subcategory: "" +description: |- + Retrieves all LiteLLM model access groups. +--- + +# litellm_access_groups (Data Source) + +Retrieves all LiteLLM model access groups configured on the proxy. + +## Example Usage + +```terraform +data "litellm_access_groups" "all" {} + +output "access_group_names" { + value = data.litellm_access_groups.all.ids +} +``` + +## Argument Reference + +This data source takes no arguments. + +## Attribute Reference + +* `access_groups` - List of access groups. Each entry exports: + * `access_group` - The access group name. + * `model_names` - List of model names in the access group. + * `deployment_count` - Number of deployments tagged with this access group. + +* `ids` - List of all access group names. diff --git a/terraform/provider/docs/data-sources/agent.md b/terraform/provider/docs/data-sources/agent.md new file mode 100644 index 00000000000..09638ddd385 --- /dev/null +++ b/terraform/provider/docs/data-sources/agent.md @@ -0,0 +1,43 @@ +# litellm_agent Data Source + +Retrieves information about an existing A2A agent on the LiteLLM proxy. + +## Example Usage + +```hcl +data "litellm_agent" "existing" { + agent_id = "123e4567-e89b-12d3-a456-426614174000" +} + +output "agent_card" { + value = jsondecode(data.litellm_agent.existing.agent_card_params) +} +``` + +## Argument Reference + +The following arguments are supported: + +* `agent_id` - (Required) Unique identifier of the agent to retrieve. + +## Attribute Reference + +In addition to all arguments above, the following attributes are exported: + +* `agent_name` - Name of the agent. +* `agent_card_params` - The A2A agent card as a JSON object string (decode with `jsondecode`). +* `object_permission` - Access control permissions as a JSON object string. +* `extra_headers` - List of incoming request header names forwarded to the agent. +* `tpm_limit` - Tokens per minute limit. +* `rpm_limit` - Requests per minute limit. +* `session_tpm_limit` - Per-session tokens per minute limit. +* `session_rpm_limit` - Per-session requests per minute limit. +* `spend` - Total spend recorded for this agent. +* `created_at` - Timestamp when the agent was created. +* `updated_at` - Timestamp when the agent was last updated. +* `created_by` - User who created the agent. +* `updated_by` - User who last updated the agent. + +## Security Note + +`litellm_params` and `static_headers` are not exposed through this data source because they may hold API keys or tokens. diff --git a/terraform/provider/docs/data-sources/agents.md b/terraform/provider/docs/data-sources/agents.md new file mode 100644 index 00000000000..5b93780f307 --- /dev/null +++ b/terraform/provider/docs/data-sources/agents.md @@ -0,0 +1,42 @@ +# litellm_agents Data Source + +Retrieves the list of A2A agents registered on the LiteLLM proxy. + +## Example Usage + +```hcl +data "litellm_agents" "all" {} + +output "agent_ids" { + value = data.litellm_agents.all.ids +} + +# Only agents whose URL is currently reachable (or that have no URL) +data "litellm_agents" "healthy" { + health_check = true +} +``` + +## Argument Reference + +The following arguments are supported: + +* `health_check` - (Optional, default `false`) When true, the proxy probes each agent's URL and only returns agents that are reachable or have no URL. + +## Attribute Reference + +The following attributes are exported: + +* `ids` - List of agent IDs. +* `agents` - List of agents. Each entry exports: + * `agent_id` - The unique agent ID. + * `agent_name` - Name of the agent. + * `tpm_limit` - Tokens per minute limit. + * `rpm_limit` - Requests per minute limit. + * `session_tpm_limit` - Per-session tokens per minute limit. + * `session_rpm_limit` - Per-session requests per minute limit. + * `spend` - Total spend recorded for the agent. + * `created_at` - Timestamp when the agent was created. + * `updated_at` - Timestamp when the agent was last updated. + * `created_by` - User who created the agent. + * `updated_by` - User who last updated the agent. diff --git a/terraform/provider/docs/data-sources/budget.md b/terraform/provider/docs/data-sources/budget.md new file mode 100644 index 00000000000..b7c33df0a02 --- /dev/null +++ b/terraform/provider/docs/data-sources/budget.md @@ -0,0 +1,31 @@ +# litellm_budget Data Source + +Retrieves information about an existing LiteLLM budget by ID + +## Example Usage + +```hcl +data "litellm_budget" "engineering" { + budget_id = "engineering-monthly" +} + +output "engineering_max_budget" { + value = data.litellm_budget.engineering.max_budget +} +``` + +## Argument Reference + +- `budget_id` (Required) - ID of the budget to retrieve + +## Attribute Reference + +- `id` - The budget ID +- `max_budget` - Hard budget limit in USD +- `soft_budget` - Soft budget limit in USD that triggers alerts +- `max_parallel_requests` - Maximum concurrent requests allowed for this budget +- `tpm_limit` - Maximum tokens per minute allowed for this budget +- `rpm_limit` - Maximum requests per minute allowed for this budget +- `budget_duration` - Budget reset period +- `model_max_budget` - JSON string of per-model budget config +- `budget_reset_at` - Datetime when the budget is reset diff --git a/terraform/provider/docs/data-sources/budgets.md b/terraform/provider/docs/data-sources/budgets.md new file mode 100644 index 00000000000..c8dff98e390 --- /dev/null +++ b/terraform/provider/docs/data-sources/budgets.md @@ -0,0 +1,31 @@ +# litellm_budgets Data Source + +Retrieves all budgets configured on the LiteLLM proxy + +## Example Usage + +```hcl +data "litellm_budgets" "all" {} + +output "budget_ids" { + value = data.litellm_budgets.all.ids +} +``` + +## Argument Reference + +This data source takes no arguments + +## Attribute Reference + +- `budgets` - All budgets configured on the proxy. Each entry has: + - `budget_id` - The budget ID + - `max_budget` - Hard budget limit in USD + - `soft_budget` - Soft budget limit in USD that triggers alerts + - `max_parallel_requests` - Maximum concurrent requests allowed for this budget + - `tpm_limit` - Maximum tokens per minute allowed for this budget + - `rpm_limit` - Maximum requests per minute allowed for this budget + - `budget_duration` - Budget reset period + - `model_max_budget` - JSON string of per-model budget config + - `budget_reset_at` - Datetime when the budget is reset +- `ids` - IDs of all budgets configured on the proxy diff --git a/terraform/provider/docs/data-sources/fallback.md b/terraform/provider/docs/data-sources/fallback.md new file mode 100644 index 00000000000..856bee6eb79 --- /dev/null +++ b/terraform/provider/docs/data-sources/fallback.md @@ -0,0 +1,38 @@ +# litellm_fallback (Data Source) + +Retrieves the fallback configuration for a LiteLLM model. Use this to reference fallbacks that were configured outside of Terraform. + +## Example Usage + +```hcl +data "litellm_fallback" "gpt4" { + model = "gpt-4" +} + +output "gpt4_fallback_models" { + value = data.litellm_fallback.gpt4.fallback_models +} +``` + +### Specific Fallback Type + +```hcl +data "litellm_fallback" "gpt4_context_window" { + model = "gpt-4" + fallback_type = "context_window" +} +``` + +## Argument Reference + +The following arguments are supported: + +* `model` - (Required) The model name to get fallbacks for. +* `fallback_type` - (Optional) Type of fallback to retrieve. One of `general` (default), `context_window`, or `content_policy`. + +## Attribute Reference + +In addition to the arguments above, the following attributes are exported: + +* `id` - The primary model name. +* `fallback_models` - List of fallback model names in order of priority. diff --git a/terraform/provider/docs/data-sources/guardrail.md b/terraform/provider/docs/data-sources/guardrail.md new file mode 100644 index 00000000000..a54c652277a --- /dev/null +++ b/terraform/provider/docs/data-sources/guardrail.md @@ -0,0 +1,27 @@ +# litellm_guardrail Data Source + +Retrieves information about an existing LiteLLM guardrail by ID. Sensitive `litellm_params` are not exposed. + +## Example Usage + +```hcl +data "litellm_guardrail" "existing" { + guardrail_id = "123e4567-e89b-12d3-a456-426614174000" +} + +output "guardrail_name" { + value = data.litellm_guardrail.existing.guardrail_name +} +``` + +## Argument Reference + +* `guardrail_id` - (Required) Unique identifier of the guardrail to retrieve. + +## Attribute Reference + +* `guardrail_name` - Human-readable name of the guardrail. +* `guardrail_info` - Map of additional metadata for the guardrail. +* `guardrail_definition_location` - Where the guardrail is defined: `config` or `db`. +* `created_at` - Timestamp when the guardrail was created. +* `updated_at` - Timestamp when the guardrail was last updated. diff --git a/terraform/provider/docs/data-sources/guardrails.md b/terraform/provider/docs/data-sources/guardrails.md new file mode 100644 index 00000000000..589690cbb52 --- /dev/null +++ b/terraform/provider/docs/data-sources/guardrails.md @@ -0,0 +1,32 @@ +# litellm_guardrails Data Source + +Retrieves the list of all guardrails configured on the LiteLLM proxy (from both config and DB). Sensitive `litellm_params` are not exposed. + +## Example Usage + +```hcl +data "litellm_guardrails" "all" {} + +output "guardrail_ids" { + value = data.litellm_guardrails.all.ids +} + +output "guardrail_names" { + value = [for g in data.litellm_guardrails.all.guardrails : g.guardrail_name] +} +``` + +## Argument Reference + +This data source takes no arguments. + +## Attribute Reference + +* `guardrails` - List of guardrails. Each entry contains: + * `guardrail_id` - Unique identifier of the guardrail. + * `guardrail_name` - Human-readable name of the guardrail. + * `guardrail_info` - Map of additional metadata for the guardrail. + * `guardrail_definition_location` - Where the guardrail is defined: `config` or `db`. + * `created_at` - Timestamp when the guardrail was created. + * `updated_at` - Timestamp when the guardrail was last updated. +* `ids` - List of all guardrail IDs. diff --git a/terraform/provider/docs/data-sources/key.md b/terraform/provider/docs/data-sources/key.md new file mode 100644 index 00000000000..c11a90c4a48 --- /dev/null +++ b/terraform/provider/docs/data-sources/key.md @@ -0,0 +1,57 @@ +--- +# generated by https://github.com/hashicorp/terraform-plugin-docs +page_title: "litellm_key Data Source - terraform-provider-litellm" +subcategory: "" +description: |- + Retrieves information about an existing LiteLLM API key. +--- + +# litellm_key (Data Source) + +Retrieves information about an existing LiteLLM API key via `/key/info`. Pass either the raw key or its hashed token. The raw key value is never written to state beyond the input you provide; the data source ID is the hashed token. + +## Example Usage + +```terraform +data "litellm_key" "ci" { + key = var.ci_key_hash +} + +output "ci_key_team" { + value = data.litellm_key.ci.team_id +} +``` + +## Argument Reference + +The following arguments are supported: + +* `key` - (Required, Sensitive) The API key (or its hash) to look up. + +## Attributes Reference + +In addition to all arguments above, the following attributes are exported: + +* `token_id` - Hashed token identifier of the key (safe to store in state). +* `key_name` - Redacted display name of the key. +* `key_alias` - User-friendly alias for the key. +* `models` - List of models this key can access. +* `spend` - Amount spent by this key. +* `max_budget` - Maximum budget for this key. +* `user_id` - User ID associated with this key. +* `team_id` - Team ID associated with this key. +* `organization_id` - Organization ID associated with this key. +* `tpm_limit` - Tokens per minute limit. +* `rpm_limit` - Requests per minute limit. +* `max_parallel_requests` - Maximum parallel requests allowed. +* `budget_duration` - Budget reset duration. +* `metadata` - Map of string metadata values for the key. +* `tags` - Tags attached to the key. +* `blocked` - Whether the key is blocked. +* `expires` - Expiry timestamp, if set. +* `created_at` - Timestamp when the key was created. +* `updated_at` - Timestamp when the key was last updated. + +## Security Note + +The raw key value is only used to perform the lookup; it is never exported as an attribute or used as the data source ID. diff --git a/terraform/provider/docs/data-sources/keys.md b/terraform/provider/docs/data-sources/keys.md new file mode 100644 index 00000000000..24e187ec541 --- /dev/null +++ b/terraform/provider/docs/data-sources/keys.md @@ -0,0 +1,62 @@ +--- +# generated by https://github.com/hashicorp/terraform-plugin-docs +page_title: "litellm_keys Data Source - terraform-provider-litellm" +subcategory: "" +description: |- + Lists LiteLLM API keys with optional server-side filters. +--- + +# litellm_keys (Data Source) + +Lists LiteLLM API keys via `/key/list`. Supports server-side filtering and pagination. Raw key values are never returned; each entry is identified by its hashed token. + +## Example Usage + +```terraform +data "litellm_keys" "team_keys" { + team_id = litellm_team.ml.id + size = 50 +} + +output "team_key_aliases" { + value = [for k in data.litellm_keys.team_keys.keys : k.key_alias] +} +``` + +## Argument Reference + +The following arguments are supported: + +* `page` - (Optional) Page number for pagination. Defaults to `1`. +* `size` - (Optional) Number of keys per page. Defaults to `100`. +* `user_id` - (Optional) Filter keys by user ID. +* `team_id` - (Optional) Filter keys by team ID. +* `organization_id` - (Optional) Filter keys by organization ID. +* `key_alias` - (Optional) Filter keys by key alias. +* `include_team_keys` - (Optional) Include all keys for teams the caller is an admin of. + +## Attributes Reference + +In addition to all arguments above, the following attributes are exported: + +* `total_count` - Total number of keys matching the filters. +* `total_pages` - Total number of pages. +* `current_page` - The page returned. +* `ids` - Hashed token identifiers of the returned keys. +* `keys` - List of key objects. Each entry exports: + * `token_id` - Hashed token identifier. + * `key_name` - Redacted display name. + * `key_alias` - User-friendly alias. + * `spend` - Amount spent by the key. + * `max_budget` - Maximum budget. + * `models` - Models the key can access. + * `user_id` - Associated user ID. + * `team_id` - Associated team ID. + * `organization_id` - Associated organization ID. + * `tpm_limit` - Tokens per minute limit. + * `rpm_limit` - Requests per minute limit. + * `budget_duration` - Budget reset duration. + * `blocked` - Whether the key is blocked. + * `expires` - Expiry timestamp, if set. + * `created_at` - Creation timestamp. + * `updated_at` - Last update timestamp. diff --git a/terraform/provider/docs/data-sources/mcp_server.md b/terraform/provider/docs/data-sources/mcp_server.md new file mode 100644 index 00000000000..412d0a77fbe --- /dev/null +++ b/terraform/provider/docs/data-sources/mcp_server.md @@ -0,0 +1,58 @@ +--- +# generated by https://github.com/hashicorp/terraform-plugin-docs +page_title: "litellm_mcp_server Data Source - terraform-provider-litellm" +subcategory: "" +description: |- + Retrieves information about an existing LiteLLM MCP server. +--- + +# litellm_mcp_server (Data Source) + +Retrieves information about an existing MCP server via `/v1/mcp/server/{server_id}`. Secret material (environment variables, credentials, and static header values) is never exposed. + +## Example Usage + +```terraform +data "litellm_mcp_server" "github" { + server_id = "srv-1234" +} + +output "github_mcp_url" { + value = data.litellm_mcp_server.github.url +} +``` + +## Argument Reference + +The following arguments are supported: + +* `server_id` - (Required) Unique identifier of the MCP server to retrieve. + +## Attributes Reference + +In addition to all arguments above, the following attributes are exported: + +* `server_name` - Name of the MCP server. +* `alias` - Alias for the MCP server. +* `description` - Description of the MCP server. +* `url` - URL of the MCP server. +* `transport` - Transport type (`http`, `sse`, `stdio`). +* `spec_version` - MCP specification version. +* `auth_type` - Authentication type (`none`, `bearer`, `basic`, ...). +* `mcp_access_groups` - Access groups for the MCP server. +* `allowed_tools` - Tools allowed on this server. +* `extra_headers` - Names of request headers forwarded to the MCP server. +* `command` - Command for stdio transport. +* `args` - Arguments for the command (stdio transport). +* `allow_all_keys` - Whether all keys can access the server. +* `status` - Health status (`healthy`, `unhealthy`, `unknown`). +* `last_health_check` - Timestamp of the last health check. +* `health_check_error` - Error message from the last health check, if any. +* `created_at` - Timestamp when the server was created. +* `created_by` - User who created the server. +* `updated_at` - Timestamp when the server was last updated. +* `updated_by` - User who last updated the server. + +## Security Note + +For security reasons, `env`, `credentials`, and `static_headers` are not exposed through this data source since they may hold secrets. diff --git a/terraform/provider/docs/data-sources/mcp_servers.md b/terraform/provider/docs/data-sources/mcp_servers.md new file mode 100644 index 00000000000..fac50d610d6 --- /dev/null +++ b/terraform/provider/docs/data-sources/mcp_servers.md @@ -0,0 +1,50 @@ +--- +# generated by https://github.com/hashicorp/terraform-plugin-docs +page_title: "litellm_mcp_servers Data Source - terraform-provider-litellm" +subcategory: "" +description: |- + Lists LiteLLM MCP servers. +--- + +# litellm_mcp_servers (Data Source) + +Lists MCP servers via `/v1/mcp/server`. Secret material is never exposed. + +## Example Usage + +```terraform +data "litellm_mcp_servers" "all" {} + +data "litellm_mcp_servers" "team_scoped" { + team_id = litellm_team.ml.id +} + +output "mcp_server_urls" { + value = [for s in data.litellm_mcp_servers.all.mcp_servers : s.url] +} +``` + +## Argument Reference + +The following arguments are supported: + +* `team_id` - (Optional) Filter to servers this team can access plus globally available (`allow_all_keys`) servers. + +## Attributes Reference + +In addition to all arguments above, the following attributes are exported: + +* `ids` - IDs of the returned MCP servers. +* `mcp_servers` - List of MCP server objects. Each entry exports: + * `server_id` - Unique identifier of the MCP server. + * `server_name` - Name of the MCP server. + * `alias` - Alias for the MCP server. + * `description` - Description of the MCP server. + * `url` - URL of the MCP server. + * `transport` - Transport type (`http`, `sse`, `stdio`). + * `spec_version` - MCP specification version. + * `auth_type` - Authentication type. + * `allow_all_keys` - Whether all keys can access the server. + * `status` - Health status (`healthy`, `unhealthy`, `unknown`). + * `created_at` - Creation timestamp. + * `updated_at` - Last update timestamp. diff --git a/terraform/provider/docs/data-sources/model.md b/terraform/provider/docs/data-sources/model.md new file mode 100644 index 00000000000..6976ff1523a --- /dev/null +++ b/terraform/provider/docs/data-sources/model.md @@ -0,0 +1,50 @@ +--- +# generated by https://github.com/hashicorp/terraform-plugin-docs +page_title: "litellm_model Data Source - terraform-provider-litellm" +subcategory: "" +description: |- + Retrieves information about a model deployment on the LiteLLM proxy. +--- + +# litellm_model (Data Source) + +Retrieves information about a single model deployment via `/v1/model/info`. Sensitive `litellm_params` fields (API keys and other credentials) are never exposed; only safe routing metadata is exported. + +## Example Usage + +```terraform +data "litellm_model" "gpt4o" { + model_id = "0e5x74fab24a7a5245d2ced3536dd8f5" +} + +output "gpt4o_provider" { + value = data.litellm_model.gpt4o.custom_llm_provider +} +``` + +## Argument Reference + +The following arguments are supported: + +* `model_id` - (Required) LiteLLM model ID (the `x-litellm-model-id` response header value). + +## Attributes Reference + +In addition to all arguments above, the following attributes are exported: + +* `model_name` - Public model name used for routing. +* `model` - The underlying `litellm_params` model, e.g. `openai/gpt-4o`. +* `custom_llm_provider` - Provider for the model. +* `model_api_base` - API base URL, if configured. +* `api_version` - API version, if configured. +* `tpm` - Tokens per minute limit for the deployment. +* `rpm` - Requests per minute limit for the deployment. +* `base_model` - Base model used for pricing and capabilities. +* `tier` - Model tier (`free` or `paid`). +* `mode` - Model mode, e.g. `chat` or `embedding`. +* `team_id` - Team the deployment is scoped to, if any. +* `db_model` - Whether the deployment is stored in the database (as opposed to config). + +## Security Note + +Credential material inside `litellm_params` (such as `api_key`, `aws_secret_access_key`, and `vertex_credentials`) is never exported by this data source. diff --git a/terraform/provider/docs/data-sources/models.md b/terraform/provider/docs/data-sources/models.md new file mode 100644 index 00000000000..7862dc30ab7 --- /dev/null +++ b/terraform/provider/docs/data-sources/models.md @@ -0,0 +1,44 @@ +--- +# generated by https://github.com/hashicorp/terraform-plugin-docs +page_title: "litellm_models Data Source - terraform-provider-litellm" +subcategory: "" +description: |- + Lists model deployments on the LiteLLM proxy. +--- + +# litellm_models (Data Source) + +Lists all model deployments via `/v1/model/info`. Sensitive `litellm_params` fields (API keys and other credentials) are never exposed. + +## Example Usage + +```terraform +data "litellm_models" "all" {} + +output "model_names" { + value = [for m in data.litellm_models.all.models : m.model_name] +} +``` + +## Argument Reference + +The following arguments are supported: + +* `team_id` - (Optional) Filter models to those accessible by this team. + +## Attributes Reference + +In addition to all arguments above, the following attributes are exported: + +* `ids` - LiteLLM model IDs of the returned models. +* `models` - List of model objects. Each entry exports: + * `id` - LiteLLM model ID. + * `model_name` - Public model name used for routing. + * `model` - The underlying `litellm_params` model. + * `custom_llm_provider` - Provider for the model. + * `model_api_base` - API base URL, if configured. + * `base_model` - Base model used for pricing and capabilities. + * `tier` - Model tier (`free` or `paid`). + * `mode` - Model mode, e.g. `chat` or `embedding`. + * `team_id` - Team the deployment is scoped to, if any. + * `db_model` - Whether the deployment is stored in the database. diff --git a/terraform/provider/docs/data-sources/organization.md b/terraform/provider/docs/data-sources/organization.md new file mode 100644 index 00000000000..acc303cdf6e --- /dev/null +++ b/terraform/provider/docs/data-sources/organization.md @@ -0,0 +1,48 @@ +--- +# generated by https://github.com/hashicorp/terraform-plugin-docs +page_title: "litellm_organization Data Source - terraform-provider-litellm" +subcategory: "" +description: |- + Retrieves information about an existing LiteLLM organization. +--- + +# litellm_organization (Data Source) + +Retrieves information about an existing LiteLLM organization via `/organization/info`, including its attached budget settings. + +## Example Usage + +```terraform +data "litellm_organization" "main" { + organization_id = "org-1234" +} + +resource "litellm_team" "ml" { + team_alias = "ml-team" + organization_id = data.litellm_organization.main.organization_id +} +``` + +## Argument Reference + +The following arguments are supported: + +* `organization_id` - (Required) Unique identifier of the organization to retrieve. + +## Attributes Reference + +In addition to all arguments above, the following attributes are exported: + +* `organization_alias` - User-friendly name of the organization. +* `budget_id` - ID of the attached budget. +* `models` - Models the organization can access. +* `spend` - Amount spent by the organization. +* `metadata` - Map of string metadata values for the organization. +* `max_budget` - Maximum budget from the attached budget. +* `soft_budget` - Soft budget alert threshold from the attached budget. +* `tpm_limit` - Tokens per minute limit from the attached budget. +* `rpm_limit` - Requests per minute limit from the attached budget. +* `max_parallel_requests` - Maximum parallel requests from the attached budget. +* `budget_duration` - Budget reset duration from the attached budget. +* `created_at` - Timestamp when the organization was created. +* `updated_at` - Timestamp when the organization was last updated. diff --git a/terraform/provider/docs/data-sources/organizations.md b/terraform/provider/docs/data-sources/organizations.md new file mode 100644 index 00000000000..72e9ff8c916 --- /dev/null +++ b/terraform/provider/docs/data-sources/organizations.md @@ -0,0 +1,45 @@ +--- +# generated by https://github.com/hashicorp/terraform-plugin-docs +page_title: "litellm_organizations Data Source - terraform-provider-litellm" +subcategory: "" +description: |- + Lists LiteLLM organizations. +--- + +# litellm_organizations (Data Source) + +Lists LiteLLM organizations via `/organization/list`. + +## Example Usage + +```terraform +data "litellm_organizations" "all" {} + +output "organization_ids" { + value = data.litellm_organizations.all.ids +} +``` + +## Argument Reference + +The following arguments are supported: + +* `org_alias` - (Optional) Filter organizations by alias. + +## Attributes Reference + +In addition to all arguments above, the following attributes are exported: + +* `ids` - IDs of the returned organizations. +* `organizations` - List of organization objects. Each entry exports: + * `organization_id` - Unique identifier of the organization. + * `organization_alias` - User-friendly name of the organization. + * `budget_id` - ID of the attached budget. + * `models` - Models the organization can access. + * `spend` - Amount spent by the organization. + * `max_budget` - Maximum budget from the attached budget. + * `tpm_limit` - Tokens per minute limit from the attached budget. + * `rpm_limit` - Requests per minute limit from the attached budget. + * `budget_duration` - Budget reset duration from the attached budget. + * `created_at` - Creation timestamp. + * `updated_at` - Last update timestamp. diff --git a/terraform/provider/docs/data-sources/project.md b/terraform/provider/docs/data-sources/project.md new file mode 100644 index 00000000000..fb46cb98e86 --- /dev/null +++ b/terraform/provider/docs/data-sources/project.md @@ -0,0 +1,43 @@ +# litellm_project (Data Source) + +Retrieves information about an existing LiteLLM project, including its budget settings + +## Example Usage + +```hcl +data "litellm_project" "ml_experiments" { + project_id = "4a422a4c-e246-4d02-a1eb-13e835cd0725" +} + +output "project_spend" { + value = data.litellm_project.ml_experiments.spend +} +``` + +## Argument Reference + +The following arguments are supported: + +* `project_id` - (Required) Unique identifier of the project to retrieve + +## Attribute Reference + +In addition to all arguments above, the following attributes are exported: + +* `project_alias` - Human-friendly name for the project +* `description` - Description of the project +* `team_id` - The team ID this project belongs to +* `budget_id` - Budget ID associated with this project +* `models` - List of models the project can access +* `max_budget` - Maximum budget for this project +* `soft_budget` - Soft budget limit for warnings +* `budget_duration` - Budget reset duration +* `tpm_limit` - Tokens per minute limit +* `rpm_limit` - Requests per minute limit +* `max_parallel_requests` - Maximum parallel requests allowed +* `blocked` - Whether the project is blocked from making requests +* `spend` - Current spend for the project +* `created_at` - Timestamp when the project was created +* `updated_at` - Timestamp when the project was last updated +* `created_by` - User that created the project +* `updated_by` - User that last updated the project diff --git a/terraform/provider/docs/data-sources/projects.md b/terraform/provider/docs/data-sources/projects.md new file mode 100644 index 00000000000..1b53b327ae4 --- /dev/null +++ b/terraform/provider/docs/data-sources/projects.md @@ -0,0 +1,40 @@ +# litellm_projects (Data Source) + +Retrieves the list of all LiteLLM projects visible to the caller + +## Example Usage + +```hcl +data "litellm_projects" "all" {} + +output "project_ids" { + value = data.litellm_projects.all.ids +} + +output "project_aliases" { + value = [for p in data.litellm_projects.all.projects : p.project_alias] +} +``` + +## Argument Reference + +This data source takes no arguments + +## Attribute Reference + +The following attributes are exported: + +* `ids` - IDs of all projects +* `projects` - List of projects. Each entry exports: + * `project_id` - The project ID + * `project_alias` - Human-friendly name for the project + * `description` - Description of the project + * `team_id` - The team ID this project belongs to + * `budget_id` - Budget ID associated with this project + * `models` - List of models the project can access + * `blocked` - Whether the project is blocked from making requests + * `spend` - Current spend for the project + * `created_at` - Timestamp when the project was created + * `updated_at` - Timestamp when the project was last updated + * `created_by` - User that created the project + * `updated_by` - User that last updated the project diff --git a/terraform/provider/docs/data-sources/prompt.md b/terraform/provider/docs/data-sources/prompt.md new file mode 100644 index 00000000000..aa9e1e27148 --- /dev/null +++ b/terraform/provider/docs/data-sources/prompt.md @@ -0,0 +1,43 @@ +# litellm_prompt Data Source + +Retrieves information about an existing LiteLLM prompt by ID. The provider API key is not exposed. + +## Example Usage + +```hcl +data "litellm_prompt" "existing" { + prompt_id = "my-langfuse-prompt" +} + +output "prompt_integration" { + value = data.litellm_prompt.existing.prompt_integration +} +``` + +### With Environment + +```hcl +data "litellm_prompt" "prod" { + prompt_id = "my-langfuse-prompt" + environment = "production" +} +``` + +## Argument Reference + +* `prompt_id` - (Required) Unique identifier of the prompt to retrieve. +* `environment` - (Optional) Environment to fetch the prompt from (e.g. `development`, `production`). + +## Attribute Reference + +* `prompt_integration` - The prompt integration provider. +* `api_base` - Base URL for the prompt provider API. +* `provider_specific_query_params` - JSON string of provider-specific query parameters. +* `ignore_prompt_manager_model` - Whether the model specified in the prompt manager is ignored. +* `ignore_prompt_manager_optional_params` - Whether optional params from the prompt manager are ignored. +* `dotprompt_content` - Content for the dotprompt integration. +* `prompt_type` - Type of prompt: `config` or `db`. +* `version` - Version number of the prompt. +* `environments` - List of environments this prompt exists in. +* `created_at` - Timestamp when the prompt was created. +* `updated_at` - Timestamp when the prompt was last updated. diff --git a/terraform/provider/docs/data-sources/prompts.md b/terraform/provider/docs/data-sources/prompts.md new file mode 100644 index 00000000000..c433750b40f --- /dev/null +++ b/terraform/provider/docs/data-sources/prompts.md @@ -0,0 +1,37 @@ +# litellm_prompts Data Source + +Retrieves the list of all prompts configured on the LiteLLM proxy. + +## Example Usage + +```hcl +data "litellm_prompts" "all" {} + +output "prompt_ids" { + value = data.litellm_prompts.all.ids +} +``` + +### Filter by Environment + +```hcl +data "litellm_prompts" "production" { + environment = "production" +} +``` + +## Argument Reference + +* `environment` - (Optional) Filter prompts by environment (e.g. `development`, `production`). + +## Attribute Reference + +* `prompts` - List of prompts. Each entry contains: + * `prompt_id` - Unique identifier of the prompt. + * `prompt_integration` - The prompt integration provider. + * `prompt_type` - Type of prompt: `config` or `db`. + * `version` - Version number of the prompt. + * `environment` - Environment the prompt belongs to. + * `created_at` - Timestamp when the prompt was created. + * `updated_at` - Timestamp when the prompt was last updated. +* `ids` - List of all prompt IDs. diff --git a/terraform/provider/docs/data-sources/search_tool.md b/terraform/provider/docs/data-sources/search_tool.md new file mode 100644 index 00000000000..42dd73500c9 --- /dev/null +++ b/terraform/provider/docs/data-sources/search_tool.md @@ -0,0 +1,34 @@ +# litellm_search_tool Data Source + +Retrieves information about an existing search tool on the LiteLLM proxy. + +## Example Usage + +```hcl +data "litellm_search_tool" "existing" { + search_tool_id = "123e4567-e89b-12d3-a456-426614174000" +} + +output "search_tool_name" { + value = data.litellm_search_tool.existing.search_tool_name +} +``` + +## Argument Reference + +The following arguments are supported: + +* `search_tool_id` - (Required) Unique identifier of the search tool to retrieve. + +## Attribute Reference + +In addition to all arguments above, the following attributes are exported: + +* `search_tool_name` - Name of the search tool. +* `search_tool_info` - Additional metadata as a JSON object string (decode with `jsondecode`). +* `created_at` - Timestamp when the search tool was created. +* `updated_at` - Timestamp when the search tool was last updated. + +## Security Note + +`litellm_params` is not exposed through this data source because it may hold provider API keys. diff --git a/terraform/provider/docs/data-sources/search_tools.md b/terraform/provider/docs/data-sources/search_tools.md new file mode 100644 index 00000000000..a7d2add19fd --- /dev/null +++ b/terraform/provider/docs/data-sources/search_tools.md @@ -0,0 +1,34 @@ +# litellm_search_tools Data Source + +Retrieves the list of search tools configured on the LiteLLM proxy, from both the database and the proxy config. + +## Example Usage + +```hcl +data "litellm_search_tools" "all" {} + +output "search_tool_ids" { + value = data.litellm_search_tools.all.ids +} +``` + +## Argument Reference + +This data source takes no arguments. + +## Attribute Reference + +The following attributes are exported: + +* `ids` - List of search tool IDs. +* `search_tools` - List of search tools. Each entry exports: + * `search_tool_id` - The unique search tool ID. + * `search_tool_name` - Name of the search tool. + * `search_tool_info` - Additional metadata as a JSON object string. + * `is_from_config` - Whether the search tool comes from the proxy config file rather than the database. + * `created_at` - Timestamp when the search tool was created. + * `updated_at` - Timestamp when the search tool was last updated. + +## Security Note + +`litellm_params` is not exposed through this data source because it may hold provider API keys. diff --git a/terraform/provider/docs/data-sources/tag.md b/terraform/provider/docs/data-sources/tag.md new file mode 100644 index 00000000000..e87b1602c75 --- /dev/null +++ b/terraform/provider/docs/data-sources/tag.md @@ -0,0 +1,38 @@ +# litellm_tag (Data Source) + +Retrieves information about an existing LiteLLM tag, including its budget settings + +## Example Usage + +```hcl +data "litellm_tag" "production" { + name = "production" +} + +output "production_tag_budget" { + value = data.litellm_tag.production.max_budget +} +``` + +## Argument Reference + +The following arguments are supported: + +* `name` - (Required) Name of the tag to retrieve + +## Attribute Reference + +In addition to all arguments above, the following attributes are exported: + +* `description` - Description of the tag +* `models` - Model IDs this tag applies to +* `budget_id` - Budget ID associated with this tag +* `max_budget` - Max budget in USD for this tag +* `soft_budget` - Soft budget in USD for this tag +* `max_parallel_requests` - Max concurrent requests allowed for this tag +* `tpm_limit` - Max tokens per minute for this tag +* `rpm_limit` - Max requests per minute for this tag +* `budget_duration` - Duration for budget reset +* `created_at` - Timestamp when the tag was created +* `updated_at` - Timestamp when the tag was last updated +* `created_by` - User that created the tag diff --git a/terraform/provider/docs/data-sources/tags.md b/terraform/provider/docs/data-sources/tags.md new file mode 100644 index 00000000000..a65dcd4a927 --- /dev/null +++ b/terraform/provider/docs/data-sources/tags.md @@ -0,0 +1,50 @@ +# litellm_tags (Data Source) + +Retrieves the list of all LiteLLM tags. This includes stored tags created via `litellm_tag` or the API, and dynamic tags that were passed on requests + +## Example Usage + +```hcl +data "litellm_tags" "all" {} + +output "tag_names" { + value = data.litellm_tags.all.ids +} +``` + +## Example Usage with Date Filter + +```hcl +# Limit dynamic tags to those active in a window; stored tags are always returned +data "litellm_tags" "january" { + start_date = "2026-01-01" + end_date = "2026-01-31" +} +``` + +## Argument Reference + +The following arguments are supported: + +* `start_date` - (Optional) Start date (YYYY-MM-DD) limiting dynamic tags to those active in the window. Must be given with `end_date` +* `end_date` - (Optional) End date (YYYY-MM-DD). Must be given with `start_date` + +## Attribute Reference + +The following attributes are exported: + +* `ids` - Names of all tags (tag names are their IDs) +* `tags` - List of tags. Each entry exports: + * `name` - The tag name + * `description` - Description of the tag + * `models` - Model IDs this tag applies to + * `budget_id` - Budget ID associated with this tag + * `max_budget` - Max budget in USD + * `soft_budget` - Soft budget in USD + * `max_parallel_requests` - Max concurrent requests allowed + * `tpm_limit` - Max tokens per minute + * `rpm_limit` - Max requests per minute + * `budget_duration` - Duration for budget reset + * `created_at` - Timestamp when the tag was created + * `updated_at` - Timestamp when the tag was last updated + * `created_by` - User that created the tag diff --git a/terraform/provider/docs/data-sources/team.md b/terraform/provider/docs/data-sources/team.md new file mode 100644 index 00000000000..2e46238713b --- /dev/null +++ b/terraform/provider/docs/data-sources/team.md @@ -0,0 +1,52 @@ +--- +# generated by https://github.com/hashicorp/terraform-plugin-docs +page_title: "litellm_team Data Source - terraform-provider-litellm" +subcategory: "" +description: |- + Retrieves information about an existing LiteLLM team. +--- + +# litellm_team (Data Source) + +Retrieves information about an existing LiteLLM team via `/team/info`. Use it to reference teams created outside of Terraform or in other configurations. + +## Example Usage + +```terraform +data "litellm_team" "ml" { + team_id = "team-1234" +} + +resource "litellm_key" "ml_key" { + team_id = data.litellm_team.ml.team_id + models = data.litellm_team.ml.models +} +``` + +## Argument Reference + +The following arguments are supported: + +* `team_id` - (Required) Unique identifier of the team to retrieve. + +## Attributes Reference + +In addition to all arguments above, the following attributes are exported: + +* `team_alias` - User-friendly name of the team. +* `organization_id` - Organization the team belongs to. +* `models` - Models the team can access. +* `metadata` - Map of string metadata values for the team. +* `tags` - Tags for spend tracking and tag-based routing. +* `soft_budget_alerting_emails` - Email addresses alerted when the team crosses `soft_budget`. +* `tpm_limit` - Tokens per minute limit. +* `rpm_limit` - Requests per minute limit. +* `max_parallel_requests` - Maximum parallel requests allowed. +* `max_budget` - Maximum budget for the team. +* `soft_budget` - Soft budget alert threshold. +* `spend` - Amount spent by the team. +* `budget_duration` - Budget reset duration. +* `blocked` - Whether the team is blocked. +* `team_member_permissions` - Permissions granted to team members. +* `created_at` - Timestamp when the team was created. +* `updated_at` - Timestamp when the team was last updated. diff --git a/terraform/provider/docs/data-sources/teams.md b/terraform/provider/docs/data-sources/teams.md new file mode 100644 index 00000000000..b7587ae74c7 --- /dev/null +++ b/terraform/provider/docs/data-sources/teams.md @@ -0,0 +1,49 @@ +--- +# generated by https://github.com/hashicorp/terraform-plugin-docs +page_title: "litellm_teams Data Source - terraform-provider-litellm" +subcategory: "" +description: |- + Lists LiteLLM teams with optional server-side filters. +--- + +# litellm_teams (Data Source) + +Lists LiteLLM teams via `/team/list`. Supports filtering by user and organization. + +## Example Usage + +```terraform +data "litellm_teams" "org_teams" { + organization_id = litellm_organization.main.id +} + +output "team_ids" { + value = data.litellm_teams.org_teams.ids +} +``` + +## Argument Reference + +The following arguments are supported: + +* `user_id` - (Optional) Only return teams this user belongs to. +* `organization_id` - (Optional) Only return teams in this organization. + +## Attributes Reference + +In addition to all arguments above, the following attributes are exported: + +* `ids` - IDs of the returned teams. +* `teams` - List of team objects. Each entry exports: + * `team_id` - Unique identifier of the team. + * `team_alias` - User-friendly name of the team. + * `organization_id` - Organization the team belongs to. + * `models` - Models the team can access. + * `spend` - Amount spent by the team. + * `max_budget` - Maximum budget for the team. + * `tpm_limit` - Tokens per minute limit. + * `rpm_limit` - Requests per minute limit. + * `budget_duration` - Budget reset duration. + * `blocked` - Whether the team is blocked. + * `created_at` - Creation timestamp. + * `updated_at` - Last update timestamp. diff --git a/terraform/provider/docs/data-sources/unified_access_group.md b/terraform/provider/docs/data-sources/unified_access_group.md new file mode 100644 index 00000000000..8ca98c2d46a --- /dev/null +++ b/terraform/provider/docs/data-sources/unified_access_group.md @@ -0,0 +1,52 @@ +--- +page_title: "litellm_unified_access_group Data Source - terraform-provider-litellm" +subcategory: "" +description: |- + Retrieves information about an existing LiteLLM unified access group. +--- + +# litellm_unified_access_group (Data Source) + +Retrieves information about an existing LiteLLM unified access group by ID. + +## Example Usage + +```terraform +data "litellm_unified_access_group" "engineering" { + access_group_id = "b6e5f9d0-..." +} + +output "engineering_models" { + value = data.litellm_unified_access_group.engineering.access_model_names +} +``` + +## Argument Reference + +* `access_group_id` - (Required) ID of the unified access group to look up. + +## Attribute Reference + +* `id` - The unified access group ID. + +* `access_group_name` - Display name of the unified access group. + +* `description` - Description of the unified access group. + +* `access_model_names` - Model names the access group grants access to. + +* `access_mcp_server_ids` - MCP server IDs the access group grants access to. + +* `access_agent_ids` - Agent IDs the access group grants access to. + +* `assigned_team_ids` - Team IDs the access group is assigned to. + +* `assigned_key_ids` - Key IDs the access group is assigned to. + +* `created_at` - Timestamp when the access group was created. + +* `created_by` - User who created the access group. + +* `updated_at` - Timestamp when the access group was last updated. + +* `updated_by` - User who last updated the access group. diff --git a/terraform/provider/docs/data-sources/unified_access_groups.md b/terraform/provider/docs/data-sources/unified_access_groups.md new file mode 100644 index 00000000000..118d003d76f --- /dev/null +++ b/terraform/provider/docs/data-sources/unified_access_groups.md @@ -0,0 +1,30 @@ +--- +page_title: "litellm_unified_access_groups Data Source - terraform-provider-litellm" +subcategory: "" +description: |- + Retrieves all LiteLLM unified access groups. +--- + +# litellm_unified_access_groups (Data Source) + +Retrieves all LiteLLM unified access groups configured on the proxy. + +## Example Usage + +```terraform +data "litellm_unified_access_groups" "all" {} + +output "unified_access_group_ids" { + value = data.litellm_unified_access_groups.all.ids +} +``` + +## Argument Reference + +This data source takes no arguments. + +## Attribute Reference + +* `access_groups` - List of unified access groups. Each entry exports the same attributes as the `litellm_unified_access_group` data source: `access_group_id`, `access_group_name`, `description`, `access_model_names`, `access_mcp_server_ids`, `access_agent_ids`, `assigned_team_ids`, `assigned_key_ids`, `created_at`, `created_by`, `updated_at`, and `updated_by`. + +* `ids` - List of all unified access group IDs. diff --git a/terraform/provider/docs/data-sources/user.md b/terraform/provider/docs/data-sources/user.md new file mode 100644 index 00000000000..2d4fc946a7d --- /dev/null +++ b/terraform/provider/docs/data-sources/user.md @@ -0,0 +1,36 @@ +# litellm_user Data Source + +Retrieves information about an existing LiteLLM user by ID + +## Example Usage + +```hcl +data "litellm_user" "alice" { + user_id = "alice-user-id" +} + +output "alice_email" { + value = data.litellm_user.alice.user_email +} +``` + +## Argument Reference + +- `user_id` (Required) - ID of the user to retrieve + +## Attribute Reference + +- `id` - The user ID +- `user_email` - Email address of the user +- `user_alias` - Descriptive name for the user +- `user_role` - Role of the user on the proxy +- `teams` - List of team IDs the user belongs to +- `models` - Models the user is allowed to call +- `max_budget` - Maximum budget in USD for the user +- `spend` - Current spend in USD for the user +- `budget_duration` - Budget reset period for the user +- `tpm_limit` - Tokens per minute limit +- `rpm_limit` - Requests per minute limit +- `max_parallel_requests` - Maximum number of parallel requests +- `metadata` - Map of metadata for the user +- `model_max_budget` - JSON string of per-model budget config diff --git a/terraform/provider/docs/data-sources/users.md b/terraform/provider/docs/data-sources/users.md new file mode 100644 index 00000000000..5cc44aee07e --- /dev/null +++ b/terraform/provider/docs/data-sources/users.md @@ -0,0 +1,47 @@ +# litellm_users Data Source + +Retrieves a page of LiteLLM users, with optional server-side filters + +## Example Usage + +```hcl +data "litellm_users" "internal" { + role = "internal_user" + page = 1 + page_size = 100 +} + +output "internal_user_ids" { + value = data.litellm_users.internal.ids +} +``` + +## Argument Reference + +- `role` (Optional) - Filter users by role +- `user_ids` (Optional) - Comma-separated list of user IDs to filter by +- `user_email` (Optional) - Filter users by partial email match +- `team` (Optional) - Filter users by team ID +- `page` (Optional, Default `1`) - Page number to fetch +- `page_size` (Optional, Default `25`) - Number of users per page, max 100 +- `sort_by` (Optional) - Column to sort by, e.g. `user_id`, `user_email`, `created_at` +- `sort_order` (Optional) - Sort order, `asc` or `desc` + +## Attribute Reference + +- `users` - Users returned for the requested page. Each entry has: + - `user_id` - The user ID + - `user_email` - Email address of the user + - `user_alias` - Descriptive name for the user + - `user_role` - Role of the user on the proxy + - `teams` - List of team IDs the user belongs to + - `models` - Models the user is allowed to call + - `max_budget` - Maximum budget in USD + - `spend` - Current spend in USD + - `tpm_limit` - Tokens per minute limit + - `rpm_limit` - Requests per minute limit + - `key_count` - Number of API keys owned by the user + - `created_at` - Timestamp when the user was created +- `ids` - IDs of the users returned for the requested page +- `total` - Total number of users matching the filters +- `total_pages` - Total number of pages available diff --git a/terraform/provider/docs/index.md b/terraform/provider/docs/index.md index c03071e7ed3..e6641782a4d 100644 --- a/terraform/provider/docs/index.md +++ b/terraform/provider/docs/index.md @@ -51,6 +51,7 @@ The LiteLLM provider supports the following resources: * [`litellm_mcp_server`](./resources/mcp_server) - Manage MCP (Model Context Protocol) servers * [`litellm_credential`](./resources/credential) - Manage credentials for various providers * [`litellm_vector_store`](./resources/vector_store) - Manage vector stores +* [`litellm_jwt_key_mapping`](./resources/jwt_key_mapping) - Map JWT claim values to virtual keys ## Available Data Sources diff --git a/terraform/provider/docs/resources/access_group.md b/terraform/provider/docs/resources/access_group.md new file mode 100644 index 00000000000..e7b05116d43 --- /dev/null +++ b/terraform/provider/docs/resources/access_group.md @@ -0,0 +1,49 @@ +--- +page_title: "litellm_access_group Resource - terraform-provider-litellm" +subcategory: "" +description: |- + Manages a LiteLLM model access group. +--- + +# litellm_access_group (Resource) + +Manages a LiteLLM model access group. Access groups bundle model deployments under one name so keys and teams can be granted access to the whole group at once. + +## Example Usage + +```terraform +resource "litellm_access_group" "production" { + access_group = "production-models" + model_names = ["gpt-4", "claude-3-sonnet"] +} + +# Target specific deployments by model ID instead of model name +resource "litellm_access_group" "pinned" { + access_group = "pinned-deployments" + model_ids = ["4dbd9f43-...", "9a1e2c77-..."] +} +``` + +## Argument Reference + +* `access_group` - (Required, Forces new resource) Name of the access group. + +* `model_names` - (Optional) List of model names (the `model_name` of each deployment) to include in the group. At least one of `model_names` or `model_ids` must be set. + +* `model_ids` - (Optional) List of specific deployment model IDs to include in the group. Takes precedence over `model_names` when both are set. + +## Attribute Reference + +In addition to the arguments above, the following attributes are exported: + +* `id` - The access group name. + +* `deployment_count` - Number of deployments currently tagged with this access group. + +## Import + +Access groups can be imported using the access group name: + +```shell +terraform import litellm_access_group.production production-models +``` diff --git a/terraform/provider/docs/resources/agent.md b/terraform/provider/docs/resources/agent.md new file mode 100644 index 00000000000..93b78197784 --- /dev/null +++ b/terraform/provider/docs/resources/agent.md @@ -0,0 +1,88 @@ +# litellm_agent Resource + +Manages an A2A (Agent-to-Agent) agent on the LiteLLM proxy. Agents are AI-powered entities that can be discovered, invoked, and composed using the A2A protocol. + +## Example Usage + +```hcl +resource "litellm_agent" "hello_world" { + agent_name = "hello-world-agent" + + agent_card_params = jsonencode({ + protocolVersion = "1.0" + name = "Hello World Agent" + description = "Just a hello world agent" + url = "http://localhost:9999/" + version = "1.0.0" + defaultInputModes = ["text"] + defaultOutputModes = ["text"] + capabilities = { + streaming = true + } + skills = [ + { + id = "hello_world" + name = "Returns hello world" + description = "just returns hello world" + tags = ["hello world"] + examples = ["hi", "hello world"] + } + ] + }) + + litellm_params = jsonencode({ + make_public = false + }) + + object_permission = jsonencode({ + models = ["gpt-4-proxy"] + mcp_servers = ["my-mcp-server-id"] + }) + + static_headers = { + "x-api-key" = var.agent_api_key + } + + extra_headers = ["x-request-id"] + + tpm_limit = 100000 + rpm_limit = 1000 + session_tpm_limit = 10000 + session_rpm_limit = 100 +} +``` + +## Argument Reference + +The following arguments are supported: + +* `agent_name` - (Required) Name of the agent. Must be unique on the proxy. +* `agent_card_params` - (Required) The A2A agent card as a JSON object string (use `jsonencode`). Supports the standard A2A card fields: `name`, `description`, `url`, `version`, `protocolVersion`, `capabilities`, `skills`, `defaultInputModes`, `defaultOutputModes`, `preferredTransport`, `iconUrl`, `provider`, `documentationUrl`, and more. The proxy merges LiteLLM-fronting fields (such as `supportedInterfaces`) into the stored card, so the value you configure stays authoritative in state. +* `litellm_params` - (Optional, Sensitive) LiteLLM-specific parameters as a JSON object string. May include secrets such as `api_key`, so the value is never read back from the API; the configured value is authoritative. +* `object_permission` - (Optional) Access control permissions as a JSON object string with keys `mcp_servers`, `mcp_access_groups`, `mcp_tool_permissions`, `models`, and `agents`. +* `static_headers` - (Optional, Sensitive) Map of static headers sent with agent requests. May hold tokens, so it is never read back from the API. +* `extra_headers` - (Optional) List of incoming request header names to forward to the agent. +* `tpm_limit` - (Optional) Tokens per minute limit for the agent. +* `rpm_limit` - (Optional) Requests per minute limit for the agent. +* `session_tpm_limit` - (Optional) Per-session tokens per minute limit. +* `session_rpm_limit` - (Optional) Per-session requests per minute limit. + +## Attribute Reference + +In addition to all arguments above, the following attributes are exported: + +* `id` - The agent ID assigned by LiteLLM. +* `created_at` - Timestamp when the agent was created. +* `updated_at` - Timestamp when the agent was last updated. +* `created_by` - User who created the agent. +* `updated_by` - User who last updated the agent. + +## Import + +Agents can be imported using the agent ID: + +```shell +terraform import litellm_agent.example +``` + +Note: `litellm_params` and `static_headers` cannot be recovered on import because the API never returns their unmasked values; re-apply after import to set them. diff --git a/terraform/provider/docs/resources/budget.md b/terraform/provider/docs/resources/budget.md new file mode 100644 index 00000000000..d635543013b --- /dev/null +++ b/terraform/provider/docs/resources/budget.md @@ -0,0 +1,48 @@ +# litellm_budget Resource + +Manages a budget object on the LiteLLM proxy. Budgets can be attached to keys, teams, organizations, and end users to enforce spend limits + +## Example Usage + +```hcl +resource "litellm_budget" "engineering" { + budget_id = "engineering-monthly" + max_budget = 500.0 + soft_budget = 400.0 + budget_duration = "30d" + tpm_limit = 500000 + rpm_limit = 5000 + max_parallel_requests = 100 + + model_max_budget = jsonencode({ + "gpt-4o" = { + max_budget = 100.0 + budget_duration = "1d" + } + }) +} +``` + +## Argument Reference + +- `budget_id` (Optional, Forces new resource) - Unique ID for the budget. Generated by the server if not provided +- `max_budget` (Optional) - Requests fail if this budget in USD is exceeded +- `soft_budget` (Optional) - Requests do not fail if this is exceeded, but alerts fire +- `max_parallel_requests` (Optional) - Maximum concurrent requests allowed for this budget +- `tpm_limit` (Optional) - Maximum tokens per minute allowed for this budget +- `rpm_limit` (Optional) - Maximum requests per minute allowed for this budget +- `budget_duration` (Optional) - Budget reset period, e.g. `1hr`, `1d`, `28d` +- `model_max_budget` (Optional) - JSON string of per-model budget config, e.g. `jsonencode({"gpt-4o" = {max_budget = 10.0}})` + +## Attribute Reference + +- `id` - The budget ID +- `budget_reset_at` - Datetime when the budget is reset + +## Import + +Budgets can be imported using the budget ID: + +```shell +terraform import litellm_budget.engineering +``` diff --git a/terraform/provider/docs/resources/fallback.md b/terraform/provider/docs/resources/fallback.md new file mode 100644 index 00000000000..7d93d4c5bb4 --- /dev/null +++ b/terraform/provider/docs/resources/fallback.md @@ -0,0 +1,48 @@ +# litellm_fallback Resource + +Manages a fallback configuration for a model in LiteLLM. Fallbacks are triggered when a call to the primary model fails after retries. + +## Example Usage + +### Basic Fallback Configuration + +```hcl +resource "litellm_fallback" "gpt4_fallbacks" { + model = "gpt-4" + fallback_models = ["claude-3-sonnet", "gpt-3.5-turbo"] +} +``` + +### Context Window Fallback + +```hcl +resource "litellm_fallback" "gpt4_context_window" { + model = "gpt-4" + fallback_models = ["claude-3-sonnet"] + fallback_type = "context_window" +} +``` + +## Argument Reference + +The following arguments are supported: + +* `model` - (Required, Forces new resource) The model name to configure fallbacks for. The model must already exist on the proxy. +* `fallback_models` - (Required) List of fallback model names in order of priority. Each model must exist on the proxy, and the primary model cannot be its own fallback. +* `fallback_type` - (Optional, Forces new resource) Type of fallback. One of `general` (default), `context_window`, or `content_policy`. + +## Attribute Reference + +In addition to the arguments above, the following attribute is exported: + +* `id` - The primary model name. + +## Import + +Fallback configurations can be imported using the primary model name: + +```shell +terraform import litellm_fallback.example gpt-4 +``` + +Note: import always reads the `general` fallback type. Fallbacks of type `context_window` or `content_policy` cannot be imported. diff --git a/terraform/provider/docs/resources/guardrail.md b/terraform/provider/docs/resources/guardrail.md new file mode 100644 index 00000000000..978ec169a34 --- /dev/null +++ b/terraform/provider/docs/resources/guardrail.md @@ -0,0 +1,57 @@ +# litellm_guardrail Resource + +Manages a guardrail in LiteLLM. Guardrails provide content filtering, PII detection, prompt injection protection, and more. + +## Example Usage + +```hcl +resource "litellm_guardrail" "bedrock_guard" { + guardrail_name = "my-bedrock-guard" + guardrail = "bedrock" + mode = "pre_call" + default_on = true + + litellm_params = jsonencode({ + guardrailIdentifier = "ff6ujrregl1q" + guardrailVersion = "DRAFT" + }) + + guardrail_info = { + description = "Bedrock content moderation guardrail" + } +} +``` + +### Multiple Modes + +```hcl +resource "litellm_guardrail" "pii_guard" { + guardrail_name = "presidio-pii" + guardrail = "presidio" + mode = jsonencode(["pre_call", "post_call"]) +} +``` + +## Argument Reference + +* `guardrail_name` - (Required) Human-readable name for the guardrail. +* `guardrail` - (Required) The guardrail integration type (e.g. `bedrock`, `lakera`, `presidio`, `openai_moderation`, `hide_secrets`). +* `mode` - (Required) When to apply the guardrail. A single value (`pre_call`, `post_call`, `during_call`, `logging_only`) or a JSON array of values. +* `default_on` - (Optional) Whether the guardrail is enabled by default for all requests. +* `litellm_params` - (Optional, Sensitive) JSON string with additional provider-specific parameters merged into `litellm_params` (may contain API keys). The API masks these values, so the configured value stays authoritative in state. +* `guardrail_info` - (Optional) Map of additional metadata for the guardrail. + +## Attribute Reference + +* `id` - The guardrail ID assigned by LiteLLM. +* `created_at` - Timestamp when the guardrail was created. + +## Import + +Guardrails can be imported using the guardrail ID: + +```shell +terraform import litellm_guardrail.example 123e4567-e89b-12d3-a456-426614174000 +``` + +Note: `guardrail`, `mode`, `default_on` and `litellm_params` are not returned unmasked by the API, so after import you must set them in configuration to match the server. diff --git a/terraform/provider/docs/resources/jwt_key_mapping.md b/terraform/provider/docs/resources/jwt_key_mapping.md new file mode 100644 index 00000000000..fbc30947113 --- /dev/null +++ b/terraform/provider/docs/resources/jwt_key_mapping.md @@ -0,0 +1,94 @@ +# litellm_jwt_key_mapping + +Maps a JWT claim value to a LiteLLM virtual key. Every JWT client identified by a claim, typically `client_id`, `azp` or `sub`, then gets the model restrictions, budgets, rate limits, guardrails and spend tracking of the virtual key it maps to, without that key ever being handed to the client. + +The mappings only take effect once JWT auth is enabled on the proxy, which is configuration rather than API state: + +```yaml +general_settings: + enable_jwt_auth: True + litellm_jwtauth: + virtual_key_claim_field: "client_id" + unregistered_jwt_client_behavior: "fallback_team_mapping" +``` + +See [JWT to virtual key mapping](https://docs.litellm.ai/docs/proxy/jwt_key_mapping) for the proxy side of the feature + +## Example Usage + +The mapped virtual key has to exist already and its value has to be known to Terraform, so it comes from a variable or a secret manager rather than from a `litellm_key` resource. `litellm_key` deliberately made its generated `key` write-only, to avoid storing raw API keys in state, so referencing it here does not merely read back null: Terraform's write-only enforcement turns `key = litellm_key.foo.key` into a static `Missing required argument` error at `terraform plan`, before any API call, in every apply ordering, including a first apply where both resources are created together: + +```hcl +variable "alice_key" { + type = string + sensitive = true +} + +resource "litellm_jwt_key_mapping" "alice" { + jwt_claim_name = "client_id" + jwt_claim_value = "dev-alice" + key = var.alice_key +} +``` + +Per-client limits live on the virtual key, so one mapping per client is how each JWT client gets its own budget and quota: + +```hcl +resource "litellm_jwt_key_mapping" "billing_service" { + jwt_claim_name = "client_id" + jwt_claim_value = "billing-service" + key = var.billing_service_key + description = "Billing service JWT client" + is_active = true +} +``` + +Several clients at once, with the key values coming from a map of secrets: + +```hcl +variable "jwt_client_keys" { + type = map(string) + sensitive = true +} + +resource "litellm_jwt_key_mapping" "developer" { + for_each = var.jwt_client_keys + + jwt_claim_name = "client_id" + jwt_claim_value = each.key + key = each.value + description = "Developer JWT client ${each.key}" +} +``` + +## Argument Reference + +- `jwt_claim_name` - (Required, ForceNew) Name of the JWT claim to match on, for example `client_id`, `azp` or `sub`. Must match `virtual_key_claim_field` in the proxy JWT config +- `jwt_claim_value` - (Required, ForceNew) Value of the claim identifying the JWT client. Unique together with `jwt_claim_name`, so a second mapping for the same pair fails with a 409 +- `key` - (Required, Sensitive) The virtual key this claim value maps to. It has to exist already, otherwise the proxy rejects the mapping with `The provided key does not match an existing virtual key` +- `description` - (Optional) Description of the mapping +- `is_active` - (Optional) Whether the mapping is active. Inactive mappings are ignored during JWT auth. Defaults to `true` + +## Attribute Reference + +- `id` - The mapping ID assigned by LiteLLM +- `created_at` - Timestamp when the mapping was created +- `updated_at` - Timestamp when the mapping was last updated +- `created_by` - User who created the mapping +- `updated_by` - User who last updated the mapping + +## Notes + +The proxy stores only a hash of `key` and never returns it, so drift on that attribute cannot be detected and Terraform tracks the value from your configuration. Changing `key` rotates the mapping onto the new virtual key in place, with no replacement. Like the other secrets this provider accepts, such as `credential_values` and `model_api_key`, the configured value is kept in state, so treat the state as sensitive + +Only proxy admins can create, update or delete mappings, so the provider `api_key` has to be a master key or an admin key + +## Import + +Mappings are imported by their mapping ID: + +```shell +terraform import litellm_jwt_key_mapping.alice 297a5536-1aeb-4cf1-b666-b3809c2750a8 +``` + +Because the API does not return the mapped key, `key` is empty in state right after an import, so the first plan shows an in-place update that pushes the configured key back to the proxy. That update is harmless, the proxy just rehashes the same value when the key has not actually changed diff --git a/terraform/provider/docs/resources/key.md b/terraform/provider/docs/resources/key.md index b48d3334c14..5094b77cbec 100644 --- a/terraform/provider/docs/resources/key.md +++ b/terraform/provider/docs/resources/key.md @@ -93,6 +93,24 @@ The following arguments are supported: * `tags` - (Optional) List of tags associated with this key. This can be used for organization and filtering of keys. +* `budget_id` - (Optional) ID of a shared budget (created via `litellm_budget`) to attach to this key. + +* `enforced_params` - (Optional) List of request parameters that callers must supply when using this key (for example `user`). + +* `allowed_routes` - (Optional) List of proxy routes this key is allowed to call. + +* `allowed_passthrough_routes` - (Optional) List of pass-through routes this key is allowed to call. + +* `rpm_limit_type` - (Optional) How the RPM limit is enforced. One of `guaranteed_throughput`, `best_effort_throughput` or `dynamic`. + +* `tpm_limit_type` - (Optional) How the TPM limit is enforced. One of `guaranteed_throughput`, `best_effort_throughput` or `dynamic`. + +* `prompts` - (Optional) List of prompt IDs this key is allowed to use. + +* `organization_id` - (Optional) ID of the organization this key belongs to. + +* `project_id` - (Optional) ID of the project this key belongs to. Changing this forces a new key to be created. + ## Attribute Reference In addition to all arguments above, the following attributes are exported: diff --git a/terraform/provider/docs/resources/key_block.md b/terraform/provider/docs/resources/key_block.md new file mode 100644 index 00000000000..42fea57c13b --- /dev/null +++ b/terraform/provider/docs/resources/key_block.md @@ -0,0 +1,40 @@ +# litellm_key_block Resource + +Manages the blocked state of an existing LiteLLM API key. Creating this resource blocks the key; destroying it unblocks the key. + +If the key is unblocked outside of Terraform (or deleted), the resource is removed from state and Terraform plans to re-block it on the next apply. + +## Example Usage + +```hcl +resource "litellm_key" "example" { + models = ["gpt-4"] +} + +resource "litellm_key_block" "example" { + key = litellm_key.example.key +} +``` + +## Argument Reference + +The following arguments are supported: + +* `key` - (Required, Forces new resource, Sensitive) The API key to block, as the raw `sk-` value or its SHA-256 token hash. The provider normalizes raw values to the hash before talking to the API, so the plaintext key never appears in request URLs, the resource ID, or plan output. + +## Attribute Reference + +In addition to the arguments above, the following attributes are exported: + +* `id` - The SHA-256 token hash of the key. +* `blocked` - Whether the key is currently blocked. Always `true` while this resource exists. + +If the same key is also managed by a `litellm_key` resource, that resource's `blocked` attribute will show drift while the block is active; either set `blocked` there instead of using this resource, or add `lifecycle { ignore_changes = [blocked] }` to the `litellm_key`. + +## Import + +Key blocks can be imported using the key's SHA-256 token hash (shown as the key's ID in `litellm_key` state and in `/key/info`): + +```shell +terraform import litellm_key_block.example 88362cbb875f4b48b4b5b56b2ea45f66465e27d55a189816bd54e5643e5410eb +``` diff --git a/terraform/provider/docs/resources/project.md b/terraform/provider/docs/resources/project.md new file mode 100644 index 00000000000..6824ffabc8c --- /dev/null +++ b/terraform/provider/docs/resources/project.md @@ -0,0 +1,71 @@ +# litellm_project Resource + +Manages a project in LiteLLM. Projects sit between teams and keys in the hierarchy, allowing fine-grained budget and model access control within a team + +## Example Usage + +```hcl +resource "litellm_team" "research" { + team_alias = "research-team" +} + +resource "litellm_project" "ml_experiments" { + team_id = litellm_team.research.id + project_alias = "ml-experiments" + description = "ML experimentation project" + models = ["gpt-5.6", "claude-opus-5"] + + max_budget = 1000.0 + soft_budget = 800.0 + budget_duration = "30d" + tpm_limit = 500000 + rpm_limit = 5000 + + tags = ["research", "gpu"] + + metadata = { + cost_center = "R&D-001" + } +} +``` + +## Argument Reference + +The following arguments are supported: + +* `team_id` - (Required, Forces new resource) The team ID this project belongs to +* `project_alias` - (Optional) Human-friendly name for the project +* `description` - (Optional) Description of the project's purpose and use case +* `models` - (Optional) List of models the project can access +* `metadata` - (Optional) Map of metadata for the project +* `tags` - (Optional) Tags associated with the project +* `max_budget` - (Optional) Maximum budget for this project +* `soft_budget` - (Optional) Soft budget limit for warnings +* `budget_duration` - (Optional) Budget reset duration, for example `1h`, `30d` +* `budget_id` - (Optional) Budget ID to associate with this project +* `tpm_limit` - (Optional) Tokens per minute limit +* `rpm_limit` - (Optional) Requests per minute limit +* `max_parallel_requests` - (Optional) Maximum parallel requests allowed +* `model_max_budget` - (Optional) Map of per-model budget limits +* `model_rpm_limit` - (Optional) Map of per-model RPM limits +* `model_tpm_limit` - (Optional) Map of per-model TPM limits +* `blocked` - (Optional) Whether the project is blocked from making requests + +## Attribute Reference + +In addition to all arguments above, the following attributes are exported: + +* `id` - The project ID assigned by LiteLLM +* `spend` - Current spend for the project +* `created_at` - Timestamp when the project was created +* `updated_at` - Timestamp when the project was last updated +* `created_by` - User that created the project +* `updated_by` - User that last updated the project + +## Import + +Projects can be imported using the project ID: + +```shell +terraform import litellm_project.example 4a422a4c-e246-4d02-a1eb-13e835cd0725 +``` diff --git a/terraform/provider/docs/resources/prompt.md b/terraform/provider/docs/resources/prompt.md new file mode 100644 index 00000000000..e9bda62ebf4 --- /dev/null +++ b/terraform/provider/docs/resources/prompt.md @@ -0,0 +1,61 @@ +# litellm_prompt Resource + +Manages a prompt in LiteLLM. Prompts let you manage prompt templates from external providers such as Langfuse, or inline dotprompt content. + +## Example Usage + +```hcl +resource "litellm_prompt" "langfuse_prompt" { + prompt_id = "my-langfuse-prompt" + prompt_integration = "langfuse" + api_base = "https://cloud.langfuse.com" + api_key = var.langfuse_api_key + prompt_type = "db" + + litellm_params = jsonencode({ + prompt_id = "prompt-name-in-langfuse" + }) +} +``` + +### Dotprompt + +```hcl +resource "litellm_prompt" "greeting" { + prompt_id = "greeting" + prompt_integration = "dotprompt" + prompt_type = "db" + + dotprompt_content = <<-EOT + --- + model: gpt-5.2 + --- + Say hello to {{name}}. + EOT +} +``` + +## Argument Reference + +* `prompt_id` - (Required, Forces new resource) Unique identifier for the prompt. +* `prompt_integration` - (Required) The prompt integration provider (e.g. `langfuse`, `dotprompt`). +* `api_base` - (Optional) Base URL for the prompt provider API. +* `api_key` - (Optional, Sensitive) API key for the prompt provider. Never read back into state. +* `provider_specific_query_params` - (Optional) JSON string of provider-specific query parameters. +* `ignore_prompt_manager_model` - (Optional) If true, ignore the model specified in the prompt manager. +* `ignore_prompt_manager_optional_params` - (Optional) If true, ignore optional params from the prompt manager. +* `dotprompt_content` - (Optional) Content for the dotprompt integration. +* `litellm_params` - (Optional, Sensitive) JSON string with additional `litellm_params` merged into the request, e.g. the integration's own `prompt_id`, `prompt_directory` or `prompt_data`. Never read back into state. +* `prompt_type` - (Optional) Type of prompt: `config` or `db`. + +## Attribute Reference + +* `id` - The prompt ID (same as `prompt_id`). + +## Import + +Prompts can be imported using the prompt ID: + +```shell +terraform import litellm_prompt.example my-langfuse-prompt +``` diff --git a/terraform/provider/docs/resources/search_tool.md b/terraform/provider/docs/resources/search_tool.md new file mode 100644 index 00000000000..9f55e143a9f --- /dev/null +++ b/terraform/provider/docs/resources/search_tool.md @@ -0,0 +1,46 @@ +# litellm_search_tool Resource + +Manages a search tool configuration on the LiteLLM proxy. Search tools connect the proxy's `/search` endpoints to an external search provider such as Tavily, Perplexity, or Exa. + +## Example Usage + +```hcl +resource "litellm_search_tool" "tavily" { + search_tool_name = "tavily-search" + + litellm_params = jsonencode({ + search_provider = "tavily" + api_key = var.tavily_api_key + }) + + search_tool_info = jsonencode({ + description = "Tavily web search" + }) +} +``` + +## Argument Reference + +The following arguments are supported: + +* `search_tool_name` - (Required) Name of the search tool. +* `litellm_params` - (Required, Sensitive) Search tool parameters as a JSON object string (use `jsonencode`). Must include `search_provider`, and typically an `api_key`; may also carry `api_base`, `timeout`, `max_retries`, and other provider options. The API only returns masked values, so this is never read back; the configured value is authoritative. +* `search_tool_info` - (Optional) Additional metadata as a JSON object string, e.g. a `description`. + +## Attribute Reference + +In addition to all arguments above, the following attributes are exported: + +* `id` - The search tool ID assigned by LiteLLM. +* `created_at` - Timestamp when the search tool was created. +* `updated_at` - Timestamp when the search tool was last updated. + +## Import + +Search tools can be imported using the search tool ID: + +```shell +terraform import litellm_search_tool.example +``` + +Note: `litellm_params` cannot be recovered on import because the API only returns masked values; re-apply after import to set it. diff --git a/terraform/provider/docs/resources/tag.md b/terraform/provider/docs/resources/tag.md new file mode 100644 index 00000000000..98b274bf9e3 --- /dev/null +++ b/terraform/provider/docs/resources/tag.md @@ -0,0 +1,49 @@ +# litellm_tag Resource + +Manages a tag in LiteLLM. Tags are used for spend tracking, budgets, and tag-based routing to specific model deployments + +## Example Usage + +```hcl +resource "litellm_tag" "production" { + name = "production" + description = "Production traffic" + models = ["4a422a4c-e246-4d02-a1eb-13e835cd0725"] + + max_budget = 500.0 + soft_budget = 400.0 + budget_duration = "30d" + tpm_limit = 100000 + rpm_limit = 1000 +} +``` + +## Argument Reference + +The following arguments are supported: + +* `name` - (Required, Forces new resource) Unique name of the tag. Also used as the resource ID +* `description` - (Optional) Description of the tag +* `models` - (Optional) List of model IDs this tag applies to +* `budget_id` - (Optional) Existing budget ID to associate with this tag. If omitted and budget fields are set, the proxy creates a budget +* `max_budget` - (Optional) Max budget in USD for this tag +* `soft_budget` - (Optional) Soft budget in USD for this tag +* `max_parallel_requests` - (Optional) Max concurrent requests allowed for this tag +* `tpm_limit` - (Optional) Max tokens per minute for this tag +* `rpm_limit` - (Optional) Max requests per minute for this tag +* `budget_duration` - (Optional) Duration for budget reset, for example `1h`, `1d`, `30d` +* `model_max_budget` - (Optional) JSON object string with per-model budget configuration + +## Attribute Reference + +In addition to all arguments above, the following attributes are exported: + +* `id` - The tag name + +## Import + +Tags can be imported using the tag name: + +```shell +terraform import litellm_tag.example production +``` diff --git a/terraform/provider/docs/resources/team.md b/terraform/provider/docs/resources/team.md index 65ab4bf82d4..821d8c1dee3 100644 --- a/terraform/provider/docs/resources/team.md +++ b/terraform/provider/docs/resources/team.md @@ -122,6 +122,32 @@ The following arguments are supported: * `team_member_permissions` - (Optional) List of permissions granted to team members. This controls what actions team members can perform within the team context. +* `model_aliases` - (Optional) Map of alias names to model names, letting the team call models under stable alias names. + +* `guardrails` - (Optional) List of guardrails applied to every request made by this team. + +* `prompts` - (Optional) List of prompt IDs the team is allowed to use. + +* `team_member_budget` - (Optional) Budget (in USD) applied to each individual team member. + +* `team_member_budget_duration` - (Optional) Reset cycle for the per-member budget (e.g. `30d`, `1mo`). + +* `team_member_rpm_limit` - (Optional) Requests per minute limit applied to each individual team member. + +* `team_member_tpm_limit` - (Optional) Tokens per minute limit applied to each individual team member. + +* `team_member_key_duration` - (Optional) Lifetime for keys created by team members (e.g. `1d`, `1w`). + +* `model_rpm_limit` - (Optional) Map of model name to requests per minute limit for that model. + +* `model_tpm_limit` - (Optional) Map of model name to tokens per minute limit for that model. + +* `allowed_passthrough_routes` - (Optional) List of pass-through routes this team is allowed to call. + +* `rpm_limit_type` - (Optional) How the RPM limit is enforced: `guaranteed_throughput` or `best_effort_throughput`. Changing this forces a new team to be created. + +* `tpm_limit_type` - (Optional) How the TPM limit is enforced: `guaranteed_throughput` or `best_effort_throughput`. Changing this forces a new team to be created. + ## Attribute Reference In addition to the arguments above, the following attributes are exported: diff --git a/terraform/provider/docs/resources/team_block.md b/terraform/provider/docs/resources/team_block.md new file mode 100644 index 00000000000..3749b6dc827 --- /dev/null +++ b/terraform/provider/docs/resources/team_block.md @@ -0,0 +1,38 @@ +# litellm_team_block Resource + +Manages the blocked state of an existing LiteLLM team. Creating this resource blocks the team (all calls from its keys are rejected); destroying it unblocks the team. + +If the team is unblocked outside of Terraform (or deleted), the resource is removed from state and Terraform plans to re-block it on the next apply. + +## Example Usage + +```hcl +resource "litellm_team" "example" { + team_alias = "suspended-team" +} + +resource "litellm_team_block" "example" { + team_id = litellm_team.example.id +} +``` + +## Argument Reference + +The following arguments are supported: + +* `team_id` - (Required, Forces new resource) The ID of the team to block. + +## Attribute Reference + +In addition to the arguments above, the following attributes are exported: + +* `id` - The team ID. +* `blocked` - Whether the team is currently blocked. Always `true` while this resource exists. + +## Import + +Team blocks can be imported using the team ID: + +```shell +terraform import litellm_team_block.example team-1234 +``` diff --git a/terraform/provider/docs/resources/unified_access_group.md b/terraform/provider/docs/resources/unified_access_group.md new file mode 100644 index 00000000000..038e0e02d77 --- /dev/null +++ b/terraform/provider/docs/resources/unified_access_group.md @@ -0,0 +1,64 @@ +--- +page_title: "litellm_unified_access_group Resource - terraform-provider-litellm" +subcategory: "" +description: |- + Manages a LiteLLM unified access group. +--- + +# litellm_unified_access_group (Resource) + +Manages a LiteLLM unified access group. Unified access groups grant access to models, MCP servers, and agents in one bundle, and can be assigned to teams and keys. + +## Example Usage + +```terraform +resource "litellm_unified_access_group" "engineering" { + access_group_name = "engineering-access" + description = "Models and tools for the engineering org" + + access_model_names = ["gpt-4", "claude-3-sonnet"] + access_mcp_server_ids = [litellm_mcp_server.github.id] + + assigned_team_ids = [litellm_team.engineering.id] +} +``` + +## Argument Reference + +* `access_group_name` - (Required) Display name of the unified access group. + +* `description` - (Optional) Description of the unified access group. + +* `access_model_names` - (Optional) Model names this access group grants access to. + +* `access_mcp_server_ids` - (Optional) MCP server IDs this access group grants access to. + +* `access_agent_ids` - (Optional) Agent IDs this access group grants access to. + +* `assigned_team_ids` - (Optional) Team IDs the access group is assigned to. + +* `assigned_key_ids` - (Optional) Key IDs (token hashes) the access group is assigned to. + +## Attribute Reference + +In addition to the arguments above, the following attributes are exported: + +* `id` - The unique identifier of the unified access group. + +* `access_group_id` - Same as `id`. + +* `created_at` - Timestamp when the access group was created. + +* `created_by` - User who created the access group. + +* `updated_at` - Timestamp when the access group was last updated. + +* `updated_by` - User who last updated the access group. + +## Import + +Unified access groups can be imported using the access group ID: + +```shell +terraform import litellm_unified_access_group.engineering +``` diff --git a/terraform/provider/docs/resources/user.md b/terraform/provider/docs/resources/user.md new file mode 100644 index 00000000000..9b537292d54 --- /dev/null +++ b/terraform/provider/docs/resources/user.md @@ -0,0 +1,66 @@ +# litellm_user Resource + +Manages an internal user on the LiteLLM proxy. Internal users can log into the Admin UI, own API keys, and belong to teams + +## Example Usage + +```hcl +resource "litellm_user" "alice" { + user_email = "alice@example.com" + user_alias = "Alice" + user_role = "internal_user" + max_budget = 100.0 + budget_duration = "30d" + tpm_limit = 100000 + rpm_limit = 1000 + teams = [litellm_team.engineering.id] + models = ["gpt-4o", "claude-sonnet-4-5"] + + metadata = { + department = "engineering" + } + + model_max_budget = jsonencode({ + "gpt-4o" = { + max_budget = 25.0 + } + }) +} +``` + +## Argument Reference + +- `user_id` (Optional, Forces new resource) - Unique ID for the user. Generated by the server if not provided +- `user_email` (Optional) - Email address of the user +- `user_alias` (Optional) - Descriptive name for the user +- `user_role` (Optional) - Role of the user. One of `proxy_admin`, `proxy_admin_viewer`, `internal_user`, `internal_user_viewer` +- `teams` (Optional) - List of team IDs the user belongs to +- `models` (Optional) - Models the user is allowed to call +- `max_budget` (Optional) - Maximum budget in USD for the user +- `budget_duration` (Optional) - Budget reset period, e.g. `30s`, `30m`, `30d` +- `tpm_limit` (Optional) - Tokens per minute limit +- `rpm_limit` (Optional) - Requests per minute limit +- `max_parallel_requests` (Optional) - Maximum number of parallel requests +- `metadata` (Optional) - Map of metadata for the user +- `auto_create_key` (Optional, Default `true`, Forces new resource) - Whether to auto-create an API key on creation +- `send_invite_email` (Optional, Default `false`, Forces new resource) - Whether to send an invite email on creation +- `key_alias` (Optional) - Alias for the auto-created API key +- `aliases` (Optional) - Map of model aliases for the user +- `config` (Optional) - Map of config values for the user +- `permissions` (Optional) - Map of permission values for the user +- `model_max_budget` (Optional) - JSON string of per-model budget config, e.g. `jsonencode({"gpt-4o" = {max_budget = 10.0}})` +- `guardrails` (Optional) - List of guardrails applied to the user's requests +- `blocked` (Optional, Default `false`) - Whether the user is blocked from making requests + +## Attribute Reference + +- `id` - The user ID +- `key` (Sensitive) - The auto-created API key for the user, populated when `auto_create_key` is `true` + +## Import + +Users can be imported using the user ID: + +```shell +terraform import litellm_user.alice +``` diff --git a/terraform/provider/litellm/client.go b/terraform/provider/litellm/client.go index e0aba61477d..0f825d85d31 100644 --- a/terraform/provider/litellm/client.go +++ b/terraform/provider/litellm/client.go @@ -61,6 +61,17 @@ func (c *Client) GetKey(keyID string) (*Key, error) { return nil, err } + // /key/info nests the key's fields under "info"; only "key" itself is + // top-level. Without unwrapping, reads map nothing back into state. + if info, ok := resp["info"].(map[string]interface{}); ok { + if _, present := info["key"]; !present { + if k, ok := resp["key"].(string); ok { + info["key"] = k + } + } + return c.parseKeyResponse(info) + } + return c.parseKeyResponse(resp) } @@ -70,7 +81,6 @@ func (c *Client) UpdateKey(key *Key) (*Key, error) { "key": key.Key, "team_id": key.TeamID, "metadata": key.Metadata, - "budget_duration": key.BudgetDuration, "key_alias": key.KeyAlias, "aliases": key.Aliases, "permissions": key.Permissions, @@ -80,6 +90,12 @@ func (c *Client) UpdateKey(key *Key) (*Key, error) { "blocked": key.Blocked, } + // The proxy rejects an empty-string budget_duration with a 400, so only + // send it when set. + if key.BudgetDuration != "" { + updateData["budget_duration"] = key.BudgetDuration + } + // Only add pointer fields if they are explicitly set if key.MaxBudget != nil { updateData["max_budget"] = *key.MaxBudget @@ -107,6 +123,30 @@ func (c *Client) UpdateKey(key *Key) (*Key, error) { if len(key.Tags) > 0 { updateData["tags"] = key.Tags } + if key.BudgetID != "" { + updateData["budget_id"] = key.BudgetID + } + if len(key.EnforcedParams) > 0 { + updateData["enforced_params"] = key.EnforcedParams + } + if len(key.AllowedRoutes) > 0 { + updateData["allowed_routes"] = key.AllowedRoutes + } + if len(key.AllowedPassthroughRoutes) > 0 { + updateData["allowed_passthrough_routes"] = key.AllowedPassthroughRoutes + } + if key.RPMLimitType != "" { + updateData["rpm_limit_type"] = key.RPMLimitType + } + if key.TPMLimitType != "" { + updateData["tpm_limit_type"] = key.TPMLimitType + } + if len(key.Prompts) > 0 { + updateData["prompts"] = key.Prompts + } + if key.OrganizationID != "" { + updateData["organization_id"] = key.OrganizationID + } resp, err := c.sendRequest("POST", "/key/update", updateData) if err != nil { @@ -251,6 +291,34 @@ func (c *Client) parseKeyResponse(resp map[string]interface{}) (*Key, error) { } } } + case "budget_id": + if s, ok := v.(string); ok { + createdKey.BudgetID = s + } + case "enforced_params": + createdKey.EnforcedParams = toStringSlice(v) + case "allowed_routes": + createdKey.AllowedRoutes = toStringSlice(v) + case "allowed_passthrough_routes": + createdKey.AllowedPassthroughRoutes = toStringSlice(v) + case "rpm_limit_type": + if s, ok := v.(string); ok { + createdKey.RPMLimitType = s + } + case "tpm_limit_type": + if s, ok := v.(string); ok { + createdKey.TPMLimitType = s + } + case "prompts": + createdKey.Prompts = toStringSlice(v) + case "organization_id": + if s, ok := v.(string); ok { + createdKey.OrganizationID = s + } + case "project_id": + if s, ok := v.(string); ok { + createdKey.ProjectID = s + } } } diff --git a/terraform/provider/litellm/data_source_access_group.go b/terraform/provider/litellm/data_source_access_group.go new file mode 100644 index 00000000000..6741a8060c3 --- /dev/null +++ b/terraform/provider/litellm/data_source_access_group.go @@ -0,0 +1,140 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "net/http" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const endpointAccessGroupList = "/access_group/list" + +type accessGroupListResponse struct { + AccessGroups []accessGroupInfoResponse `json:"access_groups"` +} + +func dataSourceLiteLLMAccessGroup() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMAccessGroupRead, + + Schema: map[string]*schema.Schema{ + "access_group": { + Type: schema.TypeString, + Required: true, + Description: "Name of the access group to retrieve", + }, + "model_names": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "deployment_count": { + Type: schema.TypeInt, + Computed: true, + }, + }, + } +} + +func dataSourceLiteLLMAccessGroupRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + name := d.Get("access_group").(string) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("/access_group/%s/info", name), nil) + if err != nil { + return fmt.Errorf("error reading access group: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + return fmt.Errorf("access group '%s' not found", name) + } + + if err := handleResponse(resp, "reading access group"); err != nil { + return err + } + + var info accessGroupInfoResponse + if err := json.NewDecoder(resp.Body).Decode(&info); err != nil { + return fmt.Errorf("error decoding access group info response: %w", err) + } + + d.SetId(GetStringValue(info.AccessGroup, name)) + d.Set("access_group", GetStringValue(info.AccessGroup, name)) + d.Set("model_names", info.ModelNames) + d.Set("deployment_count", info.DeploymentCount) + + return nil +} + +func dataSourceLiteLLMAccessGroups() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMAccessGroupsRead, + + Schema: map[string]*schema.Schema{ + "access_groups": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "access_group": { + Type: schema.TypeString, + Computed: true, + }, + "model_names": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "deployment_count": { + Type: schema.TypeInt, + Computed: true, + }, + }, + }, + }, + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + }, + } +} + +func dataSourceLiteLLMAccessGroupsRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + resp, err := MakeRequest(client, "GET", endpointAccessGroupList, nil) + if err != nil { + return fmt.Errorf("error listing access groups: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing access groups"); err != nil { + return err + } + + var listResp accessGroupListResponse + if err := json.NewDecoder(resp.Body).Decode(&listResp); err != nil { + return fmt.Errorf("error decoding access group list response: %w", err) + } + + groups := make([]map[string]interface{}, 0, len(listResp.AccessGroups)) + ids := make([]string, 0, len(listResp.AccessGroups)) + for _, group := range listResp.AccessGroups { + groups = append(groups, map[string]interface{}{ + "access_group": group.AccessGroup, + "model_names": group.ModelNames, + "deployment_count": group.DeploymentCount, + }) + ids = append(ids, group.AccessGroup) + } + + d.SetId("access_groups") + d.Set("access_groups", groups) + d.Set("ids", ids) + + return nil +} diff --git a/terraform/provider/litellm/data_source_access_group_test.go b/terraform/provider/litellm/data_source_access_group_test.go new file mode 100644 index 00000000000..e07788d823b --- /dev/null +++ b/terraform/provider/litellm/data_source_access_group_test.go @@ -0,0 +1,97 @@ +package litellm + +import ( + "net/http" + "net/http/httptest" + "reflect" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestAccessGroupDataSourceRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" || r.URL.Path != "/access_group/prod-models/info" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + return + } + w.Write(accessGroupInfoJSON("prod-models", []string{"gpt-4", "claude-3"}, 2)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMAccessGroup().Schema, map[string]interface{}{ + "access_group": "prod-models", + }) + + if err := dataSourceLiteLLMAccessGroupRead(d, client); err != nil { + t.Fatalf("data source read failed: %v", err) + } + + if d.Id() != "prod-models" { + t.Fatalf("expected ID 'prod-models', got %q", d.Id()) + } + wantModels := []interface{}{"gpt-4", "claude-3"} + if !reflect.DeepEqual(d.Get("model_names"), wantModels) { + t.Fatalf("expected model_names %v, got %v", wantModels, d.Get("model_names")) + } + if d.Get("deployment_count").(int) != 2 { + t.Fatalf("expected deployment_count 2, got %v", d.Get("deployment_count")) + } +} + +func TestAccessGroupDataSourceReadNotFound(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMAccessGroup().Schema, map[string]interface{}{ + "access_group": "missing", + }) + + if err := dataSourceLiteLLMAccessGroupRead(d, client); err == nil { + t.Fatal("expected error for missing access group, got nil") + } +} + +func TestAccessGroupsDataSourceRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" || r.URL.Path != "/access_group/list" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + return + } + w.Write([]byte(`{"access_groups": [` + + `{"access_group": "group-a", "model_names": ["gpt-4"], "deployment_count": 1},` + + `{"access_group": "group-b", "model_names": ["claude-3"], "deployment_count": 2}]}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMAccessGroups().Schema, map[string]interface{}{}) + + if err := dataSourceLiteLLMAccessGroupsRead(d, client); err != nil { + t.Fatalf("data source read failed: %v", err) + } + + groups := d.Get("access_groups").([]interface{}) + if len(groups) != 2 { + t.Fatalf("expected 2 access groups, got %d", len(groups)) + } + first := groups[0].(map[string]interface{}) + if first["access_group"] != "group-a" { + t.Fatalf("expected first access_group 'group-a', got %v", first["access_group"]) + } + if !reflect.DeepEqual(first["model_names"], []interface{}{"gpt-4"}) { + t.Fatalf("expected first model_names [gpt-4], got %v", first["model_names"]) + } + if first["deployment_count"].(int) != 1 { + t.Fatalf("expected first deployment_count 1, got %v", first["deployment_count"]) + } + if !reflect.DeepEqual(d.Get("ids"), []interface{}{"group-a", "group-b"}) { + t.Fatalf("expected ids [group-a group-b], got %v", d.Get("ids")) + } +} diff --git a/terraform/provider/litellm/data_source_agent.go b/terraform/provider/litellm/data_source_agent.go new file mode 100644 index 00000000000..8e3f12d0d55 --- /dev/null +++ b/terraform/provider/litellm/data_source_agent.go @@ -0,0 +1,281 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "net/http" + "strconv" + "time" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func dataSourceLiteLLMAgent() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMAgentRead, + + Schema: map[string]*schema.Schema{ + "agent_id": { + Type: schema.TypeString, + Required: true, + Description: "Unique identifier of the agent to retrieve.", + }, + "agent_name": { + Type: schema.TypeString, + Computed: true, + }, + "agent_card_params": { + Type: schema.TypeString, + Computed: true, + Description: "A2A agent card as a JSON object string.", + }, + "object_permission": { + Type: schema.TypeString, + Computed: true, + Description: "Access control permissions as a JSON object string.", + }, + "extra_headers": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "tpm_limit": { + Type: schema.TypeInt, + Computed: true, + }, + "rpm_limit": { + Type: schema.TypeInt, + Computed: true, + }, + "session_tpm_limit": { + Type: schema.TypeInt, + Computed: true, + }, + "session_rpm_limit": { + Type: schema.TypeInt, + Computed: true, + }, + "spend": { + Type: schema.TypeFloat, + Computed: true, + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + "created_by": { + Type: schema.TypeString, + Computed: true, + }, + "updated_by": { + Type: schema.TypeString, + Computed: true, + }, + }, + } +} + +func dataSourceLiteLLMAgentRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + agentID := d.Get("agent_id").(string) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf(endpointAgentByID, agentID), nil) + if err != nil { + return fmt.Errorf("error reading agent: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + return fmt.Errorf("agent '%s' not found", agentID) + } + + if err := handleResponse(resp, "reading agent"); err != nil { + return err + } + + var agentResp agentAPIResponse + if err := json.NewDecoder(resp.Body).Decode(&agentResp); err != nil { + return fmt.Errorf("error decoding agent info response: %w", err) + } + + d.SetId(agentResp.AgentID) + d.Set("agent_name", agentResp.AgentName) + + if agentResp.AgentCardParams != nil { + cardJSON, err := json.Marshal(agentResp.AgentCardParams) + if err != nil { + return fmt.Errorf("error encoding agent_card_params: %w", err) + } + d.Set("agent_card_params", string(cardJSON)) + } + if agentResp.ObjectPermission != nil { + permJSON, err := json.Marshal(agentResp.ObjectPermission) + if err != nil { + return fmt.Errorf("error encoding object_permission: %w", err) + } + d.Set("object_permission", string(permJSON)) + } + + if agentResp.ExtraHeaders != nil { + d.Set("extra_headers", agentResp.ExtraHeaders) + } + if agentResp.TPMLimit != nil { + d.Set("tpm_limit", *agentResp.TPMLimit) + } + if agentResp.RPMLimit != nil { + d.Set("rpm_limit", *agentResp.RPMLimit) + } + if agentResp.SessionTPMLimit != nil { + d.Set("session_tpm_limit", *agentResp.SessionTPMLimit) + } + if agentResp.SessionRPMLimit != nil { + d.Set("session_rpm_limit", *agentResp.SessionRPMLimit) + } + if agentResp.Spend != nil { + d.Set("spend", *agentResp.Spend) + } + d.Set("created_at", agentResp.CreatedAt) + d.Set("updated_at", agentResp.UpdatedAt) + d.Set("created_by", agentResp.CreatedBy) + d.Set("updated_by", agentResp.UpdatedBy) + + return nil +} + +func dataSourceLiteLLMAgents() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMAgentsRead, + + Schema: map[string]*schema.Schema{ + "health_check": { + Type: schema.TypeBool, + Optional: true, + Default: false, + Description: "When true, the proxy probes each agent's URL and only returns agents that are " + + "reachable or have no URL.", + }, + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "agents": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "agent_id": { + Type: schema.TypeString, + Computed: true, + }, + "agent_name": { + Type: schema.TypeString, + Computed: true, + }, + "tpm_limit": { + Type: schema.TypeInt, + Computed: true, + }, + "rpm_limit": { + Type: schema.TypeInt, + Computed: true, + }, + "session_tpm_limit": { + Type: schema.TypeInt, + Computed: true, + }, + "session_rpm_limit": { + Type: schema.TypeInt, + Computed: true, + }, + "spend": { + Type: schema.TypeFloat, + Computed: true, + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + "created_by": { + Type: schema.TypeString, + Computed: true, + }, + "updated_by": { + Type: schema.TypeString, + Computed: true, + }, + }, + }, + }, + }, + } +} + +func dataSourceLiteLLMAgentsRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + endpoint := endpointAgents + if d.Get("health_check").(bool) { + endpoint = fmt.Sprintf("%s?health_check=true", endpointAgents) + } + + resp, err := MakeRequest(client, "GET", endpoint, nil) + if err != nil { + return fmt.Errorf("error listing agents: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing agents"); err != nil { + return err + } + + var agentResps []agentAPIResponse + if err := json.NewDecoder(resp.Body).Decode(&agentResps); err != nil { + return fmt.Errorf("error decoding agents list response: %w", err) + } + + ids := make([]string, 0, len(agentResps)) + agents := make([]map[string]interface{}, 0, len(agentResps)) + for _, agentResp := range agentResps { + ids = append(ids, agentResp.AgentID) + + agent := map[string]interface{}{ + "agent_id": agentResp.AgentID, + "agent_name": agentResp.AgentName, + "created_at": agentResp.CreatedAt, + "updated_at": agentResp.UpdatedAt, + "created_by": agentResp.CreatedBy, + "updated_by": agentResp.UpdatedBy, + } + if agentResp.TPMLimit != nil { + agent["tpm_limit"] = *agentResp.TPMLimit + } + if agentResp.RPMLimit != nil { + agent["rpm_limit"] = *agentResp.RPMLimit + } + if agentResp.SessionTPMLimit != nil { + agent["session_tpm_limit"] = *agentResp.SessionTPMLimit + } + if agentResp.SessionRPMLimit != nil { + agent["session_rpm_limit"] = *agentResp.SessionRPMLimit + } + if agentResp.Spend != nil { + agent["spend"] = *agentResp.Spend + } + agents = append(agents, agent) + } + + d.SetId(strconv.FormatInt(time.Now().UnixNano(), 10)) + d.Set("ids", ids) + d.Set("agents", agents) + + return nil +} diff --git a/terraform/provider/litellm/data_source_agent_test.go b/terraform/provider/litellm/data_source_agent_test.go new file mode 100644 index 00000000000..0474cf0017e --- /dev/null +++ b/terraform/provider/litellm/data_source_agent_test.go @@ -0,0 +1,94 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestDataSourceLiteLLMAgentRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/v1/agents/agent-123" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write(agentReadResponseBody()) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMAgent().Schema, map[string]interface{}{ + "agent_id": "agent-123", + }) + + if err := dataSourceLiteLLMAgentRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "agent-123" { + t.Fatalf("expected ID 'agent-123', got %q", d.Id()) + } + if d.Get("agent_name").(string) != "my-agent" { + t.Errorf("expected agent_name 'my-agent', got %q", d.Get("agent_name").(string)) + } + var card map[string]interface{} + if err := json.Unmarshal([]byte(d.Get("agent_card_params").(string)), &card); err != nil { + t.Fatalf("agent_card_params not populated as JSON: %v", err) + } + if card["url"] != "http://agent.local:9999/" { + t.Errorf("expected card url, got %v", card["url"]) + } + if d.Get("spend").(float64) != 1.5 { + t.Errorf("expected spend 1.5, got %v", d.Get("spend")) + } + if d.Get("tpm_limit").(int) != 1000 { + t.Errorf("expected tpm_limit 1000, got %d", d.Get("tpm_limit").(int)) + } +} + +func TestDataSourceLiteLLMAgentsRead(t *testing.T) { + var gotQuery string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/v1/agents" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + gotQuery = r.URL.RawQuery + w.Header().Set("Content-Type", "application/json") + body, _ := json.Marshal([]map[string]interface{}{ + {"agent_id": "agent-1", "agent_name": "first", "tpm_limit": 100, "spend": 0.5}, + {"agent_id": "agent-2", "agent_name": "second"}, + }) + w.Write(body) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMAgents().Schema, map[string]interface{}{ + "health_check": true, + }) + + if err := dataSourceLiteLLMAgentsRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if gotQuery != "health_check=true" { + t.Errorf("expected health_check=true query, got %q", gotQuery) + } + + ids := d.Get("ids").([]interface{}) + if len(ids) != 2 || ids[0] != "agent-1" || ids[1] != "agent-2" { + t.Fatalf("expected ids [agent-1 agent-2], got %v", ids) + } + agents := d.Get("agents").([]interface{}) + if len(agents) != 2 { + t.Fatalf("expected 2 agents, got %d", len(agents)) + } + first := agents[0].(map[string]interface{}) + if first["agent_name"] != "first" || first["tpm_limit"] != 100 || first["spend"] != 0.5 { + t.Errorf("unexpected first agent entry: %v", first) + } + if d.Id() == "" { + t.Fatal("expected data source ID to be set") + } +} diff --git a/terraform/provider/litellm/data_source_budget.go b/terraform/provider/litellm/data_source_budget.go new file mode 100644 index 00000000000..6c493dbedcb --- /dev/null +++ b/terraform/provider/litellm/data_source_budget.go @@ -0,0 +1,195 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "net/http" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const endpointBudgetList = "/budget/list" + +func dataSourceLiteLLMBudget() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMBudgetRead, + + Schema: map[string]*schema.Schema{ + "budget_id": { + Type: schema.TypeString, + Required: true, + Description: "ID of the budget to retrieve", + }, + "max_budget": { + Type: schema.TypeFloat, + Computed: true, + Description: "Hard budget limit in USD", + }, + "soft_budget": { + Type: schema.TypeFloat, + Computed: true, + Description: "Soft budget limit in USD that triggers alerts", + }, + "max_parallel_requests": { + Type: schema.TypeInt, + Computed: true, + Description: "Maximum concurrent requests allowed for this budget", + }, + "tpm_limit": { + Type: schema.TypeInt, + Computed: true, + Description: "Maximum tokens per minute allowed for this budget", + }, + "rpm_limit": { + Type: schema.TypeInt, + Computed: true, + Description: "Maximum requests per minute allowed for this budget", + }, + "budget_duration": { + Type: schema.TypeString, + Computed: true, + Description: "Budget reset period", + }, + "model_max_budget": { + Type: schema.TypeString, + Computed: true, + Description: "JSON string of per-model budget config", + }, + "budget_reset_at": { + Type: schema.TypeString, + Computed: true, + Description: "Datetime when the budget is reset", + }, + }, + } +} + +func dataSourceLiteLLMBudgetRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + budgetID := d.Get("budget_id").(string) + + resp, err := MakeRequest(client, "POST", endpointBudgetInfo, map[string]interface{}{ + "budgets": []string{budgetID}, + }) + if err != nil { + return fmt.Errorf("failed to read budget: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + return fmt.Errorf("budget '%s' not found", budgetID) + } + + if err := handleResponse(resp, "reading budget"); err != nil { + return err + } + + var budgetResps []budgetResponse + if err := json.NewDecoder(resp.Body).Decode(&budgetResps); err != nil { + return fmt.Errorf("error decoding budget info response: %w", err) + } + if len(budgetResps) == 0 { + return fmt.Errorf("budget '%s' not found", budgetID) + } + + d.SetId(budgetID) + setBudgetState(d, budgetResps[0]) + + return nil +} + +func dataSourceLiteLLMBudgets() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMBudgetsRead, + + Schema: map[string]*schema.Schema{ + "budgets": { + Type: schema.TypeList, + Computed: true, + Description: "All budgets configured on the proxy", + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "budget_id": {Type: schema.TypeString, Computed: true}, + "max_budget": {Type: schema.TypeFloat, Computed: true}, + "soft_budget": {Type: schema.TypeFloat, Computed: true}, + "max_parallel_requests": {Type: schema.TypeInt, Computed: true}, + "tpm_limit": {Type: schema.TypeInt, Computed: true}, + "rpm_limit": {Type: schema.TypeInt, Computed: true}, + "budget_duration": {Type: schema.TypeString, Computed: true}, + "model_max_budget": {Type: schema.TypeString, Computed: true}, + "budget_reset_at": {Type: schema.TypeString, Computed: true}, + }, + }, + }, + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "IDs of all budgets configured on the proxy", + }, + }, + } +} + +func budgetListEntry(budgetResp budgetResponse) map[string]interface{} { + entry := map[string]interface{}{ + "budget_id": budgetResp.BudgetID, + } + if budgetResp.MaxBudget != nil { + entry["max_budget"] = *budgetResp.MaxBudget + } + if budgetResp.SoftBudget != nil { + entry["soft_budget"] = *budgetResp.SoftBudget + } + if budgetResp.MaxParallelRequests != nil { + entry["max_parallel_requests"] = *budgetResp.MaxParallelRequests + } + if budgetResp.TPMLimit != nil { + entry["tpm_limit"] = *budgetResp.TPMLimit + } + if budgetResp.RPMLimit != nil { + entry["rpm_limit"] = *budgetResp.RPMLimit + } + if budgetResp.BudgetDuration != nil { + entry["budget_duration"] = *budgetResp.BudgetDuration + } + if encoded, ok := budgetModelMaxBudgetString(budgetResp.ModelMaxBudget); ok { + entry["model_max_budget"] = encoded + } + if budgetResp.BudgetResetAt != nil { + entry["budget_reset_at"] = *budgetResp.BudgetResetAt + } + return entry +} + +func dataSourceLiteLLMBudgetsRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + resp, err := MakeRequest(client, "GET", endpointBudgetList, nil) + if err != nil { + return fmt.Errorf("failed to list budgets: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing budgets"); err != nil { + return err + } + + var budgetResps []budgetResponse + if err := json.NewDecoder(resp.Body).Decode(&budgetResps); err != nil { + return fmt.Errorf("error decoding budget list response: %w", err) + } + + budgets := make([]map[string]interface{}, 0, len(budgetResps)) + ids := make([]string, 0, len(budgetResps)) + for _, budgetResp := range budgetResps { + budgets = append(budgets, budgetListEntry(budgetResp)) + ids = append(ids, budgetResp.BudgetID) + } + + d.SetId("budgets") + d.Set("budgets", budgets) + d.Set("ids", ids) + + return nil +} diff --git a/terraform/provider/litellm/data_source_budget_test.go b/terraform/provider/litellm/data_source_budget_test.go new file mode 100644 index 00000000000..7a4fe0529cb --- /dev/null +++ b/terraform/provider/litellm/data_source_budget_test.go @@ -0,0 +1,107 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestDataSourceBudgetRead_MapsFields(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/budget/info" || r.Method != http.MethodPost { + t.Errorf("expected POST /budget/info, got %s %s", r.Method, r.URL.Path) + } + var payload map[string]interface{} + if err := json.NewDecoder(r.Body).Decode(&payload); err != nil { + t.Fatalf("failed to decode info payload: %v", err) + } + budgets, ok := payload["budgets"].([]interface{}) + if !ok || len(budgets) != 1 || budgets[0] != "bud-ds" { + t.Errorf("expected budgets ['bud-ds'], got %v", payload["budgets"]) + } + w.Write(budgetInfoBody("bud-ds")) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMBudget().Schema, map[string]interface{}{ + "budget_id": "bud-ds", + }) + + if err := dataSourceLiteLLMBudgetRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Id() != "bud-ds" { + t.Fatalf("expected ID 'bud-ds', got %q", d.Id()) + } + if got := d.Get("max_budget").(float64); got != 100.0 { + t.Errorf("expected max_budget 100.0, got %v", got) + } + if got := d.Get("budget_duration").(string); got != "30d" { + t.Errorf("expected budget_duration '30d', got %q", got) + } + if got := d.Get("budget_reset_at").(string); got != "2026-09-01T00:00:00Z" { + t.Errorf("expected budget_reset_at set, got %q", got) + } +} + +func TestDataSourceBudgetsRead_MapsList(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/budget/list" || r.Method != http.MethodGet { + t.Errorf("expected GET /budget/list, got %s %s", r.Method, r.URL.Path) + } + body, _ := json.Marshal([]map[string]interface{}{ + { + "budget_id": "bud-1", + "max_budget": 10.0, + "tpm_limit": 500, + "model_max_budget": map[string]interface{}{"gpt-4o": map[string]interface{}{"max_budget": 1.0}}, + }, + { + "budget_id": "bud-2", + "soft_budget": 5.0, + }, + }) + w.Write(body) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMBudgets().Schema, map[string]interface{}{}) + + if err := dataSourceLiteLLMBudgetsRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + budgets := d.Get("budgets").([]interface{}) + if len(budgets) != 2 { + t.Fatalf("expected 2 budgets, got %d", len(budgets)) + } + first := budgets[0].(map[string]interface{}) + if got := first["budget_id"].(string); got != "bud-1" { + t.Errorf("expected first budget_id 'bud-1', got %q", got) + } + if got := first["max_budget"].(float64); got != 10.0 { + t.Errorf("expected first max_budget 10.0, got %v", got) + } + if got := first["tpm_limit"].(int); got != 500 { + t.Errorf("expected first tpm_limit 500, got %d", got) + } + var mmb map[string]interface{} + if err := json.Unmarshal([]byte(first["model_max_budget"].(string)), &mmb); err != nil { + t.Fatalf("model_max_budget is not valid JSON: %v", err) + } + if _, ok := mmb["gpt-4o"]; !ok { + t.Errorf("expected gpt-4o key in model_max_budget, got %v", mmb) + } + second := budgets[1].(map[string]interface{}) + if got := second["soft_budget"].(float64); got != 5.0 { + t.Errorf("expected second soft_budget 5.0, got %v", got) + } + ids := d.Get("ids").([]interface{}) + if len(ids) != 2 || ids[0] != "bud-1" || ids[1] != "bud-2" { + t.Errorf("expected ids [bud-1 bud-2], got %v", ids) + } +} diff --git a/terraform/provider/litellm/data_source_fallback.go b/terraform/provider/litellm/data_source_fallback.go new file mode 100644 index 00000000000..60cec19851a --- /dev/null +++ b/terraform/provider/litellm/data_source_fallback.go @@ -0,0 +1,71 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "net/http" + "net/url" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation" +) + +func dataSourceLiteLLMFallback() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMFallbackRead, + + Schema: map[string]*schema.Schema{ + "model": { + Type: schema.TypeString, + Required: true, + Description: "The model name to get fallbacks for", + }, + "fallback_type": { + Type: schema.TypeString, + Optional: true, + Default: "general", + ValidateFunc: validation.StringInSlice([]string{"general", "context_window", "content_policy"}, false), + Description: "Type of fallback: 'general' (default), 'context_window', or 'content_policy'", + }, + "fallback_models": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "List of fallback model names in order of priority", + }, + }, + } +} + +func dataSourceLiteLLMFallbackRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + model := d.Get("model").(string) + fallbackType := GetStringValue(d.Get("fallback_type").(string), "general") + + endpoint := fmt.Sprintf("/fallback/%s?fallback_type=%s", url.PathEscape(model), url.QueryEscape(fallbackType)) + resp, err := MakeRequest(client, "GET", endpoint, nil) + if err != nil { + return fmt.Errorf("failed to read fallback: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + return fmt.Errorf("no %s fallbacks configured for model '%s'", fallbackType, model) + } + + if err := handleResponse(resp, "reading fallback"); err != nil { + return err + } + + var fallbackResp FallbackGetResponse + if err := json.NewDecoder(resp.Body).Decode(&fallbackResp); err != nil { + return fmt.Errorf("error decoding fallback response: %w", err) + } + + d.SetId(model) + d.Set("model", GetStringValue(fallbackResp.Model, model)) + d.Set("fallback_models", fallbackResp.FallbackModels) + d.Set("fallback_type", GetStringValue(fallbackResp.FallbackType, fallbackType)) + + return nil +} diff --git a/terraform/provider/litellm/data_source_fallback_test.go b/terraform/provider/litellm/data_source_fallback_test.go new file mode 100644 index 00000000000..12aa879619d --- /dev/null +++ b/terraform/provider/litellm/data_source_fallback_test.go @@ -0,0 +1,63 @@ +package litellm + +import ( + "net/http" + "net/http/httptest" + "reflect" + "strings" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestDataSourceLiteLLMFallbackRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/fallback/gpt-4" { + t.Errorf("expected path /fallback/gpt-4, got %s", r.URL.Path) + } + if got := r.URL.Query().Get("fallback_type"); got != "general" { + t.Errorf("expected fallback_type query 'general', got %q", got) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"model":"gpt-4","fallback_models":["claude-3","gpt-3.5-turbo"],"fallback_type":"general"}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMFallback().Schema, map[string]interface{}{ + "model": "gpt-4", + "fallback_type": "general", + }) + + if err := dataSourceLiteLLMFallbackRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "gpt-4" { + t.Fatalf("expected ID 'gpt-4', got %q", d.Id()) + } + got := d.Get("fallback_models").([]interface{}) + if !reflect.DeepEqual(got, []interface{}{"claude-3", "gpt-3.5-turbo"}) { + t.Fatalf("expected fallback_models [claude-3 gpt-3.5-turbo], got %+v", got) + } +} + +func TestDataSourceLiteLLMFallbackRead_NotFoundErrors(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMFallback().Schema, map[string]interface{}{ + "model": "missing-model", + "fallback_type": "general", + }) + + err := dataSourceLiteLLMFallbackRead(d, client) + if err == nil { + t.Fatal("expected error for missing fallback, got nil") + } + if !strings.Contains(err.Error(), "missing-model") { + t.Fatalf("expected error to name the model, got: %v", err) + } +} diff --git a/terraform/provider/litellm/data_source_guardrail.go b/terraform/provider/litellm/data_source_guardrail.go new file mode 100644 index 00000000000..567221b71e7 --- /dev/null +++ b/terraform/provider/litellm/data_source_guardrail.go @@ -0,0 +1,178 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "net/http" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const endpointGuardrailList = "/guardrails/list" + +func dataSourceLiteLLMGuardrail() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMGuardrailRead, + + Schema: map[string]*schema.Schema{ + "guardrail_id": { + Type: schema.TypeString, + Required: true, + Description: "Unique identifier of the guardrail to retrieve", + }, + "guardrail_name": { + Type: schema.TypeString, + Computed: true, + }, + "guardrail_info": { + Type: schema.TypeMap, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "guardrail_definition_location": { + Type: schema.TypeString, + Computed: true, + Description: "Where the guardrail is defined: 'config' or 'db'", + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + }, + } +} + +type guardrailListItemAPIResponse struct { + GuardrailID string `json:"guardrail_id"` + GuardrailName string `json:"guardrail_name"` + GuardrailInfo map[string]interface{} `json:"guardrail_info"` + GuardrailDefinitionLocation string `json:"guardrail_definition_location"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` +} + +func dataSourceLiteLLMGuardrailRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + guardrailID := d.Get("guardrail_id").(string) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf(endpointGuardrailInfo, guardrailID), nil) + if err != nil { + return fmt.Errorf("failed to read guardrail: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + return fmt.Errorf("guardrail '%s' not found", guardrailID) + } + + if err := handleResponse(resp, "reading guardrail"); err != nil { + return err + } + + var info guardrailListItemAPIResponse + if err := json.NewDecoder(resp.Body).Decode(&info); err != nil { + return fmt.Errorf("error decoding guardrail info response: %w", err) + } + + d.SetId(guardrailID) + d.Set("guardrail_name", info.GuardrailName) + d.Set("guardrail_info", guardrailInfoToStringMap(info.GuardrailInfo)) + d.Set("guardrail_definition_location", info.GuardrailDefinitionLocation) + d.Set("created_at", info.CreatedAt) + d.Set("updated_at", info.UpdatedAt) + // litellm_params is intentionally not exposed: it can carry API keys. + + return nil +} + +func dataSourceLiteLLMGuardrails() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMGuardrailsRead, + + Schema: map[string]*schema.Schema{ + "guardrails": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "guardrail_id": { + Type: schema.TypeString, + Computed: true, + }, + "guardrail_name": { + Type: schema.TypeString, + Computed: true, + }, + "guardrail_info": { + Type: schema.TypeMap, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "guardrail_definition_location": { + Type: schema.TypeString, + Computed: true, + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + }, + }, + }, + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + }, + } +} + +func dataSourceLiteLLMGuardrailsRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + resp, err := MakeRequest(client, "GET", endpointGuardrailList, nil) + if err != nil { + return fmt.Errorf("failed to list guardrails: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing guardrails"); err != nil { + return err + } + + var listResp struct { + Guardrails []guardrailListItemAPIResponse `json:"guardrails"` + } + if err := json.NewDecoder(resp.Body).Decode(&listResp); err != nil { + return fmt.Errorf("error decoding guardrails list response: %w", err) + } + + guardrails := make([]map[string]interface{}, 0, len(listResp.Guardrails)) + ids := make([]string, 0, len(listResp.Guardrails)) + for _, g := range listResp.Guardrails { + guardrails = append(guardrails, map[string]interface{}{ + "guardrail_id": g.GuardrailID, + "guardrail_name": g.GuardrailName, + "guardrail_info": guardrailInfoToStringMap(g.GuardrailInfo), + "guardrail_definition_location": g.GuardrailDefinitionLocation, + "created_at": g.CreatedAt, + "updated_at": g.UpdatedAt, + }) + ids = append(ids, g.GuardrailID) + } + + d.SetId("guardrails") + d.Set("guardrails", guardrails) + d.Set("ids", ids) + + return nil +} diff --git a/terraform/provider/litellm/data_source_guardrail_test.go b/terraform/provider/litellm/data_source_guardrail_test.go new file mode 100644 index 00000000000..4e854f58229 --- /dev/null +++ b/terraform/provider/litellm/data_source_guardrail_test.go @@ -0,0 +1,83 @@ +package litellm + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestDataSourceGuardrailRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" || r.URL.Path != "/guardrails/gid-1/info" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "guardrail_id": "gid-1", + "guardrail_name": "guard1", + "guardrail_info": {"description": "pii guard"}, + "guardrail_definition_location": "db", + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-02T00:00:00Z" + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMGuardrail().Schema, map[string]interface{}{ + "guardrail_id": "gid-1", + }) + + if err := dataSourceLiteLLMGuardrailRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "gid-1" { + t.Fatalf("expected ID 'gid-1', got %q", d.Id()) + } + if got := d.Get("guardrail_name").(string); got != "guard1" { + t.Errorf("expected guardrail_name 'guard1', got %q", got) + } + if got := d.Get("guardrail_definition_location").(string); got != "db" { + t.Errorf("expected guardrail_definition_location 'db', got %q", got) + } + info := d.Get("guardrail_info").(map[string]interface{}) + if info["description"] != "pii guard" { + t.Errorf("expected guardrail_info from API, got: %v", info) + } +} + +func TestDataSourceGuardrailsRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" || r.URL.Path != "/guardrails/list" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"guardrails": [ + {"guardrail_id": "gid-1", "guardrail_name": "guard1", "guardrail_definition_location": "db"}, + {"guardrail_id": "gid-2", "guardrail_name": "guard2", "guardrail_definition_location": "config"} + ]}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMGuardrails().Schema, map[string]interface{}{}) + + if err := dataSourceLiteLLMGuardrailsRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + + guardrails := d.Get("guardrails").([]interface{}) + if len(guardrails) != 2 { + t.Fatalf("expected 2 guardrails, got %d", len(guardrails)) + } + first := guardrails[0].(map[string]interface{}) + if first["guardrail_id"] != "gid-1" || first["guardrail_name"] != "guard1" { + t.Errorf("unexpected first guardrail: %v", first) + } + ids := d.Get("ids").([]interface{}) + if len(ids) != 2 || ids[0] != "gid-1" || ids[1] != "gid-2" { + t.Errorf("unexpected ids: %v", ids) + } +} diff --git a/terraform/provider/litellm/data_source_key.go b/terraform/provider/litellm/data_source_key.go new file mode 100644 index 00000000000..2407e82211c --- /dev/null +++ b/terraform/provider/litellm/data_source_key.go @@ -0,0 +1,384 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "log" + "net/url" + "strconv" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const ( + endpointKeyInfo = "/key/info" + endpointKeyList = "/key/list" +) + +type keyInfoDetail struct { + Token string `json:"token"` + KeyName string `json:"key_name"` + KeyAlias string `json:"key_alias"` + Spend float64 `json:"spend"` + MaxBudget *float64 `json:"max_budget"` + Models []string `json:"models"` + UserID string `json:"user_id"` + TeamID string `json:"team_id"` + OrgID string `json:"org_id"` + TPMLimit *int `json:"tpm_limit"` + RPMLimit *int `json:"rpm_limit"` + MaxParallelRequests *int `json:"max_parallel_requests"` + BudgetDuration string `json:"budget_duration"` + Metadata map[string]interface{} `json:"metadata"` + Blocked *bool `json:"blocked"` + Expires string `json:"expires"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` +} + +type keyInfoEnvelope struct { + Key string `json:"key"` + Info keyInfoDetail `json:"info"` +} + +type keyListEnvelope struct { + Keys []keyInfoDetail `json:"keys"` + TotalCount int `json:"total_count"` + CurrentPage int `json:"current_page"` + TotalPages int `json:"total_pages"` +} + +func dataSourceLiteLLMKey() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMKeyRead, + + Schema: map[string]*schema.Schema{ + "key": { + Type: schema.TypeString, + Required: true, + Sensitive: true, + Description: "The API key (or its hash) to look up", + }, + "token_id": { + Type: schema.TypeString, + Computed: true, + Description: "Hashed token identifier of the key", + }, + "key_name": { + Type: schema.TypeString, + Computed: true, + Description: "Redacted display name of the key", + }, + "key_alias": { + Type: schema.TypeString, + Computed: true, + }, + "models": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "spend": { + Type: schema.TypeFloat, + Computed: true, + }, + "max_budget": { + Type: schema.TypeFloat, + Computed: true, + }, + "user_id": { + Type: schema.TypeString, + Computed: true, + }, + "team_id": { + Type: schema.TypeString, + Computed: true, + }, + "organization_id": { + Type: schema.TypeString, + Computed: true, + }, + "tpm_limit": { + Type: schema.TypeInt, + Computed: true, + }, + "rpm_limit": { + Type: schema.TypeInt, + Computed: true, + }, + "max_parallel_requests": { + Type: schema.TypeInt, + Computed: true, + }, + "budget_duration": { + Type: schema.TypeString, + Computed: true, + }, + "metadata": { + Type: schema.TypeMap, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "tags": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "blocked": { + Type: schema.TypeBool, + Computed: true, + }, + "expires": { + Type: schema.TypeString, + Computed: true, + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + }, + } +} + +func dataSourceLiteLLMKeyRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + // Look up by the SHA-256 token hash so the raw key never appears in the + // request URL, where reverse-proxy access logs could record it. + key := hashedKeyToken(d.Get("key").(string)) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("%s?key=%s", endpointKeyInfo, url.QueryEscape(key)), nil) + if err != nil { + return fmt.Errorf("failed to read key info: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "reading key info"); err != nil { + return err + } + + var envelope keyInfoEnvelope + if err := json.NewDecoder(resp.Body).Decode(&envelope); err != nil { + return fmt.Errorf("failed to decode key info response: %w", err) + } + info := envelope.Info + + // Never persist the raw key as the ID; the hashed token is safe to store. + d.SetId(GetStringValue(info.Token, "key")) + d.Set("token_id", info.Token) + d.Set("key_name", info.KeyName) + d.Set("key_alias", info.KeyAlias) + d.Set("models", info.Models) + d.Set("spend", info.Spend) + if info.MaxBudget != nil { + d.Set("max_budget", *info.MaxBudget) + } + d.Set("user_id", info.UserID) + d.Set("team_id", info.TeamID) + d.Set("organization_id", info.OrgID) + if info.TPMLimit != nil { + d.Set("tpm_limit", *info.TPMLimit) + } + if info.RPMLimit != nil { + d.Set("rpm_limit", *info.RPMLimit) + } + if info.MaxParallelRequests != nil { + d.Set("max_parallel_requests", *info.MaxParallelRequests) + } + d.Set("budget_duration", info.BudgetDuration) + + metadata := map[string]string{} + for k, v := range info.Metadata { + if s, ok := v.(string); ok { + metadata[k] = s + } + } + d.Set("metadata", metadata) + d.Set("tags", toStringSlice(info.Metadata["tags"])) + + if info.Blocked != nil { + d.Set("blocked", *info.Blocked) + } + d.Set("expires", info.Expires) + d.Set("created_at", info.CreatedAt) + d.Set("updated_at", info.UpdatedAt) + + log.Printf("[INFO] Successfully read key info for token: %s", info.Token) + return nil +} + +func dataSourceLiteLLMKeys() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMKeysRead, + + Schema: map[string]*schema.Schema{ + "page": { + Type: schema.TypeInt, + Optional: true, + Default: 1, + Description: "Page number for pagination", + }, + "size": { + Type: schema.TypeInt, + Optional: true, + Default: 100, + Description: "Number of keys per page", + }, + "user_id": { + Type: schema.TypeString, + Optional: true, + Description: "Filter keys by user ID", + }, + "team_id": { + Type: schema.TypeString, + Optional: true, + Description: "Filter keys by team ID", + }, + "organization_id": { + Type: schema.TypeString, + Optional: true, + Description: "Filter keys by organization ID", + }, + "key_alias": { + Type: schema.TypeString, + Optional: true, + Description: "Filter keys by key alias", + }, + "include_team_keys": { + Type: schema.TypeBool, + Optional: true, + Description: "Include all keys for teams the caller is an admin of", + }, + "total_count": { + Type: schema.TypeInt, + Computed: true, + }, + "total_pages": { + Type: schema.TypeInt, + Computed: true, + }, + "current_page": { + Type: schema.TypeInt, + Computed: true, + }, + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Hashed token identifiers of the returned keys", + }, + "keys": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "token_id": {Type: schema.TypeString, Computed: true}, + "key_name": {Type: schema.TypeString, Computed: true}, + "key_alias": {Type: schema.TypeString, Computed: true}, + "spend": {Type: schema.TypeFloat, Computed: true}, + "max_budget": {Type: schema.TypeFloat, Computed: true}, + "models": {Type: schema.TypeList, Computed: true, Elem: &schema.Schema{Type: schema.TypeString}}, + "user_id": {Type: schema.TypeString, Computed: true}, + "team_id": {Type: schema.TypeString, Computed: true}, + "organization_id": {Type: schema.TypeString, Computed: true}, + "tpm_limit": {Type: schema.TypeInt, Computed: true}, + "rpm_limit": {Type: schema.TypeInt, Computed: true}, + "budget_duration": {Type: schema.TypeString, Computed: true}, + "blocked": {Type: schema.TypeBool, Computed: true}, + "expires": {Type: schema.TypeString, Computed: true}, + "created_at": {Type: schema.TypeString, Computed: true}, + "updated_at": {Type: schema.TypeString, Computed: true}, + }, + }, + }, + }, + } +} + +func dataSourceLiteLLMKeysRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + query := url.Values{} + query.Set("return_full_object", "true") + query.Set("page", strconv.Itoa(d.Get("page").(int))) + query.Set("size", strconv.Itoa(d.Get("size").(int))) + for param, attr := range map[string]string{ + "user_id": "user_id", + "team_id": "team_id", + "organization_id": "organization_id", + "key_alias": "key_alias", + } { + if v, ok := d.GetOk(attr); ok { + query.Set(param, v.(string)) + } + } + if d.Get("include_team_keys").(bool) { + query.Set("include_team_keys", "true") + } + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("%s?%s", endpointKeyList, query.Encode()), nil) + if err != nil { + return fmt.Errorf("failed to list keys: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing keys"); err != nil { + return err + } + + var envelope keyListEnvelope + if err := json.NewDecoder(resp.Body).Decode(&envelope); err != nil { + return fmt.Errorf("failed to decode key list response: %w", err) + } + + ids := make([]string, 0, len(envelope.Keys)) + keys := make([]map[string]interface{}, 0, len(envelope.Keys)) + for _, k := range envelope.Keys { + ids = append(ids, k.Token) + keys = append(keys, map[string]interface{}{ + "token_id": k.Token, + "key_name": k.KeyName, + "key_alias": k.KeyAlias, + "spend": k.Spend, + "max_budget": keyDerefFloat(k.MaxBudget), + "models": k.Models, + "user_id": k.UserID, + "team_id": k.TeamID, + "organization_id": k.OrgID, + "tpm_limit": keyDerefInt(k.TPMLimit), + "rpm_limit": keyDerefInt(k.RPMLimit), + "budget_duration": k.BudgetDuration, + "blocked": k.Blocked != nil && *k.Blocked, + "expires": k.Expires, + "created_at": k.CreatedAt, + "updated_at": k.UpdatedAt, + }) + } + + d.SetId(query.Encode()) + d.Set("total_count", envelope.TotalCount) + d.Set("total_pages", envelope.TotalPages) + d.Set("current_page", envelope.CurrentPage) + d.Set("ids", ids) + d.Set("keys", keys) + + log.Printf("[INFO] Successfully listed %d keys", len(keys)) + return nil +} + +func keyDerefFloat(v *float64) float64 { + if v == nil { + return 0 + } + return *v +} + +func keyDerefInt(v *int) int { + if v == nil { + return 0 + } + return *v +} diff --git a/terraform/provider/litellm/data_source_key_test.go b/terraform/provider/litellm/data_source_key_test.go new file mode 100644 index 00000000000..5f13e385c00 --- /dev/null +++ b/terraform/provider/litellm/data_source_key_test.go @@ -0,0 +1,198 @@ +package litellm + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestDataSourceKeyRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/key/info" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + if got := r.URL.Query().Get("key"); got != "43d0a3c1b9dc2739952a8ffc4ee4f41ea34da6587cbc717c3a51185b9fac611c" { + t.Errorf("expected key query param to be the token hash, got %q", got) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "key": "sk-raw-secret", + "info": { + "token": "hashed-token-123", + "key_name": "sk-...cret", + "key_alias": "ci-key", + "spend": 12.5, + "max_budget": 100, + "models": ["gpt-4o", "claude-3"], + "user_id": "user-1", + "team_id": "team-1", + "org_id": "org-1", + "tpm_limit": 1000, + "rpm_limit": 60, + "max_parallel_requests": 5, + "budget_duration": "30d", + "metadata": {"env": "prod", "tags": ["alpha", "beta"]}, + "blocked": true, + "expires": "2027-01-01T00:00:00Z", + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-02-01T00:00:00Z" + } + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMKey().Schema, map[string]interface{}{ + "key": "sk-raw-secret", + }) + + if err := dataSourceLiteLLMKeyRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Id() != "hashed-token-123" { + t.Fatalf("expected ID 'hashed-token-123', got %q", d.Id()) + } + checks := map[string]interface{}{ + "token_id": "hashed-token-123", + "key_name": "sk-...cret", + "key_alias": "ci-key", + "spend": 12.5, + "max_budget": 100.0, + "user_id": "user-1", + "team_id": "team-1", + "organization_id": "org-1", + "tpm_limit": 1000, + "rpm_limit": 60, + "max_parallel_requests": 5, + "budget_duration": "30d", + "blocked": true, + "expires": "2027-01-01T00:00:00Z", + } + for attr, want := range checks { + if got := d.Get(attr); got != want { + t.Errorf("attr %s: expected %v, got %v", attr, want, got) + } + } + models := d.Get("models").([]interface{}) + if len(models) != 2 || models[0] != "gpt-4o" { + t.Errorf("unexpected models: %v", models) + } + tags := d.Get("tags").([]interface{}) + if len(tags) != 2 || tags[0] != "alpha" { + t.Errorf("unexpected tags: %v", tags) + } + metadata := d.Get("metadata").(map[string]interface{}) + if metadata["env"] != "prod" { + t.Errorf("unexpected metadata: %v", metadata) + } + if _, hasTags := metadata["tags"]; hasTags { + t.Errorf("non-string metadata value should not be in the metadata map: %v", metadata) + } +} + +func TestDataSourceKeyReadNotFound(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + w.Write([]byte(`{"detail": {"error": "key not found"}}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMKey().Schema, map[string]interface{}{ + "key": "sk-missing", + }) + + if err := dataSourceLiteLLMKeyRead(d, client); err == nil { + t.Fatal("expected error for missing key, got nil") + } +} + +func TestDataSourceKeysRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/key/list" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + query := r.URL.Query() + if query.Get("return_full_object") != "true" { + t.Errorf("expected return_full_object=true, got %q", query.Get("return_full_object")) + } + if query.Get("team_id") != "team-1" { + t.Errorf("expected team_id=team-1, got %q", query.Get("team_id")) + } + if query.Get("page") != "2" || query.Get("size") != "10" { + t.Errorf("expected page=2 size=10, got page=%q size=%q", query.Get("page"), query.Get("size")) + } + if query.Get("include_team_keys") != "true" { + t.Errorf("expected include_team_keys=true, got %q", query.Get("include_team_keys")) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "keys": [ + {"token": "tok-1", "key_alias": "a", "team_id": "team-1", "spend": 1.5, "max_budget": 10, "models": ["m1"], "blocked": false}, + {"token": "tok-2", "key_alias": "b", "team_id": "team-1", "spend": 0, "blocked": true} + ], + "total_count": 2, + "current_page": 2, + "total_pages": 1 + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMKeys().Schema, map[string]interface{}{ + "team_id": "team-1", + "page": 2, + "size": 10, + "include_team_keys": true, + }) + + if err := dataSourceLiteLLMKeysRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Id() == "" { + t.Fatal("expected data source ID to be set") + } + if got := d.Get("total_count").(int); got != 2 { + t.Errorf("expected total_count 2, got %d", got) + } + ids := d.Get("ids").([]interface{}) + if len(ids) != 2 || ids[0] != "tok-1" || ids[1] != "tok-2" { + t.Errorf("unexpected ids: %v", ids) + } + keys := d.Get("keys").([]interface{}) + if len(keys) != 2 { + t.Fatalf("expected 2 keys, got %d", len(keys)) + } + first := keys[0].(map[string]interface{}) + if first["token_id"] != "tok-1" || first["key_alias"] != "a" || first["max_budget"] != 10.0 { + t.Errorf("unexpected first key: %v", first) + } + second := keys[1].(map[string]interface{}) + if second["blocked"] != true || second["max_budget"] != 0.0 { + t.Errorf("unexpected second key: %v", second) + } +} + +// Regression for the security review finding: the singular key data source +// must query /key/info by the SHA-256 token hash, never the raw sk- value. +func TestDataSourceKeyQueriesByTokenHash(t *testing.T) { + var gotQuery string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotQuery = r.URL.Query().Get("key") + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"key": "hash", "info": {"token": "hash", "key_alias": "a"}}`)) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMKey().Schema, map[string]interface{}{"key": "sk-test-123"}) + if err := dataSourceLiteLLMKeyRead(d, NewClient(srv.URL, "master-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + if gotQuery != keyBlockTestHash { + t.Fatalf("query key = %q, want the token hash %q", gotQuery, keyBlockTestHash) + } +} diff --git a/terraform/provider/litellm/data_source_mcp_server.go b/terraform/provider/litellm/data_source_mcp_server.go new file mode 100644 index 00000000000..0918605b61b --- /dev/null +++ b/terraform/provider/litellm/data_source_mcp_server.go @@ -0,0 +1,271 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "log" + "net/url" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +// mcpServerDetail intentionally omits env, credentials, and static_headers: +// those may hold secrets and must never reach data source state. +type mcpServerDetail struct { + ServerID string `json:"server_id"` + ServerName string `json:"server_name"` + Alias string `json:"alias"` + Description string `json:"description"` + URL string `json:"url"` + Transport string `json:"transport"` + SpecVersion string `json:"spec_version"` + AuthType string `json:"auth_type"` + MCPAccessGroups []string `json:"mcp_access_groups"` + AllowedTools []string `json:"allowed_tools"` + ExtraHeaders []string `json:"extra_headers"` + Command string `json:"command"` + Args []string `json:"args"` + AllowAllKeys bool `json:"allow_all_keys"` + Status string `json:"status"` + LastHealthCheck string `json:"last_health_check"` + HealthCheckError string `json:"health_check_error"` + CreatedAt string `json:"created_at"` + CreatedBy string `json:"created_by"` + UpdatedAt string `json:"updated_at"` + UpdatedBy string `json:"updated_by"` +} + +func dataSourceLiteLLMMCPServer() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMMCPServerRead, + + Schema: map[string]*schema.Schema{ + "server_id": { + Type: schema.TypeString, + Required: true, + Description: "Unique identifier of the MCP server to retrieve", + }, + "server_name": { + Type: schema.TypeString, + Computed: true, + }, + "alias": { + Type: schema.TypeString, + Computed: true, + }, + "description": { + Type: schema.TypeString, + Computed: true, + }, + "url": { + Type: schema.TypeString, + Computed: true, + }, + "transport": { + Type: schema.TypeString, + Computed: true, + }, + "spec_version": { + Type: schema.TypeString, + Computed: true, + }, + "auth_type": { + Type: schema.TypeString, + Computed: true, + }, + "mcp_access_groups": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "allowed_tools": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "extra_headers": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Names of request headers forwarded to the MCP server", + }, + "command": { + Type: schema.TypeString, + Computed: true, + }, + "args": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "allow_all_keys": { + Type: schema.TypeBool, + Computed: true, + }, + "status": { + Type: schema.TypeString, + Computed: true, + }, + "last_health_check": { + Type: schema.TypeString, + Computed: true, + }, + "health_check_error": { + Type: schema.TypeString, + Computed: true, + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "created_by": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_by": { + Type: schema.TypeString, + Computed: true, + }, + }, + } +} + +func dataSourceLiteLLMMCPServerRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + serverID := d.Get("server_id").(string) + + endpoint := fmt.Sprintf("%s/%s", endpointMCPServerRead, serverID) + resp, err := MakeRequest(client, "GET", endpoint, nil) + if err != nil { + return fmt.Errorf("failed to read MCP server: %w", err) + } + defer resp.Body.Close() + + var server mcpServerDetail + if err := handleMCPAPIResponse(resp, &server, client); err != nil { + if err.Error() == "mcp_server_not_found" { + return fmt.Errorf("MCP server %q not found", serverID) + } + return fmt.Errorf("failed to read MCP server: %w", err) + } + + d.SetId(GetStringValue(server.ServerID, serverID)) + d.Set("server_name", server.ServerName) + d.Set("alias", server.Alias) + d.Set("description", server.Description) + d.Set("url", server.URL) + d.Set("transport", server.Transport) + d.Set("spec_version", server.SpecVersion) + d.Set("auth_type", server.AuthType) + d.Set("mcp_access_groups", server.MCPAccessGroups) + d.Set("allowed_tools", server.AllowedTools) + d.Set("extra_headers", server.ExtraHeaders) + d.Set("command", server.Command) + d.Set("args", server.Args) + d.Set("allow_all_keys", server.AllowAllKeys) + d.Set("status", server.Status) + d.Set("last_health_check", server.LastHealthCheck) + d.Set("health_check_error", server.HealthCheckError) + d.Set("created_at", server.CreatedAt) + d.Set("created_by", server.CreatedBy) + d.Set("updated_at", server.UpdatedAt) + d.Set("updated_by", server.UpdatedBy) + + log.Printf("[INFO] Successfully read MCP server with ID: %s", serverID) + return nil +} + +func dataSourceLiteLLMMCPServers() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMMCPServersRead, + + Schema: map[string]*schema.Schema{ + "team_id": { + Type: schema.TypeString, + Optional: true, + Description: "Filter to servers this team can access plus globally available servers", + }, + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "IDs of the returned MCP servers", + }, + "mcp_servers": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "server_id": {Type: schema.TypeString, Computed: true}, + "server_name": {Type: schema.TypeString, Computed: true}, + "alias": {Type: schema.TypeString, Computed: true}, + "description": {Type: schema.TypeString, Computed: true}, + "url": {Type: schema.TypeString, Computed: true}, + "transport": {Type: schema.TypeString, Computed: true}, + "spec_version": {Type: schema.TypeString, Computed: true}, + "auth_type": {Type: schema.TypeString, Computed: true}, + "allow_all_keys": {Type: schema.TypeBool, Computed: true}, + "status": {Type: schema.TypeString, Computed: true}, + "created_at": {Type: schema.TypeString, Computed: true}, + "updated_at": {Type: schema.TypeString, Computed: true}, + }, + }, + }, + }, + } +} + +func dataSourceLiteLLMMCPServersRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + endpoint := endpointMCPServerRead + if v, ok := d.GetOk("team_id"); ok { + endpoint = fmt.Sprintf("%s?team_id=%s", endpointMCPServerRead, url.QueryEscape(v.(string))) + } + + resp, err := MakeRequest(client, "GET", endpoint, nil) + if err != nil { + return fmt.Errorf("failed to list MCP servers: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing MCP servers"); err != nil { + return err + } + + var serverList []mcpServerDetail + if err := json.NewDecoder(resp.Body).Decode(&serverList); err != nil { + return fmt.Errorf("failed to decode MCP server list response: %w", err) + } + + ids := make([]string, 0, len(serverList)) + servers := make([]map[string]interface{}, 0, len(serverList)) + for _, server := range serverList { + ids = append(ids, server.ServerID) + servers = append(servers, map[string]interface{}{ + "server_id": server.ServerID, + "server_name": server.ServerName, + "alias": server.Alias, + "description": server.Description, + "url": server.URL, + "transport": server.Transport, + "spec_version": server.SpecVersion, + "auth_type": server.AuthType, + "allow_all_keys": server.AllowAllKeys, + "status": server.Status, + "created_at": server.CreatedAt, + "updated_at": server.UpdatedAt, + }) + } + + d.SetId(GetStringValue(d.Get("team_id").(string), "all")) + d.Set("ids", ids) + d.Set("mcp_servers", servers) + + log.Printf("[INFO] Successfully listed %d MCP servers", len(servers)) + return nil +} diff --git a/terraform/provider/litellm/data_source_mcp_server_test.go b/terraform/provider/litellm/data_source_mcp_server_test.go new file mode 100644 index 00000000000..e061d7ffb56 --- /dev/null +++ b/terraform/provider/litellm/data_source_mcp_server_test.go @@ -0,0 +1,150 @@ +package litellm + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestDataSourceMCPServerRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/v1/mcp/server/srv-123" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "server_id": "srv-123", + "server_name": "github-mcp", + "alias": "gh", + "description": "GitHub MCP server", + "url": "https://mcp.example.com", + "transport": "http", + "spec_version": "2024-11-05", + "auth_type": "bearer", + "mcp_access_groups": ["dev"], + "allowed_tools": ["list_repos"], + "extra_headers": ["x-request-id"], + "command": "", + "args": [], + "env": {"SECRET_TOKEN": "should-never-surface"}, + "static_headers": {"Authorization": "Bearer should-never-surface"}, + "allow_all_keys": true, + "status": "healthy", + "last_health_check": "2026-02-01T00:00:00Z", + "health_check_error": "", + "created_at": "2026-01-01T00:00:00Z", + "created_by": "admin", + "updated_at": "2026-02-01T00:00:00Z", + "updated_by": "admin" + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMMCPServer().Schema, map[string]interface{}{ + "server_id": "srv-123", + }) + + if err := dataSourceLiteLLMMCPServerRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Id() != "srv-123" { + t.Fatalf("expected ID 'srv-123', got %q", d.Id()) + } + checks := map[string]interface{}{ + "server_name": "github-mcp", + "alias": "gh", + "description": "GitHub MCP server", + "url": "https://mcp.example.com", + "transport": "http", + "spec_version": "2024-11-05", + "auth_type": "bearer", + "allow_all_keys": true, + "status": "healthy", + "last_health_check": "2026-02-01T00:00:00Z", + "created_by": "admin", + } + for attr, want := range checks { + if got := d.Get(attr); got != want { + t.Errorf("attr %s: expected %v, got %v", attr, want, got) + } + } + groups := d.Get("mcp_access_groups").([]interface{}) + if len(groups) != 1 || groups[0] != "dev" { + t.Errorf("unexpected access groups: %v", groups) + } + tools := d.Get("allowed_tools").([]interface{}) + if len(tools) != 1 || tools[0] != "list_repos" { + t.Errorf("unexpected allowed tools: %v", tools) + } + headers := d.Get("extra_headers").([]interface{}) + if len(headers) != 1 || headers[0] != "x-request-id" { + t.Errorf("unexpected extra headers: %v", headers) + } +} + +func TestDataSourceMCPServerReadNotFound(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + w.Write([]byte(`{"detail": {"error": "MCP server not found"}}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMMCPServer().Schema, map[string]interface{}{ + "server_id": "srv-missing", + }) + + if err := dataSourceLiteLLMMCPServerRead(d, client); err == nil { + t.Fatal("expected error for missing MCP server, got nil") + } +} + +func TestDataSourceMCPServersRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/v1/mcp/server" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + if got := r.URL.Query().Get("team_id"); got != "team-1" { + t.Errorf("expected team_id 'team-1', got %q", got) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`[ + {"server_id": "srv-1", "server_name": "one", "url": "https://one.example.com", "transport": "http", "status": "healthy", "allow_all_keys": false}, + {"server_id": "srv-2", "server_name": "two", "url": "https://two.example.com", "transport": "sse", "status": "unknown", "allow_all_keys": true} + ]`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMMCPServers().Schema, map[string]interface{}{ + "team_id": "team-1", + }) + + if err := dataSourceLiteLLMMCPServersRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Id() != "team-1" { + t.Fatalf("expected ID 'team-1', got %q", d.Id()) + } + ids := d.Get("ids").([]interface{}) + if len(ids) != 2 || ids[0] != "srv-1" || ids[1] != "srv-2" { + t.Errorf("unexpected ids: %v", ids) + } + servers := d.Get("mcp_servers").([]interface{}) + if len(servers) != 2 { + t.Fatalf("expected 2 servers, got %d", len(servers)) + } + first := servers[0].(map[string]interface{}) + if first["server_name"] != "one" || first["transport"] != "http" || first["allow_all_keys"] != false { + t.Errorf("unexpected first server: %v", first) + } + second := servers[1].(map[string]interface{}) + if second["status"] != "unknown" || second["allow_all_keys"] != true { + t.Errorf("unexpected second server: %v", second) + } +} diff --git a/terraform/provider/litellm/data_source_model.go b/terraform/provider/litellm/data_source_model.go new file mode 100644 index 00000000000..78af04ac160 --- /dev/null +++ b/terraform/provider/litellm/data_source_model.go @@ -0,0 +1,260 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "log" + "net/url" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const endpointModelInfoV1 = "/v1/model/info" + +// modelInfoParams intentionally maps only the non-sensitive litellm_params fields; +// credentials (api_key, aws_secret_access_key, ...) must never reach state. +type modelInfoParams struct { + Model string `json:"model"` + CustomLLMProvider string `json:"custom_llm_provider"` + APIBase string `json:"api_base"` + APIVersion string `json:"api_version"` + TPM int `json:"tpm"` + RPM int `json:"rpm"` +} + +type modelInfoMeta struct { + ID string `json:"id"` + DBModel bool `json:"db_model"` + BaseModel string `json:"base_model"` + Tier string `json:"tier"` + Mode string `json:"mode"` + TeamID string `json:"team_id"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` +} + +type modelInfoEntry struct { + ModelName string `json:"model_name"` + LiteLLMParams modelInfoParams `json:"litellm_params"` + ModelInfo modelInfoMeta `json:"model_info"` +} + +type modelInfoEnvelope struct { + Data json.RawMessage `json:"data"` +} + +// /v1/model/info returns data as a single object on the DB path and as a +// one-element list on the config path, so both shapes must be handled. +func modelDecodeInfoEntries(raw json.RawMessage) ([]modelInfoEntry, error) { + var single modelInfoEntry + if err := json.Unmarshal(raw, &single); err == nil { + return []modelInfoEntry{single}, nil + } + var list []modelInfoEntry + if err := json.Unmarshal(raw, &list); err != nil { + return nil, fmt.Errorf("failed to decode model info data: %w", err) + } + return list, nil +} + +func dataSourceLiteLLMModel() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMModelRead, + + Schema: map[string]*schema.Schema{ + "model_id": { + Type: schema.TypeString, + Required: true, + Description: "LiteLLM model ID (the x-litellm-model-id response header value)", + }, + "model_name": { + Type: schema.TypeString, + Computed: true, + }, + "model": { + Type: schema.TypeString, + Computed: true, + Description: "The underlying litellm_params model, e.g. openai/gpt-4o", + }, + "custom_llm_provider": { + Type: schema.TypeString, + Computed: true, + }, + "model_api_base": { + Type: schema.TypeString, + Computed: true, + }, + "api_version": { + Type: schema.TypeString, + Computed: true, + }, + "tpm": { + Type: schema.TypeInt, + Computed: true, + }, + "rpm": { + Type: schema.TypeInt, + Computed: true, + }, + "base_model": { + Type: schema.TypeString, + Computed: true, + }, + "tier": { + Type: schema.TypeString, + Computed: true, + }, + "mode": { + Type: schema.TypeString, + Computed: true, + }, + "team_id": { + Type: schema.TypeString, + Computed: true, + }, + "db_model": { + Type: schema.TypeBool, + Computed: true, + }, + }, + } +} + +func dataSourceLiteLLMModelRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + modelID := d.Get("model_id").(string) + + endpoint := fmt.Sprintf("%s?litellm_model_id=%s", endpointModelInfoV1, url.QueryEscape(modelID)) + resp, err := MakeRequest(client, "GET", endpoint, nil) + if err != nil { + return fmt.Errorf("failed to read model info: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "reading model info"); err != nil { + return err + } + + var envelope modelInfoEnvelope + if err := json.NewDecoder(resp.Body).Decode(&envelope); err != nil { + return fmt.Errorf("failed to decode model info response: %w", err) + } + + entries, err := modelDecodeInfoEntries(envelope.Data) + if err != nil { + return err + } + if len(entries) == 0 { + return fmt.Errorf("model with id %q not found", modelID) + } + entry := entries[0] + + d.SetId(GetStringValue(entry.ModelInfo.ID, modelID)) + d.Set("model_name", entry.ModelName) + d.Set("model", entry.LiteLLMParams.Model) + d.Set("custom_llm_provider", entry.LiteLLMParams.CustomLLMProvider) + d.Set("model_api_base", entry.LiteLLMParams.APIBase) + d.Set("api_version", entry.LiteLLMParams.APIVersion) + d.Set("tpm", entry.LiteLLMParams.TPM) + d.Set("rpm", entry.LiteLLMParams.RPM) + d.Set("base_model", entry.ModelInfo.BaseModel) + d.Set("tier", entry.ModelInfo.Tier) + d.Set("mode", entry.ModelInfo.Mode) + d.Set("team_id", entry.ModelInfo.TeamID) + d.Set("db_model", entry.ModelInfo.DBModel) + + log.Printf("[INFO] Successfully read model with ID: %s", modelID) + return nil +} + +func dataSourceLiteLLMModels() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMModelsRead, + + Schema: map[string]*schema.Schema{ + "team_id": { + Type: schema.TypeString, + Optional: true, + Description: "Filter models to those accessible by this team", + }, + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "LiteLLM model IDs of the returned models", + }, + "models": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "id": {Type: schema.TypeString, Computed: true}, + "model_name": {Type: schema.TypeString, Computed: true}, + "model": {Type: schema.TypeString, Computed: true}, + "custom_llm_provider": {Type: schema.TypeString, Computed: true}, + "model_api_base": {Type: schema.TypeString, Computed: true}, + "base_model": {Type: schema.TypeString, Computed: true}, + "tier": {Type: schema.TypeString, Computed: true}, + "mode": {Type: schema.TypeString, Computed: true}, + "team_id": {Type: schema.TypeString, Computed: true}, + "db_model": {Type: schema.TypeBool, Computed: true}, + }, + }, + }, + }, + } +} + +func dataSourceLiteLLMModelsRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + endpoint := endpointModelInfoV1 + if v, ok := d.GetOk("team_id"); ok { + endpoint = fmt.Sprintf("%s?teamId=%s", endpointModelInfoV1, url.QueryEscape(v.(string))) + } + + resp, err := MakeRequest(client, "GET", endpoint, nil) + if err != nil { + return fmt.Errorf("failed to list models: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing models"); err != nil { + return err + } + + var envelope modelInfoEnvelope + if err := json.NewDecoder(resp.Body).Decode(&envelope); err != nil { + return fmt.Errorf("failed to decode model list response: %w", err) + } + + entries, err := modelDecodeInfoEntries(envelope.Data) + if err != nil { + return err + } + + ids := make([]string, 0, len(entries)) + models := make([]map[string]interface{}, 0, len(entries)) + for _, entry := range entries { + ids = append(ids, entry.ModelInfo.ID) + models = append(models, map[string]interface{}{ + "id": entry.ModelInfo.ID, + "model_name": entry.ModelName, + "model": entry.LiteLLMParams.Model, + "custom_llm_provider": entry.LiteLLMParams.CustomLLMProvider, + "model_api_base": entry.LiteLLMParams.APIBase, + "base_model": entry.ModelInfo.BaseModel, + "tier": entry.ModelInfo.Tier, + "mode": entry.ModelInfo.Mode, + "team_id": entry.ModelInfo.TeamID, + "db_model": entry.ModelInfo.DBModel, + }) + } + + d.SetId(GetStringValue(d.Get("team_id").(string), "all")) + d.Set("ids", ids) + d.Set("models", models) + + log.Printf("[INFO] Successfully listed %d models", len(models)) + return nil +} diff --git a/terraform/provider/litellm/data_source_model_test.go b/terraform/provider/litellm/data_source_model_test.go new file mode 100644 index 00000000000..97d7f07dcd8 --- /dev/null +++ b/terraform/provider/litellm/data_source_model_test.go @@ -0,0 +1,149 @@ +package litellm + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestDataSourceModelReadSingleObject(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/v1/model/info" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + if got := r.URL.Query().Get("litellm_model_id"); got != "model-abc" { + t.Errorf("expected litellm_model_id 'model-abc', got %q", got) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "data": { + "model_name": "gpt-4o-alias", + "litellm_params": { + "model": "openai/gpt-4o", + "custom_llm_provider": "openai", + "api_base": "https://api.openai.com/v1", + "api_version": "2024-06-01", + "api_key": "sk-should-never-surface", + "tpm": 100000, + "rpm": 500 + }, + "model_info": { + "id": "model-abc", + "db_model": true, + "base_model": "gpt-4o", + "tier": "paid", + "mode": "chat", + "team_id": "team-1" + } + } + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMModel().Schema, map[string]interface{}{ + "model_id": "model-abc", + }) + + if err := dataSourceLiteLLMModelRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Id() != "model-abc" { + t.Fatalf("expected ID 'model-abc', got %q", d.Id()) + } + checks := map[string]interface{}{ + "model_name": "gpt-4o-alias", + "model": "openai/gpt-4o", + "custom_llm_provider": "openai", + "model_api_base": "https://api.openai.com/v1", + "api_version": "2024-06-01", + "tpm": 100000, + "rpm": 500, + "base_model": "gpt-4o", + "tier": "paid", + "mode": "chat", + "team_id": "team-1", + "db_model": true, + } + for attr, want := range checks { + if got := d.Get(attr); got != want { + t.Errorf("attr %s: expected %v, got %v", attr, want, got) + } + } +} + +func TestDataSourceModelReadListShape(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "data": [{ + "model_name": "claude-alias", + "litellm_params": {"model": "anthropic/claude-opus-4", "custom_llm_provider": "anthropic"}, + "model_info": {"id": "model-xyz", "mode": "chat"} + }] + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMModel().Schema, map[string]interface{}{ + "model_id": "model-xyz", + }) + + if err := dataSourceLiteLLMModelRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + if d.Id() != "model-xyz" { + t.Fatalf("expected ID 'model-xyz', got %q", d.Id()) + } + if got := d.Get("model").(string); got != "anthropic/claude-opus-4" { + t.Errorf("expected model 'anthropic/claude-opus-4', got %q", got) + } +} + +func TestDataSourceModelsRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/v1/model/info" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + if got := r.URL.Query().Get("teamId"); got != "team-1" { + t.Errorf("expected teamId 'team-1', got %q", got) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "data": [ + {"model_name": "a", "litellm_params": {"model": "openai/a", "custom_llm_provider": "openai"}, "model_info": {"id": "id-1", "db_model": true}}, + {"model_name": "b", "litellm_params": {"model": "anthropic/b", "custom_llm_provider": "anthropic"}, "model_info": {"id": "id-2"}} + ] + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMModels().Schema, map[string]interface{}{ + "team_id": "team-1", + }) + + if err := dataSourceLiteLLMModelsRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Id() != "team-1" { + t.Fatalf("expected ID 'team-1', got %q", d.Id()) + } + ids := d.Get("ids").([]interface{}) + if len(ids) != 2 || ids[0] != "id-1" || ids[1] != "id-2" { + t.Errorf("unexpected ids: %v", ids) + } + models := d.Get("models").([]interface{}) + if len(models) != 2 { + t.Fatalf("expected 2 models, got %d", len(models)) + } + first := models[0].(map[string]interface{}) + if first["model_name"] != "a" || first["custom_llm_provider"] != "openai" || first["db_model"] != true { + t.Errorf("unexpected first model: %v", first) + } +} diff --git a/terraform/provider/litellm/data_source_organization.go b/terraform/provider/litellm/data_source_organization.go new file mode 100644 index 00000000000..43ad869f3a1 --- /dev/null +++ b/terraform/provider/litellm/data_source_organization.go @@ -0,0 +1,270 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "log" + "net/url" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const endpointOrganizationList = "/organization/list" + +type organizationBudget struct { + MaxBudget *float64 `json:"max_budget"` + SoftBudget *float64 `json:"soft_budget"` + TPMLimit *int `json:"tpm_limit"` + RPMLimit *int `json:"rpm_limit"` + MaxParallelRequests *int `json:"max_parallel_requests"` + BudgetDuration string `json:"budget_duration"` +} + +type organizationDetail struct { + OrganizationID string `json:"organization_id"` + OrganizationAlias string `json:"organization_alias"` + BudgetID string `json:"budget_id"` + Models []string `json:"models"` + Spend float64 `json:"spend"` + Metadata map[string]interface{} `json:"metadata"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` + Budget *organizationBudget `json:"litellm_budget_table"` +} + +func dataSourceLiteLLMOrganization() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMOrganizationRead, + + Schema: map[string]*schema.Schema{ + "organization_id": { + Type: schema.TypeString, + Required: true, + Description: "Unique identifier of the organization to retrieve", + }, + "organization_alias": { + Type: schema.TypeString, + Computed: true, + }, + "budget_id": { + Type: schema.TypeString, + Computed: true, + }, + "models": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "spend": { + Type: schema.TypeFloat, + Computed: true, + }, + "metadata": { + Type: schema.TypeMap, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "max_budget": { + Type: schema.TypeFloat, + Computed: true, + }, + "soft_budget": { + Type: schema.TypeFloat, + Computed: true, + }, + "tpm_limit": { + Type: schema.TypeInt, + Computed: true, + }, + "rpm_limit": { + Type: schema.TypeInt, + Computed: true, + }, + "max_parallel_requests": { + Type: schema.TypeInt, + Computed: true, + }, + "budget_duration": { + Type: schema.TypeString, + Computed: true, + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + }, + } +} + +func dataSourceLiteLLMOrganizationRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + orgID := d.Get("organization_id").(string) + + endpoint := fmt.Sprintf("%s?organization_id=%s", endpointOrganizationInfo, url.QueryEscape(orgID)) + resp, err := MakeRequest(client, "GET", endpoint, nil) + if err != nil { + return fmt.Errorf("failed to read organization: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "reading organization info"); err != nil { + return err + } + + var org organizationDetail + if err := json.NewDecoder(resp.Body).Decode(&org); err != nil { + return fmt.Errorf("failed to decode organization info response: %w", err) + } + + d.SetId(GetStringValue(org.OrganizationID, orgID)) + organizationSetDetail(d, org) + + log.Printf("[INFO] Successfully read organization with ID: %s", orgID) + return nil +} + +func organizationSetDetail(d *schema.ResourceData, org organizationDetail) { + d.Set("organization_alias", org.OrganizationAlias) + d.Set("budget_id", org.BudgetID) + d.Set("models", org.Models) + d.Set("spend", org.Spend) + + metadata := map[string]string{} + for k, v := range org.Metadata { + if s, ok := v.(string); ok { + metadata[k] = s + } + } + d.Set("metadata", metadata) + + if org.Budget != nil { + if org.Budget.MaxBudget != nil { + d.Set("max_budget", *org.Budget.MaxBudget) + } + if org.Budget.SoftBudget != nil { + d.Set("soft_budget", *org.Budget.SoftBudget) + } + if org.Budget.TPMLimit != nil { + d.Set("tpm_limit", *org.Budget.TPMLimit) + } + if org.Budget.RPMLimit != nil { + d.Set("rpm_limit", *org.Budget.RPMLimit) + } + if org.Budget.MaxParallelRequests != nil { + d.Set("max_parallel_requests", *org.Budget.MaxParallelRequests) + } + d.Set("budget_duration", org.Budget.BudgetDuration) + } + d.Set("created_at", org.CreatedAt) + d.Set("updated_at", org.UpdatedAt) +} + +func dataSourceLiteLLMOrganizations() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMOrganizationsRead, + + Schema: map[string]*schema.Schema{ + "org_alias": { + Type: schema.TypeString, + Optional: true, + Description: "Filter organizations by alias", + }, + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "IDs of the returned organizations", + }, + "organizations": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "organization_id": {Type: schema.TypeString, Computed: true}, + "organization_alias": {Type: schema.TypeString, Computed: true}, + "budget_id": {Type: schema.TypeString, Computed: true}, + "models": {Type: schema.TypeList, Computed: true, Elem: &schema.Schema{Type: schema.TypeString}}, + "spend": {Type: schema.TypeFloat, Computed: true}, + "max_budget": {Type: schema.TypeFloat, Computed: true}, + "tpm_limit": {Type: schema.TypeInt, Computed: true}, + "rpm_limit": {Type: schema.TypeInt, Computed: true}, + "budget_duration": {Type: schema.TypeString, Computed: true}, + "created_at": {Type: schema.TypeString, Computed: true}, + "updated_at": {Type: schema.TypeString, Computed: true}, + }, + }, + }, + }, + } +} + +func dataSourceLiteLLMOrganizationsRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + endpoint := endpointOrganizationList + if v, ok := d.GetOk("org_alias"); ok { + endpoint = fmt.Sprintf("%s?org_alias=%s", endpointOrganizationList, url.QueryEscape(v.(string))) + } + + resp, err := MakeRequest(client, "GET", endpoint, nil) + if err != nil { + return fmt.Errorf("failed to list organizations: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing organizations"); err != nil { + return err + } + + var orgList []organizationDetail + if err := json.NewDecoder(resp.Body).Decode(&orgList); err != nil { + return fmt.Errorf("failed to decode organization list response: %w", err) + } + + ids := make([]string, 0, len(orgList)) + orgs := make([]map[string]interface{}, 0, len(orgList)) + for _, org := range orgList { + ids = append(ids, org.OrganizationID) + item := map[string]interface{}{ + "organization_id": org.OrganizationID, + "organization_alias": org.OrganizationAlias, + "budget_id": org.BudgetID, + "models": org.Models, + "spend": org.Spend, + "created_at": org.CreatedAt, + "updated_at": org.UpdatedAt, + } + if org.Budget != nil { + item["max_budget"] = organizationDerefFloat(org.Budget.MaxBudget) + item["tpm_limit"] = organizationDerefInt(org.Budget.TPMLimit) + item["rpm_limit"] = organizationDerefInt(org.Budget.RPMLimit) + item["budget_duration"] = org.Budget.BudgetDuration + } + orgs = append(orgs, item) + } + + d.SetId(GetStringValue(d.Get("org_alias").(string), "all")) + d.Set("ids", ids) + d.Set("organizations", orgs) + + log.Printf("[INFO] Successfully listed %d organizations", len(orgs)) + return nil +} + +func organizationDerefFloat(v *float64) float64 { + if v == nil { + return 0 + } + return *v +} + +func organizationDerefInt(v *int) int { + if v == nil { + return 0 + } + return *v +} diff --git a/terraform/provider/litellm/data_source_organization_test.go b/terraform/provider/litellm/data_source_organization_test.go new file mode 100644 index 00000000000..23e3e75eaee --- /dev/null +++ b/terraform/provider/litellm/data_source_organization_test.go @@ -0,0 +1,120 @@ +package litellm + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestDataSourceOrganizationRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/organization/info" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + if got := r.URL.Query().Get("organization_id"); got != "org-123" { + t.Errorf("expected organization_id 'org-123', got %q", got) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "organization_id": "org-123", + "organization_alias": "acme-org", + "budget_id": "budget-1", + "models": ["gpt-4o"], + "spend": 77.5, + "metadata": {"env": "prod"}, + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-02-01T00:00:00Z", + "litellm_budget_table": { + "max_budget": 1000, + "soft_budget": 800, + "tpm_limit": 50000, + "rpm_limit": 500, + "max_parallel_requests": 20, + "budget_duration": "30d" + } + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMOrganization().Schema, map[string]interface{}{ + "organization_id": "org-123", + }) + + if err := dataSourceLiteLLMOrganizationRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Id() != "org-123" { + t.Fatalf("expected ID 'org-123', got %q", d.Id()) + } + checks := map[string]interface{}{ + "organization_alias": "acme-org", + "budget_id": "budget-1", + "spend": 77.5, + "max_budget": 1000.0, + "soft_budget": 800.0, + "tpm_limit": 50000, + "rpm_limit": 500, + "max_parallel_requests": 20, + "budget_duration": "30d", + "created_at": "2026-01-01T00:00:00Z", + } + for attr, want := range checks { + if got := d.Get(attr); got != want { + t.Errorf("attr %s: expected %v, got %v", attr, want, got) + } + } + metadata := d.Get("metadata").(map[string]interface{}) + if metadata["env"] != "prod" { + t.Errorf("unexpected metadata: %v", metadata) + } +} + +func TestDataSourceOrganizationsRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/organization/list" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + if got := r.URL.Query().Get("org_alias"); got != "acme" { + t.Errorf("expected org_alias 'acme', got %q", got) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`[ + {"organization_id": "org-1", "organization_alias": "acme", "spend": 1.5, "litellm_budget_table": {"max_budget": 100, "tpm_limit": 10, "rpm_limit": 5, "budget_duration": "7d"}}, + {"organization_id": "org-2", "organization_alias": "acme-eu", "spend": 0} + ]`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMOrganizations().Schema, map[string]interface{}{ + "org_alias": "acme", + }) + + if err := dataSourceLiteLLMOrganizationsRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Id() != "acme" { + t.Fatalf("expected ID 'acme', got %q", d.Id()) + } + ids := d.Get("ids").([]interface{}) + if len(ids) != 2 || ids[0] != "org-1" || ids[1] != "org-2" { + t.Errorf("unexpected ids: %v", ids) + } + orgs := d.Get("organizations").([]interface{}) + if len(orgs) != 2 { + t.Fatalf("expected 2 organizations, got %d", len(orgs)) + } + first := orgs[0].(map[string]interface{}) + if first["organization_alias"] != "acme" || first["max_budget"] != 100.0 || first["budget_duration"] != "7d" { + t.Errorf("unexpected first organization: %v", first) + } + second := orgs[1].(map[string]interface{}) + if second["organization_id"] != "org-2" || second["max_budget"] != 0.0 { + t.Errorf("unexpected second organization: %v", second) + } +} diff --git a/terraform/provider/litellm/data_source_project.go b/terraform/provider/litellm/data_source_project.go new file mode 100644 index 00000000000..d30ce346d38 --- /dev/null +++ b/terraform/provider/litellm/data_source_project.go @@ -0,0 +1,255 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "net/http" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const endpointProjectList = "/project/list" + +func dataSourceLiteLLMProject() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMProjectRead, + + Schema: map[string]*schema.Schema{ + "project_id": { + Type: schema.TypeString, + Required: true, + Description: "Unique identifier of the project to retrieve", + }, + "project_alias": { + Type: schema.TypeString, + Computed: true, + Description: "Human-friendly name for the project", + }, + "description": { + Type: schema.TypeString, + Computed: true, + Description: "Description of the project", + }, + "team_id": { + Type: schema.TypeString, + Computed: true, + Description: "The team ID this project belongs to", + }, + "budget_id": { + Type: schema.TypeString, + Computed: true, + Description: "Budget ID associated with this project", + }, + "models": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "List of models the project can access", + }, + "max_budget": { + Type: schema.TypeFloat, + Computed: true, + Description: "Maximum budget for this project", + }, + "soft_budget": { + Type: schema.TypeFloat, + Computed: true, + Description: "Soft budget limit for warnings", + }, + "budget_duration": { + Type: schema.TypeString, + Computed: true, + Description: "Budget reset duration", + }, + "tpm_limit": { + Type: schema.TypeInt, + Computed: true, + Description: "Tokens per minute limit", + }, + "rpm_limit": { + Type: schema.TypeInt, + Computed: true, + Description: "Requests per minute limit", + }, + "max_parallel_requests": { + Type: schema.TypeInt, + Computed: true, + Description: "Maximum parallel requests allowed", + }, + "blocked": { + Type: schema.TypeBool, + Computed: true, + Description: "Whether the project is blocked from making requests", + }, + "spend": { + Type: schema.TypeFloat, + Computed: true, + Description: "Current spend for the project", + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + Description: "Timestamp when the project was created", + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + Description: "Timestamp when the project was last updated", + }, + "created_by": { + Type: schema.TypeString, + Computed: true, + Description: "User that created the project", + }, + "updated_by": { + Type: schema.TypeString, + Computed: true, + Description: "User that last updated the project", + }, + }, + } +} + +func dataSourceLiteLLMProjectRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + projectID := d.Get("project_id").(string) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("%s?project_id=%s", endpointProjectInfo, projectID), nil) + if err != nil { + return fmt.Errorf("failed to read project: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + return fmt.Errorf("project '%s' not found", projectID) + } + + if err := handleResponse(resp, "reading project"); err != nil { + return err + } + + var projResp projectResponse + if err := json.NewDecoder(resp.Body).Decode(&projResp); err != nil { + return fmt.Errorf("error decoding project info response: %w", err) + } + + d.SetId(projResp.ProjectID) + d.Set("project_id", projResp.ProjectID) + d.Set("project_alias", projResp.ProjectAlias) + d.Set("description", projResp.Description) + d.Set("team_id", projResp.TeamID) + d.Set("budget_id", projResp.BudgetID) + d.Set("models", projResp.Models) + d.Set("blocked", projResp.Blocked) + d.Set("spend", projResp.Spend) + d.Set("created_at", projResp.CreatedAt) + d.Set("updated_at", projResp.UpdatedAt) + d.Set("created_by", projResp.CreatedBy) + d.Set("updated_by", projResp.UpdatedBy) + + if bt := projResp.LitellmBudgetTable; bt != nil { + if bt.MaxBudget != nil { + d.Set("max_budget", *bt.MaxBudget) + } + if bt.SoftBudget != nil { + d.Set("soft_budget", *bt.SoftBudget) + } + if bt.MaxParallelRequests != nil { + d.Set("max_parallel_requests", *bt.MaxParallelRequests) + } + if bt.TPMLimit != nil { + d.Set("tpm_limit", *bt.TPMLimit) + } + if bt.RPMLimit != nil { + d.Set("rpm_limit", *bt.RPMLimit) + } + d.Set("budget_duration", bt.BudgetDuration) + } + + return nil +} + +func dataSourceLiteLLMProjects() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMProjectsRead, + + Schema: map[string]*schema.Schema{ + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "IDs of all projects", + }, + "projects": { + Type: schema.TypeList, + Computed: true, + Description: "List of projects", + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "project_id": {Type: schema.TypeString, Computed: true}, + "project_alias": {Type: schema.TypeString, Computed: true}, + "description": {Type: schema.TypeString, Computed: true}, + "team_id": {Type: schema.TypeString, Computed: true}, + "budget_id": {Type: schema.TypeString, Computed: true}, + "models": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "blocked": {Type: schema.TypeBool, Computed: true}, + "spend": {Type: schema.TypeFloat, Computed: true}, + "created_at": {Type: schema.TypeString, Computed: true}, + "updated_at": {Type: schema.TypeString, Computed: true}, + "created_by": {Type: schema.TypeString, Computed: true}, + "updated_by": {Type: schema.TypeString, Computed: true}, + }, + }, + }, + }, + } +} + +func dataSourceLiteLLMProjectsRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + resp, err := MakeRequest(client, "GET", endpointProjectList, nil) + if err != nil { + return fmt.Errorf("failed to list projects: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing projects"); err != nil { + return err + } + + var projResps []projectResponse + if err := json.NewDecoder(resp.Body).Decode(&projResps); err != nil { + return fmt.Errorf("error decoding project list response: %w", err) + } + + ids := make([]string, 0, len(projResps)) + projects := make([]map[string]interface{}, 0, len(projResps)) + for _, projResp := range projResps { + ids = append(ids, projResp.ProjectID) + projects = append(projects, map[string]interface{}{ + "project_id": projResp.ProjectID, + "project_alias": projResp.ProjectAlias, + "description": projResp.Description, + "team_id": projResp.TeamID, + "budget_id": projResp.BudgetID, + "models": projResp.Models, + "blocked": projResp.Blocked, + "spend": projResp.Spend, + "created_at": projResp.CreatedAt, + "updated_at": projResp.UpdatedAt, + "created_by": projResp.CreatedBy, + "updated_by": projResp.UpdatedBy, + }) + } + + d.SetId("litellm-projects") + d.Set("ids", ids) + d.Set("projects", projects) + + return nil +} diff --git a/terraform/provider/litellm/data_source_project_test.go b/terraform/provider/litellm/data_source_project_test.go new file mode 100644 index 00000000000..0224655f79c --- /dev/null +++ b/terraform/provider/litellm/data_source_project_test.go @@ -0,0 +1,104 @@ +package litellm + +import ( + "net/http" + "net/http/httptest" + "reflect" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestDataSourceLiteLLMProjectRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/project/info" || r.Method != http.MethodGet { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + if got := r.URL.Query().Get("project_id"); got != "proj-123" { + t.Errorf("expected project_id query 'proj-123', got %q", got) + } + w.Write([]byte(projectInfoBody)) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMProject().Schema, map[string]interface{}{ + "project_id": "proj-123", + }) + + if err := dataSourceLiteLLMProjectRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Id() != "proj-123" { + t.Fatalf("expected ID 'proj-123', got %q", d.Id()) + } + checks := map[string]interface{}{ + "project_alias": "ml-experiments", + "description": "ML experimentation project", + "team_id": "team-1", + "budget_id": "bud-9", + "spend": 12.5, + "max_budget": 100.0, + "tpm_limit": 5000, + "budget_duration": "30d", + "created_by": "admin", + } + for key, want := range checks { + if got := d.Get(key); got != want { + t.Errorf("expected %s %v, got %v", key, want, got) + } + } + if !reflect.DeepEqual(d.Get("models"), []interface{}{"gpt-4"}) { + t.Errorf("expected models ['gpt-4'], got %v", d.Get("models")) + } +} + +func TestDataSourceLiteLLMProjectRead_NotFound(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMProject().Schema, map[string]interface{}{ + "project_id": "gone", + }) + + if err := dataSourceLiteLLMProjectRead(d, NewClient(srv.URL, "test-key", true)); err == nil { + t.Fatal("expected error for missing project, got nil") + } +} + +func TestDataSourceLiteLLMProjectsRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/project/list" || r.Method != http.MethodGet { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Write([]byte(`[ + ` + projectInfoBody + `, + {"project_id": "proj-456", "project_alias": "second", "team_id": "team-2", "models": [], "spend": 0.0} + ]`)) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMProjects().Schema, map[string]interface{}{}) + + if err := dataSourceLiteLLMProjectsRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + if !reflect.DeepEqual(d.Get("ids"), []interface{}{"proj-123", "proj-456"}) { + t.Errorf("expected ids ['proj-123', 'proj-456'], got %v", d.Get("ids")) + } + if got := d.Get("projects.#").(int); got != 2 { + t.Fatalf("expected 2 projects, got %d", got) + } + if got := d.Get("projects.0.project_alias").(string); got != "ml-experiments" { + t.Errorf("expected projects.0.project_alias 'ml-experiments', got %q", got) + } + if got := d.Get("projects.0.spend").(float64); got != 12.5 { + t.Errorf("expected projects.0.spend 12.5, got %v", got) + } + if got := d.Get("projects.1.team_id").(string); got != "team-2" { + t.Errorf("expected projects.1.team_id 'team-2', got %q", got) + } +} diff --git a/terraform/provider/litellm/data_source_prompt.go b/terraform/provider/litellm/data_source_prompt.go new file mode 100644 index 00000000000..0a42c951a40 --- /dev/null +++ b/terraform/provider/litellm/data_source_prompt.go @@ -0,0 +1,243 @@ +package litellm + +import ( + "encoding/json" + "fmt" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func dataSourceLiteLLMPrompt() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMPromptRead, + + Schema: map[string]*schema.Schema{ + "prompt_id": { + Type: schema.TypeString, + Required: true, + Description: "Unique identifier of the prompt to retrieve", + }, + "environment": { + Type: schema.TypeString, + Optional: true, + Description: "Environment to fetch the prompt from (e.g. 'development', 'production')", + }, + "prompt_integration": { + Type: schema.TypeString, + Computed: true, + }, + "api_base": { + Type: schema.TypeString, + Computed: true, + }, + "provider_specific_query_params": { + Type: schema.TypeString, + Computed: true, + }, + "ignore_prompt_manager_model": { + Type: schema.TypeBool, + Computed: true, + }, + "ignore_prompt_manager_optional_params": { + Type: schema.TypeBool, + Computed: true, + }, + "dotprompt_content": { + Type: schema.TypeString, + Computed: true, + }, + "prompt_type": { + Type: schema.TypeString, + Computed: true, + }, + "version": { + Type: schema.TypeInt, + Computed: true, + }, + "environments": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + }, + } +} + +func dataSourceLiteLLMPromptRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + promptID := d.Get("prompt_id").(string) + + endpoint := fmt.Sprintf(endpointPromptInfo, promptID) + if env := d.Get("environment").(string); env != "" { + endpoint = fmt.Sprintf("/prompts/%s/info?environment=%s", promptID, env) + } + + resp, err := MakeRequest(client, "GET", endpoint, nil) + if err != nil { + return fmt.Errorf("failed to read prompt: %w", err) + } + defer resp.Body.Close() + + if promptIsNotFoundResponse(resp) { + return fmt.Errorf("prompt '%s' not found", promptID) + } + + if err := handleResponse(resp, "reading prompt"); err != nil { + return err + } + + var info struct { + PromptSpec promptSpecAPIResponse `json:"prompt_spec"` + Environments []string `json:"environments"` + } + if err := json.NewDecoder(resp.Body).Decode(&info); err != nil { + return fmt.Errorf("error decoding prompt info response: %w", err) + } + + d.SetId(info.PromptSpec.PromptID) + d.Set("prompt_id", info.PromptSpec.PromptID) + d.Set("version", info.PromptSpec.Version) + d.Set("environments", info.Environments) + d.Set("created_at", info.PromptSpec.CreatedAt) + d.Set("updated_at", info.PromptSpec.UpdatedAt) + + params := info.PromptSpec.LitellmParams + if v, ok := params["prompt_integration"].(string); ok { + d.Set("prompt_integration", v) + } + if v, ok := params["api_base"].(string); ok { + d.Set("api_base", v) + } + if v, ok := params["dotprompt_content"].(string); ok { + d.Set("dotprompt_content", v) + } + if v, ok := params["ignore_prompt_manager_model"].(bool); ok { + d.Set("ignore_prompt_manager_model", v) + } + if v, ok := params["ignore_prompt_manager_optional_params"].(bool); ok { + d.Set("ignore_prompt_manager_optional_params", v) + } + if v, ok := params["provider_specific_query_params"].(map[string]interface{}); ok { + if encoded, err := json.Marshal(v); err == nil { + d.Set("provider_specific_query_params", string(encoded)) + } + } + if v, ok := info.PromptSpec.PromptInfo["prompt_type"].(string); ok { + d.Set("prompt_type", v) + } + // api_key is intentionally not exposed. + + return nil +} + +func dataSourceLiteLLMPrompts() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMPromptsRead, + + Schema: map[string]*schema.Schema{ + "environment": { + Type: schema.TypeString, + Optional: true, + Description: "Filter prompts by environment (e.g. 'development', 'production')", + }, + "prompts": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "prompt_id": { + Type: schema.TypeString, + Computed: true, + }, + "prompt_integration": { + Type: schema.TypeString, + Computed: true, + }, + "prompt_type": { + Type: schema.TypeString, + Computed: true, + }, + "version": { + Type: schema.TypeInt, + Computed: true, + }, + "environment": { + Type: schema.TypeString, + Computed: true, + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + }, + }, + }, + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + }, + } +} + +func dataSourceLiteLLMPromptsRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + endpoint := endpointPromptList + if env := d.Get("environment").(string); env != "" { + endpoint = fmt.Sprintf("/prompts/list?environment=%s", env) + } + + resp, err := MakeRequest(client, "GET", endpoint, nil) + if err != nil { + return fmt.Errorf("failed to list prompts: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing prompts"); err != nil { + return err + } + + var listResp struct { + Prompts []promptSpecAPIResponse `json:"prompts"` + } + if err := json.NewDecoder(resp.Body).Decode(&listResp); err != nil { + return fmt.Errorf("error decoding prompts list response: %w", err) + } + + prompts := make([]map[string]interface{}, 0, len(listResp.Prompts)) + ids := make([]string, 0, len(listResp.Prompts)) + for _, p := range listResp.Prompts { + integration, _ := p.LitellmParams["prompt_integration"].(string) + promptType, _ := p.PromptInfo["prompt_type"].(string) + prompts = append(prompts, map[string]interface{}{ + "prompt_id": p.PromptID, + "prompt_integration": integration, + "prompt_type": promptType, + "version": p.Version, + "environment": p.Environment, + "created_at": p.CreatedAt, + "updated_at": p.UpdatedAt, + }) + ids = append(ids, p.PromptID) + } + + d.SetId("prompts") + d.Set("prompts", prompts) + d.Set("ids", ids) + + return nil +} diff --git a/terraform/provider/litellm/data_source_prompt_test.go b/terraform/provider/litellm/data_source_prompt_test.go new file mode 100644 index 00000000000..ded71c5549a --- /dev/null +++ b/terraform/provider/litellm/data_source_prompt_test.go @@ -0,0 +1,92 @@ +package litellm + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestDataSourcePromptRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" || r.URL.Path != "/prompts/p1/info" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(promptInfoJSON("p1"))) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMPrompt().Schema, map[string]interface{}{ + "prompt_id": "p1", + }) + + if err := dataSourceLiteLLMPromptRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "p1" { + t.Fatalf("expected ID 'p1', got %q", d.Id()) + } + if got := d.Get("prompt_integration").(string); got != "langfuse" { + t.Errorf("expected prompt_integration 'langfuse', got %q", got) + } + if got := d.Get("prompt_type").(string); got != "db" { + t.Errorf("expected prompt_type 'db', got %q", got) + } + if got := d.Get("version").(int); got != 3 { + t.Errorf("expected version 3, got %d", got) + } + envs := d.Get("environments").([]interface{}) + if len(envs) != 1 || envs[0] != "development" { + t.Errorf("unexpected environments: %v", envs) + } +} + +func TestDataSourcePromptsRead_WithEnvironmentFilter(t *testing.T) { + var gotQuery string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" || r.URL.Path != "/prompts/list" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + gotQuery = r.URL.RawQuery + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"prompts": [ + { + "prompt_id": "p1", + "litellm_params": {"prompt_integration": "langfuse"}, + "prompt_info": {"prompt_type": "db"}, + "version": 2, + "environment": "production" + } + ]}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMPrompts().Schema, map[string]interface{}{ + "environment": "production", + }) + + if err := dataSourceLiteLLMPromptsRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if gotQuery != "environment=production" { + t.Fatalf("expected environment filter in query, got %q", gotQuery) + } + + prompts := d.Get("prompts").([]interface{}) + if len(prompts) != 1 { + t.Fatalf("expected 1 prompt, got %d", len(prompts)) + } + first := prompts[0].(map[string]interface{}) + if first["prompt_id"] != "p1" || first["prompt_integration"] != "langfuse" || + first["prompt_type"] != "db" || first["version"] != 2 || first["environment"] != "production" { + t.Errorf("unexpected prompt item: %v", first) + } + ids := d.Get("ids").([]interface{}) + if len(ids) != 1 || ids[0] != "p1" { + t.Errorf("unexpected ids: %v", ids) + } +} diff --git a/terraform/provider/litellm/data_source_search_tool.go b/terraform/provider/litellm/data_source_search_tool.go new file mode 100644 index 00000000000..2050b87281b --- /dev/null +++ b/terraform/provider/litellm/data_source_search_tool.go @@ -0,0 +1,179 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "net/http" + "strconv" + "time" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func dataSourceLiteLLMSearchTool() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMSearchToolRead, + + Schema: map[string]*schema.Schema{ + "search_tool_id": { + Type: schema.TypeString, + Required: true, + Description: "Unique identifier of the search tool to retrieve.", + }, + "search_tool_name": { + Type: schema.TypeString, + Computed: true, + }, + "search_tool_info": { + Type: schema.TypeString, + Computed: true, + Description: "Additional metadata as a JSON object string.", + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + }, + } +} + +func dataSourceLiteLLMSearchToolRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + searchToolID := d.Get("search_tool_id").(string) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf(endpointSearchToolByID, searchToolID), nil) + if err != nil { + return fmt.Errorf("error reading search tool: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + return fmt.Errorf("search tool '%s' not found", searchToolID) + } + + if err := handleResponse(resp, "reading search tool"); err != nil { + return err + } + + var searchToolResp searchToolAPIResponse + if err := json.NewDecoder(resp.Body).Decode(&searchToolResp); err != nil { + return fmt.Errorf("error decoding search tool info response: %w", err) + } + + // litellm_params is intentionally never exposed: it may hold provider API keys. + d.SetId(searchToolResp.SearchToolID) + d.Set("search_tool_name", searchToolResp.SearchToolName) + if searchToolResp.SearchToolInfo != nil { + infoJSON, err := json.Marshal(searchToolResp.SearchToolInfo) + if err != nil { + return fmt.Errorf("error encoding search_tool_info: %w", err) + } + d.Set("search_tool_info", string(infoJSON)) + } + d.Set("created_at", searchToolResp.CreatedAt) + d.Set("updated_at", searchToolResp.UpdatedAt) + + return nil +} + +func dataSourceLiteLLMSearchTools() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMSearchToolsRead, + + Schema: map[string]*schema.Schema{ + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "search_tools": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "search_tool_id": { + Type: schema.TypeString, + Computed: true, + }, + "search_tool_name": { + Type: schema.TypeString, + Computed: true, + }, + "search_tool_info": { + Type: schema.TypeString, + Computed: true, + Description: "Additional metadata as a JSON object string.", + }, + "is_from_config": { + Type: schema.TypeBool, + Computed: true, + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + }, + }, + }, + }, + } +} + +func dataSourceLiteLLMSearchToolsRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + resp, err := MakeRequest(client, "GET", endpointSearchToolsList, nil) + if err != nil { + return fmt.Errorf("error listing search tools: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing search tools"); err != nil { + return err + } + + var listResp struct { + SearchTools []searchToolAPIResponse `json:"search_tools"` + } + if err := json.NewDecoder(resp.Body).Decode(&listResp); err != nil { + return fmt.Errorf("error decoding search tools list response: %w", err) + } + + ids := make([]string, 0, len(listResp.SearchTools)) + searchTools := make([]map[string]interface{}, 0, len(listResp.SearchTools)) + for _, searchToolResp := range listResp.SearchTools { + ids = append(ids, searchToolResp.SearchToolID) + + searchTool := map[string]interface{}{ + "search_tool_id": searchToolResp.SearchToolID, + "search_tool_name": searchToolResp.SearchToolName, + "created_at": searchToolResp.CreatedAt, + "updated_at": searchToolResp.UpdatedAt, + } + if searchToolResp.SearchToolInfo != nil { + infoJSON, err := json.Marshal(searchToolResp.SearchToolInfo) + if err != nil { + return fmt.Errorf("error encoding search_tool_info: %w", err) + } + searchTool["search_tool_info"] = string(infoJSON) + } + if searchToolResp.IsFromConfig != nil { + searchTool["is_from_config"] = *searchToolResp.IsFromConfig + } + searchTools = append(searchTools, searchTool) + } + + d.SetId(strconv.FormatInt(time.Now().UnixNano(), 10)) + d.Set("ids", ids) + d.Set("search_tools", searchTools) + + return nil +} diff --git a/terraform/provider/litellm/data_source_search_tool_test.go b/terraform/provider/litellm/data_source_search_tool_test.go new file mode 100644 index 00000000000..03dc692695b --- /dev/null +++ b/terraform/provider/litellm/data_source_search_tool_test.go @@ -0,0 +1,95 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestDataSourceLiteLLMSearchToolRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/search_tools/st-123" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write(searchToolReadResponseBody()) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMSearchTool().Schema, map[string]interface{}{ + "search_tool_id": "st-123", + }) + + if err := dataSourceLiteLLMSearchToolRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "st-123" { + t.Fatalf("expected ID 'st-123', got %q", d.Id()) + } + if d.Get("search_tool_name").(string) != "my-search" { + t.Errorf("expected search_tool_name 'my-search', got %q", d.Get("search_tool_name").(string)) + } + var info map[string]interface{} + if err := json.Unmarshal([]byte(d.Get("search_tool_info").(string)), &info); err != nil { + t.Fatalf("search_tool_info not populated as JSON: %v", err) + } + if info["description"] != "Tavily search" { + t.Errorf("expected description 'Tavily search', got %v", info["description"]) + } +} + +func TestDataSourceLiteLLMSearchToolsRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/search_tools/list" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + body, _ := json.Marshal(map[string]interface{}{ + "search_tools": []map[string]interface{}{ + { + "search_tool_id": "st-1", + "search_tool_name": "first", + "search_tool_info": map[string]interface{}{"description": "first tool"}, + "is_from_config": true, + }, + {"search_tool_id": "st-2", "search_tool_name": "second"}, + }, + }) + w.Write(body) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMSearchTools().Schema, map[string]interface{}{}) + + if err := dataSourceLiteLLMSearchToolsRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + + ids := d.Get("ids").([]interface{}) + if len(ids) != 2 || ids[0] != "st-1" || ids[1] != "st-2" { + t.Fatalf("expected ids [st-1 st-2], got %v", ids) + } + searchTools := d.Get("search_tools").([]interface{}) + if len(searchTools) != 2 { + t.Fatalf("expected 2 search tools, got %d", len(searchTools)) + } + first := searchTools[0].(map[string]interface{}) + if first["search_tool_name"] != "first" || first["is_from_config"] != true { + t.Errorf("unexpected first search tool entry: %v", first) + } + var info map[string]interface{} + if err := json.Unmarshal([]byte(first["search_tool_info"].(string)), &info); err != nil { + t.Fatalf("search_tool_info not JSON-encoded in list: %v", err) + } + if info["description"] != "first tool" { + t.Errorf("expected description 'first tool', got %v", info["description"]) + } + if d.Id() == "" { + t.Fatal("expected data source ID to be set") + } +} diff --git a/terraform/provider/litellm/data_source_tag.go b/terraform/provider/litellm/data_source_tag.go new file mode 100644 index 00000000000..55af2ac56f2 --- /dev/null +++ b/terraform/provider/litellm/data_source_tag.go @@ -0,0 +1,246 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "net/url" + "strings" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const endpointTagList = "/tag/list" + +func dataSourceLiteLLMTag() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMTagRead, + + Schema: map[string]*schema.Schema{ + "name": { + Type: schema.TypeString, + Required: true, + Description: "Name of the tag to retrieve", + }, + "description": { + Type: schema.TypeString, + Computed: true, + Description: "Description of the tag", + }, + "models": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Model IDs this tag applies to", + }, + "budget_id": { + Type: schema.TypeString, + Computed: true, + Description: "Budget ID associated with this tag", + }, + "max_budget": { + Type: schema.TypeFloat, + Computed: true, + Description: "Max budget in USD for this tag", + }, + "soft_budget": { + Type: schema.TypeFloat, + Computed: true, + Description: "Soft budget in USD for this tag", + }, + "max_parallel_requests": { + Type: schema.TypeInt, + Computed: true, + Description: "Max concurrent requests allowed for this tag", + }, + "tpm_limit": { + Type: schema.TypeInt, + Computed: true, + Description: "Max tokens per minute for this tag", + }, + "rpm_limit": { + Type: schema.TypeInt, + Computed: true, + Description: "Max requests per minute for this tag", + }, + "budget_duration": { + Type: schema.TypeString, + Computed: true, + Description: "Duration for budget reset", + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + Description: "Timestamp when the tag was created", + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + Description: "Timestamp when the tag was last updated", + }, + "created_by": { + Type: schema.TypeString, + Computed: true, + Description: "User that created the tag", + }, + }, + } +} + +func dataSourceLiteLLMTagRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + name := d.Get("name").(string) + + entry, gone, err := fetchTagInfo(client, name) + if err != nil { + return fmt.Errorf("failed to read tag: %w", err) + } + if gone { + return fmt.Errorf("tag '%s' not found", name) + } + + d.SetId(name) + d.Set("description", entry.Description) + d.Set("models", entry.Models) + d.Set("created_at", entry.CreatedAt) + d.Set("updated_at", entry.UpdatedAt) + d.Set("created_by", entry.CreatedBy) + + if bt := entry.LitellmBudgetTable; bt != nil { + d.Set("budget_id", bt.BudgetID) + if bt.MaxBudget != nil { + d.Set("max_budget", *bt.MaxBudget) + } + if bt.SoftBudget != nil { + d.Set("soft_budget", *bt.SoftBudget) + } + if bt.MaxParallelRequests != nil { + d.Set("max_parallel_requests", *bt.MaxParallelRequests) + } + if bt.TPMLimit != nil { + d.Set("tpm_limit", *bt.TPMLimit) + } + if bt.RPMLimit != nil { + d.Set("rpm_limit", *bt.RPMLimit) + } + d.Set("budget_duration", bt.BudgetDuration) + } + + return nil +} + +func dataSourceLiteLLMTags() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMTagsRead, + + Schema: map[string]*schema.Schema{ + "start_date": { + Type: schema.TypeString, + Optional: true, + Description: "Optional start date (YYYY-MM-DD) limiting dynamic tags to those active in the window", + }, + "end_date": { + Type: schema.TypeString, + Optional: true, + Description: "Optional end date (YYYY-MM-DD), must be given with start_date", + }, + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Names of all tags (tag names are their IDs)", + }, + "tags": { + Type: schema.TypeList, + Computed: true, + Description: "List of tags", + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "name": {Type: schema.TypeString, Computed: true}, + "description": {Type: schema.TypeString, Computed: true}, + "models": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "budget_id": {Type: schema.TypeString, Computed: true}, + "max_budget": {Type: schema.TypeFloat, Computed: true}, + "soft_budget": {Type: schema.TypeFloat, Computed: true}, + "max_parallel_requests": {Type: schema.TypeInt, Computed: true}, + "tpm_limit": {Type: schema.TypeInt, Computed: true}, + "rpm_limit": {Type: schema.TypeInt, Computed: true}, + "budget_duration": {Type: schema.TypeString, Computed: true}, + "created_at": {Type: schema.TypeString, Computed: true}, + "updated_at": {Type: schema.TypeString, Computed: true}, + "created_by": {Type: schema.TypeString, Computed: true}, + }, + }, + }, + }, + } +} + +func dataSourceLiteLLMTagsRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + endpoint := endpointTagList + if startDate, ok := d.GetOk("start_date"); ok { + endpoint = fmt.Sprintf("%s?start_date=%s&end_date=%s", endpointTagList, + url.QueryEscape(startDate.(string)), url.QueryEscape(d.Get("end_date").(string))) + } + + resp, err := MakeRequest(client, "GET", endpoint, nil) + if err != nil { + return fmt.Errorf("failed to list tags: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing tags"); err != nil { + return err + } + + var entries []tagInfoEntry + if err := json.NewDecoder(resp.Body).Decode(&entries); err != nil { + return fmt.Errorf("error decoding tag list response: %w", err) + } + + ids := make([]string, 0, len(entries)) + tags := make([]map[string]interface{}, 0, len(entries)) + for _, entry := range entries { + ids = append(ids, entry.Name) + + tag := map[string]interface{}{ + "name": entry.Name, + "description": entry.Description, + "models": entry.Models, + "created_at": entry.CreatedAt, + "updated_at": entry.UpdatedAt, + "created_by": entry.CreatedBy, + } + if bt := entry.LitellmBudgetTable; bt != nil { + tag["budget_id"] = bt.BudgetID + tag["budget_duration"] = bt.BudgetDuration + if bt.MaxBudget != nil { + tag["max_budget"] = *bt.MaxBudget + } + if bt.SoftBudget != nil { + tag["soft_budget"] = *bt.SoftBudget + } + if bt.MaxParallelRequests != nil { + tag["max_parallel_requests"] = *bt.MaxParallelRequests + } + if bt.TPMLimit != nil { + tag["tpm_limit"] = *bt.TPMLimit + } + if bt.RPMLimit != nil { + tag["rpm_limit"] = *bt.RPMLimit + } + } + tags = append(tags, tag) + } + + d.SetId(strings.Join([]string{"litellm-tags", d.Get("start_date").(string), d.Get("end_date").(string)}, "-")) + d.Set("ids", ids) + d.Set("tags", tags) + + return nil +} diff --git a/terraform/provider/litellm/data_source_tag_test.go b/terraform/provider/litellm/data_source_tag_test.go new file mode 100644 index 00000000000..4d279bcd562 --- /dev/null +++ b/terraform/provider/litellm/data_source_tag_test.go @@ -0,0 +1,116 @@ +package litellm + +import ( + "net/http" + "net/http/httptest" + "reflect" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestDataSourceLiteLLMTagRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/tag/info" || r.Method != http.MethodPost { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Write([]byte(tagInfoBody("prod"))) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMTag().Schema, map[string]interface{}{"name": "prod"}) + + if err := dataSourceLiteLLMTagRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Id() != "prod" { + t.Fatalf("expected ID 'prod', got %q", d.Id()) + } + checks := map[string]interface{}{ + "description": "Production traffic", + "budget_id": "bud-1", + "max_budget": 50.5, + "tpm_limit": 1000, + "created_at": "2026-01-01T00:00:00", + "created_by": "admin", + } + for key, want := range checks { + if got := d.Get(key); got != want { + t.Errorf("expected %s %v, got %v", key, want, got) + } + } +} + +func TestDataSourceLiteLLMTagRead_NotFound(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMTag().Schema, map[string]interface{}{"name": "gone"}) + + if err := dataSourceLiteLLMTagRead(d, NewClient(srv.URL, "test-key", true)); err == nil { + t.Fatal("expected error for missing tag, got nil") + } +} + +func TestDataSourceLiteLLMTagsRead(t *testing.T) { + var gotQuery string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/tag/list" || r.Method != http.MethodGet { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + gotQuery = r.URL.RawQuery + w.Write([]byte(`[ + { + "name": "prod", + "description": "Production traffic", + "models": ["model-1"], + "created_at": "2026-01-01T00:00:00", + "updated_at": "2026-01-02T00:00:00", + "created_by": "admin", + "litellm_budget_table": {"budget_id": "bud-1", "max_budget": 50.5} + }, + { + "name": "dynamic-tag", + "description": "This is just a spend tag that was passed dynamically in a request.", + "models": null, + "created_at": "2026-02-01T00:00:00", + "updated_at": "2026-02-02T00:00:00" + } + ]`)) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMTags().Schema, map[string]interface{}{ + "start_date": "2026-01-01", + "end_date": "2026-03-01", + }) + + if err := dataSourceLiteLLMTagsRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + if gotQuery != "start_date=2026-01-01&end_date=2026-03-01" { + t.Errorf("expected date filter query params, got %q", gotQuery) + } + if !reflect.DeepEqual(d.Get("ids"), []interface{}{"prod", "dynamic-tag"}) { + t.Errorf("expected ids ['prod', 'dynamic-tag'], got %v", d.Get("ids")) + } + if got := d.Get("tags.#").(int); got != 2 { + t.Fatalf("expected 2 tags, got %d", got) + } + if got := d.Get("tags.0.name").(string); got != "prod" { + t.Errorf("expected tags.0.name 'prod', got %q", got) + } + if got := d.Get("tags.0.max_budget").(float64); got != 50.5 { + t.Errorf("expected tags.0.max_budget 50.5, got %v", got) + } + if got := d.Get("tags.1.name").(string); got != "dynamic-tag" { + t.Errorf("expected tags.1.name 'dynamic-tag', got %q", got) + } + if got := d.Get("tags.1.budget_id").(string); got != "" { + t.Errorf("expected empty budget_id for dynamic tag, got %q", got) + } +} diff --git a/terraform/provider/litellm/data_source_team.go b/terraform/provider/litellm/data_source_team.go new file mode 100644 index 00000000000..a484c10246a --- /dev/null +++ b/terraform/provider/litellm/data_source_team.go @@ -0,0 +1,294 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "log" + "net/url" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const endpointTeamList = "/team/list" + +type teamDetail struct { + TeamID string `json:"team_id"` + TeamAlias string `json:"team_alias"` + OrganizationID string `json:"organization_id"` + Models []string `json:"models"` + Metadata map[string]interface{} `json:"metadata"` + TPMLimit *int `json:"tpm_limit"` + RPMLimit *int `json:"rpm_limit"` + MaxParallelRequests *int `json:"max_parallel_requests"` + MaxBudget *float64 `json:"max_budget"` + SoftBudget *float64 `json:"soft_budget"` + Spend *float64 `json:"spend"` + BudgetDuration string `json:"budget_duration"` + Blocked bool `json:"blocked"` + TeamMemberPermissions []string `json:"team_member_permissions"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` +} + +type teamInfoEnvelope struct { + TeamID string `json:"team_id"` + TeamInfo teamDetail `json:"team_info"` +} + +func dataSourceLiteLLMTeam() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMTeamRead, + + Schema: map[string]*schema.Schema{ + "team_id": { + Type: schema.TypeString, + Required: true, + Description: "Unique identifier of the team to retrieve", + }, + "team_alias": { + Type: schema.TypeString, + Computed: true, + }, + "organization_id": { + Type: schema.TypeString, + Computed: true, + }, + "models": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "metadata": { + Type: schema.TypeMap, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "tags": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "soft_budget_alerting_emails": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "tpm_limit": { + Type: schema.TypeInt, + Computed: true, + }, + "rpm_limit": { + Type: schema.TypeInt, + Computed: true, + }, + "max_parallel_requests": { + Type: schema.TypeInt, + Computed: true, + }, + "max_budget": { + Type: schema.TypeFloat, + Computed: true, + }, + "soft_budget": { + Type: schema.TypeFloat, + Computed: true, + }, + "spend": { + Type: schema.TypeFloat, + Computed: true, + }, + "budget_duration": { + Type: schema.TypeString, + Computed: true, + }, + "blocked": { + Type: schema.TypeBool, + Computed: true, + }, + "team_member_permissions": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + }, + } +} + +func dataSourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + teamID := d.Get("team_id").(string) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("%s?team_id=%s", endpointTeamInfo, url.QueryEscape(teamID)), nil) + if err != nil { + return fmt.Errorf("failed to read team: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "reading team info"); err != nil { + return err + } + + var envelope teamInfoEnvelope + if err := json.NewDecoder(resp.Body).Decode(&envelope); err != nil { + return fmt.Errorf("failed to decode team info response: %w", err) + } + team := envelope.TeamInfo + + d.SetId(teamID) + d.Set("team_alias", team.TeamAlias) + d.Set("organization_id", team.OrganizationID) + d.Set("models", team.Models) + + metadata, tags, alertEmails := splitTeamMetadata(team.Metadata) + d.Set("metadata", metadata) + d.Set("tags", tags) + d.Set("soft_budget_alerting_emails", alertEmails) + + if team.TPMLimit != nil { + d.Set("tpm_limit", *team.TPMLimit) + } + if team.RPMLimit != nil { + d.Set("rpm_limit", *team.RPMLimit) + } + if team.MaxParallelRequests != nil { + d.Set("max_parallel_requests", *team.MaxParallelRequests) + } + if team.MaxBudget != nil { + d.Set("max_budget", *team.MaxBudget) + } + if team.SoftBudget != nil { + d.Set("soft_budget", *team.SoftBudget) + } + if team.Spend != nil { + d.Set("spend", *team.Spend) + } + d.Set("budget_duration", team.BudgetDuration) + d.Set("blocked", team.Blocked) + d.Set("team_member_permissions", team.TeamMemberPermissions) + d.Set("created_at", team.CreatedAt) + d.Set("updated_at", team.UpdatedAt) + + log.Printf("[INFO] Successfully read team with ID: %s", teamID) + return nil +} + +func dataSourceLiteLLMTeams() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMTeamsRead, + + Schema: map[string]*schema.Schema{ + "user_id": { + Type: schema.TypeString, + Optional: true, + Description: "Only return teams this user belongs to", + }, + "organization_id": { + Type: schema.TypeString, + Optional: true, + Description: "Only return teams in this organization", + }, + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "IDs of the returned teams", + }, + "teams": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "team_id": {Type: schema.TypeString, Computed: true}, + "team_alias": {Type: schema.TypeString, Computed: true}, + "organization_id": {Type: schema.TypeString, Computed: true}, + "models": {Type: schema.TypeList, Computed: true, Elem: &schema.Schema{Type: schema.TypeString}}, + "spend": {Type: schema.TypeFloat, Computed: true}, + "max_budget": {Type: schema.TypeFloat, Computed: true}, + "tpm_limit": {Type: schema.TypeInt, Computed: true}, + "rpm_limit": {Type: schema.TypeInt, Computed: true}, + "budget_duration": {Type: schema.TypeString, Computed: true}, + "blocked": {Type: schema.TypeBool, Computed: true}, + "created_at": {Type: schema.TypeString, Computed: true}, + "updated_at": {Type: schema.TypeString, Computed: true}, + }, + }, + }, + }, + } +} + +func dataSourceLiteLLMTeamsRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + query := url.Values{} + if v, ok := d.GetOk("user_id"); ok { + query.Set("user_id", v.(string)) + } + if v, ok := d.GetOk("organization_id"); ok { + query.Set("organization_id", v.(string)) + } + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("%s?%s", endpointTeamList, query.Encode()), nil) + if err != nil { + return fmt.Errorf("failed to list teams: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing teams"); err != nil { + return err + } + + var teamList []teamDetail + if err := json.NewDecoder(resp.Body).Decode(&teamList); err != nil { + return fmt.Errorf("failed to decode team list response: %w", err) + } + + ids := make([]string, 0, len(teamList)) + teams := make([]map[string]interface{}, 0, len(teamList)) + for _, team := range teamList { + ids = append(ids, team.TeamID) + teams = append(teams, map[string]interface{}{ + "team_id": team.TeamID, + "team_alias": team.TeamAlias, + "organization_id": team.OrganizationID, + "models": team.Models, + "spend": teamDerefFloat(team.Spend), + "max_budget": teamDerefFloat(team.MaxBudget), + "tpm_limit": teamDerefInt(team.TPMLimit), + "rpm_limit": teamDerefInt(team.RPMLimit), + "budget_duration": team.BudgetDuration, + "blocked": team.Blocked, + "created_at": team.CreatedAt, + "updated_at": team.UpdatedAt, + }) + } + + d.SetId(GetStringValue(query.Encode(), "all")) + d.Set("ids", ids) + d.Set("teams", teams) + + log.Printf("[INFO] Successfully listed %d teams", len(teams)) + return nil +} + +func teamDerefFloat(v *float64) float64 { + if v == nil { + return 0 + } + return *v +} + +func teamDerefInt(v *int) int { + if v == nil { + return 0 + } + return *v +} diff --git a/terraform/provider/litellm/data_source_team_test.go b/terraform/provider/litellm/data_source_team_test.go new file mode 100644 index 00000000000..e40f5d95a6f --- /dev/null +++ b/terraform/provider/litellm/data_source_team_test.go @@ -0,0 +1,145 @@ +package litellm + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestDataSourceTeamRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/team/info" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + if got := r.URL.Query().Get("team_id"); got != "team-123" { + t.Errorf("expected team_id 'team-123', got %q", got) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "team_id": "team-123", + "team_info": { + "team_id": "team-123", + "team_alias": "ml-team", + "organization_id": "org-1", + "models": ["gpt-4o"], + "metadata": {"env": "prod", "tags": ["ml"], "soft_budget_alerting_emails": ["ops@example.com"]}, + "tpm_limit": 5000, + "rpm_limit": 100, + "max_budget": 250.5, + "soft_budget": 200, + "spend": 42.25, + "budget_duration": "30d", + "blocked": true, + "team_member_permissions": ["/key/generate"], + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-02-01T00:00:00Z" + } + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMTeam().Schema, map[string]interface{}{ + "team_id": "team-123", + }) + + if err := dataSourceLiteLLMTeamRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Id() != "team-123" { + t.Fatalf("expected ID 'team-123', got %q", d.Id()) + } + checks := map[string]interface{}{ + "team_alias": "ml-team", + "organization_id": "org-1", + "tpm_limit": 5000, + "rpm_limit": 100, + "max_budget": 250.5, + "soft_budget": 200.0, + "spend": 42.25, + "budget_duration": "30d", + "blocked": true, + } + for attr, want := range checks { + if got := d.Get(attr); got != want { + t.Errorf("attr %s: expected %v, got %v", attr, want, got) + } + } + tags := d.Get("tags").([]interface{}) + if len(tags) != 1 || tags[0] != "ml" { + t.Errorf("unexpected tags: %v", tags) + } + emails := d.Get("soft_budget_alerting_emails").([]interface{}) + if len(emails) != 1 || emails[0] != "ops@example.com" { + t.Errorf("unexpected alerting emails: %v", emails) + } + metadata := d.Get("metadata").(map[string]interface{}) + if metadata["env"] != "prod" || len(metadata) != 1 { + t.Errorf("unexpected metadata: %v", metadata) + } + perms := d.Get("team_member_permissions").([]interface{}) + if len(perms) != 1 || perms[0] != "/key/generate" { + t.Errorf("unexpected permissions: %v", perms) + } +} + +func TestDataSourceTeamsRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/team/list" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + if got := r.URL.Query().Get("organization_id"); got != "org-1" { + t.Errorf("expected organization_id 'org-1', got %q", got) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`[ + {"team_id": "team-1", "team_alias": "alpha", "organization_id": "org-1", "spend": 5, "max_budget": 50, "tpm_limit": 100, "rpm_limit": 10, "models": ["m1"], "blocked": false}, + {"team_id": "team-2", "team_alias": "beta", "organization_id": "org-1", "blocked": true} + ]`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMTeams().Schema, map[string]interface{}{ + "organization_id": "org-1", + }) + + if err := dataSourceLiteLLMTeamsRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + + ids := d.Get("ids").([]interface{}) + if len(ids) != 2 || ids[0] != "team-1" || ids[1] != "team-2" { + t.Errorf("unexpected ids: %v", ids) + } + teams := d.Get("teams").([]interface{}) + if len(teams) != 2 { + t.Fatalf("expected 2 teams, got %d", len(teams)) + } + first := teams[0].(map[string]interface{}) + if first["team_alias"] != "alpha" || first["max_budget"] != 50.0 || first["tpm_limit"] != 100 { + t.Errorf("unexpected first team: %v", first) + } + second := teams[1].(map[string]interface{}) + if second["blocked"] != true || second["max_budget"] != 0.0 { + t.Errorf("unexpected second team: %v", second) + } +} + +func TestDataSourceTeamsReadError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"error": "boom"}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMTeams().Schema, map[string]interface{}{}) + + if err := dataSourceLiteLLMTeamsRead(d, client); err == nil { + t.Fatal("expected error on server failure, got nil") + } +} diff --git a/terraform/provider/litellm/data_source_unified_access_group.go b/terraform/provider/litellm/data_source_unified_access_group.go new file mode 100644 index 00000000000..0153fa380ce --- /dev/null +++ b/terraform/provider/litellm/data_source_unified_access_group.go @@ -0,0 +1,189 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "net/http" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const endpointUnifiedAccessGroupList = "/v1/unified_access_group" + +func unifiedAccessGroupComputedSchema() map[string]*schema.Schema { + return map[string]*schema.Schema{ + "access_group_name": { + Type: schema.TypeString, + Computed: true, + }, + "description": { + Type: schema.TypeString, + Computed: true, + }, + "access_model_names": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "access_mcp_server_ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "access_agent_ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "assigned_team_ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "assigned_key_ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "created_by": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_by": { + Type: schema.TypeString, + Computed: true, + }, + } +} + +func dataSourceLiteLLMUnifiedAccessGroup() *schema.Resource { + dsSchema := unifiedAccessGroupComputedSchema() + dsSchema["access_group_id"] = &schema.Schema{ + Type: schema.TypeString, + Required: true, + Description: "ID of the unified access group to retrieve", + } + + return &schema.Resource{ + Read: dataSourceLiteLLMUnifiedAccessGroupRead, + Schema: dsSchema, + } +} + +func dataSourceLiteLLMUnifiedAccessGroupRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + groupID := d.Get("access_group_id").(string) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("/v1/unified_access_group/%s", groupID), nil) + if err != nil { + return fmt.Errorf("error reading unified access group: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + return fmt.Errorf("unified access group '%s' not found", groupID) + } + + if err := handleResponse(resp, "reading unified access group"); err != nil { + return err + } + + var group unifiedAccessGroupResponse + if err := json.NewDecoder(resp.Body).Decode(&group); err != nil { + return fmt.Errorf("error decoding unified access group info response: %w", err) + } + + d.SetId(GetStringValue(group.AccessGroupID, groupID)) + setUnifiedAccessGroupFields(d, group) + + return nil +} + +func dataSourceLiteLLMUnifiedAccessGroups() *schema.Resource { + itemSchema := unifiedAccessGroupComputedSchema() + itemSchema["access_group_id"] = &schema.Schema{ + Type: schema.TypeString, + Computed: true, + } + + return &schema.Resource{ + Read: dataSourceLiteLLMUnifiedAccessGroupsRead, + + Schema: map[string]*schema.Schema{ + "access_groups": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Resource{Schema: itemSchema}, + }, + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + }, + } +} + +func dataSourceLiteLLMUnifiedAccessGroupsRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + resp, err := MakeRequest(client, "GET", endpointUnifiedAccessGroupList, nil) + if err != nil { + return fmt.Errorf("error listing unified access groups: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing unified access groups"); err != nil { + return err + } + + var groups []unifiedAccessGroupResponse + if err := json.NewDecoder(resp.Body).Decode(&groups); err != nil { + return fmt.Errorf("error decoding unified access group list response: %w", err) + } + + items := make([]map[string]interface{}, 0, len(groups)) + ids := make([]string, 0, len(groups)) + for _, group := range groups { + items = append(items, unifiedAccessGroupFlatten(group)) + ids = append(ids, group.AccessGroupID) + } + + d.SetId("unified_access_groups") + d.Set("access_groups", items) + d.Set("ids", ids) + + return nil +} + +func unifiedAccessGroupFlatten(group unifiedAccessGroupResponse) map[string]interface{} { + item := map[string]interface{}{ + "access_group_id": group.AccessGroupID, + "access_group_name": group.AccessGroupName, + "access_model_names": group.AccessModelNames, + "access_mcp_server_ids": group.AccessMCPServerIDs, + "access_agent_ids": group.AccessAgentIDs, + "assigned_team_ids": group.AssignedTeamIDs, + "assigned_key_ids": group.AssignedKeyIDs, + "created_at": group.CreatedAt, + "updated_at": group.UpdatedAt, + } + if group.Description != nil { + item["description"] = *group.Description + } + if group.CreatedBy != nil { + item["created_by"] = *group.CreatedBy + } + if group.UpdatedBy != nil { + item["updated_by"] = *group.UpdatedBy + } + return item +} diff --git a/terraform/provider/litellm/data_source_unified_access_group_test.go b/terraform/provider/litellm/data_source_unified_access_group_test.go new file mode 100644 index 00000000000..f1567be36af --- /dev/null +++ b/terraform/provider/litellm/data_source_unified_access_group_test.go @@ -0,0 +1,112 @@ +package litellm + +import ( + "net/http" + "net/http/httptest" + "reflect" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestUnifiedAccessGroupDataSourceRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" || r.URL.Path != "/v1/unified_access_group/uag-123" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + return + } + w.Write(unifiedAccessGroupJSON("uag-123")) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMUnifiedAccessGroup().Schema, map[string]interface{}{ + "access_group_id": "uag-123", + }) + + if err := dataSourceLiteLLMUnifiedAccessGroupRead(d, client); err != nil { + t.Fatalf("data source read failed: %v", err) + } + + if d.Id() != "uag-123" { + t.Fatalf("expected ID 'uag-123', got %q", d.Id()) + } + if d.Get("access_group_name").(string) != "prod-group" { + t.Fatalf("expected access_group_name 'prod-group', got %v", d.Get("access_group_name")) + } + if d.Get("description").(string) != "prod access" { + t.Fatalf("expected description 'prod access', got %v", d.Get("description")) + } + if !reflect.DeepEqual(d.Get("access_model_names"), []interface{}{"gpt-4"}) { + t.Fatalf("expected access_model_names [gpt-4], got %v", d.Get("access_model_names")) + } + if !reflect.DeepEqual(d.Get("assigned_team_ids"), []interface{}{"team-1"}) { + t.Fatalf("expected assigned_team_ids [team-1], got %v", d.Get("assigned_team_ids")) + } +} + +func TestUnifiedAccessGroupDataSourceReadNotFound(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMUnifiedAccessGroup().Schema, map[string]interface{}{ + "access_group_id": "missing", + }) + + if err := dataSourceLiteLLMUnifiedAccessGroupRead(d, client); err == nil { + t.Fatal("expected error for missing unified access group, got nil") + } +} + +func TestUnifiedAccessGroupsDataSourceRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" || r.URL.Path != "/v1/unified_access_group" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + return + } + w.Write([]byte(`[` + + `{"access_group_id": "uag-1", "access_group_name": "group-one", "description": "first",` + + ` "access_model_names": ["gpt-4"], "access_mcp_server_ids": [], "access_agent_ids": [],` + + ` "assigned_team_ids": ["team-1"], "assigned_key_ids": [],` + + ` "created_at": "2026-01-01T00:00:00Z", "updated_at": "2026-01-02T00:00:00Z"},` + + `{"access_group_id": "uag-2", "access_group_name": "group-two",` + + ` "access_model_names": [], "access_mcp_server_ids": ["mcp-1"], "access_agent_ids": [],` + + ` "assigned_team_ids": [], "assigned_key_ids": [],` + + ` "created_at": "2026-01-03T00:00:00Z", "updated_at": "2026-01-04T00:00:00Z"}]`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMUnifiedAccessGroups().Schema, map[string]interface{}{}) + + if err := dataSourceLiteLLMUnifiedAccessGroupsRead(d, client); err != nil { + t.Fatalf("data source read failed: %v", err) + } + + groups := d.Get("access_groups").([]interface{}) + if len(groups) != 2 { + t.Fatalf("expected 2 unified access groups, got %d", len(groups)) + } + first := groups[0].(map[string]interface{}) + if first["access_group_id"] != "uag-1" { + t.Fatalf("expected first access_group_id 'uag-1', got %v", first["access_group_id"]) + } + if first["access_group_name"] != "group-one" { + t.Fatalf("expected first access_group_name 'group-one', got %v", first["access_group_name"]) + } + if first["description"] != "first" { + t.Fatalf("expected first description 'first', got %v", first["description"]) + } + second := groups[1].(map[string]interface{}) + if !reflect.DeepEqual(second["access_mcp_server_ids"], []interface{}{"mcp-1"}) { + t.Fatalf("expected second access_mcp_server_ids [mcp-1], got %v", second["access_mcp_server_ids"]) + } + if !reflect.DeepEqual(d.Get("ids"), []interface{}{"uag-1", "uag-2"}) { + t.Fatalf("expected ids [uag-1 uag-2], got %v", d.Get("ids")) + } +} diff --git a/terraform/provider/litellm/data_source_user.go b/terraform/provider/litellm/data_source_user.go new file mode 100644 index 00000000000..460415b37c6 --- /dev/null +++ b/terraform/provider/litellm/data_source_user.go @@ -0,0 +1,307 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "net/http" + "net/url" + "strconv" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const endpointUserList = "/user/list" + +func dataSourceLiteLLMUser() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMUserRead, + + Schema: map[string]*schema.Schema{ + "user_id": { + Type: schema.TypeString, + Required: true, + Description: "ID of the user to retrieve", + }, + "user_email": { + Type: schema.TypeString, + Computed: true, + Description: "Email address of the user", + }, + "user_alias": { + Type: schema.TypeString, + Computed: true, + Description: "Descriptive name for the user", + }, + "user_role": { + Type: schema.TypeString, + Computed: true, + Description: "Role of the user on the proxy", + }, + "teams": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "List of team IDs the user belongs to", + }, + "models": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Models the user is allowed to call", + }, + "max_budget": { + Type: schema.TypeFloat, + Computed: true, + Description: "Maximum budget in USD for the user", + }, + "spend": { + Type: schema.TypeFloat, + Computed: true, + Description: "Current spend in USD for the user", + }, + "budget_duration": { + Type: schema.TypeString, + Computed: true, + Description: "Budget reset period for the user", + }, + "tpm_limit": { + Type: schema.TypeInt, + Computed: true, + Description: "Tokens per minute limit for the user", + }, + "rpm_limit": { + Type: schema.TypeInt, + Computed: true, + Description: "Requests per minute limit for the user", + }, + "max_parallel_requests": { + Type: schema.TypeInt, + Computed: true, + Description: "Maximum number of parallel requests for the user", + }, + "metadata": { + Type: schema.TypeMap, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Metadata for the user", + }, + "model_max_budget": { + Type: schema.TypeString, + Computed: true, + Description: "JSON string of per-model budget config", + }, + }, + } +} + +func dataSourceLiteLLMUserRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + userID := d.Get("user_id").(string) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("%s?user_id=%s", endpointUserInfo, url.QueryEscape(userID)), nil) + if err != nil { + return fmt.Errorf("failed to read user: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + return fmt.Errorf("user '%s' not found", userID) + } + + if err := handleResponse(resp, "reading user"); err != nil { + return err + } + + var infoResp userInfoResponse + if err := json.NewDecoder(resp.Body).Decode(&infoResp); err != nil { + return fmt.Errorf("error decoding user info response: %w", err) + } + if infoResp.UserInfo == nil { + return fmt.Errorf("user '%s' not found", userID) + } + + d.SetId(userID) + setUserStateFromInfo(d, infoResp.UserInfo) + if v, ok := infoResp.UserInfo["spend"].(float64); ok { + d.Set("spend", v) + } + + return nil +} + +func dataSourceLiteLLMUsers() *schema.Resource { + return &schema.Resource{ + Read: dataSourceLiteLLMUsersRead, + + Schema: map[string]*schema.Schema{ + "role": { + Type: schema.TypeString, + Optional: true, + Description: "Filter users by role", + }, + "user_ids": { + Type: schema.TypeString, + Optional: true, + Description: "Comma-separated list of user IDs to filter by", + }, + "user_email": { + Type: schema.TypeString, + Optional: true, + Description: "Filter users by partial email match", + }, + "team": { + Type: schema.TypeString, + Optional: true, + Description: "Filter users by team ID", + }, + "page": { + Type: schema.TypeInt, + Optional: true, + Default: 1, + Description: "Page number to fetch", + }, + "page_size": { + Type: schema.TypeInt, + Optional: true, + Default: 25, + Description: "Number of users per page (max 100)", + }, + "sort_by": { + Type: schema.TypeString, + Optional: true, + Description: "Column to sort by (e.g. 'user_id', 'user_email', 'created_at')", + }, + "sort_order": { + Type: schema.TypeString, + Optional: true, + Description: "Sort order, 'asc' or 'desc'", + }, + "users": { + Type: schema.TypeList, + Computed: true, + Description: "Users returned for the requested page", + Elem: &schema.Resource{ + Schema: map[string]*schema.Schema{ + "user_id": {Type: schema.TypeString, Computed: true}, + "user_email": {Type: schema.TypeString, Computed: true}, + "user_alias": {Type: schema.TypeString, Computed: true}, + "user_role": {Type: schema.TypeString, Computed: true}, + "teams": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "models": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "max_budget": {Type: schema.TypeFloat, Computed: true}, + "spend": {Type: schema.TypeFloat, Computed: true}, + "tpm_limit": {Type: schema.TypeInt, Computed: true}, + "rpm_limit": {Type: schema.TypeInt, Computed: true}, + "key_count": {Type: schema.TypeInt, Computed: true}, + "created_at": {Type: schema.TypeString, Computed: true}, + }, + }, + }, + "ids": { + Type: schema.TypeList, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "IDs of the users returned for the requested page", + }, + "total": { + Type: schema.TypeInt, + Computed: true, + Description: "Total number of users matching the filters", + }, + "total_pages": { + Type: schema.TypeInt, + Computed: true, + Description: "Total number of pages available", + }, + }, + } +} + +type userListResponse struct { + Users []map[string]interface{} `json:"users"` + Total int `json:"total"` + TotalPages int `json:"total_pages"` +} + +func userListQuery(d *schema.ResourceData) string { + query := url.Values{} + for _, key := range []string{"role", "user_ids", "user_email", "team", "sort_by", "sort_order"} { + if v, ok := d.GetOk(key); ok { + query.Set(key, v.(string)) + } + } + query.Set("page", strconv.Itoa(d.Get("page").(int))) + query.Set("page_size", strconv.Itoa(d.Get("page_size").(int))) + return query.Encode() +} + +func userListEntry(user map[string]interface{}) map[string]interface{} { + entry := map[string]interface{}{} + for _, key := range []string{"user_id", "user_email", "user_alias", "user_role", "created_at"} { + if v, ok := user[key].(string); ok { + entry[key] = v + } + } + for _, key := range []string{"max_budget", "spend"} { + if v, ok := user[key].(float64); ok { + entry[key] = v + } + } + for _, key := range []string{"tpm_limit", "rpm_limit", "key_count"} { + if v, ok := user[key].(float64); ok { + entry[key] = int(v) + } + } + for _, key := range []string{"teams", "models"} { + if v, ok := user[key].([]interface{}); ok { + entry[key] = v + } + } + return entry +} + +func dataSourceLiteLLMUsersRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + query := userListQuery(d) + resp, err := MakeRequest(client, "GET", fmt.Sprintf("%s?%s", endpointUserList, query), nil) + if err != nil { + return fmt.Errorf("failed to list users: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "listing users"); err != nil { + return err + } + + var listResp userListResponse + if err := json.NewDecoder(resp.Body).Decode(&listResp); err != nil { + return fmt.Errorf("error decoding user list response: %w", err) + } + + users := make([]map[string]interface{}, 0, len(listResp.Users)) + ids := make([]string, 0, len(listResp.Users)) + for _, user := range listResp.Users { + entry := userListEntry(user) + if id, ok := entry["user_id"].(string); ok { + ids = append(ids, id) + } + users = append(users, entry) + } + + d.SetId(fmt.Sprintf("users?%s", query)) + d.Set("users", users) + d.Set("ids", ids) + d.Set("total", listResp.Total) + d.Set("total_pages", listResp.TotalPages) + + return nil +} diff --git a/terraform/provider/litellm/data_source_user_test.go b/terraform/provider/litellm/data_source_user_test.go new file mode 100644 index 00000000000..ece532ed8fc --- /dev/null +++ b/terraform/provider/litellm/data_source_user_test.go @@ -0,0 +1,144 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func TestDataSourceUserRead_MapsFields(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/user/info" || r.Method != http.MethodGet { + t.Errorf("expected GET /user/info, got %s %s", r.Method, r.URL.Path) + } + if got := r.URL.Query().Get("user_id"); got != "u-ds" { + t.Errorf("expected user_id query 'u-ds', got %q", got) + } + w.Write(userInfoBody("u-ds", map[string]interface{}{ + "user_email": "carol@example.com", + "user_role": "internal_user", + "max_budget": 42.0, + "spend": 1.5, + "models": []interface{}{"gpt-4o"}, + "model_max_budget": map[string]interface{}{"gpt-4o": map[string]interface{}{"max_budget": 2.0}}, + })) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMUser().Schema, map[string]interface{}{ + "user_id": "u-ds", + }) + + if err := dataSourceLiteLLMUserRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Id() != "u-ds" { + t.Fatalf("expected ID 'u-ds', got %q", d.Id()) + } + if got := d.Get("user_email").(string); got != "carol@example.com" { + t.Errorf("expected user_email 'carol@example.com', got %q", got) + } + if got := d.Get("spend").(float64); got != 1.5 { + t.Errorf("expected spend 1.5, got %v", got) + } + models := d.Get("models").([]interface{}) + if len(models) != 1 || models[0] != "gpt-4o" { + t.Errorf("expected models [gpt-4o], got %v", models) + } + var mmb map[string]interface{} + if err := json.Unmarshal([]byte(d.Get("model_max_budget").(string)), &mmb); err != nil { + t.Fatalf("model_max_budget in state is not valid JSON: %v", err) + } + if _, ok := mmb["gpt-4o"]; !ok { + t.Errorf("expected gpt-4o key in model_max_budget state, got %v", mmb) + } +} + +func TestDataSourceUsersRead_FiltersAndMapsList(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/user/list" || r.Method != http.MethodGet { + t.Errorf("expected GET /user/list, got %s %s", r.Method, r.URL.Path) + } + query := r.URL.Query() + if got := query.Get("role"); got != "internal_user" { + t.Errorf("expected role query 'internal_user', got %q", got) + } + if got := query.Get("page"); got != "2" { + t.Errorf("expected page query '2', got %q", got) + } + if got := query.Get("page_size"); got != "50" { + t.Errorf("expected page_size query '50', got %q", got) + } + body, _ := json.Marshal(map[string]interface{}{ + "users": []map[string]interface{}{ + { + "user_id": "u-1", + "user_email": "one@example.com", + "user_role": "internal_user", + "max_budget": 10.0, + "spend": 2.0, + "tpm_limit": 100, + "key_count": 3, + }, + { + "user_id": "u-2", + "user_email": "two@example.com", + "teams": []string{"team-x"}, + }, + }, + "total": 52, + "page": 2, + "page_size": 50, + "total_pages": 2, + }) + w.Write(body) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, dataSourceLiteLLMUsers().Schema, map[string]interface{}{ + "role": "internal_user", + "page": 2, + "page_size": 50, + }) + + if err := dataSourceLiteLLMUsersRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + users := d.Get("users").([]interface{}) + if len(users) != 2 { + t.Fatalf("expected 2 users, got %d", len(users)) + } + first := users[0].(map[string]interface{}) + if got := first["user_id"].(string); got != "u-1" { + t.Errorf("expected first user_id 'u-1', got %q", got) + } + if got := first["spend"].(float64); got != 2.0 { + t.Errorf("expected first spend 2.0, got %v", got) + } + if got := first["tpm_limit"].(int); got != 100 { + t.Errorf("expected first tpm_limit 100, got %d", got) + } + if got := first["key_count"].(int); got != 3 { + t.Errorf("expected first key_count 3, got %d", got) + } + second := users[1].(map[string]interface{}) + teams := second["teams"].([]interface{}) + if len(teams) != 1 || teams[0] != "team-x" { + t.Errorf("expected second user teams [team-x], got %v", teams) + } + ids := d.Get("ids").([]interface{}) + if len(ids) != 2 || ids[0] != "u-1" || ids[1] != "u-2" { + t.Errorf("expected ids [u-1 u-2], got %v", ids) + } + if got := d.Get("total").(int); got != 52 { + t.Errorf("expected total 52, got %d", got) + } + if got := d.Get("total_pages").(int); got != 2 { + t.Errorf("expected total_pages 2, got %d", got) + } +} diff --git a/terraform/provider/litellm/provider.go b/terraform/provider/litellm/provider.go index 57f9cc24183..0afbbe9a464 100644 --- a/terraform/provider/litellm/provider.go +++ b/terraform/provider/litellm/provider.go @@ -19,10 +19,55 @@ func Provider() *schema.Provider { "litellm_mcp_server": resourceLiteLLMMCPServer(), "litellm_credential": resourceLiteLLMCredential(), "litellm_vector_store": resourceLiteLLMVectorStore(), + "litellm_jwt_key_mapping": resourceLiteLLMJWTKeyMapping(), + "litellm_fallback": resourceLiteLLMFallback(), + "litellm_key_block": resourceLiteLLMKeyBlock(), + "litellm_team_block": resourceLiteLLMTeamBlock(), + "litellm_access_group": resourceLiteLLMAccessGroup(), + "litellm_unified_access_group": resourceLiteLLMUnifiedAccessGroup(), + "litellm_guardrail": resourceLiteLLMGuardrail(), + "litellm_prompt": resourceLiteLLMPrompt(), + "litellm_agent": resourceLiteLLMAgent(), + "litellm_search_tool": resourceLiteLLMSearchTool(), + "litellm_user": resourceLiteLLMUser(), + "litellm_budget": resourceLiteLLMBudget(), + "litellm_tag": resourceLiteLLMTag(), + "litellm_project": resourceLiteLLMProject(), }, DataSourcesMap: map[string]*schema.Resource{ - "litellm_credential": dataSourceLiteLLMCredential(), - "litellm_vector_store": dataSourceLiteLLMVectorStore(), + "litellm_credential": dataSourceLiteLLMCredential(), + "litellm_vector_store": dataSourceLiteLLMVectorStore(), + "litellm_fallback": dataSourceLiteLLMFallback(), + "litellm_access_group": dataSourceLiteLLMAccessGroup(), + "litellm_access_groups": dataSourceLiteLLMAccessGroups(), + "litellm_unified_access_group": dataSourceLiteLLMUnifiedAccessGroup(), + "litellm_unified_access_groups": dataSourceLiteLLMUnifiedAccessGroups(), + "litellm_guardrail": dataSourceLiteLLMGuardrail(), + "litellm_guardrails": dataSourceLiteLLMGuardrails(), + "litellm_prompt": dataSourceLiteLLMPrompt(), + "litellm_prompts": dataSourceLiteLLMPrompts(), + "litellm_agent": dataSourceLiteLLMAgent(), + "litellm_agents": dataSourceLiteLLMAgents(), + "litellm_search_tool": dataSourceLiteLLMSearchTool(), + "litellm_search_tools": dataSourceLiteLLMSearchTools(), + "litellm_user": dataSourceLiteLLMUser(), + "litellm_users": dataSourceLiteLLMUsers(), + "litellm_budget": dataSourceLiteLLMBudget(), + "litellm_budgets": dataSourceLiteLLMBudgets(), + "litellm_tag": dataSourceLiteLLMTag(), + "litellm_tags": dataSourceLiteLLMTags(), + "litellm_project": dataSourceLiteLLMProject(), + "litellm_projects": dataSourceLiteLLMProjects(), + "litellm_key": dataSourceLiteLLMKey(), + "litellm_keys": dataSourceLiteLLMKeys(), + "litellm_team": dataSourceLiteLLMTeam(), + "litellm_teams": dataSourceLiteLLMTeams(), + "litellm_model": dataSourceLiteLLMModel(), + "litellm_models": dataSourceLiteLLMModels(), + "litellm_organization": dataSourceLiteLLMOrganization(), + "litellm_organizations": dataSourceLiteLLMOrganizations(), + "litellm_mcp_server": dataSourceLiteLLMMCPServer(), + "litellm_mcp_servers": dataSourceLiteLLMMCPServers(), }, Schema: map[string]*schema.Schema{ "api_base": { diff --git a/terraform/provider/litellm/resource_access_group.go b/terraform/provider/litellm/resource_access_group.go new file mode 100644 index 00000000000..d28f3dd2c31 --- /dev/null +++ b/terraform/provider/litellm/resource_access_group.go @@ -0,0 +1,161 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "log" + "net/http" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const endpointAccessGroupNew = "/access_group/new" + +type accessGroupInfoResponse struct { + AccessGroup string `json:"access_group"` + ModelNames []string `json:"model_names"` + DeploymentCount int `json:"deployment_count"` +} + +func resourceLiteLLMAccessGroup() *schema.Resource { + return &schema.Resource{ + Create: resourceLiteLLMAccessGroupCreate, + Read: resourceLiteLLMAccessGroupRead, + Update: resourceLiteLLMAccessGroupUpdate, + Delete: resourceLiteLLMAccessGroupDelete, + + Importer: &schema.ResourceImporter{StateContext: schema.ImportStatePassthroughContext}, + + Schema: map[string]*schema.Schema{ + "access_group": { + Type: schema.TypeString, + Required: true, + ForceNew: true, + }, + "model_names": { + Type: schema.TypeList, + Optional: true, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "model_ids": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "deployment_count": { + Type: schema.TypeInt, + Computed: true, + }, + }, + } +} + +func buildAccessGroupData(d *schema.ResourceData) map[string]interface{} { + data := map[string]interface{}{} + for _, key := range []string{"model_names", "model_ids"} { + if v, ok := d.GetOk(key); ok { + data[key] = v + } + } + return data +} + +func resourceLiteLLMAccessGroupCreate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + name := d.Get("access_group").(string) + groupData := buildAccessGroupData(d) + groupData["access_group"] = name + + log.Printf("[DEBUG] Create access group request payload: %+v", groupData) + + resp, err := MakeRequest(client, "POST", endpointAccessGroupNew, groupData) + if err != nil { + return fmt.Errorf("error creating access group: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "creating access group"); err != nil { + return err + } + + d.SetId(name) + log.Printf("[INFO] Access group created with name: %s", name) + + return resourceLiteLLMAccessGroupRead(d, m) +} + +func resourceLiteLLMAccessGroupRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Reading access group: %s", d.Id()) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("/access_group/%s/info", d.Id()), nil) + if err != nil { + return fmt.Errorf("error reading access group: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + log.Printf("[WARN] Access group %s not found, removing from state", d.Id()) + d.SetId("") + return nil + } + + if err := handleResponse(resp, "reading access group"); err != nil { + return err + } + + var info accessGroupInfoResponse + if err := json.NewDecoder(resp.Body).Decode(&info); err != nil { + return fmt.Errorf("error decoding access group info response: %w", err) + } + + d.Set("access_group", GetStringValue(info.AccessGroup, d.Id())) + d.Set("model_names", info.ModelNames) + d.Set("deployment_count", info.DeploymentCount) + + log.Printf("[INFO] Successfully read access group: %s", d.Id()) + return nil +} + +func resourceLiteLLMAccessGroupUpdate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + groupData := buildAccessGroupData(d) + log.Printf("[DEBUG] Update access group request payload: %+v", groupData) + + resp, err := MakeRequest(client, "PUT", fmt.Sprintf("/access_group/%s/update", d.Id()), groupData) + if err != nil { + return fmt.Errorf("error updating access group: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "updating access group"); err != nil { + return err + } + + log.Printf("[INFO] Successfully updated access group: %s", d.Id()) + return resourceLiteLLMAccessGroupRead(d, m) +} + +func resourceLiteLLMAccessGroupDelete(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Deleting access group: %s", d.Id()) + + resp, err := MakeRequest(client, "DELETE", fmt.Sprintf("/access_group/%s/delete", d.Id()), nil) + if err != nil { + return fmt.Errorf("error deleting access group: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "deleting access group"); err != nil { + return err + } + + log.Printf("[INFO] Successfully deleted access group: %s", d.Id()) + d.SetId("") + return nil +} diff --git a/terraform/provider/litellm/resource_access_group_test.go b/terraform/provider/litellm/resource_access_group_test.go new file mode 100644 index 00000000000..56ead47949a --- /dev/null +++ b/terraform/provider/litellm/resource_access_group_test.go @@ -0,0 +1,185 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "reflect" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func accessGroupTestData(t *testing.T, raw map[string]interface{}) *schema.ResourceData { + t.Helper() + return schema.TestResourceDataRaw(t, resourceLiteLLMAccessGroup().Schema, raw) +} + +func accessGroupInfoJSON(name string, modelNames []string, deploymentCount int) []byte { + body, _ := json.Marshal(accessGroupInfoResponse{ + AccessGroup: name, + ModelNames: modelNames, + DeploymentCount: deploymentCount, + }) + return body +} + +func TestAccessGroupCreate(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.Method + " " + r.URL.Path { + case "POST /access_group/new": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(`{"access_group": "prod-models", "models_updated": 2}`)) + case "GET /access_group/prod-models/info": + w.Write(accessGroupInfoJSON("prod-models", []string{"gpt-4", "claude-3"}, 2)) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := accessGroupTestData(t, map[string]interface{}{ + "access_group": "prod-models", + "model_names": []interface{}{"gpt-4", "claude-3"}, + }) + + if err := resourceLiteLLMAccessGroupCreate(d, client); err != nil { + t.Fatalf("create failed: %v", err) + } + + if createPayload["access_group"] != "prod-models" { + t.Fatalf("expected access_group 'prod-models' in payload, got %v", createPayload["access_group"]) + } + wantModels := []interface{}{"gpt-4", "claude-3"} + if !reflect.DeepEqual(createPayload["model_names"], wantModels) { + t.Fatalf("expected model_names %v in payload, got %v", wantModels, createPayload["model_names"]) + } + if d.Id() != "prod-models" { + t.Fatalf("expected ID 'prod-models', got %q", d.Id()) + } + if d.Get("deployment_count").(int) != 2 { + t.Fatalf("expected deployment_count 2, got %v", d.Get("deployment_count")) + } +} + +func TestAccessGroupRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" || r.URL.Path != "/access_group/prod-models/info" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + return + } + w.Write(accessGroupInfoJSON("prod-models", []string{"gpt-4"}, 1)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := accessGroupTestData(t, map[string]interface{}{"access_group": "prod-models"}) + d.SetId("prod-models") + + if err := resourceLiteLLMAccessGroupRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Get("access_group").(string) != "prod-models" { + t.Fatalf("expected access_group 'prod-models', got %v", d.Get("access_group")) + } + wantModels := []interface{}{"gpt-4"} + if !reflect.DeepEqual(d.Get("model_names"), wantModels) { + t.Fatalf("expected model_names %v, got %v", wantModels, d.Get("model_names")) + } + if d.Get("deployment_count").(int) != 1 { + t.Fatalf("expected deployment_count 1, got %v", d.Get("deployment_count")) + } +} + +func TestAccessGroupReadNotFound(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := accessGroupTestData(t, map[string]interface{}{"access_group": "gone"}) + d.SetId("gone") + + if err := resourceLiteLLMAccessGroupRead(d, client); err != nil { + t.Fatalf("expected nil error on 404, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID to be cleared on 404, got %q", d.Id()) + } +} + +func TestAccessGroupUpdate(t *testing.T) { + var updatePayload map[string]interface{} + var updatePath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case "PUT": + updatePath = r.URL.Path + if err := json.NewDecoder(r.Body).Decode(&updatePayload); err != nil { + t.Errorf("failed to decode update payload: %v", err) + } + w.Write([]byte(`{"access_group": "prod-models", "models_updated": 1}`)) + case "GET": + w.Write(accessGroupInfoJSON("prod-models", []string{"gpt-4o"}, 1)) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := accessGroupTestData(t, map[string]interface{}{ + "access_group": "prod-models", + "model_names": []interface{}{"gpt-4o"}, + }) + d.SetId("prod-models") + + if err := resourceLiteLLMAccessGroupUpdate(d, client); err != nil { + t.Fatalf("update failed: %v", err) + } + + if updatePath != "/access_group/prod-models/update" { + t.Fatalf("expected update path '/access_group/prod-models/update', got %q", updatePath) + } + wantModels := []interface{}{"gpt-4o"} + if !reflect.DeepEqual(updatePayload["model_names"], wantModels) { + t.Fatalf("expected model_names %v in payload, got %v", wantModels, updatePayload["model_names"]) + } + if _, ok := updatePayload["access_group"]; ok { + t.Fatalf("update payload must not include access_group, got %v", updatePayload["access_group"]) + } +} + +func TestAccessGroupDelete(t *testing.T) { + var deleteMethod, deletePath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + deleteMethod = r.Method + deletePath = r.URL.Path + w.Write([]byte(`{"access_group": "prod-models", "models_updated": 2, "message": "deleted"}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := accessGroupTestData(t, map[string]interface{}{"access_group": "prod-models"}) + d.SetId("prod-models") + + if err := resourceLiteLLMAccessGroupDelete(d, client); err != nil { + t.Fatalf("delete failed: %v", err) + } + + if deleteMethod != "DELETE" || deletePath != "/access_group/prod-models/delete" { + t.Fatalf("expected DELETE /access_group/prod-models/delete, got %s %s", deleteMethod, deletePath) + } + if d.Id() != "" { + t.Fatalf("expected ID to be cleared after delete, got %q", d.Id()) + } +} diff --git a/terraform/provider/litellm/resource_agent.go b/terraform/provider/litellm/resource_agent.go new file mode 100644 index 00000000000..4d141595fff --- /dev/null +++ b/terraform/provider/litellm/resource_agent.go @@ -0,0 +1,320 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "log" + "net/http" + "reflect" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const ( + endpointAgents = "/v1/agents" + endpointAgentByID = "/v1/agents/%s" +) + +type agentAPIResponse struct { + AgentID string `json:"agent_id"` + AgentName string `json:"agent_name"` + AgentCardParams map[string]interface{} `json:"agent_card_params"` + ObjectPermission map[string]interface{} `json:"object_permission"` + ExtraHeaders []string `json:"extra_headers"` + TPMLimit *int `json:"tpm_limit"` + RPMLimit *int `json:"rpm_limit"` + SessionTPMLimit *int `json:"session_tpm_limit"` + SessionRPMLimit *int `json:"session_rpm_limit"` + Spend *float64 `json:"spend"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` + CreatedBy string `json:"created_by"` + UpdatedBy string `json:"updated_by"` +} + +func agentSuppressEquivalentJSON(k, oldValue, newValue string, d *schema.ResourceData) bool { + var oldObj, newObj interface{} + if err := json.Unmarshal([]byte(oldValue), &oldObj); err != nil { + return false + } + if err := json.Unmarshal([]byte(newValue), &newObj); err != nil { + return false + } + return reflect.DeepEqual(oldObj, newObj) +} + +func agentParseJSONObject(raw, field string) (map[string]interface{}, error) { + var obj map[string]interface{} + if err := json.Unmarshal([]byte(raw), &obj); err != nil { + return nil, fmt.Errorf("%s must be a JSON object: %w", field, err) + } + return obj, nil +} + +func resourceLiteLLMAgent() *schema.Resource { + return &schema.Resource{ + Create: resourceLiteLLMAgentCreate, + Read: resourceLiteLLMAgentRead, + Update: resourceLiteLLMAgentUpdate, + Delete: resourceLiteLLMAgentDelete, + + Importer: &schema.ResourceImporter{StateContext: schema.ImportStatePassthroughContext}, + + Schema: map[string]*schema.Schema{ + "agent_name": { + Type: schema.TypeString, + Required: true, + Description: "Name of the agent.", + }, + "agent_card_params": { + Type: schema.TypeString, + Required: true, + DiffSuppressFunc: agentSuppressEquivalentJSON, + Description: "A2A agent card as a JSON object string (name, description, url, version, " + + "capabilities, skills, ...). The proxy merges in LiteLLM-fronting fields, so the configured " + + "value stays authoritative in state.", + }, + "litellm_params": { + Type: schema.TypeString, + Optional: true, + Sensitive: true, + DiffSuppressFunc: agentSuppressEquivalentJSON, + Description: "LiteLLM-specific parameters as a JSON object string (may include model, api_key, ...). " + + "Never read back from the API.", + }, + "object_permission": { + Type: schema.TypeString, + Optional: true, + DiffSuppressFunc: agentSuppressEquivalentJSON, + Description: "Access control permissions as a JSON object string " + + "(mcp_servers, mcp_access_groups, mcp_tool_permissions, models, agents).", + }, + "static_headers": { + Type: schema.TypeMap, + Optional: true, + Sensitive: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Static headers sent with agent requests (may hold tokens). Never read back from the API.", + }, + "extra_headers": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Names of incoming request headers to forward to the agent.", + }, + "tpm_limit": { + Type: schema.TypeInt, + Optional: true, + }, + "rpm_limit": { + Type: schema.TypeInt, + Optional: true, + }, + "session_tpm_limit": { + Type: schema.TypeInt, + Optional: true, + }, + "session_rpm_limit": { + Type: schema.TypeInt, + Optional: true, + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + "created_by": { + Type: schema.TypeString, + Computed: true, + }, + "updated_by": { + Type: schema.TypeString, + Computed: true, + }, + }, + } +} + +func buildAgentData(d *schema.ResourceData) (map[string]interface{}, error) { + card, err := agentParseJSONObject(d.Get("agent_card_params").(string), "agent_card_params") + if err != nil { + return nil, err + } + + agentData := map[string]interface{}{ + "agent_name": d.Get("agent_name").(string), + "agent_card_params": card, + } + + for _, key := range []string{"litellm_params", "object_permission"} { + raw, ok := d.GetOk(key) + if !ok || raw.(string) == "" { + continue + } + obj, err := agentParseJSONObject(raw.(string), key) + if err != nil { + return nil, err + } + agentData[key] = obj + } + + for _, key := range []string{"static_headers", "extra_headers", "tpm_limit", "rpm_limit", "session_tpm_limit", "session_rpm_limit"} { + if v, ok := d.GetOk(key); ok { + agentData[key] = v + } + } + + return agentData, nil +} + +func resourceLiteLLMAgentCreate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + agentData, err := buildAgentData(d) + if err != nil { + return err + } + + log.Printf("[DEBUG] Create agent request for: %s", d.Get("agent_name").(string)) + + resp, err := MakeRequest(client, "POST", endpointAgents, agentData) + if err != nil { + return fmt.Errorf("error creating agent: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "creating agent"); err != nil { + return err + } + + var agentResp agentAPIResponse + if err := json.NewDecoder(resp.Body).Decode(&agentResp); err != nil { + return fmt.Errorf("error decoding create agent response: %w", err) + } + if agentResp.AgentID == "" { + return fmt.Errorf("create agent response did not contain an agent_id") + } + + d.SetId(agentResp.AgentID) + log.Printf("[INFO] Agent created with ID: %s", agentResp.AgentID) + + return resourceLiteLLMAgentRead(d, m) +} + +func resourceLiteLLMAgentRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Reading agent with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf(endpointAgentByID, d.Id()), nil) + if err != nil { + return fmt.Errorf("error reading agent: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + log.Printf("[WARN] Agent with ID %s not found, removing from state", d.Id()) + d.SetId("") + return nil + } + + if err := handleResponse(resp, "reading agent"); err != nil { + return err + } + + var agentResp agentAPIResponse + if err := json.NewDecoder(resp.Body).Decode(&agentResp); err != nil { + return fmt.Errorf("error decoding agent info response: %w", err) + } + + d.Set("agent_name", agentResp.AgentName) + + // The proxy merges LiteLLM-fronting fields into the stored card, so the configured + // JSON stays authoritative; only populate from the API when importing. + if d.Get("agent_card_params").(string) == "" && agentResp.AgentCardParams != nil { + cardJSON, err := json.Marshal(agentResp.AgentCardParams) + if err != nil { + return fmt.Errorf("error encoding agent_card_params: %w", err) + } + d.Set("agent_card_params", string(cardJSON)) + } + if d.Get("object_permission").(string) == "" && agentResp.ObjectPermission != nil { + permJSON, err := json.Marshal(agentResp.ObjectPermission) + if err != nil { + return fmt.Errorf("error encoding object_permission: %w", err) + } + d.Set("object_permission", string(permJSON)) + } + + if agentResp.ExtraHeaders != nil { + d.Set("extra_headers", agentResp.ExtraHeaders) + } + if agentResp.TPMLimit != nil { + d.Set("tpm_limit", *agentResp.TPMLimit) + } + if agentResp.RPMLimit != nil { + d.Set("rpm_limit", *agentResp.RPMLimit) + } + if agentResp.SessionTPMLimit != nil { + d.Set("session_tpm_limit", *agentResp.SessionTPMLimit) + } + if agentResp.SessionRPMLimit != nil { + d.Set("session_rpm_limit", *agentResp.SessionRPMLimit) + } + d.Set("created_at", agentResp.CreatedAt) + d.Set("updated_at", agentResp.UpdatedAt) + d.Set("created_by", agentResp.CreatedBy) + d.Set("updated_by", agentResp.UpdatedBy) + + log.Printf("[INFO] Successfully read agent with ID: %s", d.Id()) + return nil +} + +func resourceLiteLLMAgentUpdate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + agentData, err := buildAgentData(d) + if err != nil { + return err + } + + log.Printf("[DEBUG] Update agent request for ID: %s", d.Id()) + + resp, err := MakeRequest(client, "PATCH", fmt.Sprintf(endpointAgentByID, d.Id()), agentData) + if err != nil { + return fmt.Errorf("error updating agent: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "updating agent"); err != nil { + return err + } + + log.Printf("[INFO] Successfully updated agent with ID: %s", d.Id()) + return resourceLiteLLMAgentRead(d, m) +} + +func resourceLiteLLMAgentDelete(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Deleting agent with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "DELETE", fmt.Sprintf(endpointAgentByID, d.Id()), nil) + if err != nil { + return fmt.Errorf("error deleting agent: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusNotFound { + if err := handleResponse(resp, "deleting agent"); err != nil { + return err + } + } + + log.Printf("[INFO] Successfully deleted agent with ID: %s", d.Id()) + d.SetId("") + return nil +} diff --git a/terraform/provider/litellm/resource_agent_test.go b/terraform/provider/litellm/resource_agent_test.go new file mode 100644 index 00000000000..fadba98fdbe --- /dev/null +++ b/terraform/provider/litellm/resource_agent_test.go @@ -0,0 +1,235 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const testAgentCardJSON = `{"name": "Hello Agent", "url": "http://agent.local:9999/", "version": "1.0.0"}` + +func newAgentTestResourceData(t *testing.T) *schema.ResourceData { + t.Helper() + return schema.TestResourceDataRaw(t, resourceLiteLLMAgent().Schema, map[string]interface{}{ + "agent_name": "my-agent", + "agent_card_params": testAgentCardJSON, + "litellm_params": `{"model": "gpt-5.2", "api_key": "sk-secret"}`, + "extra_headers": []interface{}{"x-request-id"}, + "tpm_limit": 1000, + }) +} + +func agentReadResponseBody() []byte { + body, _ := json.Marshal(map[string]interface{}{ + "agent_id": "agent-123", + "agent_name": "my-agent", + "agent_card_params": map[string]interface{}{ + "name": "Hello Agent", + "url": "http://agent.local:9999/", + "version": "1.0.0", + "supportedInterfaces": []string{"http://proxy/a2a/agent-123"}, + }, + "litellm_params": map[string]interface{}{"model": "gpt-5.2", "api_key": "sk-1****"}, + "extra_headers": []string{"x-request-id"}, + "tpm_limit": 1000, + "spend": 1.5, + "created_at": "2026-01-01T00:00:00", + "updated_at": "2026-01-02T00:00:00", + "created_by": "admin", + "updated_by": "admin", + }) + return body +} + +func TestResourceLiteLLMAgentCreate(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch { + case r.Method == http.MethodPost && r.URL.Path == "/v1/agents": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(`{"agent_id": "agent-123", "agent_name": "my-agent", "agent_card_params": {}}`)) + case r.Method == http.MethodGet && r.URL.Path == "/v1/agents/agent-123": + w.Write(agentReadResponseBody()) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newAgentTestResourceData(t) + + if err := resourceLiteLLMAgentCreate(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "agent-123" { + t.Fatalf("expected ID 'agent-123', got %q", d.Id()) + } + + if createPayload["agent_name"] != "my-agent" { + t.Errorf("expected agent_name 'my-agent' in payload, got %v", createPayload["agent_name"]) + } + card, ok := createPayload["agent_card_params"].(map[string]interface{}) + if !ok || card["url"] != "http://agent.local:9999/" { + t.Errorf("expected agent_card_params sent as JSON object with url, got %v", createPayload["agent_card_params"]) + } + params, ok := createPayload["litellm_params"].(map[string]interface{}) + if !ok || params["api_key"] != "sk-secret" { + t.Errorf("expected litellm_params sent as JSON object, got %v", createPayload["litellm_params"]) + } + if createPayload["tpm_limit"] != float64(1000) { + t.Errorf("expected tpm_limit 1000 in payload, got %v", createPayload["tpm_limit"]) + } + + if d.Get("created_at").(string) != "2026-01-01T00:00:00" { + t.Errorf("expected created_at from read-back, got %q", d.Get("created_at").(string)) + } + if got := d.Get("agent_card_params").(string); got != testAgentCardJSON { + t.Errorf("expected configured agent_card_params to stay authoritative, got %q", got) + } + if got := d.Get("litellm_params").(string); got != `{"model": "gpt-5.2", "api_key": "sk-secret"}` { + t.Errorf("expected litellm_params to keep configured value, got %q", got) + } +} + +func TestResourceLiteLLMAgentCreateInvalidCardJSON(t *testing.T) { + d := schema.TestResourceDataRaw(t, resourceLiteLLMAgent().Schema, map[string]interface{}{ + "agent_name": "my-agent", + "agent_card_params": "not-json", + }) + client := NewClient("http://unused.invalid", "test-key", true) + + if err := resourceLiteLLMAgentCreate(d, client); err == nil { + t.Fatal("expected error for invalid agent_card_params JSON, got nil") + } +} + +func TestResourceLiteLLMAgentReadPopulatesStateOnImport(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/v1/agents/agent-123" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write(agentReadResponseBody()) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMAgent().Schema, map[string]interface{}{}) + d.SetId("agent-123") + + if err := resourceLiteLLMAgentRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Get("agent_name").(string) != "my-agent" { + t.Errorf("expected agent_name 'my-agent', got %q", d.Get("agent_name").(string)) + } + var card map[string]interface{} + if err := json.Unmarshal([]byte(d.Get("agent_card_params").(string)), &card); err != nil { + t.Fatalf("agent_card_params not populated as JSON on import: %v", err) + } + if card["name"] != "Hello Agent" { + t.Errorf("expected card name 'Hello Agent', got %v", card["name"]) + } + if d.Get("tpm_limit").(int) != 1000 { + t.Errorf("expected tpm_limit 1000, got %d", d.Get("tpm_limit").(int)) + } + headers := d.Get("extra_headers").([]interface{}) + if len(headers) != 1 || headers[0] != "x-request-id" { + t.Errorf("expected extra_headers ['x-request-id'], got %v", headers) + } + if d.Get("litellm_params").(string) != "" { + t.Errorf("expected litellm_params to never be read back, got %q", d.Get("litellm_params").(string)) + } +} + +func TestResourceLiteLLMAgentRead404ClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newAgentTestResourceData(t) + d.SetId("agent-123") + + if err := resourceLiteLLMAgentRead(d, client); err != nil { + t.Fatalf("expected nil error on 404, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared on 404, got %q", d.Id()) + } +} + +func TestResourceLiteLLMAgentUpdate(t *testing.T) { + var updateMethod, updatePath string + var updatePayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == http.MethodGet { + w.Write(agentReadResponseBody()) + return + } + updateMethod = r.Method + updatePath = r.URL.Path + if err := json.NewDecoder(r.Body).Decode(&updatePayload); err != nil { + t.Errorf("failed to decode update payload: %v", err) + } + w.Write([]byte(`{}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newAgentTestResourceData(t) + d.SetId("agent-123") + + if err := resourceLiteLLMAgentUpdate(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if updateMethod != http.MethodPatch { + t.Errorf("expected PATCH, got %s", updateMethod) + } + if updatePath != "/v1/agents/agent-123" { + t.Errorf("expected path '/v1/agents/agent-123', got %q", updatePath) + } + if updatePayload["agent_name"] != "my-agent" { + t.Errorf("expected agent_name in update payload, got %v", updatePayload["agent_name"]) + } + if updatePayload["tpm_limit"] != float64(1000) { + t.Errorf("expected tpm_limit 1000 in update payload, got %v", updatePayload["tpm_limit"]) + } +} + +func TestResourceLiteLLMAgentDelete(t *testing.T) { + var deleteMethod, deletePath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + deleteMethod = r.Method + deletePath = r.URL.Path + w.Write([]byte(`{}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newAgentTestResourceData(t) + d.SetId("agent-123") + + if err := resourceLiteLLMAgentDelete(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if deleteMethod != http.MethodDelete { + t.Errorf("expected DELETE, got %s", deleteMethod) + } + if deletePath != "/v1/agents/agent-123" { + t.Errorf("expected path '/v1/agents/agent-123', got %q", deletePath) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared after delete, got %q", d.Id()) + } +} diff --git a/terraform/provider/litellm/resource_budget.go b/terraform/provider/litellm/resource_budget.go new file mode 100644 index 00000000000..8d56ccfd9a3 --- /dev/null +++ b/terraform/provider/litellm/resource_budget.go @@ -0,0 +1,287 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "log" + "net/http" + "reflect" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation" +) + +const ( + endpointBudgetNew = "/budget/new" + endpointBudgetInfo = "/budget/info" + endpointBudgetUpdate = "/budget/update" + endpointBudgetDelete = "/budget/delete" +) + +func budgetSuppressEquivalentJSON(k, oldValue, newValue string, d *schema.ResourceData) bool { + var oldParsed, newParsed interface{} + if err := json.Unmarshal([]byte(oldValue), &oldParsed); err != nil { + return false + } + if err := json.Unmarshal([]byte(newValue), &newParsed); err != nil { + return false + } + return reflect.DeepEqual(oldParsed, newParsed) +} + +func resourceLiteLLMBudget() *schema.Resource { + return &schema.Resource{ + Create: resourceLiteLLMBudgetCreate, + Read: resourceLiteLLMBudgetRead, + Update: resourceLiteLLMBudgetUpdate, + Delete: resourceLiteLLMBudgetDelete, + + Importer: &schema.ResourceImporter{ + StateContext: schema.ImportStatePassthroughContext, + }, + + Schema: map[string]*schema.Schema{ + "budget_id": { + Type: schema.TypeString, + Optional: true, + Computed: true, + ForceNew: true, + Description: "Unique ID for the budget. Generated by the server if not provided", + }, + "max_budget": { + Type: schema.TypeFloat, + Optional: true, + Description: "Requests fail if this budget in USD is exceeded", + }, + "soft_budget": { + Type: schema.TypeFloat, + Optional: true, + Description: "Requests do not fail if this is exceeded, but alerts fire", + }, + "max_parallel_requests": { + Type: schema.TypeInt, + Optional: true, + Description: "Maximum concurrent requests allowed for this budget", + }, + "tpm_limit": { + Type: schema.TypeInt, + Optional: true, + Description: "Maximum tokens per minute allowed for this budget", + }, + "rpm_limit": { + Type: schema.TypeInt, + Optional: true, + Description: "Maximum requests per minute allowed for this budget", + }, + "budget_duration": { + Type: schema.TypeString, + Optional: true, + Description: "Budget reset period (e.g. '1hr', '1d', '28d')", + }, + "model_max_budget": { + Type: schema.TypeString, + Optional: true, + ValidateFunc: validation.StringIsJSON, + DiffSuppressFunc: budgetSuppressEquivalentJSON, + Description: "JSON string of per-model budget config (e.g. '{\"gpt-4o\": {\"max_budget\": 10.0}}')", + }, + "budget_reset_at": { + Type: schema.TypeString, + Computed: true, + Description: "Datetime when the budget is reset", + }, + }, + } +} + +type budgetResponse struct { + BudgetID string `json:"budget_id"` + MaxBudget *float64 `json:"max_budget"` + SoftBudget *float64 `json:"soft_budget"` + MaxParallelRequests *int `json:"max_parallel_requests"` + TPMLimit *int `json:"tpm_limit"` + RPMLimit *int `json:"rpm_limit"` + BudgetDuration *string `json:"budget_duration"` + ModelMaxBudget interface{} `json:"model_max_budget"` + BudgetResetAt *string `json:"budget_reset_at"` +} + +func budgetModelMaxBudgetString(v interface{}) (string, bool) { + switch typed := v.(type) { + case string: + return typed, typed != "" + case map[string]interface{}: + if len(typed) == 0 { + return "", false + } + encoded, err := json.Marshal(typed) + return string(encoded), err == nil + } + return "", false +} + +func setBudgetState(d *schema.ResourceData, budgetResp budgetResponse) { + if budgetResp.MaxBudget != nil { + d.Set("max_budget", *budgetResp.MaxBudget) + } + if budgetResp.SoftBudget != nil { + d.Set("soft_budget", *budgetResp.SoftBudget) + } + if budgetResp.MaxParallelRequests != nil { + d.Set("max_parallel_requests", *budgetResp.MaxParallelRequests) + } + if budgetResp.TPMLimit != nil { + d.Set("tpm_limit", *budgetResp.TPMLimit) + } + if budgetResp.RPMLimit != nil { + d.Set("rpm_limit", *budgetResp.RPMLimit) + } + if budgetResp.BudgetDuration != nil { + d.Set("budget_duration", *budgetResp.BudgetDuration) + } + if encoded, ok := budgetModelMaxBudgetString(budgetResp.ModelMaxBudget); ok { + d.Set("model_max_budget", encoded) + } + if budgetResp.BudgetResetAt != nil { + d.Set("budget_reset_at", *budgetResp.BudgetResetAt) + } +} + +func resourceLiteLLMBudgetCreate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + budgetData := buildBudgetData(d) + if v, ok := d.GetOk("budget_id"); ok { + budgetData["budget_id"] = v.(string) + } + + log.Printf("[DEBUG] Create budget request payload: %+v", budgetData) + + resp, err := MakeRequest(client, "POST", endpointBudgetNew, budgetData) + if err != nil { + return fmt.Errorf("error creating budget: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "creating budget"); err != nil { + return err + } + + var budgetResp budgetResponse + if err := json.NewDecoder(resp.Body).Decode(&budgetResp); err != nil { + return fmt.Errorf("error decoding create budget response: %w", err) + } + if budgetResp.BudgetID == "" { + return fmt.Errorf("create budget response did not contain a budget_id") + } + + d.SetId(budgetResp.BudgetID) + log.Printf("[INFO] Budget created with ID: %s", budgetResp.BudgetID) + + return resourceLiteLLMBudgetRead(d, m) +} + +func resourceLiteLLMBudgetRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Reading budget with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "POST", endpointBudgetInfo, map[string]interface{}{ + "budgets": []string{d.Id()}, + }) + if err != nil { + return fmt.Errorf("error reading budget: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + log.Printf("[WARN] Budget with ID %s not found, removing from state", d.Id()) + d.SetId("") + return nil + } + + if err := handleResponse(resp, "reading budget"); err != nil { + return err + } + + var budgetResps []budgetResponse + if err := json.NewDecoder(resp.Body).Decode(&budgetResps); err != nil { + return fmt.Errorf("error decoding budget info response: %w", err) + } + if len(budgetResps) == 0 { + log.Printf("[WARN] Budget with ID %s not found in response, removing from state", d.Id()) + d.SetId("") + return nil + } + + d.Set("budget_id", d.Id()) + setBudgetState(d, budgetResps[0]) + + log.Printf("[INFO] Successfully read budget with ID: %s", d.Id()) + return nil +} + +func resourceLiteLLMBudgetUpdate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + budgetData := buildBudgetData(d) + budgetData["budget_id"] = d.Id() + + log.Printf("[DEBUG] Update budget request payload: %+v", budgetData) + + resp, err := MakeRequest(client, "POST", endpointBudgetUpdate, budgetData) + if err != nil { + return fmt.Errorf("error updating budget: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "updating budget"); err != nil { + return err + } + + log.Printf("[INFO] Successfully updated budget with ID: %s", d.Id()) + return resourceLiteLLMBudgetRead(d, m) +} + +func resourceLiteLLMBudgetDelete(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Deleting budget with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "POST", endpointBudgetDelete, map[string]interface{}{ + "id": d.Id(), + }) + if err != nil { + return fmt.Errorf("error deleting budget: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "deleting budget"); err != nil { + return err + } + + log.Printf("[INFO] Successfully deleted budget with ID: %s", d.Id()) + d.SetId("") + return nil +} + +func buildBudgetData(d *schema.ResourceData) map[string]interface{} { + budgetData := map[string]interface{}{} + + for _, key := range []string{ + "max_budget", "soft_budget", "max_parallel_requests", "tpm_limit", "rpm_limit", "budget_duration", + } { + if v, ok := d.GetOk(key); ok { + budgetData[key] = v + } + } + + if v, ok := d.GetOk("model_max_budget"); ok { + var parsed map[string]interface{} + if err := json.Unmarshal([]byte(v.(string)), &parsed); err == nil { + budgetData["model_max_budget"] = parsed + } + } + + return budgetData +} diff --git a/terraform/provider/litellm/resource_budget_test.go b/terraform/provider/litellm/resource_budget_test.go new file mode 100644 index 00000000000..d5108520da1 --- /dev/null +++ b/terraform/provider/litellm/resource_budget_test.go @@ -0,0 +1,268 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func budgetInfoBody(budgetID string) []byte { + body, _ := json.Marshal([]map[string]interface{}{{ + "budget_id": budgetID, + "max_budget": 100.0, + "soft_budget": 80.0, + "max_parallel_requests": 10, + "tpm_limit": 1000, + "rpm_limit": 60, + "budget_duration": "30d", + "model_max_budget": map[string]interface{}{"gpt-4o": map[string]interface{}{"max_budget": 5.0}}, + "budget_reset_at": "2026-09-01T00:00:00Z", + }}) + return body +} + +func TestResourceBudgetCreate_ServerGeneratedID(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/budget/new": + if r.Method != http.MethodPost { + t.Errorf("expected POST /budget/new, got %s", r.Method) + } + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Fatalf("failed to decode create payload: %v", err) + } + w.Write([]byte(`{"budget_id": "bud-generated", "max_budget": 100.0}`)) + case "/budget/info": + var infoPayload map[string]interface{} + if err := json.NewDecoder(r.Body).Decode(&infoPayload); err != nil { + t.Fatalf("failed to decode info payload: %v", err) + } + budgets, ok := infoPayload["budgets"].([]interface{}) + if !ok || len(budgets) != 1 || budgets[0] != "bud-generated" { + t.Errorf("expected budgets ['bud-generated'], got %v", infoPayload["budgets"]) + } + w.Write(budgetInfoBody("bud-generated")) + default: + t.Errorf("unexpected request to %s", r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMBudget().Schema, map[string]interface{}{ + "max_budget": 100.0, + "soft_budget": 80.0, + "tpm_limit": 1000, + "model_max_budget": `{"gpt-4o": {"max_budget": 5.0}}`, + }) + + if err := resourceLiteLLMBudgetCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + if d.Id() != "bud-generated" { + t.Fatalf("expected ID 'bud-generated', got %q", d.Id()) + } + if _, ok := createPayload["budget_id"]; ok { + t.Errorf("budget_id must be omitted when not configured, got %v", createPayload["budget_id"]) + } + if got := createPayload["max_budget"]; got != 100.0 { + t.Errorf("expected max_budget 100.0 in payload, got %v", got) + } + if got := createPayload["soft_budget"]; got != 80.0 { + t.Errorf("expected soft_budget 80.0 in payload, got %v", got) + } + mmb, ok := createPayload["model_max_budget"].(map[string]interface{}) + if !ok { + t.Fatalf("expected model_max_budget object in payload, got %v", createPayload["model_max_budget"]) + } + if _, ok := mmb["gpt-4o"]; !ok { + t.Errorf("expected gpt-4o key in model_max_budget, got %v", mmb) + } + if got := d.Get("budget_reset_at").(string); got != "2026-09-01T00:00:00Z" { + t.Errorf("expected budget_reset_at from read, got %q", got) + } +} + +func TestResourceBudgetCreate_ConfiguredID(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/budget/new": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Fatalf("failed to decode create payload: %v", err) + } + w.Write([]byte(`{"budget_id": "my-budget"}`)) + case "/budget/info": + w.Write(budgetInfoBody("my-budget")) + default: + t.Errorf("unexpected request to %s", r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMBudget().Schema, map[string]interface{}{ + "budget_id": "my-budget", + "max_budget": 100.0, + }) + + if err := resourceLiteLLMBudgetCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + if d.Id() != "my-budget" { + t.Fatalf("expected ID 'my-budget', got %q", d.Id()) + } + if got := createPayload["budget_id"]; got != "my-budget" { + t.Errorf("expected budget_id 'my-budget' in payload, got %v", got) + } +} + +func TestResourceBudgetRead_MapsFields(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write(budgetInfoBody("bud-1")) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMBudget().Schema, map[string]interface{}{}) + d.SetId("bud-1") + + if err := resourceLiteLLMBudgetRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + if got := d.Get("max_budget").(float64); got != 100.0 { + t.Errorf("expected max_budget 100.0, got %v", got) + } + if got := d.Get("soft_budget").(float64); got != 80.0 { + t.Errorf("expected soft_budget 80.0, got %v", got) + } + if got := d.Get("max_parallel_requests").(int); got != 10 { + t.Errorf("expected max_parallel_requests 10, got %d", got) + } + if got := d.Get("tpm_limit").(int); got != 1000 { + t.Errorf("expected tpm_limit 1000, got %d", got) + } + if got := d.Get("rpm_limit").(int); got != 60 { + t.Errorf("expected rpm_limit 60, got %d", got) + } + if got := d.Get("budget_duration").(string); got != "30d" { + t.Errorf("expected budget_duration '30d', got %q", got) + } + var mmb map[string]interface{} + if err := json.Unmarshal([]byte(d.Get("model_max_budget").(string)), &mmb); err != nil { + t.Fatalf("model_max_budget in state is not valid JSON: %v", err) + } + if _, ok := mmb["gpt-4o"]; !ok { + t.Errorf("expected gpt-4o key in model_max_budget state, got %v", mmb) + } +} + +func TestResourceBudgetRead_EmptyListClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write([]byte(`[]`)) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMBudget().Schema, map[string]interface{}{}) + d.SetId("gone-budget") + + if err := resourceLiteLLMBudgetRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("expected nil error on empty response, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID to be cleared, got %q", d.Id()) + } +} + +func TestResourceBudgetRead_404ClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMBudget().Schema, map[string]interface{}{}) + d.SetId("gone-budget") + + if err := resourceLiteLLMBudgetRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("expected nil error on 404, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID to be cleared on 404, got %q", d.Id()) + } +} + +func TestResourceBudgetUpdate_SendsPayload(t *testing.T) { + var updatePayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/budget/update": + if r.Method != http.MethodPost { + t.Errorf("expected POST /budget/update, got %s", r.Method) + } + if err := json.NewDecoder(r.Body).Decode(&updatePayload); err != nil { + t.Fatalf("failed to decode update payload: %v", err) + } + w.Write([]byte(`{"budget_id": "bud-1"}`)) + case "/budget/info": + w.Write(budgetInfoBody("bud-1")) + default: + t.Errorf("unexpected request to %s", r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMBudget().Schema, map[string]interface{}{ + "max_budget": 200.0, + "rpm_limit": 120, + }) + d.SetId("bud-1") + + if err := resourceLiteLLMBudgetUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + + if got := updatePayload["budget_id"]; got != "bud-1" { + t.Errorf("expected budget_id 'bud-1' in payload, got %v", got) + } + if got := updatePayload["max_budget"]; got != 200.0 { + t.Errorf("expected max_budget 200.0 in payload, got %v", got) + } + if got := updatePayload["rpm_limit"]; got != 120.0 { + t.Errorf("expected rpm_limit 120 in payload, got %v", got) + } +} + +func TestResourceBudgetDelete_SendsID(t *testing.T) { + var deletePayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/budget/delete" || r.Method != http.MethodPost { + t.Errorf("expected POST /budget/delete, got %s %s", r.Method, r.URL.Path) + } + if err := json.NewDecoder(r.Body).Decode(&deletePayload); err != nil { + t.Fatalf("failed to decode delete payload: %v", err) + } + w.Write([]byte(`{}`)) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMBudget().Schema, map[string]interface{}{}) + d.SetId("bud-del") + + if err := resourceLiteLLMBudgetDelete(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("delete failed: %v", err) + } + + if got := deletePayload["id"]; got != "bud-del" { + t.Fatalf("expected id 'bud-del' in payload, got %v", got) + } + if d.Id() != "" { + t.Fatalf("expected ID to be cleared after delete, got %q", d.Id()) + } +} diff --git a/terraform/provider/litellm/resource_fallback.go b/terraform/provider/litellm/resource_fallback.go new file mode 100644 index 00000000000..680e051e60e --- /dev/null +++ b/terraform/provider/litellm/resource_fallback.go @@ -0,0 +1,155 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "log" + "net/http" + "net/url" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation" +) + +const endpointFallbackCreate = "/fallback" + +type FallbackGetResponse struct { + Model string `json:"model"` + FallbackModels []string `json:"fallback_models"` + FallbackType string `json:"fallback_type"` +} + +func resourceLiteLLMFallback() *schema.Resource { + return &schema.Resource{ + Create: resourceLiteLLMFallbackCreate, + Read: resourceLiteLLMFallbackRead, + Update: resourceLiteLLMFallbackUpdate, + Delete: resourceLiteLLMFallbackDelete, + + Importer: &schema.ResourceImporter{ + StateContext: schema.ImportStatePassthroughContext, + }, + + Schema: map[string]*schema.Schema{ + "model": { + Type: schema.TypeString, + Required: true, + ForceNew: true, + Description: "The model name to configure fallbacks for", + }, + "fallback_models": { + Type: schema.TypeList, + Required: true, + MinItems: 1, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "List of fallback model names in order of priority", + }, + "fallback_type": { + Type: schema.TypeString, + Optional: true, + ForceNew: true, + Default: "general", + ValidateFunc: validation.StringInSlice([]string{"general", "context_window", "content_policy"}, false), + Description: "Type of fallback: 'general' (default), 'context_window', or 'content_policy'", + }, + }, + } +} + +func fallbackTypeFromState(d *schema.ResourceData) string { + return GetStringValue(d.Get("fallback_type").(string), "general") +} + +func buildFallbackData(d *schema.ResourceData) map[string]interface{} { + return map[string]interface{}{ + "model": d.Get("model").(string), + "fallback_models": d.Get("fallback_models"), + "fallback_type": fallbackTypeFromState(d), + } +} + +func upsertLiteLLMFallback(d *schema.ResourceData, m interface{}, action string) error { + client := m.(*Client) + + fallbackData := buildFallbackData(d) + log.Printf("[DEBUG] %s fallback request payload: %+v", action, fallbackData) + + resp, err := MakeRequest(client, "POST", endpointFallbackCreate, fallbackData) + if err != nil { + return fmt.Errorf("error %s fallback: %w", action, err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, action+" fallback"); err != nil { + return err + } + + d.SetId(d.Get("model").(string)) + return resourceLiteLLMFallbackRead(d, m) +} + +func resourceLiteLLMFallbackCreate(d *schema.ResourceData, m interface{}) error { + return upsertLiteLLMFallback(d, m, "creating") +} + +func resourceLiteLLMFallbackRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Reading fallback for model: %s", d.Id()) + + endpoint := fmt.Sprintf("/fallback/%s?fallback_type=%s", + url.PathEscape(d.Id()), url.QueryEscape(fallbackTypeFromState(d))) + resp, err := MakeRequest(client, "GET", endpoint, nil) + if err != nil { + return fmt.Errorf("error reading fallback: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + log.Printf("[WARN] Fallback for model %s not found, removing from state", d.Id()) + d.SetId("") + return nil + } + + if err := handleResponse(resp, "reading fallback"); err != nil { + return err + } + + var fallbackResp FallbackGetResponse + if err := json.NewDecoder(resp.Body).Decode(&fallbackResp); err != nil { + return fmt.Errorf("error decoding fallback response: %w", err) + } + + d.Set("model", GetStringValue(fallbackResp.Model, d.Id())) + d.Set("fallback_models", fallbackResp.FallbackModels) + d.Set("fallback_type", GetStringValue(fallbackResp.FallbackType, fallbackTypeFromState(d))) + + return nil +} + +func resourceLiteLLMFallbackUpdate(d *schema.ResourceData, m interface{}) error { + return upsertLiteLLMFallback(d, m, "updating") +} + +func resourceLiteLLMFallbackDelete(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Deleting fallback for model: %s", d.Id()) + + endpoint := fmt.Sprintf("/fallback/%s?fallback_type=%s", + url.PathEscape(d.Id()), url.QueryEscape(fallbackTypeFromState(d))) + resp, err := MakeRequest(client, "DELETE", endpoint, nil) + if err != nil { + return fmt.Errorf("error deleting fallback: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusNotFound { + if err := handleResponse(resp, "deleting fallback"); err != nil { + return err + } + } + + d.SetId("") + return nil +} diff --git a/terraform/provider/litellm/resource_fallback_test.go b/terraform/provider/litellm/resource_fallback_test.go new file mode 100644 index 00000000000..e2c90d25424 --- /dev/null +++ b/terraform/provider/litellm/resource_fallback_test.go @@ -0,0 +1,180 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "reflect" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func newFallbackTestResourceData(t *testing.T, model string, fallbackModels []interface{}, fallbackType string) *schema.ResourceData { + t.Helper() + d := schema.TestResourceDataRaw(t, resourceLiteLLMFallback().Schema, map[string]interface{}{ + "model": model, + "fallback_models": fallbackModels, + "fallback_type": fallbackType, + }) + return d +} + +func fallbackGetHandler(t *testing.T, wantPath string, resp FallbackGetResponse) http.HandlerFunc { + t.Helper() + return func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + t.Errorf("expected GET, got %s", r.Method) + } + if r.URL.Path != wantPath { + t.Errorf("expected path %s, got %s", wantPath, r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(resp) + } +} + +func TestResourceLiteLLMFallbackCreate(t *testing.T) { + var createPayload map[string]interface{} + mux := http.NewServeMux() + mux.HandleFunc("/fallback", func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("expected POST, got %s", r.Method) + } + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Fatalf("failed to decode create payload: %v", err) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"model":"gpt-4","fallback_models":["claude-3","gpt-3.5-turbo"],"fallback_type":"general","message":"ok"}`)) + }) + mux.Handle("/fallback/gpt-4", fallbackGetHandler(t, "/fallback/gpt-4", FallbackGetResponse{ + Model: "gpt-4", + FallbackModels: []string{"claude-3", "gpt-3.5-turbo"}, + FallbackType: "general", + })) + srv := httptest.NewServer(mux) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newFallbackTestResourceData(t, "gpt-4", []interface{}{"claude-3", "gpt-3.5-turbo"}, "general") + + if err := resourceLiteLLMFallbackCreate(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "gpt-4" { + t.Fatalf("expected ID 'gpt-4', got %q", d.Id()) + } + want := map[string]interface{}{ + "model": "gpt-4", + "fallback_models": []interface{}{"claude-3", "gpt-3.5-turbo"}, + "fallback_type": "general", + } + if !reflect.DeepEqual(createPayload, want) { + t.Fatalf("unexpected create payload: %+v, want %+v", createPayload, want) + } +} + +func TestResourceLiteLLMFallbackRead_MapsFields(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/fallback/gpt-4" { + t.Errorf("expected path /fallback/gpt-4, got %s", r.URL.Path) + } + if got := r.URL.Query().Get("fallback_type"); got != "context_window" { + t.Errorf("expected fallback_type query 'context_window', got %q", got) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"model":"gpt-4","fallback_models":["claude-3"],"fallback_type":"context_window"}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newFallbackTestResourceData(t, "gpt-4", []interface{}{"stale-model"}, "context_window") + d.SetId("gpt-4") + + if err := resourceLiteLLMFallbackRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + got := d.Get("fallback_models").([]interface{}) + if !reflect.DeepEqual(got, []interface{}{"claude-3"}) { + t.Fatalf("expected fallback_models [claude-3], got %+v", got) + } + if d.Get("fallback_type").(string) != "context_window" { + t.Fatalf("expected fallback_type 'context_window', got %q", d.Get("fallback_type")) + } + if d.Get("model").(string) != "gpt-4" { + t.Fatalf("expected model 'gpt-4', got %q", d.Get("model")) + } +} + +func TestResourceLiteLLMFallbackRead_404ClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newFallbackTestResourceData(t, "gpt-4", []interface{}{"claude-3"}, "general") + d.SetId("gpt-4") + + if err := resourceLiteLLMFallbackRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared on 404, got %q", d.Id()) + } +} + +func TestResourceLiteLLMFallbackUpdate_SendsChangedModels(t *testing.T) { + var updatePayload map[string]interface{} + mux := http.NewServeMux() + mux.HandleFunc("/fallback", func(w http.ResponseWriter, r *http.Request) { + json.NewDecoder(r.Body).Decode(&updatePayload) + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"model":"gpt-4","fallback_models":["new-model"],"fallback_type":"general","message":"ok"}`)) + }) + mux.Handle("/fallback/gpt-4", fallbackGetHandler(t, "/fallback/gpt-4", FallbackGetResponse{ + Model: "gpt-4", + FallbackModels: []string{"new-model"}, + FallbackType: "general", + })) + srv := httptest.NewServer(mux) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newFallbackTestResourceData(t, "gpt-4", []interface{}{"new-model"}, "general") + d.SetId("gpt-4") + + if err := resourceLiteLLMFallbackUpdate(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if !reflect.DeepEqual(updatePayload["fallback_models"], []interface{}{"new-model"}) { + t.Fatalf("expected updated fallback_models [new-model], got %+v", updatePayload["fallback_models"]) + } +} + +func TestResourceLiteLLMFallbackDelete(t *testing.T) { + var gotMethod, gotPath, gotType string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotMethod = r.Method + gotPath = r.URL.Path + gotType = r.URL.Query().Get("fallback_type") + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"model":"gpt-4","fallback_type":"general","message":"deleted"}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newFallbackTestResourceData(t, "gpt-4", []interface{}{"claude-3"}, "general") + d.SetId("gpt-4") + + if err := resourceLiteLLMFallbackDelete(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if gotMethod != http.MethodDelete || gotPath != "/fallback/gpt-4" || gotType != "general" { + t.Fatalf("expected DELETE /fallback/gpt-4?fallback_type=general, got %s %s?fallback_type=%s", + gotMethod, gotPath, gotType) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared after delete, got %q", d.Id()) + } +} diff --git a/terraform/provider/litellm/resource_guardrail.go b/terraform/provider/litellm/resource_guardrail.go new file mode 100644 index 00000000000..5d8f7a92e13 --- /dev/null +++ b/terraform/provider/litellm/resource_guardrail.go @@ -0,0 +1,255 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "log" + "net/http" + "reflect" + "strings" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const ( + endpointGuardrailCreate = "/guardrails" + endpointGuardrailByID = "/guardrails/%s" + endpointGuardrailInfo = "/guardrails/%s/info" +) + +func resourceLiteLLMGuardrail() *schema.Resource { + return &schema.Resource{ + Create: resourceLiteLLMGuardrailCreate, + Read: resourceLiteLLMGuardrailRead, + Update: resourceLiteLLMGuardrailUpdate, + Delete: resourceLiteLLMGuardrailDelete, + + Importer: &schema.ResourceImporter{StateContext: schema.ImportStatePassthroughContext}, + + Schema: map[string]*schema.Schema{ + "guardrail_name": { + Type: schema.TypeString, + Required: true, + Description: "Human-readable name for the guardrail", + }, + "guardrail": { + Type: schema.TypeString, + Required: true, + Description: "The guardrail integration type (e.g. 'bedrock', 'lakera', 'presidio', 'hide_secrets')", + }, + "mode": { + Type: schema.TypeString, + Required: true, + Description: "When to apply the guardrail: a single value ('pre_call', 'post_call', 'during_call', " + + "'logging_only') or a JSON array of values (e.g. '[\"pre_call\", \"post_call\"]')", + }, + "default_on": { + Type: schema.TypeBool, + Optional: true, + Description: "Whether the guardrail is enabled by default for all requests", + }, + "litellm_params": { + Type: schema.TypeString, + Optional: true, + Sensitive: true, + DiffSuppressFunc: guardrailSuppressJSONDiff, + Description: "JSON string with additional provider-specific litellm_params (may contain API keys)", + }, + "guardrail_info": { + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Additional metadata for the guardrail", + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + }, + } +} + +func guardrailSuppressJSONDiff(k, oldValue, newValue string, d *schema.ResourceData) bool { + var oldParsed, newParsed interface{} + if json.Unmarshal([]byte(oldValue), &oldParsed) != nil || json.Unmarshal([]byte(newValue), &newParsed) != nil { + return false + } + return reflect.DeepEqual(oldParsed, newParsed) +} + +func guardrailParseMode(mode string) interface{} { + if strings.HasPrefix(strings.TrimSpace(mode), "[") { + var modes []string + if err := json.Unmarshal([]byte(mode), &modes); err == nil { + return modes + } + } + return mode +} + +func buildGuardrailData(d *schema.ResourceData, guardrailID string) (map[string]interface{}, error) { + litellmParams := map[string]interface{}{ + "guardrail": d.Get("guardrail").(string), + "mode": guardrailParseMode(d.Get("mode").(string)), + "default_on": d.Get("default_on").(bool), + } + + if raw := d.Get("litellm_params").(string); raw != "" { + var extra map[string]interface{} + if err := json.Unmarshal([]byte(raw), &extra); err != nil { + return nil, fmt.Errorf("litellm_params is not valid JSON: %w", err) + } + for k, v := range extra { + litellmParams[k] = v + } + } + + guardrail := map[string]interface{}{ + "guardrail_name": d.Get("guardrail_name").(string), + "litellm_params": litellmParams, + } + + if guardrailID != "" { + guardrail["guardrail_id"] = guardrailID + } + + if v, ok := d.GetOk("guardrail_info"); ok { + guardrail["guardrail_info"] = v + } + + return map[string]interface{}{"guardrail": guardrail}, nil +} + +type guardrailInfoAPIResponse struct { + GuardrailID string `json:"guardrail_id"` + GuardrailName string `json:"guardrail_name"` + GuardrailInfo map[string]interface{} `json:"guardrail_info"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` +} + +func resourceLiteLLMGuardrailCreate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + guardrailData, err := buildGuardrailData(d, "") + if err != nil { + return err + } + + log.Printf("[DEBUG] Create guardrail request for: %s", d.Get("guardrail_name").(string)) + + resp, err := MakeRequest(client, "POST", endpointGuardrailCreate, guardrailData) + if err != nil { + return fmt.Errorf("error creating guardrail: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "creating guardrail"); err != nil { + return err + } + + var created guardrailInfoAPIResponse + if err := json.NewDecoder(resp.Body).Decode(&created); err != nil { + return fmt.Errorf("error decoding create guardrail response: %w", err) + } + if created.GuardrailID == "" { + return fmt.Errorf("create guardrail response did not contain a guardrail_id") + } + + d.SetId(created.GuardrailID) + log.Printf("[INFO] Guardrail created with ID: %s", created.GuardrailID) + + return resourceLiteLLMGuardrailRead(d, m) +} + +func resourceLiteLLMGuardrailRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Reading guardrail with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf(endpointGuardrailInfo, d.Id()), nil) + if err != nil { + return fmt.Errorf("error reading guardrail: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + log.Printf("[WARN] Guardrail with ID %s not found, removing from state", d.Id()) + d.SetId("") + return nil + } + + if err := handleResponse(resp, "reading guardrail"); err != nil { + return err + } + + var info guardrailInfoAPIResponse + if err := json.NewDecoder(resp.Body).Decode(&info); err != nil { + return fmt.Errorf("error decoding guardrail info response: %w", err) + } + + d.Set("guardrail_name", info.GuardrailName) + d.Set("created_at", info.CreatedAt) + if len(info.GuardrailInfo) > 0 { + d.Set("guardrail_info", guardrailInfoToStringMap(info.GuardrailInfo)) + } + // guardrail, mode, default_on and litellm_params are intentionally not read + // back: the API masks litellm_params values, so state keeps the configured + // values authoritative. + + return nil +} + +func resourceLiteLLMGuardrailUpdate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + guardrailData, err := buildGuardrailData(d, d.Id()) + if err != nil { + return err + } + + log.Printf("[DEBUG] Update guardrail request for ID: %s", d.Id()) + + resp, err := MakeRequest(client, "PUT", fmt.Sprintf(endpointGuardrailByID, d.Id()), guardrailData) + if err != nil { + return fmt.Errorf("error updating guardrail: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "updating guardrail"); err != nil { + return err + } + + log.Printf("[INFO] Successfully updated guardrail with ID: %s", d.Id()) + return resourceLiteLLMGuardrailRead(d, m) +} + +func resourceLiteLLMGuardrailDelete(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Deleting guardrail with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "DELETE", fmt.Sprintf(endpointGuardrailByID, d.Id()), nil) + if err != nil { + return fmt.Errorf("error deleting guardrail: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusNotFound { + if err := handleResponse(resp, "deleting guardrail"); err != nil { + return err + } + } + + log.Printf("[INFO] Successfully deleted guardrail with ID: %s", d.Id()) + d.SetId("") + return nil +} + +func guardrailInfoToStringMap(info map[string]interface{}) map[string]string { + result := make(map[string]string, len(info)) + for k, v := range info { + result[k] = fmt.Sprintf("%v", v) + } + return result +} diff --git a/terraform/provider/litellm/resource_guardrail_test.go b/terraform/provider/litellm/resource_guardrail_test.go new file mode 100644 index 00000000000..d2173f5d223 --- /dev/null +++ b/terraform/provider/litellm/resource_guardrail_test.go @@ -0,0 +1,271 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "reflect" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func newGuardrailTestData(t *testing.T, raw map[string]interface{}) *schema.ResourceData { + t.Helper() + return schema.TestResourceDataRaw(t, resourceLiteLLMGuardrail().Schema, raw) +} + +func guardrailInfoJSON(id, name string) string { + body, _ := json.Marshal(map[string]interface{}{ + "guardrail_id": id, + "guardrail_name": name, + "guardrail_info": map[string]interface{}{"description": "test guardrail"}, + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-02T00:00:00Z", + }) + return string(body) +} + +func TestGuardrailCreate_SendsPayloadAndSetsID(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch { + case r.Method == "POST" && r.URL.Path == "/guardrails": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(guardrailInfoJSON("gid-123", "guard1"))) + case r.Method == "GET" && r.URL.Path == "/guardrails/gid-123/info": + w.Write([]byte(guardrailInfoJSON("gid-123", "guard1"))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newGuardrailTestData(t, map[string]interface{}{ + "guardrail_name": "guard1", + "guardrail": "bedrock", + "mode": "pre_call", + "default_on": true, + "litellm_params": `{"api_key": "sk-123", "guardrailIdentifier": "abc"}`, + "guardrail_info": map[string]interface{}{"description": "test guardrail"}, + }) + + if err := resourceLiteLLMGuardrailCreate(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "gid-123" { + t.Fatalf("expected ID 'gid-123', got %q", d.Id()) + } + + guardrail, ok := createPayload["guardrail"].(map[string]interface{}) + if !ok { + t.Fatalf("expected payload wrapped in 'guardrail' key, got: %v", createPayload) + } + if guardrail["guardrail_name"] != "guard1" { + t.Errorf("expected guardrail_name 'guard1', got %v", guardrail["guardrail_name"]) + } + params, ok := guardrail["litellm_params"].(map[string]interface{}) + if !ok { + t.Fatalf("expected litellm_params object, got: %v", guardrail["litellm_params"]) + } + if params["guardrail"] != "bedrock" || params["mode"] != "pre_call" || params["default_on"] != true { + t.Errorf("unexpected base litellm_params: %v", params) + } + if params["api_key"] != "sk-123" || params["guardrailIdentifier"] != "abc" { + t.Errorf("expected merged extra litellm_params, got: %v", params) + } + info, ok := guardrail["guardrail_info"].(map[string]interface{}) + if !ok || info["description"] != "test guardrail" { + t.Errorf("expected guardrail_info to be sent, got: %v", guardrail["guardrail_info"]) + } +} + +func TestGuardrailCreate_ModeJSONArray(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == "POST" { + json.NewDecoder(r.Body).Decode(&createPayload) + } + w.Write([]byte(guardrailInfoJSON("gid-456", "guard2"))) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newGuardrailTestData(t, map[string]interface{}{ + "guardrail_name": "guard2", + "guardrail": "lakera", + "mode": `["pre_call", "post_call"]`, + }) + + if err := resourceLiteLLMGuardrailCreate(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + + params := createPayload["guardrail"].(map[string]interface{})["litellm_params"].(map[string]interface{}) + mode, ok := params["mode"].([]interface{}) + if !ok { + t.Fatalf("expected mode to be a JSON array, got: %v", params["mode"]) + } + if !reflect.DeepEqual(mode, []interface{}{"pre_call", "post_call"}) { + t.Errorf("unexpected mode array: %v", mode) + } +} + +func TestGuardrailCreate_InvalidLitellmParamsJSON(t *testing.T) { + client := NewClient("http://unused.invalid", "test-key", true) + d := newGuardrailTestData(t, map[string]interface{}{ + "guardrail_name": "guard1", + "guardrail": "bedrock", + "mode": "pre_call", + "litellm_params": "{not json", + }) + + if err := resourceLiteLLMGuardrailCreate(d, client); err == nil { + t.Fatal("expected error for invalid litellm_params JSON, got nil") + } +} + +func TestGuardrailRead_MapsFieldsAndKeepsConfiguredParams(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" || r.URL.Path != "/guardrails/gid-1/info" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(guardrailInfoJSON("gid-1", "renamed-guard"))) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newGuardrailTestData(t, map[string]interface{}{ + "guardrail_name": "old-name", + "guardrail": "bedrock", + "mode": "pre_call", + "litellm_params": `{"api_key": "sk-123"}`, + }) + d.SetId("gid-1") + + if err := resourceLiteLLMGuardrailRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if got := d.Get("guardrail_name").(string); got != "renamed-guard" { + t.Errorf("expected guardrail_name 'renamed-guard', got %q", got) + } + if got := d.Get("created_at").(string); got != "2026-01-01T00:00:00Z" { + t.Errorf("expected created_at to be set, got %q", got) + } + if got := d.Get("litellm_params").(string); got != `{"api_key": "sk-123"}` { + t.Errorf("expected configured litellm_params to stay authoritative, got %q", got) + } + info := d.Get("guardrail_info").(map[string]interface{}) + if info["description"] != "test guardrail" { + t.Errorf("expected guardrail_info from API, got: %v", info) + } +} + +func TestGuardrailRead_404ClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newGuardrailTestData(t, map[string]interface{}{ + "guardrail_name": "guard1", + "guardrail": "bedrock", + "mode": "pre_call", + }) + d.SetId("gid-gone") + + if err := resourceLiteLLMGuardrailRead(d, client); err != nil { + t.Fatalf("expected nil error on 404, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID to be cleared on 404, got %q", d.Id()) + } +} + +func TestGuardrailUpdate_SendsPUTToGuardrailEndpoint(t *testing.T) { + var updateMethod, updatePath string + var updatePayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == "PUT" { + updateMethod, updatePath = r.Method, r.URL.Path + json.NewDecoder(r.Body).Decode(&updatePayload) + } + w.Write([]byte(guardrailInfoJSON("gid-1", "new-name"))) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newGuardrailTestData(t, map[string]interface{}{ + "guardrail_name": "new-name", + "guardrail": "bedrock", + "mode": "post_call", + }) + d.SetId("gid-1") + + if err := resourceLiteLLMGuardrailUpdate(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if updateMethod != "PUT" || updatePath != "/guardrails/gid-1" { + t.Fatalf("expected PUT /guardrails/gid-1, got %s %s", updateMethod, updatePath) + } + guardrail := updatePayload["guardrail"].(map[string]interface{}) + if guardrail["guardrail_name"] != "new-name" { + t.Errorf("expected updated guardrail_name, got %v", guardrail["guardrail_name"]) + } + if guardrail["guardrail_id"] != "gid-1" { + t.Errorf("expected guardrail_id in update payload, got %v", guardrail["guardrail_id"]) + } + params := guardrail["litellm_params"].(map[string]interface{}) + if params["mode"] != "post_call" { + t.Errorf("expected updated mode 'post_call', got %v", params["mode"]) + } +} + +func TestGuardrailDelete_CallsDeleteEndpoint(t *testing.T) { + var deleteMethod, deletePath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + deleteMethod, deletePath = r.Method, r.URL.Path + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"message": "deleted"}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newGuardrailTestData(t, map[string]interface{}{ + "guardrail_name": "guard1", + "guardrail": "bedrock", + "mode": "pre_call", + }) + d.SetId("gid-1") + + if err := resourceLiteLLMGuardrailDelete(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if deleteMethod != "DELETE" || deletePath != "/guardrails/gid-1" { + t.Fatalf("expected DELETE /guardrails/gid-1, got %s %s", deleteMethod, deletePath) + } + if d.Id() != "" { + t.Fatalf("expected ID to be cleared after delete, got %q", d.Id()) + } +} + +func TestGuardrailSuppressJSONDiff(t *testing.T) { + if !guardrailSuppressJSONDiff("", `{"a": 1, "b": "x"}`, `{"b":"x","a":1}`, nil) { + t.Error("expected semantically equal JSON to be suppressed") + } + if guardrailSuppressJSONDiff("", `{"a": 1}`, `{"a": 2}`, nil) { + t.Error("expected different JSON not to be suppressed") + } + if guardrailSuppressJSONDiff("", "", `{"a": 1}`, nil) { + t.Error("expected empty old value not to be suppressed") + } +} diff --git a/terraform/provider/litellm/resource_jwt_key_mapping.go b/terraform/provider/litellm/resource_jwt_key_mapping.go new file mode 100644 index 00000000000..e606e865737 --- /dev/null +++ b/terraform/provider/litellm/resource_jwt_key_mapping.go @@ -0,0 +1,70 @@ +package litellm + +import ( + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func resourceLiteLLMJWTKeyMapping() *schema.Resource { + return &schema.Resource{ + Create: resourceLiteLLMJWTKeyMappingCreate, + Read: resourceLiteLLMJWTKeyMappingRead, + Update: resourceLiteLLMJWTKeyMappingUpdate, + Delete: resourceLiteLLMJWTKeyMappingDelete, + + Importer: &schema.ResourceImporter{ + StateContext: schema.ImportStatePassthroughContext, + }, + + Schema: map[string]*schema.Schema{ + "jwt_claim_name": { + Type: schema.TypeString, + Required: true, + ForceNew: true, + Description: "Name of the JWT claim to match on, for example client_id, azp or sub. Must match virtual_key_claim_field in the proxy JWT config", + }, + "jwt_claim_value": { + Type: schema.TypeString, + Required: true, + ForceNew: true, + Description: "Value of the claim identifying the JWT client. Unique together with jwt_claim_name", + }, + "key": { + Type: schema.TypeString, + Required: true, + Sensitive: true, + Description: "The virtual key this claim value maps to. The proxy stores only a hash of it and never returns it, so drift on this attribute cannot be detected and Terraform tracks the configured value", + }, + "description": { + Type: schema.TypeString, + Optional: true, + Description: "Description of the mapping", + }, + "is_active": { + Type: schema.TypeBool, + Optional: true, + Default: true, + Description: "Whether the mapping is active. Inactive mappings are ignored during JWT auth", + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + Description: "Timestamp when the mapping was created", + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + Description: "Timestamp when the mapping was last updated", + }, + "created_by": { + Type: schema.TypeString, + Computed: true, + Description: "User who created the mapping", + }, + "updated_by": { + Type: schema.TypeString, + Computed: true, + Description: "User who last updated the mapping", + }, + }, + } +} diff --git a/terraform/provider/litellm/resource_jwt_key_mapping_crud.go b/terraform/provider/litellm/resource_jwt_key_mapping_crud.go new file mode 100644 index 00000000000..725235305f6 --- /dev/null +++ b/terraform/provider/litellm/resource_jwt_key_mapping_crud.go @@ -0,0 +1,186 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "net/url" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const jwtKeyMappingNotFound = "jwt_key_mapping_not_found" + +func resourceLiteLLMJWTKeyMappingCreate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + createRequest := JWTKeyMappingRequest{ + JWTClaimName: d.Get("jwt_claim_name").(string), + JWTClaimValue: d.Get("jwt_claim_value").(string), + Key: d.Get("key").(string), + Description: d.Get("description").(string), + } + + resp, err := MakeRequest(client, "POST", "/jwt/key/mapping/new", createRequest) + if err != nil { + return fmt.Errorf("failed to create JWT key mapping: %w", err) + } + defer resp.Body.Close() + + var mapping JWTKeyMappingResponse + if err := handleJWTKeyMappingAPIResponse(resp, &mapping, client); err != nil { + return fmt.Errorf("failed to create JWT key mapping: %w", err) + } + + if mapping.ID == "" { + return fmt.Errorf("failed to create JWT key mapping: the proxy returned no mapping id") + } + + d.SetId(mapping.ID) + + // The create endpoint has no is_active field and always activates the + // mapping, so a JWT client matching this claim can authenticate during + // the gap before the deactivation call below runs. If deactivation + // itself fails, delete the mapping rather than leaving it active and + // unmanaged indefinitely. + if !d.Get("is_active").(bool) { + if err := updateJWTKeyMapping(d, client); err != nil { + if deleteErr := deleteJWTKeyMapping(mapping.ID, client); deleteErr != nil { + return fmt.Errorf( + "JWT key mapping %s was created active and could not be deactivated (%v); it also could not be deleted and remains active on the proxy, remove it manually via POST /jwt/key/mapping/delete: %v", + mapping.ID, err, deleteErr, + ) + } + d.SetId("") + return fmt.Errorf("JWT key mapping was created active but could not be deactivated, so it was deleted instead: %w", err) + } + } + + return resourceLiteLLMJWTKeyMappingRead(d, m) +} + +func resourceLiteLLMJWTKeyMappingRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("/jwt/key/mapping/info?id=%s", url.QueryEscape(d.Id())), nil) + if err != nil { + return fmt.Errorf("failed to read JWT key mapping: %w", err) + } + defer resp.Body.Close() + + var mapping JWTKeyMappingResponse + if err := handleJWTKeyMappingAPIResponse(resp, &mapping, client); err != nil { + if err.Error() == jwtKeyMappingNotFound { + d.SetId("") + return nil + } + return fmt.Errorf("failed to read JWT key mapping: %w", err) + } + + d.SetId(mapping.ID) + d.Set("jwt_claim_name", mapping.JWTClaimName) + d.Set("jwt_claim_value", mapping.JWTClaimValue) + d.Set("description", mapping.Description) + d.Set("is_active", mapping.IsActive) + d.Set("created_at", mapping.CreatedAt) + d.Set("updated_at", mapping.UpdatedAt) + d.Set("created_by", mapping.CreatedBy) + d.Set("updated_by", mapping.UpdatedBy) + + return nil +} + +func resourceLiteLLMJWTKeyMappingUpdate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + oldKey, _ := d.GetChange("key") + oldDescription, _ := d.GetChange("description") + oldIsActive, _ := d.GetChange("is_active") + + if err := updateJWTKeyMapping(d, client); err != nil { + // The update is a single atomic API call: on failure nothing changed + // server-side. Revert every field the update could have changed before + // attempting to resync, so a failed refresh can't leave the rejected + // values persisted into state. + d.Set("key", oldKey) + d.Set("description", oldDescription) + d.Set("is_active", oldIsActive) + if readErr := resourceLiteLLMJWTKeyMappingRead(d, m); readErr != nil { + return fmt.Errorf("failed to update JWT key mapping: %w (and failed to refresh state afterward: %v)", err, readErr) + } + return fmt.Errorf("failed to update JWT key mapping: %w", err) + } + + return resourceLiteLLMJWTKeyMappingRead(d, m) +} + +func resourceLiteLLMJWTKeyMappingDelete(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + if err := deleteJWTKeyMapping(d.Id(), client); err != nil { + return fmt.Errorf("failed to delete JWT key mapping: %w", err) + } + + d.SetId("") + return nil +} + +func deleteJWTKeyMapping(id string, client *Client) error { + resp, err := MakeRequest(client, "POST", "/jwt/key/mapping/delete", JWTKeyMappingDeleteRequest{ID: id}) + if err != nil { + return err + } + defer resp.Body.Close() + + if err := handleJWTKeyMappingAPIResponse(resp, nil, client); err != nil { + if err.Error() != jwtKeyMappingNotFound { + return err + } + } + + return nil +} + +func updateJWTKeyMapping(d *schema.ResourceData, client *Client) error { + updateRequest := JWTKeyMappingUpdateRequest{ + ID: d.Id(), + Key: d.Get("key").(string), + Description: d.Get("description").(string), + IsActive: d.Get("is_active").(bool), + } + + resp, err := MakeRequest(client, "POST", "/jwt/key/mapping/update", updateRequest) + if err != nil { + return err + } + defer resp.Body.Close() + + return handleJWTKeyMappingAPIResponse(resp, nil, client) +} + +func handleJWTKeyMappingAPIResponse(resp *http.Response, result interface{}, client *Client) error { + bodyBytes, err := io.ReadAll(resp.Body) + if err != nil { + return fmt.Errorf("failed to read response body: %v", err) + } + + if resp.StatusCode == http.StatusNotFound { + return fmt.Errorf(jwtKeyMappingNotFound) + } + + if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusCreated { + return fmt.Errorf("API request failed: Status: %s, Response: %s", + resp.Status, client.redactSensitiveData(string(bodyBytes))) + } + + if result == nil { + return nil + } + + if err := json.Unmarshal(bodyBytes, result); err != nil { + return fmt.Errorf("failed to parse response: %v", err) + } + + return nil +} diff --git a/terraform/provider/litellm/resource_jwt_key_mapping_crud_test.go b/terraform/provider/litellm/resource_jwt_key_mapping_crud_test.go new file mode 100644 index 00000000000..8007d1d4e08 --- /dev/null +++ b/terraform/provider/litellm/resource_jwt_key_mapping_crud_test.go @@ -0,0 +1,630 @@ +package litellm + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-sdk/v2/terraform" +) + +// resourceDataWithChange builds a ResourceData carrying a real diff between +// prior state and new config, so d.GetChange reflects true old/new values. +// schema.TestResourceDataRaw diffs against a nil prior state, which collapses +// GetChange's old side to the zero value and can't exercise this. +func resourceDataWithChange(t *testing.T, oldAttrs map[string]string, newRaw map[string]interface{}) *schema.ResourceData { + t.Helper() + + sm := schema.InternalMap(resourceLiteLLMJWTKeyMapping().Schema) + state := &terraform.InstanceState{ID: oldAttrs["id"], Attributes: oldAttrs} + config := terraform.NewResourceConfigRaw(newRaw) + + diff, err := sm.Diff(context.Background(), state, config, nil, nil, true) + if err != nil { + t.Fatalf("diff: %v", err) + } + d, err := sm.Data(state, diff) + if err != nil { + t.Fatalf("data: %v", err) + } + return d +} + +type jwtKeyMappingCall struct { + Method string + Path string + Query string + Body map[string]interface{} +} + +func jwtKeyMappingTestServer(t *testing.T, mapping JWTKeyMappingResponse) (*httptest.Server, *[]jwtKeyMappingCall) { + t.Helper() + + calls := make([]jwtKeyMappingCall, 0) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body := map[string]interface{}{} + if r.Body != nil { + _ = json.NewDecoder(r.Body).Decode(&body) + } + calls = append(calls, jwtKeyMappingCall{Method: r.Method, Path: r.URL.Path, Query: r.URL.RawQuery, Body: body}) + + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/jwt/key/mapping/delete": + _ = json.NewEncoder(w).Encode(map[string]string{"status": "success"}) + default: + _ = json.NewEncoder(w).Encode(mapping) + } + })) + + return srv, &calls +} + +func jwtKeyMappingFixture() JWTKeyMappingResponse { + return JWTKeyMappingResponse{ + ID: "map-abc-123", + JWTClaimName: "client_id", + JWTClaimValue: "dev-alice", + Description: "dev-alice", + IsActive: true, + CreatedAt: "2026-08-06T10:00:00Z", + UpdatedAt: "2026-08-06T11:00:00Z", + CreatedBy: "admin", + UpdatedBy: "admin", + } +} + +func TestJWTKeyMappingCreateSendsClaimAndKey(t *testing.T) { + srv, calls := jwtKeyMappingTestServer(t, jwtKeyMappingFixture()) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-abc123", + "description": "dev-alice", + "is_active": true, + }) + + if err := resourceLiteLLMJWTKeyMappingCreate(d, client); err != nil { + t.Fatalf("create failed: %v", err) + } + + if d.Id() != "map-abc-123" { + t.Fatalf("expected id from the API response, got %q", d.Id()) + } + + create := (*calls)[0] + if create.Method != "POST" || create.Path != "/jwt/key/mapping/new" { + t.Fatalf("expected POST /jwt/key/mapping/new, got %s %s", create.Method, create.Path) + } + if create.Body["jwt_claim_name"] != "client_id" || create.Body["jwt_claim_value"] != "dev-alice" { + t.Fatalf("claim fields not sent: %v", create.Body) + } + if create.Body["key"] != "sk-abc123" { + t.Fatalf("virtual key not sent: %v", create.Body["key"]) + } + if create.Body["description"] != "dev-alice" { + t.Fatalf("description not sent: %v", create.Body["description"]) + } + if _, sent := create.Body["is_active"]; sent { + t.Fatalf("is_active is not accepted by /jwt/key/mapping/new but was sent: %v", create.Body) + } + + for _, call := range (*calls)[1:] { + if call.Path == "/jwt/key/mapping/update" { + t.Fatalf("an active mapping must not trigger a follow-up update") + } + } +} + +func TestJWTKeyMappingCreateOmitsEmptyDescription(t *testing.T) { + srv, calls := jwtKeyMappingTestServer(t, jwtKeyMappingFixture()) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-abc123", + "is_active": true, + }) + + if err := resourceLiteLLMJWTKeyMappingCreate(d, client); err != nil { + t.Fatalf("create failed: %v", err) + } + + if _, sent := (*calls)[0].Body["description"]; sent { + t.Fatalf("unset description should be omitted: %v", (*calls)[0].Body) + } +} + +func TestJWTKeyMappingCreateDeactivatesWhenNotActive(t *testing.T) { + mapping := jwtKeyMappingFixture() + mapping.IsActive = false + srv, calls := jwtKeyMappingTestServer(t, mapping) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-abc123", + "is_active": false, + }) + + if err := resourceLiteLLMJWTKeyMappingCreate(d, client); err != nil { + t.Fatalf("create failed: %v", err) + } + + var update *jwtKeyMappingCall + for i := range *calls { + if (*calls)[i].Path == "/jwt/key/mapping/update" { + update = &(*calls)[i] + break + } + } + if update == nil { + t.Fatal("expected a follow-up update, since the create endpoint always starts a mapping active") + } + if update.Body["id"] != "map-abc-123" { + t.Fatalf("update must target the new mapping, got %v", update.Body["id"]) + } + if update.Body["is_active"] != false { + t.Fatalf("expected is_active false in the follow-up update, got %v", update.Body["is_active"]) + } + if d.Get("is_active").(bool) { + t.Fatal("state should reflect the inactive mapping after create") + } +} + +func TestJWTKeyMappingCreateDeletesMappingWhenDeactivationFails(t *testing.T) { + // Regression test: the create endpoint has no is_active field and always + // activates the mapping, so a failed deactivation used to leave that + // mapping active and unmanaged indefinitely. It must be deleted instead. + calls := make([]jwtKeyMappingCall, 0) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body := map[string]interface{}{} + if r.Body != nil { + _ = json.NewDecoder(r.Body).Decode(&body) + } + calls = append(calls, jwtKeyMappingCall{Method: r.Method, Path: r.URL.Path, Query: r.URL.RawQuery, Body: body}) + + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/jwt/key/mapping/new": + _ = json.NewEncoder(w).Encode(jwtKeyMappingFixture()) + case "/jwt/key/mapping/update": + w.WriteHeader(http.StatusInternalServerError) + _ = json.NewEncoder(w).Encode(map[string]string{"detail": "proxy unavailable"}) + case "/jwt/key/mapping/delete": + _ = json.NewEncoder(w).Encode(map[string]string{"status": "success"}) + default: + t.Fatalf("unexpected request to %s", r.URL.Path) + } + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-abc123", + "is_active": false, + }) + + err := resourceLiteLLMJWTKeyMappingCreate(d, client) + if err == nil { + t.Fatal("expected the failed deactivation to surface as an error") + } + if !strings.Contains(err.Error(), "deleted instead") { + t.Fatalf("expected the error to explain the mapping was deleted, got %v", err) + } + + deleteCalls := 0 + for _, c := range calls { + if c.Path == "/jwt/key/mapping/delete" { + deleteCalls++ + if c.Body["id"] != "map-abc-123" { + t.Fatalf("delete must target the mapping that could not be deactivated, got %v", c.Body["id"]) + } + } + } + if deleteCalls != 1 { + t.Fatalf("expected exactly one cleanup delete call, got %d", deleteCalls) + } + + if d.Id() != "" { + t.Fatalf("a successfully deleted mapping must not remain in state, got id %q", d.Id()) + } +} + +func TestJWTKeyMappingCreateReportsWhenDeactivationAndDeleteBothFail(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/jwt/key/mapping/new": + _ = json.NewEncoder(w).Encode(jwtKeyMappingFixture()) + default: + w.WriteHeader(http.StatusInternalServerError) + _ = json.NewEncoder(w).Encode(map[string]string{"detail": "proxy unavailable"}) + } + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-abc123", + "is_active": false, + }) + + err := resourceLiteLLMJWTKeyMappingCreate(d, client) + if err == nil { + t.Fatal("expected an error when both deactivation and the cleanup delete fail") + } + if !strings.Contains(err.Error(), "remove it manually") { + t.Fatalf("expected the error to demand manual cleanup, got %v", err) + } + + // The mapping is still active on the proxy since neither call succeeded, so + // the id must stay in state: the next apply taints and retries the delete, + // rather than Terraform losing track of a live, active mapping entirely. + if d.Id() != "map-abc-123" { + t.Fatalf("expected the id to remain in state so a retry can find it, got %q", d.Id()) + } +} + +func TestJWTKeyMappingUpdateRevertsDescriptionAndIsActiveWhenTheRecoveryReadAlsoFails(t *testing.T) { + // Regression test: on a failed update, only `key` was being reverted + // before Read ran. If Read itself then failed too (network blip, proxy + // hiccup), description/is_active kept the rejected, never-applied values, + // and Terraform could persist them as if the update had succeeded. + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/jwt/key/mapping/update": + w.WriteHeader(http.StatusBadRequest) + _ = json.NewEncoder(w).Encode(map[string]string{"detail": "rejected"}) + case "/jwt/key/mapping/info": + w.WriteHeader(http.StatusInternalServerError) + _ = json.NewEncoder(w).Encode(map[string]string{"detail": "proxy unavailable"}) + default: + t.Fatalf("unexpected request to %s", r.URL.Path) + } + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + + d := resourceDataWithChange(t, + map[string]string{ + "id": "map-abc-123", + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-old-key-0000000000", + "description": "old description", + "is_active": "true", + }, + map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-old-key-0000000000", + "description": "attempted new description", + "is_active": false, + }, + ) + d.SetId("map-abc-123") + + err := resourceLiteLLMJWTKeyMappingUpdate(d, client) + if err == nil { + t.Fatal("expected the update failure to surface as an error") + } + if !strings.Contains(err.Error(), "failed to refresh state afterward") { + t.Fatalf("expected the error to mention the failed recovery read, got %v", err) + } + + if d.Get("description").(string) != "old description" { + t.Fatalf("a rejected description must not survive when the recovery read also fails, got %q", d.Get("description").(string)) + } + if d.Get("is_active").(bool) != true { + t.Fatalf("a rejected is_active must not survive when the recovery read also fails, got %v", d.Get("is_active").(bool)) + } +} + +func TestJWTKeyMappingReadPopulatesStateAndKeepsKey(t *testing.T) { + srv, calls := jwtKeyMappingTestServer(t, jwtKeyMappingFixture()) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-configured-value", + }) + d.SetId("map-abc-123") + + if err := resourceLiteLLMJWTKeyMappingRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + + read := (*calls)[0] + if read.Method != "GET" || read.Path != "/jwt/key/mapping/info" { + t.Fatalf("expected GET /jwt/key/mapping/info, got %s %s", read.Method, read.Path) + } + if read.Query != "id=map-abc-123" { + t.Fatalf("expected the mapping id in the query, got %q", read.Query) + } + + if d.Get("jwt_claim_value").(string) != "dev-alice" { + t.Fatalf("claim value not populated: %q", d.Get("jwt_claim_value").(string)) + } + if d.Get("description").(string) != "dev-alice" { + t.Fatalf("description not populated: %q", d.Get("description").(string)) + } + if !d.Get("is_active").(bool) { + t.Fatal("is_active not populated") + } + if d.Get("created_at").(string) != "2026-08-06T10:00:00Z" || d.Get("created_by").(string) != "admin" { + t.Fatalf("computed audit fields not populated: %v", d.State().Attributes) + } + if d.Get("key").(string) != "sk-configured-value" { + t.Fatalf("the API never returns the key, so the configured value must survive a read, got %q", d.Get("key").(string)) + } +} + +func TestJWTKeyMappingReadClearsIDWhenMappingIsGone(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotFound) + _ = json.NewEncoder(w).Encode(map[string]string{"detail": "Mapping not found"}) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-abc123", + }) + d.SetId("map-gone") + + if err := resourceLiteLLMJWTKeyMappingRead(d, client); err != nil { + t.Fatalf("a deleted mapping must not fail the read: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected the id to be cleared so Terraform plans a recreate, got %q", d.Id()) + } +} + +func TestJWTKeyMappingUpdateClearsDescriptionAndSendsKey(t *testing.T) { + mapping := jwtKeyMappingFixture() + mapping.Description = "" + srv, calls := jwtKeyMappingTestServer(t, mapping) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-rotated", + "is_active": true, + }) + d.SetId("map-abc-123") + + if err := resourceLiteLLMJWTKeyMappingUpdate(d, client); err != nil { + t.Fatalf("update failed: %v", err) + } + + update := (*calls)[0] + if update.Method != "POST" || update.Path != "/jwt/key/mapping/update" { + t.Fatalf("expected POST /jwt/key/mapping/update, got %s %s", update.Method, update.Path) + } + if update.Body["id"] != "map-abc-123" { + t.Fatalf("update must carry the mapping id, got %v", update.Body["id"]) + } + if update.Body["key"] != "sk-rotated" { + t.Fatalf("rotated key not sent: %v", update.Body["key"]) + } + description, sent := update.Body["description"] + if !sent || description != "" { + t.Fatalf("a dropped description must be sent as an empty string, since the proxy ignores absent fields: %v", update.Body) + } + if d.Get("description").(string) != "" { + t.Fatalf("description should be cleared in state, got %q", d.Get("description").(string)) + } +} + +func TestJWTKeyMappingUpdateRevertsKeyOnFailureAndResyncsRest(t *testing.T) { + // Regression test for a live-verified bug: Terraform's classic SDKv2 CRUD + // model persists ResourceData's diff-applied (attempted) values to state + // even when the callback returns an error, unless the provider reverts + // them explicitly. Confirmed live: a rejected key rotation left the new, + // never-applied key in `terraform state pull` while the proxy kept the + // old one, so the next plan falsely reported convergence. + calls := make([]jwtKeyMappingCall, 0) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body := map[string]interface{}{} + if r.Body != nil { + _ = json.NewDecoder(r.Body).Decode(&body) + } + calls = append(calls, jwtKeyMappingCall{Method: r.Method, Path: r.URL.Path, Query: r.URL.RawQuery, Body: body}) + + w.Header().Set("Content-Type", "application/json") + switch r.URL.Path { + case "/jwt/key/mapping/update": + w.WriteHeader(http.StatusBadRequest) + _ = json.NewEncoder(w).Encode(map[string]string{ + "detail": "The provided key does not match an existing virtual key.", + }) + case "/jwt/key/mapping/info": + // Server truth: unchanged, since the rejected update above never applied. + _ = json.NewEncoder(w).Encode(jwtKeyMappingFixture()) + default: + t.Fatalf("unexpected request to %s", r.URL.Path) + } + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + + d := resourceDataWithChange(t, + map[string]string{ + "id": "map-abc-123", + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-old-key-0000000000", + "description": "dev-alice", + "is_active": "true", + }, + map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-rejected-new-key-00", + "description": "attempted new description", + "is_active": false, + }, + ) + d.SetId("map-abc-123") + + err := resourceLiteLLMJWTKeyMappingUpdate(d, client) + if err == nil { + t.Fatal("expected the rejected key to fail the update") + } + if !strings.Contains(err.Error(), "does not match an existing virtual key") { + t.Fatalf("expected the proxy's rejection reason in the error, got %v", err) + } + + if d.Get("key").(string) != "sk-old-key-0000000000" { + t.Fatalf("a failed update must not persist the rejected key into state, got %q", d.Get("key").(string)) + } + if d.Get("description").(string) != "dev-alice" { + t.Fatalf("a failed update must resync description from the server, got %q", d.Get("description").(string)) + } + if d.Get("is_active").(bool) != true { + t.Fatalf("a failed update must resync is_active from the server, got %v", d.Get("is_active").(bool)) + } + + readCalls := 0 + for _, c := range calls { + if c.Path == "/jwt/key/mapping/info" { + readCalls++ + } + } + if readCalls != 1 { + t.Fatalf("expected exactly one read to resync state after the failed update, got %d", readCalls) + } +} + +func TestJWTKeyMappingUpdateOmitsMissingKeyRatherThanBlankingIt(t *testing.T) { + srv, calls := jwtKeyMappingTestServer(t, jwtKeyMappingFixture()) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "description": "dev-alice", + "is_active": true, + }) + d.SetId("map-abc-123") + + if err := resourceLiteLLMJWTKeyMappingUpdate(d, client); err != nil { + t.Fatalf("update failed: %v", err) + } + + if _, sent := (*calls)[0].Body["key"]; sent { + t.Fatalf("a missing key must be omitted rather than blanking the mapping token: %v", (*calls)[0].Body) + } +} + +func TestJWTKeyMappingDeleteToleratesMissingMapping(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusNotFound) + _ = json.NewEncoder(w).Encode(map[string]string{"detail": "Mapping not found"}) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-abc123", + }) + d.SetId("map-already-gone") + + if err := resourceLiteLLMJWTKeyMappingDelete(d, client); err != nil { + t.Fatalf("deleting an already deleted mapping must succeed: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected the id to be cleared after delete, got %q", d.Id()) + } +} + +func TestJWTKeyMappingCreateSurfacesDuplicateClaimError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusConflict) + _ = json.NewEncoder(w).Encode(map[string]string{ + "detail": "A mapping for claim 'client_id' = 'dev-alice' already exists.", + }) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-abc123", + "is_active": true, + }) + + err := resourceLiteLLMJWTKeyMappingCreate(d, client) + if err == nil { + t.Fatal("expected a duplicate claim pair to fail") + } + if !strings.Contains(err.Error(), "already exists") { + t.Fatalf("the proxy explanation must reach the user, got %v", err) + } + if d.Id() != "" { + t.Fatalf("no id should be recorded for a failed create, got %q", d.Id()) + } +} + +func TestJWTKeyMappingCreateDoesNotLeakKeyInErrors(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + _ = json.NewEncoder(w).Encode(map[string]string{ + "key": "sk-super-secret", + "detail": "The provided key does not match an existing virtual key.", + }) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMJWTKeyMapping().Schema, map[string]interface{}{ + "jwt_claim_name": "client_id", + "jwt_claim_value": "dev-alice", + "key": "sk-super-secret", + "is_active": true, + }) + + err := resourceLiteLLMJWTKeyMappingCreate(d, client) + if err == nil { + t.Fatal("expected an unknown virtual key to fail") + } + if !strings.Contains(err.Error(), "does not match an existing virtual key") { + t.Fatalf("the proxy explanation must reach the user, got %v", err) + } + if strings.Contains(err.Error(), "sk-super-secret") { + t.Fatalf("the virtual key must be redacted in errors, got %v", err) + } +} diff --git a/terraform/provider/litellm/resource_key.go b/terraform/provider/litellm/resource_key.go index 5c80198cf6a..0d8674f2d4c 100644 --- a/terraform/provider/litellm/resource_key.go +++ b/terraform/provider/litellm/resource_key.go @@ -4,6 +4,7 @@ import ( "context" "fmt" + "github.com/hashicorp/go-cty/cty" "github.com/hashicorp/terraform-plugin-sdk/v2/diag" "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" ) @@ -136,6 +137,49 @@ func resourceKey() *schema.Resource { Type: schema.TypeFloat, Computed: true, }, + "budget_id": { + Type: schema.TypeString, + Optional: true, + }, + "enforced_params": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "allowed_routes": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "allowed_passthrough_routes": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "rpm_limit_type": { + Type: schema.TypeString, + Optional: true, + Description: "One of 'guaranteed_throughput', 'best_effort_throughput' or 'dynamic'", + }, + "tpm_limit_type": { + Type: schema.TypeString, + Optional: true, + Description: "One of 'guaranteed_throughput', 'best_effort_throughput' or 'dynamic'", + }, + "prompts": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "organization_id": { + Type: schema.TypeString, + Optional: true, + }, + "project_id": { + Type: schema.TypeString, + Optional: true, + ForceNew: true, + }, }, } } @@ -145,6 +189,14 @@ func resourceKeyCreate(ctx context.Context, d *schema.ResourceData, m interface{ key := &Key{} mapResourceDataToKey(d, key) + // A config-supplied key value becomes the key itself; when absent the + // proxy generates one. Write-only attributes are invisible to d.Get in + // real Terraform runs, so read the raw config first. + if raw, err := d.GetRawConfigAt(cty.GetAttrPath("key")); err == nil && !raw.IsNull() && raw.Type() == cty.String && raw.AsString() != "" { + key.Key = raw.AsString() + } else if v := d.Get("key").(string); v != "" { + key.Key = v + } createdKey, err := c.CreateKey(key) if err != nil { @@ -239,6 +291,15 @@ func mapResourceDataToKey(d *schema.ResourceData, key *Key) { key.Guardrails = expandStringList(d.Get("guardrails").([]interface{})) key.Blocked = d.Get("blocked").(bool) key.Tags = expandStringList(d.Get("tags").([]interface{})) + key.BudgetID = d.Get("budget_id").(string) + key.EnforcedParams = expandStringList(d.Get("enforced_params").([]interface{})) + key.AllowedRoutes = expandStringList(d.Get("allowed_routes").([]interface{})) + key.AllowedPassthroughRoutes = expandStringList(d.Get("allowed_passthrough_routes").([]interface{})) + key.RPMLimitType = d.Get("rpm_limit_type").(string) + key.TPMLimitType = d.Get("tpm_limit_type").(string) + key.Prompts = expandStringList(d.Get("prompts").([]interface{})) + key.OrganizationID = d.Get("organization_id").(string) + key.ProjectID = d.Get("project_id").(string) } func mapKeyToResourceData(d *schema.ResourceData, key *Key) { @@ -316,4 +377,31 @@ func mapKeyToResourceData(d *schema.ResourceData, key *Key) { if key.Spend != 0 { d.Set("spend", key.Spend) } + if key.BudgetID != "" { + d.Set("budget_id", key.BudgetID) + } + if len(key.EnforcedParams) > 0 { + d.Set("enforced_params", key.EnforcedParams) + } + if len(key.AllowedRoutes) > 0 { + d.Set("allowed_routes", key.AllowedRoutes) + } + if len(key.AllowedPassthroughRoutes) > 0 { + d.Set("allowed_passthrough_routes", key.AllowedPassthroughRoutes) + } + if key.RPMLimitType != "" { + d.Set("rpm_limit_type", key.RPMLimitType) + } + if key.TPMLimitType != "" { + d.Set("tpm_limit_type", key.TPMLimitType) + } + if len(key.Prompts) > 0 { + d.Set("prompts", key.Prompts) + } + if key.OrganizationID != "" { + d.Set("organization_id", key.OrganizationID) + } + if key.ProjectID != "" { + d.Set("project_id", key.ProjectID) + } } diff --git a/terraform/provider/litellm/resource_key_block.go b/terraform/provider/litellm/resource_key_block.go new file mode 100644 index 00000000000..7aa41f832bb --- /dev/null +++ b/terraform/provider/litellm/resource_key_block.go @@ -0,0 +1,135 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "log" + "net/http" + "net/url" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const ( + endpointKeyBlock = "/key/block" + endpointKeyUnblock = "/key/unblock" +) + +type KeyBlockInfoResponse struct { + Info struct { + Blocked *bool `json:"blocked"` + } `json:"info"` +} + +func resourceLiteLLMKeyBlock() *schema.Resource { + return &schema.Resource{ + Create: resourceLiteLLMKeyBlockCreate, + Read: resourceLiteLLMKeyBlockRead, + Delete: resourceLiteLLMKeyBlockDelete, + + Importer: &schema.ResourceImporter{ + StateContext: schema.ImportStatePassthroughContext, + }, + + Schema: map[string]*schema.Schema{ + "key": { + Type: schema.TypeString, + Required: true, + ForceNew: true, + Sensitive: true, + Description: "The API key to block, as the raw sk- value or its SHA-256 token hash. Destroying this resource unblocks the key", + DiffSuppressFunc: func(k, old, new string, d *schema.ResourceData) bool { + return old != "" && hashedKeyToken(old) == hashedKeyToken(new) + }, + }, + "blocked": { + Type: schema.TypeBool, + Computed: true, + Description: "Whether the key is currently blocked", + }, + }, + } +} + +func resourceLiteLLMKeyBlockCreate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + // Block by the SHA-256 token hash so the raw key never appears in the + // request, the resource ID, or Terraform plan output. + token := hashedKeyToken(d.Get("key").(string)) + + log.Printf("[INFO] Blocking key") + + resp, err := MakeRequest(client, "POST", endpointKeyBlock, map[string]interface{}{"key": token}) + if err != nil { + return fmt.Errorf("error blocking key: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "blocking key"); err != nil { + return err + } + + d.SetId(token) + return resourceLiteLLMKeyBlockRead(d, m) +} + +func resourceLiteLLMKeyBlockRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + key := d.Id() + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("/key/info?key=%s", url.QueryEscape(key)), nil) + if err != nil { + return fmt.Errorf("error reading key info: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + log.Printf("[WARN] Key not found, removing key block from state") + d.SetId("") + return nil + } + + if err := handleResponse(resp, "reading key info"); err != nil { + return err + } + + var infoResp KeyBlockInfoResponse + if err := json.NewDecoder(resp.Body).Decode(&infoResp); err != nil { + return fmt.Errorf("error decoding key info response: %w", err) + } + + if infoResp.Info.Blocked == nil || !*infoResp.Info.Blocked { + log.Printf("[WARN] Key is no longer blocked, removing key block from state") + d.SetId("") + return nil + } + + // Keep the configured key value; only fill it from the hashed ID when + // importing, where no configured value exists yet. + if _, ok := d.GetOk("key"); !ok { + d.Set("key", key) + } + d.Set("blocked", true) + return nil +} + +func resourceLiteLLMKeyBlockDelete(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Unblocking key") + + resp, err := MakeRequest(client, "POST", endpointKeyUnblock, map[string]interface{}{"key": d.Id()}) + if err != nil { + return fmt.Errorf("error unblocking key: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusNotFound { + if err := handleResponse(resp, "unblocking key"); err != nil { + return err + } + } + + d.SetId("") + return nil +} diff --git a/terraform/provider/litellm/resource_key_block_test.go b/terraform/provider/litellm/resource_key_block_test.go new file mode 100644 index 00000000000..3de7d3494a5 --- /dev/null +++ b/terraform/provider/litellm/resource_key_block_test.go @@ -0,0 +1,160 @@ +package litellm + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +// SHA-256 of "sk-test-123", the token hash the proxy stores for that key. +const keyBlockTestHash = "e0dbaa0c6455768bf812d8345ec96a2677d1e3bf17dbb0020b115c80092811e6" + +func newKeyBlockTestResourceData(t *testing.T, key string) *schema.ResourceData { + t.Helper() + return schema.TestResourceDataRaw(t, resourceLiteLLMKeyBlock().Schema, map[string]interface{}{ + "key": key, + }) +} + +func TestResourceLiteLLMKeyBlockCreate(t *testing.T) { + var blockPayload map[string]interface{} + mux := http.NewServeMux() + mux.HandleFunc("/key/block", func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("expected POST, got %s", r.Method) + } + if err := json.NewDecoder(r.Body).Decode(&blockPayload); err != nil { + t.Fatalf("failed to decode block payload: %v", err) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"blocked":true}`)) + }) + mux.HandleFunc("/key/info", func(w http.ResponseWriter, r *http.Request) { + if got := r.URL.Query().Get("key"); got != keyBlockTestHash { + t.Errorf("expected key query to be the token hash %q, got %q", keyBlockTestHash, got) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"key":"sk-test-123","info":{"blocked":true}}`)) + }) + srv := httptest.NewServer(mux) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newKeyBlockTestResourceData(t, "sk-test-123") + + if err := resourceLiteLLMKeyBlockCreate(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != keyBlockTestHash { + t.Fatalf("expected ID to be the token hash %q, got %q", keyBlockTestHash, d.Id()) + } + if blockPayload["key"] != keyBlockTestHash { + t.Fatalf("expected block payload to carry the token hash, got %+v", blockPayload) + } + if !d.Get("blocked").(bool) { + t.Fatal("expected blocked=true in state") + } +} + +func TestResourceLiteLLMKeyBlockRead_UnblockedClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"key":"sk-test-123","info":{"blocked":false}}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newKeyBlockTestResourceData(t, "sk-test-123") + d.SetId(keyBlockTestHash) + + if err := resourceLiteLLMKeyBlockRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared for unblocked key, got %q", d.Id()) + } +} + +func TestResourceLiteLLMKeyBlockRead_404ClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newKeyBlockTestResourceData(t, "sk-test-123") + d.SetId(keyBlockTestHash) + + if err := resourceLiteLLMKeyBlockRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared on 404, got %q", d.Id()) + } +} + +func TestResourceLiteLLMKeyBlockDelete(t *testing.T) { + var gotPath string + var unblockPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + json.NewDecoder(r.Body).Decode(&unblockPayload) + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"blocked":false}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newKeyBlockTestResourceData(t, "sk-test-123") + d.SetId(keyBlockTestHash) + + if err := resourceLiteLLMKeyBlockDelete(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if gotPath != "/key/unblock" { + t.Fatalf("expected path /key/unblock, got %s", gotPath) + } + if unblockPayload["key"] != keyBlockTestHash { + t.Fatalf("expected unblock payload to carry the token hash, got %+v", unblockPayload) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared after delete, got %q", d.Id()) + } +} + +// Regression for the security review finding: a raw sk- key must never leave +// the provider in a URL, request body, or resource ID; only its SHA-256 token +// hash may. +func TestKeyBlockNeverSendsRawKey(t *testing.T) { + var seen []string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + seen = append(seen, r.URL.String()+" "+string(body)) + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"key":"x","info":{"blocked":true}}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "master-key", true) + d := newKeyBlockTestResourceData(t, "sk-test-123") + if err := resourceLiteLLMKeyBlockCreate(d, client); err != nil { + t.Fatalf("create failed: %v", err) + } + if err := resourceLiteLLMKeyBlockRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + if err := resourceLiteLLMKeyBlockDelete(d, client); err != nil { + t.Fatalf("delete failed: %v", err) + } + + for _, req := range seen { + if strings.Contains(req, "sk-test-123") { + t.Fatalf("raw key leaked to the API: %s", req) + } + } +} diff --git a/terraform/provider/litellm/resource_key_test.go b/terraform/provider/litellm/resource_key_test.go new file mode 100644 index 00000000000..91f0061a9ef --- /dev/null +++ b/terraform/provider/litellm/resource_key_test.go @@ -0,0 +1,256 @@ +package litellm + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func newKeyResourceData(t *testing.T, raw map[string]interface{}) *schema.ResourceData { + t.Helper() + return schema.TestResourceDataRaw(t, resourceKey().Schema, raw) +} + +func TestMapResourceDataToKeyNewFields(t *testing.T) { + d := newKeyResourceData(t, map[string]interface{}{ + "budget_id": "budget-1", + "enforced_params": []interface{}{"user"}, + "allowed_routes": []interface{}{"/chat/completions"}, + "allowed_passthrough_routes": []interface{}{"/vertex-ai"}, + "rpm_limit_type": "guaranteed_throughput", + "tpm_limit_type": "best_effort_throughput", + "prompts": []interface{}{"prompt-1"}, + "organization_id": "org-1", + "project_id": "proj-1", + }) + + key := &Key{} + mapResourceDataToKey(d, key) + + if key.BudgetID != "budget-1" { + t.Errorf("BudgetID = %q, want budget-1", key.BudgetID) + } + if len(key.EnforcedParams) != 1 || key.EnforcedParams[0] != "user" { + t.Errorf("EnforcedParams = %v, want [user]", key.EnforcedParams) + } + if len(key.AllowedRoutes) != 1 || key.AllowedRoutes[0] != "/chat/completions" { + t.Errorf("AllowedRoutes = %v", key.AllowedRoutes) + } + if len(key.AllowedPassthroughRoutes) != 1 || key.AllowedPassthroughRoutes[0] != "/vertex-ai" { + t.Errorf("AllowedPassthroughRoutes = %v", key.AllowedPassthroughRoutes) + } + if key.RPMLimitType != "guaranteed_throughput" { + t.Errorf("RPMLimitType = %q", key.RPMLimitType) + } + if key.TPMLimitType != "best_effort_throughput" { + t.Errorf("TPMLimitType = %q", key.TPMLimitType) + } + if len(key.Prompts) != 1 || key.Prompts[0] != "prompt-1" { + t.Errorf("Prompts = %v", key.Prompts) + } + if key.OrganizationID != "org-1" { + t.Errorf("OrganizationID = %q", key.OrganizationID) + } + if key.ProjectID != "proj-1" { + t.Errorf("ProjectID = %q", key.ProjectID) + } +} + +func TestUpdateKeySendsNewFields(t *testing.T) { + var captured map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + json.Unmarshal(body, &captured) + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"key": "sk-test"}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + _, err := client.UpdateKey(&Key{ + Key: "sk-test", + BudgetID: "budget-1", + EnforcedParams: []string{"user"}, + AllowedRoutes: []string{"/chat/completions"}, + AllowedPassthroughRoutes: []string{"/vertex-ai"}, + RPMLimitType: "guaranteed_throughput", + TPMLimitType: "dynamic", + Prompts: []string{"prompt-1"}, + OrganizationID: "org-1", + }) + if err != nil { + t.Fatalf("UpdateKey returned error: %v", err) + } + + want := map[string]interface{}{ + "budget_id": "budget-1", + "rpm_limit_type": "guaranteed_throughput", + "tpm_limit_type": "dynamic", + "organization_id": "org-1", + } + for k, v := range want { + if captured[k] != v { + t.Errorf("update payload %s = %v, want %v", k, captured[k], v) + } + } + for _, k := range []string{"enforced_params", "allowed_routes", "allowed_passthrough_routes", "prompts"} { + list, ok := captured[k].([]interface{}) + if !ok || len(list) != 1 { + t.Errorf("update payload %s = %v, want single-element list", k, captured[k]) + } + } +} + +func TestUpdateKeyOmitsUnsetNewFields(t *testing.T) { + var captured map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + json.Unmarshal(body, &captured) + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"key": "sk-test"}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + if _, err := client.UpdateKey(&Key{Key: "sk-test"}); err != nil { + t.Fatalf("UpdateKey returned error: %v", err) + } + + for _, k := range []string{ + "budget_id", "enforced_params", "allowed_routes", "allowed_passthrough_routes", + "rpm_limit_type", "tpm_limit_type", "prompts", "organization_id", + } { + if _, present := captured[k]; present { + t.Errorf("update payload unexpectedly contains %s", k) + } + } +} + +func TestParseKeyResponseNewFields(t *testing.T) { + client := NewClient("http://localhost:4000", "test-key", true) + resp := map[string]interface{}{ + "key": "sk-test", + "budget_id": "budget-1", + "enforced_params": []interface{}{"user"}, + "allowed_routes": []interface{}{"/chat/completions"}, + "allowed_passthrough_routes": []interface{}{"/vertex-ai"}, + "rpm_limit_type": "guaranteed_throughput", + "tpm_limit_type": "best_effort_throughput", + "prompts": []interface{}{"prompt-1"}, + "organization_id": "org-1", + "project_id": "proj-1", + } + + key, err := client.parseKeyResponse(resp) + if err != nil { + t.Fatalf("parseKeyResponse returned error: %v", err) + } + if key.BudgetID != "budget-1" || key.OrganizationID != "org-1" || key.ProjectID != "proj-1" { + t.Errorf("string fields not parsed: %+v", key) + } + if key.RPMLimitType != "guaranteed_throughput" || key.TPMLimitType != "best_effort_throughput" { + t.Errorf("limit types not parsed: %+v", key) + } + if len(key.EnforcedParams) != 1 || len(key.AllowedRoutes) != 1 || len(key.AllowedPassthroughRoutes) != 1 || len(key.Prompts) != 1 { + t.Errorf("list fields not parsed: %+v", key) + } +} + +// A config-supplied key value must be forwarded to /key/generate; previously +// it was silently dropped and the proxy generated a random key instead. +func TestCreateKeySendsConfigSuppliedKey(t *testing.T) { + var captured map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/key/generate" { + body, _ := io.ReadAll(r.Body) + json.Unmarshal(body, &captured) + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"key": "sk-custom", "token_id": "hash-1"}`)) + return + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"key": "sk-custom", "token_id": "hash-1"}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newKeyResourceData(t, map[string]interface{}{"key": "sk-custom"}) + + diags := resourceKeyCreate(context.Background(), d, client) + if diags.HasError() { + t.Fatalf("create returned error: %v", diags) + } + if captured["key"] != "sk-custom" { + t.Errorf("create payload key = %v, want sk-custom", captured["key"]) + } + if d.Id() != "hash-1" { + t.Errorf("resource ID = %q, want hash-1", d.Id()) + } +} + +// The proxy 400s on budget_duration: "", so an unset duration must be +// omitted from the update payload entirely. +func TestUpdateKeyOmitsEmptyBudgetDuration(t *testing.T) { + var captured map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + json.Unmarshal(body, &captured) + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"key": "sk-test"}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + if _, err := client.UpdateKey(&Key{Key: "sk-test"}); err != nil { + t.Fatalf("UpdateKey returned error: %v", err) + } + if _, present := captured["budget_duration"]; present { + t.Errorf("update payload contains empty budget_duration: %v", captured["budget_duration"]) + } + + if _, err := client.UpdateKey(&Key{Key: "sk-test", BudgetDuration: "30d"}); err != nil { + t.Fatalf("UpdateKey returned error: %v", err) + } + if captured["budget_duration"] != "30d" { + t.Errorf("budget_duration = %v, want 30d", captured["budget_duration"]) + } +} + +// /key/info nests the key's fields under "info"; GetKey must unwrap that +// envelope or reads map nothing back into state. +func TestGetKeyUnwrapsInfoEnvelope(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{ + "key": "hash-1", + "info": { + "key_alias": "envelope-alias", + "models": ["gpt-4o-mini"], + "budget_id": "budget-1", + "team_id": "team-1", + "rpm_limit": 100 + } + }`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + key, err := client.GetKey("hash-1") + if err != nil { + t.Fatalf("GetKey returned error: %v", err) + } + if key.KeyAlias != "envelope-alias" { + t.Errorf("KeyAlias = %q, want envelope-alias (info envelope not unwrapped)", key.KeyAlias) + } + if key.BudgetID != "budget-1" || key.TeamID != "team-1" { + t.Errorf("nested fields not parsed: %+v", key) + } + if key.RPMLimit == nil || *key.RPMLimit != 100 { + t.Errorf("RPMLimit not parsed: %+v", key.RPMLimit) + } +} diff --git a/terraform/provider/litellm/resource_mcp_server.go b/terraform/provider/litellm/resource_mcp_server.go index b3eaef4a468..318925c4367 100644 --- a/terraform/provider/litellm/resource_mcp_server.go +++ b/terraform/provider/litellm/resource_mcp_server.go @@ -11,6 +11,9 @@ func resourceLiteLLMMCPServer() *schema.Resource { Read: resourceLiteLLMMCPServerRead, Update: resourceLiteLLMMCPServerUpdate, Delete: resourceLiteLLMMCPServerDelete, + Importer: &schema.ResourceImporter{ + StateContext: schema.ImportStatePassthroughContext, + }, Schema: map[string]*schema.Schema{ "server_name": { diff --git a/terraform/provider/litellm/resource_model.go b/terraform/provider/litellm/resource_model.go index 4bad057871d..b0a7304718b 100644 --- a/terraform/provider/litellm/resource_model.go +++ b/terraform/provider/litellm/resource_model.go @@ -11,6 +11,9 @@ func resourceLiteLLMModel() *schema.Resource { Read: resourceLiteLLMModelRead, Update: resourceLiteLLMModelUpdate, Delete: resourceLiteLLMModelDelete, + Importer: &schema.ResourceImporter{ + StateContext: schema.ImportStatePassthroughContext, + }, Schema: map[string]*schema.Schema{ "model_name": { diff --git a/terraform/provider/litellm/resource_organization.go b/terraform/provider/litellm/resource_organization.go index 30e7feba1ec..d0908434b1e 100644 --- a/terraform/provider/litellm/resource_organization.go +++ b/terraform/provider/litellm/resource_organization.go @@ -23,6 +23,9 @@ func resourceLiteLLMOrganization() *schema.Resource { Read: resourceLiteLLMOrganizationRead, Update: resourceLiteLLMOrganizationUpdate, Delete: resourceLiteLLMOrganizationDelete, + Importer: &schema.ResourceImporter{ + StateContext: schema.ImportStatePassthroughContext, + }, Schema: map[string]*schema.Schema{ "organization_alias": { diff --git a/terraform/provider/litellm/resource_project.go b/terraform/provider/litellm/resource_project.go new file mode 100644 index 00000000000..ae6b372c72c --- /dev/null +++ b/terraform/provider/litellm/resource_project.go @@ -0,0 +1,352 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "io" + "log" + "net/http" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const ( + endpointProjectNew = "/project/new" + endpointProjectInfo = "/project/info" + endpointProjectUpdate = "/project/update" + endpointProjectDelete = "/project/delete" +) + +type projectBudgetTable struct { + MaxBudget *float64 `json:"max_budget"` + SoftBudget *float64 `json:"soft_budget"` + MaxParallelRequests *int `json:"max_parallel_requests"` + TPMLimit *int `json:"tpm_limit"` + RPMLimit *int `json:"rpm_limit"` + BudgetDuration string `json:"budget_duration"` +} + +type projectResponse struct { + ProjectID string `json:"project_id"` + ProjectAlias string `json:"project_alias"` + Description string `json:"description"` + TeamID string `json:"team_id"` + BudgetID string `json:"budget_id"` + Metadata map[string]interface{} `json:"metadata"` + Models []string `json:"models"` + Spend float64 `json:"spend"` + Blocked bool `json:"blocked"` + CreatedBy string `json:"created_by"` + UpdatedBy string `json:"updated_by"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` + LitellmBudgetTable *projectBudgetTable `json:"litellm_budget_table"` +} + +func resourceLiteLLMProject() *schema.Resource { + return &schema.Resource{ + Create: resourceLiteLLMProjectCreate, + Read: resourceLiteLLMProjectRead, + Update: resourceLiteLLMProjectUpdate, + Delete: resourceLiteLLMProjectDelete, + + Importer: &schema.ResourceImporter{StateContext: schema.ImportStatePassthroughContext}, + + Schema: map[string]*schema.Schema{ + "team_id": { + Type: schema.TypeString, + Required: true, + ForceNew: true, + Description: "The team ID this project belongs to.", + }, + "project_alias": { + Type: schema.TypeString, + Optional: true, + Description: "Human-friendly name for the project.", + }, + "description": { + Type: schema.TypeString, + Optional: true, + Description: "Description of the project's purpose and use case.", + }, + "models": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "List of models the project can access.", + }, + "metadata": { + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Metadata for the project.", + }, + "tags": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Tags associated with the project.", + }, + "max_budget": { + Type: schema.TypeFloat, + Optional: true, + Description: "Maximum budget for this project.", + }, + "soft_budget": { + Type: schema.TypeFloat, + Optional: true, + Description: "Soft budget limit for warnings.", + }, + "budget_duration": { + Type: schema.TypeString, + Optional: true, + Description: "Budget reset duration (e.g. '30d', '1h').", + }, + "budget_id": { + Type: schema.TypeString, + Optional: true, + Description: "Budget ID to associate with this project.", + }, + "tpm_limit": { + Type: schema.TypeInt, + Optional: true, + Description: "Tokens per minute limit.", + }, + "rpm_limit": { + Type: schema.TypeInt, + Optional: true, + Description: "Requests per minute limit.", + }, + "max_parallel_requests": { + Type: schema.TypeInt, + Optional: true, + Description: "Maximum parallel requests allowed.", + }, + "model_max_budget": { + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeFloat}, + Description: "Per-model budget limits.", + }, + "model_rpm_limit": { + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeInt}, + Description: "Per-model RPM limits.", + }, + "model_tpm_limit": { + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeInt}, + Description: "Per-model TPM limits.", + }, + "blocked": { + Type: schema.TypeBool, + Optional: true, + Description: "Whether the project is blocked from making requests.", + }, + "spend": { + Type: schema.TypeFloat, + Computed: true, + Description: "Current spend for the project.", + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + Description: "Timestamp when the project was created.", + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + Description: "Timestamp when the project was last updated.", + }, + "created_by": { + Type: schema.TypeString, + Computed: true, + Description: "User that created the project.", + }, + "updated_by": { + Type: schema.TypeString, + Computed: true, + Description: "User that last updated the project.", + }, + }, + } +} + +func buildProjectData(d *schema.ResourceData) map[string]interface{} { + projectData := map[string]interface{}{ + "team_id": d.Get("team_id").(string), + } + + for _, key := range []string{"project_alias", "description", "models", "metadata", "tags", + "max_budget", "soft_budget", "budget_duration", "budget_id", "tpm_limit", "rpm_limit", + "max_parallel_requests", "model_max_budget", "model_rpm_limit", "model_tpm_limit", "blocked"} { + if v, ok := d.GetOk(key); ok { + projectData[key] = v + } + } + + return projectData +} + +func resourceLiteLLMProjectCreate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + projectData := buildProjectData(d) + log.Printf("[DEBUG] Create project request payload: %+v", projectData) + + resp, err := MakeRequest(client, "POST", endpointProjectNew, projectData) + if err != nil { + return fmt.Errorf("error creating project: %w", err) + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return fmt.Errorf("error reading create project response: %w", err) + } + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("error creating project: %s - %s", resp.Status, string(body)) + } + + var projResp projectResponse + if err := json.Unmarshal(body, &projResp); err != nil { + return fmt.Errorf("error decoding create project response: %w", err) + } + if projResp.ProjectID == "" { + return fmt.Errorf("create project response did not contain a project_id: %s", string(body)) + } + + d.SetId(projResp.ProjectID) + log.Printf("[INFO] Project created with ID: %s", projResp.ProjectID) + + return resourceLiteLLMProjectRead(d, m) +} + +func resourceLiteLLMProjectRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Reading project with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("%s?project_id=%s", endpointProjectInfo, d.Id()), nil) + if err != nil { + return fmt.Errorf("error reading project: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + log.Printf("[WARN] Project with ID %s not found, removing from state", d.Id()) + d.SetId("") + return nil + } + + if err := handleResponse(resp, "reading project"); err != nil { + return err + } + + var projResp projectResponse + if err := json.NewDecoder(resp.Body).Decode(&projResp); err != nil { + return fmt.Errorf("error decoding project info response: %w", err) + } + + d.Set("team_id", GetStringValue(projResp.TeamID, d.Get("team_id").(string))) + d.Set("project_alias", GetStringValue(projResp.ProjectAlias, d.Get("project_alias").(string))) + d.Set("description", GetStringValue(projResp.Description, d.Get("description").(string))) + d.Set("budget_id", GetStringValue(projResp.BudgetID, d.Get("budget_id").(string))) + if projResp.Models != nil { + d.Set("models", projResp.Models) + } + setProjectMetadataAndTags(d, projResp.Metadata) + + d.Set("blocked", projResp.Blocked) + d.Set("spend", projResp.Spend) + d.Set("created_at", projResp.CreatedAt) + d.Set("updated_at", projResp.UpdatedAt) + d.Set("created_by", projResp.CreatedBy) + d.Set("updated_by", projResp.UpdatedBy) + + if bt := projResp.LitellmBudgetTable; bt != nil { + if bt.MaxBudget != nil { + d.Set("max_budget", *bt.MaxBudget) + } + if bt.SoftBudget != nil { + d.Set("soft_budget", *bt.SoftBudget) + } + if bt.MaxParallelRequests != nil { + d.Set("max_parallel_requests", *bt.MaxParallelRequests) + } + if bt.TPMLimit != nil { + d.Set("tpm_limit", *bt.TPMLimit) + } + if bt.RPMLimit != nil { + d.Set("rpm_limit", *bt.RPMLimit) + } + d.Set("budget_duration", GetStringValue(bt.BudgetDuration, d.Get("budget_duration").(string))) + } + + log.Printf("[INFO] Successfully read project with ID: %s", d.Id()) + return nil +} + +// The proxy stores project tags inside metadata; split them back out so state matches the config shape. +func setProjectMetadataAndTags(d *schema.ResourceData, metadata map[string]interface{}) { + if metadata == nil { + return + } + + if tags, ok := metadata["tags"].([]interface{}); ok { + d.Set("tags", tags) + } + + stringMetadata := map[string]interface{}{} + for k, v := range metadata { + if s, ok := v.(string); ok { + stringMetadata[k] = s + } + } + d.Set("metadata", stringMetadata) +} + +func resourceLiteLLMProjectUpdate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + projectData := buildProjectData(d) + projectData["project_id"] = d.Id() + log.Printf("[DEBUG] Update project request payload: %+v", projectData) + + resp, err := MakeRequest(client, "POST", endpointProjectUpdate, projectData) + if err != nil { + return fmt.Errorf("error updating project: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "updating project"); err != nil { + return err + } + + log.Printf("[INFO] Successfully updated project with ID: %s", d.Id()) + return resourceLiteLLMProjectRead(d, m) +} + +func resourceLiteLLMProjectDelete(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Deleting project with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "DELETE", endpointProjectDelete, map[string]interface{}{ + "project_ids": []string{d.Id()}, + }) + if err != nil { + return fmt.Errorf("error deleting project: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "deleting project"); err != nil { + return err + } + + log.Printf("[INFO] Successfully deleted project with ID: %s", d.Id()) + d.SetId("") + return nil +} diff --git a/terraform/provider/litellm/resource_project_test.go b/terraform/provider/litellm/resource_project_test.go new file mode 100644 index 00000000000..0c9538976df --- /dev/null +++ b/terraform/provider/litellm/resource_project_test.go @@ -0,0 +1,236 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "reflect" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const projectInfoBody = `{ + "project_id": "proj-123", + "project_alias": "ml-experiments", + "description": "ML experimentation project", + "team_id": "team-1", + "budget_id": "bud-9", + "metadata": {"env": "prod", "tags": ["research", "gpu"]}, + "models": ["gpt-4"], + "spend": 12.5, + "blocked": false, + "created_by": "admin", + "updated_by": "admin", + "created_at": "2026-01-01T00:00:00", + "updated_at": "2026-01-02T00:00:00", + "litellm_budget_table": { + "max_budget": 100.0, + "soft_budget": 80.0, + "max_parallel_requests": 10, + "tpm_limit": 5000, + "rpm_limit": 500, + "budget_duration": "30d" + } +}` + +func TestResourceLiteLLMProjectCreate(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/project/new": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(projectInfoBody)) + case "/project/info": + if got := r.URL.Query().Get("project_id"); got != "proj-123" { + t.Errorf("expected project_id query 'proj-123', got %q", got) + } + w.Write([]byte(projectInfoBody)) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMProject().Schema, map[string]interface{}{ + "team_id": "team-1", + "project_alias": "ml-experiments", + "description": "ML experimentation project", + "models": []interface{}{"gpt-4"}, + "metadata": map[string]interface{}{"env": "prod"}, + "tags": []interface{}{"research", "gpu"}, + "max_budget": 100.0, + "tpm_limit": 5000, + }) + + if err := resourceLiteLLMProjectCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + if d.Id() != "proj-123" { + t.Fatalf("expected ID 'proj-123', got %q", d.Id()) + } + if createPayload["team_id"] != "team-1" { + t.Errorf("expected payload team_id 'team-1', got %v", createPayload["team_id"]) + } + if createPayload["project_alias"] != "ml-experiments" { + t.Errorf("expected payload project_alias, got %v", createPayload["project_alias"]) + } + if !reflect.DeepEqual(createPayload["models"], []interface{}{"gpt-4"}) { + t.Errorf("expected payload models ['gpt-4'], got %v", createPayload["models"]) + } + if !reflect.DeepEqual(createPayload["tags"], []interface{}{"research", "gpu"}) { + t.Errorf("expected payload tags, got %v", createPayload["tags"]) + } + if createPayload["max_budget"] != 100.0 { + t.Errorf("expected payload max_budget 100.0, got %v", createPayload["max_budget"]) + } + if createPayload["tpm_limit"] != float64(5000) { + t.Errorf("expected payload tpm_limit 5000, got %v", createPayload["tpm_limit"]) + } + if _, ok := createPayload["project_id"]; ok { + t.Errorf("create payload must not contain project_id, got %v", createPayload["project_id"]) + } +} + +func TestResourceLiteLLMProjectRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/project/info" || r.Method != http.MethodGet { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Write([]byte(projectInfoBody)) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMProject().Schema, map[string]interface{}{ + "team_id": "team-1", + }) + d.SetId("proj-123") + + if err := resourceLiteLLMProjectRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + checks := map[string]interface{}{ + "project_alias": "ml-experiments", + "description": "ML experimentation project", + "team_id": "team-1", + "budget_id": "bud-9", + "spend": 12.5, + "max_budget": 100.0, + "soft_budget": 80.0, + "max_parallel_requests": 10, + "tpm_limit": 5000, + "rpm_limit": 500, + "budget_duration": "30d", + "created_by": "admin", + "created_at": "2026-01-01T00:00:00", + } + for key, want := range checks { + if got := d.Get(key); got != want { + t.Errorf("expected %s %v, got %v", key, want, got) + } + } + if !reflect.DeepEqual(d.Get("tags"), []interface{}{"research", "gpu"}) { + t.Errorf("expected tags extracted from metadata, got %v", d.Get("tags")) + } + wantMetadata := map[string]interface{}{"env": "prod"} + if !reflect.DeepEqual(d.Get("metadata"), wantMetadata) { + t.Errorf("expected metadata without injected tags key, got %v", d.Get("metadata")) + } +} + +func TestResourceLiteLLMProjectRead_404ClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMProject().Schema, map[string]interface{}{ + "team_id": "team-1", + }) + d.SetId("gone") + + if err := resourceLiteLLMProjectRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("expected nil error on 404, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared on 404, got %q", d.Id()) + } +} + +func TestResourceLiteLLMProjectUpdate(t *testing.T) { + var updatePayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/project/update": + if r.Method != http.MethodPost { + t.Errorf("expected POST for update, got %s", r.Method) + } + if err := json.NewDecoder(r.Body).Decode(&updatePayload); err != nil { + t.Errorf("failed to decode update payload: %v", err) + } + w.Write([]byte(projectInfoBody)) + case "/project/info": + w.Write([]byte(projectInfoBody)) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMProject().Schema, map[string]interface{}{ + "team_id": "team-1", + "project_alias": "renamed-project", + "rpm_limit": 900, + }) + d.SetId("proj-123") + + if err := resourceLiteLLMProjectUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + + if updatePayload["project_id"] != "proj-123" { + t.Errorf("expected update payload project_id 'proj-123', got %v", updatePayload["project_id"]) + } + if updatePayload["project_alias"] != "renamed-project" { + t.Errorf("expected updated project_alias in payload, got %v", updatePayload["project_alias"]) + } + if updatePayload["rpm_limit"] != float64(900) { + t.Errorf("expected rpm_limit 900 in payload, got %v", updatePayload["rpm_limit"]) + } +} + +func TestResourceLiteLLMProjectDelete(t *testing.T) { + var deletePayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/project/delete" || r.Method != http.MethodDelete { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + if err := json.NewDecoder(r.Body).Decode(&deletePayload); err != nil { + t.Errorf("failed to decode delete payload: %v", err) + } + w.Write([]byte(`[]`)) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMProject().Schema, map[string]interface{}{ + "team_id": "team-1", + }) + d.SetId("proj-123") + + if err := resourceLiteLLMProjectDelete(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("delete failed: %v", err) + } + + if !reflect.DeepEqual(deletePayload["project_ids"], []interface{}{"proj-123"}) { + t.Errorf("expected delete payload project_ids ['proj-123'], got %v", deletePayload["project_ids"]) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared after delete, got %q", d.Id()) + } +} diff --git a/terraform/provider/litellm/resource_prompt.go b/terraform/provider/litellm/resource_prompt.go new file mode 100644 index 00000000000..b7d138227e0 --- /dev/null +++ b/terraform/provider/litellm/resource_prompt.go @@ -0,0 +1,304 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "io" + "log" + "net/http" + "reflect" + "strings" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const ( + endpointPromptCreate = "/prompts" + endpointPromptByID = "/prompts/%s" + endpointPromptInfo = "/prompts/%s/info" + endpointPromptList = "/prompts/list" +) + +func resourceLiteLLMPrompt() *schema.Resource { + return &schema.Resource{ + Create: resourceLiteLLMPromptCreate, + Read: resourceLiteLLMPromptRead, + Update: resourceLiteLLMPromptUpdate, + Delete: resourceLiteLLMPromptDelete, + + Importer: &schema.ResourceImporter{StateContext: schema.ImportStatePassthroughContext}, + + Schema: map[string]*schema.Schema{ + "prompt_id": { + Type: schema.TypeString, + Required: true, + ForceNew: true, + Description: "Unique identifier for the prompt", + }, + "prompt_integration": { + Type: schema.TypeString, + Required: true, + Description: "The prompt integration provider (e.g. 'langfuse', 'dotprompt')", + }, + "api_base": { + Type: schema.TypeString, + Optional: true, + Description: "Base URL for the prompt provider API", + }, + "api_key": { + Type: schema.TypeString, + Optional: true, + Sensitive: true, + Description: "API key for the prompt provider", + }, + "provider_specific_query_params": { + Type: schema.TypeString, + Optional: true, + DiffSuppressFunc: promptSuppressJSONDiff, + Description: "JSON string of provider-specific query parameters", + }, + "ignore_prompt_manager_model": { + Type: schema.TypeBool, + Optional: true, + Description: "If true, ignore the model specified in the prompt manager", + }, + "ignore_prompt_manager_optional_params": { + Type: schema.TypeBool, + Optional: true, + Description: "If true, ignore optional params from the prompt manager", + }, + "dotprompt_content": { + Type: schema.TypeString, + Optional: true, + Description: "Content for dotprompt integration", + }, + "litellm_params": { + Type: schema.TypeString, + Optional: true, + Sensitive: true, + DiffSuppressFunc: promptSuppressJSONDiff, + Description: "JSON string with additional litellm_params merged into the request " + + "(e.g. the integration's own prompt_id, prompt_directory, prompt_data; may contain secrets)", + }, + "prompt_type": { + Type: schema.TypeString, + Optional: true, + Description: "Type of prompt: 'config' or 'db'", + }, + }, + } +} + +func promptSuppressJSONDiff(k, oldValue, newValue string, d *schema.ResourceData) bool { + var oldParsed, newParsed interface{} + if json.Unmarshal([]byte(oldValue), &oldParsed) != nil || json.Unmarshal([]byte(newValue), &newParsed) != nil { + return false + } + return reflect.DeepEqual(oldParsed, newParsed) +} + +func buildPromptData(d *schema.ResourceData) (map[string]interface{}, error) { + litellmParams := map[string]interface{}{ + "prompt_integration": d.Get("prompt_integration").(string), + } + + for tfKey, apiKey := range map[string]string{ + "api_base": "api_base", + "api_key": "api_key", + "dotprompt_content": "dotprompt_content", + } { + if v := d.Get(tfKey).(string); v != "" { + litellmParams[apiKey] = v + } + } + + if v := d.Get("provider_specific_query_params").(string); v != "" { + var params map[string]interface{} + if err := json.Unmarshal([]byte(v), ¶ms); err != nil { + return nil, fmt.Errorf("provider_specific_query_params is not valid JSON: %w", err) + } + litellmParams["provider_specific_query_params"] = params + } + + litellmParams["ignore_prompt_manager_model"] = d.Get("ignore_prompt_manager_model").(bool) + litellmParams["ignore_prompt_manager_optional_params"] = d.Get("ignore_prompt_manager_optional_params").(bool) + + if raw := d.Get("litellm_params").(string); raw != "" { + var extra map[string]interface{} + if err := json.Unmarshal([]byte(raw), &extra); err != nil { + return nil, fmt.Errorf("litellm_params is not valid JSON: %w", err) + } + for k, v := range extra { + litellmParams[k] = v + } + } + + promptData := map[string]interface{}{ + "prompt_id": d.Get("prompt_id").(string), + "litellm_params": litellmParams, + } + + if v := d.Get("prompt_type").(string); v != "" { + promptData["prompt_info"] = map[string]interface{}{"prompt_type": v} + } + + return promptData, nil +} + +type promptSpecAPIResponse struct { + PromptID string `json:"prompt_id"` + LitellmParams map[string]interface{} `json:"litellm_params"` + PromptInfo map[string]interface{} `json:"prompt_info"` + Version int `json:"version"` + Environment string `json:"environment"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` +} + +func resourceLiteLLMPromptCreate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + promptData, err := buildPromptData(d) + if err != nil { + return err + } + + promptID := d.Get("prompt_id").(string) + log.Printf("[DEBUG] Create prompt request for: %s", promptID) + + resp, err := MakeRequest(client, "POST", endpointPromptCreate, promptData) + if err != nil { + return fmt.Errorf("error creating prompt: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "creating prompt"); err != nil { + return err + } + + d.SetId(promptID) + log.Printf("[INFO] Prompt created with ID: %s", promptID) + + return resourceLiteLLMPromptRead(d, m) +} + +func promptIsNotFoundResponse(resp *http.Response) bool { + if resp.StatusCode == http.StatusNotFound { + return true + } + if resp.StatusCode != http.StatusBadRequest { + return false + } + body, err := io.ReadAll(resp.Body) + if err != nil { + return false + } + resp.Body = io.NopCloser(strings.NewReader(string(body))) + return strings.Contains(string(body), "not found") +} + +func resourceLiteLLMPromptRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Reading prompt with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf(endpointPromptInfo, d.Id()), nil) + if err != nil { + return fmt.Errorf("error reading prompt: %w", err) + } + defer resp.Body.Close() + + if promptIsNotFoundResponse(resp) { + log.Printf("[WARN] Prompt with ID %s not found, removing from state", d.Id()) + d.SetId("") + return nil + } + + if err := handleResponse(resp, "reading prompt"); err != nil { + return err + } + + var info struct { + PromptSpec promptSpecAPIResponse `json:"prompt_spec"` + } + if err := json.NewDecoder(resp.Body).Decode(&info); err != nil { + return fmt.Errorf("error decoding prompt info response: %w", err) + } + + d.Set("prompt_id", info.PromptSpec.PromptID) + + params := info.PromptSpec.LitellmParams + if v, ok := params["prompt_integration"].(string); ok { + d.Set("prompt_integration", v) + } + if v, ok := params["api_base"].(string); ok { + d.Set("api_base", v) + } + if v, ok := params["dotprompt_content"].(string); ok { + d.Set("dotprompt_content", v) + } + if v, ok := params["ignore_prompt_manager_model"].(bool); ok { + d.Set("ignore_prompt_manager_model", v) + } + if v, ok := params["ignore_prompt_manager_optional_params"].(bool); ok { + d.Set("ignore_prompt_manager_optional_params", v) + } + if v, ok := params["provider_specific_query_params"].(map[string]interface{}); ok { + if encoded, err := json.Marshal(v); err == nil { + d.Set("provider_specific_query_params", string(encoded)) + } + } + if v, ok := info.PromptSpec.PromptInfo["prompt_type"].(string); ok { + d.Set("prompt_type", v) + } + // api_key and the litellm_params catch-all are intentionally not read back: + // they can carry secrets, so state keeps the configured values authoritative. + + return nil +} + +func resourceLiteLLMPromptUpdate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + promptData, err := buildPromptData(d) + if err != nil { + return err + } + + log.Printf("[DEBUG] Update prompt request for ID: %s", d.Id()) + + resp, err := MakeRequest(client, "PUT", fmt.Sprintf(endpointPromptByID, d.Id()), promptData) + if err != nil { + return fmt.Errorf("error updating prompt: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "updating prompt"); err != nil { + return err + } + + log.Printf("[INFO] Successfully updated prompt with ID: %s", d.Id()) + return resourceLiteLLMPromptRead(d, m) +} + +func resourceLiteLLMPromptDelete(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Deleting prompt with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "DELETE", fmt.Sprintf(endpointPromptByID, d.Id()), nil) + if err != nil { + return fmt.Errorf("error deleting prompt: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusNotFound { + if err := handleResponse(resp, "deleting prompt"); err != nil { + return err + } + } + + log.Printf("[INFO] Successfully deleted prompt with ID: %s", d.Id()) + d.SetId("") + return nil +} diff --git a/terraform/provider/litellm/resource_prompt_test.go b/terraform/provider/litellm/resource_prompt_test.go new file mode 100644 index 00000000000..5d25ad0cffa --- /dev/null +++ b/terraform/provider/litellm/resource_prompt_test.go @@ -0,0 +1,238 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func newPromptTestData(t *testing.T, raw map[string]interface{}) *schema.ResourceData { + t.Helper() + return schema.TestResourceDataRaw(t, resourceLiteLLMPrompt().Schema, raw) +} + +func promptInfoJSON(promptID string) string { + body, _ := json.Marshal(map[string]interface{}{ + "prompt_spec": map[string]interface{}{ + "prompt_id": promptID, + "litellm_params": map[string]interface{}{ + "prompt_integration": "langfuse", + "api_base": "https://langfuse.example.com", + "ignore_prompt_manager_model": true, + "provider_specific_query_params": map[string]interface{}{"label": "prod"}, + }, + "prompt_info": map[string]interface{}{"prompt_type": "db"}, + "version": 3, + "environment": "development", + "created_at": "2026-01-01T00:00:00Z", + "updated_at": "2026-01-02T00:00:00Z", + }, + "environments": []string{"development"}, + }) + return string(body) +} + +func TestPromptCreate_SendsPayloadAndSetsID(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch { + case r.Method == "POST" && r.URL.Path == "/prompts": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(`{"prompt_id": "p1"}`)) + case r.Method == "GET" && r.URL.Path == "/prompts/p1/info": + w.Write([]byte(promptInfoJSON("p1"))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newPromptTestData(t, map[string]interface{}{ + "prompt_id": "p1", + "prompt_integration": "langfuse", + "api_key": "sk-langfuse", + "litellm_params": `{"prompt_id": "external-prompt", "prompt_directory": "/prompts"}`, + "prompt_type": "db", + }) + + if err := resourceLiteLLMPromptCreate(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "p1" { + t.Fatalf("expected ID 'p1', got %q", d.Id()) + } + + if createPayload["prompt_id"] != "p1" { + t.Errorf("expected prompt_id 'p1', got %v", createPayload["prompt_id"]) + } + params, ok := createPayload["litellm_params"].(map[string]interface{}) + if !ok { + t.Fatalf("expected litellm_params object, got: %v", createPayload["litellm_params"]) + } + if params["prompt_integration"] != "langfuse" || params["api_key"] != "sk-langfuse" { + t.Errorf("unexpected litellm_params: %v", params) + } + if params["prompt_id"] != "external-prompt" || params["prompt_directory"] != "/prompts" { + t.Errorf("expected merged extra litellm_params, got: %v", params) + } + info, ok := createPayload["prompt_info"].(map[string]interface{}) + if !ok || info["prompt_type"] != "db" { + t.Errorf("expected prompt_info with prompt_type 'db', got: %v", createPayload["prompt_info"]) + } +} + +func TestPromptRead_MapsFieldsAndKeepsAPIKey(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" || r.URL.Path != "/prompts/p1/info" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(promptInfoJSON("p1"))) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newPromptTestData(t, map[string]interface{}{ + "prompt_id": "p1", + "prompt_integration": "old-integration", + "api_key": "sk-configured", + }) + d.SetId("p1") + + if err := resourceLiteLLMPromptRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if got := d.Get("prompt_integration").(string); got != "langfuse" { + t.Errorf("expected prompt_integration 'langfuse', got %q", got) + } + if got := d.Get("api_base").(string); got != "https://langfuse.example.com" { + t.Errorf("expected api_base from API, got %q", got) + } + if got := d.Get("ignore_prompt_manager_model").(bool); !got { + t.Error("expected ignore_prompt_manager_model true from API") + } + if got := d.Get("provider_specific_query_params").(string); got != `{"label":"prod"}` { + t.Errorf("expected provider_specific_query_params JSON, got %q", got) + } + if got := d.Get("prompt_type").(string); got != "db" { + t.Errorf("expected prompt_type 'db', got %q", got) + } + if got := d.Get("api_key").(string); got != "sk-configured" { + t.Errorf("expected configured api_key to stay authoritative, got %q", got) + } +} + +func TestPromptRead_NotFound400ClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusBadRequest) + w.Write([]byte(`{"detail": "Prompt p-gone not found"}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newPromptTestData(t, map[string]interface{}{ + "prompt_id": "p-gone", + "prompt_integration": "langfuse", + }) + d.SetId("p-gone") + + if err := resourceLiteLLMPromptRead(d, client); err != nil { + t.Fatalf("expected nil error on not-found 400, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID to be cleared, got %q", d.Id()) + } +} + +func TestPromptRead_Other400ReturnsError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusBadRequest) + w.Write([]byte(`{"detail": "invalid environment"}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newPromptTestData(t, map[string]interface{}{ + "prompt_id": "p1", + "prompt_integration": "langfuse", + }) + d.SetId("p1") + + if err := resourceLiteLLMPromptRead(d, client); err == nil { + t.Fatal("expected error for non-not-found 400, got nil") + } + if d.Id() != "p1" { + t.Fatalf("expected ID to be kept, got %q", d.Id()) + } +} + +func TestPromptUpdate_SendsPUTToPromptEndpoint(t *testing.T) { + var updateMethod, updatePath string + var updatePayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == "PUT" { + updateMethod, updatePath = r.Method, r.URL.Path + json.NewDecoder(r.Body).Decode(&updatePayload) + w.Write([]byte(`{"prompt_id": "p1"}`)) + return + } + w.Write([]byte(promptInfoJSON("p1"))) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newPromptTestData(t, map[string]interface{}{ + "prompt_id": "p1", + "prompt_integration": "langfuse", + "api_base": "https://new-base.example.com", + }) + d.SetId("p1") + + if err := resourceLiteLLMPromptUpdate(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if updateMethod != "PUT" || updatePath != "/prompts/p1" { + t.Fatalf("expected PUT /prompts/p1, got %s %s", updateMethod, updatePath) + } + params := updatePayload["litellm_params"].(map[string]interface{}) + if params["api_base"] != "https://new-base.example.com" { + t.Errorf("expected updated api_base in payload, got %v", params["api_base"]) + } +} + +func TestPromptDelete_CallsDeleteEndpoint(t *testing.T) { + var deleteMethod, deletePath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + deleteMethod, deletePath = r.Method, r.URL.Path + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"message": "deleted"}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newPromptTestData(t, map[string]interface{}{ + "prompt_id": "p1", + "prompt_integration": "langfuse", + }) + d.SetId("p1") + + if err := resourceLiteLLMPromptDelete(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if deleteMethod != "DELETE" || deletePath != "/prompts/p1" { + t.Fatalf("expected DELETE /prompts/p1, got %s %s", deleteMethod, deletePath) + } + if d.Id() != "" { + t.Fatalf("expected ID to be cleared after delete, got %q", d.Id()) + } +} diff --git a/terraform/provider/litellm/resource_search_tool.go b/terraform/provider/litellm/resource_search_tool.go new file mode 100644 index 00000000000..367bc9a9523 --- /dev/null +++ b/terraform/provider/litellm/resource_search_tool.go @@ -0,0 +1,237 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "log" + "net/http" + "reflect" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const ( + endpointSearchTools = "/search_tools" + endpointSearchToolByID = "/search_tools/%s" + endpointSearchToolsList = "/search_tools/list" +) + +type searchToolAPIResponse struct { + SearchToolID string `json:"search_tool_id"` + SearchToolName string `json:"search_tool_name"` + SearchToolInfo map[string]interface{} `json:"search_tool_info"` + IsFromConfig *bool `json:"is_from_config"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` +} + +func searchToolSuppressEquivalentJSON(k, oldValue, newValue string, d *schema.ResourceData) bool { + var oldObj, newObj interface{} + if err := json.Unmarshal([]byte(oldValue), &oldObj); err != nil { + return false + } + if err := json.Unmarshal([]byte(newValue), &newObj); err != nil { + return false + } + return reflect.DeepEqual(oldObj, newObj) +} + +func searchToolParseJSONObject(raw, field string) (map[string]interface{}, error) { + var obj map[string]interface{} + if err := json.Unmarshal([]byte(raw), &obj); err != nil { + return nil, fmt.Errorf("%s must be a JSON object: %w", field, err) + } + return obj, nil +} + +func resourceLiteLLMSearchTool() *schema.Resource { + return &schema.Resource{ + Create: resourceLiteLLMSearchToolCreate, + Read: resourceLiteLLMSearchToolRead, + Update: resourceLiteLLMSearchToolUpdate, + Delete: resourceLiteLLMSearchToolDelete, + + Importer: &schema.ResourceImporter{StateContext: schema.ImportStatePassthroughContext}, + + Schema: map[string]*schema.Schema{ + "search_tool_name": { + Type: schema.TypeString, + Required: true, + Description: "Name of the search tool.", + }, + "litellm_params": { + Type: schema.TypeString, + Required: true, + Sensitive: true, + DiffSuppressFunc: searchToolSuppressEquivalentJSON, + Description: "Search tool parameters as a JSON object string (search_provider, api_key, " + + "api_base, timeout, max_retries, ...). The API only returns masked values, so this is " + + "never read back.", + }, + "search_tool_info": { + Type: schema.TypeString, + Optional: true, + DiffSuppressFunc: searchToolSuppressEquivalentJSON, + Description: "Additional metadata as a JSON object string (e.g. description).", + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + }, + } +} + +func buildSearchToolData(d *schema.ResourceData) (map[string]interface{}, error) { + litellmParams, err := searchToolParseJSONObject(d.Get("litellm_params").(string), "litellm_params") + if err != nil { + return nil, err + } + + searchToolData := map[string]interface{}{ + "search_tool_name": d.Get("search_tool_name").(string), + "litellm_params": litellmParams, + } + + if raw, ok := d.GetOk("search_tool_info"); ok && raw.(string) != "" { + info, err := searchToolParseJSONObject(raw.(string), "search_tool_info") + if err != nil { + return nil, err + } + searchToolData["search_tool_info"] = info + } + + return searchToolData, nil +} + +func resourceLiteLLMSearchToolCreate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + searchToolData, err := buildSearchToolData(d) + if err != nil { + return err + } + + log.Printf("[DEBUG] Create search tool request for: %s", d.Get("search_tool_name").(string)) + + resp, err := MakeRequest(client, "POST", endpointSearchTools, map[string]interface{}{ + "search_tool": searchToolData, + }) + if err != nil { + return fmt.Errorf("error creating search tool: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "creating search tool"); err != nil { + return err + } + + var searchToolResp searchToolAPIResponse + if err := json.NewDecoder(resp.Body).Decode(&searchToolResp); err != nil { + return fmt.Errorf("error decoding create search tool response: %w", err) + } + if searchToolResp.SearchToolID == "" { + return fmt.Errorf("create search tool response did not contain a search_tool_id") + } + + d.SetId(searchToolResp.SearchToolID) + log.Printf("[INFO] Search tool created with ID: %s", searchToolResp.SearchToolID) + + return resourceLiteLLMSearchToolRead(d, m) +} + +func resourceLiteLLMSearchToolRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Reading search tool with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf(endpointSearchToolByID, d.Id()), nil) + if err != nil { + return fmt.Errorf("error reading search tool: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + log.Printf("[WARN] Search tool with ID %s not found, removing from state", d.Id()) + d.SetId("") + return nil + } + + if err := handleResponse(resp, "reading search tool"); err != nil { + return err + } + + var searchToolResp searchToolAPIResponse + if err := json.NewDecoder(resp.Body).Decode(&searchToolResp); err != nil { + return fmt.Errorf("error decoding search tool info response: %w", err) + } + + d.Set("search_tool_name", searchToolResp.SearchToolName) + + // litellm_params is intentionally not read back: the API masks its values and it may hold secrets. + if searchToolResp.SearchToolInfo != nil { + infoJSON, err := json.Marshal(searchToolResp.SearchToolInfo) + if err != nil { + return fmt.Errorf("error encoding search_tool_info: %w", err) + } + d.Set("search_tool_info", string(infoJSON)) + } + d.Set("created_at", searchToolResp.CreatedAt) + d.Set("updated_at", searchToolResp.UpdatedAt) + + log.Printf("[INFO] Successfully read search tool with ID: %s", d.Id()) + return nil +} + +func resourceLiteLLMSearchToolUpdate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + searchToolData, err := buildSearchToolData(d) + if err != nil { + return err + } + searchToolData["search_tool_id"] = d.Id() + + log.Printf("[DEBUG] Update search tool request for ID: %s", d.Id()) + + resp, err := MakeRequest(client, "PUT", fmt.Sprintf(endpointSearchToolByID, d.Id()), map[string]interface{}{ + "search_tool": searchToolData, + }) + if err != nil { + return fmt.Errorf("error updating search tool: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "updating search tool"); err != nil { + return err + } + + log.Printf("[INFO] Successfully updated search tool with ID: %s", d.Id()) + return resourceLiteLLMSearchToolRead(d, m) +} + +func resourceLiteLLMSearchToolDelete(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Deleting search tool with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "DELETE", fmt.Sprintf(endpointSearchToolByID, d.Id()), nil) + if err != nil { + return fmt.Errorf("error deleting search tool: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusNotFound { + if err := handleResponse(resp, "deleting search tool"); err != nil { + return err + } + } + + log.Printf("[INFO] Successfully deleted search tool with ID: %s", d.Id()) + d.SetId("") + return nil +} diff --git a/terraform/provider/litellm/resource_search_tool_test.go b/terraform/provider/litellm/resource_search_tool_test.go new file mode 100644 index 00000000000..4435289ac86 --- /dev/null +++ b/terraform/provider/litellm/resource_search_tool_test.go @@ -0,0 +1,221 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const testSearchToolParamsJSON = `{"search_provider": "tavily", "api_key": "sk-secret"}` + +func newSearchToolTestResourceData(t *testing.T) *schema.ResourceData { + t.Helper() + return schema.TestResourceDataRaw(t, resourceLiteLLMSearchTool().Schema, map[string]interface{}{ + "search_tool_name": "my-search", + "litellm_params": testSearchToolParamsJSON, + "search_tool_info": `{"description": "Tavily search"}`, + }) +} + +func searchToolReadResponseBody() []byte { + body, _ := json.Marshal(map[string]interface{}{ + "search_tool_id": "st-123", + "search_tool_name": "my-search", + "litellm_params": map[string]interface{}{"search_provider": "tavily", "api_key": "sk-s****"}, + "search_tool_info": map[string]interface{}{"description": "Tavily search"}, + "created_at": "2026-01-01T00:00:00", + "updated_at": "2026-01-02T00:00:00", + }) + return body +} + +func TestResourceLiteLLMSearchToolCreate(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch { + case r.Method == http.MethodPost && r.URL.Path == "/search_tools": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(`{"search_tool_id": "st-123", "search_tool_name": "my-search"}`)) + case r.Method == http.MethodGet && r.URL.Path == "/search_tools/st-123": + w.Write(searchToolReadResponseBody()) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newSearchToolTestResourceData(t) + + if err := resourceLiteLLMSearchToolCreate(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "st-123" { + t.Fatalf("expected ID 'st-123', got %q", d.Id()) + } + + wrapped, ok := createPayload["search_tool"].(map[string]interface{}) + if !ok { + t.Fatalf("expected payload wrapped in 'search_tool', got %v", createPayload) + } + if wrapped["search_tool_name"] != "my-search" { + t.Errorf("expected search_tool_name 'my-search', got %v", wrapped["search_tool_name"]) + } + params, ok := wrapped["litellm_params"].(map[string]interface{}) + if !ok || params["search_provider"] != "tavily" || params["api_key"] != "sk-secret" { + t.Errorf("expected litellm_params sent as JSON object, got %v", wrapped["litellm_params"]) + } + info, ok := wrapped["search_tool_info"].(map[string]interface{}) + if !ok || info["description"] != "Tavily search" { + t.Errorf("expected search_tool_info sent as JSON object, got %v", wrapped["search_tool_info"]) + } + + if got := d.Get("litellm_params").(string); got != testSearchToolParamsJSON { + t.Errorf("expected litellm_params to keep configured value (masked API value not read back), got %q", got) + } + if d.Get("created_at").(string) != "2026-01-01T00:00:00" { + t.Errorf("expected created_at from read-back, got %q", d.Get("created_at").(string)) + } +} + +func TestResourceLiteLLMSearchToolCreateInvalidParamsJSON(t *testing.T) { + d := schema.TestResourceDataRaw(t, resourceLiteLLMSearchTool().Schema, map[string]interface{}{ + "search_tool_name": "my-search", + "litellm_params": "not-json", + }) + client := NewClient("http://unused.invalid", "test-key", true) + + if err := resourceLiteLLMSearchToolCreate(d, client); err == nil { + t.Fatal("expected error for invalid litellm_params JSON, got nil") + } +} + +func TestResourceLiteLLMSearchToolReadMapsFields(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet || r.URL.Path != "/search_tools/st-123" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.Write(searchToolReadResponseBody()) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMSearchTool().Schema, map[string]interface{}{}) + d.SetId("st-123") + + if err := resourceLiteLLMSearchToolRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Get("search_tool_name").(string) != "my-search" { + t.Errorf("expected search_tool_name 'my-search', got %q", d.Get("search_tool_name").(string)) + } + var info map[string]interface{} + if err := json.Unmarshal([]byte(d.Get("search_tool_info").(string)), &info); err != nil { + t.Fatalf("search_tool_info not populated as JSON: %v", err) + } + if info["description"] != "Tavily search" { + t.Errorf("expected description 'Tavily search', got %v", info["description"]) + } + if d.Get("litellm_params").(string) != "" { + t.Errorf("expected litellm_params to never be read back, got %q", d.Get("litellm_params").(string)) + } + if d.Get("updated_at").(string) != "2026-01-02T00:00:00" { + t.Errorf("expected updated_at from response, got %q", d.Get("updated_at").(string)) + } +} + +func TestResourceLiteLLMSearchToolRead404ClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newSearchToolTestResourceData(t) + d.SetId("st-123") + + if err := resourceLiteLLMSearchToolRead(d, client); err != nil { + t.Fatalf("expected nil error on 404, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared on 404, got %q", d.Id()) + } +} + +func TestResourceLiteLLMSearchToolUpdate(t *testing.T) { + var updateMethod, updatePath string + var updatePayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.Method == http.MethodGet { + w.Write(searchToolReadResponseBody()) + return + } + updateMethod = r.Method + updatePath = r.URL.Path + if err := json.NewDecoder(r.Body).Decode(&updatePayload); err != nil { + t.Errorf("failed to decode update payload: %v", err) + } + w.Write([]byte(`{}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newSearchToolTestResourceData(t) + d.SetId("st-123") + + if err := resourceLiteLLMSearchToolUpdate(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if updateMethod != http.MethodPut { + t.Errorf("expected PUT, got %s", updateMethod) + } + if updatePath != "/search_tools/st-123" { + t.Errorf("expected path '/search_tools/st-123', got %q", updatePath) + } + wrapped, ok := updatePayload["search_tool"].(map[string]interface{}) + if !ok { + t.Fatalf("expected payload wrapped in 'search_tool', got %v", updatePayload) + } + if wrapped["search_tool_id"] != "st-123" { + t.Errorf("expected search_tool_id in update payload, got %v", wrapped["search_tool_id"]) + } + if wrapped["search_tool_name"] != "my-search" { + t.Errorf("expected search_tool_name in update payload, got %v", wrapped["search_tool_name"]) + } +} + +func TestResourceLiteLLMSearchToolDelete(t *testing.T) { + var deleteMethod, deletePath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + deleteMethod = r.Method + deletePath = r.URL.Path + w.Write([]byte(`{}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newSearchToolTestResourceData(t) + d.SetId("st-123") + + if err := resourceLiteLLMSearchToolDelete(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if deleteMethod != http.MethodDelete { + t.Errorf("expected DELETE, got %s", deleteMethod) + } + if deletePath != "/search_tools/st-123" { + t.Errorf("expected path '/search_tools/st-123', got %q", deletePath) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared after delete, got %q", d.Id()) + } +} diff --git a/terraform/provider/litellm/resource_tag.go b/terraform/provider/litellm/resource_tag.go new file mode 100644 index 00000000000..dd1505541cb --- /dev/null +++ b/terraform/provider/litellm/resource_tag.go @@ -0,0 +1,285 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "io" + "log" + "net/http" + "strings" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const ( + endpointTagNew = "/tag/new" + endpointTagInfo = "/tag/info" + endpointTagUpdate = "/tag/update" + endpointTagDelete = "/tag/delete" +) + +type tagBudgetTable struct { + BudgetID string `json:"budget_id"` + MaxBudget *float64 `json:"max_budget"` + SoftBudget *float64 `json:"soft_budget"` + MaxParallelRequests *int `json:"max_parallel_requests"` + TPMLimit *int `json:"tpm_limit"` + RPMLimit *int `json:"rpm_limit"` + BudgetDuration string `json:"budget_duration"` +} + +type tagInfoEntry struct { + Name string `json:"name"` + Description string `json:"description"` + Models []string `json:"models"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` + CreatedBy string `json:"created_by"` + LitellmBudgetTable *tagBudgetTable `json:"litellm_budget_table"` +} + +func resourceLiteLLMTag() *schema.Resource { + return &schema.Resource{ + Create: resourceLiteLLMTagCreate, + Read: resourceLiteLLMTagRead, + Update: resourceLiteLLMTagUpdate, + Delete: resourceLiteLLMTagDelete, + + Importer: &schema.ResourceImporter{StateContext: schema.ImportStatePassthroughContext}, + + Schema: map[string]*schema.Schema{ + "name": { + Type: schema.TypeString, + Required: true, + ForceNew: true, + Description: "Unique name of the tag. Also used as the resource ID.", + }, + "description": { + Type: schema.TypeString, + Optional: true, + Description: "Description of the tag.", + }, + "models": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "List of model IDs this tag applies to.", + }, + "budget_id": { + Type: schema.TypeString, + Optional: true, + Description: "Existing budget ID to associate with this tag.", + }, + "max_budget": { + Type: schema.TypeFloat, + Optional: true, + Description: "Max budget in USD for this tag.", + }, + "soft_budget": { + Type: schema.TypeFloat, + Optional: true, + Description: "Soft budget in USD for this tag.", + }, + "max_parallel_requests": { + Type: schema.TypeInt, + Optional: true, + Description: "Max concurrent requests allowed for this tag.", + }, + "tpm_limit": { + Type: schema.TypeInt, + Optional: true, + Description: "Max tokens per minute for this tag.", + }, + "rpm_limit": { + Type: schema.TypeInt, + Optional: true, + Description: "Max requests per minute for this tag.", + }, + "budget_duration": { + Type: schema.TypeString, + Optional: true, + Description: "Duration for budget reset (e.g. '1h', '1d', '30d').", + }, + "model_max_budget": { + Type: schema.TypeString, + Optional: true, + Description: "JSON object string with per-model budget configuration.", + }, + }, + } +} + +func buildTagData(d *schema.ResourceData, name string) (map[string]interface{}, error) { + tagData := map[string]interface{}{ + "name": name, + } + + for _, key := range []string{"description", "models", "budget_id", "max_budget", "soft_budget", + "max_parallel_requests", "tpm_limit", "rpm_limit", "budget_duration"} { + if v, ok := d.GetOk(key); ok { + tagData[key] = v + } + } + + if v, ok := d.GetOk("model_max_budget"); ok { + var modelMaxBudget map[string]interface{} + if err := json.Unmarshal([]byte(v.(string)), &modelMaxBudget); err != nil { + return nil, fmt.Errorf("model_max_budget must be a JSON object: %w", err) + } + tagData["model_max_budget"] = modelMaxBudget + } + + return tagData, nil +} + +// fetchTagInfo returns the tag entry, or gone=true when the proxy reports the tag missing. +func fetchTagInfo(client *Client, name string) (*tagInfoEntry, bool, error) { + resp, err := MakeRequest(client, "POST", endpointTagInfo, map[string]interface{}{ + "names": []string{name}, + }) + if err != nil { + return nil, false, err + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + if err != nil { + return nil, false, fmt.Errorf("failed to read tag info response: %w", err) + } + + if resp.StatusCode == http.StatusNotFound || + (resp.StatusCode != http.StatusOK && strings.Contains(string(body), "Tags not found")) { + return nil, true, nil + } + if resp.StatusCode != http.StatusOK { + return nil, false, fmt.Errorf("error reading tag: %s - %s", resp.Status, string(body)) + } + + var tags map[string]tagInfoEntry + if err := json.Unmarshal(body, &tags); err != nil { + return nil, false, fmt.Errorf("error decoding tag info response: %w", err) + } + + entry, ok := tags[name] + if !ok { + return nil, true, nil + } + return &entry, false, nil +} + +func resourceLiteLLMTagCreate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + name := d.Get("name").(string) + tagData, err := buildTagData(d, name) + if err != nil { + return err + } + + log.Printf("[DEBUG] Create tag request payload: %+v", tagData) + + resp, err := MakeRequest(client, "POST", endpointTagNew, tagData) + if err != nil { + return fmt.Errorf("error creating tag: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "creating tag"); err != nil { + return err + } + + d.SetId(name) + log.Printf("[INFO] Tag created with name: %s", name) + + return resourceLiteLLMTagRead(d, m) +} + +func resourceLiteLLMTagRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Reading tag with name: %s", d.Id()) + + entry, gone, err := fetchTagInfo(client, d.Id()) + if err != nil { + return err + } + if gone { + log.Printf("[WARN] Tag %s not found, removing from state", d.Id()) + d.SetId("") + return nil + } + + d.Set("name", d.Id()) + d.Set("description", GetStringValue(entry.Description, d.Get("description").(string))) + if entry.Models != nil { + d.Set("models", entry.Models) + } + + if bt := entry.LitellmBudgetTable; bt != nil { + d.Set("budget_id", GetStringValue(bt.BudgetID, d.Get("budget_id").(string))) + if bt.MaxBudget != nil { + d.Set("max_budget", *bt.MaxBudget) + } + if bt.SoftBudget != nil { + d.Set("soft_budget", *bt.SoftBudget) + } + if bt.MaxParallelRequests != nil { + d.Set("max_parallel_requests", *bt.MaxParallelRequests) + } + if bt.TPMLimit != nil { + d.Set("tpm_limit", *bt.TPMLimit) + } + if bt.RPMLimit != nil { + d.Set("rpm_limit", *bt.RPMLimit) + } + d.Set("budget_duration", GetStringValue(bt.BudgetDuration, d.Get("budget_duration").(string))) + } + + log.Printf("[INFO] Successfully read tag with name: %s", d.Id()) + return nil +} + +func resourceLiteLLMTagUpdate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + tagData, err := buildTagData(d, d.Id()) + if err != nil { + return err + } + log.Printf("[DEBUG] Update tag request payload: %+v", tagData) + + resp, err := MakeRequest(client, "POST", endpointTagUpdate, tagData) + if err != nil { + return fmt.Errorf("error updating tag: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "updating tag"); err != nil { + return err + } + + log.Printf("[INFO] Successfully updated tag with name: %s", d.Id()) + return resourceLiteLLMTagRead(d, m) +} + +func resourceLiteLLMTagDelete(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Deleting tag with name: %s", d.Id()) + + resp, err := MakeRequest(client, "POST", endpointTagDelete, map[string]interface{}{ + "name": d.Id(), + }) + if err != nil { + return fmt.Errorf("error deleting tag: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "deleting tag"); err != nil { + return err + } + + log.Printf("[INFO] Successfully deleted tag with name: %s", d.Id()) + d.SetId("") + return nil +} diff --git a/terraform/provider/litellm/resource_tag_test.go b/terraform/provider/litellm/resource_tag_test.go new file mode 100644 index 00000000000..f6bcc6d74ab --- /dev/null +++ b/terraform/provider/litellm/resource_tag_test.go @@ -0,0 +1,245 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "reflect" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func tagInfoBody(name string) string { + return `{"` + name + `": { + "name": "` + name + `", + "description": "Production traffic", + "models": ["model-1", "model-2"], + "created_at": "2026-01-01T00:00:00", + "updated_at": "2026-01-02T00:00:00", + "created_by": "admin", + "litellm_budget_table": { + "budget_id": "bud-1", + "max_budget": 50.5, + "soft_budget": 40.0, + "max_parallel_requests": 5, + "tpm_limit": 1000, + "rpm_limit": 100, + "budget_duration": "30d" + } + }}` +} + +func TestResourceLiteLLMTagCreate(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/tag/new": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write([]byte(`{"message": "created"}`)) + case "/tag/info": + w.Write([]byte(tagInfoBody("prod"))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMTag().Schema, map[string]interface{}{ + "name": "prod", + "description": "Production traffic", + "models": []interface{}{"model-1", "model-2"}, + "max_budget": 50.5, + "tpm_limit": 1000, + "model_max_budget": `{"gpt-4": {"budget_limit": 10}}`, + }) + + if err := resourceLiteLLMTagCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + if d.Id() != "prod" { + t.Fatalf("expected ID 'prod', got %q", d.Id()) + } + if createPayload["name"] != "prod" { + t.Errorf("expected payload name 'prod', got %v", createPayload["name"]) + } + if createPayload["description"] != "Production traffic" { + t.Errorf("expected payload description, got %v", createPayload["description"]) + } + if !reflect.DeepEqual(createPayload["models"], []interface{}{"model-1", "model-2"}) { + t.Errorf("expected payload models, got %v", createPayload["models"]) + } + if createPayload["max_budget"] != 50.5 { + t.Errorf("expected payload max_budget 50.5, got %v", createPayload["max_budget"]) + } + if createPayload["tpm_limit"] != float64(1000) { + t.Errorf("expected payload tpm_limit 1000, got %v", createPayload["tpm_limit"]) + } + modelMaxBudget, ok := createPayload["model_max_budget"].(map[string]interface{}) + if !ok || modelMaxBudget["gpt-4"] == nil { + t.Errorf("expected model_max_budget sent as JSON object, got %v", createPayload["model_max_budget"]) + } + if got := d.Get("budget_id").(string); got != "bud-1" { + t.Errorf("expected budget_id 'bud-1' from read, got %q", got) + } +} + +func TestResourceLiteLLMTagCreate_InvalidModelMaxBudget(t *testing.T) { + d := schema.TestResourceDataRaw(t, resourceLiteLLMTag().Schema, map[string]interface{}{ + "name": "prod", + "model_max_budget": "not-json", + }) + + if err := resourceLiteLLMTagCreate(d, NewClient("http://127.0.0.1:1", "test-key", true)); err == nil { + t.Fatal("expected error for invalid model_max_budget JSON, got nil") + } +} + +func TestResourceLiteLLMTagRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/tag/info" { + t.Errorf("unexpected request path: %s", r.URL.Path) + } + var payload map[string]interface{} + json.NewDecoder(r.Body).Decode(&payload) + if !reflect.DeepEqual(payload["names"], []interface{}{"prod"}) { + t.Errorf("expected names ['prod'], got %v", payload["names"]) + } + w.Write([]byte(tagInfoBody("prod"))) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMTag().Schema, map[string]interface{}{"name": "prod"}) + d.SetId("prod") + + if err := resourceLiteLLMTagRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + checks := map[string]interface{}{ + "description": "Production traffic", + "budget_id": "bud-1", + "max_budget": 50.5, + "soft_budget": 40.0, + "max_parallel_requests": 5, + "tpm_limit": 1000, + "rpm_limit": 100, + "budget_duration": "30d", + } + for key, want := range checks { + if got := d.Get(key); got != want { + t.Errorf("expected %s %v, got %v", key, want, got) + } + } + if !reflect.DeepEqual(d.Get("models"), []interface{}{"model-1", "model-2"}) { + t.Errorf("expected models in state, got %v", d.Get("models")) + } +} + +func TestResourceLiteLLMTagRead_404ClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMTag().Schema, map[string]interface{}{"name": "gone"}) + d.SetId("gone") + + if err := resourceLiteLLMTagRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("expected nil error on 404, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared on 404, got %q", d.Id()) + } +} + +// The proxy wraps its internal 404 into a 500 whose detail mentions "Tags not found". +func TestResourceLiteLLMTagRead_WrappedNotFoundClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"detail": "404: Tags not found: ['gone']"}`)) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMTag().Schema, map[string]interface{}{"name": "gone"}) + d.SetId("gone") + + if err := resourceLiteLLMTagRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("expected nil error on wrapped not-found, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared on wrapped not-found, got %q", d.Id()) + } +} + +func TestResourceLiteLLMTagUpdate(t *testing.T) { + var updatePayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/tag/update": + if err := json.NewDecoder(r.Body).Decode(&updatePayload); err != nil { + t.Errorf("failed to decode update payload: %v", err) + } + w.Write([]byte(`{"message": "updated"}`)) + case "/tag/info": + w.Write([]byte(tagInfoBody("prod"))) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMTag().Schema, map[string]interface{}{ + "name": "prod", + "description": "Updated description", + "rpm_limit": 200, + }) + d.SetId("prod") + + if err := resourceLiteLLMTagUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + + if updatePayload["name"] != "prod" { + t.Errorf("expected update payload name 'prod', got %v", updatePayload["name"]) + } + if updatePayload["description"] != "Updated description" { + t.Errorf("expected updated description in payload, got %v", updatePayload["description"]) + } + if updatePayload["rpm_limit"] != float64(200) { + t.Errorf("expected rpm_limit 200 in payload, got %v", updatePayload["rpm_limit"]) + } +} + +func TestResourceLiteLLMTagDelete(t *testing.T) { + var deletePayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/tag/delete" || r.Method != http.MethodPost { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + } + if err := json.NewDecoder(r.Body).Decode(&deletePayload); err != nil { + t.Errorf("failed to decode delete payload: %v", err) + } + w.Write([]byte(`{"message": "deleted"}`)) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMTag().Schema, map[string]interface{}{"name": "prod"}) + d.SetId("prod") + + if err := resourceLiteLLMTagDelete(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("delete failed: %v", err) + } + + if deletePayload["name"] != "prod" { + t.Errorf("expected delete payload name 'prod', got %v", deletePayload["name"]) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared after delete, got %q", d.Id()) + } +} diff --git a/terraform/provider/litellm/resource_team.go b/terraform/provider/litellm/resource_team.go index 2a167a1b5c4..24c47843cd1 100644 --- a/terraform/provider/litellm/resource_team.go +++ b/terraform/provider/litellm/resource_team.go @@ -26,6 +26,9 @@ func ResourceLiteLLMTeam() *schema.Resource { Read: resourceLiteLLMTeamRead, Update: resourceLiteLLMTeamUpdate, Delete: resourceLiteLLMTeamDelete, + Importer: &schema.ResourceImporter{ + StateContext: schema.ImportStatePassthroughContext, + }, Schema: map[string]*schema.Schema{ "team_alias": { @@ -89,6 +92,69 @@ func ResourceLiteLLMTeam() *schema.Resource { Elem: &schema.Schema{Type: schema.TypeString}, Description: "Email addresses alerted when the team crosses soft_budget", }, + "model_aliases": { + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "guardrails": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "prompts": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "team_member_budget": { + Type: schema.TypeFloat, + Optional: true, + Description: "Budget applied to every team member", + }, + "team_member_budget_duration": { + Type: schema.TypeString, + Optional: true, + }, + "team_member_rpm_limit": { + Type: schema.TypeInt, + Optional: true, + }, + "team_member_tpm_limit": { + Type: schema.TypeInt, + Optional: true, + }, + "team_member_key_duration": { + Type: schema.TypeString, + Optional: true, + }, + "model_rpm_limit": { + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeInt}, + }, + "model_tpm_limit": { + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeInt}, + }, + "allowed_passthrough_routes": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "rpm_limit_type": { + Type: schema.TypeString, + Optional: true, + ForceNew: true, + Description: "One of 'guaranteed_throughput' or 'best_effort_throughput'; only settable at creation", + }, + "tpm_limit_type": { + Type: schema.TypeString, + Optional: true, + ForceNew: true, + Description: "One of 'guaranteed_throughput' or 'best_effort_throughput'; only settable at creation", + }, }, } } @@ -99,6 +165,13 @@ func resourceLiteLLMTeamCreate(d *schema.ResourceData, m interface{}) error { teamID := uuid.New().String() teamData := buildTeamData(d, teamID) + // Throughput limit types are only accepted by /team/new, not /team/update. + for _, key := range []string{"rpm_limit_type", "tpm_limit_type"} { + if v, ok := d.GetOk(key); ok { + teamData[key] = v + } + } + log.Printf("[DEBUG] Create team request payload: %+v", teamData) resp, err := MakeRequest(client, "POST", endpointTeamNew, teamData) @@ -170,6 +243,36 @@ func resourceLiteLLMTeamRead(d *schema.ResourceData, m interface{}) error { d.Set("blocked", GetBoolValue(teamResp.Blocked, d.Get("blocked").(bool))) + if teamResp.ModelAliases != nil { + d.Set("model_aliases", teamResp.ModelAliases) + } + if teamResp.Guardrails != nil { + d.Set("guardrails", teamResp.Guardrails) + } + if teamResp.Prompts != nil { + d.Set("prompts", teamResp.Prompts) + } + if teamResp.TeamMemberBudget != nil { + d.Set("team_member_budget", *teamResp.TeamMemberBudget) + } + d.Set("team_member_budget_duration", GetStringValue(teamResp.TeamMemberBudgetDuration, d.Get("team_member_budget_duration").(string))) + if teamResp.TeamMemberRPMLimit != nil { + d.Set("team_member_rpm_limit", *teamResp.TeamMemberRPMLimit) + } + if teamResp.TeamMemberTPMLimit != nil { + d.Set("team_member_tpm_limit", *teamResp.TeamMemberTPMLimit) + } + d.Set("team_member_key_duration", GetStringValue(teamResp.TeamMemberKeyDuration, d.Get("team_member_key_duration").(string))) + if teamResp.ModelRPMLimit != nil { + d.Set("model_rpm_limit", teamResp.ModelRPMLimit) + } + if teamResp.ModelTPMLimit != nil { + d.Set("model_tpm_limit", teamResp.ModelTPMLimit) + } + if teamResp.AllowedPassthroughRoutes != nil { + d.Set("allowed_passthrough_routes", teamResp.AllowedPassthroughRoutes) + } + // Explicitly fetch the current permissions from the API permResp, err := getTeamPermissions(client, d.Id()) if err != nil { @@ -257,7 +360,13 @@ func buildTeamData(d *schema.ResourceData, teamID string) map[string]interface{} "team_alias": d.Get("team_alias").(string), } - for _, key := range []string{"organization_id", "tpm_limit", "rpm_limit", "max_budget", "budget_duration", "models", "blocked", "team_member_permissions"} { + for _, key := range []string{ + "organization_id", "tpm_limit", "rpm_limit", "max_budget", "budget_duration", "models", + "blocked", "team_member_permissions", "model_aliases", "guardrails", "prompts", + "team_member_budget", "team_member_budget_duration", "team_member_rpm_limit", + "team_member_tpm_limit", "team_member_key_duration", "model_rpm_limit", + "model_tpm_limit", "allowed_passthrough_routes", + } { if v, ok := d.GetOk(key); ok { teamData[key] = v } diff --git a/terraform/provider/litellm/resource_team_block.go b/terraform/provider/litellm/resource_team_block.go new file mode 100644 index 00000000000..e3e35520257 --- /dev/null +++ b/terraform/provider/litellm/resource_team_block.go @@ -0,0 +1,127 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "log" + "net/http" + "net/url" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const ( + endpointTeamBlock = "/team/block" + endpointTeamUnblock = "/team/unblock" +) + +type TeamBlockInfoResponse struct { + TeamInfo struct { + Blocked *bool `json:"blocked"` + } `json:"team_info"` +} + +func resourceLiteLLMTeamBlock() *schema.Resource { + return &schema.Resource{ + Create: resourceLiteLLMTeamBlockCreate, + Read: resourceLiteLLMTeamBlockRead, + Delete: resourceLiteLLMTeamBlockDelete, + + Importer: &schema.ResourceImporter{ + StateContext: schema.ImportStatePassthroughContext, + }, + + Schema: map[string]*schema.Schema{ + "team_id": { + Type: schema.TypeString, + Required: true, + ForceNew: true, + Description: "The ID of the team to block. Destroying this resource unblocks the team", + }, + "blocked": { + Type: schema.TypeBool, + Computed: true, + Description: "Whether the team is currently blocked", + }, + }, + } +} + +func resourceLiteLLMTeamBlockCreate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + teamID := d.Get("team_id").(string) + + log.Printf("[INFO] Blocking team with ID: %s", teamID) + + resp, err := MakeRequest(client, "POST", endpointTeamBlock, map[string]interface{}{"team_id": teamID}) + if err != nil { + return fmt.Errorf("error blocking team: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "blocking team"); err != nil { + return err + } + + d.SetId(teamID) + return resourceLiteLLMTeamBlockRead(d, m) +} + +func resourceLiteLLMTeamBlockRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + teamID := d.Id() + + log.Printf("[INFO] Reading block state for team with ID: %s", teamID) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("/team/info?team_id=%s", url.QueryEscape(teamID)), nil) + if err != nil { + return fmt.Errorf("error reading team info: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + log.Printf("[WARN] Team with ID %s not found, removing team block from state", teamID) + d.SetId("") + return nil + } + + if err := handleResponse(resp, "reading team info"); err != nil { + return err + } + + var infoResp TeamBlockInfoResponse + if err := json.NewDecoder(resp.Body).Decode(&infoResp); err != nil { + return fmt.Errorf("error decoding team info response: %w", err) + } + + if infoResp.TeamInfo.Blocked == nil || !*infoResp.TeamInfo.Blocked { + log.Printf("[WARN] Team with ID %s is no longer blocked, removing team block from state", teamID) + d.SetId("") + return nil + } + + d.Set("team_id", teamID) + d.Set("blocked", true) + return nil +} + +func resourceLiteLLMTeamBlockDelete(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Unblocking team with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "POST", endpointTeamUnblock, map[string]interface{}{"team_id": d.Id()}) + if err != nil { + return fmt.Errorf("error unblocking team: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusNotFound { + if err := handleResponse(resp, "unblocking team"); err != nil { + return err + } + } + + d.SetId("") + return nil +} diff --git a/terraform/provider/litellm/resource_team_block_test.go b/terraform/provider/litellm/resource_team_block_test.go new file mode 100644 index 00000000000..7c37e1a5af8 --- /dev/null +++ b/terraform/provider/litellm/resource_team_block_test.go @@ -0,0 +1,123 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func newTeamBlockTestResourceData(t *testing.T, teamID string) *schema.ResourceData { + t.Helper() + return schema.TestResourceDataRaw(t, resourceLiteLLMTeamBlock().Schema, map[string]interface{}{ + "team_id": teamID, + }) +} + +func TestResourceLiteLLMTeamBlockCreate(t *testing.T) { + var blockPayload map[string]interface{} + mux := http.NewServeMux() + mux.HandleFunc("/team/block", func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + t.Errorf("expected POST, got %s", r.Method) + } + if err := json.NewDecoder(r.Body).Decode(&blockPayload); err != nil { + t.Fatalf("failed to decode block payload: %v", err) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"team_id":"team-123","blocked":true}`)) + }) + mux.HandleFunc("/team/info", func(w http.ResponseWriter, r *http.Request) { + if got := r.URL.Query().Get("team_id"); got != "team-123" { + t.Errorf("expected team_id query 'team-123', got %q", got) + } + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"team_id":"team-123","team_info":{"blocked":true}}`)) + }) + srv := httptest.NewServer(mux) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newTeamBlockTestResourceData(t, "team-123") + + if err := resourceLiteLLMTeamBlockCreate(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "team-123" { + t.Fatalf("expected ID 'team-123', got %q", d.Id()) + } + if blockPayload["team_id"] != "team-123" { + t.Fatalf("expected block payload team_id 'team-123', got %+v", blockPayload) + } + if !d.Get("blocked").(bool) { + t.Fatal("expected blocked=true in state") + } +} + +func TestResourceLiteLLMTeamBlockRead_UnblockedClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"team_id":"team-123","team_info":{"blocked":false}}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newTeamBlockTestResourceData(t, "team-123") + d.SetId("team-123") + + if err := resourceLiteLLMTeamBlockRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared for unblocked team, got %q", d.Id()) + } +} + +func TestResourceLiteLLMTeamBlockRead_404ClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newTeamBlockTestResourceData(t, "team-123") + d.SetId("team-123") + + if err := resourceLiteLLMTeamBlockRead(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared on 404, got %q", d.Id()) + } +} + +func TestResourceLiteLLMTeamBlockDelete(t *testing.T) { + var gotPath string + var unblockPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + json.NewDecoder(r.Body).Decode(&unblockPayload) + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"team_id":"team-123","blocked":false}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newTeamBlockTestResourceData(t, "team-123") + d.SetId("team-123") + + if err := resourceLiteLLMTeamBlockDelete(d, client); err != nil { + t.Fatalf("expected nil error, got: %v", err) + } + if gotPath != "/team/unblock" { + t.Fatalf("expected path /team/unblock, got %s", gotPath) + } + if unblockPayload["team_id"] != "team-123" { + t.Fatalf("expected unblock payload team_id 'team-123', got %+v", unblockPayload) + } + if d.Id() != "" { + t.Fatalf("expected ID cleared after delete, got %q", d.Id()) + } +} diff --git a/terraform/provider/litellm/resource_team_test.go b/terraform/provider/litellm/resource_team_test.go index 1f74be4819d..9638378cdfe 100644 --- a/terraform/provider/litellm/resource_team_test.go +++ b/terraform/provider/litellm/resource_team_test.go @@ -182,3 +182,102 @@ func TestTeamReadClearsSoftBudgetWhenProxyReturnsNull(t *testing.T) { t.Fatalf("soft_budget = %v, want cleared after the proxy returned null", got) } } + +func newTeamResourceData(t *testing.T, raw map[string]interface{}) *schema.ResourceData { + t.Helper() + return schema.TestResourceDataRaw(t, ResourceLiteLLMTeam().Schema, raw) +} + +func TestBuildTeamDataIncludesNewFields(t *testing.T) { + d := newTeamResourceData(t, map[string]interface{}{ + "team_alias": "eng", + "model_aliases": map[string]interface{}{"gpt": "gpt-5.2"}, + "guardrails": []interface{}{"pii-mask"}, + "prompts": []interface{}{"prompt-1"}, + "team_member_budget": 5.0, + "team_member_budget_duration": "30d", + "team_member_rpm_limit": 10, + "team_member_tpm_limit": 1000, + "team_member_key_duration": "7d", + "allowed_passthrough_routes": []interface{}{"/vertex-ai"}, + }) + + data := buildTeamData(d, "team-1") + + for _, k := range []string{ + "model_aliases", "guardrails", "prompts", "team_member_budget", + "team_member_budget_duration", "team_member_rpm_limit", "team_member_tpm_limit", + "team_member_key_duration", "allowed_passthrough_routes", + } { + if _, ok := data[k]; !ok { + t.Errorf("buildTeamData missing %s", k) + } + } + if data["team_id"] != "team-1" || data["team_alias"] != "eng" { + t.Errorf("identity fields wrong: %v", data) + } +} + +func TestTeamReadMapsNewFields(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{ + "team_id": "team-1", + "team_info": { + "team_id": "team-1", + "team_alias": "eng", + "guardrails": ["pii-mask"], + "team_member_budget": 5.0, + "team_member_rpm_limit": 10 + } + }`) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newTeamResourceData(t, map[string]interface{}{"team_alias": "config-alias"}) + d.SetId("team-1") + + if err := resourceLiteLLMTeamRead(d, client); err != nil { + t.Fatalf("read returned error: %v", err) + } + if got := d.Get("guardrails").([]interface{}); len(got) != 1 || got[0] != "pii-mask" { + t.Errorf("guardrails = %v, want [pii-mask]", got) + } + if got := d.Get("team_member_budget").(float64); got != 5.0 { + t.Errorf("team_member_budget = %v, want 5.0", got) + } + if got := d.Get("team_member_rpm_limit").(int); got != 10 { + t.Errorf("team_member_rpm_limit = %v, want 10", got) + } +} + +// rpm_limit_type / tpm_limit_type are accepted by /team/new but not +// /team/update, so create must send them and update must not. +func TestTeamLimitTypesSentOnCreateOnly(t *testing.T) { + var captured map[string]interface{} + srv := newTeamTestServer(t, &captured, `{"team_id": "x", "team_info": {"team_alias": "eng"}}`) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := newTeamResourceData(t, map[string]interface{}{ + "team_alias": "eng", + "rpm_limit_type": "guaranteed_throughput", + "tpm_limit_type": "best_effort_throughput", + }) + + if err := resourceLiteLLMTeamCreate(d, client); err != nil { + t.Fatalf("create returned error: %v", err) + } + if captured["rpm_limit_type"] != "guaranteed_throughput" || captured["tpm_limit_type"] != "best_effort_throughput" { + t.Errorf("create payload missing limit types: %v", captured) + } + + captured = nil + if err := resourceLiteLLMTeamUpdate(d, client); err != nil { + t.Fatalf("update returned error: %v", err) + } + for _, k := range []string{"rpm_limit_type", "tpm_limit_type"} { + if _, present := captured[k]; present { + t.Errorf("update payload unexpectedly contains %s", k) + } + } +} diff --git a/terraform/provider/litellm/resource_unified_access_group.go b/terraform/provider/litellm/resource_unified_access_group.go new file mode 100644 index 00000000000..0b2a67ebf23 --- /dev/null +++ b/terraform/provider/litellm/resource_unified_access_group.go @@ -0,0 +1,246 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "io" + "log" + "net/http" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +const endpointUnifiedAccessGroupCreate = "/v1/unified_access_group" + +var unifiedAccessGroupListFields = []string{ + "access_model_names", + "access_mcp_server_ids", + "access_agent_ids", + "assigned_team_ids", + "assigned_key_ids", +} + +type unifiedAccessGroupResponse struct { + AccessGroupID string `json:"access_group_id"` + AccessGroupName string `json:"access_group_name"` + Description *string `json:"description"` + AccessModelNames []string `json:"access_model_names"` + AccessMCPServerIDs []string `json:"access_mcp_server_ids"` + AccessAgentIDs []string `json:"access_agent_ids"` + AssignedTeamIDs []string `json:"assigned_team_ids"` + AssignedKeyIDs []string `json:"assigned_key_ids"` + CreatedAt string `json:"created_at"` + CreatedBy *string `json:"created_by"` + UpdatedAt string `json:"updated_at"` + UpdatedBy *string `json:"updated_by"` +} + +func resourceLiteLLMUnifiedAccessGroup() *schema.Resource { + return &schema.Resource{ + Create: resourceLiteLLMUnifiedAccessGroupCreate, + Read: resourceLiteLLMUnifiedAccessGroupRead, + Update: resourceLiteLLMUnifiedAccessGroupUpdate, + Delete: resourceLiteLLMUnifiedAccessGroupDelete, + + Importer: &schema.ResourceImporter{StateContext: schema.ImportStatePassthroughContext}, + + Schema: map[string]*schema.Schema{ + "access_group_name": { + Type: schema.TypeString, + Required: true, + }, + "description": { + Type: schema.TypeString, + Optional: true, + }, + "access_model_names": { + Type: schema.TypeList, + Optional: true, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "access_mcp_server_ids": { + Type: schema.TypeList, + Optional: true, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "access_agent_ids": { + Type: schema.TypeList, + Optional: true, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "assigned_team_ids": { + Type: schema.TypeList, + Optional: true, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "assigned_key_ids": { + Type: schema.TypeList, + Optional: true, + Computed: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, + "access_group_id": { + Type: schema.TypeString, + Computed: true, + }, + "created_at": { + Type: schema.TypeString, + Computed: true, + }, + "created_by": { + Type: schema.TypeString, + Computed: true, + }, + "updated_at": { + Type: schema.TypeString, + Computed: true, + }, + "updated_by": { + Type: schema.TypeString, + Computed: true, + }, + }, + } +} + +func buildUnifiedAccessGroupData(d *schema.ResourceData) map[string]interface{} { + data := map[string]interface{}{ + "access_group_name": d.Get("access_group_name").(string), + } + if v, ok := d.GetOk("description"); ok { + data["description"] = v + } + for _, key := range unifiedAccessGroupListFields { + data[key] = d.Get(key) + } + return data +} + +func setUnifiedAccessGroupFields(d *schema.ResourceData, group unifiedAccessGroupResponse) { + d.Set("access_group_id", group.AccessGroupID) + d.Set("access_group_name", group.AccessGroupName) + if group.Description != nil { + d.Set("description", *group.Description) + } + d.Set("access_model_names", group.AccessModelNames) + d.Set("access_mcp_server_ids", group.AccessMCPServerIDs) + d.Set("access_agent_ids", group.AccessAgentIDs) + d.Set("assigned_team_ids", group.AssignedTeamIDs) + d.Set("assigned_key_ids", group.AssignedKeyIDs) + d.Set("created_at", group.CreatedAt) + if group.CreatedBy != nil { + d.Set("created_by", *group.CreatedBy) + } + d.Set("updated_at", group.UpdatedAt) + if group.UpdatedBy != nil { + d.Set("updated_by", *group.UpdatedBy) + } +} + +func resourceLiteLLMUnifiedAccessGroupCreate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + groupData := buildUnifiedAccessGroupData(d) + log.Printf("[DEBUG] Create unified access group request payload: %+v", groupData) + + resp, err := MakeRequest(client, "POST", endpointUnifiedAccessGroupCreate, groupData) + if err != nil { + return fmt.Errorf("error creating unified access group: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "creating unified access group"); err != nil { + return err + } + + var group unifiedAccessGroupResponse + if err := json.NewDecoder(resp.Body).Decode(&group); err != nil { + return fmt.Errorf("error decoding unified access group create response: %w", err) + } + + if group.AccessGroupID == "" { + return fmt.Errorf("unified access group create response missing access_group_id") + } + + d.SetId(group.AccessGroupID) + log.Printf("[INFO] Unified access group created with ID: %s", group.AccessGroupID) + + return resourceLiteLLMUnifiedAccessGroupRead(d, m) +} + +func resourceLiteLLMUnifiedAccessGroupRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Reading unified access group with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("/v1/unified_access_group/%s", d.Id()), nil) + if err != nil { + return fmt.Errorf("error reading unified access group: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + log.Printf("[WARN] Unified access group with ID %s not found, removing from state", d.Id()) + d.SetId("") + return nil + } + + if err := handleResponse(resp, "reading unified access group"); err != nil { + return err + } + + var group unifiedAccessGroupResponse + if err := json.NewDecoder(resp.Body).Decode(&group); err != nil { + return fmt.Errorf("error decoding unified access group info response: %w", err) + } + + setUnifiedAccessGroupFields(d, group) + + log.Printf("[INFO] Successfully read unified access group with ID: %s", d.Id()) + return nil +} + +func resourceLiteLLMUnifiedAccessGroupUpdate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + groupData := buildUnifiedAccessGroupData(d) + log.Printf("[DEBUG] Update unified access group request payload: %+v", groupData) + + resp, err := MakeRequest(client, "PUT", fmt.Sprintf("/v1/unified_access_group/%s", d.Id()), groupData) + if err != nil { + return fmt.Errorf("error updating unified access group: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "updating unified access group"); err != nil { + return err + } + + log.Printf("[INFO] Successfully updated unified access group with ID: %s", d.Id()) + return resourceLiteLLMUnifiedAccessGroupRead(d, m) +} + +func resourceLiteLLMUnifiedAccessGroupDelete(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Deleting unified access group with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "DELETE", fmt.Sprintf("/v1/unified_access_group/%s", d.Id()), nil) + if err != nil { + return fmt.Errorf("error deleting unified access group: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusNoContent { + body, _ := io.ReadAll(resp.Body) + return fmt.Errorf("error deleting unified access group: %s - %s", resp.Status, string(body)) + } + + log.Printf("[INFO] Successfully deleted unified access group with ID: %s", d.Id()) + d.SetId("") + return nil +} diff --git a/terraform/provider/litellm/resource_unified_access_group_test.go b/terraform/provider/litellm/resource_unified_access_group_test.go new file mode 100644 index 00000000000..39ff2d24f74 --- /dev/null +++ b/terraform/provider/litellm/resource_unified_access_group_test.go @@ -0,0 +1,209 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "reflect" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func unifiedAccessGroupTestData(t *testing.T, raw map[string]interface{}) *schema.ResourceData { + t.Helper() + return schema.TestResourceDataRaw(t, resourceLiteLLMUnifiedAccessGroup().Schema, raw) +} + +func unifiedAccessGroupJSON(id string) []byte { + description := "prod access" + createdBy := "admin" + body, _ := json.Marshal(unifiedAccessGroupResponse{ + AccessGroupID: id, + AccessGroupName: "prod-group", + Description: &description, + AccessModelNames: []string{"gpt-4"}, + AccessMCPServerIDs: []string{"mcp-1"}, + AccessAgentIDs: []string{"agent-1"}, + AssignedTeamIDs: []string{"team-1"}, + AssignedKeyIDs: []string{"key-1"}, + CreatedAt: "2026-01-01T00:00:00Z", + CreatedBy: &createdBy, + UpdatedAt: "2026-01-02T00:00:00Z", + }) + return body +} + +func TestUnifiedAccessGroupCreate(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.Method + " " + r.URL.Path { + case "POST /v1/unified_access_group": + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Errorf("failed to decode create payload: %v", err) + } + w.Write(unifiedAccessGroupJSON("uag-123")) + case "GET /v1/unified_access_group/uag-123": + w.Write(unifiedAccessGroupJSON("uag-123")) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := unifiedAccessGroupTestData(t, map[string]interface{}{ + "access_group_name": "prod-group", + "description": "prod access", + "access_model_names": []interface{}{"gpt-4"}, + "assigned_team_ids": []interface{}{"team-1"}, + }) + + if err := resourceLiteLLMUnifiedAccessGroupCreate(d, client); err != nil { + t.Fatalf("create failed: %v", err) + } + + if createPayload["access_group_name"] != "prod-group" { + t.Fatalf("expected access_group_name 'prod-group' in payload, got %v", createPayload["access_group_name"]) + } + if createPayload["description"] != "prod access" { + t.Fatalf("expected description 'prod access' in payload, got %v", createPayload["description"]) + } + if !reflect.DeepEqual(createPayload["access_model_names"], []interface{}{"gpt-4"}) { + t.Fatalf("expected access_model_names [gpt-4] in payload, got %v", createPayload["access_model_names"]) + } + if !reflect.DeepEqual(createPayload["assigned_team_ids"], []interface{}{"team-1"}) { + t.Fatalf("expected assigned_team_ids [team-1] in payload, got %v", createPayload["assigned_team_ids"]) + } + if d.Id() != "uag-123" { + t.Fatalf("expected ID 'uag-123', got %q", d.Id()) + } + if d.Get("access_group_id").(string) != "uag-123" { + t.Fatalf("expected access_group_id 'uag-123', got %v", d.Get("access_group_id")) + } + if d.Get("created_by").(string) != "admin" { + t.Fatalf("expected created_by 'admin', got %v", d.Get("created_by")) + } +} + +func TestUnifiedAccessGroupRead(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != "GET" || r.URL.Path != "/v1/unified_access_group/uag-123" { + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + return + } + w.Write(unifiedAccessGroupJSON("uag-123")) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := unifiedAccessGroupTestData(t, map[string]interface{}{}) + d.SetId("uag-123") + + if err := resourceLiteLLMUnifiedAccessGroupRead(d, client); err != nil { + t.Fatalf("read failed: %v", err) + } + + if d.Get("access_group_name").(string) != "prod-group" { + t.Fatalf("expected access_group_name 'prod-group', got %v", d.Get("access_group_name")) + } + if d.Get("description").(string) != "prod access" { + t.Fatalf("expected description 'prod access', got %v", d.Get("description")) + } + if !reflect.DeepEqual(d.Get("access_mcp_server_ids"), []interface{}{"mcp-1"}) { + t.Fatalf("expected access_mcp_server_ids [mcp-1], got %v", d.Get("access_mcp_server_ids")) + } + if !reflect.DeepEqual(d.Get("assigned_key_ids"), []interface{}{"key-1"}) { + t.Fatalf("expected assigned_key_ids [key-1], got %v", d.Get("assigned_key_ids")) + } + if d.Get("created_at").(string) != "2026-01-01T00:00:00Z" { + t.Fatalf("expected created_at '2026-01-01T00:00:00Z', got %v", d.Get("created_at")) + } +} + +func TestUnifiedAccessGroupReadNotFound(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := unifiedAccessGroupTestData(t, map[string]interface{}{}) + d.SetId("uag-gone") + + if err := resourceLiteLLMUnifiedAccessGroupRead(d, client); err != nil { + t.Fatalf("expected nil error on 404, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID to be cleared on 404, got %q", d.Id()) + } +} + +func TestUnifiedAccessGroupUpdate(t *testing.T) { + var updatePayload map[string]interface{} + var updatePath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.Method { + case "PUT": + updatePath = r.URL.Path + if err := json.NewDecoder(r.Body).Decode(&updatePayload); err != nil { + t.Errorf("failed to decode update payload: %v", err) + } + w.Write(unifiedAccessGroupJSON("uag-123")) + case "GET": + w.Write(unifiedAccessGroupJSON("uag-123")) + default: + t.Errorf("unexpected request: %s %s", r.Method, r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := unifiedAccessGroupTestData(t, map[string]interface{}{ + "access_group_name": "renamed-group", + "access_model_names": []interface{}{"gpt-4", "claude-3"}, + }) + d.SetId("uag-123") + + if err := resourceLiteLLMUnifiedAccessGroupUpdate(d, client); err != nil { + t.Fatalf("update failed: %v", err) + } + + if updatePath != "/v1/unified_access_group/uag-123" { + t.Fatalf("expected update path '/v1/unified_access_group/uag-123', got %q", updatePath) + } + if updatePayload["access_group_name"] != "renamed-group" { + t.Fatalf("expected access_group_name 'renamed-group' in payload, got %v", updatePayload["access_group_name"]) + } + if !reflect.DeepEqual(updatePayload["access_model_names"], []interface{}{"gpt-4", "claude-3"}) { + t.Fatalf("expected access_model_names [gpt-4 claude-3] in payload, got %v", updatePayload["access_model_names"]) + } +} + +func TestUnifiedAccessGroupDelete(t *testing.T) { + var deleteMethod, deletePath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + deleteMethod = r.Method + deletePath = r.URL.Path + w.WriteHeader(http.StatusNoContent) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := unifiedAccessGroupTestData(t, map[string]interface{}{}) + d.SetId("uag-123") + + if err := resourceLiteLLMUnifiedAccessGroupDelete(d, client); err != nil { + t.Fatalf("delete failed: %v", err) + } + + if deleteMethod != "DELETE" || deletePath != "/v1/unified_access_group/uag-123" { + t.Fatalf("expected DELETE /v1/unified_access_group/uag-123, got %s %s", deleteMethod, deletePath) + } + if d.Id() != "" { + t.Fatalf("expected ID to be cleared after delete, got %q", d.Id()) + } +} diff --git a/terraform/provider/litellm/resource_user.go b/terraform/provider/litellm/resource_user.go new file mode 100644 index 00000000000..c1aa9d7e9fa --- /dev/null +++ b/terraform/provider/litellm/resource_user.go @@ -0,0 +1,362 @@ +package litellm + +import ( + "encoding/json" + "fmt" + "log" + "net/http" + "net/url" + "reflect" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation" +) + +const ( + endpointUserNew = "/user/new" + endpointUserInfo = "/user/info" + endpointUserUpdate = "/user/update" + endpointUserDelete = "/user/delete" +) + +func userSuppressEquivalentJSON(k, oldValue, newValue string, d *schema.ResourceData) bool { + var oldParsed, newParsed interface{} + if err := json.Unmarshal([]byte(oldValue), &oldParsed); err != nil { + return false + } + if err := json.Unmarshal([]byte(newValue), &newParsed); err != nil { + return false + } + return reflect.DeepEqual(oldParsed, newParsed) +} + +func resourceLiteLLMUser() *schema.Resource { + return &schema.Resource{ + Create: resourceLiteLLMUserCreate, + Read: resourceLiteLLMUserRead, + Update: resourceLiteLLMUserUpdate, + Delete: resourceLiteLLMUserDelete, + + Importer: &schema.ResourceImporter{ + StateContext: schema.ImportStatePassthroughContext, + }, + + Schema: map[string]*schema.Schema{ + "user_id": { + Type: schema.TypeString, + Optional: true, + Computed: true, + ForceNew: true, + Description: "Unique ID for the user. Generated by the server if not provided", + }, + "user_email": { + Type: schema.TypeString, + Optional: true, + Description: "Email address of the user", + }, + "user_alias": { + Type: schema.TypeString, + Optional: true, + Description: "Descriptive name for the user", + }, + "user_role": { + Type: schema.TypeString, + Optional: true, + ValidateFunc: validation.StringInSlice([]string{ + "proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer", + }, false), + Description: "Role of the user on the proxy", + }, + "teams": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "List of team IDs the user belongs to", + }, + "models": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Models the user is allowed to call", + }, + "max_budget": { + Type: schema.TypeFloat, + Optional: true, + Description: "Maximum budget in USD for the user", + }, + "budget_duration": { + Type: schema.TypeString, + Optional: true, + Description: "Budget reset period (e.g. '30s', '30m', '30d')", + }, + "tpm_limit": { + Type: schema.TypeInt, + Optional: true, + Description: "Tokens per minute limit for the user", + }, + "rpm_limit": { + Type: schema.TypeInt, + Optional: true, + Description: "Requests per minute limit for the user", + }, + "max_parallel_requests": { + Type: schema.TypeInt, + Optional: true, + Description: "Maximum number of parallel requests for the user", + }, + "metadata": { + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Metadata for the user", + }, + "auto_create_key": { + Type: schema.TypeBool, + Optional: true, + Default: true, + ForceNew: true, + Description: "Whether to auto-create an API key for the user on creation", + }, + "send_invite_email": { + Type: schema.TypeBool, + Optional: true, + Default: false, + ForceNew: true, + Description: "Whether to send an invite email to the user on creation", + }, + "key_alias": { + Type: schema.TypeString, + Optional: true, + Description: "Alias for the auto-created API key", + }, + "aliases": { + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Model aliases for the user", + }, + "config": { + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Config values for the user", + }, + "permissions": { + Type: schema.TypeMap, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Permission values for the user", + }, + "model_max_budget": { + Type: schema.TypeString, + Optional: true, + ValidateFunc: validation.StringIsJSON, + DiffSuppressFunc: userSuppressEquivalentJSON, + Description: "JSON string of per-model budget config (e.g. '{\"gpt-4o\": {\"max_budget\": 10.0}}')", + }, + "guardrails": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + Description: "Guardrails applied to the user's requests", + }, + "blocked": { + Type: schema.TypeBool, + Optional: true, + Default: false, + Description: "Whether the user is blocked from making requests", + }, + "key": { + Type: schema.TypeString, + Computed: true, + Sensitive: true, + Description: "Auto-created API key for the user (when auto_create_key is true)", + }, + }, + } +} + +type userNewResponse struct { + UserID string `json:"user_id"` + Key string `json:"key"` +} + +type userInfoResponse struct { + UserID string `json:"user_id"` + UserInfo map[string]interface{} `json:"user_info"` +} + +func resourceLiteLLMUserCreate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + userData := buildUserData(d) + if v, ok := d.GetOk("user_id"); ok { + userData["user_id"] = v.(string) + } + userData["auto_create_key"] = d.Get("auto_create_key").(bool) + userData["send_invite_email"] = d.Get("send_invite_email").(bool) + + log.Printf("[DEBUG] Create user request payload: %+v", userData) + + resp, err := MakeRequest(client, "POST", endpointUserNew, userData) + if err != nil { + return fmt.Errorf("error creating user: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "creating user"); err != nil { + return err + } + + var userResp userNewResponse + if err := json.NewDecoder(resp.Body).Decode(&userResp); err != nil { + return fmt.Errorf("error decoding create user response: %w", err) + } + if userResp.UserID == "" { + return fmt.Errorf("create user response did not contain a user_id") + } + + d.SetId(userResp.UserID) + if userResp.Key != "" { + d.Set("key", userResp.Key) + } + log.Printf("[INFO] User created with ID: %s", userResp.UserID) + + return resourceLiteLLMUserRead(d, m) +} + +func resourceLiteLLMUserRead(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Reading user with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "GET", fmt.Sprintf("%s?user_id=%s", endpointUserInfo, url.QueryEscape(d.Id())), nil) + if err != nil { + return fmt.Errorf("error reading user: %w", err) + } + defer resp.Body.Close() + + if resp.StatusCode == http.StatusNotFound { + log.Printf("[WARN] User with ID %s not found, removing from state", d.Id()) + d.SetId("") + return nil + } + + if err := handleResponse(resp, "reading user"); err != nil { + return err + } + + var infoResp userInfoResponse + if err := json.NewDecoder(resp.Body).Decode(&infoResp); err != nil { + return fmt.Errorf("error decoding user info response: %w", err) + } + if infoResp.UserInfo == nil { + log.Printf("[WARN] User with ID %s has no user_info, removing from state", d.Id()) + d.SetId("") + return nil + } + + d.Set("user_id", d.Id()) + setUserStateFromInfo(d, infoResp.UserInfo) + + log.Printf("[INFO] Successfully read user with ID: %s", d.Id()) + return nil +} + +func setUserStateFromInfo(d *schema.ResourceData, info map[string]interface{}) { + for _, key := range []string{"user_email", "user_alias", "user_role", "budget_duration"} { + if v, ok := info[key].(string); ok && v != "" { + d.Set(key, v) + } + } + if v, ok := info["max_budget"].(float64); ok { + d.Set("max_budget", v) + } + for _, key := range []string{"tpm_limit", "rpm_limit", "max_parallel_requests"} { + if v, ok := info[key].(float64); ok { + d.Set(key, int(v)) + } + } + for _, key := range []string{"teams", "models"} { + if v, ok := info[key].([]interface{}); ok && len(v) > 0 { + d.Set(key, v) + } + } + if v, ok := info["metadata"].(map[string]interface{}); ok && len(v) > 0 { + d.Set("metadata", v) + } + if v, ok := info["model_max_budget"].(map[string]interface{}); ok && len(v) > 0 { + if encoded, err := json.Marshal(v); err == nil { + d.Set("model_max_budget", string(encoded)) + } + } +} + +func resourceLiteLLMUserUpdate(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + userData := buildUserData(d) + userData["user_id"] = d.Id() + + log.Printf("[DEBUG] Update user request payload: %+v", userData) + + resp, err := MakeRequest(client, "POST", endpointUserUpdate, userData) + if err != nil { + return fmt.Errorf("error updating user: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "updating user"); err != nil { + return err + } + + log.Printf("[INFO] Successfully updated user with ID: %s", d.Id()) + return resourceLiteLLMUserRead(d, m) +} + +func resourceLiteLLMUserDelete(d *schema.ResourceData, m interface{}) error { + client := m.(*Client) + + log.Printf("[INFO] Deleting user with ID: %s", d.Id()) + + resp, err := MakeRequest(client, "POST", endpointUserDelete, map[string]interface{}{ + "user_ids": []string{d.Id()}, + }) + if err != nil { + return fmt.Errorf("error deleting user: %w", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "deleting user"); err != nil { + return err + } + + log.Printf("[INFO] Successfully deleted user with ID: %s", d.Id()) + d.SetId("") + return nil +} + +func buildUserData(d *schema.ResourceData) map[string]interface{} { + userData := map[string]interface{}{ + "blocked": d.Get("blocked").(bool), + } + + for _, key := range []string{ + "user_email", "user_alias", "user_role", "teams", "models", "max_budget", + "budget_duration", "tpm_limit", "rpm_limit", "max_parallel_requests", + "metadata", "key_alias", "aliases", "config", "permissions", "guardrails", + } { + if v, ok := d.GetOk(key); ok { + userData[key] = v + } + } + + if v, ok := d.GetOk("model_max_budget"); ok { + var parsed map[string]interface{} + if err := json.Unmarshal([]byte(v.(string)), &parsed); err == nil { + userData["model_max_budget"] = parsed + } + } + + return userData +} diff --git a/terraform/provider/litellm/resource_user_test.go b/terraform/provider/litellm/resource_user_test.go new file mode 100644 index 00000000000..c254c0c5929 --- /dev/null +++ b/terraform/provider/litellm/resource_user_test.go @@ -0,0 +1,241 @@ +package litellm + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" +) + +func userInfoBody(userID string, info map[string]interface{}) []byte { + body, _ := json.Marshal(map[string]interface{}{ + "user_id": userID, + "user_info": info, + }) + return body +} + +func TestResourceUserCreate_SendsPayloadAndSetsID(t *testing.T) { + var createPayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/user/new": + if r.Method != http.MethodPost { + t.Errorf("expected POST /user/new, got %s", r.Method) + } + if err := json.NewDecoder(r.Body).Decode(&createPayload); err != nil { + t.Fatalf("failed to decode create payload: %v", err) + } + w.Write([]byte(`{"user_id": "u-123", "key": "sk-generated"}`)) + case "/user/info": + if got := r.URL.Query().Get("user_id"); got != "u-123" { + t.Errorf("expected user_id query 'u-123', got %q", got) + } + w.Write(userInfoBody("u-123", map[string]interface{}{ + "user_email": "alice@example.com", + "user_role": "internal_user", + "max_budget": 50.5, + "tpm_limit": float64(1000), + "teams": []interface{}{"team-1"}, + })) + default: + t.Errorf("unexpected request to %s", r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMUser().Schema, map[string]interface{}{ + "user_email": "alice@example.com", + "user_role": "internal_user", + "max_budget": 50.5, + "tpm_limit": 1000, + "auto_create_key": true, + "teams": []interface{}{"team-1"}, + "model_max_budget": `{"gpt-4o": {"max_budget": 10.0}}`, + }) + + if err := resourceLiteLLMUserCreate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("create failed: %v", err) + } + + if d.Id() != "u-123" { + t.Fatalf("expected ID 'u-123', got %q", d.Id()) + } + if got := d.Get("key").(string); got != "sk-generated" { + t.Fatalf("expected key 'sk-generated', got %q", got) + } + if got := createPayload["user_email"]; got != "alice@example.com" { + t.Errorf("expected user_email in payload, got %v", got) + } + if got := createPayload["user_role"]; got != "internal_user" { + t.Errorf("expected user_role in payload, got %v", got) + } + if got := createPayload["max_budget"]; got != 50.5 { + t.Errorf("expected max_budget 50.5 in payload, got %v", got) + } + if got := createPayload["auto_create_key"]; got != true { + t.Errorf("expected auto_create_key true in payload, got %v", got) + } + mmb, ok := createPayload["model_max_budget"].(map[string]interface{}) + if !ok { + t.Fatalf("expected model_max_budget object in payload, got %v", createPayload["model_max_budget"]) + } + if _, ok := mmb["gpt-4o"]; !ok { + t.Errorf("expected gpt-4o key in model_max_budget, got %v", mmb) + } + if got := d.Get("user_email").(string); got != "alice@example.com" { + t.Errorf("expected user_email in state, got %q", got) + } +} + +func TestResourceUserRead_MapsFields(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Write(userInfoBody("u-42", map[string]interface{}{ + "user_email": "bob@example.com", + "user_alias": "bob", + "user_role": "proxy_admin", + "max_budget": 100.0, + "budget_duration": "30d", + "tpm_limit": float64(5000), + "rpm_limit": float64(60), + "teams": []interface{}{"team-a", "team-b"}, + "models": []interface{}{"gpt-4o"}, + "model_max_budget": map[string]interface{}{"gpt-4o": map[string]interface{}{"max_budget": 5.0}}, + })) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMUser().Schema, map[string]interface{}{}) + d.SetId("u-42") + + if err := resourceLiteLLMUserRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("read failed: %v", err) + } + + if got := d.Get("user_email").(string); got != "bob@example.com" { + t.Errorf("expected user_email 'bob@example.com', got %q", got) + } + if got := d.Get("user_alias").(string); got != "bob" { + t.Errorf("expected user_alias 'bob', got %q", got) + } + if got := d.Get("user_role").(string); got != "proxy_admin" { + t.Errorf("expected user_role 'proxy_admin', got %q", got) + } + if got := d.Get("max_budget").(float64); got != 100.0 { + t.Errorf("expected max_budget 100.0, got %v", got) + } + if got := d.Get("budget_duration").(string); got != "30d" { + t.Errorf("expected budget_duration '30d', got %q", got) + } + if got := d.Get("tpm_limit").(int); got != 5000 { + t.Errorf("expected tpm_limit 5000, got %d", got) + } + if got := d.Get("rpm_limit").(int); got != 60 { + t.Errorf("expected rpm_limit 60, got %d", got) + } + teams := d.Get("teams").([]interface{}) + if len(teams) != 2 || teams[0] != "team-a" { + t.Errorf("expected teams [team-a team-b], got %v", teams) + } + var mmb map[string]interface{} + if err := json.Unmarshal([]byte(d.Get("model_max_budget").(string)), &mmb); err != nil { + t.Fatalf("model_max_budget in state is not valid JSON: %v", err) + } + if _, ok := mmb["gpt-4o"]; !ok { + t.Errorf("expected gpt-4o key in model_max_budget state, got %v", mmb) + } +} + +func TestResourceUserRead_404ClearsID(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNotFound) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMUser().Schema, map[string]interface{}{}) + d.SetId("gone-user") + + if err := resourceLiteLLMUserRead(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("expected nil error on 404, got: %v", err) + } + if d.Id() != "" { + t.Fatalf("expected ID to be cleared on 404, got %q", d.Id()) + } +} + +func TestResourceUserUpdate_SendsPayload(t *testing.T) { + var updatePayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/user/update": + if r.Method != http.MethodPost { + t.Errorf("expected POST /user/update, got %s", r.Method) + } + if err := json.NewDecoder(r.Body).Decode(&updatePayload); err != nil { + t.Fatalf("failed to decode update payload: %v", err) + } + w.Write([]byte(`{"user_id": "u-7"}`)) + case "/user/info": + w.Write(userInfoBody("u-7", map[string]interface{}{"user_role": "internal_user_viewer"})) + default: + t.Errorf("unexpected request to %s", r.URL.Path) + w.WriteHeader(http.StatusNotFound) + } + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMUser().Schema, map[string]interface{}{ + "user_role": "internal_user_viewer", + "max_budget": 25.0, + }) + d.SetId("u-7") + + if err := resourceLiteLLMUserUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + + if got := updatePayload["user_id"]; got != "u-7" { + t.Errorf("expected user_id 'u-7' in payload, got %v", got) + } + if got := updatePayload["user_role"]; got != "internal_user_viewer" { + t.Errorf("expected user_role in payload, got %v", got) + } + if got := updatePayload["max_budget"]; got != 25.0 { + t.Errorf("expected max_budget 25.0 in payload, got %v", got) + } + if _, ok := updatePayload["auto_create_key"]; ok { + t.Errorf("auto_create_key must not be sent on update, got %v", updatePayload["auto_create_key"]) + } +} + +func TestResourceUserDelete_SendsUserIDs(t *testing.T) { + var deletePayload map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/user/delete" || r.Method != http.MethodPost { + t.Errorf("expected POST /user/delete, got %s %s", r.Method, r.URL.Path) + } + if err := json.NewDecoder(r.Body).Decode(&deletePayload); err != nil { + t.Fatalf("failed to decode delete payload: %v", err) + } + w.Write([]byte(`{}`)) + })) + defer srv.Close() + + d := schema.TestResourceDataRaw(t, resourceLiteLLMUser().Schema, map[string]interface{}{}) + d.SetId("u-del") + + if err := resourceLiteLLMUserDelete(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("delete failed: %v", err) + } + + ids, ok := deletePayload["user_ids"].([]interface{}) + if !ok || len(ids) != 1 || ids[0] != "u-del" { + t.Fatalf("expected user_ids ['u-del'], got %v", deletePayload["user_ids"]) + } + if d.Id() != "" { + t.Fatalf("expected ID to be cleared after delete, got %q", d.Id()) + } +} diff --git a/terraform/provider/litellm/resource_vector_store.go b/terraform/provider/litellm/resource_vector_store.go index f77ba18c6d4..a3faf9673c3 100644 --- a/terraform/provider/litellm/resource_vector_store.go +++ b/terraform/provider/litellm/resource_vector_store.go @@ -10,6 +10,9 @@ func resourceLiteLLMVectorStore() *schema.Resource { Read: resourceLiteLLMVectorStoreRead, Update: resourceLiteLLMVectorStoreUpdate, Delete: resourceLiteLLMVectorStoreDelete, + Importer: &schema.ResourceImporter{ + StateContext: schema.ImportStatePassthroughContext, + }, Schema: map[string]*schema.Schema{ "vector_store_id": { diff --git a/terraform/provider/litellm/types.go b/terraform/provider/litellm/types.go index 66d1f6a8ba9..7bef44409fd 100644 --- a/terraform/provider/litellm/types.go +++ b/terraform/provider/litellm/types.go @@ -40,18 +40,29 @@ type TeamInfoResponse struct { // TeamResponse represents a response from the API containing team information. type TeamResponse struct { - TeamID string `json:"team_id,omitempty"` - TeamAlias string `json:"team_alias,omitempty"` - OrganizationID string `json:"organization_id,omitempty"` - Metadata map[string]interface{} `json:"metadata,omitempty"` - TPMLimit *int `json:"tpm_limit,omitempty"` - RPMLimit *int `json:"rpm_limit,omitempty"` - MaxBudget *float64 `json:"max_budget,omitempty"` - SoftBudget *float64 `json:"soft_budget,omitempty"` - BudgetDuration string `json:"budget_duration,omitempty"` - Models []string `json:"models"` - Blocked bool `json:"blocked,omitempty"` - TeamMemberPermissions []string `json:"team_member_permissions,omitempty"` + TeamID string `json:"team_id,omitempty"` + TeamAlias string `json:"team_alias,omitempty"` + OrganizationID string `json:"organization_id,omitempty"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + TPMLimit *int `json:"tpm_limit,omitempty"` + RPMLimit *int `json:"rpm_limit,omitempty"` + MaxBudget *float64 `json:"max_budget,omitempty"` + SoftBudget *float64 `json:"soft_budget,omitempty"` + BudgetDuration string `json:"budget_duration,omitempty"` + Models []string `json:"models"` + Blocked bool `json:"blocked,omitempty"` + TeamMemberPermissions []string `json:"team_member_permissions,omitempty"` + ModelAliases map[string]interface{} `json:"model_aliases,omitempty"` + Guardrails []string `json:"guardrails,omitempty"` + Prompts []string `json:"prompts,omitempty"` + TeamMemberBudget *float64 `json:"team_member_budget,omitempty"` + TeamMemberBudgetDuration string `json:"team_member_budget_duration,omitempty"` + TeamMemberRPMLimit *int `json:"team_member_rpm_limit,omitempty"` + TeamMemberTPMLimit *int `json:"team_member_tpm_limit,omitempty"` + TeamMemberKeyDuration string `json:"team_member_key_duration,omitempty"` + ModelRPMLimit map[string]interface{} `json:"model_rpm_limit,omitempty"` + ModelTPMLimit map[string]interface{} `json:"model_tpm_limit,omitempty"` + AllowedPassthroughRoutes []string `json:"allowed_passthrough_routes,omitempty"` } // OrganizationResponse represents a response from the API containing organization information. @@ -107,31 +118,40 @@ type ModelInfo struct { // Key represents a LiteLLM API key. type Key struct { - Key string `json:"key,omitempty"` - TokenID string `json:"token_id,omitempty"` - Models []string `json:"models"` - Spend float64 `json:"spend,omitempty"` - MaxBudget *float64 `json:"max_budget,omitempty"` - UserID string `json:"user_id,omitempty"` - TeamID string `json:"team_id,omitempty"` - MaxParallelRequests *int `json:"max_parallel_requests,omitempty"` - Metadata map[string]interface{} `json:"metadata,omitempty"` - TPMLimit *int `json:"tpm_limit,omitempty"` - RPMLimit *int `json:"rpm_limit,omitempty"` - BudgetDuration string `json:"budget_duration,omitempty"` - AllowedCacheControls []string `json:"allowed_cache_controls,omitempty"` - SoftBudget *float64 `json:"soft_budget,omitempty"` - KeyAlias string `json:"key_alias,omitempty"` - Duration string `json:"duration,omitempty"` - Aliases map[string]interface{} `json:"aliases,omitempty"` - Config map[string]interface{} `json:"config,omitempty"` - Permissions map[string]interface{} `json:"permissions,omitempty"` - ModelMaxBudget map[string]interface{} `json:"model_max_budget,omitempty"` - ModelRPMLimit map[string]interface{} `json:"model_rpm_limit,omitempty"` - ModelTPMLimit map[string]interface{} `json:"model_tpm_limit,omitempty"` - Guardrails []string `json:"guardrails,omitempty"` - Blocked bool `json:"blocked"` - Tags []string `json:"tags,omitempty"` + Key string `json:"key,omitempty"` + TokenID string `json:"token_id,omitempty"` + Models []string `json:"models"` + Spend float64 `json:"spend,omitempty"` + MaxBudget *float64 `json:"max_budget,omitempty"` + UserID string `json:"user_id,omitempty"` + TeamID string `json:"team_id,omitempty"` + MaxParallelRequests *int `json:"max_parallel_requests,omitempty"` + Metadata map[string]interface{} `json:"metadata,omitempty"` + TPMLimit *int `json:"tpm_limit,omitempty"` + RPMLimit *int `json:"rpm_limit,omitempty"` + BudgetDuration string `json:"budget_duration,omitempty"` + AllowedCacheControls []string `json:"allowed_cache_controls,omitempty"` + SoftBudget *float64 `json:"soft_budget,omitempty"` + KeyAlias string `json:"key_alias,omitempty"` + Duration string `json:"duration,omitempty"` + Aliases map[string]interface{} `json:"aliases,omitempty"` + Config map[string]interface{} `json:"config,omitempty"` + Permissions map[string]interface{} `json:"permissions,omitempty"` + ModelMaxBudget map[string]interface{} `json:"model_max_budget,omitempty"` + ModelRPMLimit map[string]interface{} `json:"model_rpm_limit,omitempty"` + ModelTPMLimit map[string]interface{} `json:"model_tpm_limit,omitempty"` + Guardrails []string `json:"guardrails,omitempty"` + Blocked bool `json:"blocked"` + Tags []string `json:"tags,omitempty"` + BudgetID string `json:"budget_id,omitempty"` + EnforcedParams []string `json:"enforced_params,omitempty"` + AllowedRoutes []string `json:"allowed_routes,omitempty"` + AllowedPassthroughRoutes []string `json:"allowed_passthrough_routes,omitempty"` + RPMLimitType string `json:"rpm_limit_type,omitempty"` + TPMLimitType string `json:"tpm_limit_type,omitempty"` + Prompts []string `json:"prompts,omitempty"` + OrganizationID string `json:"organization_id,omitempty"` + ProjectID string `json:"project_id,omitempty"` } // KeyResponse represents a response from the API containing key information. @@ -252,3 +272,33 @@ type VectorStoreDeleteRequest struct { type VectorStoreInfoRequest struct { VectorStoreID string `json:"vector_store_id"` } + +type JWTKeyMappingRequest struct { + JWTClaimName string `json:"jwt_claim_name"` + JWTClaimValue string `json:"jwt_claim_value"` + Key string `json:"key"` + Description string `json:"description,omitempty"` +} + +type JWTKeyMappingUpdateRequest struct { + ID string `json:"id"` + Key string `json:"key,omitempty"` + Description string `json:"description"` + IsActive bool `json:"is_active"` +} + +type JWTKeyMappingDeleteRequest struct { + ID string `json:"id"` +} + +type JWTKeyMappingResponse struct { + ID string `json:"id"` + JWTClaimName string `json:"jwt_claim_name"` + JWTClaimValue string `json:"jwt_claim_value"` + Description string `json:"description,omitempty"` + IsActive bool `json:"is_active"` + CreatedAt string `json:"created_at,omitempty"` + UpdatedAt string `json:"updated_at,omitempty"` + CreatedBy string `json:"created_by,omitempty"` + UpdatedBy string `json:"updated_by,omitempty"` +} diff --git a/terraform/provider/litellm/utils.go b/terraform/provider/litellm/utils.go index 01d8045300c..5e81766d3f3 100644 --- a/terraform/provider/litellm/utils.go +++ b/terraform/provider/litellm/utils.go @@ -2,6 +2,8 @@ package litellm import ( "bytes" + "crypto/sha256" + "encoding/hex" "encoding/json" "fmt" "io" @@ -60,6 +62,18 @@ func handleAPIResponse(resp *http.Response, reqBody interface{}, client *Client) return &modelResp, nil } +// hashedKeyToken normalizes a raw sk- API key to its SHA-256 token hash, the +// identifier the proxy stores and accepts, so the plaintext key never lands +// in request URLs, resource IDs, or proxy access logs. Values that are +// already hashed pass through unchanged. +func hashedKeyToken(key string) string { + if !strings.HasPrefix(key, "sk-") { + return key + } + sum := sha256.Sum256([]byte(key)) + return hex.EncodeToString(sum[:]) +} + // MakeRequest is a helper function to make HTTP requests func MakeRequest(client *Client, method, endpoint string, body interface{}) (*http.Response, error) { var req *http.Request diff --git a/terraform/provider/tools/endpointaudit/coverage.go b/terraform/provider/tools/endpointaudit/coverage.go new file mode 100644 index 00000000000..671758d3477 --- /dev/null +++ b/terraform/provider/tools/endpointaudit/coverage.go @@ -0,0 +1,106 @@ +package main + +import ( + "bufio" + "encoding/json" + "fmt" + "os" + "sort" + "strings" +) + +var managementPrefixes = map[string]bool{ + "access_group": true, + "agent": true, + "budget": true, + "cache": true, + "config": true, + "coordination_redis": true, + "credentials": true, + "customer": true, + "fallback": true, + "guardrails": true, + "jwt": true, + "key": true, + "model": true, + "organization": true, + "project": true, + "prompts": true, + "router": true, + "search_tools": true, + "tag": true, + "team": true, + "user": true, + "vector_store": true, +} + +func isManagementPath(path string) bool { + segments := strings.SplitN(strings.TrimPrefix(path, "/"), "/", 2) + return len(segments) > 0 && managementPrefixes[segments[0]] +} + +func parseAllowlist(path string) (map[string]bool, error) { + file, err := os.Open(path) + if err != nil { + return nil, err + } + defer file.Close() + entries := make(map[string]bool) + scanner := bufio.NewScanner(file) + line := 0 + for scanner.Scan() { + line++ + text := strings.TrimSpace(scanner.Text()) + if text == "" || strings.HasPrefix(text, "#") { + continue + } + if idx := strings.Index(text, "#"); idx >= 0 { + text = strings.TrimSpace(text[:idx]) + } + fields := strings.Fields(text) + if len(fields) != 2 || !strings.HasPrefix(fields[1], "/") { + return nil, fmt.Errorf("%s:%d: allowlist entries must be \"METHOD /path\", got %q", path, line, text) + } + entries[strings.ToUpper(fields[0])+" "+fields[1]] = true + } + return entries, scanner.Err() +} + +func specCallCovered(calls []endpointCall, specMethod, specPath string) bool { + for _, call := range calls { + if strings.EqualFold(call.Method, specMethod) && pathMatches(call.Path, specPath) { + return true + } + } + return false +} + +func auditCoverage(calls []endpointCall, specPaths map[string]map[string]json.RawMessage, allowlist map[string]bool) []string { + var violations []string + seen := make(map[string]bool) + for specPath, operations := range specPaths { + if !isManagementPath(specPath) { + continue + } + for method := range operations { + entry := strings.ToUpper(method) + " " + specPath + covered := specCallCovered(calls, method, specPath) + switch { + case allowlist[entry]: + seen[entry] = true + if covered { + violations = append(violations, fmt.Sprintf("stale allowlist entry: %s is covered by the provider; remove it from the allowlist", entry)) + } + case !covered: + violations = append(violations, fmt.Sprintf("uncovered management endpoint: %s has no provider resource or data source; add coverage or allowlist it with a reason", entry)) + } + } + } + for entry := range allowlist { + if !seen[entry] { + violations = append(violations, fmt.Sprintf("stale allowlist entry: %s is not a management endpoint in the proxy schema; remove it from the allowlist", entry)) + } + } + sort.Strings(violations) + return violations +} diff --git a/terraform/provider/tools/endpointaudit/coverage_allowlist.txt b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt new file mode 100644 index 00000000000..d10be89b90c --- /dev/null +++ b/terraform/provider/tools/endpointaudit/coverage_allowlist.txt @@ -0,0 +1,128 @@ +# Management endpoints deliberately not covered by a Terraform resource or data source. +# +# Format: one "METHOD /path" per line, matching the proxy OpenAPI schema exactly; +# "#" starts a comment. The coverage gate (endpointaudit -coverage-allowlist) fails +# when a management endpoint is neither covered nor listed here, and also when an +# entry goes stale (the provider now covers it, or the endpoint left the schema), +# so this file can only shrink relative to the schema over time. +# +# Every entry needs a reason. Endpoints that are analytics, UI helpers, or +# imperative one-shot operations never get a resource. Entries marked "known gap" +# are real coverage gaps awaiting a resource; remove them when the resource lands. + +# Read-only analytics and spend reporting; observability, not Terraform-managed state +GET /agent/daily/activity +GET /customer/daily/activity +GET /guardrails/usage/detail/{guardrail_id} +GET /guardrails/usage/logs +GET /guardrails/usage/overview +GET /key/spend/report +GET /organization/daily/activity +GET /organization/spend/report +GET /tag/daily/activity +GET /tag/dau +GET /tag/distinct +GET /tag/mau +GET /tag/summary +GET /tag/user-agent/per-user-analytics +GET /tag/wau +GET /team/daily/activity +GET /team/daily/activity/aggregated +GET /team/spend/report +GET /user/daily/activity +GET /user/daily/activity/aggregated +GET /user/spend/report + +# Admin UI helper endpoints; serve UI forms and caller-scoped views, not desired state +GET /budget/settings +GET /router/fields +GET /guardrails/ui/add_guardrail_settings +GET /guardrails/ui/category_yaml/{category_name} +GET /guardrails/ui/major_airlines +GET /guardrails/ui/provider_specific_params +GET /key/aliases +GET /model/deprecations +GET /search_tools/ui/available_providers +GET /team/available +GET /team/metadata_schema +GET /team/{team_id}/members/me +GET /user/available_users + +# Imperative one-shot operations: bulk edits, rotation, health probes, test hooks, +# migrations, and approval workflows; procedural, not declarative state +GET /cache/ping +GET /cache/redis/info +GET /credentials/migrate-encryption/check +POST /cache/delete +POST /cache/flushall +POST /cache/settings/test +POST /coordination_redis/settings/test +GET /guardrails/submissions +GET /guardrails/submissions/{guardrail_id} +POST /credentials/migrate-encryption +POST /customer/block +POST /customer/unblock +POST /guardrails/apply_guardrail +POST /guardrails/register +POST /guardrails/submissions/{guardrail_id}/approve +POST /guardrails/submissions/{guardrail_id}/reject +POST /guardrails/test_custom_code +POST /guardrails/validate_blocked_words_file +POST /key/bulk_update +POST /key/health +POST /key/regenerate +POST /key/service-account/generate +POST /key/{key}/regenerate +POST /key/{key}/reset_spend +POST /model/block +POST /model/unblock +POST /prompts/test +POST /search_tools/test_connection +POST /team/bulk_member_add +POST /team/{team_id}/member/{user_id}/reset_spend +POST /team/key/bulk_update +POST /team/permissions_bulk_update +POST /team/{team_id}/disable_logging +POST /user/bulk_update + +# Alternate method or path for functionality the provider already manages elsewhere +GET /credentials/by_model/{model_id} +GET /guardrails/{guardrail_id} +GET /prompts/{prompt_id} +GET /prompts/{prompt_id}/versions +PATCH /guardrails/{guardrail_id} +PATCH /model/{model_id}/update +PATCH /prompts/{prompt_id} +PATCH /team/{team_id} +POST /team/model/add +POST /team/model/delete + +# Known gaps awaiting a resource or data source; remove the entry when it lands +GET /credentials # known gap: plural credentials data source +GET /cache/settings # known gap: cache settings resource +POST /cache/settings # known gap: cache settings resource +GET /coordination_redis/settings # known gap: coordination redis settings resource +POST /coordination_redis/settings # known gap: coordination redis settings resource +GET /router/settings # known gap: router settings data source +GET /router/fields # known gap: router settings data source +GET /config/block_requests_for_models_without_pricing # known gap: proxy config resource +PATCH /config/block_requests_for_models_without_pricing # known gap: proxy config resource +GET /config/cost_discount_config # known gap: proxy config resource +PATCH /config/cost_discount_config # known gap: proxy config resource +GET /config/cost_margin_config # known gap: proxy config resource +PATCH /config/cost_margin_config # known gap: proxy config resource +GET /config/pass_through_endpoint # known gap: pass-through endpoint resource +POST /config/pass_through_endpoint # known gap: pass-through endpoint resource +DELETE /config/pass_through_endpoint # known gap: pass-through endpoint resource +POST /config/pass_through_endpoint/{endpoint_id} # known gap: pass-through endpoint resource +GET /config/pass_through_endpoint/team/{team_id} # known gap: pass-through endpoint resource +GET /vector_store/list # known gap: plural vector stores data source +GET /jwt/key/mapping/list # known gap: plural jwt key mappings data source +GET /customer/info # known gap: litellm_customer resource +GET /customer/list # known gap: litellm_customer resource +POST /customer/new # known gap: litellm_customer resource +POST /customer/update # known gap: litellm_customer resource +POST /customer/delete # known gap: litellm_customer resource +GET /team/{team_id}/callback # known gap: team callback resource +POST /team/{team_id}/callback # known gap: team callback resource +DELETE /team/{team_id}/callback/{callback_name} # known gap: team callback resource diff --git a/terraform/provider/tools/endpointaudit/coverage_test.go b/terraform/provider/tools/endpointaudit/coverage_test.go new file mode 100644 index 00000000000..30fa31a480f --- /dev/null +++ b/terraform/provider/tools/endpointaudit/coverage_test.go @@ -0,0 +1,136 @@ +package main + +import ( + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" +) + +func coverageSpecFixture(paths map[string][]string) map[string]map[string]json.RawMessage { + spec := make(map[string]map[string]json.RawMessage) + for path, methods := range paths { + operations := make(map[string]json.RawMessage) + for _, method := range methods { + operations[method] = json.RawMessage(`{}`) + } + spec[path] = operations + } + return spec +} + +func writeAllowlist(t *testing.T, body string) string { + t.Helper() + path := filepath.Join(t.TempDir(), "allowlist.txt") + if err := os.WriteFile(path, []byte(body), 0o644); err != nil { + t.Fatal(err) + } + return path +} + +func TestParseAllowlist(t *testing.T) { + path := writeAllowlist(t, `# comment +GET /team/spend/report + +post /key/regenerate # inline reason +`) + entries, err := parseAllowlist(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 || !entries["GET /team/spend/report"] || !entries["POST /key/regenerate"] { + t.Fatalf("unexpected entries: %v", entries) + } +} + +func TestParseAllowlistRejectsMalformedLines(t *testing.T) { + path := writeAllowlist(t, "GET\n") + if _, err := parseAllowlist(path); err == nil { + t.Fatal("expected error for malformed line") + } +} + +func TestAuditCoverageFailsOnUncoveredManagementEndpoint(t *testing.T) { + spec := coverageSpecFixture(map[string][]string{ + "/team/new": {"post"}, + "/team/spend/report": {"get"}, + "/chat/completions": {"post"}, + "/health/liveliness": {"get"}, + "/v1/chat/completions": {"post"}, + }) + calls := []endpointCall{{Method: "POST", Path: "/team/new"}} + violations := auditCoverage(calls, spec, nil) + if len(violations) != 1 || !strings.Contains(violations[0], "GET /team/spend/report") { + t.Fatalf("unexpected violations: %v", violations) + } +} + +func TestAuditCoverageAllowlistSuppressesUncovered(t *testing.T) { + spec := coverageSpecFixture(map[string][]string{"/team/spend/report": {"get"}}) + violations := auditCoverage(nil, spec, map[string]bool{"GET /team/spend/report": true}) + if len(violations) != 0 { + t.Fatalf("unexpected violations: %v", violations) + } +} + +func TestAuditCoverageFailsOnStaleCoveredEntry(t *testing.T) { + spec := coverageSpecFixture(map[string][]string{"/team/new": {"post"}}) + calls := []endpointCall{{Method: "POST", Path: "/team/new"}} + violations := auditCoverage(calls, spec, map[string]bool{"POST /team/new": true}) + if len(violations) != 1 || !strings.Contains(violations[0], "stale allowlist entry: POST /team/new is covered") { + t.Fatalf("unexpected violations: %v", violations) + } +} + +func TestAuditCoverageFailsOnEntryMissingFromSchema(t *testing.T) { + spec := coverageSpecFixture(map[string][]string{"/team/new": {"post"}}) + calls := []endpointCall{{Method: "POST", Path: "/team/new"}} + violations := auditCoverage(calls, spec, map[string]bool{"POST /team/removed": true}) + if len(violations) != 1 || !strings.Contains(violations[0], "POST /team/removed is not a management endpoint") { + t.Fatalf("unexpected violations: %v", violations) + } +} + +func TestAuditCoverageMatchesPathParams(t *testing.T) { + spec := coverageSpecFixture(map[string][]string{"/team/{team_id}/callback": {"get"}}) + calls := []endpointCall{{Method: "GET", Path: "/team/{param}/callback"}} + violations := auditCoverage(calls, spec, nil) + if len(violations) != 0 { + t.Fatalf("unexpected violations: %v", violations) + } +} + +func TestMountedDeclarativeAPIsAreManagementPaths(t *testing.T) { + for _, path := range []string{ + "/cache/settings", + "/config/cost_discount_config", + "/coordination_redis/settings", + "/router/settings", + } { + if !isManagementPath(path) { + t.Fatalf("%s should be classified as a management path", path) + } + } + for _, path := range []string{"/chat/completions", "/health/liveliness"} { + if isManagementPath(path) { + t.Fatalf("%s should not be classified as a management path", path) + } + } +} + +func TestBundledAllowlistEntriesAreManagementPaths(t *testing.T) { + entries, err := parseAllowlist("coverage_allowlist.txt") + if err != nil { + t.Fatal(err) + } + if len(entries) == 0 { + t.Fatal("bundled allowlist parsed to zero entries") + } + for entry := range entries { + fields := strings.Fields(entry) + if !isManagementPath(fields[1]) { + t.Fatalf("allowlist entry %q is not under a management prefix", entry) + } + } +} diff --git a/terraform/provider/tools/endpointaudit/main.go b/terraform/provider/tools/endpointaudit/main.go index ebc011ee910..71d452c9816 100644 --- a/terraform/provider/tools/endpointaudit/main.go +++ b/terraform/provider/tools/endpointaudit/main.go @@ -306,7 +306,7 @@ func auditCalls(calls []endpointCall, specPaths map[string]map[string]json.RawMe return violations } -func run(providerDir, specPath string) error { +func run(providerDir, specPath, coverageAllowlistPath string) error { extracted, err := extractProviderCalls(providerDir) if err != nil { return err @@ -326,6 +326,16 @@ func run(providerDir, specPath string) error { sort.Strings(violations) return fmt.Errorf("provider/proxy endpoint drift:\n %s", strings.Join(violations, "\n ")) } + if coverageAllowlistPath != "" { + allowlist, err := parseAllowlist(coverageAllowlistPath) + if err != nil { + return err + } + coverageViolations := auditCoverage(extracted.Calls, specPaths, allowlist) + if len(coverageViolations) > 0 { + return fmt.Errorf("provider coverage gaps:\n %s", strings.Join(coverageViolations, "\n ")) + } + } fmt.Printf("OK: %d request call sites verified against %d proxy OpenAPI paths\n", len(extracted.Calls), len(specPaths)) return nil } @@ -333,12 +343,13 @@ func run(providerDir, specPath string) error { func main() { providerDir := flag.String("provider-dir", "./litellm", "directory containing the provider Go source") specPath := flag.String("spec", "", "path to the proxy OpenAPI schema JSON") + coverageAllowlist := flag.String("coverage-allowlist", "", "path to the coverage allowlist; when set, also fail on management endpoints with no provider coverage") flag.Parse() if *specPath == "" { fmt.Fprintln(os.Stderr, "error: -spec is required") os.Exit(2) } - if err := run(*providerDir, *specPath); err != nil { + if err := run(*providerDir, *specPath, *coverageAllowlist); err != nil { fmt.Fprintf(os.Stderr, "error: %v\n", err) os.Exit(1) } diff --git a/tests/e2e/coverage_registry/guardrail.yaml b/tests/e2e/coverage_registry/guardrail.yaml index f66a73e7daf..c64fd6150af 100644 --- a/tests/e2e/coverage_registry/guardrail.yaml +++ b/tests/e2e/coverage_registry/guardrail.yaml @@ -32,3 +32,4 @@ - {id: guardrail.mcp_security.pre_call.blocks, module: guardrail, tier: P2, hook_point: pre_call, assertions: [blocks], exercised_on: [mcp_operations], source: "guardrail_hooks/mcp_security", rationale: "MCP protocol security"} - {id: guardrail.llm_as_a_judge.pre_call.blocks, module: guardrail, tier: P2, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/llm_as_a_judge", rationale: "LLM-based judgment guardrail"} - {id: guardrail.litellm_content_filter.pre_mcp_call.blocks, module: guardrail, tier: P1, hook_point: pre_mcp_call, assertions: [blocks], exercised_on: [mcp_operations], source: "guardrail_hooks/litellm_content_filter/content_filter.py:_scan_mcp_tool_call_arguments", rationale: "A general content-filter guardrail configured mode=pre_mcp_call blocks a banned keyword in an MCP tool call's arguments before it reaches the upstream MCP server; a clean argument passes"} +- {id: guardrail.dispatch.pre_call.rejects_unknown_name, module: guardrail, tier: P1, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "proxy guardrail dispatch (per-request `guardrails` selector)", rationale: "A request naming a guardrail this proxy does not serve must fail closed with a 4xx; today it is silently served unguarded, so a typo'd name drops the protection the caller asked for"} diff --git a/tests/e2e/coverage_registry/logging.yaml b/tests/e2e/coverage_registry/logging.yaml index 856636c3dbc..1f2f1d64711 100644 --- a/tests/e2e/coverage_registry/logging.yaml +++ b/tests/e2e/coverage_registry/logging.yaml @@ -24,3 +24,4 @@ - {id: logging.focus.success.writes_object, module: logging, tier: P1, event: success, assertions: [writes_object], exercised_on: [chat_completions, messages], source: "integrations/focus/focus_logger.py", rationale: "Cost mgmt multi-destination export"} - {id: logging.niche_integrations.success.logs_spend, module: logging, tier: P2, event: success, assertions: [logs_spend], exercised_on: [chat_completions], source: grammar, rationale: "SMOKE cohort: athina/galileo/deepeval/langtrace/weave/lunary/humanloop/traceloop/helicone/argilla/newrelic/sqs/supabase/dynamodb/agentops/lago/etc"} - {id: logging.niche_integrations.failure.logs_spend, module: logging, tier: P2, event: failure, assertions: [logs_spend], exercised_on: [chat_completions], source: grammar, rationale: "SMOKE niche failure path"} +- {id: logging.langfuse.success.logs_spend, module: logging, tier: P1, event: success, assertions: [logs_spend], exercised_on: [chat_completions], source: "integrations/langfuse/langfuse_otel.py", rationale: "Team-scoped Langfuse delivery via /team/callback; LangChain-ecosystem evals spend"} diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index c158fc89c81..f03e70df84a 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -68,11 +68,27 @@ class BlockCodeExecutionParamsBody(GuardrailParamsBase): guardrail: Literal["block_code_execution"] = "block_code_execution" +class PresidioParamsBody(GuardrailParamsBase): + """Presidio PII guardrail params. `presidio_filter_scope="input"` keeps the + registration to a single callback on the configured mode; the default + ("both") also registers a second post_call output-masking callback, which a + pre_call- or logging_only-scoped test must not drag in. `output_parse_pii` + stays unset/False: True would unmask the response back to the caller.""" + + guardrail: Literal["presidio"] = "presidio" + presidio_analyzer_api_base: str + presidio_anonymizer_api_base: str + presidio_filter_scope: Literal["input", "output", "both"] | None = None + presidio_language: str | None = None + output_parse_pii: bool | None = None + + GuardrailParamsBody = ( ContentFilterParamsBody | BedrockGuardrailParamsBody | OpenAIModerationParamsBody | BlockCodeExecutionParamsBody + | PresidioParamsBody ) @@ -253,6 +269,30 @@ class GuardrailsClient: ), ) + def chat_stream_raw( + self, + key: str, + model: str, + text: str, + *, + guardrails: list[str] | None = None, + max_tokens: int = 64, + ) -> StreamingResponse: + """Drive /chat/completions with stream=true, returning the raw HTTP + outcome (status, headers, SSE events) via the shared ProxyClient stream + sender - a streamed guardrail block is judged on status and stream + shape, not a typed body.""" + return self.proxy.chat_stream( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content=text)], + max_tokens=max_tokens, + stream=True, + guardrails=guardrails, + ), + ) + def messages( self, key: str, @@ -318,7 +358,7 @@ def build_client(proxy: ProxyClient) -> GuardrailsClient: return GuardrailsClient(proxy=proxy) -def poll_until_blocked(call: Callable[[], Result[ChatResponse]]) -> Result[ChatResponse]: +def poll_until_blocked[R: BaseModel](call: Callable[[], Result[R]]) -> Result[R]: """Retry a call that a guardrail should reject until it is, returning the last result. Registering a guardrail is a control-plane write; the data-plane worker that @@ -337,3 +377,25 @@ def poll_until_blocked(call: Callable[[], Result[ChatResponse]]) -> Result[ChatR time.sleep(POLL_INTERVAL) last = call() return last + + +#: Statuses a stream poll keeps retrying through instead of returning as "the +#: block": network failures (-1), key propagation (401), rate limits (429) - +#: transient rig noise, not a guardrail verdict. +_TRANSIENT_STREAM_STATUSES = frozenset({-1, 401, 429}) + + +def poll_until_blocked_stream(call: Callable[[], StreamingResponse]) -> StreamingResponse: + """poll_until_blocked for raw/streamed sends, which return a StreamingResponse + instead of a Result: retry while the call still succeeds (the data-plane worker + has not picked the new guardrail up yet) or fails with a transient status, + returning the first guardrail-shaped non-2xx outcome or the last result at + the deadline.""" + deadline = time.monotonic() + POLL_TIMEOUT + last = call() + while time.monotonic() < deadline: + if not last.ok and last.status_code not in _TRANSIENT_STREAM_STATUSES: + return last + time.sleep(POLL_INTERVAL) + last = call() + return last diff --git a/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py b/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py index dd61e630d7d..449803f3c80 100644 --- a/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_bedrock_guardrail_e2e.py @@ -1,9 +1,12 @@ -"""Live e2e: Bedrock ApplyGuardrail pre_call blocks denied input on chat. +"""Live e2e: Bedrock ApplyGuardrail blocks on chat, pre_call and post_call. -Registers a default-on bedrock guardrail via POST /guardrails with identifier/ +pre_call registers a bedrock guardrail via POST /guardrails with identifier/ version from env, then sends a prompt the guardrail's configured policy denies. HTTP 400 (or other non-2xx block) with a guardrail-shaped body is the contract; -a 200 means the guardrail never ran. +a 200 means the guardrail never ran. post_call scans the MODEL OUTPUT only, so +its test makes the model echo the word the guardrail's word policy denies +(BEDROCK_GUARDRAIL_BLOCKED_WORD, default FORBIDDENWORD) and the block must +arrive without leaking the model's text. No AWS keys are passed: the gateway signs ApplyGuardrail with its own pod-identity role, since the static AWS_ACCESS_KEY_ID / AWS_SECRET_ACCESS_KEY @@ -12,17 +15,37 @@ env vars are deliberately absent from the gateway (they hijack RDS IAM auth). from __future__ import annotations +import json import os +from typing import Final import pytest - from e2e_config import unique_marker from e2e_http import UnknownApiError -from guardrails_client import GuardrailsClient, poll_until_blocked +from guardrails_client import ( + BedrockGuardrailParamsBody, + GuardrailsClient, + poll_until_blocked, +) from lifecycle import ResourceManager +from pydantic import JsonValue, TypeAdapter pytestmark = pytest.mark.e2e +_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) + + +def _without_assessments(value: JsonValue) -> JsonValue: + """The assessments echo guardrail CONFIG, not content: the stage guardrail's + topic policy is itself named after the denied word, so its label lands in + every assessment listing and would trip a leak check aimed at model output.""" + if isinstance(value, dict): + return {key: _without_assessments(child) for key, child in value.items() if key != "assessments"} + if isinstance(value, list): + return [_without_assessments(item) for item in value] + return value + + MODEL = "gemini-2.5-flash" # Matches the word/topic policy the guardrail this suite points at actually denies. # Content filters are not assumed: the guardrail resource carries no contentPolicy, @@ -42,23 +65,17 @@ class TestBedrockGuardrail: version = os.environ["BEDROCK_GUARDRAIL_VERSION"] name = f"e2e-bedrock-guard-{unique_marker()}" - guardrail_id = client.create_bedrock_guardrail( - name, identifier=identifier, version=version - ) + guardrail_id = client.create_bedrock_guardrail(name, identifier=identifier, version=version) resources.defer(lambda: client.delete_guardrail(guardrail_id)) # Selected per request rather than registered default_on, so an upstream # ApplyGuardrail failure surfaces here instead of 403ing every other suite # running against this proxy. - result = poll_until_blocked( - lambda: client.chat(scoped_key, MODEL, BLOCKED_PROMPT, guardrails=[name]) - ) + result = poll_until_blocked(lambda: client.chat(scoped_key, MODEL, BLOCKED_PROMPT, guardrails=[name])) match result: case UnknownApiError(status_code=status, body=body): - assert status in {400, 403}, ( - f"expected a guardrail block status, got {status}: {body[:400]}" - ) + assert status in {400, 403}, f"expected a guardrail block status, got {status}: {body[:400]}" body_lower = body.lower() assert any( token in body_lower @@ -72,6 +89,49 @@ class TestBedrockGuardrail: ) ), f"block body should name the guardrail reason; got: {body[:400]}" case _: - pytest.fail( - f"bedrock default-on guardrail did not block harmful prompt; got {result}" + pytest.fail(f"bedrock default-on guardrail did not block harmful prompt; got {result}") + + @pytest.mark.covers( + "guardrail.bedrock.post_call.blocks", + exercised_on=["chat_completions"], + ) + def test_bedrock_post_call_blocks_denied_model_output( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + identifier = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"] + version = os.environ["BEDROCK_GUARDRAIL_VERSION"] + blocked_word = os.environ.get("BEDROCK_GUARDRAIL_BLOCKED_WORD", "FORBIDDENWORD") + + name = f"e2e-bedrock-post-{unique_marker()}" + guardrail_id = client.register( + name, + BedrockGuardrailParamsBody( + mode="post_call", + default_on=False, + guardrailIdentifier=identifier, + guardrailVersion=version, + ), + ) + resources.defer(lambda: client.delete_guardrail(guardrail_id)) + + # post_call scans OUTPUT only, so the denied word has to come out of the + # model: ask it to echo the word verbatim. The word in the prompt itself + # is not scanned in this mode. + prompt = f"Reply with exactly this one word and nothing else: {blocked_word}" + result = poll_until_blocked(lambda: client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128)) + + match result: + case UnknownApiError(status_code=status, body=body): + # A policy block is a 400 naming the verdict; a failed + # ApplyGuardrail call surfaces as 403 "guardrail request + # failed", which must not count as a block. + assert status == 400, f"expected the guardrail block status 400, got {status}: {body[:400]}" + body_lower = body.lower() + assert any(token in body_lower for token in ("violated", "blocked", "intervened")), ( + f"block body should name the guardrail verdict; got: {body[:400]}" ) + assert blocked_word not in json.dumps(_without_assessments(_JSON.validate_json(body))), ( + f"the blocked model output must not leak into the error body; got: {body[:400]}" + ) + case _: + pytest.fail(f"bedrock post_call guardrail did not block denied model output; got {result}") diff --git a/tests/e2e/guardrails/test_guardrail_dispatch_e2e.py b/tests/e2e/guardrails/test_guardrail_dispatch_e2e.py new file mode 100644 index 00000000000..793974ccdb1 --- /dev/null +++ b/tests/e2e/guardrails/test_guardrail_dispatch_e2e.py @@ -0,0 +1,41 @@ +"""Live e2e: the per-request `guardrails` selector must fail closed. + +A request that names a guardrail is a caller asking for protection. When the +proxy does not serve that name (a typo, a deleted guardrail, or a worker that +never loaded it), answering 200 silently drops the protection the caller asked +for; the contract this test pins is a 4xx naming the unknown guardrail. +""" + +from __future__ import annotations + +import pytest +from e2e_config import unique_marker +from e2e_http import UnknownApiError, ValidationError +from guardrails_client import GuardrailsClient + +pytestmark = pytest.mark.e2e + +MODEL = "gemini-2.5-flash" + + +@pytest.mark.skip( + reason=( + "stage red: product gap, a request naming a guardrail the proxy does not " + "serve is silently served unguarded (200) instead of failing closed" + ) +) +@pytest.mark.covers( + "guardrail.dispatch.pre_call.rejects_unknown_name", + exercised_on=["chat_completions"], +) +def test_request_naming_an_unknown_guardrail_fails_closed(client: GuardrailsClient, scoped_key: str) -> None: + result = client.chat(scoped_key, MODEL, "say hi", guardrails=[f"e2e-no-such-guardrail-{unique_marker()}"]) + + match result: + case UnknownApiError(status_code=status, body=body): + assert status == 400, f"expected a 400 for an unknown guardrail name, got {status}: {body[:400]}" + assert "guardrail" in body.lower(), f"the rejection should name the guardrail; got: {body[:400]}" + case ValidationError(message=message): + assert "guardrail" in message.lower(), f"the rejection should name the guardrail; got: {message[:400]}" + case _: + pytest.fail(f"a request naming an unknown guardrail must fail closed with a 4xx; got {result}") diff --git a/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py b/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py index d117832221d..43deb279bc8 100644 --- a/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py +++ b/tests/e2e/guardrails/test_openai_moderation_guardrail_e2e.py @@ -7,7 +7,9 @@ before the upstream model runs; a prompt that trips the policy must be rejected with HTTP 400 naming the moderation policy, and the same guardrail must let a benign prompt through. The chat backend is a gemini deployment created for the test (and torn down); moderation runs independently of it, so the block is -attributable to the guardrail, not the model. +attributable to the guardrail, not the model. The same pre_call contract is +also exercised through /v1/messages (Anthropic format): a flagged prompt is +rejected with a 400 naming moderation and a benign one passes. """ from __future__ import annotations @@ -69,3 +71,46 @@ class TestOpenAIModerationGuardrail: "the same moderation guardrail must let a benign prompt through, but the " f"call returned no choices: {allowed}" ) + + @pytest.mark.covers( + "guardrail.openai_moderations.pre_call.blocks", + exercised_on=["messages"], + ) + def test_moderation_blocks_flagged_input_on_messages( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + model = client.create_backend_model(resources, prefix="e2e-moderation-msg-backend") + + name = f"e2e-openai-moderation-msg-{unique_marker()}" + guardrail_id = client.register( + name, + OpenAIModerationParamsBody( + mode="pre_call", default_on=False, api_key="os.environ/OPENAI_API_KEY" + ), + ) + resources.defer(lambda: client.delete_guardrail(guardrail_id)) + + blocked = poll_until_blocked( + lambda: client.messages(scoped_key, model, FLAGGED_PROMPT, guardrails=[name]) + ) + match blocked: + case UnknownApiError(status_code=400, body=body): + assert "moderation" in body.lower(), ( + f"the block body must name the moderation policy, got: {body[:400]}" + ) + case UnknownApiError(status_code=status, body=body): + pytest.fail( + f"expected a 400 moderation block on /v1/messages, got {status}: {body[:400]}" + ) + case _: + pytest.fail( + f"openai moderation did not block a flagged /v1/messages prompt; got {blocked}" + ) + + allowed = unwrap( + client.messages(scoped_key, model, BENIGN_PROMPT, guardrails=[name], max_tokens=64) + ) + assert allowed.content or allowed.choices, ( + "the same moderation guardrail must let a benign /v1/messages prompt through, but " + f"the response carried neither content nor choices: {allowed}" + ) diff --git a/tests/e2e/guardrails/test_presidio_masking_e2e.py b/tests/e2e/guardrails/test_presidio_masking_e2e.py new file mode 100644 index 00000000000..6d927292975 --- /dev/null +++ b/tests/e2e/guardrails/test_presidio_masking_e2e.py @@ -0,0 +1,184 @@ +"""Live e2e: the Presidio PII guardrail masks, per its configured hook point. + +pre_call: the guardrail calls the Presidio analyzer/anonymizer on the request +messages BEFORE the model runs, so the model only ever sees placeholders like +. A prompt asking the model to repeat a fake email + phone back +must come back with the placeholders echoed and the raw PII absent, on +/chat/completions and on /v1/messages (Anthropic format). + +The analyzer/anonymizer endpoints come from PRESIDIO_ANALYZER_API_BASE / +PRESIDIO_ANONYMIZER_API_BASE; missing env is a hard failure, never a skip. +Each guardrail registers with presidio_filter_scope="input" so only the +configured hook's callback exists (the default "both" adds a second post_call +output masker), and is deleted on teardown. +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Callable + +import pytest +from pydantic import BaseModel + +from e2e_config import unique_marker +from e2e_http import Result, Success +from guardrails_client import GuardrailsClient, PresidioParamsBody +from lifecycle import ResourceManager +from models import AnthropicMessagesResponse, ChatResponse + +pytestmark = pytest.mark.e2e + +MODEL = "gemini-2.5-flash" + +# A guardrail created via POST /guardrails reaches the worker that served the +# create immediately, but every other worker only picks it up on its next +# periodic DB sync (~30s), so the first requests can be served unguarded. +GUARDRAIL_PROPAGATION_DEADLINE_SECONDS = 40.0 +GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS = 5.0 + +# Presidio's anonymizer replaces a detected entity with its unnumbered type +# placeholder, e.g. . The pre_call assertions match on the bare +# token because the model is echoing the masked prompt and may not preserve the +# angle brackets; the logged payload keeps the placeholder verbatim. +MASKED_EMAIL_TOKEN = "EMAIL_ADDRESS" +MASKED_PHONE_TOKEN = "PHONE_NUMBER" + +# Fictional NANP 555 number; a standard format Presidio's phone recognizer detects. +FAKE_PHONE = "+1 415-555-0134" + + +def _presidio_bases() -> tuple[str, str]: + analyzer = os.environ.get("PRESIDIO_ANALYZER_API_BASE", "").strip() + anonymizer = os.environ.get("PRESIDIO_ANONYMIZER_API_BASE", "").strip() + if not analyzer or not anonymizer: + pytest.fail( + "Presidio e2e requires PRESIDIO_ANALYZER_API_BASE and PRESIDIO_ANONYMIZER_API_BASE " + "(the running Presidio analyzer/anonymizer services); missing env is a hard failure, not a skip" + ) + return analyzer, anonymizer + + +def _register_presidio( + client: GuardrailsClient, + resources: ResourceManager, + *, + name: str, +) -> None: + analyzer, anonymizer = _presidio_bases() + guardrail_id = client.register( + name, + PresidioParamsBody( + mode="pre_call", + default_on=False, + presidio_analyzer_api_base=analyzer, + presidio_anonymizer_api_base=anonymizer, + presidio_filter_scope="input", + ), + ) + resources.defer(lambda: client.delete_guardrail(guardrail_id)) + + +def _fake_email() -> str: + return f"jane.doe.{unique_marker()}@example.com" + + +def _pii_prompt(marker: str, email: str) -> str: + return ( + f"{marker} Repeat this sentence back to me exactly, word for word: " + f"My email address is {email} and my phone number is {FAKE_PHONE}." + ) + + +def _first_content(response: ChatResponse) -> str: + if not response.choices: + return "" + message = response.choices[0].message + return (message.content if message else None) or "" + + +def _messages_text(response: AnthropicMessagesResponse) -> str: + """The text of a /v1/messages answer, whichever shape the proxy produced + (Anthropic-native content blocks or OpenAI-normalized choices).""" + parts: list[str] = [] + for block in response.content or []: + if block.text: + parts.append(block.text) + for choice in response.choices or []: + if choice.message and choice.message.content: + parts.append(choice.message.content) + return "\n".join(parts) + + +def _assert_eventually_masked[R: BaseModel]( + fetch: Callable[[], Result[R]], extract: Callable[[R], str], *, email: str +) -> None: + """Retry the call until the response comes back masked, to the propagation + deadline. An unmasked early response is in-flight guardrail propagation, not + a failure, and neither is a transient non-Success (a replica that has not + reloaded the guardrail answers 404, the live model can rate-limit) - only a + response that still carries the raw PII at the deadline is.""" + deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS + last: str = "" + while True: + result = fetch() + match result: + case Success(data=data): + content = extract(data) + last = content + masked = MASKED_EMAIL_TOKEN in content and MASKED_PHONE_TOKEN in content and email not in content + if masked: + assert FAKE_PHONE not in content, ( + f"the raw phone number must be masked before the model sees it, but the " + f"response echoed it: {content[:300]!r}" + ) + return + case _: + last = f"" + if time.monotonic() >= deadline: + pytest.fail( + f"presidio pre_call guardrail never masked the PII within " + f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}" + ) + time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS) + + +class TestPresidioPreCallMasking: + @pytest.mark.covers( + "guardrail.presidio.pre_call.masks", + exercised_on=["chat_completions"], + ) + def test_pre_call_masks_pii_on_chat_completions( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + name = f"e2e-presidio-pre-chat-{unique_marker()}" + _register_presidio(client, resources, name=name) + + email = _fake_email() + prompt = _pii_prompt(unique_marker(), email) + + _assert_eventually_masked( + lambda: client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128), + _first_content, + email=email, + ) + + @pytest.mark.covers( + "guardrail.presidio.pre_call.masks", + exercised_on=["messages"], + ) + def test_pre_call_masks_pii_on_messages( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + name = f"e2e-presidio-pre-msg-{unique_marker()}" + _register_presidio(client, resources, name=name) + + email = _fake_email() + prompt = _pii_prompt(unique_marker(), email) + + _assert_eventually_masked( + lambda: client.messages(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=128), + _messages_text, + email=email, + ) diff --git a/tests/e2e/guardrails/test_streaming_guardrail_e2e.py b/tests/e2e/guardrails/test_streaming_guardrail_e2e.py new file mode 100644 index 00000000000..911ddf9304b --- /dev/null +++ b/tests/e2e/guardrails/test_streaming_guardrail_e2e.py @@ -0,0 +1,87 @@ +"""Live e2e: a Bedrock guardrail in during_call mode blocks a streamed chat. + +during_call runs the Bedrock ApplyGuardrail INPUT scan in an asyncio.gather +alongside the LLM call (common_request_processing.py); when the scan flags the +prompt, the raised block cancels the LLM task before the stream ever starts, so +the client sees a non-2xx JSON error - not an SSE stream, not an in-stream +error frame - and zero content chunks are delivered. + +The prompt deliberately contains the exact word the guardrail's word policy +denies (BEDROCK_GUARDRAIL_BLOCKED_WORD, default FORBIDDENWORD), so the INPUT +scan intervenes deterministically. Identifier/version come from +BEDROCK_GUARDRAIL_IDENTIFIER / BEDROCK_GUARDRAIL_VERSION like the rest of the +bedrock suite; no AWS keys are passed (the gateway signs with pod identity). +The guardrail registers default_on=False and is selected per request, so an +upstream ApplyGuardrail failure surfaces here instead of 403ing other suites. +""" + +from __future__ import annotations + +import os + +import pytest + +from e2e_config import unique_marker +from guardrails_client import ( + BedrockGuardrailParamsBody, + GuardrailsClient, + poll_until_blocked_stream, +) +from lifecycle import ResourceManager + +pytestmark = pytest.mark.e2e + +MODEL = "gemini-2.5-flash" + + +class TestBedrockDuringCallStreaming: + @pytest.mark.covers( + "guardrail.bedrock.during.blocks", + exercised_on=["chat_completions"], + ) + def test_during_call_blocks_stream_before_first_chunk( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + identifier = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"] + version = os.environ["BEDROCK_GUARDRAIL_VERSION"] + blocked_word = os.environ.get("BEDROCK_GUARDRAIL_BLOCKED_WORD", "FORBIDDENWORD") + + name = f"e2e-bedrock-during-{unique_marker()}" + guardrail_id = client.register( + name, + BedrockGuardrailParamsBody( + mode="during_call", + default_on=False, + guardrailIdentifier=identifier, + guardrailVersion=version, + ), + ) + resources.defer(lambda: client.delete_guardrail(guardrail_id)) + + # The denied word sits in the INPUT: during_call scans the request + # messages while the model call runs, and the flag must win the race + # by cancelling the stream outright. + prompt = f"Please use the word {blocked_word} in a sentence." + result = poll_until_blocked_stream( + lambda: client.chat_stream_raw(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=64) + ) + + assert not result.ok, ( + f"the during_call guardrail never blocked the streamed request; got a " + f"{result.status_code} with {result.chunks} chunks" + ) + assert result.status_code == 400, ( + f"a during_call block surfaces as HTTP 400 before the stream starts, got " + f"{result.status_code}: {result.body[:400]}" + ) + assert result.chunks == 0 and not result.stream_events, ( + f"no content chunk may be delivered on a during_call block, but " + f"{result.chunks} chunks arrived: {result.stream_events[:3]}" + ) + assert "text/event-stream" not in (result.content_type or ""), ( + f"the block must be a JSON error response, not an SSE stream; got content-type {result.content_type!r}" + ) + body_lower = result.body.lower() + assert any(token in body_lower for token in ("guardrail", "violated", "blocked", "bedrock", "intervened")), ( + f"block body should name the guardrail reason; got: {result.body[:400]}" + ) diff --git a/tests/e2e/logging/conftest.py b/tests/e2e/logging/conftest.py index 60536ea01d4..621595e6b46 100644 --- a/tests/e2e/logging/conftest.py +++ b/tests/e2e/logging/conftest.py @@ -47,6 +47,4 @@ def dd_logs() -> DdLogsReader: def datadog_creds() -> None: """Require Datadog shipping credentials. Hard-fail when absent; never skip.""" if not (os.getenv("DD_API_KEY") and os.getenv("DD_SITE")): - pytest.fail( - "Datadog e2e requires DD_API_KEY and DD_SITE; missing credentials is a hard failure, not a skip" - ) + pytest.fail("Datadog e2e requires DD_API_KEY and DD_SITE; missing credentials is a hard failure, not a skip") diff --git a/tests/e2e/logging/datadog_reader.py b/tests/e2e/logging/datadog_reader.py index 7d882a7fa81..d0f478185c2 100644 --- a/tests/e2e/logging/datadog_reader.py +++ b/tests/e2e/logging/datadog_reader.py @@ -97,15 +97,22 @@ class DdLogsReader: indexed ``message`` empty, so a plain full-text query matches nothing; ``*:`` extends the scan to every attribute (the marker sits in the prompt, e.g. ``messages.content``, wherever the route's payload puts - it). More than one hit for one call IS the duplicate-delivery bug, so - this never collapses to a single event. A 429 backs off and retries - - the search budget is org-wide, so another consumer can empty it under - us - while any other failure stays a hard fail.""" + it).""" + return self.events_for_query(f"*:*{marker}*") + + def events_for_query(self, query: str) -> list[DdLogEvent]: + """Every ingested event the search query matches (failure payloads + carry no prompt to mark, so failure scenarios query indexed attributes + like ``@model_group:...`` instead of a body marker). More than one hit + for one call IS the duplicate-delivery bug, so this never collapses to + a single event. A 429 backs off and retries - the search budget is + org-wide, so another consumer can empty it under us - while any other + failure stays a hard fail.""" for _ in range(_RATE_LIMIT_RETRIES): result = post( URL(f"https://api.{self.site}/api/v2/logs/events/search"), headers=_DdAuthHeaders(api_key=self.api_key, app_key=self.app_key), - json=_SearchRequest(filter=_SearchFilter(query=f"*:*{marker}*")), + json=_SearchRequest(filter=_SearchFilter(query=query)), response_type=_SearchResponse, timeout=30.0, ) @@ -123,6 +130,10 @@ class DdLogsReader: ) def poll_events_for_marker(self, marker: str) -> list[DdLogEvent]: + """``poll_events_for_query`` over the every-attribute marker scan.""" + return self.poll_events_for_query(f"*:*{marker}*") + + def poll_events_for_query(self, query: str) -> list[DdLogEvent]: """Poll until at least one matching event is searchable (the callback flushes in periodic batches and DataDog ingestion adds seconds of lag), then keep re-reading for DD_SETTLE_SECONDS so a late duplicate cannot @@ -132,15 +143,13 @@ class DdLogsReader: request budget. At the deadline the last result is returned as-is.""" deadline = time.monotonic() + POLL_TIMEOUT while time.monotonic() < deadline: - events = self.events_for_marker(marker) + events = self.events_for_query(query) if events: - return self._settled_events_for_marker(marker, events) + return self._settled_events_for_query(query, events) time.sleep(DD_SEARCH_INTERVAL) - return self.events_for_marker(marker) + return self.events_for_query(query) - def _settled_events_for_marker( - self, marker: str, events: list[DdLogEvent] - ) -> list[DdLogEvent]: + def _settled_events_for_query(self, query: str, events: list[DdLogEvent]) -> list[DdLogEvent]: """Re-read at every search interval until the settle window closes; a duplicate ends the watch early because more waiting cannot clear it. @@ -151,7 +160,7 @@ class DdLogsReader: last_nonempty = events while time.monotonic() < settle_deadline: time.sleep(DD_SEARCH_INTERVAL) - latest = self.events_for_marker(marker) + latest = self.events_for_query(query) if not latest: continue if len(latest) > 1: diff --git a/tests/e2e/logging/gcs_reader.py b/tests/e2e/logging/gcs_reader.py new file mode 100644 index 00000000000..60622c121ac --- /dev/null +++ b/tests/e2e/logging/gcs_reader.py @@ -0,0 +1,220 @@ +"""Read-back for the gcs_bucket logging test against the real GCS bucket. + +The proxy ships StandardLoggingPayload objects with its own service account +(litellm_settings.callbacks: ["gcs_bucket"] + GCS_BUCKET_NAME), and the test +reads them back through the GCS JSON API. Auth is a self-signed service-account +JWT (RS256 via PyJWT + cryptography, both litellm proxy dependencies the +runner installs) minted per request and sent directly as the Bearer token - +Google accepts that for storage.googleapis.com with no token exchange, which +keeps every HTTP read inside ``e2e_http``. + +The default gcs_bucket mode batches payloads into ``{date}/batch-{id}.ndjson`` +objects; unbatched mode writes ``{date}/{response_id}`` per call. The reader +handles both: it polls the day's listing, downloads the direct object when +present, and otherwise scans batch objects fresh enough to hold the call. +Missing configuration is a hard failure, never a skip. +""" + +from __future__ import annotations + +import os +import time +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from pathlib import Path +from urllib.parse import quote + +import jwt +import pytest +from pydantic import BaseModel, ConfigDict, Field + +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT +from e2e_http import URL, Headers, probe + +_GCS_API = "https://storage.googleapis.com" +#: Tolerance for clock skew between this host and GCS object timestamps. +_SKEW = timedelta(seconds=120) +#: How long to keep re-reading after the first match before trusting the +#: exactly-one assertion: past one full gcs_bucket flush interval (~20s), so +#: a duplicate shipped by a later flush is seen, plus listing-latency margin. +GCS_SETTLE_SECONDS = 45.0 + + +class _ServiceAccount(BaseModel): + model_config = ConfigDict(extra="ignore") + + client_email: str + private_key: str + + +class _GcsAuthHeaders(Headers): + authorization: str = Field(serialization_alias="Authorization") + + +class _GcsObject(BaseModel): + model_config = ConfigDict(extra="ignore") + + name: str + updated: datetime | None = None + + +class _GcsListResponse(BaseModel): + model_config = ConfigDict(extra="ignore") + + items: list[_GcsObject] = [] + next_page_token: str | None = Field(default=None, validation_alias="nextPageToken") + + +class _GcsListParams(BaseModel): + prefix: str + max_results: int = Field(default=1000, serialization_alias="maxResults") + page_token: str | None = Field(default=None, serialization_alias="pageToken") + + +class _GcsMediaParams(BaseModel): + alt: str = "media" + + +class GcsLogRecord(BaseModel): + """The StandardLoggingPayload fields the gcs scenario pins.""" + + model_config = ConfigDict(extra="ignore") + + id: str + status: str + model_group: str | None = None + response_cost: float | None = None + total_tokens: int | None = None + error_str: str | None = None + + +def _mint_bearer(account: _ServiceAccount) -> str: + """Self-signed service-account JWT: for Google APIs a token whose ``aud`` + is the service endpoint authorizes directly, no oauth2 token exchange. + Minted per request so a long session never outlives one token's expiry.""" + now = int(time.time()) + claims: dict[str, str | int] = { + "iss": account.client_email, + "sub": account.client_email, + "aud": f"{_GCS_API}/", + "iat": now, + "exp": now + 3600, + } + return jwt.encode(claims, account.private_key, algorithm="RS256") + + +@dataclass(frozen=True, slots=True) +class GcsLogReader: + bucket: str + account: _ServiceAccount + + def _headers(self) -> _GcsAuthHeaders: + return _GcsAuthHeaders(authorization=f"Bearer {_mint_bearer(self.account)}") + + def _list(self, prefix: str) -> list[_GcsObject]: + """Every object under ``prefix``, following ``nextPageToken`` - the + shared day prefix accumulates all of the proxy's traffic, and a fresh + record past the 1000-object page cap must still be seen.""" + items: list[_GcsObject] = [] + page_token: str | None = None + while True: + result = probe( + URL(f"{_GCS_API}/storage/v1/b/{self.bucket}/o"), + headers=self._headers(), + params=_GcsListParams(prefix=prefix, page_token=page_token), + ) + if result.status_code != 200: + pytest.fail( + f"GCS object listing for gs://{self.bucket}/{prefix} failed " + f"({result.status_code}): {result.body[:300]}" + ) + page = _GcsListResponse.model_validate_json(result.body) + items.extend(page.items) + page_token = page.next_page_token + if not page_token: + return items + + def _download(self, name: str) -> str: + result = probe( + URL(f"{_GCS_API}/storage/v1/b/{self.bucket}/o/{quote(name, safe='')}"), + headers=self._headers(), + params=_GcsMediaParams(), + ) + if result.status_code != 200: + pytest.fail( + f"GCS object download gs://{self.bucket}/{name} failed ({result.status_code}): {result.body[:300]}" + ) + return result.body + + def records_for_response_id(self, response_id: str, *, since: datetime) -> list[GcsLogRecord]: + """Every payload written for ``response_id``: the direct + ``{date}/{response_id}`` object plus any hit inside batch NDJSON + objects updated after ``since``. More than one hit is the + duplicate-delivery bug, so this never collapses to a single record.""" + records: list[GcsLogRecord] = [] + window_start = since - _SKEW + for day_offset in (-1, 0, 1): + day = (since + timedelta(days=day_offset)).strftime("%Y-%m-%d") + for obj in self._list(f"{day}/"): + if obj.name == f"{day}/{response_id}": + records.append(GcsLogRecord.model_validate_json(self._download(obj.name))) + continue + is_fresh_batch = f"{day}/batch-" in obj.name and obj.updated is not None and obj.updated >= window_start + if is_fresh_batch: + records.extend( + GcsLogRecord.model_validate_json(line) + for line in self._download(obj.name).splitlines() + if response_id in line + ) + return records + + def poll_records_for_response_id(self, response_id: str, *, since: datetime) -> list[GcsLogRecord]: + """Poll until the payload is readable (the gcs_bucket callback flushes + on a ~20s timer), then keep re-reading for GCS_SETTLE_SECONDS - past a + full flush interval - so a duplicate shipped by a later flush cannot + hide from the exactly-one assertion. A duplicate ends the settle early + because more waiting cannot clear it.""" + deadline = time.monotonic() + POLL_TIMEOUT + while time.monotonic() < deadline: + records = self.records_for_response_id(response_id, since=since) + if records: + return self._settled_records(response_id, since=since, first=records) + time.sleep(POLL_INTERVAL) + return [] + + def _settled_records(self, response_id: str, *, since: datetime, first: list[GcsLogRecord]) -> list[GcsLogRecord]: + """Re-read at every poll interval until the settle window closes; a + transiently empty re-read never downgrades what was already seen.""" + settle_deadline = time.monotonic() + GCS_SETTLE_SECONDS + latest = first + while time.monotonic() < settle_deadline and len(latest) <= 1: + time.sleep(POLL_INTERVAL) + latest = self.records_for_response_id(response_id, since=since) or latest + return latest + + +def utc_now() -> datetime: + return datetime.now(timezone.utc) + + +def build_gcs_reader() -> GcsLogReader: + bucket = os.environ.get("GCS_BUCKET_NAME", "") + if not bucket: + pytest.fail( + "GCS_BUCKET_NAME must be set: the gcs test reads the proxy's gcs_bucket " + "delivery back from the real bucket (the cluster secret manager injects " + "it; locally set it in tests/e2e/.env)" + ) + raw = "" + credentials_path = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", "") + if credentials_path and Path(credentials_path).is_file(): + raw = Path(credentials_path).read_text() + else: + raw = os.environ.get("VERTEXAI_CREDENTIALS", "") + if not raw: + pytest.fail( + "GCS read-back needs a service-account key: set " + "GOOGLE_APPLICATION_CREDENTIALS (path) or VERTEXAI_CREDENTIALS (JSON), " + "as the cluster secret manager does" + ) + return GcsLogReader(bucket=bucket, account=_ServiceAccount.model_validate_json(raw)) diff --git a/tests/e2e/logging/logging_client.py b/tests/e2e/logging/logging_client.py index d76f7b356b2..f0f7ad7eaa4 100644 --- a/tests/e2e/logging/logging_client.py +++ b/tests/e2e/logging/logging_client.py @@ -480,12 +480,8 @@ class LoggingClient: stream=True if stream else None, ) if stream: - return self.proxy.transport.stream( - "/v1/messages", headers=self.proxy.transport.bearer(key), json=body - ) - return self.proxy.transport.send( - "/v1/messages", headers=self.proxy.transport.bearer(key), json=body - ) + return self.proxy.transport.stream("/v1/messages", headers=self.proxy.transport.bearer(key), json=body) + return self.proxy.transport.send("/v1/messages", headers=self.proxy.transport.bearer(key), json=body) def responses_raw( self, key: str, model: str, text: str, *, max_output_tokens: int = 64, stream: bool = False @@ -499,12 +495,8 @@ class LoggingClient: model=model, input=text, max_output_tokens=max_output_tokens, stream=True if stream else None ) if stream: - return self.proxy.transport.stream( - "/v1/responses", headers=self.proxy.transport.bearer(key), json=body - ) - return self.proxy.transport.send( - "/v1/responses", headers=self.proxy.transport.bearer(key), json=body - ) + return self.proxy.transport.stream("/v1/responses", headers=self.proxy.transport.bearer(key), json=body) + return self.proxy.transport.send("/v1/responses", headers=self.proxy.transport.bearer(key), json=body) def scrape_metrics(self) -> str: return self.proxy.probe("/metrics", params=NoBody()).body @@ -530,9 +522,7 @@ class LoggingClient: return False return True - rows = self.proxy.poll_logs_for_key( - key, min_rows=1, predicate=lambda rs: any(_matches(r) for r in rs) - ) + rows = self.proxy.poll_logs_for_key(key, min_rows=1, predicate=lambda rs: any(_matches(r) for r in rs)) for row in rows: if _matches(row): return row @@ -593,9 +583,7 @@ class LoggingClient: deadline = time.monotonic() + POLL_TIMEOUT last: LangfuseObservation | None = None while time.monotonic() < deadline: - last = self.find_langfuse_observation( - creds, key_alias=key_alias, prompt_marker=prompt_marker - ) + last = self.find_langfuse_observation(creds, key_alias=key_alias, prompt_marker=prompt_marker) if last is not None: cost = observation_spend(last) if not require_positive_cost or (cost is not None and cost > 0): @@ -611,9 +599,7 @@ class LoggingClient: prompt_marker: str, ) -> list[LangfuseObservation]: """Generation plus any sibling/child observations (guardrail spans, etc.).""" - gen = self.poll_langfuse_observation( - creds, key_alias=key_alias, prompt_marker=prompt_marker - ) + gen = self.poll_langfuse_observation(creds, key_alias=key_alias, prompt_marker=prompt_marker) if gen is None or not gen.trace_id: return [] if gen is None else [gen] return self.list_langfuse_observations(creds, trace_id=gen.trace_id) or [gen] @@ -636,3 +622,15 @@ def first_ok(client: LoggingClient, send: Callable[[], StreamingResponse]) -> St def build_logging_client(proxy: ProxyClient) -> LoggingClient: return LoggingClient(proxy=proxy) + + +def readiness_details_body(client: LoggingClient) -> str: + """/health/readiness/details, tolerating the 503 it serves while the + ephemeral stack's DB leg blips: the recorded state the logging suites check + here is the callback list, which the body carries either way.""" + result = client.proxy.probe("/health/readiness/details", params=NoBody()) + db_blip = result.status_code == 503 and '"db":"disconnected"' in result.body + assert result.status_code == 200 or db_blip, ( + f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}" + ) + return result.body diff --git a/tests/e2e/logging/s3_reader.py b/tests/e2e/logging/s3_reader.py new file mode 100644 index 00000000000..d605dec6096 --- /dev/null +++ b/tests/e2e/logging/s3_reader.py @@ -0,0 +1,115 @@ +"""Read-back for the s3 logging tests against the real S3 bucket the proxy +ships StandardLoggingPayload objects to (litellm_settings.callbacks: ["s3_v2"]). + +Delivery is judged on what actually landed in the bucket: the proxy writes +with its own credentials exactly as in production, and the tests list and +download the objects back with boto3 (already a litellm proxy dependency, so +the e2e runner image carries it; it is an AWS SDK, not a raw HTTP client, so +the e2e_http-only transport rule is untouched). The bucket comes from +S3_LOGS_BUCKET_NAME - on the cluster the secret manager injects it, locally +tests/e2e/.env provides it. Missing configuration is a hard failure, never a +skip. +""" + +from __future__ import annotations + +import os +import time +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import boto3 +import pytest +from pydantic import BaseModel, ConfigDict + +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT + +if TYPE_CHECKING: + from types_boto3_s3.client import S3Client + +#: How long to keep re-reading after the first match before trusting the +#: exactly-one assertion: past one full s3_v2 flush interval (~10s), so a +#: duplicate shipped by a LATER flush is seen, plus listing-latency margin. +#: The DataDog reader settles the same way (DD_SETTLE_SECONDS). +S3_SETTLE_SECONDS = 25.0 + + +class S3LogRecord(BaseModel): + """The StandardLoggingPayload fields the s3 scenarios pin.""" + + model_config = ConfigDict(extra="ignore") + + id: str + status: str + model_group: str | None = None + response_cost: float | None = None + total_tokens: int | None = None + error_str: str | None = None + + +@dataclass(frozen=True, slots=True) +class S3LogReader: + bucket: str + client: S3Client + + def list_keys(self, prefix: str) -> list[str]: + response = self.client.list_objects_v2(Bucket=self.bucket, Prefix=prefix) + return [obj["Key"] for obj in response.get("Contents", []) if "Key" in obj] + + def read_record(self, key: str) -> S3LogRecord: + body = self.client.get_object(Bucket=self.bucket, Key=key)["Body"].read() + return S3LogRecord.model_validate_json(body) + + def records_matching(self, *, prefix: str, predicate: Callable[[S3LogRecord], bool]) -> list[S3LogRecord]: + return [record for record in map(self.read_record, self.list_keys(prefix)) if predicate(record)] + + def poll_records(self, *, prefix: str, predicate: Callable[[S3LogRecord], bool]) -> list[S3LogRecord]: + """Poll until at least one matching object is listed (the s3_v2 + callback flushes on a ~10s timer), then keep re-reading for + S3_SETTLE_SECONDS - past a full flush interval - so a duplicate + shipped by a later flush cannot hide from the exactly-one assertion. + One blind spot is inherent: a duplicate write that reuses the exact + same object key overwrites the first object and no listing can see + it; distinct-key duplicates are what this catches. At the deadline an + empty list is returned and the caller's assertion carries the failure + message.""" + deadline = time.monotonic() + POLL_TIMEOUT + while time.monotonic() < deadline: + records = self.records_matching(prefix=prefix, predicate=predicate) + if records: + return self._settled_records(prefix=prefix, predicate=predicate, first=records) + time.sleep(POLL_INTERVAL) + return [] + + def _settled_records( + self, *, prefix: str, predicate: Callable[[S3LogRecord], bool], first: list[S3LogRecord] + ) -> list[S3LogRecord]: + """Re-read at every poll interval until the settle window closes; a + duplicate ends the watch early because more waiting cannot clear it. + A transiently empty re-read never downgrades what was already seen.""" + settle_deadline = time.monotonic() + S3_SETTLE_SECONDS + latest = first + while time.monotonic() < settle_deadline and len(latest) <= 1: + time.sleep(POLL_INTERVAL) + latest = self.records_matching(prefix=prefix, predicate=predicate) or latest + return latest + + +def build_s3_reader() -> S3LogReader: + bucket = os.environ.get("S3_LOGS_BUCKET_NAME", "") + if not bucket: + pytest.fail( + "S3_LOGS_BUCKET_NAME must be set: the s3 tests read the proxy's s3_v2 " + "delivery back from the real bucket (the cluster secret manager injects " + "it; locally set it in tests/e2e/.env to the same bucket " + "s3_callback_params.s3_bucket_name names)" + ) + region = os.environ.get("AWS_REGION_NAME") or os.environ.get("AWS_REGION") or "us-east-1" + return S3LogReader( + bucket=bucket, + # boto3.client's overload set covers every AWS service; the ones without + # installed stubs type as Unknown, so the member is "partially unknown" + # even though the s3 overload itself resolves to S3Client. + client=boto3.client("s3", region_name=region), # pyright: ignore[reportUnknownMemberType] + ) diff --git a/tests/e2e/logging/test_datadog_log_e2e.py b/tests/e2e/logging/test_datadog_log_e2e.py index 94811c6217e..a4821ed058b 100644 --- a/tests/e2e/logging/test_datadog_log_e2e.py +++ b/tests/e2e/logging/test_datadog_log_e2e.py @@ -19,15 +19,16 @@ received). from __future__ import annotations import math +import time import pytest from pydantic import BaseModel, ConfigDict from datadog_reader import DdLogEvent, DdLogsReader from e2e_config import CHEAP_ANTHROPIC_MODEL, CHEAP_OPENAI_MODEL, unique_marker -from e2e_http import NoBody from lifecycle import ResourceManager -from logging_client import LoggingClient, first_ok +from logging_client import INVALID_UPSTREAM_API_KEY, LoggingClient, first_ok, readiness_details_body +from models import LiteLLMParamsBody pytestmark = pytest.mark.e2e @@ -46,19 +47,17 @@ class _DdMessagePayload(BaseModel): status: str call_type: str stream: bool | None = None + error_str: str | None = None def _assert_datadog_configured(client: LoggingClient) -> None: """Recorded state: the proxy reports the DataDog callback among its active callbacks, so a missing destination config fails here, before any delivery-based assertion can time out confusingly.""" - result = client.proxy.probe("/health/readiness/details", params=NoBody()) - assert result.status_code == 200, ( - f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}" - ) - assert DD_LOGGER_NAME in result.body, ( + body = readiness_details_body(client) + assert DD_LOGGER_NAME in body, ( f"the proxy must report the {DD_LOGGER_NAME} callback active " - f"(callbacks + DD_* env in the compose config); got: {result.body[:400]}" + f"(callbacks + DD_* env in the compose config); got: {body[:400]}" ) @@ -89,18 +88,14 @@ def _assert_exactly_one_event( # indexed event status from the parsed payload's status attribute # ("success") and normalizes it to its OK severity - so "ok" is what a # successfully ingested success event looks like on the search API. - assert event.status == "ok", ( - f"success events must index at DataDog's ok severity, got {event.status!r}" - ) + assert event.status == "ok", f"success events must index at DataDog's ok severity, got {event.status!r}" payload = _DdMessagePayload.model_validate(event.attributes) assert payload.status == "success", f"payload status must be success, got {payload.status!r}" assert payload.model_group == model_group, ( f"payload model_group must be {model_group!r}, got {payload.model_group!r}" ) - assert payload.call_type == call_type, ( - f"payload call_type must be {call_type!r}, got {payload.call_type!r}" - ) + assert payload.call_type == call_type, f"payload call_type must be {call_type!r}, got {payload.call_type!r}" assert payload.total_tokens > 0, f"payload must count real tokens, got {payload.total_tokens}" # Relative tolerance, not bit-equality: the cost round-trips through # DataDog's attribute indexing, whose float serialization may drift in the @@ -109,9 +104,7 @@ def _assert_exactly_one_event( f"payload response_cost {payload.response_cost} must equal the anchor cost {cost_anchor}" ) if expect_stream: - assert payload.stream is True, ( - f"a streamed call's payload must record stream=true, got {payload.stream!r}" - ) + assert payload.stream is True, f"a streamed call's payload must record stream=true, got {payload.stream!r}" return payload @@ -211,7 +204,9 @@ class TestDataDogLogDelivery: marker = unique_marker() outcome = first_ok( client, - lambda: client.chat_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", stream=True, max_tokens=16), + lambda: client.chat_raw( + key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", stream=True, max_tokens=16 + ), ) assert outcome.is_streaming, f"response must be an event stream, got content-type {outcome.content_type!r}" assert outcome.chunks > 0, "the stream must deliver at least one event" @@ -231,9 +226,7 @@ class TestDataDogLogDelivery: cost_anchor=spend_row.spend, expect_stream=True, ) - assert spend_row.total_tokens is not None, ( - "the spend row must record total_tokens for the token cross-check" - ) + assert spend_row.total_tokens is not None, "the spend row must record total_tokens for the token cross-check" assert spend_row.total_tokens == payload.total_tokens, ( f"the spend row and the DataDog event must agree on tokens: " f"{spend_row.total_tokens} vs {payload.total_tokens}" @@ -255,7 +248,9 @@ class TestDataDogLogDelivery: marker = unique_marker() outcome = first_ok( client, - lambda: client.messages_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=16, stream=True), + lambda: client.messages_raw( + key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=16, stream=True + ), ) assert outcome.is_streaming, f"response must be an event stream, got content-type {outcome.content_type!r}" assert outcome.chunks > 0, "the stream must deliver at least one event" @@ -275,9 +270,7 @@ class TestDataDogLogDelivery: cost_anchor=spend_row.spend, expect_stream=True, ) - assert spend_row.total_tokens is not None, ( - "the spend row must record total_tokens for the token cross-check" - ) + assert spend_row.total_tokens is not None, "the spend row must record total_tokens for the token cross-check" assert spend_row.total_tokens == payload.total_tokens, ( f"the spend row and the DataDog event must agree on tokens: " f"{spend_row.total_tokens} vs {payload.total_tokens}" @@ -319,10 +312,89 @@ class TestDataDogLogDelivery: cost_anchor=spend_row.spend, expect_stream=True, ) - assert spend_row.total_tokens is not None, ( - "the spend row must record total_tokens for the token cross-check" - ) + assert spend_row.total_tokens is not None, "the spend row must record total_tokens for the token cross-check" assert spend_row.total_tokens == payload.total_tokens, ( f"the spend row and the DataDog event must agree on tokens: " f"{spend_row.total_tokens} vs {payload.total_tokens}" ) + + +def _assert_exactly_one_failure_event(events: list[DdLogEvent], *, model_group: str) -> _DdMessagePayload: + """The enforced behavior for a failed call: the intake holds exactly one + event for the deployment, sourced from litellm, indexed at an error-grade + severity (DataDog derives it from the payload's status="failure"; observed + as its "emergency" bucket), whose payload carries the provider error and + no cost.""" + assert events, "no DataDog log event for the failed call reached the intake within the deadline" + assert len(events) == 1, ( + f"expected exactly ONE DataDog log event for the failed call, got {len(events)} - " + "more than one event for one call is the duplicate-delivery bug" + ) + event = events[0] + assert "source:litellm" in event.tags, ( + f"the ingested event must carry the litellm source (shipped as ddsource), got tags {event.tags!r}" + ) + assert event.status in ("error", "emergency"), ( + f"failure events must index at an error-grade severity, got {event.status!r}" + ) + payload = _DdMessagePayload.model_validate(event.attributes) + assert payload.status == "failure", f"payload status must be failure, got {payload.status!r}" + assert payload.model_group == model_group, ( + f"payload model_group must be {model_group!r}, got {payload.model_group!r}" + ) + assert not payload.response_cost, f"a failed call must not be billed, got response_cost={payload.response_cost!r}" + return payload + + +class TestDataDogFailureDelivery: + @pytest.mark.covers("logging.datadog.failure.exports_metric", exercised_on=["chat_completions"]) + def test_failed_chat_completions_emits_one_error_event( + self, client: LoggingClient, dd_logs: DdLogsReader, resources: ResourceManager + ) -> None: + """A /chat/completions call that fails at the provider must reach the + DataDog logs intake as exactly one error-grade event carrying the + provider error - failure metrics drive alerting and SLOs, so a dropped + failure event is an invisible outage. + + A deployment with an invalid upstream key lets the request pass proxy + auth and fail at the provider (the same lever as the OTEL error test). + Failure payloads carry no prompt to mark, so the read-back queries the + indexed @model_group attribute of the per-run unique deployment name; + proxy-side 401s during key propagation never reach the provider and + ship no payload, so exactly one provider failure exists for it.""" + _assert_datadog_configured(client) + + model_name = f"dd-err-{unique_marker()}" + model_id = client.create_model( + model_name, + LiteLLMParamsBody(model="anthropic/claude-haiku-4-5", api_key=INVALID_UPSTREAM_API_KEY), + ) + resources.defer(lambda: client.delete_model(model_id)) + key = client.key_with_alias(f"dd-err-key-{unique_marker()}", models=[model_name]) + resources.defer(lambda: client.delete_key(key)) + + deadline = time.monotonic() + client.proxy.poll_timeout + while True: + outcome = client.chat_raw(key, model_name, "trigger an upstream auth failure", max_tokens=16) + assert not outcome.ok, "the call must fail; the deployment's upstream key is invalid" + assert outcome.status_code != -1, ( + "network failure between the test and the proxy while provoking the provider " + "failure; retrying now could double-log the failure payload and falsely trip " + f"the exactly-one assertion - fix the rig connectivity first: {outcome.body[:200]}" + ) + if "AnthropicException" in outcome.body or time.monotonic() >= deadline: + break + time.sleep(client.proxy.poll_interval) + assert "AnthropicException" in outcome.body, ( + "never saw the upstream provider failure before the deadline; the key may still be " + f"propagating - last outcome {outcome.status_code}: {outcome.body[:200]}" + ) + assert outcome.status_code == 401, ( + f"an upstream auth failure must map to 401, got {outcome.status_code}: {outcome.body[:200]}" + ) + + events = dd_logs.poll_events_for_query(f"@model_group:{model_name}") + payload = _assert_exactly_one_failure_event(events, model_group=model_name) + assert payload.error_str is not None and "AnthropicException" in payload.error_str, ( + f"the event must carry the provider error, got error_str={payload.error_str!r}" + ) diff --git a/tests/e2e/logging/test_gcs_log_e2e.py b/tests/e2e/logging/test_gcs_log_e2e.py new file mode 100644 index 00000000000..17ad1507049 --- /dev/null +++ b/tests/e2e/logging/test_gcs_log_e2e.py @@ -0,0 +1,97 @@ +"""Live e2e: gcs_bucket log delivery for successful calls. + +Covers logging.gcs_bucket.success.writes_object: one successful +/chat/completions call must land in the real GCS bucket as exactly one +StandardLoggingPayload record (GCS is the audit-trail parallel to S3 for GCP +deployments). Delivery is judged on what is actually readable in the bucket: +the proxy writes with its production service account, and the test reads the +record back through the GCS JSON API - covering both the batched NDJSON layout +(the default) and the per-request object layout. + +Both halves of the contract are asserted: the recorded state (the proxy +reports the GCSBucketLogger callback active via /health/readiness/details - +note gcs_bucket is enterprise-gated, so this also requires a license) and the +enforced behavior (the record in the bucket, cost cross-checked against the +x-litellm-response-cost header of the very response the caller received). +""" + +from __future__ import annotations + +import math + +import pytest + +from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker +from gcs_reader import GcsLogReader, build_gcs_reader, utc_now +from lifecycle import ResourceManager +from logging_client import LoggingClient, completion_response_id, first_ok, readiness_details_body + +pytestmark = pytest.mark.e2e + +#: The active gcs_bucket callback's name in /health/readiness/details success_callbacks. +GCS_LOGGER_NAME = "GCSBucketLogger" + + +@pytest.fixture(scope="session") +def gcs_logs() -> GcsLogReader: + return build_gcs_reader() + + +def _assert_gcs_configured(client: LoggingClient) -> None: + """Recorded state: the proxy reports the gcs_bucket callback among its + active callbacks, so a missing destination config (or a missing enterprise + license - gcs_bucket refuses to initialize without one) fails here, before + any delivery-based assertion can time out confusingly.""" + body = readiness_details_body(client) + assert GCS_LOGGER_NAME in body, ( + f"the proxy must report the {GCS_LOGGER_NAME} callback active " + f"(litellm_settings.callbacks: ['gcs_bucket'] + GCS_BUCKET_NAME env + enterprise license); " + f"got: {body[:400]}" + ) + + +class TestGcsLogDelivery: + @pytest.mark.covers("logging.gcs_bucket.success.writes_object", exercised_on=["chat_completions"]) + def test_chat_completions_writes_one_success_record( + self, client: LoggingClient, gcs_logs: GcsLogReader, resources: ResourceManager + ) -> None: + """One successful non-streaming /chat/completions call must be + readable back from the bucket as exactly one payload record carrying + the model group, the token counts, and the same cost the caller's + response header reported.""" + _assert_gcs_configured(client) + + alias = f"gcs-chat-{unique_marker()}" + key = client.key_with_alias(alias, models=[CHEAP_ANTHROPIC_MODEL]) + resources.defer(lambda: client.delete_key(key)) + + since = utc_now() + marker = unique_marker() + outcome = first_ok( + client, + lambda: client.chat_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=16), + ) + assert outcome.response_cost is not None and outcome.response_cost > 0, ( + f"the response must report x-litellm-response-cost, got {outcome.response_cost!r}" + ) + body_id = completion_response_id(outcome.body) + assert body_id is not None, "the completion body must carry an id (it names the gcs record)" + + records = gcs_logs.poll_records_for_response_id(body_id, since=since) + assert records, f"no gcs record for response {body_id} was readable from the bucket within the deadline" + assert len(records) == 1, ( + f"expected exactly ONE gcs record for the call, got {len(records)} - " + "more than one record for one call is the duplicate-delivery bug" + ) + record = records[0] + assert record.id == body_id, f"record id must be the response id, got {record.id!r}" + assert record.status == "success", f"payload status must be success, got {record.status!r}" + assert record.model_group == CHEAP_ANTHROPIC_MODEL, ( + f"payload model_group must be {CHEAP_ANTHROPIC_MODEL!r}, got {record.model_group!r}" + ) + assert record.total_tokens is not None and record.total_tokens > 0, ( + f"payload must count real tokens, got {record.total_tokens!r}" + ) + assert record.response_cost is not None and math.isclose( + record.response_cost, outcome.response_cost, rel_tol=1e-9 + ), f"payload response_cost {record.response_cost!r} must equal the header cost {outcome.response_cost}" diff --git a/tests/e2e/logging/test_otel_trace_e2e.py b/tests/e2e/logging/test_otel_trace_e2e.py index 52cb691e2b7..9f08fa6c4e7 100644 --- a/tests/e2e/logging/test_otel_trace_e2e.py +++ b/tests/e2e/logging/test_otel_trace_e2e.py @@ -23,9 +23,8 @@ import pytest from pydantic import BaseModel, ConfigDict, ValidationError from e2e_config import CHEAP_ANTHROPIC_MODEL, CHEAP_OPENAI_MODEL, unique_marker -from e2e_http import NoBody from lifecycle import ResourceManager -from logging_client import INVALID_UPSTREAM_API_KEY, LoggingClient, first_ok +from logging_client import INVALID_UPSTREAM_API_KEY, LoggingClient, first_ok, readiness_details_body from models import LiteLLMParamsBody from otel_client import JaegerSpan, JaegerTrace, OtelReader @@ -48,11 +47,7 @@ def _assert_otel_destination_configured(client: LoggingClient) -> None: """Recorded state: the proxy reports the OTEL v2 logger among its active callbacks, so a missing/failed destination config fails here, before any traffic-based assertion can time out confusingly.""" - result = client.proxy.probe("/health/readiness/details", params=NoBody()) - assert result.status_code == 200, ( - f"/health/readiness/details must answer 200, got {result.status_code}: {result.body[:300]}" - ) - details = _ReadinessDetails.model_validate_json(result.body) + details = _ReadinessDetails.model_validate_json(readiness_details_body(client)) assert OTEL_V2_LOGGER_NAME in details.success_callbacks, ( f"the proxy must report the {OTEL_V2_LOGGER_NAME} callback active " f"(LITELLM_OTEL_V2 + arize_phoenix preset in the compose config); got: {details.success_callbacks}" @@ -164,17 +159,14 @@ def served_genai_spans(trace: JaegerTrace, genai_span: str) -> list[JaegerSpan]: these tests fail whenever the upstream 429s, 529s, or hands back a stale credential on the first try.""" return [ - span - for span in trace.spans - if span.operation_name == genai_span and _tag(span, ERROR_STATUS_TAG) != "ERROR" + span for span in trace.spans if span.operation_name == genai_span and _tag(span, ERROR_STATUS_TAG) != "ERROR" ] def one_served_genai_span(trace: JaegerTrace, genai_span: str) -> JaegerSpan: served = served_genai_spans(trace, genai_span) assert len(served) == 1, ( - f"a streamed call must produce exactly ONE served gen-AI span, got {len(served)}; " - f"spans: {trace.span_names()}" + f"a streamed call must produce exactly ONE served gen-AI span, got {len(served)}; spans: {trace.span_names()}" ) return served[0] @@ -190,8 +182,7 @@ def _assert_real_ttft(hits: list[JaegerTrace], *, genai_span: str) -> None: "(nothing tagged with its call id was found)" ) assert len(hits) == 1, ( - f"expected exactly ONE trace for the call, got {len(hits)}: " - f"{[(t.trace_id, t.span_names()) for t in hits]}" + f"expected exactly ONE trace for the call, got {len(hits)}: {[(t.trace_id, t.span_names()) for t in hits]}" ) trace = hits[0] span = one_served_genai_span(trace, genai_span) @@ -280,9 +271,7 @@ def _assert_error_span_contract(span: JaegerSpan) -> None: "the span status description must carry the same untruncated message as error.message" ) stack = _tag(span, "litellm.provider.error.stack_trace") - assert isinstance(stack, str) and stack, ( - "the error span must carry a non-empty litellm.provider.error.stack_trace" - ) + assert isinstance(stack, str) and stack, "the error span must carry a non-empty litellm.provider.error.stack_trace" class TestOtelTraceCompleteness: @@ -313,9 +302,7 @@ class TestOtelTraceCompleteness: resources.defer(lambda: client.delete_key(key)) marker = unique_marker() - outcome = first_ok( - client, lambda: client.chat_raw(key, MODEL, f"reply with one word {marker}", max_tokens=16) - ) + outcome = first_ok(client, lambda: client.chat_raw(key, MODEL, f"reply with one word {marker}", max_tokens=16)) assert outcome.call_id is not None, "success response must carry x-litellm-call-id" hits = otel_reader.poll_traces_for_call( @@ -520,9 +507,7 @@ class TestOtelTraceCompleteness: route = "/v1/responses" _assert_otel_destination_configured(client) - key = client.key_with_alias( - f"otel-stream-responses-{unique_marker()}", models=[CHEAP_OPENAI_MODEL] - ) + key = client.key_with_alias(f"otel-stream-responses-{unique_marker()}", models=[CHEAP_OPENAI_MODEL]) resources.defer(lambda: client.delete_key(key)) marker = unique_marker() @@ -660,9 +645,7 @@ class TestOtelTraceCompleteness: route = "/v1/responses" _assert_otel_destination_configured(client) - key = client.key_with_alias( - f"otel-ttft-responses-{unique_marker()}", models=[CHEAP_OPENAI_MODEL] - ) + key = client.key_with_alias(f"otel-ttft-responses-{unique_marker()}", models=[CHEAP_OPENAI_MODEL]) resources.defer(lambda: client.delete_key(key)) marker = unique_marker() diff --git a/tests/e2e/logging/test_s3_log_e2e.py b/tests/e2e/logging/test_s3_log_e2e.py new file mode 100644 index 00000000000..7a1ee1e6536 --- /dev/null +++ b/tests/e2e/logging/test_s3_log_e2e.py @@ -0,0 +1,170 @@ +"""Live e2e: s3_v2 log delivery for successful and failed calls. + +Covers logging.s3.success.writes_object and logging.s3.failure.writes_object: +one /chat/completions call must land in the real S3 bucket as exactly one +StandardLoggingPayload object (the primary audit trail; the batch flush must +neither drop nor duplicate it), and a failed call must be persisted the same +way for compliance. Delivery is judged on what is actually in the bucket: the +proxy writes with its production credentials and the test lists and reads the +objects back. + +Both halves of the contract are asserted: the recorded state (the proxy +reports the S3Logger callback active via /health/readiness/details) and the +enforced behavior (the object in the bucket, with the cost cross-checked +against the x-litellm-response-cost header of the very response the caller +received). + +The suite requires ``s3_callback_params.s3_use_key_prefix: true`` on the proxy, +which keys objects as ``{key_alias}/{date}/time-..._{id}.json`` - a unique key +alias per test turns the poll into a cheap prefix listing. +""" + +from __future__ import annotations + +import math +import time + +import pytest + +from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker +from lifecycle import ResourceManager +from logging_client import ( + INVALID_UPSTREAM_API_KEY, + LoggingClient, + completion_response_id, + first_ok, + readiness_details_body, +) +from models import LiteLLMParamsBody +from s3_reader import S3LogReader, build_s3_reader + +pytestmark = pytest.mark.e2e + +#: The active s3_v2 callback's name in /health/readiness/details success_callbacks. +S3_LOGGER_NAME = "S3Logger" + + +@pytest.fixture(scope="session") +def s3_logs() -> S3LogReader: + return build_s3_reader() + + +def _assert_s3_configured(client: LoggingClient) -> None: + """Recorded state: the proxy reports the s3_v2 callback among its active + callbacks, so a missing destination config fails here, before any + delivery-based assertion can time out confusingly.""" + body = readiness_details_body(client) + assert S3_LOGGER_NAME in body, ( + f"the proxy must report the {S3_LOGGER_NAME} callback active " + f"(litellm_settings.callbacks: ['s3_v2'] + s3_callback_params in the proxy config); " + f"got: {body[:400]}" + ) + + +class TestS3LogDelivery: + @pytest.mark.covers("logging.s3.success.writes_object", exercised_on=["chat_completions"]) + def test_chat_completions_writes_one_success_object( + self, client: LoggingClient, s3_logs: S3LogReader, resources: ResourceManager + ) -> None: + """One successful non-streaming /chat/completions call must land in + the bucket as exactly one payload object carrying the model group, the + token counts, and the same cost the caller's response header reported.""" + _assert_s3_configured(client) + + alias = f"s3-chat-{unique_marker()}" + key = client.key_with_alias(alias, models=[CHEAP_ANTHROPIC_MODEL]) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + outcome = first_ok( + client, + lambda: client.chat_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=16), + ) + assert outcome.response_cost is not None and outcome.response_cost > 0, ( + f"the response must report x-litellm-response-cost, got {outcome.response_cost!r}" + ) + body_id = completion_response_id(outcome.body) + assert body_id is not None, "the completion body must carry an id (it names the s3 object)" + + records = s3_logs.poll_records(prefix=f"{alias}/", predicate=lambda r: r.id == body_id) + assert records, ( + f"no s3 object for response {body_id} under prefix {alias}/ reached the bucket within the deadline" + ) + assert len(records) == 1, ( + f"expected exactly ONE s3 object for the call, got {len(records)} - " + "more than one object for one call is the duplicate-delivery bug" + ) + record = records[0] + assert record.status == "success", f"payload status must be success, got {record.status!r}" + assert record.model_group == CHEAP_ANTHROPIC_MODEL, ( + f"payload model_group must be {CHEAP_ANTHROPIC_MODEL!r}, got {record.model_group!r}" + ) + assert record.total_tokens is not None and record.total_tokens > 0, ( + f"payload must count real tokens, got {record.total_tokens!r}" + ) + assert record.response_cost is not None and math.isclose( + record.response_cost, outcome.response_cost, rel_tol=1e-9 + ), f"payload response_cost {record.response_cost!r} must equal the header cost {outcome.response_cost}" + + @pytest.mark.covers("logging.s3.failure.writes_object", exercised_on=["chat_completions"]) + def test_chat_completions_failure_writes_one_object( + self, client: LoggingClient, s3_logs: S3LogReader, resources: ResourceManager + ) -> None: + """A call that fails at the provider must be persisted to the bucket as + exactly one failure payload carrying the provider error - failed calls + are part of the audit trail, not an exemption from it. + + A deployment with an invalid upstream key lets the request pass proxy + auth and fail at the provider (the same lever as the OTEL error test). + Proxy-side rejections during key/model propagation can also ship + failure payloads under this alias, but without a model_group and + without the provider error, so the read-back keys on both: only + provider-reaching calls carry them, and with this key every one of + those is the AnthropicException that ends the send loop.""" + _assert_s3_configured(client) + + model_name = f"s3-err-{unique_marker()}" + model_id = client.create_model( + model_name, + LiteLLMParamsBody(model="anthropic/claude-haiku-4-5", api_key=INVALID_UPSTREAM_API_KEY), + ) + resources.defer(lambda: client.delete_model(model_id)) + alias = f"s3-err-key-{unique_marker()}" + key = client.key_with_alias(alias, models=[model_name]) + resources.defer(lambda: client.delete_key(key)) + + deadline = time.monotonic() + client.proxy.poll_timeout + while True: + outcome = client.chat_raw(key, model_name, "trigger an upstream auth failure", max_tokens=16) + assert not outcome.ok, "the call must fail; the deployment's upstream key is invalid" + assert outcome.status_code != -1, ( + "network failure between the test and the proxy while provoking the provider " + "failure; retrying now could double-log the failure payload and falsely trip " + f"the exactly-one assertion - fix the rig connectivity first: {outcome.body[:200]}" + ) + if "AnthropicException" in outcome.body or time.monotonic() >= deadline: + break + time.sleep(client.proxy.poll_interval) + assert "AnthropicException" in outcome.body, ( + "never saw the upstream provider failure before the deadline; the key may still be " + f"propagating - last outcome {outcome.status_code}: {outcome.body[:200]}" + ) + assert outcome.status_code == 401, ( + f"an upstream auth failure must map to 401, got {outcome.status_code}: {outcome.body[:200]}" + ) + + records = s3_logs.poll_records( + prefix=f"{alias}/", + predicate=lambda r: ( + r.status == "failure" and r.model_group == model_name and "AnthropicException" in (r.error_str or "") + ), + ) + assert records, ( + f"no failure object for {model_name} under prefix {alias}/ reached the bucket within the deadline" + ) + assert len(records) == 1, f"expected exactly ONE failure object for the call, got {len(records)}" + record = records[0] + assert record.error_str is not None and "AnthropicException" in record.error_str, ( + f"the persisted failure must carry the provider error, got error_str={record.error_str!r}" + ) + assert not record.response_cost, f"a failed call must not be billed, got response_cost={record.response_cost!r}" diff --git a/tests/e2e/logging/test_team_langfuse_callback_e2e.py b/tests/e2e/logging/test_team_langfuse_callback_e2e.py new file mode 100644 index 00000000000..89cd45c9f16 --- /dev/null +++ b/tests/e2e/logging/test_team_langfuse_callback_e2e.py @@ -0,0 +1,123 @@ +"""Live e2e: team-scoped Langfuse callback delivery and isolation. + +Covers logging.langfuse.success.logs_spend: a team configured with a Langfuse +callback via POST /team/{id}/callback must deliver its members' calls to the +real Langfuse project (generation readable back through Langfuse's own API, +with the cost agreeing with the x-litellm-response-cost header), while traffic +from keys outside the team must NOT reach that project - the isolation is the +point of team-scoped callbacks. + +Both halves of the contract are asserted: the recorded state (the /team/callback +registration itself answers success) and the enforced behavior (the generation +at the destination for the team key, and its absence for the non-team key). +""" + +from __future__ import annotations + +import time + +import pytest + +from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker +from lifecycle import ResourceManager +from logging_client import ( + LangfuseCreds, + LoggingClient, + costs_agree, + first_ok, + load_langfuse_creds, + observation_spend, +) + +pytestmark = pytest.mark.e2e + +#: How long to keep re-checking that the non-team call never surfaces in +#: Langfuse after the team call's generation has already been ingested; the +#: positive observation bounds the pipeline's latency, so a wrong delivery +#: would be visible within the same order of magnitude. +ISOLATION_SETTLE_SECONDS = 30.0 +ISOLATION_CHECK_INTERVAL_SECONDS = 5.0 + + +@pytest.fixture(scope="session") +def langfuse_creds() -> LangfuseCreds: + return load_langfuse_creds() + + +class TestTeamLangfuseCallback: + @pytest.mark.covers("logging.langfuse.success.logs_spend", exercised_on=["chat_completions"]) + def test_team_callback_delivers_and_isolates( + self, client: LoggingClient, langfuse_creds: LangfuseCreds, resources: ResourceManager + ) -> None: + team_id = client.create_team(f"lf-team-{unique_marker()}", models=[CHEAP_ANTHROPIC_MODEL]) + resources.defer(lambda: client.delete_team(team_id)) + # Recorded state: the registration endpoint itself must answer success + # (add_team_langfuse_callback asserts it). + client.add_team_langfuse_callback(team_id, langfuse_creds) + + team_alias = f"lf-team-key-{unique_marker()}" + team_key = client.key_with_alias(team_alias, models=[CHEAP_ANTHROPIC_MODEL], team_id=team_id) + resources.defer(lambda: client.delete_key(team_key)) + solo_alias = f"lf-solo-key-{unique_marker()}" + solo_key = client.key_with_alias(solo_alias, models=[CHEAP_ANTHROPIC_MODEL]) + resources.defer(lambda: client.delete_key(solo_key)) + + # Enforced behavior, positive half, with one propagation retry: a + # worker still holding the pre-callback team object can serve the + # first call without shipping it, and by the time the first Langfuse + # poll has timed out the team cache TTL has lapsed, so a second call + # must deliver. + team_marker = "" + team_outcome = None + observation = None + for _attempt in range(2): + team_marker = unique_marker() + team_outcome = first_ok( + client, + lambda marker=team_marker: client.chat_raw( + team_key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=16 + ), + ) + assert team_outcome.response_cost is not None and team_outcome.response_cost > 0, ( + f"the response must report x-litellm-response-cost, got {team_outcome.response_cost!r}" + ) + observation = client.poll_langfuse_observation( + langfuse_creds, + key_alias=team_alias, + prompt_marker=team_marker, + require_positive_cost=True, + ) + if observation is not None: + break + solo_marker = unique_marker() + _ = first_ok( + client, + lambda: client.chat_raw( + solo_key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {solo_marker}", max_tokens=16 + ), + ) + + assert observation is not None, ( + f"the team key's call (marker {team_marker}) never reached Langfuse within the deadline, " + "even after a fresh call past the team-object cache TTL" + ) + assert team_outcome is not None and team_outcome.response_cost is not None + cost = observation_spend(observation) + assert cost is not None and costs_agree(team_outcome.response_cost, cost), ( + f"Langfuse calculatedTotalCost {cost!r} must agree with the header cost {team_outcome.response_cost}" + ) + + # Enforced behavior, negative half: the non-team call must never show + # up in this project. The positive generation above has already been + # ingested, which bounds the pipeline latency, so keep re-checking for + # a settle window rather than trusting a single instant. + settle_deadline = time.monotonic() + ISOLATION_SETTLE_SECONDS + while True: + leaked = client.find_langfuse_observation(langfuse_creds, key_alias=solo_alias, prompt_marker=solo_marker) + assert leaked is None, ( + f"a non-team key's call (marker {solo_marker}) reached the team's Langfuse " + f"project: {leaked.id} - team callbacks must not apply outside the team" + ) + if time.monotonic() >= settle_deadline: + break + time.sleep(ISOLATION_CHECK_INTERVAL_SECONDS) diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 2d845a445b5..e7d7fdaef81 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -2635,6 +2635,93 @@ async def test_list_batches_unparseable_row_does_not_truncate_pagination(): assert len(seen) == len(set(seen)) +@pytest.mark.asyncio +async def test_list_batches_fills_a_page_past_a_full_page_of_unparseable_rows(): + """A page whose rows all fail to parse must still let the caller advance. + + ``has_more`` came from the raw fetch while ``last_id`` came from the parsed + survivors, so a full page of corrupt rows answered ``data: []``, + ``last_id: None``, ``has_more: True``, and a client following ``last_id`` + could not move past them. + """ + from litellm.proxy._types import UserAPIKeyAuth + + rows = [_managed_batch_row(i) for i in range(5)] + for corrupt_row in rows[2:4]: + corrupt_row.file_object = "{ not valid json" + prisma_client = _fake_managed_object_table(rows) + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + pages = await _walk_batch_pages( + proxy_managed_files, UserAPIKeyAuth(user_id="test-user"), limit=1 + ) + + assert [[batch.id for batch in page["data"]] for page in pages] == [ + [rows[4].unified_object_id], + [rows[1].unified_object_id], + [rows[0].unified_object_id], + ] + assert [page["has_more"] for page in pages] == [True, True, False] + + +_DEEP_BATCH_SCAN_ROW_COUNT = 2000 +_DEEP_BATCH_SCAN_QUERY_BUDGET = 10 + + +@pytest.mark.asyncio +async def test_list_batches_bounds_the_queries_a_deep_unparseable_run_costs(): + """A tiny limit behind thousands of corrupt rows must not turn one request into thousands of queries.""" + from litellm.proxy._types import UserAPIKeyAuth + + rows = [_managed_batch_row(0)] + [ + _managed_batch_row(index, file_object="{ not valid json") + for index in range(1, _DEEP_BATCH_SCAN_ROW_COUNT + 1) + ] + prisma_client = _fake_managed_object_table(rows) + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + page = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=1 + ) + + assert [batch.id for batch in page["data"]] == [rows[0].unified_object_id] + assert page["has_more"] is False + assert ( + prisma_client.db.litellm_managedobjecttable.find_many.call_count + <= _DEEP_BATCH_SCAN_QUERY_BUDGET + ) + + +@pytest.mark.asyncio +async def test_list_batches_reads_one_chunk_when_the_first_one_fills_the_page(): + """The widened chunk must stay off the common path, where the newest rows already fill the page.""" + from litellm.proxy._types import UserAPIKeyAuth + + rows = [_managed_batch_row(index) for index in range(_DEEP_BATCH_SCAN_ROW_COUNT)] + prisma_client = _fake_managed_object_table(rows) + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + page = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), limit=2 + ) + + assert [batch.id for batch in page["data"]] == [ + rows[-1].unified_object_id, + rows[-2].unified_object_id, + ] + assert page["has_more"] is True + assert prisma_client.db.litellm_managedobjecttable.find_many.call_count == 1 + + @pytest.mark.asyncio async def test_return_unified_file_id_includes_expires_at(): from litellm.types.llms.openai import OpenAIFileObject diff --git a/tests/llm_responses_api_testing/test_openai_responses_api.py b/tests/llm_responses_api_testing/test_openai_responses_api.py index 5f77d5a5477..05bb9113835 100644 --- a/tests/llm_responses_api_testing/test_openai_responses_api.py +++ b/tests/llm_responses_api_testing/test_openai_responses_api.py @@ -1816,7 +1816,7 @@ async def test_extra_body_merges_with_request_data(extra_body_mock_response_data await litellm.aresponses( model="gpt-5.5", input="Test", - temperature=0.7, + temperature=1, max_output_tokens=20, extra_body={ "custom_field": "custom_value", diff --git a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py index be565972b94..1a7fb1f3e41 100644 --- a/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py +++ b/tests/openai_endpoints_tests/test_e2e_openai_responses_api.py @@ -115,9 +115,9 @@ def test_bad_request_error(): def test_bad_request_bad_param_error(): client = get_test_client() with pytest.raises(BadRequestError): - # Trigger error with invalid model name + # Out-of-range temperature on a non-reasoning model, so drop_params forwards it client.responses.create( - model="gpt-5.5", input="This should fail", temperature=2000 + model="gpt-4.1", input="This should fail", temperature=2000 ) diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index a1864c5e480..b4fe9347581 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -1113,8 +1113,12 @@ class TestCheckBatchCost: @pytest.mark.asyncio @pytest.mark.parametrize( "request_counts", - [MagicMock(completed=7, failed=0, total=7), None], - ids=["lagging_output_id", "unknown_counts"], + [ + MagicMock(completed=7, failed=0, total=7), + None, + MagicMock(completed=0, failed=0, total=0), + ], + ids=["lagging_output_id", "unknown_counts", "synthesized_zero_counts"], ) async def test_completed_with_lagging_output_file_left_for_next_cycle( self, diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 0f307450b50..51671d5101e 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -1856,3 +1856,32 @@ async def test_download_percent_encodes_reserved_characters_in_object_key(s3_obj body=None, headers=call.kwargs["headers"], ) + + +def _s3_logger_for_region(region_name: str) -> S3Logger: + logger = S3Logger.__new__(S3Logger) + logger.s3_endpoint_url = None + logger.s3_bucket_name = "my-litellm-audit" + logger.s3_region_name = region_name + return logger + + +@pytest.mark.parametrize( + "region_name,expected_url", + [ + ( + "cn-northwest-1", + "https://my-litellm-audit.s3.cn-northwest-1.amazonaws.com.cn/2025-01-01/key.json", + ), + ( + "us-gov-west-1", + "https://my-litellm-audit.s3.us-gov-west-1.amazonaws.com/2025-01-01/key.json", + ), + ( + "us-east-1", + "https://my-litellm-audit.s3.us-east-1.amazonaws.com/2025-01-01/key.json", + ), + ], +) +def test_build_object_url_uses_partition_dns_suffix(region_name: str, expected_url: str) -> None: + assert _s3_logger_for_region(region_name)._build_object_url("2025-01-01/key.json") == expected_url diff --git a/tests/test_litellm/litellm_core_utils/test_aws_partition.py b/tests/test_litellm/litellm_core_utils/test_aws_partition.py new file mode 100644 index 00000000000..3594d3c354c --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_aws_partition.py @@ -0,0 +1,202 @@ +import ast +from pathlib import Path +from typing import Final +from urllib.parse import urlparse + +import pytest + +import litellm +from litellm.integrations.s3_v2 import S3Logger +from litellm.litellm_core_utils.aws_partition import ( + AwsPartition, + contains_aws_arn, + contains_bedrock_arn, + get_aws_arn_prefix, + get_aws_dns_suffix, + get_aws_partition, + is_bedrock_arn, +) +from litellm.llms.aws_polly.text_to_speech.transformation import AWSPollyTextToSpeechConfig +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig +from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig +from litellm.llms.bedrock.common_utils import init_bedrock_client +from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig + + +@pytest.mark.parametrize( + "region,partition,dns_suffix", + [ + ("us-east-1", "aws", "amazonaws.com"), + ("eu-central-1", "aws", "amazonaws.com"), + ("ap-southeast-1", "aws", "amazonaws.com"), + ("sa-east-1", "aws", "amazonaws.com"), + ("cn-north-1", "aws-cn", "amazonaws.com.cn"), + ("cn-northwest-1", "aws-cn", "amazonaws.com.cn"), + ("us-gov-west-1", "aws-us-gov", "amazonaws.com"), + ("us-gov-east-1", "aws-us-gov", "amazonaws.com"), + ("us-iso-east-1", "aws-iso", "c2s.ic.gov"), + ("us-isob-east-1", "aws-iso-b", "sc2s.sgov.gov"), + ("us-isof-south-1", "aws-iso-f", "csp.hci.ic.gov"), + ("eu-isoe-west-1", "aws-iso-e", "cloud.adc-e.uk"), + (None, "aws", "amazonaws.com"), + ("", "aws", "amazonaws.com"), + ], +) +def test_partition_lookup(region: str | None, partition: str, dns_suffix: str) -> None: + assert get_aws_partition(region) == AwsPartition(partition=partition, dns_suffix=dns_suffix) + assert get_aws_dns_suffix(region) == dns_suffix + assert get_aws_arn_prefix(region) == f"arn:{partition}:" + + +@pytest.mark.parametrize( + "value,expected", + [ + ("arn:aws:bedrock:us-west-2:123456789012:foundation-model/anthropic.claude-3", True), + ("arn:aws-cn:bedrock:cn-north-1:123456789012:inference-profile/p", True), + ("arn:aws-us-gov:bedrock:us-gov-west-1:123456789012:foundation-model/m", True), + ("bedrock/arn:aws-cn:bedrock:cn-north-1:123456789012:application-inference-profile/p", True), + ("arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/r", True), + ("anthropic.claude-3", False), + ("arn:aws:iam::123456789012:role/foo", False), + ], +) +def test_contains_bedrock_arn(value: str, expected: bool) -> None: + assert contains_bedrock_arn(value) is expected + + +@pytest.mark.parametrize( + "value,expected", + [ + ("arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/j", True), + ("arn:aws-cn:bedrock:cn-north-1:123456789012:model-invocation-job/j", True), + ("arn:aws-us-gov:bedrock:us-gov-west-1:123456789012:model-invocation-job/j", True), + ("abc1234567", False), + ("bedrock/arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/j", False), + ("arn:aws:iam::123456789012:role/foo", False), + ], +) +def test_is_bedrock_arn(value: str, expected: bool) -> None: + assert is_bedrock_arn(value) is expected + + +@pytest.mark.parametrize( + "value,expected", + [ + ("model/arn:aws:bedrock:us-east-1:123456789012:foundation-model/m/converse", True), + ("model/arn:aws-cn:bedrock:cn-north-1:123456789012:foundation-model/m/converse", True), + ("arn:aws-us-gov:bedrock:us-gov-west-1:123456789012:inference-profile/p", True), + ("model/anthropic.claude-3/converse", False), + ("arnaws:bedrock", False), + ], +) +def test_contains_aws_arn(value: str, expected: bool) -> None: + assert contains_aws_arn(value) is expected + + +def _agentcore_model(region: str) -> str: + return f"agentcore/{get_aws_arn_prefix(region)}bedrock-agentcore:{region}:111122223333:runtime/my-agent" + + +def _s3_object_url(region: str) -> str: + logger = S3Logger.__new__(S3Logger) + logger.s3_endpoint_url = None + logger.s3_bucket_name = "audit-bucket" + logger.s3_region_name = region + return logger._build_object_url("2025-01-01/key.json") + + +ENDPOINT_BUILDERS: Final = { + "bedrock_runtime_default": lambda region: BaseAWSLLM()._select_default_endpoint_url("runtime", region), + "bedrock_agent_default": lambda region: BaseAWSLLM()._select_default_endpoint_url("agent", region), + "bedrock_agentcore_default": lambda region: BaseAWSLLM()._select_default_endpoint_url("agentcore", region), + "bedrock_get_runtime_endpoint": lambda region: BaseAWSLLM().get_runtime_endpoint(None, None, region)[0], + "bedrock_legacy_client": lambda region: init_bedrock_client( + region_name=region, + aws_access_key_id="test-key", + aws_secret_access_key="test-secret", + ).meta.endpoint_url, + "bedrock_batches": lambda region: BedrockBatchesConfig().get_complete_batch_url( + api_base=None, + api_key=None, + model="anthropic.claude-3", + optional_params={"aws_region_name": region}, + litellm_params={}, + data={"input_file_id": "s3://bucket/key.jsonl"}, + ), + "bedrock_agentcore_invoke": lambda region: AmazonAgentCoreConfig().get_complete_url( + api_base=None, + api_key=None, + model=_agentcore_model(region), + optional_params={}, + litellm_params={}, + ), + "polly": lambda region: AWSPollyTextToSpeechConfig().get_complete_url( + model="polly/neural", + api_base=None, + litellm_params={"aws_region_name": region}, + ), + "sagemaker_chat": lambda region: SagemakerChatConfig().get_complete_url( + api_base=None, + api_key=None, + model="my-endpoint", + optional_params={"aws_region_name": region}, + litellm_params={}, + stream=False, + ), + "sagemaker_chat_stream": lambda region: SagemakerChatConfig().get_complete_url( + api_base=None, + api_key=None, + model="my-endpoint", + optional_params={"aws_region_name": region}, + litellm_params={}, + stream=True, + ), + "s3_object_url": _s3_object_url, +} + + +@pytest.fixture(autouse=True) +def _clear_aws_env(monkeypatch: pytest.MonkeyPatch) -> None: + for env_var in ("AWS_BEDROCK_RUNTIME_ENDPOINT", "AWS_REGION", "AWS_DEFAULT_REGION", "AWS_REGION_NAME"): + monkeypatch.delenv(env_var, raising=False) + + +@pytest.mark.parametrize("region", ["cn-north-1", "cn-northwest-1"]) +@pytest.mark.parametrize("builder_name", sorted(ENDPOINT_BUILDERS)) +def test_every_endpoint_builder_respects_cn_partition(builder_name: str, region: str) -> None: + url = ENDPOINT_BUILDERS[builder_name](region) + hostname = urlparse(url).hostname + assert hostname is not None + assert hostname.endswith(".amazonaws.com.cn"), url + assert not hostname.endswith("amazonaws.com"), url + assert "arn:aws:" not in url, url + + +@pytest.mark.parametrize("region", ["us-east-1", "us-gov-west-1"]) +@pytest.mark.parametrize("builder_name", sorted(ENDPOINT_BUILDERS)) +def test_every_endpoint_builder_keeps_amazonaws_com_outside_cn(builder_name: str, region: str) -> None: + url = ENDPOINT_BUILDERS[builder_name](region) + hostname = urlparse(url).hostname + assert hostname is not None + assert hostname.endswith(".amazonaws.com"), url + + +def _fstring_literal_offenders(needle: str) -> list[str]: + litellm_root = Path(litellm.__file__).parent + return [ + f"{path.relative_to(litellm_root)}: {part.value!r}" + for path in sorted(litellm_root.rglob("*.py")) + for node in ast.walk(ast.parse(path.read_text(encoding="utf-8"))) + if isinstance(node, ast.JoinedStr) + for part in node.values + if isinstance(part, ast.Constant) and isinstance(part.value, str) and needle in part.value + ] + + +def test_no_fstring_hardcodes_the_commercial_dns_suffix() -> None: + assert _fstring_literal_offenders("amazonaws.com") == [] + + +def test_no_fstring_hardcodes_the_commercial_arn_prefix() -> None: + assert _fstring_literal_offenders("arn:aws:") == [] diff --git a/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py b/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py index 3a09702de45..765c47547ce 100644 --- a/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_llm_request_utils.py @@ -1,4 +1,5 @@ import httpx +import pytest from litellm.litellm_core_utils.llm_request_utils import ( flatten_form_field_values, @@ -80,9 +81,7 @@ def test_flatten_form_field_values_later_source_wins_on_collision(): def test_flatten_form_field_values_keeps_scalar_lists_as_repeated_fields(): - assert flatten_form_field_values( - {"loras": ["a", "b", "c"], "generation_config": {"tags": [1, 2]}, "seed": 42} - ) == ( + assert flatten_form_field_values({"loras": ["a", "b", "c"], "generation_config": {"tags": [1, 2]}, "seed": 42}) == ( ("loras", ("a", "b", "c")), ("generation_config[tags]", ("1", "2")), ("seed", "42"), @@ -97,3 +96,12 @@ def test_flatten_form_field_values_scalar_list_survives_update_into_multipart(): assert names.count("loras") == 2 assert names.count("model") == 1 + + +def test_flatten_form_field_values_rejects_over_deep_nesting(): + nested: object = "leaf" + for _ in range(102): + nested = {"k": nested} + assert isinstance(nested, dict) + with pytest.raises(ValueError, match="max depth"): + flatten_form_field_values(nested) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index cb31280c2d5..b8ce11db8d1 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -6,6 +6,9 @@ import pytest from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.anthropic.experimental_pass_through.messages import ( + streaming_iterator as streaming_iterator_module, +) from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( INCOMPLETE_STREAM_ERROR_MESSAGE, AnthropicMessagesStreamHiddenParams, @@ -26,7 +29,7 @@ class _RecordingLoggingIterator(BaseAnthropicMessagesStreamingIterator): self.logged_chunks: list = [] self.logging_call_count: int = 0 - async def _handle_streaming_logging(self, collected_chunks): + async def _handle_streaming_logging(self, collected_chunks, *, stream_teardown=False): self.logged_chunks = list(collected_chunks) self.logging_call_count += 1 @@ -543,3 +546,86 @@ def test_anthropic_messages_response_as_sse_events_no_content_blocks(): response = {"id": "msg_4", "content": [], "stop_reason": "end_turn"} decoded = _decode_sse_events(anthropic_messages_response_as_sse_events(response)) assert [event_type for event_type, _ in decoded] == ["message_start", "message_delta", "message_stop"] + + +class _RecordingLoggingWorker: + def __init__(self): + self.enqueued = [] + + def ensure_initialized_and_enqueue(self, async_coroutine): + self.enqueued.append(async_coroutine) + + def close_enqueued(self): + for coroutine in self.enqueued: + coroutine.close() + + +async def _noop_deferred_dispatch(logging_coroutine): + logging_coroutine.close() + + +async def _stream_of(events): + for event in events: + yield event + + +COMPLETE_STREAM_EVENTS = TRUNCATED_TOOL_USE_EVENTS + ({"type": "message_stop"},) + + +@pytest.mark.asyncio +async def test_normal_end_with_deferred_dispatch_armed_parks_logging_coroutine(monkeypatch): + """ + Regression test for LIT-6409: with post_call guardrails active the proxy + arms logging_obj._on_deferred_stream_complete, and the native /v1/messages + iterator must park its logging coroutine instead of enqueueing it at + upstream exhaustion, otherwise the spend log is built before the + guardrail end-of-stream scan writes its post_call entry. + """ + worker = _RecordingLoggingWorker() + monkeypatch.setattr(streaming_iterator_module, "GLOBAL_LOGGING_WORKER", worker) + iterator = _make_iterator("test_deferred_parks_logging_coroutine") + iterator.litellm_logging_obj._on_deferred_stream_complete = _noop_deferred_dispatch + + await _collect(iterator, _stream_of(COMPLETE_STREAM_EVENTS)) + + parked = getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) + assert worker.enqueued == [] + assert parked is not None + assert len(parked) == 1 + assert asyncio.iscoroutine(parked[0]) + parked[0].close() + + +@pytest.mark.asyncio +async def test_client_disconnect_enqueues_immediately_even_when_deferred_dispatch_armed(monkeypatch): + """ + On client disconnect the guardrail end-of-stream scan never runs, so + deferral would strand the spend log; the teardown path must keep + enqueueing immediately (LIT-5839) even when the deferred callback is armed. + """ + worker = _RecordingLoggingWorker() + monkeypatch.setattr(streaming_iterator_module, "GLOBAL_LOGGING_WORKER", worker) + iterator = _make_iterator("test_disconnect_enqueues_when_armed") + iterator.litellm_logging_obj._on_deferred_stream_complete = _noop_deferred_dispatch + + wrapped = iterator.async_sse_wrapper(_events_then_hang(TRUNCATED_TOOL_USE_EVENTS)) + for _ in range(len(TRUNCATED_TOOL_USE_EVENTS)): + await wrapped.__anext__() + await wrapped.aclose() + + assert len(worker.enqueued) == 1 + assert getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) is None + worker.close_enqueued() + + +@pytest.mark.asyncio +async def test_normal_end_without_deferred_dispatch_enqueues_immediately(monkeypatch): + worker = _RecordingLoggingWorker() + monkeypatch.setattr(streaming_iterator_module, "GLOBAL_LOGGING_WORKER", worker) + iterator = _make_iterator("test_unarmed_enqueues_at_stream_end") + + await _collect(iterator, _stream_of(COMPLETE_STREAM_EVENTS)) + + assert len(worker.enqueued) == 1 + assert getattr(iterator.litellm_logging_obj, "_deferred_stream_complete_args", None) is None + worker.close_enqueued() diff --git a/tests/test_litellm/llms/bedrock/batches/test_handler.py b/tests/test_litellm/llms/bedrock/batches/test_handler.py index d2dc89a7492..2a9b7a6d138 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_handler.py +++ b/tests/test_litellm/llms/bedrock/batches/test_handler.py @@ -150,14 +150,64 @@ def test_handle_model_invocation_job_status_completed(patched_boto3): assert batch.completed_at == int(END_TIME.timestamp()) assert batch.failed_at is None assert batch.cancelled_at is None - # Per-record counts aren't reported by GetModelInvocationJob, so we leave - # them zeroed; consumers should parse manifest.json.out for accurate counts. - assert batch.request_counts.total == 0 + assert batch.request_counts is None assert batch.metadata["job_arn"] == JOB_ARN assert batch.metadata["output_file_uri"] == expected_out assert batch.metadata["output_s3_uri"] == OUTPUT_PREFIX +@pytest.mark.parametrize("success_count,error_count", [(100, 0), (86, 14)]) +def test_completed_job_maps_provider_record_counts(patched_boto3, success_count, error_count): + fake_client, _ = patched_boto3 + fake_client.get_model_invocation_job.return_value = { + **_fake_boto3_response(), + "totalRecordCount": 100, + "successRecordCount": success_count, + "errorRecordCount": error_count, + } + + batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN) + + assert batch.request_counts is not None + assert (batch.request_counts.total, batch.request_counts.completed, batch.request_counts.failed) == ( + 100, + success_count, + error_count, + ) + + +def test_missing_record_counts_leave_request_counts_none(patched_boto3): + fake_client, _ = patched_boto3 + fake_client.get_model_invocation_job.return_value = _fake_boto3_response() + + batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN) + + assert batch.request_counts is None + + +def test_total_without_success_count_leaves_request_counts_none(patched_boto3): + fake_client, _ = patched_boto3 + fake_client.get_model_invocation_job.return_value = {**_fake_boto3_response(), "totalRecordCount": 100} + + batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN) + + assert batch.request_counts is None + + +def test_missing_error_count_maps_to_zero_failed(patched_boto3): + fake_client, _ = patched_boto3 + fake_client.get_model_invocation_job.return_value = { + **_fake_boto3_response(), + "totalRecordCount": 100, + "successRecordCount": 100, + } + + batch = BedrockBatchesHandler._handle_model_invocation_job_status(batch_id=JOB_ARN) + + assert batch.request_counts is not None + assert (batch.request_counts.total, batch.request_counts.completed, batch.request_counts.failed) == (100, 100, 0) + + @pytest.mark.parametrize( "bedrock_status,openai_status", [ diff --git a/tests/test_litellm/llms/bedrock/batches/test_transformation.py b/tests/test_litellm/llms/bedrock/batches/test_transformation.py index 87f9c506857..7e5716a7495 100644 --- a/tests/test_litellm/llms/bedrock/batches/test_transformation.py +++ b/tests/test_litellm/llms/bedrock/batches/test_transformation.py @@ -785,3 +785,37 @@ class TestBedrockBatchesContract(BatchesConfigContractTests): expected_retrieve_batch_id = ARN expected_retrieve_status = "completed" + + +def test_get_complete_batch_url_cn_partition(config: BedrockBatchesConfig) -> None: + url = config.get_complete_batch_url( + api_base=None, + api_key=None, + model="anthropic.claude-3", + optional_params={"aws_region_name": "cn-north-1"}, + litellm_params={}, + data={"input_file_id": "s3://b/k"}, + ) + assert url == "https://bedrock.cn-north-1.amazonaws.com.cn/model-invocation-job" + + +@pytest.mark.parametrize( + "arn,expected_prefix", + [ + ( + "arn:aws-cn:bedrock:cn-north-1:123456789012:model-invocation-job/abc1234567", + "https://bedrock.cn-north-1.amazonaws.com.cn/model-invocation-job/", + ), + ( + "arn:aws-us-gov:bedrock:us-gov-west-1:123456789012:model-invocation-job/abc1234567", + "https://bedrock.us-gov-west-1.amazonaws.com/model-invocation-job/", + ), + ], +) +def test_retrieve_request_accepts_partition_arns(config: BedrockBatchesConfig, arn: str, expected_prefix: str) -> None: + with patch.object(config.common_utils, "sign_aws_request") as mock_sign: + mock_sign.return_value = ({"Authorization": "signed"}, b"") + result = config.transform_retrieve_batch_request( + batch_id=arn, optional_params={}, litellm_params={} + ) + assert result["url"].startswith(expected_prefix) diff --git a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py index 114e473be98..65ae719d021 100644 --- a/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/test_litellm/llms/bedrock/embed/test_bedrock_embedding.py @@ -985,3 +985,51 @@ def test_bedrock_cohere_embedding_types_wrapped_as_list( assert "embedding_types" in request_body assert request_body["embedding_types"] == expected_embedding_types assert isinstance(request_body["embedding_types"], list) + + +def test_load_credentials_assumes_role_with_external_id(monkeypatch): + """A trust policy requiring sts:ExternalId must be satisfied by the deployment's aws_external_id.""" + import datetime + + import boto3 + from botocore.exceptions import ClientError + + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + + monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if params.get("ExternalId") != "external-id-embed": + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIAEMBEDROLEKEY", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_access_key_id": "AKIAEMBEDCALLERKEY", + "aws_secret_access_key": "pod-caller-secret", + "aws_region_name": "us-east-1", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-embed-role", + "aws_session_name": "litellm-embed-session", + "aws_external_id": "external-id-embed", + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + credentials, aws_region_name = BedrockEmbedding()._load_credentials(optional_params) + + assert credentials.access_key == "ASIAEMBEDROLEKEY" + assert credentials.token == "assumed-session-token" + assert aws_region_name == "us-east-1" + assert "aws_external_id" not in optional_params diff --git a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py index 50e2b53c2b3..7d07ac947b1 100644 --- a/tests/test_litellm/llms/bedrock/test_base_aws_llm.py +++ b/tests/test_litellm/llms/bedrock/test_base_aws_llm.py @@ -1223,7 +1223,7 @@ def test_different_roles_without_session_names_should_not_share_cache(): ({}, {"verify": True}), ( {"aws_region_name": "us-east-1"}, - {"verify": True}, + {"verify": True, "region_name": "us-east-1"}, ), ( {"aws_sts_endpoint": "https://sts.eu-west-1.amazonaws.com"}, @@ -1234,7 +1234,7 @@ def test_different_roles_without_session_names_should_not_share_cache(): }, ), ], - ids=["no_region_or_endpoint", "bedrock_region_ignored_for_sts", "explicit_sts_endpoint"], + ids=["no_region_or_endpoint", "configured_region_is_sts_fallback", "explicit_sts_endpoint"], ) def test_eks_irsa_ambient_credentials_used(role_kwargs, expected_client_kwargs): """ @@ -1418,6 +1418,135 @@ def test_build_sts_client_kwargs(env, aws_sts_endpoint, ssl_verify, expected): ) +@pytest.mark.parametrize( + "env,aws_sts_endpoint,aws_region_name,expected_region", + [ + ({}, None, "cn-north-1", "cn-north-1"), + ({"AWS_REGION": "eu-west-1"}, None, "cn-north-1", "eu-west-1"), + ({"AWS_DEFAULT_REGION": "ap-southeast-1"}, None, "cn-north-1", "ap-southeast-1"), + ({}, "https://sts.cn-north-1.amazonaws.com.cn", "us-east-1", "cn-north-1"), + ({}, None, None, None), + ], + ids=[ + "configured_region_fallback", + "env_region_beats_configured", + "env_default_region_beats_configured", + "cn_endpoint_beats_configured", + "nothing_configured", + ], +) +def test_resolve_sts_region_configured_region_fallback( + env: dict[str, str], + aws_sts_endpoint: str | None, + aws_region_name: str | None, + expected_region: str | None, +) -> None: + with patch.dict(os.environ, env, clear=True): + assert ( + BaseAWSLLM._resolve_sts_region( + aws_sts_endpoint=aws_sts_endpoint, + aws_region_name=aws_region_name, + ) + == expected_region + ) + + +def test_build_sts_client_kwargs_configured_region_fallback() -> None: + base_aws_llm = BaseAWSLLM() + with patch.dict(os.environ, {}, clear=True): + assert base_aws_llm._build_sts_client_kwargs(aws_region_name="cn-north-1") == { + "verify": True, + "region_name": "cn-north-1", + } + with patch.dict(os.environ, {"AWS_REGION": "eu-west-1"}, clear=True): + assert base_aws_llm._build_sts_client_kwargs(aws_region_name="cn-north-1") == { + "verify": True, + "region_name": "eu-west-1", + } + + +def test_assume_role_sts_client_uses_configured_cn_region() -> None: + """arn:aws-cn roles must resolve against a cn STS endpoint, not the commercial default.""" + base_aws_llm = BaseAWSLLM() + mock_expiry = MagicMock() + mock_expiry.tzinfo = timezone.utc + time_diff = MagicMock() + time_diff.total_seconds.return_value = 3600 + mock_expiry.__sub__ = MagicMock(return_value=time_diff) + mock_sts_client = MagicMock() + mock_sts_client.assume_role.return_value = { + "Credentials": { + "AccessKeyId": "assumed-access-key", + "SecretAccessKey": "assumed-secret-key", + "SessionToken": "assumed-session-token", + "Expiration": mock_expiry, + } + } + + with patch.dict(os.environ, {}, clear=True): + with patch("boto3.client", return_value=mock_sts_client) as mock_boto3_client: + credentials, ttl = base_aws_llm._auth_with_aws_role( + aws_access_key_id=None, + aws_secret_access_key=None, + aws_session_token=None, + aws_role_name="arn:aws-cn:iam::2222222222222:role/LitellmBedrockRole", + aws_session_name="test-session", + aws_region_name="cn-north-1", + ) + mock_boto3_client.assert_called_with( + "sts", + region_name="cn-north-1", + verify=True, + ) + assert credentials.access_key == "assumed-access-key" + assert credentials.secret_key == "assumed-secret-key" + assert credentials.token == "assumed-session-token" + assert ttl is not None + + +@pytest.mark.parametrize( + "model,expected_region", + [ + ( + "arn:aws-cn:bedrock:cn-north-1:123456789012:application-inference-profile/p", + "cn-north-1", + ), + ( + "arn:aws-us-gov:bedrock:us-gov-west-1:123456789012:foundation-model/m", + "us-gov-west-1", + ), + ( + "bedrock/arn:aws-cn:bedrock:cn-northwest-1:123456789012:inference-profile/p", + "cn-northwest-1", + ), + ("anthropic.claude-3", None), + ], +) +def test_get_aws_region_from_model_arn_partition_arns(model: str, expected_region: str | None) -> None: + assert BaseAWSLLM()._get_aws_region_from_model_arn(model) == expected_region + + +@pytest.mark.parametrize( + "endpoint_type,region,expected", + [ + ("runtime", "cn-north-1", "https://bedrock-runtime.cn-north-1.amazonaws.com.cn"), + ("agent", "cn-north-1", "https://bedrock-agent-runtime.cn-north-1.amazonaws.com.cn"), + ("agentcore", "cn-north-1", "https://bedrock-agentcore.cn-north-1.amazonaws.com.cn"), + ("runtime", "us-east-1", "https://bedrock-runtime.us-east-1.amazonaws.com"), + ("agent", "us-east-1", "https://bedrock-agent-runtime.us-east-1.amazonaws.com"), + ("agentcore", "us-east-1", "https://bedrock-agentcore.us-east-1.amazonaws.com"), + ("runtime", "us-gov-west-1", "https://bedrock-runtime.us-gov-west-1.amazonaws.com"), + ], +) +def test_select_default_endpoint_url_partitions(endpoint_type: str, region: str, expected: str) -> None: + assert ( + BaseAWSLLM()._select_default_endpoint_url( + endpoint_type=endpoint_type, aws_region_name=region + ) + == expected + ) + + def test_irsa_cross_account_sts_client_uses_resolved_region(): """IRSA cross-account path must use _build_sts_client_kwargs (env region, not Bedrock).""" base_aws_llm = BaseAWSLLM() @@ -1612,6 +1741,7 @@ def test_sts_endpoint_region_matches_bedrock_region_param(): "aws_secret_access_key": "explicit-secret-key", "aws_session_token": "assumed-session-token", "verify": True, + "region_name": "us-east-1", }, ), ( @@ -1626,7 +1756,7 @@ def test_sts_endpoint_region_matches_bedrock_region_param(): }, ), ], - ids=["no_region_or_endpoint", "bedrock_region_ignored_for_sts", "explicit_sts_endpoint"], + ids=["no_region_or_endpoint", "configured_region_is_sts_fallback", "explicit_sts_endpoint"], ) def test_explicit_credentials_used_when_provided(role_kwargs, expected_client_kwargs): """ diff --git a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py index 86fdd89acf6..39198bb20f3 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py +++ b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py @@ -245,6 +245,24 @@ class TestOAuthM2M: assert "/serving-endpoints" not in call_url assert call_url == "https://adb-123.azuredatabricks.net/oidc/v1/token" + def test_oauth_m2m_strips_ai_gateway_path(self): + """OAuth M2M derives the token URL from the workspace origin.""" + databricks_base = DatabricksBase() + + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = {"access_token": "token"} + + with patch("requests.post", return_value=mock_response) as mock_post: + databricks_base._get_oauth_m2m_token( + api_base="https://adb-123.azuredatabricks.net/ai-gateway/mlflow/v1", + client_id="id", + client_secret="secret", + ) + + call_url = mock_post.call_args[0][0] + assert call_url == "https://adb-123.azuredatabricks.net/oidc/v1/token" + class TestValidateEnvironmentWithOAuth: """Test OAuth M2M is used when credentials are available.""" diff --git a/tests/test_litellm/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py b/tests/test_litellm/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py new file mode 100644 index 00000000000..eb90430f303 --- /dev/null +++ b/tests/test_litellm/llms/hosted_vllm/videos/test_hosted_vllm_video_transformation.py @@ -0,0 +1,312 @@ +"""Tests for hosted_vllm video generation (vLLM-Omni /v1/videos).""" + +import json +from io import BytesIO + +import httpx +import pytest + +import litellm +from litellm.litellm_core_utils.url_utils import SSRFError +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.hosted_vllm.videos import get_hosted_vllm_video_config +from litellm.llms.hosted_vllm.videos.transformation import ( + HostedVLLMVideoConfig, + _serialize_form_value, +) +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders +from litellm.types.videos.main import VideoObject +from litellm.utils import ProviderConfigManager + + +def test_provider_config_registration(): + config = ProviderConfigManager.get_provider_video_config( + model="hosted_vllm/MiniMax-H3", + provider=LlmProviders.HOSTED_VLLM, + ) + + assert config is not None + assert isinstance(config, HostedVLLMVideoConfig) + assert isinstance(get_hosted_vllm_video_config("MiniMax-H3"), HostedVLLMVideoConfig) + + +def test_get_complete_url_appends_videos(): + config = HostedVLLMVideoConfig() + + assert ( + config.get_complete_url(model="MiniMax-H3", api_base="http://localhost:8091", litellm_params={}) + == "http://localhost:8091/v1/videos" + ) + assert ( + config.get_complete_url(model="MiniMax-H3", api_base="http://localhost:8091/v1", litellm_params={}) + == "http://localhost:8091/v1/videos" + ) + assert ( + config.get_complete_url(model="MiniMax-H3", api_base="http://localhost:8091/v1/", litellm_params={}) + == "http://localhost:8091/v1/videos" + ) + + +def test_get_complete_url_requires_api_base(): + config = HostedVLLMVideoConfig() + + with pytest.raises(ValueError, match="api_base not set"): + config.get_complete_url(model="MiniMax-H3", api_base=None, litellm_params={}) + + +def test_validate_environment_defaults_to_fake_api_key(): + config = HostedVLLMVideoConfig() + + headers = config.validate_environment( + headers={}, + model="MiniMax-H3", + litellm_params=GenericLiteLLMParams(), + ) + + assert headers.get("Authorization") == "Bearer fake-api-key" + + +def test_validate_environment_uses_provided_api_key(): + config = HostedVLLMVideoConfig() + + headers = config.validate_environment( + headers={"X-Test": "1"}, + model="MiniMax-H3", + litellm_params=GenericLiteLLMParams(api_key="my-custom-key"), + ) + + assert headers.get("Authorization") == "Bearer my-custom-key" + assert headers.get("X-Test") == "1" + + +def test_transform_video_create_request_uses_multipart_form_fields(): + """vLLM-Omni rejects JSON create bodies. Extra Omni fields must be form parts.""" + config = HostedVLLMVideoConfig() + extra_params = {"task": "t2va", "duration": 10.0, "audio_flow_shift": 3.0} + + data, files, url = config.transform_video_create_request( + model="MiniMax-H3", + prompt="three cats march into a bedroom playing tiny brass instruments", + api_base="http://localhost:8091/v1/videos", + video_create_optional_request_params={ + "width": 1280, + "height": 720, + "fps": 24, + "num_inference_steps": 20, + "flow_shift": 12, + "seed": 1101, + "aspect_ratio": "16:9", + "extra_params": extra_params, + "extra_headers": {"X-Ignored": "yes"}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert url == "http://localhost:8091/v1/videos" + assert files == () + assert data["model"] == "MiniMax-H3" + assert data["prompt"] == "three cats march into a bedroom playing tiny brass instruments" + assert data["width"] == "1280" + assert data["height"] == "720" + assert data["fps"] == "24" + assert data["num_inference_steps"] == "20" + assert data["flow_shift"] == "12" + assert data["seed"] == "1101" + assert data["aspect_ratio"] == "16:9" + assert json.loads(data["extra_params"]) == extra_params + assert "extra_headers" not in data + + +def test_transform_video_create_request_keeps_openai_size_and_seconds(): + config = HostedVLLMVideoConfig() + + data, files, _ = config.transform_video_create_request( + model="Wan2.2", + prompt="a mountain lake at sunrise", + api_base="http://localhost:8091/v1/videos", + video_create_optional_request_params={"seconds": "8", "size": "1280x720"}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert files == () + assert data["seconds"] == "8" + assert data["size"] == "1280x720" + + +def test_transform_video_create_request_attaches_input_reference_file(): + config = HostedVLLMVideoConfig() + reference = BytesIO(b"fake-png") + reference.name = "input.png" + + data, files, _ = config.transform_video_create_request( + model="Wan2.2", + prompt="animate this image", + api_base="http://localhost:8091/v1/videos", + video_create_optional_request_params={"input_reference": reference, "width": 832}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert data["width"] == "832" + assert "input_reference" not in data + reference_parts = [value for name, value in files if name == "input_reference"] + assert len(reference_parts) == 1 + filename, content, content_type = reference_parts[0] + assert filename == "input_reference.png" + assert content is reference + assert content_type == "image/png" + + +def test_serialize_form_value_does_not_quote_plain_strings(): + assert _serialize_form_value("16:9") == "16:9" + assert _serialize_form_value(True) == "true" + assert _serialize_form_value({"task": "t2va"}) == json.dumps({"task": "t2va"}) + + +def test_map_openai_params_passes_through_omni_fields(): + config = HostedVLLMVideoConfig() + + mapped = config.map_openai_params( + video_create_optional_params={ + "width": 1280, + "extra_params": {"task": "t2va"}, + "aspect_ratio": "16:9", + "extra_body": None, + }, + model="MiniMax-H3", + drop_params=False, + ) + + assert mapped["width"] == 1280 + assert mapped["extra_params"] == {"task": "t2va"} + assert mapped["aspect_ratio"] == "16:9" + assert "extra_body" not in mapped + + +def test_get_supported_openai_params_includes_omni_extensions(): + config = HostedVLLMVideoConfig() + supported = config.get_supported_openai_params("MiniMax-H3") + + assert "prompt" in supported + assert "input_reference" in supported + assert "width" in supported + assert "extra_params" in supported + assert "aspect_ratio" in supported + assert "image_reference" in supported + assert "audio_reference" in supported + + +def _http_handler_for(handler) -> HTTPHandler: + return HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))) + + +def test_video_generation_posts_multipart_not_json(): + captured: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured.append(request) + return httpx.Response( + 200, + json={ + "id": "video-123", + "object": "video", + "status": "queued", + "created_at": 1701234567, + }, + ) + + response = litellm.video_generation( + model="hosted_vllm/MiniMax-H3", + prompt="three cats march into a bedroom playing tiny brass instruments", + api_base="http://localhost:8091", + api_key="test-key", + client=_http_handler_for(handler), + extra_body={ + "width": 1280, + "height": 720, + "fps": 24, + "extra_params": {"task": "t2va", "duration": 10.0}, + }, + ) + + assert isinstance(response, VideoObject) + assert response.status == "queued" + assert len(captured) == 1 + request = captured[0] + assert str(request.url) == "http://localhost:8091/v1/videos" + assert request.headers["authorization"] == "Bearer test-key" + body = request.content + assert b'name="prompt"' in body + assert b"three cats march into a bedroom playing tiny brass instruments" in body + assert b'name="width"' in body + assert b"1280" in body + assert b'name="extra_params"' in body + assert b"t2va" in body + assert request.headers.get("content-type", "").startswith("multipart/form-data") + + +def test_http_image_reference_is_forwarded_not_downloaded(): + config = HostedVLLMVideoConfig() + data, files, _ = config.transform_video_create_request( + model="MiniMax-H3", + prompt="a person singing", + api_base="http://localhost:8091/v1/videos", + video_create_optional_request_params={ + "image_reference": {"image_url": "http://1.1.1.1/face.png"}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert files == () + payload = json.loads(data["image_reference"]) + assert payload["image_url"] == "http://1.1.1.1/face.png" + + +def test_data_url_image_reference_is_forwarded(): + data_url = "data:image/png;base64,AAAA" + config = HostedVLLMVideoConfig() + data, files, _ = config.transform_video_create_request( + model="MiniMax-H3", + prompt="a person singing", + api_base="http://localhost:8091/v1/videos", + video_create_optional_request_params={"image_reference": {"image_url": data_url}}, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert files == () + assert json.loads(data["image_reference"])["image_url"] == data_url + + +def test_metadata_url_in_image_reference_is_rejected(): + config = HostedVLLMVideoConfig() + with pytest.raises(SSRFError, match="blocked address"): + config.transform_video_create_request( + model="MiniMax-H3", + prompt="a person singing", + api_base="http://localhost:8091/v1/videos", + video_create_optional_request_params={ + "image_reference": {"image_url": "http://169.254.169.254/latest/meta-data/"}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + +def test_file_scheme_media_reference_is_rejected(): + config = HostedVLLMVideoConfig() + with pytest.raises(SSRFError, match="scheme"): + config.transform_video_create_request( + model="MiniMax-H3", + prompt="a person singing", + api_base="http://localhost:8091/v1/videos", + video_create_optional_request_params={ + "video_reference": {"video_url": "file:///etc/passwd"}, + }, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py new file mode 100644 index 00000000000..aa1a59d0e5c --- /dev/null +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_handler.py @@ -0,0 +1,48 @@ +import datetime +from unittest.mock import patch + +import boto3 +from botocore.exceptions import ClientError + +from litellm.llms.sagemaker.chat.handler import SagemakerChatHandler + + +def test_load_credentials_assumes_role_with_external_id(monkeypatch): + """A trust policy requiring sts:ExternalId must be satisfied by the deployment's aws_external_id.""" + monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if params.get("ExternalId") != "external-id-sm-chat": + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIASMCHATROLEKEY", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_access_key_id": "AKIASMCHATCALLERKEY", + "aws_secret_access_key": "pod-caller-secret", + "aws_region_name": "us-east-1", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-sm-chat-role", + "aws_session_name": "litellm-sm-chat-session", + "aws_external_id": "external-id-sm-chat", + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + credentials, aws_region_name = SagemakerChatHandler()._load_credentials(optional_params) + + assert credentials.access_key == "ASIASMCHATROLEKEY" + assert credentials.token == "assumed-session-token" + assert aws_region_name == "us-east-1" + assert "aws_external_id" not in optional_params diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py index da6caca4f05..697f5a7ff59 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_chat_transformation.py @@ -317,3 +317,55 @@ def test_body_model_stays_the_endpoint_name_when_hf_model_name_is_unset(monkeypa client = _invoke_sagemaker_chat(monkeypatch) assert client.request_body["model"] == "my-endpoint" + + +@pytest.mark.parametrize( + "region,stream,expected_url", + [ + ( + "cn-north-1", + False, + "https://runtime.sagemaker.cn-north-1.amazonaws.com.cn/endpoints/my-endpoint/invocations", + ), + ( + "cn-north-1", + True, + "https://runtime.sagemaker.cn-north-1.amazonaws.com.cn/endpoints/my-endpoint/invocations-response-stream", + ), + ( + "us-gov-west-1", + False, + "https://runtime.sagemaker.us-gov-west-1.amazonaws.com/endpoints/my-endpoint/invocations", + ), + ( + "us-west-2", + False, + "https://runtime.sagemaker.us-west-2.amazonaws.com/endpoints/my-endpoint/invocations", + ), + ], +) +def test_get_complete_url_uses_partition_dns_suffix(region: str, stream: bool, expected_url: str) -> None: + url = SagemakerChatConfig().get_complete_url( + api_base=None, + api_key=None, + model="my-endpoint", + optional_params={"aws_region_name": region}, + litellm_params={}, + stream=stream, + ) + assert url == expected_url + + +def test_get_complete_url_sagemaker_base_url_override_wins() -> None: + url = SagemakerChatConfig().get_complete_url( + api_base=None, + api_key=None, + model="my-endpoint", + optional_params={ + "aws_region_name": "cn-north-1", + "sagemaker_base_url": "https://my-private-endpoint.example.com/invocations", + }, + litellm_params={}, + stream=False, + ) + assert url == "https://my-private-endpoint.example.com/invocations" diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py index 1cb27b7cf5f..881bac096b1 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_completion_handler.py @@ -172,3 +172,50 @@ async def test_async_native_streaming_forwards_each_frame_incrementally(): assert texts == [f"token{i} " for i in range(len(frames))] assert consumed_at_token == list(range(1, len(frames) + 1)) + + +def test_load_credentials_assumes_role_with_external_id(monkeypatch): + """A trust policy requiring sts:ExternalId must be satisfied by the deployment's aws_external_id.""" + import datetime + + import boto3 + from botocore.exceptions import ClientError + from unittest.mock import patch + + monkeypatch.delenv("AWS_EXTERNAL_ID", raising=False) + + class FakeSTSClient: + def get_caller_identity(self): + return {"Arn": "arn:aws:iam::111111111111:user/litellm-proxy-pod"} + + def assume_role(self, **params): + if params.get("ExternalId") != "external-id-sm-completion": + raise ClientError( + {"Error": {"Code": "AccessDenied", "Message": "is not authorized to perform: sts:AssumeRole"}}, + "AssumeRole", + ) + return { + "Credentials": { + "AccessKeyId": "ASIASMCOMPROLEKEY", + "SecretAccessKey": "assumed-secret", + "SessionToken": "assumed-session-token", + "Expiration": datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(minutes=30), + } + } + + optional_params = { + "aws_access_key_id": "AKIASMCOMPCALLERKEY", + "aws_secret_access_key": "pod-caller-secret", + "aws_region_name": "us-east-1", + "aws_role_name": "arn:aws:iam::999999999999:role/litellm-sm-completion-role", + "aws_session_name": "litellm-sm-completion-session", + "aws_external_id": "external-id-sm-completion", + } + + with patch.object(boto3, "client", return_value=FakeSTSClient()): + credentials, aws_region_name = SagemakerLLM()._load_credentials(optional_params) + + assert credentials.access_key == "ASIASMCOMPROLEKEY" + assert credentials.token == "assumed-session-token" + assert aws_region_name == "us-east-1" + assert "aws_external_id" not in optional_params diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py index 8fa7c15d2d3..00ff06ea082 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_credentials.py @@ -3,6 +3,9 @@ from datetime import datetime, timedelta, timezone import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa +from pydantic import SecretStr from litellm.proxy._experimental.mcp_server.outbound_credentials.bridge_credentials import ( envelope_keys_from_master_key, @@ -13,17 +16,22 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_credent SessionBearerInvalid, SessionRefreshInvalid, SessionRefreshOpened, + SessionSigningConfigError, is_session_bearer_shaped, open_session_refresh_bearer, resolve_session_bearer, + resolve_session_signing_keys, session_keys_from_master_key, ) from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( SESSION_TTL_SECONDS, + AsymmetricSessionKeys, MintedSessionToken, + SessionKeys, SessionPrincipal, mint_session_refresh_token, mint_session_token, + session_public_key_pem, ) NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc) @@ -133,3 +141,86 @@ def test_refresh_grant_rejects_a_different_client(): def test_refresh_grant_rejects_access_token_presented_as_refresh(): result = open_session_refresh_bearer(_access_token(), KEYS, NOW, expected_client_id="llm_client_abc") assert isinstance(result, SessionRefreshInvalid) + + +def _rsa_private_pem() -> str: + key = rsa.generate_private_key(public_exponent=65537, key_size=2048) + return key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ).decode() + + +def test_absent_signing_setting_keeps_the_master_key_hs256_default(): + resolved = resolve_session_signing_keys(MASTER_KEY, None) + assert isinstance(resolved, SessionKeys) + assert resolved.signing_key.get_secret_value() == KEYS.signing_key.get_secret_value() + + +def test_rs256_signing_setting_resolves_inline_pem_material(): + pem = _rsa_private_pem() + resolved = resolve_session_signing_keys( + MASTER_KEY, + {"algorithm": "RS256", "kid": "2026-01", "private_key": pem}, + ) + assert isinstance(resolved, AsymmetricSessionKeys) + assert resolved.kid == "2026-01" + minted = mint_session_token(PRINCIPAL, resolved, NOW) + assert isinstance(minted, MintedSessionToken) + admitted = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", resolved, NOW) + assert isinstance(admitted, SessionBearerAdmitted) + + +def test_rs256_signing_setting_resolves_env_reference(monkeypatch): + monkeypatch.setenv("MCP_SESSION_PRIVATE_KEY", _rsa_private_pem()) + resolved = resolve_session_signing_keys( + MASTER_KEY, + {"algorithm": "RS256", "kid": "2026-01", "private_key": "os.environ/MCP_SESSION_PRIVATE_KEY"}, + ) + assert isinstance(resolved, AsymmetricSessionKeys) + + +def test_rs256_signing_setting_resolves_previous_public_keys(): + old_pem = _rsa_private_pem() + old_keys = AsymmetricSessionKeys(private_key_pem=SecretStr(old_pem), kid="2025-06") + resolved = resolve_session_signing_keys( + MASTER_KEY, + { + "algorithm": "RS256", + "kid": "2026-01", + "private_key": _rsa_private_pem(), + "previous_public_keys": [{"kid": "2025-06", "public_key": session_public_key_pem(old_keys)}], + }, + ) + assert isinstance(resolved, AsymmetricSessionKeys) + minted = mint_session_token(PRINCIPAL, old_keys, NOW) + assert isinstance(minted, MintedSessionToken) + admitted = resolve_session_bearer(f"Bearer {minted.token.get_secret_value()}", resolved, NOW) + assert isinstance(admitted, SessionBearerAdmitted) + + +@pytest.mark.parametrize( + "raw", + [ + {"algorithm": "HS512", "kid": "k", "private_key": "irrelevant"}, + {"algorithm": "RS256", "kid": "k"}, + {"algorithm": "RS256", "kid": "k", "private_key": "not a pem"}, + {"algorithm": "RS256", "kid": "k", "private_key": "os.environ/UNSET_MCP_SESSION_KEY_VAR"}, + {"algorithm": "RS256", "kid": "k", "private_key": "x", "unexpected": True}, + "not-a-mapping", + ], +) +def test_defective_signing_setting_fails_closed_never_falls_back_to_hs256(raw): + resolved = resolve_session_signing_keys(MASTER_KEY, raw) + assert isinstance(resolved, SessionSigningConfigError) + + +def test_signing_config_error_detail_never_leaks_key_material(): + pem = _rsa_private_pem() + resolved = resolve_session_signing_keys( + MASTER_KEY, + {"algorithm": "RS256", "kid": "k", "private_key": pem, "unexpected": True}, + ) + assert isinstance(resolved, SessionSigningConfigError) + assert pem.splitlines()[1] not in resolved.detail diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py index 36280530eac..2a59e6c1baa 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_session_token.py @@ -4,6 +4,8 @@ from datetime import datetime, timedelta, timezone import jwt import pytest +from cryptography.hazmat.primitives import serialization +from cryptography.hazmat.primitives.asymmetric import rsa from pydantic import SecretStr, ValidationError from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token import ( @@ -13,6 +15,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i SESSION_REFRESH_TTL_SECONDS, SESSION_TOKEN_PREFIX, SESSION_TTL_SECONDS, + AsymmetricSessionKeys, MintedSessionToken, NotASessionToken, OpenedSessionToken, @@ -21,6 +24,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i SessionKeys, SessionMalformed, SessionPrincipal, + SessionRotatedPublicKey, SessionTokenTooLarge, is_session_refresh_token, is_session_token, @@ -28,8 +32,22 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.session_token i mint_session_token, open_session_refresh_token, open_session_token, + session_public_key_pem, ) + +def _rsa_private_pem(bits: int = 2048) -> str: + key = rsa.generate_private_key(public_exponent=65537, key_size=bits) + return key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ).decode() + + +_RSA_PEM_A = _rsa_private_pem() +_RSA_PEM_B = _rsa_private_pem() + NOW = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc) KEYS = SessionKeys(signing_key=SecretStr("k" * 32)) OTHER_KEYS = SessionKeys(signing_key=SecretStr("x" * 32)) @@ -264,3 +282,172 @@ def test_signed_claims_with_a_non_string_team_are_rejected(): def test_principal_rejects_an_unknown_audience_at_construction(): with pytest.raises(ValidationError): SessionPrincipal(user_id="user-123", client_id="llm_client_abc", audience="mcp") + + +RSA_KEYS = AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_A), kid="2026-01") +OTHER_RSA_KEYS = AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_B), kid="2025-06") + + +def test_rs256_access_round_trip_with_kid_and_alg_pinned_in_header(): + minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + token = minted.token.get_secret_value() + header = jwt.get_unverified_header(token.removeprefix(SESSION_TOKEN_PREFIX)) + assert header["alg"] == "RS256" + assert header["kid"] == "2026-01" + opened = open_session_token(token, RSA_KEYS, NOW) + assert isinstance(opened, OpenedSessionToken) + assert opened.principal == PRINCIPAL + + +def test_rs256_refresh_round_trip(): + minted = mint_session_refresh_token(PRINCIPAL, RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + token = minted.token.get_secret_value() + opened = open_session_refresh_token(token, RSA_KEYS, NOW) + assert isinstance(opened, OpenedSessionToken) + assert opened.principal == PRINCIPAL + + +def test_rs256_token_verifies_with_public_key_only(): + minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + public_pem = session_public_key_pem(RSA_KEYS) + assert "PUBLIC KEY" in public_pem + assert "PRIVATE" not in public_pem + claims = jwt.decode( + minted.token.get_secret_value().removeprefix(SESSION_TOKEN_PREFIX), + public_pem, + algorithms=["RS256"], + issuer=SESSION_ISSUER, + options={"verify_exp": False}, + ) + assert claims["user_id"] == "user-123" + + +def test_rs256_tampered_signature_is_bad_signature(): + minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + token = minted.token.get_secret_value() + tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb") + assert isinstance(open_session_token(tampered, RSA_KEYS, NOW), SessionBadSignature) + + +def test_rs256_expired_token_is_expired(): + minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) + assert isinstance(open_session_token(minted.token.get_secret_value(), RSA_KEYS, after), SessionExpired) + + +def test_hs256_token_is_rejected_in_rs256_mode(): + assert isinstance(open_session_token(_mint_access(), RSA_KEYS, NOW), SessionBadSignature) + + +def test_hs256_token_claiming_the_current_kid_is_rejected_by_alg_pinning(): + token = SESSION_TOKEN_PREFIX + jwt.encode( + _valid_claims(), + KEYS.signing_key.get_secret_value(), + algorithm="HS256", + headers={"kid": RSA_KEYS.kid}, + ) + assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionMalformed) + + +def test_rs256_token_is_rejected_in_hs256_mode(): + minted = mint_session_token(PRINCIPAL, RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + assert isinstance(open_session_token(minted.token.get_secret_value(), KEYS, NOW), SessionMalformed) + + +def test_rs256_token_from_an_unknown_kid_is_bad_signature(): + minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + assert isinstance(open_session_token(minted.token.get_secret_value(), RSA_KEYS, NOW), SessionBadSignature) + + +def test_rs256_token_signed_by_a_foreign_key_claiming_the_current_kid_is_bad_signature(): + token = SESSION_TOKEN_PREFIX + jwt.encode( + _valid_claims(), + _RSA_PEM_B, + algorithm="RS256", + headers={"kid": RSA_KEYS.kid}, + ) + assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionBadSignature) + + +def test_alg_none_token_with_the_current_kid_is_rejected_in_rs256_mode(): + unsigned = jwt.api_jws.encode( + b'{"iss":"litellm-mcp-gateway"}', key=None, algorithm="none", headers={"kid": RSA_KEYS.kid} + ) + assert isinstance(open_session_token(SESSION_TOKEN_PREFIX + unsigned, RSA_KEYS, NOW), SessionMalformed) + + +def test_rotation_previous_public_key_still_verifies_until_removed(): + minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + token = minted.token.get_secret_value() + rotated = AsymmetricSessionKeys( + private_key_pem=SecretStr(_RSA_PEM_A), + kid="2026-01", + previous_public_keys=( + SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS)), + ), + ) + opened = open_session_token(token, rotated, NOW) + assert isinstance(opened, OpenedSessionToken) + assert opened.principal == PRINCIPAL + assert isinstance(open_session_token(token, RSA_KEYS, NOW), SessionBadSignature) + + +def test_rotation_window_still_enforces_expiry_and_tamper_on_the_previous_key(): + minted = mint_session_token(PRINCIPAL, OTHER_RSA_KEYS, NOW) + assert isinstance(minted, MintedSessionToken) + token = minted.token.get_secret_value() + rotated = AsymmetricSessionKeys( + private_key_pem=SecretStr(_RSA_PEM_A), + kid="2026-01", + previous_public_keys=( + SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS)), + ), + ) + after = NOW + timedelta(seconds=SESSION_TTL_SECONDS + 1) + assert isinstance(open_session_token(token, rotated, after), SessionExpired) + tampered = token[:-2] + ("aa" if not token.endswith("aa") else "bb") + assert isinstance(open_session_token(tampered, rotated, NOW), SessionBadSignature) + + +def test_weak_or_garbage_private_key_pem_rejected_at_construction(): + with pytest.raises(ValidationError): + AsymmetricSessionKeys(private_key_pem=SecretStr(_rsa_private_pem(bits=1024)), kid="weak") + with pytest.raises(ValidationError): + AsymmetricSessionKeys(private_key_pem=SecretStr("not a pem"), kid="junk") + with pytest.raises(ValidationError): + SessionRotatedPublicKey(kid="junk", public_key_pem="not a pem") + with pytest.raises(ValidationError): + SessionRotatedPublicKey(kid="private-half", public_key_pem=_RSA_PEM_A) + + +def test_weak_rotated_public_key_rejected_at_construction(): + weak_public = ( + serialization.load_pem_private_key(_rsa_private_pem(bits=1024).encode(), password=None) + .public_key() + .public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo) + .decode() + ) + with pytest.raises(ValidationError): + SessionRotatedPublicKey(kid="2024-01", public_key_pem=weak_public) + + +def test_duplicate_kids_rejected_at_construction(): + previous = SessionRotatedPublicKey(kid="2025-06", public_key_pem=session_public_key_pem(OTHER_RSA_KEYS)) + with pytest.raises(ValidationError): + AsymmetricSessionKeys(private_key_pem=SecretStr(_RSA_PEM_A), kid="2025-06", previous_public_keys=(previous,)) + with pytest.raises(ValidationError): + AsymmetricSessionKeys( + private_key_pem=SecretStr(_RSA_PEM_A), kid="2026-01", previous_public_keys=(previous, previous) + ) + + +def test_asymmetric_keys_repr_never_leaks_the_private_key(): + assert _RSA_PEM_A not in repr(RSA_KEYS) diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index dcf122745d2..e3f71692c78 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -1,10 +1,10 @@ +import asyncio import json import time from datetime import datetime, timedelta from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch - import httpx import pytest import respx @@ -15,7 +15,6 @@ from prisma.errors import ClientNotConnectedError, HTTPClientClosedError, Prisma import litellm import litellm.proxy.health_endpoints._health_endpoints as _health_endpoints_module from litellm.litellm_core_utils.health_check_helpers import TEST_IMAGE_BASE64 - from litellm.proxy._types import LitellmUserRoles, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.health_endpoints._health_endpoints import ( @@ -145,7 +144,11 @@ async def test_db_health_transport_error_never_raises(transport_error): result = await _db_health_readiness_check() assert result["status"] == "disconnected" - mock_prisma.attempt_db_reconnect.assert_called_once_with(reason="health_readiness_check") + mock_prisma.attempt_db_reconnect.assert_called_once_with( + reason="health_readiness_check", + timeout_seconds=_health_endpoints_module.DB_READINESS_CHECK_TIMEOUT_SECONDS, + lock_timeout_seconds=_health_endpoints_module.DB_READINESS_CHECK_TIMEOUT_SECONDS, + ) @pytest.mark.asyncio @@ -175,7 +178,11 @@ async def test_db_health_transport_error_reconnect_succeeds(transport_error): result = await _db_health_readiness_check() assert result["status"] == "connected" - mock_prisma.attempt_db_reconnect.assert_called_once_with(reason="health_readiness_check") + mock_prisma.attempt_db_reconnect.assert_called_once_with( + reason="health_readiness_check", + timeout_seconds=_health_endpoints_module.DB_READINESS_CHECK_TIMEOUT_SECONDS, + lock_timeout_seconds=_health_endpoints_module.DB_READINESS_CHECK_TIMEOUT_SECONDS, + ) assert mock_prisma.health_check.call_count == 2 @@ -2276,6 +2283,159 @@ async def test_health_readiness_returns_503_when_db_disconnected(): assert result == {"status": "healthy", "db": "disconnected"} +@pytest.mark.asyncio +async def test_health_readiness_returns_200_when_db_down_and_allow_requests_on_db_unavailable(): + """ + Regression test for https://github.com/BerriAI/litellm/issues/34934. + + allow_requests_on_db_unavailable keeps the proxy serving through a DB + outage, so the readiness probe must keep the pod in rotation (200) and + report the DB state through the body, not the status code. Otherwise + K8s pulls every replica before the request-layer fail-open can run. + """ + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import health_readiness + + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock(side_effect=PrismaError("nope")) + mock_prisma.attempt_db_reconnect = AsyncMock(side_effect=Exception("still nope")) + + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + + response = Response() + with ( + patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), + patch.dict( # test-quality-ok: the fail-open flag lives in the proxy-global general_settings; no injection seam + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": True}, + ), + ): + result = await health_readiness(response=response) + + assert response.status_code == 200 + assert result == {"status": "healthy", "db": "disconnected"} + + +@pytest.mark.asyncio +async def test_health_readiness_details_returns_200_when_db_down_and_allow_requests_on_db_unavailable(): + """ + The detailed readiness payload (public via + allow_public_health_readiness_details, or /health/readiness/details) + must honor the same flag so probes pointed at it also stay 200. + """ + from fastapi import Response + + from litellm.proxy.health_endpoints._health_endpoints import ( + _get_health_readiness_details, + ) + + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock(side_effect=PrismaError("nope")) + mock_prisma.attempt_db_reconnect = AsyncMock(side_effect=Exception("still nope")) + + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + + response = Response() + with ( + patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), + patch.dict( # test-quality-ok: the fail-open flag lives in the proxy-global general_settings; no injection seam + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": True}, + ), + ): + result = await _get_health_readiness_details(response=response) + + assert response.status_code == 200 + assert result["db"] == "disconnected" + + +@pytest.mark.asyncio +async def test_db_health_readiness_check_bounds_hung_health_check(): + """ + A connection that hangs mid-failover must not stall the probe past the + kubelet's timeoutSeconds; the DB round-trip is bounded and reported as + disconnected instead. + """ + from litellm.proxy.health_endpoints._health_endpoints import ( + _db_health_readiness_check, + ) + + async def hang(): + await asyncio.sleep(60) + + mock_prisma = MagicMock() + mock_prisma.health_check = hang + mock_prisma.attempt_db_reconnect = AsyncMock(side_effect=Exception("still down")) + + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + + with patch( # test-quality-ok: lowers the module-level probe timeout so the hung-call test finishes fast + "litellm.proxy.health_endpoints._health_endpoints.DB_READINESS_CHECK_TIMEOUT_SECONDS", + 0.05, + ): + start = time.monotonic() + with patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ): + result = await _db_health_readiness_check() + elapsed = time.monotonic() - start + + assert result["status"] == "disconnected" + assert elapsed < 5 + + +@pytest.mark.asyncio +async def test_db_health_readiness_check_overall_deadline_bounds_hung_reconnect(): + """ + The whole probe-path DB check (initial check + reconnect + re-check, + including reconnect lock waits) runs under one deadline, so a reconnect + that hangs on the lock still returns disconnected within the deadline. + """ + from litellm.proxy.health_endpoints._health_endpoints import ( + _db_health_readiness_check, + ) + + async def hang(**kwargs): + await asyncio.sleep(60) + + mock_prisma = MagicMock() + mock_prisma.health_check = AsyncMock(side_effect=httpx.ConnectError("down")) + mock_prisma.attempt_db_reconnect = hang + + _health_endpoints_module.db_health_cache = { + "status": "unknown", + "last_updated": datetime.now() - timedelta(seconds=60), + } + + with patch( # test-quality-ok: lowers the module-level probe timeout so the hung-call test finishes fast + "litellm.proxy.health_endpoints._health_endpoints.DB_READINESS_PROBE_DEADLINE_SECONDS", + 0.05, + ): + start = time.monotonic() + with patch( # test-quality-ok: the readiness path reads the proxy-global DB client; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ): + result = await _db_health_readiness_check() + elapsed = time.monotonic() - start + + assert result["status"] == "disconnected" + assert elapsed < 5 + + @pytest.mark.asyncio async def test_health_readiness_returns_200_when_db_connected(): """Happy path: connected DB keeps the legacy 200.""" @@ -2746,13 +2906,13 @@ def test_test_model_connection_accepts_image_edit_mode(monkeypatch): app = FastAPI() app.include_router(_health_endpoints_module.router) - app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN - ) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) client = TestClient(app) with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: the endpoint reads the proxy-global DB client and 500s when it is None; it has no injection seam + patch( # test-quality-ok: the endpoint reads the proxy-global DB client and 500s when it is None; it has no injection seam + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), respx.mock(assert_all_called=True) as respx_mock, ): respx_mock.post(host="api.openai.com", path="/v1/images/edits").respond( diff --git a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py index e3893a66094..93dc429168f 100644 --- a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py +++ b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_ai_policy_suggester.py @@ -7,6 +7,9 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +import litellm + +from litellm.proxy._types import ProxyException from litellm.proxy.management_endpoints.policy_endpoints.ai_policy_suggester import ( SUGGEST_TOOL, AiPolicySuggester, @@ -234,6 +237,7 @@ class TestAiPolicySuggester: call_kwargs = mock_acompletion.call_args.kwargs assert call_kwargs["model"] == "gpt-4o-mini" assert call_kwargs["temperature"] == 0.2 + assert call_kwargs["drop_params"] is True assert len(call_kwargs["tools"]) == 1 assert call_kwargs["tools"][0]["function"]["name"] == "select_policy_templates" assert ( @@ -242,3 +246,76 @@ class TestAiPolicySuggester: assert len(call_kwargs["messages"]) == 2 assert call_kwargs["messages"][0]["role"] == "system" assert call_kwargs["messages"][1]["role"] == "user" + + +class TestSuggesterRejectsModelsWithoutToolCalling: + @pytest.mark.asyncio + async def test_a_tools_less_model_is_rejected(self, local_model_cost_map): + with pytest.raises(ProxyException) as exc: + await AiPolicySuggester().suggest( + templates=SAMPLE_TEMPLATES, + attack_examples=["Ignore all previous instructions"], + description="Block prompt injection attempts", + model="perplexity/sonar", + ) + + assert int(exc.value.code) == 400 + assert exc.value.param == "model" + assert "tool calling" in exc.value.message + + def test_a_model_without_forced_tool_choice_support_remains_eligible(self, local_model_cost_map): + supported_params = litellm.get_supported_openai_params( + model="amazon.nova-pro-v1:0", + custom_llm_provider="bedrock", + ) + + assert supported_params is not None + assert "tools" in supported_params + assert "tool_choice" not in supported_params + + +class TestSuggesterToleratesAModelThatRefusesItsSamplingParams: + """The model is operator-supplied, so it can be a reasoning model whose only accepted + temperature is 1. This call pins temperature=0.2 for tool-selection determinism, which such + a model rejects outright: without drop_params litellm raises UnsupportedParamsError and the + whole suggestion fails rather than degrading. Every other internal LLM call in the proxy + already opts in through judge_acompletion; this one was the exception. + """ + + @pytest.mark.asyncio + async def test_a_reasoning_model_gets_past_param_mapping(self, monkeypatch, local_model_cost_map): + """Drives the real entry point with no patching and no network. Which exception escapes is + the discriminator: param mapping runs before any credential check, so UnsupportedParamsError + means the call died on the pinned temperature, while AuthenticationError means it survived + that and got as far as needing a key. Asserting the latter is what the caller observes. + """ + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + + with pytest.raises(litellm.AuthenticationError): + await AiPolicySuggester().suggest( + templates=SAMPLE_TEMPLATES, + attack_examples=["My SSN is 123-45-6789"], + description="", + model="gpt-5.6-terra", + ) + + def test_the_pinned_temperature_is_what_such_a_model_refuses(self, local_model_cost_map): + """The other half of the discriminator above: the same temperature this call pins is + exactly what the model rejects, and drop_params is what removes it.""" + from litellm.utils import get_optional_params + + optional_params = get_optional_params( + model="gpt-5.6-terra", + custom_llm_provider="openai", + temperature=0.2, + tools=[SUGGEST_TOOL], + tool_choice={"type": "function", "function": {"name": "select_policy_templates"}}, + drop_params=True, + ) + + assert "temperature" not in optional_params + assert optional_params["tools"] == [SUGGEST_TOOL] + assert optional_params["tool_choice"] == { + "type": "function", + "function": {"name": "select_policy_templates"}, + } 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 d47a88a0a4c..726e09f3162 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 @@ -2385,9 +2385,11 @@ async def test_start_shadow_eval_seeds_a_zero_funnel_row_per_leg(monkeypatch: py separates 'nothing was skipped' from a job predating the funnel.""" import litellm.proxy.proxy_server as proxy_server + _configure_anthropic_sdk_judge(monkeypatch) prisma = _shadow_prisma(legs=[]) monkeypatch.setattr(proxy_server, "prisma_client", prisma) monkeypatch.setattr(proxy_server, "llm_router", _shadow_router()) + _configure_anthropic_sdk_judge(monkeypatch) await start_shadow_eval(_start_request(api_key_ids=("key-hash", "key-hash-2")), ADMIN) diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py index 7f84407f8b3..87cd2aaff1f 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py @@ -430,13 +430,15 @@ def test_add_internal_model_credentials_survives_a_failing_deployment_lookup(): assert data == {"batch_id": "unified-batch-id"} +from openai.types.batch import BatchRequestCounts + from litellm.proxy.openai_files_endpoints.common_utils import ( _completed_batch_safe_to_retire, ) def _completed_batch_for_retire( - output_file_id: str | None, completed: int | None = None + output_file_id: str | None, counts: BatchRequestCounts | None = None ) -> LiteLLMBatch: kwargs = dict( id="batch-1", @@ -449,26 +451,30 @@ def _completed_batch_for_retire( output_file_id=output_file_id, error_file_id=None, ) - if completed is not None: - kwargs["request_counts"] = {"total": completed, "completed": completed, "failed": 0} + if counts is not None: + kwargs["request_counts"] = counts return LiteLLMBatch(**kwargs) class TestCompletedBatchSafeToRetire: """A completed batch is only safe to retire from cost recovery once its output - file has arrived or the provider proves no successful lines (#37713).""" + file has arrived or the provider proves it enumerated a positive total of + request lines and none succeeded (#37713, LIT-6360).""" def test_output_file_present_is_safe(self): assert _completed_batch_safe_to_retire(_completed_batch_for_retire("file-out")) is True - def test_no_output_and_no_successful_lines_is_safe(self): - # Every request line errored -> nothing left to recover. - assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, completed=0)) is True + def test_no_output_and_synthesized_zero_counts_is_not_safe(self): + counts = BatchRequestCounts(total=0, completed=0, failed=0) + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, counts)) is False def test_no_output_but_successful_lines_is_not_safe(self): - # The bug: output_file_id is lagging; retiring here loses the spend record. - assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, completed=5)) is False + counts = BatchRequestCounts(total=100, completed=100, failed=0) + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, counts)) is False + + def test_no_output_and_all_lines_failed_is_safe(self): + counts = BatchRequestCounts(total=100, completed=0, failed=100) + assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None, counts)) is True def test_no_output_and_unknown_counts_is_not_safe(self): - # Counts unknown -> stay eligible so the next poller pass revisits it. assert _completed_batch_safe_to_retire(_completed_batch_for_retire(None)) is False diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 23552f2fa31..87e0319f6a1 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -4476,3 +4476,82 @@ def test_scoped_list_files_still_resolves_deployment_credentials( provider_list.assert_awaited_once() assert provider_list.await_args.kwargs["custom_llm_provider"] == "openai" assert provider_list.await_args.kwargs["api_key"] == "openai_api_key" + + +def _post_user_data_file() -> httpx.Response: + return client.post( + "/v1/files", + files={"file": ("labels.jsonl", b'{"label": "restricted"}', "application/json")}, + data={"purpose": "user_data"}, + headers={"Authorization": "Bearer test-key"}, + ) + + +def _setup_create_file_over_pre_call_hook(monkeypatch, llm_router, hook): + setup_proxy_logging_object(monkeypatch, llm_router) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + monkeypatch.setattr(litellm, "callbacks", [hook]) + monkeypatch.setattr( + "litellm.proxy.openai_files_endpoints.files_endpoints.files_config", + [{"custom_llm_provider": "openai", "api_key": "sk-test"}], + ) + return respx.post("https://api.openai.com/v1/files").mock( + return_value=respx.MockResponse( + status_code=200, + json={ + "id": "file-hooked", + "object": "file", + "bytes": 23, + "created_at": 1234567890, + "filename": "labels.jsonl", + "purpose": "user_data", + "status": "uploaded", + }, + ) + ) + + +@respx.mock +def test_create_file_triggers_async_pre_call_hook(monkeypatch, llm_router: Router): + """`POST /v1/files` must run `async_pre_call_hook` so a hook can inspect the upload + before it reaches the provider (LIT-5916).""" + from litellm.integrations.custom_logger import CustomLogger + + recorded: dict = {} + + class RecordingHook(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + recorded["call_type"] = call_type + recorded["purpose"] = data.get("purpose") + recorded["file"] = data.get("file") + + provider_route = _setup_create_file_over_pre_call_hook(monkeypatch, llm_router, RecordingHook()) + + response = _post_user_data_file() + + assert response.status_code == 200, response.text + assert recorded["call_type"] == "acreate_file" + assert recorded["purpose"] == "user_data" + assert recorded["file"]["filename"] == "labels.jsonl" + assert provider_route.call_count == 1 + forwarded_body = provider_route.calls.last.request.content + assert b"user_data" in forwarded_body + assert b"labels.jsonl" in forwarded_body + + +@respx.mock +def test_create_file_async_pre_call_hook_rejection_blocks_upload(monkeypatch, llm_router: Router): + """A hook rejecting the upload must 400 before the file reaches the provider.""" + from litellm.integrations.custom_logger import CustomLogger + + class RejectingHook(CustomLogger): + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + return "file upload not allowed" + + provider_route = _setup_create_file_over_pre_call_hook(monkeypatch, llm_router, RejectingHook()) + + response = _post_user_data_file() + + assert response.status_code == 400, response.text + assert "file upload not allowed" in response.text + assert provider_route.call_count == 0 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 3b506324ad7..09bab1dc416 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1870,7 +1870,10 @@ class TestBedrockAgentRuntimePassthroughToggle: with ( patch("litellm.proxy.proxy_server.general_settings", general_settings), - patch("litellm.utils.get_secret", return_value="us-east-1"), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_secret_str", + return_value="us-east-1", + ), patch("litellm.llms.bedrock.chat.BedrockConverseLLM", return_value=bedrock_llm), patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_request_copy", @@ -1920,7 +1923,10 @@ class TestBedrockAgentRuntimePassthroughToggle: async def test_model_invoke_still_routed_when_agent_runtime_disabled(self): with ( patch("litellm.proxy.proxy_server.general_settings", self.DISABLED), - patch("litellm.utils.get_secret", return_value="us-east-1"), + patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_secret_str", + return_value="us-east-1", + ), patch( "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_request_copy", Mock(), diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py index 37d2141e460..078bd4dd402 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py @@ -92,3 +92,13 @@ def test_encode_bedrock_runtime_modelid_arn_edge_cases(): expected = "model/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile%2Ftest-profile.v1/invoke" result = CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) assert result == expected + + +def test_encode_bedrock_runtime_modelid_arn_partition_arns() -> None: + endpoint = "model/arn:aws-cn:bedrock:cn-north-1:123456789012:application-inference-profile/r742sbn2zckd/converse" + expected = "model/arn:aws-cn:bedrock:cn-north-1:123456789012:application-inference-profile%2Fr742sbn2zckd/converse" + assert CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) == expected + + endpoint = "model/arn:aws-us-gov:bedrock:us-gov-west-1:123456789012:inference-profile/test-profile/invoke" + expected = "model/arn:aws-us-gov:bedrock:us-gov-west-1:123456789012:inference-profile%2Ftest-profile/invoke" + assert CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint) == expected diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index c51c1c3f73b..9fb31d2a6db 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -1009,11 +1009,10 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "_guardrail_pipelines", "_pipeline_managed_guardrails", } - for metadata_key in ("metadata", "litellm_metadata"): - cleaned_metadata = updated.get(metadata_key) or {} - for stripped_key in stripped_keys: - assert stripped_key not in cleaned_metadata - assert cleaned_metadata.get("safe_user_metadata") == "kept" + assert "litellm_metadata" not in updated + for stripped_key in stripped_keys: + assert stripped_key not in updated["metadata"] + assert updated["metadata"]["safe_user_metadata"] == "kept" requester_metadata = updated["metadata"]["requester_metadata"] for stripped_key in stripped_keys: @@ -1576,10 +1575,7 @@ async def test_add_litellm_data_to_request_allows_redaction_opt_out_with_admin_o header.lower() for header in updated["proxy_server_request"]["body"]["metadata"]["headers"] } - assert "litellm-disable-message-redaction" in { - header.lower() - for header in (updated.get("litellm_metadata") or {}).get("headers", {}) - } + assert "litellm_metadata" not in updated @pytest.mark.asyncio @@ -6658,9 +6654,9 @@ async def test_add_litellm_data_to_request_strips_caller_supplied_callback_crede assert "gcs_bucket_name" not in updated assert updated["dd_api_key"] == "team-dd-key" assert updated[TRUSTED_CALLBACK_VARS_FIELD] == {"dd_api_key": "team-dd-key"} - for metadata_key in ("metadata", "litellm_metadata"): - assert "dd_site" not in updated[metadata_key] - assert "dd_agent_host" not in updated[metadata_key] + assert "litellm_metadata" not in updated + assert "dd_site" not in updated["metadata"] + assert "dd_agent_host" not in updated["metadata"] assert "dd_site" not in updated["litellm_params"]["metadata"] assert updated["metadata"]["safe_user_metadata"] == "kept" @@ -7510,10 +7506,10 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_bo version="test-version", ) - for bucket in ("metadata", "litellm_metadata"): - assert "attempted_fallbacks" not in updated[bucket] - assert "original_model_group" not in updated[bucket] - assert updated[bucket]["client_key"] == "client_value" + assert "litellm_metadata" not in updated + assert "attempted_fallbacks" not in updated["metadata"] + assert "original_model_group" not in updated["metadata"] + assert updated["metadata"]["client_key"] == "client_value" @pytest.mark.asyncio @@ -7535,10 +7531,10 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_from_js version="test-version", ) - assert isinstance(updated["litellm_metadata"], dict) - assert "attempted_fallbacks" not in updated["litellm_metadata"] - assert "original_model_group" not in updated["litellm_metadata"] - assert updated["litellm_metadata"]["client_key"] == "client_value" + assert "litellm_metadata" not in updated + assert "attempted_fallbacks" not in updated["metadata"] + assert "original_model_group" not in updated["metadata"] + assert updated["metadata"]["client_key"] == "client_value" @pytest.mark.asyncio @@ -7562,9 +7558,10 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_despite version="test-version", ) - assert updated["litellm_metadata"]["model_info"] == {"input_cost_per_token": 0.0} - assert "attempted_fallbacks" not in updated["litellm_metadata"] - assert "original_model_group" not in updated["litellm_metadata"] + assert "litellm_metadata" not in updated + assert updated["metadata"]["model_info"] == {"input_cost_per_token": 0.0} + assert "attempted_fallbacks" not in updated["metadata"] + assert "original_model_group" not in updated["metadata"] @pytest.mark.asyncio @@ -7601,7 +7598,8 @@ async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_ litellm_metadata made the router hand downstream a scrubbed copy, so the proxy's post_call write-backs (guardrail telemetry, applied guardrails) landed in a dict the spend row never read. After the boundary strip plus the in-place scrub, the object the - router forwards is the proxy's own request_data bucket.""" + router forwards is the proxy's own request_data bucket; on chat routes that bucket is + ``metadata``, since the boundary folds client ``litellm_metadata`` into it.""" from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request data = { @@ -7617,7 +7615,9 @@ async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_ general_settings={}, version="test-version", ) - proxy_bucket = request_data["litellm_metadata"] + proxy_bucket = request_data["metadata"] + assert "attempted_fallbacks" not in proxy_bucket + assert "original_model_group" not in proxy_bucket router = litellm.Router( model_list=[ { @@ -7630,7 +7630,7 @@ async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_ original_acompletion = router._acompletion async def _spy(*args, **spy_kwargs): - forwarded_buckets.append(spy_kwargs["litellm_metadata"]) + forwarded_buckets.append(spy_kwargs["metadata"]) return await original_acompletion(*args, **spy_kwargs) router._acompletion = _spy @@ -7639,7 +7639,45 @@ async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_ assert forwarded_buckets == [proxy_bucket] assert forwarded_buckets[0] is proxy_bucket - assert "attempted_fallbacks" not in proxy_bucket - assert "original_model_group" not in proxy_bucket + assert proxy_bucket["attempted_fallbacks"] == 0 + assert proxy_bucket.get("original_model_group") != "spoofed-group" proxy_bucket["standard_logging_guardrail_information"] = [{"guardrail_name": "postcall-guard"}] assert forwarded_buckets[0]["standard_logging_guardrail_information"] == [{"guardrail_name": "postcall-guard"}] + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_folds_litellm_metadata_into_metadata_on_chat_routes(): + data = { + "model": "gpt-3.5-turbo", + "metadata": {"tags": ["from-metadata"]}, + "litellm_metadata": {"trace_id": "abc", "tags": ["from-litellm-metadata"]}, + } + + updated = await add_litellm_data_to_request( + data=data, + request=_make_chat_request_mock(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert "litellm_metadata" not in updated + assert updated["metadata"]["trace_id"] == "abc" + assert updated["metadata"]["tags"] == ["from-metadata", "from-litellm-metadata"] + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_keeps_litellm_metadata_on_litellm_metadata_routes(): + data = {"model": "claude-sonnet-5", "litellm_metadata": {"trace_id": "abc"}} + + updated = await add_litellm_data_to_request( + data=data, + request=_make_request_mock("/v1/messages", {"Content-Type": "application/json"}), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated["litellm_metadata"]["trace_id"] == "abc" diff --git a/tests/test_litellm/proxy/test_pricing_field_strip.py b/tests/test_litellm/proxy/test_pricing_field_strip.py index 1707f5bbc05..a84c6ba2b8a 100644 --- a/tests/test_litellm/proxy/test_pricing_field_strip.py +++ b/tests/test_litellm/proxy/test_pricing_field_strip.py @@ -285,8 +285,8 @@ async def test_add_litellm_data_to_request_skips_strip_with_key_opt_in(): async def test_add_litellm_data_to_request_strips_json_string_litellm_metadata(): """``litellm_metadata`` may arrive as a JSON-encoded string (multipart/ form-data or ``extra_body``). The strip has to run after the proxy parses - it into a dict; otherwise the ``isinstance(dict)`` guard skips the field - and ``model_info`` survives the strip via the string path. + it into a dict but before the chat-route fold into ``metadata``; otherwise + ``model_info`` survives via the string path and lands in the folded bucket. """ import json @@ -305,9 +305,8 @@ async def test_add_litellm_data_to_request_strips_json_string_litellm_metadata() version="test-version", ) - parsed_metadata = updated.get("litellm_metadata") - assert isinstance(parsed_metadata, dict) - assert "model_info" not in parsed_metadata + assert "litellm_metadata" not in updated + assert "model_info" not in updated["metadata"] @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py index 719d7cc73f5..41b0eb3cf95 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_reconnect.py @@ -89,9 +89,7 @@ async def test_run_reconnect_cycle_direct_path_recreates_when_probe_fails( prisma_client._cleanup_engine_watcher = MagicMock() writer = MagicMock() - writer.query_raw = AsyncMock( - side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]] - ) + writer.query_raw = AsyncMock(side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]]) monkeypatch.setattr( PrismaClient, "writer_db", @@ -171,9 +169,7 @@ async def test_run_reconnect_cycle_passes_writer_generation_to_recreate( writer = MagicMock() writer._engine_generation = 7 - writer.query_raw = AsyncMock( - side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]] - ) + writer.query_raw = AsyncMock(side_effect=[ConnectionError("probe failed"), [{"?column?": 1}]]) monkeypatch.setattr( PrismaClient, "writer_db", @@ -229,9 +225,7 @@ async def test_attempt_reconnect_inside_lock_runs_cycle_and_resets_counter( prisma_client._consecutive_reconnect_failures = 2 prisma_client._run_reconnect_cycle = AsyncMock() - ok = await prisma_client._attempt_reconnect_inside_lock( - force=True, reason="test", timeout_seconds=1 - ) + ok = await prisma_client._attempt_reconnect_inside_lock(force=True, reason="test", timeout_seconds=1) pinned = { "returned": ok, "cycle_called": prisma_client._run_reconnect_cycle.await_count, @@ -254,9 +248,7 @@ async def test_attempt_reconnect_inside_lock_skips_when_in_cooldown( prisma_client._db_last_reconnect_attempt_ts = time.time() prisma_client._run_reconnect_cycle = AsyncMock() - ok = await prisma_client._attempt_reconnect_inside_lock( - force=False, reason="test", timeout_seconds=1 - ) + ok = await prisma_client._attempt_reconnect_inside_lock(force=False, reason="test", timeout_seconds=1) assert ok is False assert prisma_client._run_reconnect_cycle.await_count == 0 @@ -269,9 +261,7 @@ async def test_attempt_reconnect_inside_lock_increments_failure_counter_on_error prisma_client._consecutive_reconnect_failures = 0 prisma_client._run_reconnect_cycle = AsyncMock(side_effect=RuntimeError("boom")) - ok = await prisma_client._attempt_reconnect_inside_lock( - force=True, reason="failing_test", timeout_seconds=1 - ) + ok = await prisma_client._attempt_reconnect_inside_lock(force=True, reason="failing_test", timeout_seconds=1) assert ok is False assert prisma_client._consecutive_reconnect_failures == 1 @@ -316,9 +306,7 @@ async def test_attempt_db_reconnect_lock_timeout_returns_false( by replacing ``asyncio.wait`` with a callable that returns the loser task as still-pending after it's already been completed elsewhere. """ - completed_task: asyncio.Task[bool] = asyncio.get_running_loop().create_task( - _no_op_returning_true() - ) + completed_task: asyncio.Task[bool] = asyncio.get_running_loop().create_task(_no_op_returning_true()) # Ensure the inner task has finished before attempt_db_reconnect sees it. await completed_task @@ -329,7 +317,7 @@ async def test_attempt_db_reconnect_lock_timeout_returns_false( monkeypatch.setattr( asyncio, "create_task", - lambda coro, *a, **kw: (coro.close() or completed_task), + lambda coro, *a, **kw: coro.close() or completed_task, ) prisma_client._db_last_reconnect_attempt_ts = 0.0 @@ -465,9 +453,7 @@ async def test_db_health_watchdog_loop_triggers_reconnect_on_timeout( await prisma_client._db_health_watchdog_loop() pinned = { "reconnect_called": prisma_client.attempt_db_reconnect.await_count, - "reconnect_reason": prisma_client.attempt_db_reconnect.await_args.kwargs[ - "reason" - ], + "reconnect_reason": prisma_client.attempt_db_reconnect.await_args.kwargs["reason"], "wait_for_calls": call_count["n"], "loop_exited_clean": True, } @@ -522,10 +508,7 @@ async def test_iam_refresh_racing_reconnect_recreates_engine_only_once( from litellm.proxy.db.prisma_client import PrismaWrapper def token_db_url(created: datetime) -> str: - token = ( - f"host/?X-Amz-Date={created.strftime('%Y%m%dT%H%M%SZ')}" - f"&X-Amz-Expires=900&X-Amz-Signature=abc" - ) + token = f"host/?X-Amz-Date={created.strftime('%Y%m%dT%H%M%SZ')}&X-Amz-Expires=900&X-Amz-Signature=abc" return f"postgresql://user:{urllib.parse.quote(token, safe='')}@host:5432/db" # Old engine (PID 111) carries an expired token; in-flight queries on it @@ -577,9 +560,7 @@ async def test_iam_refresh_racing_reconnect_recreates_engine_only_once( # In-flight transport-error path fires while the refresh holds the # wrapper's reconnection lock mid-recreate. reconnect_task = asyncio.create_task( - prisma_client.attempt_db_reconnect( - reason="in_flight_transport_error", force=True - ) + prisma_client.attempt_db_reconnect(reason="in_flight_transport_error", force=True) ) await asyncio.sleep(0.05) release_connect.set() @@ -1096,3 +1077,27 @@ async def test_unrelated_reconnect_failure_does_not_erase_the_burst_record( "cycles_after": prisma_client._run_reconnect_cycle.await_count, } assert pinned == {"cycles_before": 2, "cycles_after": 2} + + +@pytest.mark.asyncio +async def test_attempt_db_reconnect_cancelled_while_waiting_does_not_strand_lock( + prisma_client: PrismaClient, +) -> None: + """A reconnect cancelled while waiting on the lock (e.g. the readiness + probe deadline firing) must abandon its lock-acquisition task instead of + leaving it to grab the lock later with no owner to release it.""" + prisma_client._db_last_reconnect_attempt_ts = 0.0 + prisma_client._attempt_reconnect_inside_lock = AsyncMock(return_value=True) + + await prisma_client._db_reconnect_lock.acquire() + waiting_reconnect: Final = asyncio.create_task( + prisma_client.attempt_db_reconnect(reason="probe_deadline", lock_timeout_seconds=30.0) + ) + await asyncio.sleep(0.05) + waiting_reconnect.cancel() + with pytest.raises(asyncio.CancelledError): + await waiting_reconnect + + prisma_client._db_reconnect_lock.release() + await asyncio.sleep(0.05) + assert prisma_client._db_reconnect_lock.locked() is False diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py index 65d3c3c8079..ec5b994f147 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -10,6 +10,7 @@ Covers ``_wrap_streaming_iterator_with_enrichment``, from __future__ import annotations import asyncio +from datetime import datetime from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock @@ -18,6 +19,10 @@ from fastapi import HTTPException import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( + BaseAnthropicMessagesStreamingIterator, +) from litellm.proxy.utils import ProxyLogging @@ -346,6 +351,134 @@ async def test_async_post_call_streaming_iterator_hook_upstream_error_raises(pro pass +# --------------------------------------------------------------------------- +# deferred native /v1/messages stream logging (LIT-6409) +# --------------------------------------------------------------------------- + + +_NATIVE_MESSAGES_STREAM_EVENTS = ( + {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 3, "output_tokens": 1}}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}}, + {"type": "message_stop"}, +) + + +def _armed_native_messages_stream(test_name: str, request_data: Dict[str, Any], events: List[Any]): + """The proxy-side setup for a native /v1/messages stream with post_call + guardrails active: a real BaseAnthropicMessagesStreamingIterator whose + logging_obj carries the deferred-dispatch callback the proxy arms in + common_request_processing. The callback records what the guardrail + metadata contained at the moment the deferred logging was dispatched.""" + logging_obj = LiteLLMLoggingObj( + model="bedrock/invoke/anthropic.claude-sonnet-4-20250514-v1:0", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id=test_name, + function_id=test_name, + ) + + async def _dispatch_deferred_logging(logging_coroutine): + events.append( + ( + "logging_dispatched", + "post_call_entry_visible", + bool(request_data.get("metadata", {}).get("standard_logging_guardrail_information")), + ) + ) + logging_coroutine.close() + + logging_obj._on_deferred_stream_complete = _dispatch_deferred_logging + request_data["litellm_logging_obj"] = logging_obj + + iterator = BaseAnthropicMessagesStreamingIterator(litellm_logging_obj=logging_obj, request_body={}) + + async def _upstream(): + for event in _NATIVE_MESSAGES_STREAM_EVENTS: + yield event + + return logging_obj, iterator.async_sse_wrapper(_upstream()) + + +@pytest.mark.asyncio +async def test_native_messages_stream_logging_fires_after_guardrail_end_of_stream_scan( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """ + Regression test for LIT-6409: on native /v1/messages streams the + end-of-stream guardrail scan writes its post_call entry AFTER the + upstream iterator is exhausted, so success logging dispatched at + upstream exhaustion never sees it. The deferred dispatch must fire + only after the guardrail chain fully drains. + """ + events: List[Any] = [] + request_data: Dict[str, Any] = {"metadata": {}} + _, native_stream = _armed_native_messages_stream( + "test_native_stream_deferred_ordering", request_data, events + ) + + class _EndOfStreamScanGuardrail(CustomLogger): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + async for chunk in response: + yield chunk + request_data.setdefault("metadata", {})["standard_logging_guardrail_information"] = [ + {"guardrail_mode": "post_call", "guardrail_status": "success"} + ] + events.append("scan_appended") + + monkeypatch.setattr(litellm, "callbacks", [_EndOfStreamScanGuardrail()]) + + async for _ in proxy_logging.async_post_call_streaming_iterator_hook( + response=native_stream, + user_api_key_dict=make_user_api_key_auth(), + request_data=request_data, + ): + pass + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert events == ["scan_appended", ("logging_dispatched", "post_call_entry_visible", True)] + + +@pytest.mark.asyncio +async def test_native_messages_stream_logging_fires_when_guardrail_blocks_after_stream_end( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """ + A guardrail block raised after upstream exhaustion (unified_guardrail + re-raises HTTPException for blocked content) must still flush the + parked deferred logging, or the blocked stream loses its spend log. + """ + events: List[Any] = [] + request_data: Dict[str, Any] = {"metadata": {}} + logging_obj, native_stream = _armed_native_messages_stream( + "test_native_stream_deferred_block", request_data, events + ) + + class _BlockingGuardrail(CustomLogger): + async def async_post_call_streaming_iterator_hook(self, user_api_key_dict, response, request_data): + async for chunk in response: + yield chunk + raise HTTPException(status_code=400, detail={"error": "Violated guardrail policy"}) + + monkeypatch.setattr(litellm, "callbacks", [_BlockingGuardrail()]) + + with pytest.raises(HTTPException): + async for _ in proxy_logging.async_post_call_streaming_iterator_hook( + response=native_stream, + user_api_key_dict=make_user_api_key_auth(), + request_data=request_data, + ): + pass + await asyncio.sleep(0) + await asyncio.sleep(0) + + assert [event[0] for event in events] == ["logging_dispatched"] + assert logging_obj._deferred_stream_complete_args is None + + # --------------------------------------------------------------------------- # _fire_deferred_stream_logging # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/router_strategy/test_router_tag_routing.py b/tests/test_litellm/router_strategy/test_router_tag_routing.py index 72bb6756d24..b33bd912be9 100644 --- a/tests/test_litellm/router_strategy/test_router_tag_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_routing.py @@ -3072,3 +3072,78 @@ async def test_non_router_tags_still_pick_the_matching_tier_deployment(): ) assert response._hidden_params["model_id"] == "tier-gemini-flash-us" + + +def _chat_completions_request_mock(): + from unittest.mock import MagicMock + + from fastapi import Request + + request_mock = MagicMock(spec=Request) + request_mock.url = MagicMock() + request_mock.url.path = "/v1/chat/completions" + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + return request_mock + + +def _team_a_and_default_router(): + return litellm.Router( + model_list=[ + { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "mock", "tags": ["team-a"]}, + "model_info": {"id": "team-a-deployment"}, + }, + { + "model_name": "gpt-5.4-mini", + "litellm_params": {"model": "openai/gpt-5.4-nano", "api_key": "mock", "tags": ["default"]}, + "model_info": {"id": "default-deployment"}, + }, + ], + enable_tag_filtering=True, + ) + + +@pytest.mark.asyncio() +@pytest.mark.parametrize( + "team_metadata,body_extra", + [ + ({"tags": ["team-a"]}, {}), + ({}, {"tags": ["team-a"]}), + ], + ids=["team-tags", "body-tags"], +) +async def test_chat_request_carrying_litellm_metadata_still_routes_on_proxy_merged_tags(team_metadata, body_extra): + from unittest.mock import MagicMock + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + + router = _team_a_and_default_router() + data = { + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": "hi"}], + "litellm_metadata": {"trace_id": "abc"}, + **body_extra, + } + + request_kwargs = await add_litellm_data_to_request( + data=data, + request=_chat_completions_request_mock(), + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", metadata={}, team_metadata=team_metadata), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + deployment = await router.async_get_available_deployment( + model="gpt-5.4-mini", + request_kwargs=request_kwargs, + messages=request_kwargs["messages"], + ) + + assert deployment["model_info"]["id"] == "team-a-deployment" diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index 94922e1a076..8336926c050 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -180,6 +180,30 @@ async def _acreate_file(*args: object, **kwargs: object) -> NoReturn: raise AssertionError("only used for its __name__") +async def _acancel_batch(*args: object, **kwargs: object) -> NoReturn: + raise AssertionError("only used for its __name__") + + +async def _acompletion(*args: object, **kwargs: object) -> NoReturn: + raise AssertionError("only used for its __name__") + + +async def _ageneric_api_call_with_fallbacks_helper(*args: object, **kwargs: object) -> NoReturn: + raise AssertionError("only used for its __name__") + + +async def acreate_fine_tuning_job(*args: object, **kwargs: object) -> NoReturn: + raise AssertionError("only used for its __name__") + + +async def aretrieve_fine_tuning_job(*args: object, **kwargs: object) -> NoReturn: + raise AssertionError("only used for its __name__") + + +async def afile_content(*args: object, **kwargs: object) -> NoReturn: + raise AssertionError("only used for its __name__") + + @pytest.mark.asyncio async def test_run_async_fallback_keeps_uploaded_file_requests_in_their_model_group(): """An input_file_id only exists under the credentials of the group it was uploaded @@ -217,6 +241,8 @@ async def test_run_async_fallback_keeps_fine_tuning_requests_in_their_model_grou fallback_depth=0, model="openai-group", training_file="file-owned-by-openai", + original_function=_ageneric_api_call_with_fallbacks_helper, + original_generic_function=acreate_fine_tuning_job, ) assert router.attempted_model_groups == [] @@ -299,6 +325,94 @@ async def test_run_async_fallback_still_crosses_model_groups_without_an_uploaded assert router.attempted_model_groups == ["azure-group"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("resource_key", "handler_kwargs"), + [ + ("batch_id", {"original_function": _acancel_batch}), + ( + "file_id", + { + "original_function": _ageneric_api_call_with_fallbacks_helper, + "original_generic_function": afile_content, + }, + ), + ( + "fine_tuning_job_id", + { + "original_function": _ageneric_api_call_with_fallbacks_helper, + "original_generic_function": aretrieve_fine_tuning_job, + }, + ), + ], +) +async def test_run_async_fallback_keeps_provider_scoped_ids_in_their_model_group( + resource_key: str, handler_kwargs: dict +): + """A batch, file, or fine-tuning job id only exists under the credentials of the group + that issued it, so a cross-group fallback asks a provider about an id it never saw. + Generic API calls carry the real handler in original_generic_function, so the pin + must recognize it there too.""" + router = AttemptRecordingRouter() + + with pytest.raises(RuntimeError, match="openai connection error"): + await run_async_fallback( + litellm_router=router, + fallback_model_group=["azure-group"], + original_model_group="openai-group", + original_exception=RuntimeError("openai connection error"), + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + **{resource_key: "owned-by-openai"}, + **handler_kwargs, + ) + + assert router.attempted_model_groups == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("resource_key", ["batch_id", "file_id", "fine_tuning_job_id"]) +async def test_run_async_fallback_ignores_stray_resource_ids_on_completion_calls(resource_key: str): + """A caller-supplied top-level field like file_id on a chat completion is application + data, never a provider resource reference, so it must not cost the request its + cross-group fallbacks.""" + router = AttemptRecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=["azure-group"], + original_model_group="openai-group", + original_exception=RuntimeError("openai connection error"), + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + original_function=_acompletion, + **{resource_key: "caller-app-data"}, + ) + + assert router.attempted_model_groups == ["azure-group"] + + +@pytest.mark.asyncio +async def test_run_async_fallback_allows_same_model_group_retry_for_batch_cancel(): + router = AttemptRecordingRouter() + + await run_async_fallback( + litellm_router=router, + fallback_model_group=[{"model": "openai-group", "_target_order": 2}], + original_model_group="openai-group", + original_exception=RuntimeError("first deployment failed"), + max_fallbacks=3, + fallback_depth=0, + model="openai-group", + batch_id="owned-by-openai", + original_function=_acancel_batch, + ) + + assert router.attempted_model_groups == ["openai-group"] + + @pytest.mark.asyncio async def test_run_async_fallback_handles_explicitly_none_metadata(): """/v1/batches always sets `metadata`, and sets it to None when the caller sent diff --git a/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py b/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py index 1e0e72c9ac6..7e655b70756 100644 --- a/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py +++ b/tests/test_litellm/secret_managers/test_aws_secret_manager_v2.py @@ -83,3 +83,60 @@ async def test_write_and_read_json_secret(): secret_name=test_secret_name ) assert delete_resp is not None + + +def _prepare_request_endpoint( + monkeypatch: pytest.MonkeyPatch, region_name: str, extra_optional_params: dict[str, str] | None = None +) -> str: + monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) + secret_manager = AWSSecretsManagerV2(aws_region_name=region_name) + endpoint_url, _headers, _body = secret_manager._prepare_request( + action="GetSecretValue", + secret_name="my-secret", + optional_params={ + "aws_access_key_id": "test-key", + "aws_secret_access_key": "test-secret", + **(extra_optional_params or {}), + }, + ) + return endpoint_url + + +@pytest.mark.parametrize( + "region_name,expected_endpoint", + [ + ("cn-north-1", "https://secretsmanager.cn-north-1.amazonaws.com.cn"), + ("cn-northwest-1", "https://secretsmanager.cn-northwest-1.amazonaws.com.cn"), + ("us-gov-west-1", "https://secretsmanager.us-gov-west-1.amazonaws.com"), + ("us-east-1", "https://secretsmanager.us-east-1.amazonaws.com"), + ], +) +def test_prepare_request_builds_partition_endpoint( + monkeypatch: pytest.MonkeyPatch, region_name: str, expected_endpoint: str +) -> None: + assert _prepare_request_endpoint(monkeypatch, region_name) == expected_endpoint + + +def test_prepare_request_explicit_bedrock_runtime_endpoint_param_still_wins(monkeypatch: pytest.MonkeyPatch) -> None: + endpoint_url = _prepare_request_endpoint( + monkeypatch, + "cn-north-1", + {"aws_bedrock_runtime_endpoint": "https://bedrock-runtime.my-vpce.example.com"}, + ) + assert endpoint_url == "https://secretsmanager.my-vpce.example.com" + + +def test_prepare_request_env_bedrock_runtime_endpoint_still_wins(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv( + "AWS_BEDROCK_RUNTIME_ENDPOINT", "https://bedrock-runtime.eu-west-1.amazonaws.com" + ) + secret_manager = AWSSecretsManagerV2(aws_region_name="cn-north-1") + endpoint_url, _headers, _body = secret_manager._prepare_request( + action="GetSecretValue", + secret_name="my-secret", + optional_params={ + "aws_access_key_id": "test-key", + "aws_secret_access_key": "test-secret", + }, + ) + assert endpoint_url == "https://secretsmanager.eu-west-1.amazonaws.com" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index ff8e568a935..97286017ffe 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -558,6 +558,59 @@ async def test_async_router_acreate_file_does_not_fall_back_across_model_groups( assert "gpt-4o-mini" not in called_models +@pytest.mark.asyncio +async def test_async_router_acancel_batch_does_not_fall_back_across_model_groups(monkeypatch: pytest.MonkeyPatch): + """The proxy cancels a managed batch by handing the router the deployment id decoded + from the unified batch id. A default (``*``) fallback matches that id like any other + model string, and the fallback provider is then asked to cancel a batch it never + issued, which can only answer not-found. The router re-raises the owner's error after + that wasted round trip, so the pin's observable is the foreign call never happening.""" + import respx + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router = litellm.Router( + model_list=[ + { + "model_name": "azure-gpt", + "litellm_params": { + "model": "azure/my-azure-deployment", + "api_base": "http://127.0.0.1:9", + "api_key": "dummy-key", + "api_version": "2024-06-01", + }, + "model_info": {"id": "azure-batch-dep"}, + }, + { + "model_name": "openai-gpt", + "litellm_params": {"model": "gpt-4o-mini", "api_key": "dummy-key"}, + }, + ], + default_fallbacks=["openai-gpt"], + ) + + with respx.mock(assert_all_called=False) as respx_mock: + azure_route = respx_mock.post(host="127.0.0.1").mock( + return_value=httpx.Response(401, json={"error": {"code": "401", "message": "invalid subscription key"}}) + ) + openai_route = respx_mock.post("https://api.openai.com/v1/batches/batch_owned_by_azure/cancel").mock( + return_value=httpx.Response( + 404, + json={ + "error": { + "message": "No batch found with id 'batch_owned_by_azure'.", + "type": "invalid_request_error", + "code": "batch_not_found", + } + }, + ) + ) + with pytest.raises(openai.AuthenticationError, match="invalid subscription key"): + await router.acancel_batch(model="azure-batch-dep", batch_id="batch_owned_by_azure") + + assert azure_route.called + assert not openai_route.called + + @pytest.mark.asyncio async def test_async_router_acreate_file_uses_deployment_custom_llm_provider(): """ diff --git a/type-discipline-budget.json b/type-discipline-budget.json index a7dec330a26..f3f1a7defe7 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22727 + "limit": 22705 }, "LIT002": { - "limit": 26873 + "limit": 26854 }, "LIT003": { "limit": 269 @@ -15,24 +15,24 @@ "limit": 0 }, "LIT006": { - "limit": 1065 + "limit": 1063 }, "LIT007": { "limit": 0 }, "LIT008": { - "limit": 948 + "limit": 945 }, "LIT009": { "limit": 0 }, "LIT010": { - "limit": 16616 + "limit": 16564 }, "LIT011": { - "limit": 5583 + "limit": 5577 }, "LIT012": { - "limit": 4509 + "limit": 4508 } } diff --git a/ui/litellm-dashboard/src/components/CodeBlock.tsx b/ui/litellm-dashboard/src/components/CodeBlock.tsx index 88ef9acf08e..a39d7417302 100644 --- a/ui/litellm-dashboard/src/components/CodeBlock.tsx +++ b/ui/litellm-dashboard/src/components/CodeBlock.tsx @@ -20,10 +20,10 @@ const CodeBlock = ({ code, language }: CodeBlockProps) => { }; return ( -
+