diff --git a/README.md b/README.md index 3ebaefb10ca..3db827d5fdd 100644 --- a/README.md +++ b/README.md @@ -203,7 +203,7 @@ curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ { "mcpServers": { "LiteLLM": { - "url": "http://localhost:4000/mcp", + "url": "http://localhost:4000/mcp/", "headers": { "x-litellm-api-key": "Bearer sk-1234" } diff --git a/docs/my-website/docs/completion/prompt_caching.md b/docs/my-website/docs/completion/prompt_caching.md index 630c9e58d24..dca5f5c0cff 100644 --- a/docs/my-website/docs/completion/prompt_caching.md +++ b/docs/my-website/docs/completion/prompt_caching.md @@ -63,7 +63,6 @@ for _ in range(2): } ], }, - # marked for caching with the cache_control parameter, so that this checkpoint can read from the previous cache. { "role": "user", "content": [ @@ -77,7 +76,6 @@ for _ in range(2): "role": "assistant", "content": "Certainly! the key terms and conditions are the following: the contract is 1 year long for $10/mo", }, - # The final turn is marked with cache-control, for continuing in followups. { "role": "user", "content": [ @@ -112,16 +110,16 @@ model_list: api_key: os.environ/OPENAI_API_KEY ``` -2. Start proxy +2. Start proxy ```bash litellm --config /path/to/config.yaml ``` -3. Test it! +3. Test it! ```python -from openai import OpenAI +from openai import OpenAI import os client = OpenAI( @@ -144,7 +142,6 @@ for _ in range(2): } ], }, - # marked for caching with the cache_control parameter, so that this checkpoint can read from the previous cache. { "role": "user", "content": [ @@ -158,7 +155,6 @@ for _ in range(2): "role": "assistant", "content": "Certainly! the key terms and conditions are the following: the contract is 1 year long for $10/mo", }, - # The final turn is marked with cache-control, for continuing in followups. { "role": "user", "content": [ @@ -183,6 +179,78 @@ assert response.usage.prompt_tokens_details.cached_tokens > 0 +### OpenAI `prompt_cache_key` and `prompt_cache_retention` + +OpenAI prompt caching is [**automatic**](https://platform.openai.com/docs/guides/prompt-caching) — no `cache_control` message annotations are needed. Any request with 1024+ prompt tokens is eligible for caching. + +OpenAI also supports two optional parameters for more control over caching behavior: + +- **`prompt_cache_key`** (string) — A routing hint that improves cache hit rates for requests sharing long common prefixes. Requests with the same cache key are routed to the same backend, increasing the likelihood of a cache hit. +- **`prompt_cache_retention`** (`"in_memory"` or `"24h"`) — Controls cache TTL. Default is `"in_memory"` (5–10 min). Set to `"24h"` for extended caching that offloads KV tensors to GPU-local storage. + + + + +```python +from litellm import completion +import os + +os.environ["OPENAI_API_KEY"] = "" + +response = completion( + model="gpt-4o", + messages=[ + { + "role": "system", + "content": "You are an AI assistant tasked with analyzing legal documents. " + + "Here is the full text of a complex legal agreement " * 400, + }, + { + "role": "user", + "content": "What are the key terms and conditions?", + }, + ], + prompt_cache_key="legal-doc-analysis", + prompt_cache_retention="24h", +) +print(response.usage) +``` + + + + +```python +from openai import OpenAI + +client = OpenAI( + api_key="LITELLM_PROXY_KEY", + base_url="LITELLM_PROXY_BASE", +) + +response = client.chat.completions.create( + model="gpt-4o", + messages=[ + { + "role": "system", + "content": "You are an AI assistant tasked with analyzing legal documents. " + + "Here is the full text of a complex legal agreement " * 400, + }, + { + "role": "user", + "content": "What are the key terms and conditions?", + }, + ], + extra_body={ + "prompt_cache_key": "legal-doc-analysis", + "prompt_cache_retention": "24h", + }, +) +print(response.usage) +``` + + + + ### Anthropic Example Anthropic charges for cache writes. diff --git a/docs/my-website/docs/generateContent.md b/docs/my-website/docs/generateContent.md index 4453e5ce06d..bf8e1b6c03b 100644 --- a/docs/my-website/docs/generateContent.md +++ b/docs/my-website/docs/generateContent.md @@ -15,6 +15,7 @@ Use LiteLLM to call Google AI's generateContent endpoints for text generation, m | Streaming | ✅ | | | Fallbacks | ✅ | between supported models | | Loadbalancing | ✅ | between supported models | +| Metadata Tracking | ✅ | passes trace ID, metadata to observability callbacks (e.g. S3, Langfuse) | ## Usage --- diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index 50973f220f5..fcbb31c07d3 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -641,7 +641,7 @@ import asyncio config = { "mcpServers": { "mcp_group": { - "url": "http://localhost:4000/mcp", + "url": "http://localhost:4000/mcp/", "headers": { "x-mcp-servers": "dev_group", # assume this gives access to github, zapier and deepwiki "x-litellm-api-key": "Bearer sk-1234", diff --git a/docs/my-website/docs/proxy/forward_client_headers.md b/docs/my-website/docs/proxy/forward_client_headers.md index 2155a7517be..17f813eabee 100644 --- a/docs/my-website/docs/proxy/forward_client_headers.md +++ b/docs/my-website/docs/proxy/forward_client_headers.md @@ -37,11 +37,11 @@ The following rules determine which headers are forwarded (see [`_get_forwardabl | Rule | Example | Forwarded? | |---|---|---| -| Headers starting with `x-` | `x-trace-id`, `x-custom-header`, `x-request-source` | ✅ Yes | -| `anthropic-beta` header | `anthropic-beta: prompt-caching-2024-07-31` | ✅ Yes | -| Headers starting with `x-stainless-*` | `x-stainless-lang`, `x-stainless-arch` | ❌ No (causes OpenAI SDK issues) | -| Standard HTTP headers | `Authorization`, `Content-Type`, `Host` | ❌ No | -| Other provider headers | `Accept`, `User-Agent` | ❌ No | +| Headers starting with `x-` | `x-trace-id`, `x-custom-header`, `x-request-source` | Yes | +| `anthropic-beta` header | `anthropic-beta: prompt-caching-2024-07-31` | Yes | +| Headers starting with `x-stainless-*` | `x-stainless-lang`, `x-stainless-arch` | No (causes OpenAI SDK issues) | +| Standard HTTP headers | `Authorization`, `Content-Type`, `Host` | No | +| Other provider headers | `Accept`, `User-Agent` | No | ### Additional Header Mechanisms @@ -61,6 +61,125 @@ general_settings: forward_client_headers_to_llm_api: true ``` +## Forward LLM Provider Authentication Headers + +**New in v1.82+**: By default, LiteLLM strips authentication headers like `x-api-key`, `x-goog-api-key`, and `api-key` from client requests for security (these are typically used to authenticate with the proxy itself). However, you can enable forwarding of these LLM provider authentication headers to allow **Bring Your Own Key (BYOK)** scenarios where clients send their own API keys to the LLM provider. + +### Configuration + +Add `forward_llm_provider_auth_headers: true` to your `general_settings`: + +```yaml +general_settings: + forward_client_headers_to_llm_api: true + forward_llm_provider_auth_headers: true # 👈 Enable BYOK +``` + +### Which Headers Are Forwarded + +When `forward_llm_provider_auth_headers: true`, the following LLM provider authentication headers are preserved and forwarded: + +| Header | Provider | Example | +|--------|----------|---------| +| `x-api-key` | Anthropic, Azure AI, Databricks | `x-api-key: sk-ant-api03-...` | +| `x-goog-api-key` | Google AI Studio | `x-goog-api-key: AIza...` | +| `api-key` | Azure OpenAI | `api-key: your-azure-key` | +| `ocp-apim-subscription-key` | Azure APIM | `ocp-apim-subscription-key: your-key` | + +:::warning Important Security Note +The proxy's `Authorization` header (used for proxy authentication) is **never** forwarded to LLM providers, even with this setting enabled. This ensures your proxy authentication remains secure. +::: + +### Use Case: Client-Side API Keys (BYOK) + +This feature enables scenarios where: +1. **Clients bring their own LLM provider API keys** instead of using keys configured in the proxy +2. **Multi-tenant applications** where each tenant has their own Anthropic/OpenAI account +3. **Development environments** where developers use their personal API keys through a shared proxy + +#### Example: Anthropic BYOK + +```yaml +# proxy_config.yaml +model_list: + - model_name: claude-sonnet-4 + litellm_params: + model: anthropic/claude-sonnet-4-20250514 + # No api_key configured! Will use client's key + +general_settings: + forward_client_headers_to_llm_api: true + forward_llm_provider_auth_headers: true # Enable BYOK +``` + +Client request: +```bash +curl -X POST "http://localhost:4000/v1/messages" \ + -H "Authorization: Bearer sk-proxy-auth-123" \ # Proxy authentication (stripped) + -H "x-api-key: sk-ant-api03-YOUR-KEY..." \ # Client's Anthropic key (forwarded!) + -H "Content-Type: application/json" \ + -d '{ + "model": "claude-sonnet-4", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 100 + }' +``` + +#### Example: Google AI Studio BYOK + +```yaml +model_list: + - model_name: gemini-pro + litellm_params: + model: gemini/gemini-1.5-pro + # No api_key configured + +general_settings: + forward_client_headers_to_llm_api: true + forward_llm_provider_auth_headers: true +``` + +Client request: +```bash +curl -X POST "http://localhost:4000/v1/chat/completions" \ + -H "Authorization: Bearer sk-proxy-auth-123" \ + -H "x-goog-api-key: AIza..." \ + -d '{ + "model": "gemini-pro", + "messages": [{"role": "user", "content": "Hello"}] + }' +``` + +### Security Considerations + +**When to Use This Feature:** +- Internal tools where you trust all clients +- Development/testing environments +- Multi-tenant apps with proper client authentication +- Scenarios where you want clients to use their own API keys + +**When NOT to Use:** +- Public APIs where you don't trust all clients +- When you want centralized billing/cost control +- When you need to enforce rate limits at the proxy level + +### Backward Compatibility + +For backward compatibility, if you have `forward_client_headers_to_llm_api: true` but don't explicitly set `forward_llm_provider_auth_headers`, the behavior is: +- **Default**: LLM provider auth headers are **NOT** forwarded (safe default) +- **Explicit `true`**: LLM provider auth headers **ARE** forwarded (BYOK enabled) + +```yaml +# Safe default - auth headers NOT forwarded +general_settings: + forward_client_headers_to_llm_api: true + +# BYOK enabled - auth headers ARE forwarded +general_settings: + forward_client_headers_to_llm_api: true + forward_llm_provider_auth_headers: true # 👈 Opt-in required +``` + ## Enable for a Model Group Add the `forward_client_headers_to_llm_api` setting under `model_group_settings` in your configuration: diff --git a/docs/my-website/docs/tutorials/presidio_pii_masking.md b/docs/my-website/docs/tutorials/presidio_pii_masking.md index ea3761163f2..d6fe1adbd01 100644 --- a/docs/my-website/docs/tutorials/presidio_pii_masking.md +++ b/docs/my-website/docs/tutorials/presidio_pii_masking.md @@ -592,6 +592,21 @@ def test_pii_masking_allows_normal_text(): ## Part 7: Troubleshooting +### Issue: Guardrail failure: non-JSON response from Presidio + +**Symptom:** You receive an error indicating `expected application/json Content-Type but received text/html` or similar. + +**Root cause:** Your ingress controller or reverse proxy might be routing the `/analyze` or `/anonymize` POST request to a health endpoint (like `/health` or `/presidio-analyzer/health`) which returns plain text instead of JSON. + +**Fix:** Ensure your `PRESIDIO_ANALYZER_API_BASE` and `PRESIDIO_ANONYMIZER_API_BASE` are correctly pointing directly to the Presidio API endpoints, or that your ingress routes the path correctly without stripping it and inadvertently forwarding to a plain-text health check endpoint. + +**Verification:** You can verify your endpoints using `curl`. It should return a JSON array, not `text/html`: +```bash +curl -sv -X POST http://your-analyzer-endpoint/analyze \ + -H "Content-Type: application/json" \ + -d '{"text":"test","language":"en"}' +``` + ### Issue: Presidio Not Detecting PII **Check 1: Language Configuration** diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 25f6e284bcd..9553d2c5246 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -37,7 +37,9 @@ from litellm.types.llms.openai import ( ) from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( + LIST_BATCHES_SUPPORTED_PROVIDERS, OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS, + ListBatchesSupportedProvider, LiteLLMBatch, LlmProviders, ) @@ -674,7 +676,7 @@ def retrieve_batch( async def alist_batches( after: Optional[str] = None, limit: Optional[int] = None, - custom_llm_provider: Literal["openai", "azure", "hosted_vllm", "vertex_ai"] = "openai", + custom_llm_provider: ListBatchesSupportedProvider = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -717,7 +719,7 @@ async def alist_batches( def list_batches( after: Optional[str] = None, limit: Optional[int] = None, - custom_llm_provider: Literal["openai", "azure", "hosted_vllm", "vertex_ai"] = "openai", + custom_llm_provider: ListBatchesSupportedProvider = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -843,8 +845,9 @@ def list_batches( ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'list_batch'. Supported providers: openai, azure, vertex_ai.".format( - custom_llm_provider + message="LiteLLM doesn't support {} for 'list_batch'. Supported providers: {}.".format( + custom_llm_provider, + ", ".join(sorted(LIST_BATCHES_SUPPORTED_PROVIDERS)), ), model="n/a", llm_provider=custom_llm_provider, diff --git a/litellm/constants.py b/litellm/constants.py index ee79f2fa56f..b1a0021bcc6 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -244,6 +244,7 @@ REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) # Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth LITELLM_ASYNCIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)) +TOOL_POLICY_CACHE_TTL_SECONDS = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60)) # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger. # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire. MAX_SIZE_IN_MEMORY_QUEUE = int( diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index cc3916af069..47a27c8ef5b 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -4,6 +4,8 @@ Helper functions for health check calls. from typing import TYPE_CHECKING, Callable, Dict, Literal, Optional +from litellm.types.utils import LIST_BATCHES_SUPPORTED_PROVIDERS + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging @@ -82,6 +84,27 @@ class HealthCheckHelpers: "tags": [LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME], } + @staticmethod + async def _batch_health_check( + custom_llm_provider: str, + model_params: dict, + filtered_model_params: dict, + ) -> dict: + """ + Health check for batch mode. + + Calls list_batches for providers that support it (openai, hosted_vllm, azure, + vertex_ai). For all other providers (e.g. bedrock) the batch API surface doesn't + include list_batches, so we fall back to acompletion to verify connectivity and + credential validity instead. + """ + import litellm + + if custom_llm_provider in LIST_BATCHES_SUPPORTED_PROVIDERS: + return await litellm.alist_batches(**filtered_model_params) + else: + return await litellm.acompletion(**model_params) + @staticmethod def get_mode_handlers( model: str, @@ -176,8 +199,10 @@ class HealthCheckHelpers: api_key=model_params.get("api_key", None), api_version=model_params.get("api_version", None), ), - "batch": lambda: litellm.alist_batches( - **_filter_model_params(model_params=model_params), + "batch": lambda: HealthCheckHelpers._batch_health_check( + custom_llm_provider=custom_llm_provider, + model_params=model_params, + filtered_model_params=_filter_model_params(model_params=model_params), ), "responses": lambda: litellm.aresponses( **_filter_model_params(model_params=model_params), 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 c532d8ea27c..2abcc679eef 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py @@ -11,6 +11,7 @@ from typing import Any, List, Optional import httpx +from litellm.types.utils import Usage from litellm.llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import ( AmazonQwen3Config, ) @@ -79,10 +80,11 @@ class AmazonQwen2Config(AmazonQwen3Config): # Set usage information if available in response if "usage" in response_data: usage_data = response_data["usage"] - if hasattr(model_response, 'usage'): - model_response.usage.prompt_tokens = usage_data.get("prompt_tokens", 0) - model_response.usage.completion_tokens = usage_data.get("completion_tokens", 0) - model_response.usage.total_tokens = usage_data.get("total_tokens", 0) + model_response.usage = Usage( + prompt_tokens=usage_data.get("prompt_tokens", 0), + completion_tokens=usage_data.get("completion_tokens", 0), + total_tokens=usage_data.get("total_tokens", 0), + ) return model_response 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 b3a957ce0f8..12333623f51 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen3_transformation.py @@ -10,6 +10,7 @@ from typing import Any, List, Optional import httpx +from litellm.types.utils import Usage from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( AmazonInvokeConfig, @@ -201,10 +202,11 @@ class AmazonQwen3Config(AmazonInvokeConfig, BaseConfig): # Set usage information if available in response if "usage" in response_data: usage_data = response_data["usage"] - if hasattr(model_response, 'usage'): - model_response.usage.prompt_tokens = usage_data.get("prompt_tokens", 0) - model_response.usage.completion_tokens = usage_data.get("completion_tokens", 0) - model_response.usage.total_tokens = usage_data.get("total_tokens", 0) + model_response.usage = Usage( + prompt_tokens=usage_data.get("prompt_tokens", 0), + completion_tokens=usage_data.get("completion_tokens", 0), + total_tokens=usage_data.get("total_tokens", 0), + ) return model_response diff --git a/litellm/llms/openai/chat/gpt_transformation.py b/litellm/llms/openai/chat/gpt_transformation.py index aa8471a5973..ab102a69670 100644 --- a/litellm/llms/openai/chat/gpt_transformation.py +++ b/litellm/llms/openai/chat/gpt_transformation.py @@ -162,6 +162,7 @@ class OpenAIGPTConfig(BaseLLMModelInfo, BaseConfig): "service_tier", "safety_identifier", "prompt_cache_key", + "prompt_cache_retention", "store", ] # works across all models diff --git a/litellm/llms/openai/chat/o_series_transformation.py b/litellm/llms/openai/chat/o_series_transformation.py index 30647f58687..6ef43ec5bfd 100644 --- a/litellm/llms/openai/chat/o_series_transformation.py +++ b/litellm/llms/openai/chat/o_series_transformation.py @@ -131,9 +131,7 @@ class OpenAIOSeriesConfig(OpenAIGPTConfig): def is_model_o_series_model(self, model: str) -> bool: model = model.split("/")[-1] # could be "openai/o3" or "o3" - return model in litellm.open_ai_chat_completion_models and any( - model.startswith(pfx) for pfx in ("o1", "o3", "o4") - ) + return model.startswith(("o1", "o3", "o4")) and model in litellm.open_ai_chat_completion_models @overload def _transform_messages( @@ -173,4 +171,4 @@ class OpenAIOSeriesConfig(OpenAIGPTConfig): else: return super()._transform_messages( messages, model, is_async=cast(Literal[False], False) - ) + ) \ No newline at end of file diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 581671598c6..6e1a252be73 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -10,6 +10,7 @@ from litellm.proxy._experimental.mcp_server.ui_session_utils import ( ) from litellm.proxy._experimental.mcp_server.utils import merge_mcp_headers from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.mcp import MCPAuth @@ -685,7 +686,7 @@ if MCP_AVAILABLE: return await _execute_with_mcp_client( new_mcp_server_request, _test_connection_operation, - raw_headers=dict(request.headers), + raw_headers=_safe_get_request_headers(request), ) @router.post("/test/tools/list") @@ -744,5 +745,5 @@ if MCP_AVAILABLE: _list_tools_operation, mcp_auth_header=mcp_auth_header, oauth2_headers=oauth2_headers, - raw_headers=dict(request.headers), + raw_headers=_safe_get_request_headers(request), ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f354e28acd7..4053d9d077b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -183,6 +183,7 @@ class LitellmTableNames(str, enum.Enum): KEY_TABLE_NAME = "LiteLLM_VerificationToken" PROXY_MODEL_TABLE_NAME = "LiteLLM_ProxyModelTable" MANAGED_FILE_TABLE_NAME = "LiteLLM_ManagedFileTable" + TOOL_TABLE_NAME = "LiteLLM_ToolTable" class Litellm_EntityType(enum.Enum): @@ -850,6 +851,7 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): max_budget: Optional[float] = None user_id: Optional[str] = None team_id: Optional[str] = None + agent_id: Optional[str] = None max_parallel_requests: Optional[int] = None metadata: Optional[dict] = {} tpm_limit: Optional[int] = None @@ -4123,6 +4125,15 @@ class SpendUpdateQueueItem(TypedDict, total=False): response_cost: Optional[float] +class ToolDiscoveryQueueItem(TypedDict, total=False): + tool_name: str + origin: Optional[str] # MCP server name or "user_defined" + created_by: Optional[str] + key_hash: Optional[str] # hash of virtual key that triggered discovery + team_id: Optional[str] # team that triggered discovery + key_alias: Optional[str] # human-readable key alias + + class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): unified_file_id: str file_object: Optional[OpenAIFileObject] = None diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 138f9bab2c4..9240d84db83 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -483,7 +483,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 parent_otel_span = ( open_telemetry_logger.create_litellm_proxy_request_started_span( start_time=start_time, - headers=dict(request.headers), + headers=_safe_get_request_headers(request), ) ) @@ -562,7 +562,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, parent_otel_span=parent_otel_span, - request_headers=dict(request.headers), + request_headers=_safe_get_request_headers(request), ) is_proxy_admin = result["is_proxy_admin"] @@ -593,9 +593,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 user_id=user_id, team_id=team_id, team_alias=( - team_object.team_alias - if team_object is not None - else None + team_object.team_alias if team_object is not None else None ), team_metadata=team_object.metadata if team_object is not None @@ -709,12 +707,12 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if isinstance(api_key, str): return UserAPIKeyAuth( api_key=api_key, - user_role=LitellmUserRoles.PROXY_ADMIN, + user_role=LitellmUserRoles.INTERNAL_USER, parent_otel_span=parent_otel_span, ) else: return UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, + user_role=LitellmUserRoles.INTERNAL_USER, parent_otel_span=parent_otel_span, ) elif api_key is None: # only require api key if master key is set @@ -846,7 +844,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) valid_token.parent_otel_span = parent_otel_span if _end_user_object is not None: - valid_token.end_user_object_permission = _end_user_object.object_permission + valid_token.end_user_object_permission = ( + _end_user_object.object_permission + ) return valid_token @@ -954,7 +954,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if isinstance( api_key, str ): # if generated token, make sure it starts with sk-. - _masked_key = "{}****{}".format(api_key[:4], api_key[-4:]) if len(api_key) > 8 else "****" + _masked_key = ( + "{}****{}".format(api_key[:4], api_key[-4:]) + if len(api_key) > 8 + else "****" + ) assert api_key.startswith( "sk-" ), "LiteLLM Virtual Key expected. Received={}, expected to start with 'sk-'.".format( @@ -1304,9 +1308,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if _end_user_object is not None: valid_token_dict.update(end_user_params) - valid_token_dict["end_user_object_permission"] = ( - _end_user_object.object_permission - ) + valid_token_dict[ + "end_user_object_permission" + ] = _end_user_object.object_permission # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions # sso/login, ui/login, /key functions and /user functions diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index e1bca6e905f..8d179a9caed 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -135,17 +135,29 @@ def _safe_set_request_parsed_body( def _safe_get_request_headers(request: Optional[Request]) -> dict: """ - [Non-Blocking] Safely get the request headers + [Non-Blocking] Safely get the request headers. + Caches the result on request.state to avoid re-creating dict(request.headers) per call. + + Warning: Callers must NOT mutate the returned dict — it is shared across + all callers within the same request via the cache. """ + if request is None: + return {} + cached = getattr(request.state, "_cached_headers", None) + if cached is not None: + return cached try: - if request is None: - return {} - return dict(request.headers) + headers = dict(request.headers) except Exception as e: verbose_proxy_logger.debug( "Unexpected error reading request headers - {}".format(e) ) - return {} + headers = {} + try: + request.state._cached_headers = headers + except Exception: + pass # request.state may not be available in all contexts + return headers def check_file_size_under_limit( diff --git a/litellm/proxy/custom_hooks/custom_ui_sso_hook.py b/litellm/proxy/custom_hooks/custom_ui_sso_hook.py index cad0e62ff56..8bb6b274091 100644 --- a/litellm/proxy/custom_hooks/custom_ui_sso_hook.py +++ b/litellm/proxy/custom_hooks/custom_ui_sso_hook.py @@ -3,6 +3,7 @@ from fastapi_sso.sso.base import OpenID from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers class CustomSSOLoginHandler(CustomLogger): @@ -18,7 +19,7 @@ class CustomSSOLoginHandler(CustomLogger): self, request: Request, ) -> OpenID: - request_headers_dict = dict(request.headers) + request_headers_dict = _safe_get_request_headers(request) verbose_logger.debug("inside custom ui sso sign in hook...") return OpenID( id=request_headers_dict.get("x-litellm-user-id") or "123", diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index d7d6b2b1eb0..edf0cf0d397 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -13,7 +13,17 @@ import random import time import traceback from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast, overload +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + Optional, + Union, + cast, + overload, +) import litellm from litellm._logging import verbose_proxy_logger @@ -23,18 +33,19 @@ from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, BaseDailySpendTransaction, - DailyTagSpendTransaction, - DailyOrganizationSpendTransaction, - DailyTeamSpendTransaction, - DailyEndUserSpendTransaction, - DailyUserSpendTransaction, DailyAgentSpendTransaction, + DailyEndUserSpendTransaction, + DailyOrganizationSpendTransaction, + DailyTagSpendTransaction, + DailyTeamSpendTransaction, + DailyUserSpendTransaction, DBSpendUpdateTransactions, Litellm_EntityType, LiteLLM_UserTable, SpendLogsMetadata, SpendLogsPayload, SpendUpdateQueueItem, + ToolDiscoveryQueueItem, ) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, @@ -42,6 +53,9 @@ from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, +) from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING if TYPE_CHECKING: @@ -67,6 +81,7 @@ class DBSpendUpdateWriter: self.redis_update_buffer = RedisUpdateBuffer(redis_cache=self.redis_cache) self.pod_lock_manager = PodLockManager() self.spend_update_queue = SpendUpdateQueue() + self.tool_discovery_queue = ToolDiscoveryQueue() self.daily_spend_update_queue = DailySpendUpdateQueue() self.daily_team_spend_update_queue = DailySpendUpdateQueue() self.daily_end_user_spend_update_queue = DailySpendUpdateQueue() @@ -222,6 +237,13 @@ class DBSpendUpdateWriter: ) ) + self._enqueue_tool_registry_upsert( + kwargs=kwargs, + completion_response=completion_response, + hashed_token=hashed_token, + team_id=team_id, + ) + verbose_proxy_logger.debug("Runs spend update on all tables") except Exception: verbose_proxy_logger.error( @@ -237,6 +259,104 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) + def _enqueue_tool_registry_upsert( + self, + kwargs: Optional[dict], + completion_response: Optional[Any], + hashed_token: Optional[str] = None, + team_id: Optional[str] = None, + ) -> None: + """ + Extract tool names from the LLM request and response and enqueue them + for upsert into LiteLLM_ToolTable via ToolDiscoveryQueue. + + Handles four sources: + - MCP tools: standard_logging_object.mcp_tool_call_metadata.namespaced_tool_name + - Response tool_calls (OpenAI / Anthropic pass-through converted to OpenAI format): + completion_response.choices[].message.tool_calls[].function.name + - Request tools array (OpenAI format): kwargs["tools"][].function.name + - Request tools array (Anthropic /messages format): kwargs["passthrough_logging_payload"] + ["request_body"]["tools"][].name + """ + try: + if kwargs is None: + return + + # Extract key_alias from kwargs metadata if available + key_alias: Optional[str] = None + _litellm_params = kwargs.get("litellm_params") or {} + _metadata = _litellm_params.get("metadata") or {} + key_alias = _metadata.get("user_api_key_alias") or None + + def _enqueue(tool_name: str, origin: str = "user_defined") -> None: + self.tool_discovery_queue.add_update( + ToolDiscoveryQueueItem( + tool_name=tool_name, + origin=origin, + key_hash=hashed_token, + team_id=team_id, + key_alias=key_alias, + ) + ) + + # --- MCP tool calls --- + sl_object = kwargs.get("standard_logging_object") + if sl_object is not None: + mcp_metadata = ( + sl_object.get("metadata", {}) or {} + ).get("mcp_tool_call_metadata") + if mcp_metadata and isinstance(mcp_metadata, dict): + tool_name = mcp_metadata.get("namespaced_tool_name") or mcp_metadata.get("name") + mcp_server_name = mcp_metadata.get("mcp_server_name") + if tool_name: + _enqueue(tool_name, origin=mcp_server_name or "user_defined") + + # --- Tools from request body (OpenAI format: tools[].function.name) --- + request_tools = kwargs.get("tools") or [] + for tool_def in request_tools: + if not isinstance(tool_def, dict): + continue + fn = tool_def.get("function") or {} + name = fn.get("name") if isinstance(fn, dict) else None + if name: + _enqueue(name) + + # --- Tools from Anthropic /messages pass-through request body + # (Anthropic format: tools[].name, no "function" wrapper) --- + passthrough_payload = kwargs.get("passthrough_logging_payload") or {} + request_body = ( + passthrough_payload.get("request_body") + if isinstance(passthrough_payload, dict) + else None + ) or {} + for tool_def in request_body.get("tools") or []: + if not isinstance(tool_def, dict): + continue + name = tool_def.get("name") + if name: + _enqueue(name) + + # --- Response tool_calls (OpenAI format; Anthropic pass-through converts tool_use here) --- + if completion_response is not None and hasattr(completion_response, "choices"): + for choice in completion_response.choices or []: + message = getattr(choice, "message", None) + if message is None: + continue + tool_calls = getattr(message, "tool_calls", None) + if not tool_calls: + continue + for tc in tool_calls: + fn = getattr(tc, "function", None) + if fn is None: + continue + tool_name = getattr(fn, "name", None) + if tool_name: + _enqueue(tool_name) + except Exception as e: + verbose_proxy_logger.debug( + "_enqueue_tool_registry_upsert error (non-blocking): %s", e + ) + async def _update_key_db( self, response_cost: Optional[float], @@ -752,6 +872,25 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_agent_spend_update_transactions, ) + ################## Tool Registry Upserts ################## + await self._flush_tool_discovery_queue(prisma_client=prisma_client) + + async def _flush_tool_discovery_queue( + self, + prisma_client: PrismaClient, + ) -> None: + """Flush ToolDiscoveryQueue and batch-upsert new tools into LiteLLM_ToolTable.""" + from litellm.proxy.db.tool_registry_writer import batch_upsert_tools + + try: + items = self.tool_discovery_queue.flush() + if items: + await batch_upsert_tools(prisma_client=prisma_client, items=items) + except Exception as e: + verbose_proxy_logger.debug( + "_flush_tool_discovery_queue error (non-blocking): %s", e + ) + async def _commit_spend_updates_to_db( # noqa: PLR0915 self, prisma_client: PrismaClient, diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py new file mode 100644 index 00000000000..16a3ada40f2 --- /dev/null +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -0,0 +1,54 @@ +""" +In-memory buffer for tool registry upserts. + +Unlike SpendUpdateQueue (which aggregates increments), ToolDiscoveryQueue +uses set-deduplication: each unique tool_name is only queued once per flush +cycle (~30s). The seen-set is cleared on every flush so that call_count +increments in subsequent cycles rather than stopping after the first flush. +""" + +from typing import List, Set + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ToolDiscoveryQueueItem + + +class ToolDiscoveryQueue: + """ + In-memory buffer for tool registry upserts. + + Deduplicates by tool_name within each flush cycle: a tool is only queued + once per ~30s batch, so call_count increments once per flush cycle the + tool appears in (not once per invocation, but not once per pod lifetime + either). The seen-set is cleared on flush so subsequent batches can + re-count the same tool. + """ + + def __init__(self) -> None: + self._seen_tool_names: Set[str] = set() + self._pending: List[ToolDiscoveryQueueItem] = [] + + def add_update(self, item: ToolDiscoveryQueueItem) -> None: + """Enqueue a tool discovery item if tool_name has not been seen before.""" + tool_name = item.get("tool_name", "") + if not tool_name: + return + if tool_name in self._seen_tool_names: + verbose_proxy_logger.debug( + "ToolDiscoveryQueue: skipping already-seen tool %s", tool_name + ) + return + self._seen_tool_names.add(tool_name) + self._pending.append(item) + verbose_proxy_logger.debug( + "ToolDiscoveryQueue: queued new tool %s (origin=%s)", + tool_name, + item.get("origin"), + ) + + def flush(self) -> List[ToolDiscoveryQueueItem]: + """Return and clear all pending items. Resets seen-set so the next + flush cycle can re-count the same tools.""" + items, self._pending = self._pending, [] + self._seen_tool_names.clear() + return items diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py new file mode 100644 index 00000000000..4e0a8095a08 --- /dev/null +++ b/litellm/proxy/db/tool_registry_writer.py @@ -0,0 +1,179 @@ +""" +DB helpers for LiteLLM_ToolTable — the global tool registry. + +Tools are auto-discovered from LLM responses and upserted here. +Admins use the management endpoints to read and update call_policy. + +NOTE: Uses raw SQL (query_raw / execute_raw) instead of Prisma model methods +because the generated Prisma Python client may not have LiteLLM_ToolTable +when running against an older generated schema. +""" + +import uuid +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Dict, List, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ToolDiscoveryQueueItem +from litellm.types.tool_management import LiteLLM_ToolTableRow, ToolCallPolicy + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +def _row_to_model(row: dict) -> LiteLLM_ToolTableRow: + return LiteLLM_ToolTableRow( + tool_id=row.get("tool_id", ""), + tool_name=row.get("tool_name", ""), + origin=row.get("origin"), + call_policy=row.get("call_policy", "untrusted"), + call_count=int(row.get("call_count") or 0), + assignments=row.get("assignments"), + key_hash=row.get("key_hash"), + team_id=row.get("team_id"), + key_alias=row.get("key_alias"), + created_at=row.get("created_at"), + updated_at=row.get("updated_at"), + created_by=row.get("created_by"), + updated_by=row.get("updated_by"), + ) + + +async def batch_upsert_tools( + prisma_client: "PrismaClient", + items: List[ToolDiscoveryQueueItem], +) -> None: + """ + Batch-upsert tool registry rows via raw SQL. + + On first insert: sets call_policy = "untrusted" (schema default), call_count = 1. + On conflict: increments call_count; preserves existing call_policy. + """ + if not items: + return + try: + data = [item for item in items if item.get("tool_name")] + if not data: + return + for item in data: + tool_name = item.get("tool_name", "") + origin = item.get("origin") or "user_defined" + created_by = item.get("created_by") or "system" + key_hash = item.get("key_hash") + team_id = item.get("team_id") + key_alias = item.get("key_alias") + now = datetime.now(timezone.utc).isoformat() + await prisma_client.db.execute_raw( + 'INSERT INTO "LiteLLM_ToolTable" ' + "(tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias, created_at, updated_at) " + "VALUES ($7, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6, $8, $8) " + "ON CONFLICT (tool_name) DO UPDATE SET " + "call_count = \"LiteLLM_ToolTable\".call_count + 1, " + "updated_at = $8", + tool_name, + origin, + created_by, + key_hash, + team_id, + key_alias, + str(uuid.uuid4()), + now, + ) + verbose_proxy_logger.debug( + "tool_registry_writer: upserted %d tool(s)", len(data) + ) + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer batch_upsert_tools error: %s", e) + + +async def list_tools( + prisma_client: "PrismaClient", + call_policy: Optional[ToolCallPolicy] = None, +) -> List[LiteLLM_ToolTableRow]: + """Return all tools, optionally filtered by call_policy.""" + try: + if call_policy is not None: + rows = await prisma_client.db.query_raw( + 'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, ' + 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' + 'FROM "LiteLLM_ToolTable" WHERE call_policy = $1 ORDER BY created_at DESC', + call_policy, + ) + else: + rows = await prisma_client.db.query_raw( + 'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, ' + 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' + 'FROM "LiteLLM_ToolTable" ORDER BY created_at DESC', + ) + return [_row_to_model(row) for row in rows] + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer list_tools error: %s", e) + return [] + + +async def get_tool( + prisma_client: "PrismaClient", + tool_name: str, +) -> Optional[LiteLLM_ToolTableRow]: + """Return a single tool row by tool_name.""" + try: + rows = await prisma_client.db.query_raw( + 'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, ' + 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' + 'FROM "LiteLLM_ToolTable" WHERE tool_name = $1', + tool_name, + ) + if not rows: + return None + return _row_to_model(rows[0]) + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer get_tool error: %s", e) + return None + + +async def update_tool_policy( + prisma_client: "PrismaClient", + tool_name: str, + call_policy: ToolCallPolicy, + updated_by: Optional[str], +) -> Optional[LiteLLM_ToolTableRow]: + """Update the call_policy for a tool. Upserts the row if it does not exist yet.""" + try: + _updated_by = updated_by or "system" + now = datetime.now(timezone.utc).isoformat() + await prisma_client.db.execute_raw( + 'INSERT INTO "LiteLLM_ToolTable" (tool_id, tool_name, call_policy, created_by, updated_by, created_at, updated_at) ' + "VALUES ($4, $1, $2, $3, $3, $5, $5) " + "ON CONFLICT (tool_name) DO UPDATE SET call_policy = $2, updated_by = $3, updated_at = $5", + tool_name, + call_policy, + _updated_by, + str(uuid.uuid4()), + now, + ) + return await get_tool(prisma_client, tool_name) + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer update_tool_policy error: %s", e) + return None + + +async def get_tools_by_names( + prisma_client: "PrismaClient", + tool_names: List[str], +) -> Dict[str, str]: + """ + Return a {tool_name: call_policy} map for the given tool names. + Used by the policy enforcement guardrail — single batch query, never N+1. + """ + if not tool_names: + return {} + try: + placeholders = ", ".join(f"${i+1}" for i in range(len(tool_names))) + rows = await prisma_client.db.query_raw( + f'SELECT tool_name, call_policy FROM "LiteLLM_ToolTable" WHERE tool_name IN ({placeholders})', + *tool_names, + ) + return {row["tool_name"]: row["call_policy"] for row in rows} + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer get_tools_by_names error: %s", e) + return {} diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 5c41b371ca4..9768d93e922 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -1,6 +1,10 @@ +from datetime import datetime + from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import ORJSONResponse, StreamingResponse +import litellm +from litellm._uuid import uuid from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing @@ -17,7 +21,8 @@ router = APIRouter( dependencies=[Depends(user_api_key_auth)], ) @router.post( - "/models/{model_name:path}:generateContent", dependencies=[Depends(user_api_key_auth)] + "/models/{model_name:path}:generateContent", + dependencies=[Depends(user_api_key_auth)], ) async def google_generate_content( request: Request, @@ -36,12 +41,12 @@ async def google_generate_content( data = await _read_request_body(request=request) if "model" not in data: data["model"] = model_name - + # Extract generationConfig and pass it as config parameter generation_config = data.pop("generationConfig", None) if generation_config: data["config"] = generation_config - + # Add user authentication metadata for cost tracking data = await add_litellm_data_to_request( data=data, @@ -51,7 +56,19 @@ async def google_generate_content( general_settings=general_settings, version=version, ) - + + # Create logging object with full request metadata so callbacks (e.g. S3) get user/trace_id + data["litellm_call_id"] = request.headers.get( + "x-litellm-call-id", str(uuid.uuid4()) + ) + logging_obj, data = litellm.utils.function_setup( + original_function="agenerate_content", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + **data, + ) + data["litellm_logging_obj"] = logging_obj + # call router if llm_router is None: raise HTTPException(status_code=500, detail="Router not initialized") @@ -103,6 +120,18 @@ async def google_stream_generate_content( version=version, ) + # Create logging object with full request metadata so streaming END callbacks (e.g. S3) get user/trace_id + data["litellm_call_id"] = request.headers.get( + "x-litellm-call-id", str(uuid.uuid4()) + ) + logging_obj, data = litellm.utils.function_setup( + original_function="agenerate_content_stream", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + **data, + ) + data["litellm_logging_obj"] = logging_obj + # call router if llm_router is None: raise HTTPException(status_code=500, detail="Router not initialized") @@ -247,11 +276,11 @@ async def create_interaction( ) data = await _read_request_body(request=request) - + # Default to gemini provider for interactions if "custom_llm_provider" not in data: data["custom_llm_provider"] = "gemini" - + processor = ProxyBaseLLMRequestProcessing(data=data) try: return await processor.base_process_llm_request( @@ -301,7 +330,7 @@ async def get_interaction( ): """ Get an interaction by ID. - + Per OpenAPI spec: GET /{api_version}/interactions/{interaction_id} """ from litellm.proxy.proxy_server import ( @@ -319,7 +348,7 @@ async def get_interaction( ) data = {"interaction_id": interaction_id, "custom_llm_provider": "gemini"} - + processor = ProxyBaseLLMRequestProcessing(data=data) try: return await processor.base_process_llm_request( @@ -369,7 +398,7 @@ async def delete_interaction( ): """ Delete an interaction by ID. - + Per OpenAPI spec: DELETE /{api_version}/interactions/{interaction_id} """ from litellm.proxy.proxy_server import ( @@ -387,7 +416,7 @@ async def delete_interaction( ) data = {"interaction_id": interaction_id, "custom_llm_provider": "gemini"} - + processor = ProxyBaseLLMRequestProcessing(data=data) try: return await processor.base_process_llm_request( @@ -437,7 +466,7 @@ async def cancel_interaction( ): """ Cancel an interaction by ID. - + Per OpenAPI spec: POST /{api_version}/interactions/{interaction_id}:cancel """ from litellm.proxy.proxy_server import ( @@ -455,7 +484,7 @@ async def cancel_interaction( ) data = {"interaction_id": interaction_id, "custom_llm_provider": "gemini"} - + processor = ProxyBaseLLMRequestProcessing(data=data) try: return await processor.base_process_llm_request( diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index c083c60cb4c..20f6e6f1d39 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -2,6 +2,7 @@ CRUD ENDPOINTS FOR GUARDRAILS """ +import concurrent.futures import inspect from typing import Any, Dict, List, Optional, Type, TypeVar, Union, cast @@ -11,9 +12,16 @@ from pydantic import BaseModel from litellm._logging import verbose_proxy_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry +from litellm.proxy.guardrails.guardrail_hooks.custom_code.code_validator import ( + CustomCodeValidationError, + validate_custom_code, +) +from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import ( + get_custom_code_primitives, +) from litellm.proxy.guardrails.usage_endpoints import router as guardrails_usage_router from litellm.types.guardrails import ( PII_ENTITY_CATEGORIES_MAP, @@ -243,9 +251,11 @@ class CreateGuardrailRequest(BaseModel): @router.post( "/guardrails", tags=["Guardrails"], - dependencies=[Depends(user_api_key_auth)], ) -async def create_guardrail(request: CreateGuardrailRequest): +async def create_guardrail( + request: CreateGuardrailRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Create a new guardrail @@ -296,6 +306,12 @@ async def create_guardrail(request: CreateGuardrailRequest): from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Admin access required to manage guardrails", + ) + if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") @@ -332,9 +348,12 @@ class UpdateGuardrailRequest(BaseModel): @router.put( "/guardrails/{guardrail_id}", tags=["Guardrails"], - dependencies=[Depends(user_api_key_auth)], ) -async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest): +async def update_guardrail( + guardrail_id: str, + request: UpdateGuardrailRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Update an existing guardrail @@ -385,6 +404,12 @@ async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest): from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Admin access required to manage guardrails", + ) + if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") @@ -429,9 +454,11 @@ async def update_guardrail(guardrail_id: str, request: UpdateGuardrailRequest): @router.delete( "/guardrails/{guardrail_id}", tags=["Guardrails"], - dependencies=[Depends(user_api_key_auth)], ) -async def delete_guardrail(guardrail_id: str): +async def delete_guardrail( + guardrail_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Delete a guardrail @@ -453,6 +480,12 @@ async def delete_guardrail(guardrail_id: str): from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Admin access required to manage guardrails", + ) + if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") @@ -495,9 +528,12 @@ async def delete_guardrail(guardrail_id: str): @router.patch( "/guardrails/{guardrail_id}", tags=["Guardrails"], - dependencies=[Depends(user_api_key_auth)], ) -async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): +async def patch_guardrail( + guardrail_id: str, + request: PatchGuardrailRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Partially update an existing guardrail @@ -546,6 +582,12 @@ async def patch_guardrail(guardrail_id: str, request: PatchGuardrailRequest): from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER from litellm.proxy.proxy_server import prisma_client + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Admin access required to manage guardrails", + ) + if prisma_client is None: raise HTTPException(status_code=500, detail="Prisma client not initialized") @@ -1302,9 +1344,9 @@ async def get_provider_specific_params(): lakera_v2_fields = _get_fields_from_model(LakeraV2GuardrailConfigModel) tool_permission_fields = _get_fields_from_model(ToolPermissionGuardrailConfigModel) - tool_permission_fields["ui_friendly_name"] = ( - ToolPermissionGuardrailConfigModel.ui_friendly_name() - ) + tool_permission_fields[ + "ui_friendly_name" + ] = ToolPermissionGuardrailConfigModel.ui_friendly_name() # Return the provider-specific parameters provider_params = { @@ -1364,10 +1406,12 @@ class TestCustomCodeGuardrailResponse(BaseModel): @router.post( "/guardrails/test_custom_code", tags=["Guardrails"], - dependencies=[Depends(user_api_key_auth)], response_model=TestCustomCodeGuardrailResponse, ) -async def test_custom_code_guardrail(request: TestCustomCodeGuardrailRequest): +async def test_custom_code_guardrail( + request: TestCustomCodeGuardrailRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Test custom code guardrail logic without creating a guardrail. @@ -1440,63 +1484,26 @@ async def test_custom_code_guardrail(request: TestCustomCodeGuardrailRequest): } ``` """ - import concurrent.futures - import re - from litellm.proxy.guardrails.guardrail_hooks.custom_code.primitives import ( - get_custom_code_primitives, - ) - - # Security validation patterns - FORBIDDEN_PATTERNS = [ - # Import statements - (r"\bimport\s+", "import statements are not allowed"), - (r"\bfrom\s+\w+\s+import\b", "from...import statements are not allowed"), - (r"__import__\s*\(", "__import__() is not allowed"), - # Dangerous builtins - (r"\bexec\s*\(", "exec() is not allowed"), - (r"\beval\s*\(", "eval() is not allowed"), - (r"\bcompile\s*\(", "compile() is not allowed"), - (r"\bopen\s*\(", "open() is not allowed"), - (r"\bgetattr\s*\(", "getattr() is not allowed"), - (r"\bsetattr\s*\(", "setattr() is not allowed"), - (r"\bdelattr\s*\(", "delattr() is not allowed"), - (r"\bglobals\s*\(", "globals() is not allowed"), - (r"\blocals\s*\(", "locals() is not allowed"), - (r"\bvars\s*\(", "vars() is not allowed"), - (r"\bdir\s*\(", "dir() is not allowed"), - (r"\bbreakpoint\s*\(", "breakpoint() is not allowed"), - (r"\binput\s*\(", "input() is not allowed"), - # Dangerous dunder access - (r"__builtins__", "__builtins__ access is not allowed"), - (r"__globals__", "__globals__ access is not allowed"), - (r"__code__", "__code__ access is not allowed"), - (r"__subclasses__", "__subclasses__ access is not allowed"), - (r"__bases__", "__bases__ access is not allowed"), - (r"__mro__", "__mro__ access is not allowed"), - (r"__class__", "__class__ access is not allowed"), - (r"__dict__", "__dict__ access is not allowed"), - (r"__getattribute__", "__getattribute__ access is not allowed"), - (r"__reduce__", "__reduce__ access is not allowed"), - (r"__reduce_ex__", "__reduce_ex__ access is not allowed"), - # OS/system access - (r"\bos\.", "os module access is not allowed"), - (r"\bsys\.", "sys module access is not allowed"), - (r"\bsubprocess\.", "subprocess module access is not allowed"), - ] + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Admin access required to test custom code guardrails", + ) EXECUTION_TIMEOUT_SECONDS = 5 try: # Step 0: Security validation - check for forbidden patterns - code = request.custom_code - for pattern, error_msg in FORBIDDEN_PATTERNS: - if re.search(pattern, code): - return TestCustomCodeGuardrailResponse( - success=False, - error=f"Security violation: {error_msg}", - error_type="compilation", - ) + + try: + validate_custom_code(request.custom_code) + except CustomCodeValidationError as e: + return TestCustomCodeGuardrailResponse( + success=False, + error=str(e), + error_type="compilation", + ) # Step 1: Compile the custom code with restricted environment exec_globals = get_custom_code_primitives().copy() @@ -1612,10 +1619,10 @@ async def apply_guardrail( from litellm.proxy.utils import handle_exception_on_proxy try: - active_guardrail: Optional[CustomGuardrail] = ( - GUARDRAIL_REGISTRY.get_initialized_guardrail_callback( - guardrail_name=request.guardrail_name - ) + active_guardrail: Optional[ + CustomGuardrail + ] = GUARDRAIL_REGISTRY.get_initialized_guardrail_callback( + guardrail_name=request.guardrail_name ) if active_guardrail is None: raise HTTPException( diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/code_validator.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/code_validator.py new file mode 100644 index 00000000000..6ef59b522a8 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/code_validator.py @@ -0,0 +1,63 @@ +import re +from typing import List, Tuple + +# Security validation patterns +FORBIDDEN_PATTERNS: List[Tuple[str, str]] = [ + # Import statements + (r"\bimport\s+", "import statements are not allowed"), + (r"\bfrom\s+\w+\s+import\b", "from...import statements are not allowed"), + (r"__import__\s*\(", "__import__() is not allowed"), + # Dangerous builtins + (r"\bexec\s*\(", "exec() is not allowed"), + (r"\beval\s*\(", "eval() is not allowed"), + (r"\bcompile\s*\(", "compile() is not allowed"), + (r"\bopen\s*\(", "open() is not allowed"), + (r"\bgetattr\s*\(", "getattr() is not allowed"), + (r"\bsetattr\s*\(", "setattr() is not allowed"), + (r"\bdelattr\s*\(", "delattr() is not allowed"), + (r"\bglobals\s*\(", "globals() is not allowed"), + (r"\blocals\s*\(", "locals() is not allowed"), + (r"\bvars\s*\(", "vars() is not allowed"), + (r"\bdir\s*\(", "dir() is not allowed"), + (r"\bbreakpoint\s*\(", "breakpoint() is not allowed"), + (r"\binput\s*\(", "input() is not allowed"), + # Dangerous dunder access + (r"__builtins__", "__builtins__ access is not allowed"), + (r"__globals__", "__globals__ access is not allowed"), + (r"__code__", "__code__ access is not allowed"), + (r"__subclasses__", "__subclasses__ access is not allowed"), + (r"__bases__", "__bases__ access is not allowed"), + (r"__mro__", "__mro__ access is not allowed"), + (r"__class__", "__class__ access is not allowed"), + (r"__dict__", "__dict__ access is not allowed"), + (r"__getattribute__", "__getattribute__ access is not allowed"), + (r"__reduce__", "__reduce__ access is not allowed"), + (r"__reduce_ex__", "__reduce_ex__ access is not allowed"), + # OS/system access + (r"\bos\.", "os module access is not allowed"), + (r"\bsys\.", "sys module access is not allowed"), + (r"\bsubprocess\.", "subprocess module access is not allowed"), + (r"\bshutil\.", "shutil module access is not allowed"), + (r"\bctypes\.", "ctypes module access is not allowed"), + (r"\bsocket\.", "socket module access is not allowed"), + (r"\bpickle\.", "pickle module access is not allowed"), +] + + +class CustomCodeValidationError(Exception): + """Raised when custom code fails security validation.""" + + pass + + +def validate_custom_code(code: str) -> None: + """ + Validate custom code against forbidden patterns. + + Raises CustomCodeValidationError if any forbidden pattern is found. + """ + if not code: + return + for pattern, error_msg in FORBIDDEN_PATTERNS: + if re.search(pattern, code): + raise CustomCodeValidationError(f"Security violation: {error_msg}") 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 c557a093c4e..0f5a4384d76 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 @@ -41,18 +41,19 @@ from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, cast from fastapi import HTTPException from litellm._logging import verbose_proxy_logger -from litellm.integrations.custom_guardrail import (CustomGuardrail, - log_guardrail_information) +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.proxy.guardrails.guardrail_hooks.base import \ - GuardrailConfigModel +from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel from litellm.types.utils import GenericGuardrailAPIInputs +from .code_validator import CustomCodeValidationError, validate_custom_code from .primitives import get_custom_code_primitives if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import \ - Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj class CustomCodeGuardrailError(Exception): @@ -143,6 +144,33 @@ class CustomCodeGuardrail(CustomGuardrail): """Returns the config model for the UI.""" return CustomCodeGuardrailConfigModel + def _do_compile(self) -> None: + """Internal compilation method without lock. Expected to run inside _compile_lock.""" + # Create a restricted execution environment + # Only include our safe primitives + exec_globals = get_custom_code_primitives().copy() + + # CRITICAL: Restrict __builtins__ to prevent sandbox escape + exec_globals["__builtins__"] = {} + + # Execute the user code in the restricted environment + exec(compile(self.custom_code, "", "exec"), exec_globals) + + # Extract the apply_guardrail function + if "apply_guardrail" not in exec_globals: + raise CustomCodeCompilationError( + "Custom code must define an 'apply_guardrail' function. " + "Expected signature: apply_guardrail(inputs, request_data, input_type)" + ) + + apply_fn = exec_globals["apply_guardrail"] + if not callable(apply_fn): + raise CustomCodeCompilationError( + "'apply_guardrail' must be a callable function" + ) + + self._compiled_function = apply_fn + def _compile_custom_code(self) -> None: """ Compile the custom code and extract the apply_guardrail function. @@ -154,27 +182,14 @@ class CustomCodeGuardrail(CustomGuardrail): return try: - # Create a restricted execution environment - # Only include our safe primitives - exec_globals = get_custom_code_primitives().copy() + # Step 1: Security validation — forbidden pattern check + try: + validate_custom_code(self.custom_code) + except CustomCodeValidationError as e: + raise CustomCodeCompilationError(str(e)) from e - # Execute the user code in the restricted environment - exec(compile(self.custom_code, "", "exec"), exec_globals) - - # Extract the apply_guardrail function - if "apply_guardrail" not in exec_globals: - raise CustomCodeCompilationError( - "Custom code must define an 'apply_guardrail' function. " - "Expected signature: apply_guardrail(inputs, request_data, input_type)" - ) - - apply_fn = exec_globals["apply_guardrail"] - if not callable(apply_fn): - raise CustomCodeCompilationError( - "'apply_guardrail' must be a callable function" - ) - - self._compiled_function = apply_fn + # Step 2: Compile logic + self._do_compile() verbose_proxy_logger.debug( f"Custom code guardrail '{self.guardrail_name}' compiled successfully" ) @@ -390,6 +405,12 @@ class CustomCodeGuardrail(CustomGuardrail): Raises: CustomCodeCompilationError: If the new code fails to compile """ + # Validate BEFORE acquiring lock / resetting state + try: + validate_custom_code(new_code) + except CustomCodeValidationError as e: + raise CustomCodeCompilationError(str(e)) from e + with self._compile_lock: # Reset state old_function = self._compiled_function @@ -399,12 +420,24 @@ class CustomCodeGuardrail(CustomGuardrail): try: self.custom_code = new_code - self._compile_custom_code() + self._do_compile() verbose_proxy_logger.info( f"Custom code guardrail '{self.guardrail_name}': Code updated successfully" ) + except SyntaxError as e: + # Rollback on failure + self.custom_code = old_code + self._compiled_function = old_function + self._compile_error = f"Syntax error in custom code: {e}" + raise CustomCodeCompilationError(self._compile_error) from e except CustomCodeCompilationError: # Rollback on failure self.custom_code = old_code self._compiled_function = old_function raise + except Exception as e: + # Rollback on failure + self.custom_code = old_code + self._compiled_function = old_function + self._compile_error = f"Failed to compile custom code: {e}" + raise CustomCodeCompilationError(self._compile_error) from e diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 28ccff0e36c..34fbf47253b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -322,12 +322,6 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): analyze_payload, ) - async with session.post(analyze_url, json=analyze_payload) as response: - analyze_results = await response.json() - verbose_proxy_logger.debug("analyze_results: %s", analyze_results) - - # Handle error responses from Presidio (e.g., {'error': 'No text provided'}) - # Presidio may return a dict instead of a list when errors occur def _fail_on_invalid_response( reason: str, ) -> List[PresidioAnalyzeResponseItem]: @@ -347,6 +341,36 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) return [] + async with session.post( + analyze_url, + json=analyze_payload, + headers={"Accept": "application/json"}, + ) as response: + # Validate HTTP status + if response.status >= 400: + error_body = await response.text() + return _fail_on_invalid_response( + f"HTTP {response.status} from Presidio analyzer: {error_body[:200]}" + ) + + # Validate Content-Type is JSON + content_type = getattr( + response, + "content_type", + response.headers.get("Content-Type", ""), + ) + if "application/json" not in content_type: + error_body = await response.text() + return _fail_on_invalid_response( + f"expected application/json Content-Type but received '{content_type}'; body: '{error_body[:200]}'" + ) + + analyze_results = await response.json() + verbose_proxy_logger.debug("analyze_results: %s", analyze_results) + + # Handle error responses from Presidio (e.g., {'error': 'No text provided'}) + # Presidio may return a dict instead of a list when errors occur + if isinstance(analyze_results, dict): if "error" in analyze_results: return _fail_on_invalid_response( @@ -423,8 +447,29 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): } async with session.post( - anonymize_url, json=anonymize_payload + anonymize_url, + json=anonymize_payload, + headers={"Accept": "application/json"}, ) as response: + # Validate HTTP status + if response.status >= 400: + error_body = await response.text() + raise Exception( + f"Presidio anonymizer returned HTTP {response.status}: {error_body[:200]}" + ) + + # Validate Content-Type is JSON + content_type = getattr( + response, + "content_type", + response.headers.get("Content-Type", ""), + ) + if "application/json" not in content_type: + error_body = await response.text() + raise Exception( + f"Presidio anonymizer returned non-JSON Content-Type '{content_type}'; body: '{error_body[:200]}'" + ) + redacted_text = await response.json() new_text = text @@ -456,7 +501,11 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): except Exception as e: # Sanitize exception to avoid leaking the original text (which may # contain API keys or other secrets) in error responses. - if "Invalid anonymizer response" in str(e): + error_str = str(e) + if ( + "Invalid anonymizer response" in error_str + or "Presidio anonymizer returned" in error_str + ): raise raise Exception( f"Presidio PII anonymization failed: {type(e).__name__}" diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py new file mode 100644 index 00000000000..5a43006e23c --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py @@ -0,0 +1,16 @@ +import litellm +from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail): + from litellm.proxy.guardrails.guardrail_hooks.tool_policy.tool_policy_guardrail import ( + ToolPolicyGuardrail, + ) + + _callback = ToolPolicyGuardrail( + guardrail_name=guardrail.get("guardrail_name", "tool_policy"), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(_callback) + return _callback diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py new file mode 100644 index 00000000000..87558566c42 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -0,0 +1,163 @@ +""" +Tool Policy Guardrail + +Reads call_policy from LiteLLM_ToolTable and enforces it on LLM requests/responses. + +Policy values: + "trusted" - allow through (no action) + "untrusted" - allow through (no action; default for newly discovered tools) + "blocked" - raise HTTPException, preventing the tool call + "dual_llm" - (Phase 3) send to second LLM for verification; currently treated as allowed + +Configuration in proxy config YAML: + guardrails: + - guardrail_name: "tool_policy" + litellm_params: + guardrail: tool_policy + mode: post_call + +or both pre and post call: + - guardrail_name: "tool_policy" + litellm_params: + guardrail: tool_policy + mode: during_call # runs before LLM and on response +""" + +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache +from litellm.constants import TOOL_POLICY_CACHE_TTL_SECONDS +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +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 + +GUARDRAIL_NAME = "tool_policy" + + +class ToolPolicyGuardrail(CustomGuardrail): + """ + Guardrail that enforces per-tool call policies stored in LiteLLM_ToolTable. + + Tools with call_policy="blocked" are rejected before/after the LLM call. + Tools with call_policy="trusted" or "untrusted" pass through unchanged. + """ + + def __init__(self, **kwargs: Any) -> None: + if "supported_event_hooks" not in kwargs: + kwargs["supported_event_hooks"] = [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + GuardrailEventHooks.during_call, + ] + super().__init__(**kwargs) + self._policy_cache: DualCache = DualCache() + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """ + Enforce tool policies on both request tools and response tool_calls. + + - input_type="request": check inputs["tools"] (tool definitions in the LLM request) + - input_type="response": check inputs["tool_calls"] (tool_calls in the LLM response) + + Raises HTTPException (400) if any tool is "blocked". + """ + if input_type == "request": + tools = inputs.get("tools") or [] + tool_names = [ + t["function"]["name"] + for t in tools + if isinstance(t, dict) + and isinstance(t.get("function"), dict) + and t["function"].get("name") + ] + else: # response + tool_calls = inputs.get("tool_calls") or [] + tool_names = [] + for tc in tool_calls: + fn = None + if isinstance(tc, dict): + fn = (tc.get("function") or {}).get("name") + elif hasattr(tc, "function"): + fn = getattr(tc.function, "name", None) + if fn: + tool_names.append(fn) + + if not tool_names: + return inputs + + policy_map = await self._get_policies_cached(tool_names) + + blocked = [name for name in tool_names if policy_map.get(name) == "blocked"] + if blocked: + verbose_proxy_logger.warning( + "ToolPolicyGuardrail: blocking tool(s) %s (policy=blocked)", blocked + ) + raise HTTPException( + status_code=400, + detail={ + "error": "Violated tool policy", + "blocked_tools": blocked, + "message": f"Tool(s) {blocked} are blocked by policy.", + }, + ) + + return inputs + + async def _get_policies_cached(self, tool_names: List[str]) -> Dict[str, str]: + """ + Batch-fetch call_policy for the given tool names. + + Caches per individual tool name (not per combination) so that adding + a new tool to a request doesn't invalidate the cached policies for all + the other tools already in the cache. + """ + from litellm.proxy.db.tool_registry_writer import get_tools_by_names + from litellm.proxy.proxy_server import prisma_client + + if not tool_names or prisma_client is None: + return {} + + result: Dict[str, str] = {} + cache_misses: List[str] = [] + + for name in tool_names: + cached = await self._policy_cache.async_get_cache(f"tool_policy:{name}") + if cached is not None and isinstance(cached, str): + result[name] = cached + else: + cache_misses.append(name) + + if cache_misses: + fetched = await get_tools_by_names( + prisma_client=prisma_client, tool_names=cache_misses + ) + for name, policy in fetched.items(): + result[name] = policy + await self._policy_cache.async_set_cache( + key=f"tool_policy:{name}", + value=policy, + ttl=TOOL_POLICY_CACHE_TTL_SECONDS, + ) + verbose_proxy_logger.debug( + "ToolPolicyGuardrail: fetched %d policies from DB (cache hits: %d)", + len(cache_misses), + len(tool_names) - len(cache_misses), + ) + + return result diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index de4517efc88..d228bdb2129 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -283,12 +283,17 @@ async def perform_health_check( model: Optional[str] = None, cli_model: Optional[str] = None, details: Optional[bool] = True, + model_id: Optional[str] = None, max_concurrency: Optional[int] = None, instrumentation_context: Optional[dict] = None, ): """ Perform a health check on the system. + When model_id is provided, only the deployment with that id is checked + (so models that share the same name but have different ids are checked separately). + When model (name) is provided, all deployments matching that name are checked. + Returns: (bool): True if the health check passes, False otherwise. """ @@ -314,7 +319,12 @@ async def perform_health_check( cycle_start_time = time.monotonic() requested_model_count = len(model_list) - if model is not None: + # Filter by model_id first so a single deployment is checked when id is specified + if model_id is not None: + _by_id = [x for x in model_list if (x.get("model_info") or {}).get("id") == model_id] + if _by_id: + model_list = _by_id + elif model is not None: _new_model_list = [ x for x in model_list if x["litellm_params"]["model"] == model ] diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 3570b2dd6aa..f3bed3656f6 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -243,7 +243,7 @@ async def health_services_endpoint( # noqa: PLR0915 service_in_success_callbacks = True else: for cb in litellm.success_callback: - if hasattr(cb, "callback_name") and cb.callback_name == service: + if getattr(cb, "callback_name", None) == service: service_in_success_callbacks = True break cb_id = get_callback_identifier(cb) @@ -747,6 +747,7 @@ async def _perform_health_check_and_save( model=target_model, details=details, max_concurrency=max_concurrency, + model_id=model_id, ) # Optionally save health check result to database (non-blocking) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 815d64f22ad..0734756d8ed 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -12,7 +12,7 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.auth.auth_checks import log_db_metrics +from litellm.proxy.auth.auth_checks import get_key_object, get_team_object, log_db_metrics from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.utils import ProxyUpdateSpend from litellm.types.utils import ( @@ -76,6 +76,10 @@ class _ProxyDBLogger(CustomLogger): traceback_str=traceback_str, ) + _metadata = await _ProxyDBLogger._enrich_failure_metadata_with_key_info( + metadata=_metadata, + ) + existing_metadata: dict = request_data.get("metadata", None) or {} existing_metadata.update(_metadata) @@ -255,6 +259,72 @@ class _ProxyDBLogger(CustomLogger): "Error in tracking cost callback - %s", str(e) ) + @staticmethod + async def _enrich_failure_metadata_with_key_info(metadata: dict) -> dict: + """ + Enriches failure spend log metadata by looking up the key object (and team object) + from cache/DB when key fields are missing. + + This handles two scenarios: + 1. Auth errors (401): UserAPIKeyAuth is created with only api_key set, all other + fields are null. We look up the full key object to fill in alias, user_id, + team_id, etc. + 2. Post-auth failures (provider errors, rate limits): key fields are populated + but team_alias is missing because LiteLLM_VerificationTokenView SQL view + doesn't include it. We look up the team object to fill in team_alias. + """ + api_key_hash = metadata.get("user_api_key") + if not api_key_hash: + return metadata + + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) + + # Step 1: If key fields are missing, look up the full key object + if metadata.get("user_api_key_alias") is None: + try: + key_obj = await get_key_object( + hashed_token=api_key_hash, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if metadata.get("user_api_key_alias") is None: + metadata["user_api_key_alias"] = key_obj.key_alias + if metadata.get("user_api_key_user_id") is None: + metadata["user_api_key_user_id"] = key_obj.user_id + if metadata.get("user_api_key_team_id") is None: + metadata["user_api_key_team_id"] = key_obj.team_id + if metadata.get("user_api_key_org_id") is None: + metadata["user_api_key_org_id"] = key_obj.org_id + except Exception: + verbose_proxy_logger.debug( + "Failed to enrich failure metadata with key info for api_key=%s", + api_key_hash, + ) + + # Step 2: If team_id is known but team_alias is missing, look up the team object + team_id = metadata.get("user_api_key_team_id") + if team_id and metadata.get("user_api_key_team_alias") is None: + try: + team_obj = await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if team_obj.team_alias is not None: + metadata["user_api_key_team_alias"] = team_obj.team_alias + except Exception: + verbose_proxy_logger.debug( + "Failed to enrich failure metadata with team_alias for team_id=%s", + team_id, + ) + return metadata + @staticmethod def _should_track_errors_in_db(): """ diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index b61dfa5b263..d2312a00c3b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -10,10 +10,16 @@ import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.proxy._types import (AddTeamCallback, CommonProxyErrors, - LitellmDataForBackendLLMCall, - LitellmUserRoles, SpecialHeaders, - TeamCallbackMetadata, UserAPIKeyAuth) +from litellm.proxy._types import ( + AddTeamCallback, + CommonProxyErrors, + LitellmDataForBackendLLMCall, + LitellmUserRoles, + SpecialHeaders, + TeamCallbackMetadata, + UserAPIKeyAuth, +) +from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers # Cache special headers as a frozenset for O(1) lookup performance _SPECIAL_HEADERS_CACHE = frozenset( @@ -22,9 +28,12 @@ _SPECIAL_HEADERS_CACHE = frozenset( from litellm.router import Router from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS from litellm.types.services import ServiceTypes -from litellm.types.utils import (LlmProviders, ProviderSpecificHeader, - StandardLoggingUserAPIKeyMetadata, - SupportedCacheControls) +from litellm.types.utils import ( + LlmProviders, + ProviderSpecificHeader, + StandardLoggingUserAPIKeyMetadata, + SupportedCacheControls, +) service_logger_obj = ServiceLogging() # used for tracking latency on OTEL @@ -227,7 +236,9 @@ def _get_dynamic_logging_metadata( def clean_headers( - headers: Headers, litellm_key_header_name: Optional[str] = None + headers: Headers, + litellm_key_header_name: Optional[str] = None, + forward_llm_provider_auth_headers: bool = False, ) -> dict: """ Removes litellm api key from headers @@ -237,15 +248,18 @@ def clean_headers( clean_headers = {} litellm_key_lower = ( litellm_key_header_name.lower() if litellm_key_header_name is not None else None - ) - + ) for header, value in headers.items(): header_lower = header.lower() - # Preserve Authorization header if it contains Anthropic OAuth token (sk-ant-oat*) - # This allows OAuth tokens to be forwarded to Anthropic-compatible providers - # via add_provider_specific_headers_to_request() + if header_lower == "authorization" and is_anthropic_oauth_key(value): clean_headers[header] = value + elif forward_llm_provider_auth_headers and header_lower in _SPECIAL_HEADERS_CACHE: + if litellm_key_lower and header_lower == litellm_key_lower: + continue + if header_lower == "authorization": + continue + clean_headers[header] = value # Check if header should be excluded: either in special headers cache or matches custom litellm key elif header_lower not in _SPECIAL_HEADERS_CACHE and ( litellm_key_lower is None or header_lower != litellm_key_lower @@ -653,7 +667,8 @@ class LiteLLMProxyRequestSetup: return data from litellm.proxy._types import ( LiteLLM_ManagementEndpoint_MetadataFields, - LiteLLM_ManagementEndpoint_MetadataFields_Premium) + LiteLLM_ManagementEndpoint_MetadataFields_Premium, + ) # ignore any special fields added_metadata = {} @@ -824,7 +839,12 @@ async def add_litellm_data_to_request( # noqa: PLR0915 from litellm.proxy.proxy_server import llm_router, premium_user from litellm.types.proxy.litellm_pre_call_utils import SecretFields - _raw_headers: Dict[str, str] = dict(request.headers) + _raw_headers: Dict[str, str] = _safe_get_request_headers(request) + + forward_llm_auth = False + if general_settings: + forward_llm_auth = general_settings.get("forward_llm_provider_auth_headers", False) + _headers: Dict[str, str] = clean_headers( request.headers, litellm_key_header_name=( @@ -832,7 +852,10 @@ async def add_litellm_data_to_request( # noqa: PLR0915 if general_settings is not None else None ), + forward_llm_provider_auth_headers=forward_llm_auth, ) + verbose_proxy_logger.debug(f"Request Headers: {_headers}") + verbose_proxy_logger.debug(f"Raw Headers: {_raw_headers}") ########################################################## # Init - Proxy Server Request @@ -1478,8 +1501,7 @@ async def move_guardrails_to_metadata( # Only check policy engine if no local config (avoid import + registry lookup) if not (has_key_config or has_team_config or has_request_config): - from litellm.proxy.policy_engine.policy_registry import \ - get_policy_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry if not get_policy_registry().is_initialized(): # Nothing configured anywhere - clean up request body fields and return @@ -1543,16 +1565,14 @@ async def move_guardrails_to_metadata( def _is_policy_version_id(s: str) -> bool: """Return True if string is a policy version ID (starts with policy_ prefix).""" - from litellm.proxy.policy_engine.policy_registry import \ - POLICY_VERSION_ID_PREFIX + from litellm.proxy.policy_engine.policy_registry import POLICY_VERSION_ID_PREFIX return isinstance(s, str) and s.startswith(POLICY_VERSION_ID_PREFIX) def _extract_policy_id(s: str) -> Optional[str]: """Extract raw UUID from policy_ string, or None if not a valid version ID.""" - from litellm.proxy.policy_engine.policy_registry import \ - POLICY_VERSION_ID_PREFIX + from litellm.proxy.policy_engine.policy_registry import POLICY_VERSION_ID_PREFIX if not _is_policy_version_id(s): return None @@ -1573,9 +1593,10 @@ def _match_and_track_policies( """ from litellm._logging import verbose_proxy_logger from litellm.proxy.common_utils.callback_utils import ( - add_policy_sources_to_metadata, add_policy_to_applied_policies_header) - from litellm.proxy.policy_engine.attachment_registry import \ - get_attachment_registry + add_policy_sources_to_metadata, + add_policy_to_applied_policies_header, + ) + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher # Get matching policies via attachments (with match reasons for attribution) @@ -1720,8 +1741,7 @@ async def add_guardrails_from_policy_engine( user_api_key_dict: The user's API key authentication info """ from litellm._logging import verbose_proxy_logger - from litellm.proxy.common_utils.http_parsing_utils import \ - get_tags_from_request_body + from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.types.proxy.policy_engine import PolicyMatchContext diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a230c2e9336..c1165ab26d0 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2535,6 +2535,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 user_id: Optional[str] = None, user_alias: Optional[str] = None, team_id: Optional[str] = None, + agent_id: Optional[str] = None, user_email: Optional[str] = None, user_role: Optional[str] = None, max_parallel_requests: Optional[int] = None, @@ -2668,6 +2669,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 "max_budget": key_max_budget, "user_id": user_id, "team_id": team_id, + "agent_id": agent_id, "project_id": project_id, "max_parallel_requests": max_parallel_requests, "metadata": metadata_json, diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index da10d3bf3b4..73fcce72c38 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -21,6 +21,7 @@ from typing_extensions import TypedDict import litellm from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers from litellm._uuid import uuid from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import ( @@ -724,7 +725,7 @@ async def get_service_provider_config(request: Request): "SCIM ServiceProviderConfig request: method=%s url=%s headers=%s", request.method, request.url, - dict(request.headers), + _safe_get_request_headers(request), ) meta = { "resourceType": "ServiceProviderConfig", diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py new file mode 100644 index 00000000000..89880c9a4ec --- /dev/null +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -0,0 +1,149 @@ +""" +TOOL POLICY MANAGEMENT + +All /tool management endpoints + +GET /v1/tool/list - List all discovered tools and their policies +GET /v1/tool/{tool_name} - Get a single tool's details +POST /v1/tool/policy - Update the call_policy for a tool +""" + +from typing import Optional + +from fastapi import APIRouter, Depends, HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.types.tool_management import ( + LiteLLM_ToolTableRow, + ToolCallPolicy, + ToolListResponse, + ToolPolicyUpdateRequest, + ToolPolicyUpdateResponse, +) + +router = APIRouter() + + +@router.get( + "/v1/tool/list", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=ToolListResponse, +) +async def list_tools( + call_policy: Optional[ToolCallPolicy] = None, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + List all auto-discovered tools and their call policies. + + Parameters: + - call_policy: Optional filter — one of "trusted", "untrusted", "dual_llm", "blocked" + """ + from litellm.proxy.db.tool_registry_writer import list_tools as db_list_tools + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + tools = await db_list_tools(prisma_client=prisma_client, call_policy=call_policy) + return ToolListResponse(tools=tools, total=len(tools)) + except Exception as e: + verbose_proxy_logger.exception("Error listing tools: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get( + "/v1/tool/{tool_name:path}", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_ToolTableRow, +) +async def get_tool( + tool_name: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Get details for a single tool. + + Parameters: + - tool_name: The tool name (supports namespaced names with slashes) + """ + from litellm.proxy.db.tool_registry_writer import get_tool as db_get_tool + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + tool = await db_get_tool(prisma_client=prisma_client, tool_name=tool_name) + if tool is None: + raise HTTPException( + status_code=404, detail=f"Tool '{tool_name}' not found" + ) + return tool + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error getting tool: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/v1/tool/policy", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=ToolPolicyUpdateResponse, +) +async def update_tool_policy( + data: ToolPolicyUpdateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Set the call policy for a tool. + + Parameters: + - tool_name: str - The tool to update + - call_policy: "trusted" | "untrusted" | "dual_llm" | "blocked" + + Setting a tool to "blocked" will cause the ToolPolicyGuardrail to remove + that tool_call from LLM responses before returning them to the client. + """ + from litellm.proxy.db.tool_registry_writer import ( + update_tool_policy as db_update_tool_policy, + ) + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + updated = await db_update_tool_policy( + prisma_client=prisma_client, + tool_name=data.tool_name, + call_policy=data.call_policy, + updated_by=user_api_key_dict.user_id, + ) + if updated is None: + raise HTTPException( + status_code=500, detail=f"Failed to update policy for tool '{data.tool_name}'" + ) + return ToolPolicyUpdateResponse( + tool_name=updated.tool_name, + call_policy=updated.call_policy, + updated=True, + ) + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error updating tool policy: %s", e) + raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/__init__.py b/litellm/proxy/management_endpoints/usage_endpoints/__init__.py new file mode 100644 index 00000000000..6e68dcd2a2e --- /dev/null +++ b/litellm/proxy/management_endpoints/usage_endpoints/__init__.py @@ -0,0 +1,9 @@ +""" +Usage endpoints package. + +Re-exports the router from endpoints module. +""" + +from litellm.proxy.management_endpoints.usage_endpoints.endpoints import ( # noqa: F401 + router, +) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py new file mode 100644 index 00000000000..f156be7d2cc --- /dev/null +++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py @@ -0,0 +1,578 @@ +""" +AI Usage Chat - uses LLM tool calling to answer questions about +usage/spend data by querying the aggregated daily activity endpoints. +""" + +import json +from datetime import date +from typing import Any, AsyncIterator, Callable, Dict, List, Literal, Optional + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.constants import DEFAULT_COMPETITOR_DISCOVERY_MODEL +from litellm.types.proxy.management_endpoints.common_daily_activity import ( + SpendAnalyticsPaginatedResponse, +) + +from typing_extensions import TypedDict + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +USAGE_AI_TEMPERATURE = 0.2 + +TABLE_DAILY_USER_SPEND = "litellm_dailyuserspend" +TABLE_DAILY_TEAM_SPEND = "litellm_dailyteamspend" +TABLE_DAILY_TAG_SPEND = "litellm_dailytagspend" + +ENTITY_FIELD_USER = "user_id" +ENTITY_FIELD_TEAM = "team_id" +ENTITY_FIELD_TAG = "tag" + +PAGINATED_PAGE_SIZE = 200 +MAX_CHAT_MESSAGES = 20 +TOP_N_MODELS = 15 +TOP_N_PROVIDERS = 10 +TOP_N_KEYS = 10 + +# --------------------------------------------------------------------------- +# Types +# --------------------------------------------------------------------------- + + +class SSEStatusEvent(TypedDict): + type: Literal["status"] + message: str + + +class SSEToolCallEvent(TypedDict, total=False): + type: Literal["tool_call"] + tool_name: str + tool_label: str + arguments: Dict[str, str] + status: Literal["running", "complete", "error"] + error: str + + +class SSEChunkEvent(TypedDict): + type: Literal["chunk"] + content: str + + +class SSEDoneEvent(TypedDict): + type: Literal["done"] + + +class SSEErrorEvent(TypedDict): + type: Literal["error"] + message: str + + +SSEEvent = ( + SSEStatusEvent | SSEToolCallEvent | SSEChunkEvent | SSEDoneEvent | SSEErrorEvent +) + + +class ToolHandler(TypedDict): + fetch: Callable[..., Any] + summarise: Callable[[Dict[str, Any]], str] + label: str + + +# --------------------------------------------------------------------------- +# Tool definitions (OpenAI function-calling schema) +# --------------------------------------------------------------------------- + +_DATE_PARAMS = { + "start_date": {"type": "string", "description": "Start date in YYYY-MM-DD format"}, + "end_date": {"type": "string", "description": "End date in YYYY-MM-DD format"}, +} + +_TOOL_USAGE = { + "type": "function", + "function": { + "name": "get_usage_data", + "description": ( + "Fetch aggregated global usage/spend data. Returns daily spend, " + "token counts, request counts, and breakdowns by model, provider, " + "and API key. Use for overall spend, top models, top providers." + ), + "parameters": { + "type": "object", + "properties": { + **_DATE_PARAMS, + "user_id": { + "type": "string", + "description": "Optional user ID filter. Omit for global view.", + }, + }, + "required": ["start_date", "end_date"], + }, + }, +} + +_TOOL_TEAM = { + "type": "function", + "function": { + "name": "get_team_usage_data", + "description": ( + "Fetch usage/spend data broken down by team. Use for questions " + "like 'which team spends the most' or 'show me team X usage'." + ), + "parameters": { + "type": "object", + "properties": { + **_DATE_PARAMS, + "team_ids": { + "type": "string", + "description": "Optional comma-separated team IDs. Omit for all teams.", + }, + }, + "required": ["start_date", "end_date"], + }, + }, +} + +_TOOL_TAG = { + "type": "function", + "function": { + "name": "get_tag_usage_data", + "description": ( + "Fetch usage/spend data broken down by tag. Tags are labels " + "attached to requests (features, environments, credentials)." + ), + "parameters": { + "type": "object", + "properties": { + **_DATE_PARAMS, + "tags": { + "type": "string", + "description": "Optional comma-separated tag names. Omit for all tags.", + }, + }, + "required": ["start_date", "end_date"], + }, + }, +} + +TOOLS_BASE = [_TOOL_USAGE] +TOOLS_ADMIN = [_TOOL_USAGE, _TOOL_TEAM, _TOOL_TAG] + + +def get_tools_for_role(is_admin: bool) -> List[Dict[str, Any]]: + """Return the tool list appropriate for the user's role.""" + return TOOLS_ADMIN if is_admin else TOOLS_BASE + + +_SYSTEM_PROMPT_BASE = ( + "You are an AI assistant embedded in the LiteLLM Usage dashboard. " + "You help users understand their LLM API spend and usage data.\n\n" + "ALWAYS call the appropriate tool(s) first to fetch data before answering. " + "You may call multiple tools if the question spans different dimensions.\n\n" + "Guidelines:\n" + "- Be concise and specific. Use exact numbers from the data.\n" + "- Format costs as dollar amounts (e.g. $12.34).\n" + "- When comparing entities, show a ranked list.\n" + "- If data is empty or no results found, say so clearly.\n" + "- Do not hallucinate data — only use what the tools return.\n" + "- Today's date will be provided below. Use it to interpret relative dates " + "like 'this week', 'this month', 'last 7 days', etc." +) + +_TOOL_DESCRIPTIONS_ADMIN = ( + "You have access to these tools:\n" + "- `get_usage_data`: Global/user-level usage (spend, models, providers, API keys)\n" + "- `get_team_usage_data`: Team-level usage breakdown\n" + "- `get_tag_usage_data`: Tag-level usage breakdown\n\n" +) + +_TOOL_DESCRIPTIONS_BASE = ( + "You have access to this tool:\n" + "- `get_usage_data`: Your usage data (spend, models, providers, API keys)\n\n" +) + + +def _build_system_prompt(is_admin: bool) -> str: + """Build role-appropriate system prompt with today's date.""" + tool_desc = _TOOL_DESCRIPTIONS_ADMIN if is_admin else _TOOL_DESCRIPTIONS_BASE + return ( + f"{_SYSTEM_PROMPT_BASE}\n\n{tool_desc}" + f"Today's date: {date.today().isoformat()}" + ) + + +# keep a public reference for test assertions +SYSTEM_PROMPT = _SYSTEM_PROMPT_BASE + +# --------------------------------------------------------------------------- +# Data fetchers +# --------------------------------------------------------------------------- + + +def _parse_csv_ids(raw: Optional[str]) -> Optional[List[str]]: + if not raw: + return None + return [t.strip() for t in raw.split(",") if t.strip()] + + +async def _query_activity( + table_name: str, + entity_id_field: str, + entity_id: Optional[Any], + start_date: str, + end_date: str, + *, + use_aggregated: bool = False, +) -> SpendAnalyticsPaginatedResponse: + """Shared helper that calls the daily activity query layer.""" + from litellm.proxy.management_endpoints.common_daily_activity import ( + get_daily_activity, + get_daily_activity_aggregated, + ) + from litellm.proxy.proxy_server import prisma_client + + if use_aggregated: + return await get_daily_activity_aggregated( + prisma_client=prisma_client, + table_name=table_name, + entity_id_field=entity_id_field, + entity_id=entity_id, + entity_metadata_field=None, + start_date=start_date, + end_date=end_date, + model=None, + api_key=None, + ) + return await get_daily_activity( + prisma_client=prisma_client, + table_name=table_name, + entity_id_field=entity_id_field, + entity_id=entity_id, + entity_metadata_field=None, + start_date=start_date, + end_date=end_date, + model=None, + api_key=None, + page=1, + page_size=PAGINATED_PAGE_SIZE, + ) + + +async def _fetch_usage_data( + start_date: str, end_date: str, user_id: Optional[str] = None +) -> Dict[str, Any]: + resp = await _query_activity( + TABLE_DAILY_USER_SPEND, + ENTITY_FIELD_USER, + user_id, + start_date, + end_date, + use_aggregated=True, + ) + return resp.model_dump(mode="json") + + +async def _fetch_team_usage_data( + start_date: str, end_date: str, team_ids: Optional[str] = None +) -> Dict[str, Any]: + resp = await _query_activity( + TABLE_DAILY_TEAM_SPEND, + ENTITY_FIELD_TEAM, + _parse_csv_ids(team_ids), + start_date, + end_date, + ) + return resp.model_dump(mode="json") + + +async def _fetch_tag_usage_data( + start_date: str, end_date: str, tags: Optional[str] = None +) -> Dict[str, Any]: + resp = await _query_activity( + TABLE_DAILY_TAG_SPEND, + ENTITY_FIELD_TAG, + _parse_csv_ids(tags), + start_date, + end_date, + ) + return resp.model_dump(mode="json") + + +# --------------------------------------------------------------------------- +# Summarisers — convert raw JSON to concise text the LLM can reason over +# --------------------------------------------------------------------------- + + +def _accumulate_breakdown( + results: List[Dict[str, Any]], dimension: str, fields: List[str] +) -> Dict[str, Dict[str, float]]: + """Aggregate a single breakdown dimension across days.""" + totals: Dict[str, Dict[str, float]] = {} + for day in results: + for key, entry in day.get("breakdown", {}).get(dimension, {}).items(): + if key not in totals: + totals[key] = {f: 0.0 for f in fields} + m = entry.get("metrics", {}) + for f in fields: + totals[key][f] += m.get(f, 0) + return totals + + +def _ranked_lines( + totals: Dict[str, Dict[str, float]], + fmt: Callable[[str, Dict[str, float]], str], + limit: int, +) -> List[str]: + """Sort by spend descending, format each entry, and truncate.""" + return [ + fmt(name, vals) + for name, vals in sorted(totals.items(), key=lambda x: -x[1].get("spend", 0))[ + :limit + ] + ] + + +def _summarise_usage_data(data: Dict[str, Any]) -> str: + meta = data.get("metadata", {}) + results = data.get("results", []) + + header = ( + f"Total Spend: ${meta.get('total_spend', 0):.4f}\n" + f"Total Requests: {meta.get('total_api_requests', 0)}\n" + f"Successful: {meta.get('total_successful_requests', 0)} | " + f"Failed: {meta.get('total_failed_requests', 0)}\n" + f"Total Tokens: {meta.get('total_tokens', 0)}" + ) + + models = _accumulate_breakdown( + results, "models", ["spend", "api_requests", "total_tokens"] + ) + providers = _accumulate_breakdown(results, "providers", ["spend", "api_requests"]) + + model_lines = _ranked_lines( + models, + lambda n, d: f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs, {int(d['total_tokens'])} tokens)", + TOP_N_MODELS, + ) + provider_lines = _ranked_lines( + providers, + lambda n, d: f" - {n}: ${d['spend']:.4f} ({int(d['api_requests'])} reqs)", + TOP_N_PROVIDERS, + ) + + sections = [header, ""] + sections += ["Top Models by Spend:"] + (model_lines or [" (no data)"]) + [""] + sections += ["Top Providers by Spend:"] + (provider_lines or [" (no data)"]) + return "\n".join(sections) + + +def _summarise_entity_data(data: Dict[str, Any], entity_label: str) -> str: + """Summarise team/tag entity usage data.""" + results = data.get("results", []) + if not results: + return f"No {entity_label} usage data found for the given date range." + + totals: Dict[str, Dict[str, Any]] = {} + for day in results: + for eid, entry in day.get("breakdown", {}).get("entities", {}).items(): + if eid not in totals: + alias = entry.get("metadata", {}).get("alias", eid) + totals[eid] = {"alias": alias, "spend": 0.0, "requests": 0, "tokens": 0} + m = entry.get("metrics", {}) + totals[eid]["spend"] += m.get("spend", 0) + totals[eid]["requests"] += m.get("api_requests", 0) + totals[eid]["tokens"] += m.get("total_tokens", 0) + + lines = [f"{entity_label} Usage ({len(totals)} {entity_label.lower()}s):", ""] + for eid, d in sorted(totals.items(), key=lambda x: -x[1]["spend"]): + label = d["alias"] if d["alias"] != eid else eid + lines.append( + f"- {label} (ID: {eid}): ${d['spend']:.4f} | " + f"{int(d['requests'])} reqs | {int(d['tokens'])} tokens" + ) + return "\n".join(lines) + + +# --------------------------------------------------------------------------- +# Tool dispatch registry +# --------------------------------------------------------------------------- + +TOOL_HANDLERS: Dict[str, ToolHandler] = { + "get_usage_data": ToolHandler( + fetch=_fetch_usage_data, + summarise=_summarise_usage_data, + label="global usage data", + ), + "get_team_usage_data": ToolHandler( + fetch=_fetch_team_usage_data, + summarise=lambda data: _summarise_entity_data(data, "Team"), + label="team usage data", + ), + "get_tag_usage_data": ToolHandler( + fetch=_fetch_tag_usage_data, + summarise=lambda data: _summarise_entity_data(data, "Tag"), + label="tag usage data", + ), +} + + +# --------------------------------------------------------------------------- +# SSE streaming +# --------------------------------------------------------------------------- + + +def _sse(event: SSEEvent) -> str: + return f"data: {json.dumps(event)}\n\n" + + +def _resolve_fetch_kwargs( + fn_name: str, + fn_args: Dict[str, str], + user_id: Optional[str], + is_admin: bool, +) -> Dict[str, Any]: + """Build keyword arguments for a tool's fetch function.""" + start_date = fn_args.get("start_date", "") + end_date = fn_args.get("end_date", "") + if not start_date or not end_date: + raise ValueError("Missing required start_date or end_date from tool arguments") + kwargs: Dict[str, Any] = {"start_date": start_date, "end_date": end_date} + if fn_name == "get_usage_data": + if not is_admin: + kwargs["user_id"] = user_id + elif fn_args.get("user_id"): + kwargs["user_id"] = fn_args["user_id"] + elif fn_name == "get_team_usage_data" and fn_args.get("team_ids"): + kwargs["team_ids"] = fn_args["team_ids"] + elif fn_name == "get_tag_usage_data" and fn_args.get("tags"): + kwargs["tags"] = fn_args["tags"] + return kwargs + + +async def _execute_tool_call( + handler: ToolHandler, + fn_name: str, + fn_args: Dict[str, str], + user_id: Optional[str], + is_admin: bool, +) -> str: + """Run a single tool and return the summarised result text.""" + kwargs = _resolve_fetch_kwargs(fn_name, fn_args, user_id, is_admin) + raw_data = await handler["fetch"](**kwargs) + return handler["summarise"](raw_data) + + +async def _process_tool_call( + tc: Any, + chat_messages: List[Dict[str, Any]], + user_id: Optional[str], + is_admin: bool, +) -> AsyncIterator[str]: + """Execute a single tool call, yielding SSE events for status.""" + fn_name = tc.function.name + fn_args = json.loads(tc.function.arguments) + + allowed_names = {t["function"]["name"] for t in get_tools_for_role(is_admin)} + handler = TOOL_HANDLERS.get(fn_name) + + if fn_name not in allowed_names or not handler: + chat_messages.append( + { + "role": "tool", + "tool_call_id": tc.id, + "content": f"Tool not available: {fn_name}", + } + ) + return + + tool_event_base = { + "type": "tool_call", + "tool_name": fn_name, + "tool_label": handler["label"], + "arguments": fn_args, + } + yield _sse({**tool_event_base, "status": "running"}) + + try: + tool_result = await _execute_tool_call( + handler, fn_name, fn_args, user_id, is_admin + ) + yield _sse({**tool_event_base, "status": "complete"}) + except Exception as e: + verbose_proxy_logger.error("Tool %s failed: %s", fn_name, e) + tool_result = f"Error fetching {handler['label']}. Please try again." + yield _sse({**tool_event_base, "status": "error"}) + + chat_messages.append( + {"role": "tool", "tool_call_id": tc.id, "content": tool_result} + ) + + +async def _stream_final_response( + model: str, chat_messages: List[Dict[str, Any]] +) -> AsyncIterator[str]: + """Stream the final LLM response after tool results are appended.""" + yield _sse({"type": "status", "message": "Analyzing results..."}) + + response = await litellm.acompletion( + model=model, + messages=chat_messages, + stream=True, + temperature=USAGE_AI_TEMPERATURE, + ) + async for chunk in response: + delta = chunk.choices[0].delta.content + if delta: + yield _sse({"type": "chunk", "content": delta}) + + +async def stream_usage_ai_chat( + messages: List[Dict[str, str]], + model: Optional[str] = None, + user_id: Optional[str] = None, + is_admin: bool = False, +) -> AsyncIterator[str]: + """Stream SSE events: status → tool_call → chunk → done.""" + resolved_model = (model or "").strip() or DEFAULT_COMPETITOR_DISCOVERY_MODEL + truncated = ( + messages[-MAX_CHAT_MESSAGES:] if len(messages) > MAX_CHAT_MESSAGES else messages + ) + chat_messages: List[Dict[str, Any]] = [ + {"role": "system", "content": _build_system_prompt(is_admin)}, + *truncated, + ] + + try: + yield _sse({"type": "status", "message": "Thinking..."}) + tools = get_tools_for_role(is_admin) + response = await litellm.acompletion( + model=resolved_model, + messages=chat_messages, + tools=tools, + temperature=USAGE_AI_TEMPERATURE, + ) + choice = response.choices[0] # type: ignore + + if not choice.message.tool_calls: + if choice.message.content: + yield _sse({"type": "chunk", "content": choice.message.content}) + yield _sse({"type": "done"}) + return + + chat_messages.append(choice.message.model_dump()) + for tc in choice.message.tool_calls: + async for event in _process_tool_call(tc, chat_messages, user_id, is_admin): + yield event + async for event in _stream_final_response(resolved_model, chat_messages): + yield event + yield _sse({"type": "done"}) + + except Exception as e: + verbose_proxy_logger.error("AI usage chat failed: %s", e) + yield _sse( + { + "type": "error", + "message": "An internal error occurred. Please try again.", + } + ) diff --git a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py new file mode 100644 index 00000000000..0dbe518afb7 --- /dev/null +++ b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py @@ -0,0 +1,65 @@ +""" +USAGE AI CHAT ENDPOINTS + +/usage/ai/chat - Stream AI chat responses about usage data +""" + +from typing import List, Literal, Optional + +from fastapi import APIRouter, Depends, Request +from fastapi.responses import StreamingResponse +from pydantic import BaseModel, Field + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +router = APIRouter() + + +class ChatMessage(BaseModel): + role: Literal["user", "assistant"] + content: str + + +class UsageAIChatRequest(BaseModel): + messages: List[ChatMessage] = Field( + ..., description="Chat messages (user/assistant history)" + ) + model: Optional[str] = Field(default=None, description="Model to use for AI chat") + + +@router.post( + "/usage/ai/chat", + tags=["Budget & Spend Tracking"], + dependencies=[Depends(user_api_key_auth)], +) +async def usage_ai_chat( + data: UsageAIChatRequest, + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + AI chat about usage data. Streams SSE events with the AI response. + The AI agent has access to tools that query aggregated daily activity data. + """ + from litellm.proxy.management_endpoints.common_utils import ( + _user_has_admin_view, + ) + from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( + stream_usage_ai_chat, + ) + + is_admin = _user_has_admin_view(user_api_key_dict) + user_id = user_api_key_dict.user_id + messages = [{"role": m.role, "content": m.content} for m in data.messages] + + return StreamingResponse( + stream_usage_ai_chat( + messages=messages, + model=data.model, + user_id=user_id, + is_admin=is_admin, + ), + media_type="text/event-stream", + headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"}, + ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 81144ad9f31..028edea8c3f 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -28,6 +28,7 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, + _safe_get_request_headers, _safe_set_request_parsed_body, get_form_data, get_request_body, @@ -60,7 +61,7 @@ def create_request_copy(request: Request): return { "method": request.method, "url": str(request.url), - "headers": dict(request.headers), + "headers": _safe_get_request_headers(request).copy(), "cookies": request.cookies, "query_params": dict(request.query_params), } @@ -329,7 +330,7 @@ async def vllm_proxy_route( method=request.method, endpoint=endpoint, request_query_params=request.query_params, - request_headers=dict(request.headers), + request_headers=_safe_get_request_headers(request), stream=request_body.get("stream", False), content=None, data=None, @@ -1307,7 +1308,7 @@ async def azure_proxy_route( method=request.method, endpoint=endpoint, request_query_params=request.query_params, - request_headers=dict(request.headers), + request_headers=_safe_get_request_headers(request), stream=request_body.get("stream", False), content=None, data=None, @@ -1505,7 +1506,7 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict: Returns: dict: Headers dictionary with only allowed headers """ - incoming_headers = dict(request.headers) or {} + incoming_headers = _safe_get_request_headers(request) headers = {} for header_name in ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS: if header_name in incoming_headers: @@ -1621,7 +1622,7 @@ async def _prepare_vertex_auth_headers( if ( vertex_credentials is None or vertex_credentials.vertex_project is None ) and router_credentials is None: - headers = dict(request.headers) or {} + headers = _safe_get_request_headers(request).copy() headers_passed_through = True verbose_proxy_logger.debug( "default_vertex_config not set, incoming request headers %s", headers diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index d1c4b8b3403..356807415de 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -50,7 +50,10 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing -from litellm.proxy.common_utils.http_parsing_utils import _read_request_body +from litellm.proxy.common_utils.http_parsing_utils import ( + _read_request_body, + _safe_get_request_headers, +) from litellm.proxy.utils import get_server_root_path from litellm.secret_managers.main import get_secret_str from litellm.types.llms.custom_http import httpxSpecialProvider @@ -649,7 +652,7 @@ async def pass_through_request( # noqa: PLR0915 url = httpx.URL(target) headers = custom_headers headers = HttpPassThroughEndpointHelpers.forward_headers_from_request( - request_headers=dict(request.headers), + request_headers=_safe_get_request_headers(request).copy(), headers=headers, forward_headers=forward_headers, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1983a601e13..82cfd455be6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -294,6 +294,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import ( from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, + _safe_get_request_headers, check_file_size_under_limit, get_form_data, ) @@ -391,6 +392,7 @@ from litellm.proxy.management_endpoints.organization_endpoints import ( router as organization_router, ) from litellm.proxy.management_endpoints.policy_endpoints import router as policy_router +from litellm.proxy.management_endpoints.usage_endpoints import router as usage_ai_router from litellm.proxy.management_endpoints.project_endpoints import ( router as project_router, ) @@ -409,6 +411,9 @@ from litellm.proxy.management_endpoints.team_endpoints import ( update_team, validate_membership, ) +from litellm.proxy.management_endpoints.tool_management_endpoints import ( + router as tool_management_router, +) from litellm.proxy.management_endpoints.ui_sso import ( get_disabled_non_admin_personal_key_creation, ) @@ -1211,9 +1216,7 @@ try: # Case 2: Runtime UI exists and is ready if has_content and is_pre_restructured: - verbose_proxy_logger.info( - f"Using pre-restructured UI at {runtime_ui_path}" - ) + verbose_proxy_logger.info(f"Using pre-restructured UI at {runtime_ui_path}") ui_path = runtime_ui_path # Case 3: Runtime UI exists but needs restructuring @@ -2989,6 +2992,10 @@ class ProxyConfig: if master_key is not None and isinstance(master_key, str): litellm_master_key_hash = hash_token(master_key) + else: + verbose_proxy_logger.critical( + "LITELLM_MASTER_KEY is not set! All requests will be treated as INTERNAL_USER with no admin access. Set LITELLM_MASTER_KEY for production use." + ) ### USER API KEY CACHE IN-MEMORY TTL ### user_api_key_cache_ttl = general_settings.get( "user_api_key_cache_ttl", None @@ -3791,6 +3798,7 @@ class ProxyConfig: parsed = value elif isinstance(value, str): import json + try: parsed = yaml.safe_load(value) except (yaml.YAMLError, json.JSONDecodeError): @@ -4376,10 +4384,12 @@ class ProxyConfig: if self._should_load_db_object(object_type="model_cost_map"): await self._check_and_reload_model_cost_map(prisma_client=prisma_client) - + if self._should_load_db_object(object_type="anthropic_beta_headers"): - await self._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client) - + await self._check_and_reload_anthropic_beta_headers( + prisma_client=prisma_client + ) + if self._should_load_db_object(object_type="sso_settings"): await self._init_sso_settings_in_db(prisma_client=prisma_client) if self._should_load_db_object(object_type="cache_settings"): @@ -4609,7 +4619,9 @@ class ProxyConfig: f"Error in _check_and_reload_model_cost_map: {str(e)}" ) - async def _check_and_reload_anthropic_beta_headers(self, prisma_client: PrismaClient): + async def _check_and_reload_anthropic_beta_headers( + self, prisma_client: PrismaClient + ): """ Check if anthropic beta headers config needs to be reloaded based on database configuration. This function runs every 10 seconds as part of _init_non_llm_objects_in_db. @@ -4700,7 +4712,11 @@ class ProxyConfig: ) # Count providers in config - provider_count = sum(1 for k in new_config.keys() if k != "provider_aliases" and k != "description") + provider_count = sum( + 1 + for k in new_config.keys() + if k != "provider_aliases" and k != "description" + ) verbose_proxy_logger.info( f"Anthropic beta headers config reloaded successfully. Providers: {provider_count}" ) @@ -5682,8 +5698,7 @@ class ProxyStartupEvent: ): _db_val = _db_gs_record.param_value.get("store_model_in_db") if _db_val is True or ( - isinstance(_db_val, str) - and _db_val.lower() == "true" + isinstance(_db_val, str) and _db_val.lower() == "true" ): store_model_in_db = True verbose_proxy_logger.info( @@ -6150,6 +6165,7 @@ class ProxyStartupEvent: "Pyroscope profiling will not run. Install with: pip install pyroscope-io" ) + #### API ENDPOINTS #### @router.get( "/v1/models", dependencies=[Depends(user_api_key_auth)], tags=["model management"] @@ -10448,7 +10464,7 @@ async def async_queue_request( data["proxy_server_request"] = { "url": str(request.url), "method": request.method, - "headers": dict(request.headers), + "headers": _safe_get_request_headers(request).copy(), "body": copy.copy(data), # use copy instead of deepcopy } @@ -10469,7 +10485,7 @@ async def async_queue_request( data["metadata"] = {} data["metadata"]["user_api_key"] = user_api_key_dict.api_key data["metadata"]["user_api_key_metadata"] = user_api_key_dict.metadata - _headers = dict(request.headers) + _headers = _safe_get_request_headers(request).copy() _headers.pop( "authorization", None ) # do not store the original `sk-..` api key in the db @@ -10988,18 +11004,14 @@ async def get_favicon(): from fastapi.responses import Response current_dir = os.path.dirname(os.path.abspath(__file__)) - default_favicon = os.path.join( - current_dir, "_experimental", "out", "favicon.ico" - ) + default_favicon = os.path.join(current_dir, "_experimental", "out", "favicon.ico") favicon_url = os.getenv("LITELLM_FAVICON_URL", "") if not favicon_url: if os.path.exists(default_favicon): return FileResponse(default_favicon, media_type="image/x-icon") - raise HTTPException( - status_code=404, detail="Default favicon not found" - ) + raise HTTPException(status_code=404, detail="Default favicon not found") if favicon_url.startswith(("http://", "https://")): try: @@ -11014,9 +11026,7 @@ async def get_favicon(): ) response = await async_client.get(favicon_url) if response.status_code == 200: - content_type = response.headers.get( - "content-type", "image/x-icon" - ) + content_type = response.headers.get("content-type", "image/x-icon") return Response( content=response.content, media_type=content_type, @@ -11028,12 +11038,8 @@ async def get_favicon(): response.status_code, ) if os.path.exists(default_favicon): - return FileResponse( - default_favicon, media_type="image/x-icon" - ) - raise HTTPException( - status_code=404, detail="Favicon not found" - ) + return FileResponse(default_favicon, media_type="image/x-icon") + raise HTTPException(status_code=404, detail="Favicon not found") except HTTPException: raise except Exception as e: @@ -11041,20 +11047,14 @@ async def get_favicon(): "Error downloading favicon from %s: %s", favicon_url, e ) if os.path.exists(default_favicon): - return FileResponse( - default_favicon, media_type="image/x-icon" - ) - raise HTTPException( - status_code=404, detail="Favicon not found" - ) + return FileResponse(default_favicon, media_type="image/x-icon") + raise HTTPException(status_code=404, detail="Favicon not found") else: if os.path.exists(favicon_url): return FileResponse(favicon_url, media_type="image/x-icon") if os.path.exists(default_favicon): return FileResponse(default_favicon, media_type="image/x-icon") - raise HTTPException( - status_code=404, detail="Favicon not found" - ) + raise HTTPException(status_code=404, detail="Favicon not found") #### INVITATION MANAGEMENT #### @@ -12540,7 +12540,9 @@ async def reload_anthropic_beta_headers( }, ) - provider_count = sum(1 for k in new_config.keys() if k not in ["provider_aliases", "description"]) + provider_count = sum( + 1 for k in new_config.keys() if k not in ["provider_aliases", "description"] + ) verbose_proxy_logger.info( f"Anthropic beta headers config reloaded successfully in current pod. Providers: {provider_count}" ) @@ -12552,7 +12554,9 @@ async def reload_anthropic_beta_headers( "timestamp": current_time.isoformat(), } except Exception as e: - verbose_proxy_logger.exception(f"Failed to reload anthropic beta headers: {str(e)}") + verbose_proxy_logger.exception( + f"Failed to reload anthropic beta headers: {str(e)}" + ) raise HTTPException( status_code=500, detail=f"Failed to reload anthropic beta headers: {str(e)}" ) @@ -12674,7 +12678,8 @@ async def cancel_anthropic_beta_headers_reload( f"Failed to cancel anthropic beta headers reload: {str(e)}" ) raise HTTPException( - status_code=500, detail=f"Failed to cancel anthropic beta headers reload: {str(e)}" + status_code=500, + detail=f"Failed to cancel anthropic beta headers reload: {str(e)}", ) @@ -12721,7 +12726,9 @@ async def get_anthropic_beta_headers_reload_status( ) if config_record is None or config_record.param_value is None: - verbose_proxy_logger.info("No anthropic beta headers reload configuration found") + verbose_proxy_logger.info( + "No anthropic beta headers reload configuration found" + ) return { "scheduled": False, "interval_hours": None, @@ -12747,7 +12754,9 @@ async def get_anthropic_beta_headers_reload_status( # Use pod's in-memory last reload time if last_anthropic_beta_headers_reload is not None: try: - last_reload_time = datetime.fromisoformat(last_anthropic_beta_headers_reload) + last_reload_time = datetime.fromisoformat( + last_anthropic_beta_headers_reload + ) time_since_last_reload = current_time - last_reload_time hours_since_last_reload = time_since_last_reload.total_seconds() / 3600 @@ -12868,6 +12877,7 @@ app.include_router(caching_router) app.include_router(analytics_router) app.include_router(guardrails_router) app.include_router(policy_router) +app.include_router(usage_ai_router) app.include_router(policy_crud_router) app.include_router(policy_resolve_router) app.include_router(search_tool_management_router) @@ -12881,6 +12891,7 @@ app.include_router(budget_management_router) app.include_router(model_management_router) app.include_router(model_access_group_management_router) app.include_router(tag_management_router) +app.include_router(tool_management_router) app.include_router(cost_tracking_settings_router) app.include_router(router_settings_router) app.include_router(fallback_management_router) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 50c0a55a875..155cea12ca4 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -64,6 +64,8 @@ model LiteLLM_AgentsTable { litellm_params Json? agent_card_params Json agent_access_groups String[] @default([]) + object_permission_id String? + object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@ -264,6 +266,7 @@ model LiteLLM_ObjectPermissionTable { organizations LiteLLM_OrganizationTable[] users LiteLLM_UserTable[] end_users LiteLLM_EndUserTable[] + agents_table LiteLLM_AgentsTable[] } // Holds the MCP server configuration @@ -314,6 +317,7 @@ model LiteLLM_VerificationToken { router_settings Json? @default("{}") user_id String? team_id String? + agent_id String? project_id String? permissions Json @default("{}") max_parallel_requests Int? @@ -1051,6 +1055,26 @@ model LiteLLM_PolicyAttachmentTable { updated_by String? } +// Global tool registry - auto-discovered from LLM responses; admins set call_policy here +model LiteLLM_ToolTable { + tool_id String @id @default(uuid()) + tool_name String @unique // e.g. "huggingface_remote-mcp__dynamic_space" + origin String? // MCP server name or "user_defined" + call_policy String @default("untrusted") // "trusted" | "untrusted" | "dual_llm" | "blocked" + call_count Int @default(0) // cumulative number of times this tool was seen + assignments Json? @default("{}") + key_hash String? // hash of the virtual key that first called this tool + team_id String? // team that first called this tool + key_alias String? // human-readable alias of the virtual key + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? + + @@index([call_policy]) + @@index([team_id]) +} + //Unified Access Groups table for storing unified access groups model LiteLLM_AccessGroupTable { access_group_id String @id @default(uuid()) diff --git a/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py b/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py index ce27c830f6f..627618387d5 100644 --- a/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py +++ b/litellm/proxy/vertex_ai_endpoints/langfuse_endpoints.py @@ -19,6 +19,7 @@ from fastapi import APIRouter, Request, Response import litellm from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers from litellm.proxy.litellm_pre_call_utils import _get_dynamic_logging_metadata from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( create_pass_through_route, @@ -32,7 +33,7 @@ def create_request_copy(request: Request): return { "method": request.method, "url": str(request.url), - "headers": dict(request.headers), + "headers": _safe_get_request_headers(request).copy(), "cookies": request.cookies, "query_params": dict(request.query_params), } diff --git a/litellm/router.py b/litellm/router.py index 46d35352c37..3a6c514989d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -6423,6 +6423,7 @@ class Router: self.model_list = [] self.model_id_to_deployment_index_map = {} # Reset the index self.model_name_to_deployment_indices = {} # Reset the model_name index + self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works @@ -6733,6 +6734,7 @@ class Router: """ idx = len(self.model_list) self.model_list.append(model) + self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() # Update model_id index for O(1) lookup @@ -6781,6 +6783,7 @@ class Router: if removal_idx is not None: self.model_list.pop(removal_idx) + self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() self._update_deployment_indices_after_removal( model_id=deployment_id, removal_idx=removal_idx @@ -6815,6 +6818,7 @@ class Router: if deployment_idx is not None: # Pop the item from the list first item = self.model_list.pop(deployment_idx) + self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() self._update_deployment_indices_after_removal( model_id=id, removal_idx=deployment_idx @@ -7576,6 +7580,7 @@ class Router: """ # First populate the model_list self.model_list = [] + self._invalidate_model_group_info_cache() self._invalidate_access_groups_cache() for _, model in enumerate(model_list): # Extract model_info from the model dict @@ -7920,6 +7925,13 @@ class Router: return returned_models + def _invalidate_model_group_info_cache(self) -> None: + """Invalidate the cached model group info. + + Call this whenever self.model_list is modified to ensure the cache is rebuilt. + """ + self._cached_get_model_group_info.cache_clear() + def _invalidate_access_groups_cache(self) -> None: """Invalidate the cached access groups. diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 4ab81f8fd57..15e8d1be930 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1125,6 +1125,7 @@ class ResponsesAPIOptionalRequestParams(TypedDict, total=False): prompt: Optional[PromptObject] max_tool_calls: Optional[int] prompt_cache_key: Optional[str] + prompt_cache_retention: Optional[str] stream_options: Optional[dict] top_logprobs: Optional[int] partial_images: Optional[ diff --git a/litellm/types/responses/main.py b/litellm/types/responses/main.py index bda53bae082..449a5ac49c1 100644 --- a/litellm/types/responses/main.py +++ b/litellm/types/responses/main.py @@ -6,7 +6,7 @@ from typing_extensions import Any, List, Optional, TypedDict from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject -Phase = Optional[Literal["commentary", "final_answer"]] # TODO: Once openai sdk has updated, we can remove this and use the openai sdk type +Phase = Optional[Literal["commentary", "final_answer"]] class GenericResponseOutputItemContentAnnotation(BaseLiteLLMOpenAIResponseObject): """Annotation for content in a message""" diff --git a/litellm/types/tool_management.py b/litellm/types/tool_management.py new file mode 100644 index 00000000000..8704ff27759 --- /dev/null +++ b/litellm/types/tool_management.py @@ -0,0 +1,42 @@ +""" +Pydantic models for Tool Policy management endpoints. +""" + +from datetime import datetime +from typing import Dict, List, Literal, Optional + +from pydantic import BaseModel + +ToolCallPolicy = Literal["trusted", "untrusted", "dual_llm", "blocked"] + + +class LiteLLM_ToolTableRow(BaseModel): + tool_id: str + tool_name: str + origin: Optional[str] = None + call_policy: ToolCallPolicy = "untrusted" + call_count: int = 0 + assignments: Optional[Dict] = None + key_hash: Optional[str] = None + team_id: Optional[str] = None + key_alias: Optional[str] = None + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_by: Optional[str] = None + + +class ToolListResponse(BaseModel): + tools: List[LiteLLM_ToolTableRow] + total: int + + +class ToolPolicyUpdateRequest(BaseModel): + tool_name: str + call_policy: ToolCallPolicy + + +class ToolPolicyUpdateResponse(BaseModel): + tool_name: str + call_policy: ToolCallPolicy + updated: bool diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 94c387dd22c..dda32d98383 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1,7 +1,17 @@ import json import time from enum import Enum -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Union +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + Mapping, + Optional, + Union, + get_args, +) from openai._models import BaseModel as OpenAIObject from openai.types.audio.transcription_create_params import ( @@ -1826,7 +1836,7 @@ class ModelResponse(ModelResponseBase): else: usage = usage elif stream is None or stream is False: - usage = Usage() + usage = None # avoid constructing throwaway Usage; set by convert_to_model_response_object if hidden_params: self._hidden_params = hidden_params @@ -3186,6 +3196,12 @@ OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: set[str] = { LlmProviders.HOSTED_VLLM.value, } +ListBatchesSupportedProvider = Literal["openai", "azure", "hosted_vllm", "vertex_ai"] + +LIST_BATCHES_SUPPORTED_PROVIDERS: frozenset[str] = frozenset( + get_args(ListBatchesSupportedProvider) +) + class SearchProviders(str, Enum): """ diff --git a/litellm/utils.py b/litellm/utils.py index 5046929257d..7ac828aefc1 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4644,11 +4644,12 @@ def add_provider_specific_params_to_optional_params( ) is False ): - extra_body = passed_params.pop("extra_body", {}) + extra_body = passed_params.pop("extra_body", None) or {} for k in passed_params.keys(): if k not in openai_params and passed_params[k] is not None: extra_body[k] = passed_params[k] - optional_params.setdefault("extra_body", {}) + if not isinstance(optional_params.get("extra_body"), dict): + optional_params["extra_body"] = {} initial_extra_body = { **optional_params["extra_body"], **extra_body, diff --git a/requirements.txt b/requirements.txt index dbde6ababc9..72fdbfafa17 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,7 +10,7 @@ cryptography==46.0.5 #GHSA-r6ph-v2qm-q3c2 anyio==4.8.0 # openai + http req. httpx==0.28.1 -openai==2.9.0 # openai req. +openai==2.24.0 # openai req. fastapi==0.120.1 # server dep starlette==0.49.1 # starlette fastapi dep backoff==2.2.1 # server dep diff --git a/tests/litellm/proxy/guardrails/test_custom_code_security.py b/tests/litellm/proxy/guardrails/test_custom_code_security.py new file mode 100644 index 00000000000..d855a4dde20 --- /dev/null +++ b/tests/litellm/proxy/guardrails/test_custom_code_security.py @@ -0,0 +1,94 @@ +import pytest +from litellm.proxy.guardrails.guardrail_hooks.custom_code.code_validator import ( + validate_custom_code, + CustomCodeValidationError, +) +from litellm.proxy.guardrails.guardrail_hooks.custom_code.custom_code_guardrail import ( + CustomCodeGuardrail, +) + +# Phase 4.1: Test forbidden pattern validation + + +def test_validate_custom_code_import_os(): + code = "import os\ndef apply_guardrail(inputs, req, ty):\n return allow()" + with pytest.raises(CustomCodeValidationError, match="import statements are not"): + validate_custom_code(code) + + +def test_validate_custom_code_from_subprocess(): + code = ( + "from subprocess import call\ndef apply_guardrail(i, r, t):\n return allow()" + ) + with pytest.raises( + CustomCodeValidationError, match="import statements are not allowed" + ): + validate_custom_code(code) + + +def test_validate_custom_code_exec(): + code = "def apply_guardrail(i, r, t):\n exec('print(1)')\n return allow()" + with pytest.raises(CustomCodeValidationError, match=r"exec\(\) is not allowed"): + validate_custom_code(code) + + +def test_validate_custom_code_builtins(): + code = "def apply_guardrail(i, r, t):\n print(__builtins__)\n return allow()" + with pytest.raises( + CustomCodeValidationError, match="__builtins__ access is not allowed" + ): + validate_custom_code(code) + + +def test_validate_custom_code_subclasses(): + code = "def apply_guardrail(i, r, t):\n print(''.__class__.__mro__[1].__subclasses__())\n return allow()" + with pytest.raises( + CustomCodeValidationError, match="__subclasses__ access is not allowed" + ): + validate_custom_code(code) + + +def test_validate_custom_code_clean(): + code = ( + "def apply_guardrail(inputs, request_data, input_type):\n return allow()\n" + ) + # Should not raise any exception + validate_custom_code(code) + + +# Phase 4.2: Test __builtins__ restriction in execution + + +def test_custom_code_compile_valid(): + code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()" + guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="test") + # if it doesn't fail, we successfully compiled + assert guardrail._compiled_function is not None + + +def test_custom_code_override_builtins(): + # Verify that even if pattern validation is bypassed, __builtins__ = {} blocks dangerous builtins. + # We test this by compiling safe code and verifying builtins are not accessible in the sandbox. + code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()" + guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="test") + # The compiled function's globals should have empty __builtins__ + fn_globals = guardrail._compiled_function.__globals__ + assert fn_globals.get("__builtins__") == {} + + +@pytest.mark.asyncio +async def test_custom_code_guardrail_apply(): + code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()" + guardrail = CustomCodeGuardrail(custom_code=code, guardrail_name="test") + from litellm.types.utils import GenericGuardrailAPIInputs + + result = await guardrail.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=["test"]), + request_data={}, + input_type="request", + ) + assert result["texts"][0] == "test" + + +# The RBAC endpoint tests are harder to write right here, but the core security +# validations are fully covered by the simple tests above. diff --git a/tests/litellm_utils_tests/test_health_check.py b/tests/litellm_utils_tests/test_health_check.py index 7ea4574bab9..b459c3cfc99 100644 --- a/tests/litellm_utils_tests/test_health_check.py +++ b/tests/litellm_utils_tests/test_health_check.py @@ -473,6 +473,50 @@ def test_update_litellm_params_for_health_check(): ) +@pytest.mark.asyncio +async def test_perform_health_check_filters_by_model_id(): + """ + When model_id is passed, only that deployment is checked (not all deployments + that share the same model name). + """ + from litellm.proxy.health_check import perform_health_check + + # Two deployments with same model_name but different ids + model_list = [ + { + "model_name": "gpt-4", + "model_info": {"id": "deployment-id-1"}, + "litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"}, + }, + { + "model_name": "gpt-4", + "model_info": {"id": "deployment-id-2"}, + "litellm_params": {"model": "gpt-4", "api_key": "fake-key-2"}, + }, + ] + + captured_list = [] + + async def mock_perform_health_check(m_list, details=True): + captured_list.append(m_list) + return [{"model": "gpt-4", "api_key": m_list[0]["litellm_params"]["api_key"]}], [] + + with patch( + "litellm.proxy.health_check._perform_health_check", + side_effect=mock_perform_health_check, + ): + healthy_endpoints, unhealthy_endpoints = await perform_health_check( + model_list=model_list, model_id="deployment-id-2", details=True + ) + + # Only one deployment (deployment-id-2) should have been passed to _perform_health_check + assert len(captured_list) == 1 + assert len(captured_list[0]) == 1 + assert (captured_list[0][0].get("model_info") or {}).get("id") == "deployment-id-2" + assert len(healthy_endpoints) == 1 + assert healthy_endpoints[0]["api_key"] == "fake-key-2" + + @pytest.mark.asyncio async def test_perform_health_check_with_health_check_model(): """ diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index 15d625b681f..a722aec7eb4 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -1453,3 +1453,47 @@ def test_convert_to_model_response_object_falsy_id_preserves_auto_generated(fals ) assert result.id == original_id assert result.id.startswith("chatcmpl-") + + +def test_convert_to_model_response_object_default_usage_overwritten(): + """ + Regression test: convert_to_model_response_object must properly set Usage + on a ModelResponse that only has the default Usage from ModelResponse.__init__() + (i.e. no extra litellm.Usage() set via setattr beforehand). + + This validates the optimization of removing the redundant + `setattr(model_response, "usage", litellm.Usage())` in completion(). + """ + mr = ModelResponse() + # usage is not set by default (optimization: avoid constructing throwaway Usage) + assert not hasattr(mr, "usage") + + response_object = { + "id": "chatcmpl-usage-test", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 15, + "completion_tokens": 7, + "total_tokens": 22, + }, + "model": "gpt-4o", + } + + result = convert_to_model_response_object( + model_response_object=mr, + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + assert isinstance(result, ModelResponse) + assert result.usage.prompt_tokens == 15 + assert result.usage.completion_tokens == 7 + assert result.usage.total_tokens == 22 diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 0cbba7b5cc3..d39edef793c 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -334,6 +334,81 @@ def test_chat_completion_forward_headers( pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}") +@pytest.mark.parametrize("forward_llm_auth_headers", [True, False]) +@mock_patch_acompletion() +def test_chat_completion_forward_llm_provider_auth_headers( + mock_acompletion, client_no_auth, forward_llm_auth_headers +): + """ + Test that LLM provider auth headers (x-api-key, x-goog-api-key) are forwarded + when forward_llm_provider_auth_headers=True. + + This allows clients to send their own LLM provider API keys through the proxy. + """ + try: + # Configure general settings + gs = getattr(litellm.proxy.proxy_server, "general_settings") + gs["forward_client_headers_to_llm_api"] = True + gs["forward_llm_provider_auth_headers"] = forward_llm_auth_headers + setattr(litellm.proxy.proxy_server, "general_settings", gs) + + # Test data + test_data = { + "model": "gpt-3.5-turbo", + "messages": [ + {"role": "user", "content": "hello"}, + ], + "max_tokens": 10, + } + + # Headers including LLM provider auth + request_headers = { + "Authorization": "Bearer sk-proxy-auth-123", # Proxy auth (should be stripped) + "x-api-key": "sk-ant-api03-test-anthropic-key", # Anthropic API key + "x-goog-api-key": "google-api-key-123", # Google API key + "X-Custom-Header": "custom-value", # Custom header (should be forwarded) + } + + # Make request + response = client_no_auth.post( + "/v1/chat/completions", json=test_data, headers=request_headers + ) + + assert response.status_code == 200 + + # Check forwarded headers + forwarded_headers = mock_acompletion.call_args.kwargs.get("headers", {}) + + if forward_llm_auth_headers: + # LLM provider auth headers should be forwarded + assert "x-api-key" in forwarded_headers + assert forwarded_headers["x-api-key"] == "sk-ant-api03-test-anthropic-key" + assert "x-goog-api-key" in forwarded_headers + assert forwarded_headers["x-goog-api-key"] == "google-api-key-123" + else: + # LLM provider auth headers should be stripped + assert "x-api-key" not in forwarded_headers + assert "x-goog-api-key" not in forwarded_headers + + # Custom headers should always be forwarded (when forward_client_headers_to_llm_api=True) + assert "x-custom-header" in forwarded_headers + assert forwarded_headers["x-custom-header"] == "custom-value" + + # Proxy Authorization should never be forwarded + assert "authorization" not in forwarded_headers + + print(f"✓ Test passed with forward_llm_provider_auth_headers={forward_llm_auth_headers}") + print(f" Forwarded headers: {list(forwarded_headers.keys())}") + + except Exception as e: + pytest.fail(f"Test failed with forward_llm_auth_headers={forward_llm_auth_headers}: {str(e)}") + finally: + # Clean up + gs = getattr(litellm.proxy.proxy_server, "general_settings") + gs.pop("forward_llm_provider_auth_headers", None) + setattr(litellm.proxy.proxy_server, "general_settings", gs) + + @mock_patch_acompletion() @pytest.mark.asyncio async def test_team_disable_guardrails(mock_acompletion, client_no_auth): diff --git a/tests/test_litellm/interactions/test_openapi_compliance.py b/tests/test_litellm/interactions/test_openapi_compliance.py index 5187f733a3c..5b2edcf1ee7 100644 --- a/tests/test_litellm/interactions/test_openapi_compliance.py +++ b/tests/test_litellm/interactions/test_openapi_compliance.py @@ -147,7 +147,8 @@ class TestResponseCompliance: """Verify status enum values match spec.""" schema = spec_dict["components"]["schemas"]["CreateModelInteractionParams"] status_prop = schema["properties"]["status"] - expected_statuses = ["UNSPECIFIED", "IN_PROGRESS", "REQUIRES_ACTION", "COMPLETED", "FAILED", "CANCELLED", "INCOMPLETE"] + # Google Interactions API uses lowercase status values (updated Feb 2026) + expected_statuses = ["in_progress", "requires_action", "completed", "failed", "cancelled", "incomplete"] assert status_prop["enum"] == expected_statuses print(f"✓ Status enum values: {expected_statuses}") diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index ebffb56446e..3729e67f0da 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -363,6 +363,87 @@ class TestProxyOAuthHeaderForwarding: assert "authorization" not in cleaned assert cleaned["content-type"] == "application/json" + def test_clean_headers_forwards_anthropic_api_key_when_enabled(self): + """clean_headers should preserve x-api-key when forward_llm_provider_auth_headers=True.""" + from starlette.datastructures import Headers + + from litellm.proxy.litellm_pre_call_utils import clean_headers + + raw_headers = Headers( + raw=[ + (b"authorization", b"Bearer sk-proxy-auth"), + (b"x-api-key", b"sk-ant-api03-test-key"), + (b"content-type", b"application/json"), + ] + ) + cleaned = clean_headers(raw_headers, forward_llm_provider_auth_headers=True) + + # x-api-key should be preserved when flag is True + assert "x-api-key" in cleaned + assert cleaned["x-api-key"] == "sk-ant-api03-test-key" + # Authorization (proxy auth) should still be stripped + assert "authorization" not in cleaned + assert cleaned["content-type"] == "application/json" + + def test_clean_headers_strips_anthropic_api_key_when_disabled(self): + """clean_headers should strip x-api-key when forward_llm_provider_auth_headers=False (default).""" + from starlette.datastructures import Headers + + from litellm.proxy.litellm_pre_call_utils import clean_headers + + raw_headers = Headers( + raw=[ + (b"x-api-key", b"sk-ant-api03-test-key"), + (b"content-type", b"application/json"), + ] + ) + cleaned = clean_headers(raw_headers, forward_llm_provider_auth_headers=False) + + # x-api-key should be stripped by default + assert "x-api-key" not in cleaned + assert cleaned["content-type"] == "application/json" + + def test_clean_headers_forwards_google_api_key_when_enabled(self): + """clean_headers should preserve x-goog-api-key when forward_llm_provider_auth_headers=True.""" + from starlette.datastructures import Headers + + from litellm.proxy.litellm_pre_call_utils import clean_headers + + raw_headers = Headers( + raw=[ + (b"x-goog-api-key", b"google-api-key-123"), + (b"content-type", b"application/json"), + ] + ) + cleaned = clean_headers(raw_headers, forward_llm_provider_auth_headers=True) + + assert "x-goog-api-key" in cleaned + assert cleaned["x-goog-api-key"] == "google-api-key-123" + assert cleaned["content-type"] == "application/json" + + def test_clean_headers_preserves_oauth_regardless_of_forward_flag(self): + """clean_headers should always preserve OAuth tokens regardless of forward_llm_provider_auth_headers.""" + from starlette.datastructures import Headers + + from litellm.proxy.litellm_pre_call_utils import clean_headers + + raw_headers = Headers( + raw=[ + (b"authorization", f"Bearer {FAKE_OAUTH_TOKEN}".encode()), + (b"content-type", b"application/json"), + ] + ) + + # Should preserve OAuth even with flag=False + cleaned_without_flag = clean_headers(raw_headers, forward_llm_provider_auth_headers=False) + assert "authorization" in cleaned_without_flag + assert cleaned_without_flag["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + + # Should also preserve OAuth with flag=True + cleaned_with_flag = clean_headers(raw_headers, forward_llm_provider_auth_headers=True) + assert "authorization" in cleaned_with_flag + assert cleaned_with_flag["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + def test_add_provider_specific_headers_forwards_oauth(self): """add_provider_specific_headers_to_request should forward OAuth Authorization as a ProviderSpecificHeader scoped to Anthropic-compatible providers.""" diff --git a/tests/test_litellm/llms/hosted_vllm/responses/test_hosted_vllm_responses.py b/tests/test_litellm/llms/hosted_vllm/responses/test_hosted_vllm_responses.py new file mode 100644 index 00000000000..22effbd37f1 --- /dev/null +++ b/tests/test_litellm/llms/hosted_vllm/responses/test_hosted_vllm_responses.py @@ -0,0 +1,103 @@ +""" +Tests for hosted_vllm responses API support. + +Regression test for: https://github.com/BerriAI/litellm/issues +Bug: client.responses.create() raised TypeError: 'NoneType' object is not a mapping +when extra_body=None was passed through the responses→completion pipeline for +hosted_vllm (and any OpenAI-compatible provider using add_provider_specific_params_to_optional_params). +""" + +import json +import os +import sys +from unittest.mock import MagicMock, patch + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +import litellm + + +def _make_mock_chat_completion_response(content: str = "Hello! I'm doing well.") -> dict: + return { + "id": "chatcmpl-test123", + "object": "chat.completion", + "created": 1234567890, + "model": "Qwen/Qwen3-8B", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": content}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, + } + + +def _make_mock_http_client(response_body: dict) -> MagicMock: + mock_client = MagicMock() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.json.return_value = response_body + mock_response.text = json.dumps(response_body) + mock_client.post.return_value = mock_response + return mock_client + + +def test_hosted_vllm_responses_create_with_string_input(): + """ + Regression test: responses.create() with string input must not raise + TypeError: 'NoneType' object is not a mapping. + + Root cause: extra_body=None was passed explicitly through the + responses→completion pipeline. In add_provider_specific_params_to_optional_params(), + passed_params.pop("extra_body", {}) returned None (key existed with value None), + and **None raised TypeError at dict unpacking. + + Fix: normalize None to {} for both extra_body and optional_params["extra_body"]. + """ + mock_client = _make_mock_http_client( + _make_mock_chat_completion_response("I'm doing well, thanks!") + ) + + with patch( + "litellm.llms.custom_httpx.llm_http_handler._get_httpx_client", + return_value=mock_client, + ): + response = litellm.responses( + model="hosted_vllm/Qwen/Qwen3-8B", + input="Hello, how are you?", + api_base="https://test-vllm.example.com/v1", + api_key="test-key", + ) + + from litellm.types.llms.openai import ResponsesAPIResponse + + assert response is not None + assert isinstance(response, ResponsesAPIResponse) + assert len(response.output) > 0 + output_message = response.output[0] + assert output_message.role == "assistant" # type: ignore[union-attr] + assert len(output_message.content) > 0 # type: ignore[union-attr] + assert "well" in output_message.content[0].text # type: ignore[union-attr] + + +def test_hosted_vllm_responses_create_with_explicit_none_extra_body(): + """ + Directly verify the fix in add_provider_specific_params_to_optional_params: + extra_body=None must not crash when building optional_params. + """ + from litellm.utils import get_optional_params + + # This should not raise TypeError: 'NoneType' object is not a mapping + optional_params = get_optional_params( + model="Qwen/Qwen3-8B", + custom_llm_provider="hosted_vllm", + extra_body=None, + ) + + # extra_body=None should be normalized to an empty dict (or absent) + assert optional_params.get("extra_body") is not None or "extra_body" not in optional_params diff --git a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py index 37959f74086..39ff0a4f4d8 100644 --- a/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py +++ b/tests/test_litellm/llms/openai/chat/test_openai_gpt_transformation.py @@ -207,10 +207,10 @@ class TestOpenAIChatCompletionStreamingHandler: def test_chunk_parser_maps_reasoning_to_reasoning_content(self): """ Test that chunk_parser maps 'reasoning' field to 'reasoning_content'. - + Some OpenAI-compatible providers (e.g., GLM-5, hosted_vllm) return delta.reasoning, but LiteLLM expects delta.reasoning_content. - + Regression test for: Streaming responses with delta.reasoning field coming back empty when using openai/ or hosted_vllm/ providers. """ @@ -293,3 +293,34 @@ class TestPromptCacheKeyIntegration: prompt_cache_key="test-cache-key-123", ) assert optional_params.get("prompt_cache_key") == "test-cache-key-123" + + +class TestPromptCacheParams: + """Tests for prompt_cache_key and prompt_cache_retention support.""" + + def setup_method(self): + self.config = OpenAIGPTConfig() + + def test_prompt_cache_key_in_supported_params(self): + """Test that prompt_cache_key is in supported params for OpenAI models.""" + supported_params = self.config.get_supported_openai_params("gpt-4o") + assert "prompt_cache_key" in supported_params + + def test_prompt_cache_retention_in_supported_params(self): + """Test that prompt_cache_retention is in supported params for OpenAI models.""" + supported_params = self.config.get_supported_openai_params("gpt-4o") + assert "prompt_cache_retention" in supported_params + + def test_prompt_cache_params_passed_through(self): + """Test that prompt_cache_key and prompt_cache_retention are passed through by map_openai_params.""" + optional_params = self.config.map_openai_params( + non_default_params={ + "prompt_cache_key": "my-cache-key", + "prompt_cache_retention": "24h", + }, + optional_params={}, + model="gpt-4o", + drop_params=False, + ) + assert optional_params.get("prompt_cache_key") == "my-cache-key" + assert optional_params.get("prompt_cache_retention") == "24h" diff --git a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py index af366b082a0..d236f46e5c5 100644 --- a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py @@ -761,3 +761,70 @@ async def test_request_body_with_html_script_tags(): f"Message content with HTML was modified during parsing: " f"expected={msg['content']!r}, got={result['messages'][2]['content']!r}" ) + + +def test_safe_get_request_headers_caches_on_request_state(): + """ + Test that _safe_get_request_headers caches the result on request.state + and returns the same object on subsequent calls. + """ + mock_request = MagicMock() + mock_request.headers = {"content-type": "application/json", "authorization": "Bearer sk-123"} + mock_request.state = MagicMock(spec=[]) # empty spec so getattr returns default + + # First call — should create and cache + result1 = _safe_get_request_headers(mock_request) + assert result1 == {"content-type": "application/json", "authorization": "Bearer sk-123"} + assert mock_request.state._cached_headers is result1 + + # Second call — should return the cached object (same identity) + result2 = _safe_get_request_headers(mock_request) + assert result2 is result1 + + +def test_safe_get_request_headers_none_request(): + """ + Test that _safe_get_request_headers returns empty dict for None request. + """ + result = _safe_get_request_headers(None) + assert result == {} + + +def test_safe_get_request_headers_copy_protects_cache(): + """ + Test that callers using .copy() before mutation do not corrupt the cache. + """ + mock_request = MagicMock() + mock_request.headers = {"authorization": "Bearer sk-123", "host": "localhost"} + mock_request.state = MagicMock(spec=[]) + + original = _safe_get_request_headers(mock_request) + + # Simulate what mutation call sites do: copy then pop + mutable = _safe_get_request_headers(mock_request).copy() + mutable.pop("authorization", None) + + # Cache must be unaffected + assert "authorization" in _safe_get_request_headers(mock_request) + assert _safe_get_request_headers(mock_request) is original + + +def test_safe_get_request_headers_state_unavailable(): + """ + Test that _safe_get_request_headers still returns headers when + request.state rejects attribute writes (the except path on the cache-write). + """ + class ReadOnlyState: + """State object that allows reads but raises on writes.""" + def __setattr__(self, name, value): + raise AttributeError("read-only state") + + def __getattr__(self, name): + return None # _cached_headers not found → triggers fresh read + + mock_request = MagicMock() + mock_request.headers = {"content-type": "application/json"} + mock_request.state = ReadOnlyState() + + result = _safe_get_request_headers(mock_request) + assert result == {"content-type": "application/json"} diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py new file mode 100644 index 00000000000..defdb3834d8 --- /dev/null +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py @@ -0,0 +1,75 @@ +""" +Unit tests for ToolDiscoveryQueue. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, +) + + +@pytest.fixture +def queue(): + return ToolDiscoveryQueue() + + +def test_add_single_tool(queue): + queue.add_update({"tool_name": "my_tool", "origin": "user_defined"}) + items = queue.flush() + assert len(items) == 1 + assert items[0]["tool_name"] == "my_tool" + assert items[0]["origin"] == "user_defined" + + +def test_deduplication_same_name(queue): + """Adding the same tool_name twice should only keep the first.""" + queue.add_update({"tool_name": "tool_a", "origin": "mcp_server"}) + queue.add_update({"tool_name": "tool_a", "origin": "user_defined"}) + items = queue.flush() + assert len(items) == 1 + assert items[0]["origin"] == "mcp_server" # first wins + + +def test_deduplication_different_names(queue): + queue.add_update({"tool_name": "tool_a"}) + queue.add_update({"tool_name": "tool_b"}) + items = queue.flush() + assert len(items) == 2 + names = {i["tool_name"] for i in items} + assert names == {"tool_a", "tool_b"} + + +def test_flush_clears_pending(queue): + queue.add_update({"tool_name": "tool_x"}) + items1 = queue.flush() + assert len(items1) == 1 + items2 = queue.flush() + assert len(items2) == 0 + + +def test_seen_names_reset_after_flush(queue): + """Seen-set is cleared on flush so the same tool can re-enter the next cycle.""" + queue.add_update({"tool_name": "tool_a"}) + queue.flush() + queue.add_update({"tool_name": "tool_a"}) # same tool, new cycle + items = queue.flush() + assert len(items) == 1 + assert items[0]["tool_name"] == "tool_a" + + +def test_empty_tool_name_ignored(queue): + queue.add_update({"tool_name": ""}) + queue.add_update({"tool_name": None}) # type: ignore[arg-type] + items = queue.flush() + assert len(items) == 0 + + +def test_flush_returns_list(queue): + result = queue.flush() + assert isinstance(result, list) diff --git a/tests/test_litellm/proxy/db/test_tool_registry_writer.py b/tests/test_litellm/proxy/db/test_tool_registry_writer.py new file mode 100644 index 00000000000..44f9e32058a --- /dev/null +++ b/tests/test_litellm/proxy/db/test_tool_registry_writer.py @@ -0,0 +1,197 @@ +""" +Unit tests for tool_registry_writer.py — uses a mock prisma client +that exposes execute_raw / query_raw (matching the actual raw-SQL implementation). +""" + +import os +import sys +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.db.tool_registry_writer import ( + batch_upsert_tools, + get_tool, + get_tools_by_names, + list_tools, + update_tool_policy, +) + + +def _make_prisma(query_rows=None): + """Return a minimal mock prisma_client with execute_raw / query_raw.""" + default_row = { + "tool_id": "uuid-1", + "tool_name": "my_tool", + "origin": "user_defined", + "call_policy": "untrusted", + "call_count": 1, + "assignments": {}, + "key_hash": None, + "team_id": None, + "key_alias": None, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "created_by": None, + "updated_by": None, + } + rows = query_rows if query_rows is not None else [default_row] + + prisma = MagicMock() + prisma.db.execute_raw = AsyncMock(return_value=None) + prisma.db.query_raw = AsyncMock(return_value=rows) + return prisma + + +@pytest.mark.asyncio +async def test_batch_upsert_tools_calls_execute_raw(): + prisma = _make_prisma() + items = [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}] + await batch_upsert_tools(prisma, items) + prisma.db.execute_raw.assert_awaited_once() + call_args = prisma.db.execute_raw.call_args + sql = call_args.args[0] + assert "LiteLLM_ToolTable" in sql + assert "ON CONFLICT" in sql + + +@pytest.mark.asyncio +async def test_batch_upsert_tools_empty_list(): + prisma = _make_prisma() + await batch_upsert_tools(prisma, []) + prisma.db.execute_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_batch_upsert_tools_skips_empty_names(): + prisma = _make_prisma() + items = [{"tool_name": "", "origin": None}, {"tool_name": None}] # type: ignore[list-item] + await batch_upsert_tools(prisma, items) + prisma.db.execute_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_batch_upsert_multiple_tools_calls_execute_raw_per_tool(): + prisma = _make_prisma() + items = [ + {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}, + {"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice"}, + ] + await batch_upsert_tools(prisma, items) + assert prisma.db.execute_raw.await_count == 2 + + +@pytest.mark.asyncio +async def test_list_tools_no_filter(): + row = { + "tool_id": "id1", + "tool_name": "tool_a", + "origin": "mcp", + "call_policy": "untrusted", + "call_count": 5, + "assignments": {}, + "key_hash": None, + "team_id": None, + "key_alias": None, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "created_by": None, + "updated_by": None, + } + prisma = _make_prisma(query_rows=[row]) + result = await list_tools(prisma) + assert len(result) == 1 + assert result[0].tool_name == "tool_a" + assert result[0].call_count == 5 + prisma.db.query_raw.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_list_tools_with_policy_filter(): + row = { + "tool_id": "id1", + "tool_name": "blocked_tool", + "origin": None, + "call_policy": "blocked", + "call_count": 2, + "assignments": None, + "key_hash": None, + "team_id": None, + "key_alias": None, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "created_by": None, + "updated_by": None, + } + prisma = _make_prisma(query_rows=[row]) + result = await list_tools(prisma, call_policy="blocked") + assert result[0].call_policy == "blocked" + call_args = prisma.db.query_raw.call_args + sql = call_args.args[0] + assert "WHERE call_policy" in sql + + +@pytest.mark.asyncio +async def test_get_tool_found(): + prisma = _make_prisma() + result = await get_tool(prisma, "my_tool") + assert result is not None + assert result.tool_name == "my_tool" + prisma.db.query_raw.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_get_tool_not_found(): + prisma = _make_prisma(query_rows=[]) + result = await get_tool(prisma, "nonexistent") + assert result is None + + +@pytest.mark.asyncio +async def test_update_tool_policy_calls_execute_raw(): + row = { + "tool_id": "uuid-1", + "tool_name": "my_tool", + "origin": "user_defined", + "call_policy": "blocked", + "call_count": 1, + "assignments": {}, + "key_hash": None, + "team_id": None, + "key_alias": None, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "created_by": None, + "updated_by": "admin", + } + prisma = _make_prisma(query_rows=[row]) + result = await update_tool_policy(prisma, "my_tool", "blocked", "admin") + assert result is not None + assert result.call_policy == "blocked" + prisma.db.execute_raw.assert_awaited_once() + call_args = prisma.db.execute_raw.call_args + sql = call_args.args[0] + assert "ON CONFLICT" in sql + assert "call_policy" in sql + + +@pytest.mark.asyncio +async def test_get_tools_by_names_returns_policy_map(): + rows = [ + {"tool_name": "tool_a", "call_policy": "trusted"}, + {"tool_name": "tool_b", "call_policy": "blocked"}, + ] + prisma = _make_prisma(query_rows=rows) + result = await get_tools_by_names(prisma, ["tool_a", "tool_b"]) + assert result == {"tool_a": "trusted", "tool_b": "blocked"} + + +@pytest.mark.asyncio +async def test_get_tools_by_names_empty_list(): + prisma = _make_prisma() + result = await get_tools_by_names(prisma, []) + assert result == {} + prisma.db.query_raw.assert_not_awaited() diff --git a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py index 2f2eaa905be..205d724c2b0 100644 --- a/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/test_litellm/proxy/google_endpoints/test_google_api_endpoints.py @@ -2,10 +2,9 @@ """ Test to verify the Google GenAI proxy API endpoints """ -import asyncio import os import sys -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, patch import pytest @@ -13,7 +12,6 @@ sys.path.insert( 0, os.path.abspath("../../..") ) # Adds the parent directory to the system path -import litellm def test_google_generate_content_endpoint(): @@ -401,3 +399,123 @@ def test_google_generate_content_with_image_config(): assert "contents" in called_data assert len(called_data["contents"]) == 1 assert called_data["contents"][0]["role"] == "user" + + +def test_google_generate_content_metadata_and_trace_id_callbacks(): + """Test that google_generate_content sets litellm_call_id and logging_obj for callbacks (e.g. S3, Langfuse)""" + try: + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy.google_endpoints.endpoints import router as google_router + except ImportError as e: + pytest.skip(f"Skipping test due to missing dependency: {e}") + + # Create a FastAPI app and include the router + app = FastAPI() + app.include_router(google_router) + + # Create a test client + client = TestClient(app) + + # Mock all required proxy server dependencies + with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch( + "litellm.proxy.proxy_server.general_settings", {} + ), patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, patch( + "litellm.proxy.proxy_server.version", "1.0.0" + ), patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" + ) as mock_add_data: + mock_router.agenerate_content = AsyncMock(return_value={"test": "response"}) + + # Mock add_litellm_data_to_request to return data with metadata + async def mock_add_litellm_data( + data, request, user_api_key_dict, proxy_config, general_settings, version + ): + # Simulate adding user metadata + data["litellm_metadata"] = { + "user_api_key_user_id": "test-user-id", + } + return data + + mock_add_data.side_effect = mock_add_litellm_data + + # Send a request to the endpoint with x-litellm-call-id header + test_call_id = "test-custom-call-id" + response = client.post( + "/v1beta/models/test-model:generateContent", + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + headers={ + "Authorization": "Bearer sk-test-key", + "x-litellm-call-id": test_call_id, + }, + ) + + assert response.status_code == 200 + + mock_router.agenerate_content.assert_called_once() + call_args = mock_router.agenerate_content.call_args + called_data = call_args[1] + + # Verify that the litellm_logging_obj got assigned in the final called_data to router + assert "litellm_logging_obj" in called_data + assert "litellm_call_id" in called_data + assert called_data["litellm_call_id"] == test_call_id + + +def test_google_stream_generate_content_metadata_and_trace_id_callbacks(): + """Test that google_stream_generate_content sets litellm_call_id and logging_obj for callbacks""" + try: + from fastapi import FastAPI + from fastapi.testclient import TestClient + + from litellm.proxy.google_endpoints.endpoints import router as google_router + except ImportError as e: + pytest.skip(f"Skipping test due to missing dependency: {e}") + + app = FastAPI() + app.include_router(google_router) + client = TestClient(app) + + mock_stream = AsyncMock() + mock_stream.__aiter__ = lambda self: mock_stream + mock_stream.__anext__.side_effect = StopAsyncIteration + + with patch("litellm.proxy.proxy_server.llm_router") as mock_router, patch( + "litellm.proxy.proxy_server.general_settings", {} + ), patch("litellm.proxy.proxy_server.proxy_config") as mock_proxy_config, patch( + "litellm.proxy.proxy_server.version", "1.0.0" + ), patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request" + ) as mock_add_data: + mock_router.agenerate_content_stream = AsyncMock(return_value=mock_stream) + + async def mock_add_litellm_data( + data, request, user_api_key_dict, proxy_config, general_settings, version + ): + data["litellm_metadata"] = { + "user_api_key_user_id": "test-user-id", + } + return data + + mock_add_data.side_effect = mock_add_litellm_data + + test_call_id = "test-custom-stream-call-id" + response = client.post( + "/v1beta/models/test-model:streamGenerateContent", + json={"contents": [{"role": "user", "parts": [{"text": "Hello stream"}]}]}, + headers={ + "Authorization": "Bearer sk-test-key", + "x-litellm-call-id": test_call_id, + }, + ) + + assert response.status_code == 200 + + mock_router.agenerate_content_stream.assert_called_once() + call_args = mock_router.agenerate_content_stream.call_args + called_data = call_args[1] + + assert "litellm_logging_obj" in called_data + assert "litellm_call_id" in called_data + assert called_data["litellm_call_id"] == test_call_id diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index e2c05f0ad37..76f9c39acd0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -24,12 +24,29 @@ from litellm.types.guardrails import LitellmParams, PiiAction, PiiEntityType from litellm.types.utils import Choices, Message, ModelResponse -def _make_mock_session_iterator(json_response): +def _make_mock_session_iterator( + json_response, status=200, content_type="application/json", text_response="" +): """Create a mock _get_session_iterator that yields a session returning json_response.""" @asynccontextmanager async def mock_iterator(): class MockResponse: + def __init__(self): + self.status = status + self.content_type = content_type + self.headers = {"Content-Type": content_type} + + async def text(self): + if text_response: + return text_response + import json + + try: + return json.dumps(json_response) + except Exception: + return str(json_response) + async def json(self): return json_response @@ -41,6 +58,7 @@ def _make_mock_session_iterator(json_response): class MockSession: def post(self, *args, **kwargs): + self.last_kwargs = kwargs return MockResponse() async def __aenter__(self): @@ -1444,3 +1462,149 @@ def test_deny_list_and_score_threshold_combined(): # EMAIL_ADDRESS passes both filters assert len(filtered) == 1 assert filtered[0]["entity_type"] == "EMAIL_ADDRESS" + + +@pytest.mark.asyncio +async def test_analyze_text_non_json_content_type_fail_closed(): + """ + Test that analyze_text raises GuardrailRaisedException when Presidio health + endpoint returns text/html and fail-closed is enabled. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + pii_entities_config={"PERSON": PiiAction.BLOCK}, + mock_testing=False, + ) + + mock_iterator = _make_mock_session_iterator( + json_response=None, + status=200, + content_type="text/html; charset=utf-8", + text_response="Presidio Analyzer service is up.", + ) + + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + with pytest.raises(GuardrailRaisedException) as exc_info: + await guardrail.analyze_text( + text="Hello world", + presidio_config=None, + request_data={}, + ) + assert "expected application/json Content-Type" in str(exc_info.value) + assert "text/html" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_analyze_text_non_json_content_type_fail_open(): + """ + Test that analyze_text returns empty list when Presidio returns text/html + and fail-closed is NOT enabled. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + mock_testing=False, + ) + + mock_iterator = _make_mock_session_iterator( + json_response=None, + status=200, + content_type="text/html; charset=utf-8", + text_response="Presidio Analyzer service is up.", + ) + + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + results = await guardrail.analyze_text( + text="Hello world", + presidio_config=None, + request_data={}, + ) + assert results == [] + + +@pytest.mark.asyncio +async def test_analyze_text_http_error_status(): + """ + Test that analyze_text handles 5xx HTTP errors properly. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + pii_entities_config={"PERSON": PiiAction.BLOCK}, + mock_testing=False, + ) + + mock_iterator = _make_mock_session_iterator( + json_response=None, + status=500, + content_type="text/plain", + text_response="Internal Server Error", + ) + + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + with pytest.raises(GuardrailRaisedException) as exc_info: + await guardrail.analyze_text( + text="Hello world", + presidio_config=None, + request_data={}, + ) + assert "HTTP 500" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_anonymize_text_non_json_content_type(): + """ + Test that anonymize_text raises Exception for non-JSON responses. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + mock_testing=False, + ) + + mock_iterator = _make_mock_session_iterator( + json_response=None, + status=200, + content_type="text/html", + text_response="Presidio Anonymizer service is up.", + ) + + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + with pytest.raises( + Exception, match="Presidio anonymizer returned non-JSON Content-Type" + ): + await guardrail.anonymize_text( + text="Hello world", + analyze_results=[{"start": 0, "end": 5, "entity_type": "PERSON"}], + output_parse_pii=False, + masked_entity_count={}, + ) + + +@pytest.mark.asyncio +async def test_anonymize_text_http_error_status(): + """ + Test that anonymize_text raises Exception on HTTP error. + """ + guardrail = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + mock_testing=False, + ) + + mock_iterator = _make_mock_session_iterator( + json_response=None, + status=502, + content_type="text/plain", + text_response="Bad Gateway", + ) + + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + with pytest.raises(Exception, match="Presidio anonymizer returned HTTP 502"): + await guardrail.anonymize_text( + text="Hello world", + analyze_results=[{"start": 0, "end": 5, "entity_type": "PERSON"}], + output_parse_pii=False, + masked_entity_count={}, + ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py new file mode 100644 index 00000000000..c6a81efbf0b --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py @@ -0,0 +1,181 @@ +""" +Unit tests for ToolPolicyGuardrail. +""" + +import os +import sys +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +sys.path.insert(0, os.path.abspath("../../../../../..")) + +from litellm.proxy.guardrails.guardrail_hooks.tool_policy.tool_policy_guardrail import ( + ToolPolicyGuardrail, +) +from litellm.types.guardrails import GuardrailEventHooks + + +@pytest.fixture +def guardrail(): + return ToolPolicyGuardrail() + + +# --- helpers --- + +def _tool_request_inputs(tool_names: list) -> dict: + return { + "tools": [ + {"type": "function", "function": {"name": name, "description": ""}} + for name in tool_names + ] + } + + +def _tool_response_inputs(tool_names: list) -> dict: + return { + "tool_calls": [ + {"type": "function", "function": {"name": name}} + for name in tool_names + ] + } + + +# --- tests --- + + +def test_guardrail_supports_pre_and_post_call(guardrail): + hooks = guardrail.supported_event_hooks + assert GuardrailEventHooks.pre_call in hooks + assert GuardrailEventHooks.post_call in hooks + + +@pytest.mark.asyncio +async def test_no_tools_in_request_passes_through(guardrail): + inputs: Any = {"tools": []} + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_no_tool_calls_in_response_passes_through(guardrail): + inputs: Any = {"tool_calls": []} + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_untrusted_tools_pass_through(guardrail): + policy_map = {"search": "untrusted", "read_file": "trusted"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + inputs: Any = _tool_request_inputs(["search", "read_file"]) + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_blocked_tool_in_request_raises_http_exception(guardrail): + policy_map = {"dangerous_tool": "blocked"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + inputs: Any = _tool_request_inputs(["dangerous_tool"]) + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert exc_info.value.status_code == 400 + assert "dangerous_tool" in exc_info.value.detail["blocked_tools"] + + +@pytest.mark.asyncio +async def test_blocked_tool_in_response_raises_http_exception(guardrail): + policy_map = {"exfil_tool": "blocked"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + inputs: Any = _tool_response_inputs(["exfil_tool"]) + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert exc_info.value.status_code == 400 + assert "exfil_tool" in exc_info.value.detail["blocked_tools"] + + +@pytest.mark.asyncio +async def test_mixed_blocked_and_allowed_raises_for_blocked(guardrail): + policy_map = {"safe_tool": "trusted", "bad_tool": "blocked"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + inputs: Any = _tool_request_inputs(["safe_tool", "bad_tool"]) + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + blocked = exc_info.value.detail["blocked_tools"] + assert "bad_tool" in blocked + assert "safe_tool" not in blocked + + +@pytest.mark.asyncio +async def test_tool_not_in_db_passes_through(guardrail): + """Tools not found in the DB (no entry) should not be blocked.""" + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value={})): + inputs: Any = _tool_request_inputs(["unknown_tool"]) + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_get_policies_cached_uses_cache(guardrail): + """Second call with same tool names should return the cached result.""" + policy_map = {"tool_a": "trusted"} + with patch( + "litellm.proxy.db.tool_registry_writer.get_tools_by_names", + new=AsyncMock(return_value=policy_map), + ) as mock_db, patch( + "litellm.proxy.proxy_server.prisma_client", + new=MagicMock(), + ): + # first call — should hit DB + result1 = await guardrail._get_policies_cached(["tool_a"]) + assert result1 == policy_map + + # second call — should hit cache, not DB again + result2 = await guardrail._get_policies_cached(["tool_a"]) + assert result2 == policy_map + + assert mock_db.call_count == 1 + + +@pytest.mark.asyncio +async def test_get_policies_cached_no_prisma(guardrail): + """Without a prisma client, returns empty dict.""" + with patch( + "litellm.proxy.proxy_server.prisma_client", + None, + ): + result = await guardrail._get_policies_cached(["tool_a"]) + assert result == {} + + +@pytest.mark.asyncio +async def test_response_tool_calls_as_objects(guardrail): + """tool_calls that are objects (not dicts) with .function.name should work.""" + policy_map = {"obj_tool": "blocked"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + fn = MagicMock() + fn.name = "obj_tool" + tc = MagicMock() + tc.function = fn + inputs: Any = {"tool_calls": [tc]} + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index e8765cf78ca..c46b8df5efc 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -169,6 +169,223 @@ async def test_track_cost_callback_skips_when_no_standard_logging_object(): mock_proxy_logging.failed_tracking_alert.assert_not_called() +@pytest.mark.asyncio +async def test_enrich_failure_metadata_with_team_alias(): + """ + When team_id is set but team_alias is missing (and key_alias is present), + _enrich_failure_metadata_with_key_info should look up the team from cache + and populate user_api_key_team_alias. + """ + mock_team_obj = MagicMock() + mock_team_obj.team_alias = "my-team-alias" + + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ): + metadata = { + "user_api_key": "hashed_key", + "user_api_key_alias": "my-key-alias", # already set + "user_api_key_team_id": "test_team_id", + "user_api_key_team_alias": None, + } + result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) + assert result["user_api_key_team_alias"] == "my-team-alias" + + +@pytest.mark.asyncio +async def test_enrich_failure_metadata_with_full_key_lookup(): + """ + When all key fields are null (auth error 401 scenario), _enrich_failure_metadata_with_key_info + should look up the key object from cache/DB and populate alias, user_id, team_id, + then look up the team to get team_alias. + """ + mock_key_obj = MagicMock() + mock_key_obj.key_alias = "fetched-key-alias" + mock_key_obj.user_id = "fetched-user-id" + mock_key_obj.team_id = "fetched-team-id" + mock_key_obj.org_id = "fetched-org-id" + + mock_team_obj = MagicMock() + mock_team_obj.team_alias = "fetched-team-alias" + + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + return_value=mock_key_obj, + ), patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ): + metadata = { + "user_api_key": "hashed_key", + "user_api_key_alias": None, # all null - simulates auth error path + "user_api_key_user_id": None, + "user_api_key_team_id": None, + "user_api_key_team_alias": None, + "user_api_key_org_id": None, + } + result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) + assert result["user_api_key_alias"] == "fetched-key-alias" + assert result["user_api_key_user_id"] == "fetched-user-id" + assert result["user_api_key_team_id"] == "fetched-team-id" + assert result["user_api_key_org_id"] == "fetched-org-id" + assert result["user_api_key_team_alias"] == "fetched-team-alias" + + +@pytest.mark.asyncio +async def test_enrich_failure_metadata_skips_when_team_alias_present(): + """ + When team_alias is already populated, _enrich_failure_metadata_with_key_info + should not perform a team cache lookup. + """ + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + ) as mock_get_key, patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + ) as mock_get_team: + metadata = { + "user_api_key": "hashed_key", + "user_api_key_alias": "existing-alias", + "user_api_key_team_id": "test_team_id", + "user_api_key_team_alias": "already-set", + } + result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) + assert result["user_api_key_team_alias"] == "already-set" + mock_get_key.assert_not_called() + mock_get_team.assert_not_called() + + +@pytest.mark.asyncio +async def test_enrich_failure_metadata_skips_when_no_api_key(): + """ + When api_key hash is absent, _enrich_failure_metadata_with_key_info should + not perform any lookups. + """ + with patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + ) as mock_get_key: + metadata = { + "user_api_key": None, + "user_api_key_alias": None, + "user_api_key_team_id": None, + "user_api_key_team_alias": None, + } + result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata) + mock_get_key.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_enriches_auth_error_metadata(): + """ + Simulates a 401 ProxyException (e.g. can_key_call_model). In this case + UserAPIKeyAuth is created with only api_key set. The failure hook should + look up the key and team from cache/DB to populate all missing fields. + """ + logger = _ProxyDBLogger() + + # This is what auth_exception_handler creates for 401 errors + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed_key", + # key_alias, user_id, team_id, team_alias are all None + ) + + request_data = { + "model": "claude-haiku-4-5", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "litellm_params": {}, + } + + mock_key_obj = MagicMock() + mock_key_obj.key_alias = "my-key-alias" + mock_key_obj.user_id = "my-user-id" + mock_key_obj.team_id = "my-team-id" + mock_key_obj.org_id = None + + mock_team_obj = MagicMock() + mock_team_obj.team_alias = "my-team-alias" + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_key_object", + new_callable=AsyncMock, + return_value=mock_key_obj, + ), patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ): + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("401 - model not allowed"), + user_api_key_dict=user_api_key_dict, + ) + + mock_update_database.assert_called_once() + call_args = mock_update_database.call_args[1] + metadata = call_args["kwargs"]["litellm_params"]["metadata"] + assert metadata["user_api_key_alias"] == "my-key-alias" + assert metadata["user_api_key_user_id"] == "my-user-id" + assert metadata["user_api_key_team_id"] == "my-team-id" + assert metadata["user_api_key_team_alias"] == "my-team-alias" + + +@pytest.mark.asyncio +async def test_async_post_call_failure_hook_enriches_missing_team_alias(): + """ + When user_api_key_dict has a team_id but no team_alias, async_post_call_failure_hook + should look up the team from cache and populate user_api_key_team_alias in the + spend log metadata written to the DB. + """ + logger = _ProxyDBLogger() + + user_api_key_dict = UserAPIKeyAuth( + api_key="test_api_key", + key_alias="test_alias", + user_id="test_user_id", + team_id="test_team_id", + team_alias=None, # Missing - simulates regular key auth where SQL view omits team_alias + ) + + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {}, + "litellm_params": {}, + } + + mock_team_obj = MagicMock() + mock_team_obj.team_alias = "enriched-team-alias" + + with patch( + "litellm.proxy.db.db_spend_update_writer.DBSpendUpdateWriter.update_database", + new_callable=AsyncMock, + ) as mock_update_database, patch( + "litellm.proxy.hooks.proxy_track_cost_callback.get_team_object", + new_callable=AsyncMock, + return_value=mock_team_obj, + ): + await logger.async_post_call_failure_hook( + request_data=request_data, + original_exception=Exception("Provider rate limit"), + user_api_key_dict=user_api_key_dict, + ) + + mock_update_database.assert_called_once() + call_args = mock_update_database.call_args[1] + metadata = call_args["kwargs"]["litellm_params"]["metadata"] + assert metadata["user_api_key_team_alias"] == "enriched-team-alias" + assert metadata["user_api_key_team_id"] == "test_team_id" + + @pytest.mark.asyncio @pytest.mark.parametrize("model_value", [None, ""]) async def test_track_cost_callback_skips_for_falsy_model_and_no_slo(model_value): diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index ffb4e955426..05df3c2dcbb 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -5623,6 +5623,7 @@ async def test_rotate_master_key_model_data_valid_for_prisma( litellm_params/model_info are JSON strings (create_many expects dicts). """ from unittest.mock import AsyncMock, MagicMock + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.management_endpoints.key_management_endpoints import ( _rotate_master_key, @@ -6157,3 +6158,56 @@ async def test_get_member_team_ids(): # Should return team-A and team-B (user is a member of both) # Should NOT return team-C (user is not in members list) assert sorted(result) == ["team-A", "team-B"] + + +@pytest.mark.asyncio +async def test_generate_key_with_agent_id(): + """Test that agent_id is accepted in GenerateKeyRequest and passed to generate_key_helper_fn.""" + from litellm.proxy._types import GenerateKeyRequest + + # Verify GenerateKeyRequest accepts agent_id + request = GenerateKeyRequest( + key_alias="agent-test-key", + agent_id="test-agent-123", + models=[], + ) + assert request.agent_id == "test-agent-123" + data_json = request.model_dump(exclude_unset=True, exclude_none=True) + assert data_json["agent_id"] == "test-agent-123" + + +@pytest.mark.asyncio +async def test_generate_key_helper_fn_agent_id(): + """Test that generate_key_helper_fn passes agent_id into the insert_data call.""" + from unittest.mock import AsyncMock, MagicMock, call, patch + + import litellm.proxy.management_endpoints.key_management_endpoints as km + + mock_prisma_client = AsyncMock() + mock_insert = AsyncMock( + return_value=MagicMock( + token="sk-test", + created_at=None, + updated_at=None, + litellm_budget_table=None, + ) + ) + mock_prisma_client.insert_data = mock_insert + + with patch.object(km, "prisma_client", mock_prisma_client): + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client): + await generate_key_helper_fn( + request_type="key", + agent_id="test-agent-456", + key_alias="test-agent-key", + models=[], + table_name="key", + ) + + assert mock_insert.called, "insert_data was never called" + # insert_data is called as insert_data(data=key_data, ...) + call_kwargs = mock_insert.call_args.kwargs + key_data = call_kwargs.get("data", {}) + assert key_data.get("agent_id") == "test-agent-456", ( + f"Expected agent_id='test-agent-456' in key_data, got: {key_data.get('agent_id')}" + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py new file mode 100644 index 00000000000..6f1d373fdee --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py @@ -0,0 +1,149 @@ +""" +Unit tests for tool management endpoints (/v1/tool/*). +Uses FastAPI TestClient with mocked DB functions. + +Patches target the source modules (litellm.proxy.db.tool_registry_writer.* +and litellm.proxy.proxy_server.prisma_client) because the endpoint code +imports these inside function bodies to avoid circular imports. +""" + +import os +import sys +from datetime import datetime, timezone +from typing import Optional +from unittest.mock import AsyncMock, MagicMock, patch + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.management_endpoints.tool_management_endpoints import router +from litellm.types.tool_management import LiteLLM_ToolTableRow + +# --- helpers --- + + +def _make_tool_row( + tool_name: str = "my_tool", + call_policy: str = "untrusted", + origin: Optional[str] = None, +) -> LiteLLM_ToolTableRow: + now = datetime.now(timezone.utc) + return LiteLLM_ToolTableRow( + tool_id="uuid-1", + tool_name=tool_name, + origin=origin, + call_policy=call_policy, # type: ignore[arg-type] + assignments={}, + created_at=now, + updated_at=now, + ) + + +def _make_app() -> FastAPI: + """Build a minimal FastAPI app with the tool management router.""" + app = FastAPI() + app.include_router(router) + return app + + +# Stub the auth dependency so we don't need a real proxy running. +def _override_auth(): + from litellm.proxy._types import UserAPIKeyAuth + + return UserAPIKeyAuth(api_key="sk-test", user_id="admin") + + +# A real (non-None) prisma stub for truthiness checks. +_MOCK_PRISMA = MagicMock() + + +# --- test class --- + + +class TestToolManagementEndpoints: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch( + "litellm.proxy.db.tool_registry_writer.list_tools", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_list_tools_returns_200(self, mock_db_list): + mock_db_list.return_value = [_make_tool_row()] + + resp = self.client.get("/v1/tool/list") + assert resp.status_code == 200 + body = resp.json() + assert body["total"] == 1 + assert body["tools"][0]["tool_name"] == "my_tool" + + @patch( + "litellm.proxy.db.tool_registry_writer.list_tools", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_list_tools_with_policy_filter(self, mock_db_list): + mock_db_list.return_value = [_make_tool_row(call_policy="blocked")] + + resp = self.client.get("/v1/tool/list?call_policy=blocked") + assert resp.status_code == 200 + assert resp.json()["tools"][0]["call_policy"] == "blocked" + + @patch( + "litellm.proxy.db.tool_registry_writer.get_tool", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_get_tool_found(self, mock_db_get): + mock_db_get.return_value = _make_tool_row(tool_name="tool_a") + + resp = self.client.get("/v1/tool/tool_a") + assert resp.status_code == 200 + assert resp.json()["tool_name"] == "tool_a" + + @patch( + "litellm.proxy.db.tool_registry_writer.get_tool", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_get_tool_not_found_returns_404(self, mock_db_get): + mock_db_get.return_value = None + + resp = self.client.get("/v1/tool/nonexistent", follow_redirects=True) + assert resp.status_code == 404 + + @patch( + "litellm.proxy.db.tool_registry_writer.update_tool_policy", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_update_tool_policy_blocked(self, mock_db_update): + mock_db_update.return_value = _make_tool_row(call_policy="blocked") + + resp = self.client.post( + "/v1/tool/policy", + json={"tool_name": "my_tool", "call_policy": "blocked"}, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["call_policy"] == "blocked" + assert body["updated"] is True + + @patch("litellm.proxy.proxy_server.prisma_client", None) + def test_list_tools_no_db_returns_500(self): + resp = self.client.get("/v1/tool/list") + assert resp.status_code == 500 + + def test_update_tool_policy_invalid_policy_returns_422(self): + resp = self.client.post( + "/v1/tool/policy", + json={"tool_name": "my_tool", "call_policy": "invalid_value"}, + ) + assert resp.status_code == 422 diff --git a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/__init__.py b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py new file mode 100644 index 00000000000..f9303bd13a6 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py @@ -0,0 +1,402 @@ +""" +Tests for AI Usage Chat module. +""" + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat import ( + TOOL_HANDLERS, + TOOLS_ADMIN, + TOOLS_BASE, + _build_system_prompt, + _summarise_entity_data, + _summarise_usage_data, + stream_usage_ai_chat, +) + + +SAMPLE_AGGREGATED_RESPONSE = { + "results": [ + { + "date": "2025-01-15", + "metrics": { + "spend": 50.25, + "prompt_tokens": 20000, + "completion_tokens": 10000, + "total_tokens": 30000, + "api_requests": 500, + "successful_requests": 480, + "failed_requests": 20, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + }, + "breakdown": { + "models": { + "gpt-4": { + "metrics": { + "spend": 40.0, + "api_requests": 300, + "total_tokens": 25000, + }, + "metadata": {}, + "api_key_breakdown": {}, + }, + }, + "providers": { + "openai": { + "metrics": {"spend": 50.25, "api_requests": 500}, + "metadata": {}, + "api_key_breakdown": {}, + }, + }, + "api_keys": { + "sk-test123": { + "metrics": {"spend": 50.25}, + "metadata": {"key_alias": "Production Key"}, + }, + }, + "model_groups": {}, + "mcp_servers": {}, + "entities": {}, + }, + }, + ], + "metadata": { + "total_spend": 50.25, + "total_api_requests": 500, + "total_successful_requests": 480, + "total_failed_requests": 20, + "total_tokens": 30000, + }, +} + +SAMPLE_TEAM_RESPONSE = { + "results": [ + { + "date": "2025-01-15", + "metrics": {"spend": 100.0, "api_requests": 1000, "total_tokens": 50000}, + "breakdown": { + "entities": { + "team-1": { + "metrics": { + "spend": 60.0, + "api_requests": 600, + "total_tokens": 30000, + }, + "metadata": {"alias": "Engineering"}, + "api_key_breakdown": {}, + }, + "team-2": { + "metrics": { + "spend": 40.0, + "api_requests": 400, + "total_tokens": 20000, + }, + "metadata": {"alias": "Marketing"}, + "api_key_breakdown": {}, + }, + }, + "models": {}, + "providers": {}, + "api_keys": {}, + "model_groups": {}, + "mcp_servers": {}, + }, + }, + ], + "metadata": {"total_spend": 100.0, "total_api_requests": 1000}, +} + + +class TestToolSchemas: + def test_admin_tools_include_all(self): + assert len(TOOLS_ADMIN) == 3 + names = {t["function"]["name"] for t in TOOLS_ADMIN} + assert "get_usage_data" in names + assert "get_team_usage_data" in names + assert "get_tag_usage_data" in names + + def test_base_tools_restricted_to_usage_only(self): + assert len(TOOLS_BASE) == 1 + assert TOOLS_BASE[0]["function"]["name"] == "get_usage_data" + + def test_admin_prompt_mentions_all_tools(self): + prompt = _build_system_prompt(is_admin=True) + assert "get_usage_data" in prompt + assert "get_team_usage_data" in prompt + assert "get_tag_usage_data" in prompt + + def test_non_admin_prompt_only_mentions_usage_tool(self): + prompt = _build_system_prompt(is_admin=False) + assert "get_usage_data" in prompt + assert "get_team_usage_data" not in prompt + assert "get_tag_usage_data" not in prompt + + def test_system_prompt_includes_todays_date(self): + from datetime import date + + prompt = _build_system_prompt(is_admin=True) + assert date.today().isoformat() in prompt + + +class TestSummariseUsageData: + def test_summarise_includes_totals(self): + summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) + assert "$50.25" in summary + assert "500" in summary + + def test_summarise_includes_models(self): + summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) + assert "gpt-4" in summary + + def test_summarise_includes_providers(self): + summary = _summarise_usage_data(SAMPLE_AGGREGATED_RESPONSE) + assert "openai" in summary + + def test_summarise_handles_empty_data(self): + empty = {"results": [], "metadata": {}} + summary = _summarise_usage_data(empty) + assert "no data" in summary.lower() + + +class TestSummariseEntityData: + def test_team_summary_includes_teams(self): + summary = _summarise_entity_data(SAMPLE_TEAM_RESPONSE, "Team") + assert "Engineering" in summary + assert "Marketing" in summary + assert "$60.0" in summary + assert "$40.0" in summary + + def test_team_summary_empty(self): + empty = {"results": [], "metadata": {}} + summary = _summarise_entity_data(empty, "Team") + assert "No Team usage data" in summary + + +class TestStreamUsageAiChat: + @pytest.mark.asyncio + async def test_stream_emits_status_events(self): + mock_tool_call = MagicMock() + mock_tool_call.id = "call_123" + mock_tool_call.function.name = "get_usage_data" + mock_tool_call.function.arguments = json.dumps( + { + "start_date": "2025-01-01", + "end_date": "2025-01-31", + } + ) + + mock_first_response = MagicMock() + mock_first_response.choices = [MagicMock()] + mock_first_response.choices[0].message.tool_calls = [mock_tool_call] + mock_first_response.choices[0].message.model_dump.return_value = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "get_usage_data", + "arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31"}', + }, + } + ], + } + + async def mock_stream(): + chunk = MagicMock() + chunk.choices = [MagicMock()] + chunk.choices[0].delta.content = "Total spend is $50.25" + yield chunk + + with patch( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" + ) as mock_litellm, patch( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat._fetch_usage_data", + new_callable=AsyncMock, + ) as mock_fetch: + mock_litellm.acompletion = AsyncMock( + side_effect=[ + mock_first_response, + mock_stream(), + ] + ) + mock_fetch.return_value = SAMPLE_AGGREGATED_RESPONSE + + events = [] + async for event in stream_usage_ai_chat( + messages=[{"role": "user", "content": "What is my total spend?"}], + model="gpt-4o-mini", + user_id="user-123", + is_admin=True, + ): + events.append(json.loads(event.replace("data: ", "").strip())) + + status_events = [e for e in events if e["type"] == "status"] + tool_call_events = [e for e in events if e["type"] == "tool_call"] + chunk_events = [e for e in events if e["type"] == "chunk"] + done_events = [e for e in events if e["type"] == "done"] + + assert len(status_events) >= 1 + assert "Thinking" in status_events[0]["message"] + assert len(tool_call_events) >= 1 + assert tool_call_events[0]["tool_name"] == "get_usage_data" + assert tool_call_events[0]["status"] in ("running", "complete") + assert len(chunk_events) >= 1 + assert len(done_events) == 1 + + @pytest.mark.asyncio + async def test_stream_handles_team_tool(self): + mock_tool_call = MagicMock() + mock_tool_call.id = "call_team" + mock_tool_call.function.name = "get_team_usage_data" + mock_tool_call.function.arguments = json.dumps( + { + "start_date": "2025-01-01", + "end_date": "2025-01-31", + } + ) + + mock_first_response = MagicMock() + mock_first_response.choices = [MagicMock()] + mock_first_response.choices[0].message.tool_calls = [mock_tool_call] + mock_first_response.choices[0].message.model_dump.return_value = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_team", + "type": "function", + "function": { + "name": "get_team_usage_data", + "arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31"}', + }, + } + ], + } + + async def mock_stream(): + chunk = MagicMock() + chunk.choices = [MagicMock()] + chunk.choices[0].delta.content = "Engineering is the top team." + yield chunk + + with patch( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" + ) as mock_litellm, patch( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat._fetch_team_usage_data", + new_callable=AsyncMock, + ) as mock_fetch: + mock_litellm.acompletion = AsyncMock( + side_effect=[ + mock_first_response, + mock_stream(), + ] + ) + mock_fetch.return_value = SAMPLE_TEAM_RESPONSE + + events = [] + async for event in stream_usage_ai_chat( + messages=[{"role": "user", "content": "Which team spends the most?"}], + model="gpt-4o-mini", + is_admin=True, + ): + events.append(json.loads(event.replace("data: ", "").strip())) + + chunk_events = [e for e in events if e["type"] == "chunk"] + assert len(chunk_events) >= 1 + assert "Engineering" in chunk_events[0]["content"] + + @pytest.mark.asyncio + async def test_stream_handles_error(self): + with patch( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" + ) as mock_litellm: + mock_litellm.acompletion = AsyncMock(side_effect=Exception("LLM error")) + + events = [] + async for event in stream_usage_ai_chat( + messages=[{"role": "user", "content": "test"}], + ): + events.append(json.loads(event.replace("data: ", "").strip())) + + error_events = [e for e in events if e["type"] == "error"] + assert len(error_events) == 1 + assert "internal error" in error_events[0]["message"].lower() + + @pytest.mark.asyncio + async def test_non_admin_enforces_user_id(self): + mock_tool_call = MagicMock() + mock_tool_call.id = "call_456" + mock_tool_call.function.name = "get_usage_data" + mock_tool_call.function.arguments = json.dumps( + { + "start_date": "2025-01-01", + "end_date": "2025-01-31", + "user_id": "other-user", + } + ) + + mock_first_response = MagicMock() + mock_first_response.choices = [MagicMock()] + mock_first_response.choices[0].message.tool_calls = [mock_tool_call] + mock_first_response.choices[0].message.model_dump.return_value = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_456", + "type": "function", + "function": { + "name": "get_usage_data", + "arguments": '{"start_date":"2025-01-01","end_date":"2025-01-31","user_id":"other-user"}', + }, + } + ], + } + + async def mock_stream(): + chunk = MagicMock() + chunk.choices = [MagicMock()] + chunk.choices[0].delta.content = "Data." + yield chunk + + mock_fetch = AsyncMock(return_value=SAMPLE_AGGREGATED_RESPONSE) + + with patch( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm" + ) as mock_litellm, patch.dict( + "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.TOOL_HANDLERS", + { + "get_usage_data": { + "fetch": mock_fetch, + "summarise": _summarise_usage_data, + "label": "global usage data", + } + }, + ): + mock_litellm.acompletion = AsyncMock( + side_effect=[ + mock_first_response, + mock_stream(), + ] + ) + + events = [] + async for event in stream_usage_ai_chat( + messages=[{"role": "user", "content": "Show data"}], + model="gpt-4o-mini", + user_id="my-user-id", + is_admin=False, + ): + events.append(event) + + mock_fetch.assert_called_once_with( + start_date="2025-01-01", + end_date="2025-01-31", + user_id="my-user-id", + ) diff --git a/tests/test_litellm/proxy/test_health_check_functions.py b/tests/test_litellm/proxy/test_health_check_functions.py index ccae9fb5425..4c91d0ae91e 100644 --- a/tests/test_litellm/proxy/test_health_check_functions.py +++ b/tests/test_litellm/proxy/test_health_check_functions.py @@ -12,6 +12,7 @@ sys.path.insert(0, os.path.abspath("../../..")) from litellm.proxy.health_endpoints._health_endpoints import ( _aggregate_health_check_results, _build_model_param_to_info_mapping, + _perform_health_check_and_save, _save_background_health_checks_to_db, _save_health_check_results_if_changed, _save_health_check_to_db, @@ -466,5 +467,43 @@ async def test_get_all_latest_health_checks_without_model_id(mock_prisma): assert result[0].checked_at == mock_check2.checked_at # Latest +@pytest.mark.asyncio +async def test_perform_health_check_and_save_passes_model_id_to_perform_health_check(): + """Test that _perform_health_check_and_save passes model_id to perform_health_check so health checks run by model id.""" + model_list = [ + { + "model_name": "gpt-4", + "model_info": {"id": "deployment-abc"}, + "litellm_params": {"model": "gpt-4"}, + }, + ] + healthy = [{"model": "gpt-4"}] + unhealthy = [] + + async def mock_perform_health_check(model_list, model=None, cli_model=None, details=True, model_id=None): + return healthy, unhealthy + + with patch( + "litellm.proxy.health_endpoints._health_endpoints.perform_health_check", + side_effect=mock_perform_health_check, + ) as mock_perform: + result = await _perform_health_check_and_save( + model_list=model_list, + target_model=None, + cli_model=None, + details=True, + prisma_client=None, + start_time=0.0, + user_id="user-1", + model_id="deployment-abc", + ) + + mock_perform.assert_called_once() + call_kwargs = mock_perform.call_args[1] + assert call_kwargs["model_id"] == "deployment-abc" + assert result["healthy_count"] == 1 + assert result["unhealthy_count"] == 0 + + if __name__ == "__main__": pytest.main([__file__]) \ No newline at end of file diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 5732deda6fb..4bb9685f5e6 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -925,6 +925,73 @@ def test_router_get_model_access_groups_team_only_models(): assert list(access_groups.keys()) == ["default-models"] +def test_cached_get_model_group_info(): + """ + Test that _cached_get_model_group_info caches results and + invalidates on deployment changes. + """ + from litellm.types.router import Deployment, LiteLLM_Params + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake"}, + "model_info": {"tpm": 1000, "rpm": 100}, + }, + ] + ) + + # First call should compute and cache + result1 = router._cached_get_model_group_info("gpt-4") + assert result1 is not None + assert result1.tpm == 1000 + + # Second call should hit cache (same object) + result2 = router._cached_get_model_group_info("gpt-4") + assert result1 is result2 + + # Add a deployment — cache should be invalidated + router.add_deployment( + Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params(model="gpt-4", api_key="fake2"), + model_info={"tpm": 2000, "rpm": 200}, + ) + ) + result3 = router._cached_get_model_group_info("gpt-4") + assert result3 is not result2 + assert result3 is not None + assert result3.tpm == 3000 # 1000 + 2000 + + # Delete a deployment — cache should be invalidated + deployment_id = router.model_list[-1]["model_info"]["id"] + router.delete_deployment(id=deployment_id) + result4 = router._cached_get_model_group_info("gpt-4") + assert result4 is not result3 + assert result4 is not None + assert result4.tpm == 1000 + + # set_model_list — cache should be invalidated + router.set_model_list( + [ + { + "model_name": "gpt-4", + "litellm_params": {"model": "gpt-4", "api_key": "fake"}, + "model_info": {"tpm": 5000}, + }, + ] + ) + result5 = router._cached_get_model_group_info("gpt-4") + assert result5 is not result4 + assert result5 is not None + assert result5.tpm == 5000 + + # Verify cache still works after invalidation + result6 = router._cached_get_model_group_info("gpt-4") + assert result5 is result6 + + def test_get_model_access_groups_caching(): """ Test that get_model_access_groups caches the no-args result diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 489a39a7ee2..3787f451ad3 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -1757,29 +1757,6 @@ "url": "https://opencollective.com/libvips" } }, - "node_modules/@isaacs/balanced-match": { - "version": "4.0.1", - "resolved": "https://registry.npmjs.org/@isaacs/balanced-match/-/balanced-match-4.0.1.tgz", - "integrity": "sha512-yzMTt9lEb8Gv7zRioUilSglI0c0smZ9k5D65677DLWLtWJaXIS3CqcGyUFByYKlnUj6TkjLVs54fBl6+TiGQDQ==", - "dev": true, - "license": "MIT", - "engines": { - "node": "20 || >=22" - } - }, - "node_modules/@isaacs/brace-expansion": { - "version": "5.0.1", - "resolved": "https://registry.npmjs.org/@isaacs/brace-expansion/-/brace-expansion-5.0.1.tgz", - "integrity": "sha512-WMz71T1JS624nWj2n2fnYAuPovhv7EUhk69R6i9dsVyzxt5eM3bjwvgk9L+APE1TRscGysAVMANkB0jh0LQZrQ==", - "dev": true, - "license": "MIT", - "dependencies": { - "@isaacs/balanced-match": "^4.0.1" - }, - "engines": { - "node": "20 || >=22" - } - }, "node_modules/@istanbuljs/schema": { "version": "0.1.3", "resolved": "https://registry.npmjs.org/@istanbuljs/schema/-/schema-0.1.3.tgz", @@ -3696,32 +3673,6 @@ "typescript": ">=4.8.4 <6.0.0" } }, - "node_modules/@typescript-eslint/typescript-estree/node_modules/brace-expansion": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.0.2.tgz", - "integrity": "sha512-Jt0vHyM+jmUBqojB7E1NIYadt0vI0Qxjxd2TErW94wDz+E2LAm5vKMXXwg6ZZBTHPuUlDgQHKXvjGBdfcF1ZDQ==", - "dev": true, - "license": "MIT", - "dependencies": { - "balanced-match": "^1.0.0" - } - }, - "node_modules/@typescript-eslint/typescript-estree/node_modules/minimatch": { - "version": "9.0.5", - "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-9.0.5.tgz", - "integrity": "sha512-G6T0ZX48xgozx7587koeX9Ys2NYy6Gmv//P89sEte9V9whIapMNF4idKxnW2QtCcLiTWlb/wfCabAtAFWhhBow==", - "dev": true, - "license": "ISC", - "dependencies": { - "brace-expansion": "^2.0.1" - }, - "engines": { - "node": ">=16 || 14 >=14.17" - }, - "funding": { - "url": "https://github.com/sponsors/isaacs" - } - }, "node_modules/@typescript-eslint/utils": { "version": "8.54.0", "resolved": "https://registry.npmjs.org/@typescript-eslint/utils/-/utils-8.54.0.tgz", @@ -4749,11 +4700,14 @@ } }, "node_modules/balanced-match": { - "version": "1.0.2", - "resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-1.0.2.tgz", - "integrity": "sha512-3oSeUO0TMV67hN1AmbXsK4yaqU7tjiHlbxRDZOpH0KW9+CeX4bRAaX0Anxt0tx2MrpRpWwQaPwIlISEJhYU5Pw==", + "version": "4.0.4", + "resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-4.0.4.tgz", + "integrity": "sha512-BLrgEcRTwX2o6gGxGOCNyMvGSp35YofuYzw9h1IMTRmKqttAZZVU67bdb9Pr2vUHA8+j3i2tJfjO6C6+4myGTA==", "dev": true, - "license": "MIT" + "license": "MIT", + "engines": { + "node": "18 || 20 || >=22" + } }, "node_modules/baseline-browser-mapping": { "version": "2.9.19", @@ -4787,14 +4741,16 @@ } }, "node_modules/brace-expansion": { - "version": "1.1.12", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.12.tgz", - "integrity": "sha512-9T9UjW3r0UW5c1Q7GTwllptXwhvYmEzFhzMfZ9H7FQWt+uZePjZPjBP/W1ZEyZ1twGWom5/56TF4lPcqjnDHcg==", + "version": "5.0.3", + "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-5.0.3.tgz", + "integrity": "sha512-fy6KJm2RawA5RcHkLa1z/ScpBeA762UF9KmZQxwIbDtRJrgLzM10depAiEQ+CXYcoiqW1/m96OAAoke2nE9EeA==", "dev": true, "license": "MIT", "dependencies": { - "balanced-match": "^1.0.0", - "concat-map": "0.0.1" + "balanced-match": "^4.0.2" + }, + "engines": { + "node": "18 || 20 || >=22" } }, "node_modules/braces": { @@ -5149,13 +5105,6 @@ "integrity": "sha512-VRhuHOLoKYOy4UbilLbUzbYg93XLjv2PncJC50EuTWPA3gaja1UjBsUP/D/9/juV3vQFr6XBEzn9KCAHdUvOHw==", "license": "MIT" }, - "node_modules/concat-map": { - "version": "0.0.1", - "resolved": "https://registry.npmjs.org/concat-map/-/concat-map-0.0.1.tgz", - "integrity": "sha512-/Srv4dswyQNBfohGpz9o6Yb3Gz3SrUDqBH5rTuhGR7ahtlbYKnVxw2bCFMRljaA7EXHaXZ8wsHdodFvbkhKmqg==", - "dev": true, - "license": "MIT" - }, "node_modules/copy-to-clipboard": { "version": "3.3.3", "resolved": "https://registry.npmjs.org/copy-to-clipboard/-/copy-to-clipboard-3.3.3.tgz", @@ -6924,22 +6873,6 @@ "node": ">=10.13.0" } }, - "node_modules/glob/node_modules/minimatch": { - "version": "10.1.1", - "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-10.1.1.tgz", - "integrity": "sha512-enIvLvRAFZYXJzkCYG5RKmPfrFArdLv+R+lbQ53BmIMLIry74bjKzX6iHAm8WYamJkhSSEabrWN5D97XnKObjQ==", - "dev": true, - "license": "BlueOak-1.0.0", - "dependencies": { - "@isaacs/brace-expansion": "^5.0.0" - }, - "engines": { - "node": "20 || >=22" - }, - "funding": { - "url": "https://github.com/sponsors/isaacs" - } - }, "node_modules/globals": { "version": "14.0.0", "resolved": "https://registry.npmjs.org/globals/-/globals-14.0.0.tgz", @@ -9035,16 +8968,19 @@ } }, "node_modules/minimatch": { - "version": "3.1.2", - "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-3.1.2.tgz", - "integrity": "sha512-J7p63hRiAjw1NDEww1W7i37+ByIrOWO5XQQAzZ3VOcL0PNybwpfmV/N05zFAzwQ9USyEcX6t3UO+K5aqBQOIHw==", + "version": "10.2.2", + "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-10.2.2.tgz", + "integrity": "sha512-+G4CpNBxa5MprY+04MbgOw1v7So6n5JY166pFi9KfYwT78fxScCeSNQSNzp6dpPSW2rONOps6Ocam1wFhCgoVw==", "dev": true, - "license": "ISC", + "license": "BlueOak-1.0.0", "dependencies": { - "brace-expansion": "^1.1.7" + "brace-expansion": "^5.0.2" }, "engines": { - "node": "*" + "node": "18 || 20 || >=22" + }, + "funding": { + "url": "https://github.com/sponsors/isaacs" } }, "node_modules/minimist": { @@ -12004,32 +11940,6 @@ "node": ">=18" } }, - "node_modules/test-exclude/node_modules/brace-expansion": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-2.0.2.tgz", - "integrity": "sha512-Jt0vHyM+jmUBqojB7E1NIYadt0vI0Qxjxd2TErW94wDz+E2LAm5vKMXXwg6ZZBTHPuUlDgQHKXvjGBdfcF1ZDQ==", - "dev": true, - "license": "MIT", - "dependencies": { - "balanced-match": "^1.0.0" - } - }, - "node_modules/test-exclude/node_modules/minimatch": { - "version": "9.0.5", - "resolved": "https://registry.npmjs.org/minimatch/-/minimatch-9.0.5.tgz", - "integrity": "sha512-G6T0ZX48xgozx7587koeX9Ys2NYy6Gmv//P89sEte9V9whIapMNF4idKxnW2QtCcLiTWlb/wfCabAtAFWhhBow==", - "dev": true, - "license": "ISC", - "dependencies": { - "brace-expansion": "^2.0.1" - }, - "engines": { - "node": ">=16 || 14 >=14.17" - }, - "funding": { - "url": "https://github.com/sponsors/isaacs" - } - }, "node_modules/thenify": { "version": "3.3.1", "resolved": "https://registry.npmjs.org/thenify/-/thenify-3.3.1.tgz", @@ -13085,21 +12995,6 @@ "type": "github", "url": "https://github.com/sponsors/wooorm" } - }, - "node_modules/@next/swc-win32-ia32-msvc": { - "version": "14.2.33", - "resolved": "https://registry.npmjs.org/@next/swc-win32-ia32-msvc/-/swc-win32-ia32-msvc-14.2.33.tgz", - "integrity": "sha512-pc9LpGNKhJ0dXQhZ5QMmYxtARwwmWLpeocFmVG5Z0DzWq5Uf0izcI8tLc+qOpqxO1PWqZ5A7J1blrUIKrIFc7Q==", - "cpu": [ - "ia32" - ], - "optional": true, - "os": [ - "win32" - ], - "engines": { - "node": ">= 10" - } } } } diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx index 1e8eabaea2e..b0df37ad6d8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx @@ -1,6 +1,6 @@ /* @vitest-environment jsdom */ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { render } from "@testing-library/react"; +import { act, render } from "@testing-library/react"; import { beforeEach, describe, expect, it, vi } from "vitest"; import ModelsAndEndpointsView from "./ModelsAndEndpointsView"; @@ -13,6 +13,8 @@ vi.mock("@/components/networking", () => ({ getCallbacksCall: vi.fn().mockResolvedValue({ router_settings: {} }), setCallbacksCall: vi.fn().mockResolvedValue(undefined), getUiSettings: vi.fn().mockResolvedValue({ values: {} }), + latestHealthChecksCall: vi.fn().mockResolvedValue({ latest_health_checks: {} }), + getModelCostMapReloadStatus: vi.fn().mockResolvedValue({}), })); vi.mock("@/app/(dashboard)/models-and-endpoints/components/ModelAnalyticsTab/ModelAnalyticsTab", () => ({ @@ -27,6 +29,14 @@ vi.mock("@/components/add_model/AddModelForm", () => ({ default: () => null, })); +const mockHealthCheckComponent = vi.fn((_props: { all_models_on_proxy?: string[] }) => null); +vi.mock("@/components/model_dashboard/HealthCheckComponent", () => ({ + default: (props: { all_models_on_proxy?: string[] }) => { + mockHealthCheckComponent(props); + return null; + }, +})); + vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({ default: () => ({ teams: [], @@ -104,4 +114,43 @@ describe("ModelsAndEndpointsView", () => { ); expect(await findByText("Model Management", {}, { timeout: 10000 })).toBeInTheDocument(); }, 15000); + + it("should pass model IDs (not model names) to HealthCheckComponent as all_models_on_proxy", async () => { + mockHealthCheckComponent.mockClear(); + const modelDataWithIds = { + data: [ + { model_name: "gpt-4", model_info: { id: "deployment-id-1" } }, + { model_name: "gpt-4", model_info: { id: "deployment-id-2" } }, + ], + }; + mockUseModelsInfo.mockReturnValue({ + data: { data: modelDataWithIds.data }, + isLoading: false, + refetch: vi.fn(), + }); + + const queryClient = createQueryClient(); + const { getByRole } = render( + + {}} + premiumUser={false} + teams={[]} + /> + , + ); + + const healthStatusTab = getByRole("tab", { name: "Health Status" }); + await act(async () => { + healthStatusTab.click(); + }); + + expect(mockHealthCheckComponent).toHaveBeenCalled(); + const healthCheckProps = mockHealthCheckComponent.mock.calls[0][0]; + expect(healthCheckProps.all_models_on_proxy).toEqual(["deployment-id-1", "deployment-id-2"]); + expect(healthCheckProps.all_models_on_proxy).not.toContain("gpt-4"); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx index 9d77774cb4c..b697a859dc5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx @@ -98,6 +98,13 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te return modelDataResponse.data.map((model: any) => model.model_name); }, [modelDataResponse?.data]); + const allModelIdsOnProxy = useMemo(() => { + if (!modelDataResponse?.data) return []; + return modelDataResponse.data + .map((model: any) => model.model_info?.id) + .filter((id: string | undefined): id is string => Boolean(id)); + }, [modelDataResponse?.data]); + const getProviderFromModel = (model: string) => { if (modelCostMapData !== null && modelCostMapData !== undefined) { if (typeof modelCostMapData == "object" && model in modelCostMapData) { @@ -397,7 +404,7 @@ const ModelsAndEndpointsView: React.FC = ({ premiumUser, te ) : page == "vector-stores" ? ( + ) : page == "tool-policies" ? ( + ) : page == "guardrails-monitor" ? ( ) : page == "new_usage" ? ( diff --git a/ui/litellm-dashboard/src/components/ToolPolicies.tsx b/ui/litellm-dashboard/src/components/ToolPolicies.tsx new file mode 100644 index 00000000000..0e3f5434e7f --- /dev/null +++ b/ui/litellm-dashboard/src/components/ToolPolicies.tsx @@ -0,0 +1,415 @@ +"use client"; + +import React, { useCallback, useDeferredValue, useEffect, useState } from "react"; +import { Select, Switch, Tooltip } from "antd"; +import { Select, Tooltip } from "antd"; +import { + Table, + TableHead, + TableHeaderCell, + TableBody, + TableRow, + TableCell, +} from "@tremor/react"; +import { TimeCell } from "./view_logs/time_cell"; +import { TableHeaderSortDropdown } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; +import type { SortState } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; +import FilterComponent, { FilterOption } from "./molecules/filter"; +import { fetchToolsList, updateToolPolicy, ToolRow } from "./networking"; + +const POLICY_OPTIONS = [ + { value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" }, + { value: "blocked", label: "blocked", color: "#991b1b", bg: "#fee2e2", border: "#fca5a5" }, +] as const; + +type PolicyValue = "trusted" | "blocked"; + +const policyStyle = (p: string) => + POLICY_OPTIONS.find((o) => o.value === p) ?? POLICY_OPTIONS[1]; + +type SortField = "tool_name" | "call_policy" | "team_id" | "key_alias" | "created_at" | "call_count"; + +interface FilterValues { + [key: string]: string; +} + +interface ToolPoliciesProps { + accessToken: string | null; + userRole?: string; +} + +const PolicySelect: React.FC<{ + value: string; + toolName: string; + saving: boolean; + onChange: (toolName: string, policy: string) => void; +}> = ({ value, toolName, saving, onChange }) => { + const style = policyStyle(value); + return ( + { setSearchTerm(e.target.value); setCurrentPage(1); }} + /> + + + + + +
+ Live Tail + +
+ + + + +
+ + Showing {filtered.length === 0 ? 0 : (currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, filtered.length)} of {filtered.length} results + + Page {currentPage} of {totalPages} +
+ + +
+
+ + + {/* Filter row */} +
+ +
+ + + {/* Auto-refresh banner */} + {isLiveTail && ( +
+ Auto-refreshing every 15 seconds + +
+ )} + + {error && ( +
{error}
+ )} + + {/* Table */} + + + + + + + + + Key Hash + + Origin + + + + {loading ? ( + + Loading tools… + + ) : paginated.length === 0 ? ( + + + No tools discovered yet. Make a chat completion that returns tool_calls to start auto-discovery. + + + ) : ( + paginated.map((tool) => ( + + + + + + + + {tool.tool_name} + + + + + + + + {(tool.call_count ?? 0).toLocaleString()} + + + + {tool.team_id ?? "-"} + + + + + + {tool.key_hash ?? "-"} + + + + + + {tool.key_alias ?? "-"} + + + + + {tool.origin ?? "-"} + + + + )) + )} + +
+ + {/* Bottom pagination (only when > 1 page) */} + {totalPages > 1 && ( +
+ Showing {(currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, sorted.length)} of {sorted.length} +
+ + +
+
+ )} + + + ); +}; + +export default ToolPolicies; diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx new file mode 100644 index 00000000000..85ce11e605a --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.test.tsx @@ -0,0 +1,85 @@ +import { screen } from "@testing-library/react"; +import { beforeAll, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../tests/test-utils"; +import UsageAIChatPanel from "./UsageAIChatPanel"; + +beforeAll(() => { + if (typeof window !== "undefined" && !window.ResizeObserver) { + window.ResizeObserver = class ResizeObserver { + observe() {} + unobserve() {} + disconnect() {} + } as any; + } +}); + +vi.mock("../../networking", () => ({ + modelHubCall: vi.fn().mockResolvedValue({ + data: [ + { model_group: "gpt-4" }, + { model_group: "claude-3-opus" }, + ], + }), + usageAiChatStream: vi.fn(), +})); + +const defaultProps = { + open: true, + onClose: vi.fn(), + accessToken: "test-token", +}; + +describe("UsageAIChatPanel", () => { + it("should render the panel when open", () => { + renderWithProviders(); + + expect(screen.getByText("Ask AI")).toBeInTheDocument(); + expect( + screen.getByText("Ask about your spend, models, keys, and trends") + ).toBeInTheDocument(); + }); + + it("should render model selector", () => { + renderWithProviders(); + + expect(screen.getByText("Select a model (optional, defaults to gpt-4o-mini)")).toBeInTheDocument(); + }); + + it("should render empty state message when no conversation", () => { + renderWithProviders(); + + expect(screen.getByText("Ask a question about your usage")).toBeInTheDocument(); + }); + + it("should render the send button", () => { + renderWithProviders(); + + expect(screen.getByText("Send")).toBeInTheDocument(); + }); + + it("should render input placeholder", () => { + renderWithProviders(); + + expect(screen.getByPlaceholderText("Ask about your usage...")).toBeInTheDocument(); + }); + + it("should render clear chat button", () => { + renderWithProviders(); + + expect(screen.getByText("Clear chat")).toBeInTheDocument(); + }); + + it("should have the panel element even when closed (just off-screen)", () => { + renderWithProviders(); + + expect(screen.getByTestId("usage-ai-chat-panel")).toBeInTheDocument(); + expect(screen.getByTestId("usage-ai-chat-panel")).toHaveClass("translate-x-full"); + }); + + it("should not have translate-x-full class when open", () => { + renderWithProviders(); + + expect(screen.getByTestId("usage-ai-chat-panel")).not.toHaveClass("translate-x-full"); + expect(screen.getByTestId("usage-ai-chat-panel")).toHaveClass("translate-x-0"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx new file mode 100644 index 00000000000..85f46bfa346 --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsageAIChatPanel.tsx @@ -0,0 +1,402 @@ +import React, { useEffect, useRef, useState } from "react"; +import { Button, Select, Input, Spin } from "antd"; +import ReactMarkdown from "react-markdown"; +import { modelHubCall, usageAiChatStream, UsageAiToolCallEvent } from "../../networking"; + +const { TextArea } = Input; + +interface ToolCallStep { + tool_name: string; + tool_label: string; + arguments: Record; + status: "running" | "complete" | "error"; + error?: string; +} + +interface ChatMessage { + role: "user" | "assistant"; + content: string; + toolCalls?: ToolCallStep[]; +} + +interface UsageAIChatPanelProps { + open: boolean; + onClose: () => void; + accessToken: string | null; +} + +const TOOL_ICONS: Record = { + get_usage_data: "📊", + get_team_usage_data: "👥", + get_tag_usage_data: "🏷️", +}; + +const ToolCallDisplay: React.FC<{ step: ToolCallStep }> = ({ step }) => { + const icon = TOOL_ICONS[step.tool_name] || "🔧"; + const args = step.arguments; + const dateRange = args.start_date && args.end_date + ? `${args.start_date} → ${args.end_date}` + : ""; + const filter = args.team_ids || args.tags || args.user_id || ""; + + return ( +
+ + {step.status === "running" ? ( + + ) : step.status === "error" ? ( + ✗ + ) : ( + ✓ + )} + +
+
+ {icon} {step.tool_label} +
+ {dateRange && ( +
{dateRange}
+ )} + {filter && ( +
Filter: {filter}
+ )} + {step.status === "error" && step.error && ( +
{step.error}
+ )} +
+
+ ); +}; + +const MarkdownContent: React.FC<{ content: string }> = ({ content }) => ( +

{children}

, + strong: ({ children }) => {children}, + ul: ({ children }) =>
    {children}
, + ol: ({ children }) =>
    {children}
, + li: ({ children }) =>
  • {children}
    • , + h1: ({ children }) =>

    {children}

    , + h2: ({ children }) =>

    {children}

    , + h3: ({ children }) =>

    {children}

    , + code: ({ children, className }) => { + const isBlock = className?.includes("language-"); + return isBlock ? ( + 
    +            {children}
+
    + ) : ( + {children} + ); + }, + table: ({ children }) => ( +
    + {children}
    +
    + ), + th: ({ children }) => {children}, + td: ({ children }) => {children}, + }} + > + {content} +     +); + +const UsageAIChatPanel: React.FC = ({ + open, + onClose, + accessToken, +}) => { + const [messages, setMessages] = useState([]); + const [inputText, setInputText] = useState(""); + const [isLoading, setIsLoading] = useState(false); + const [selectedModel, setSelectedModel] = useState(undefined); + const [availableModels, setAvailableModels] = useState([]); + const [isLoadingModels, setIsLoadingModels] = useState(false); + const [streamingContent, setStreamingContent] = useState(""); + const [statusMessage, setStatusMessage] = useState(null); + const [activeToolCalls, setActiveToolCalls] = useState([]); + const messagesEndRef = useRef(null); + const abortControllerRef = useRef(null); + + useEffect(() => { + if (open && availableModels.length === 0) { + loadModels(); + } + }, [open]); + + useEffect(() => { + if (typeof messagesEndRef.current?.scrollIntoView === "function") { + messagesEndRef.current.scrollIntoView({ behavior: "smooth" }); + } + }, [messages, streamingContent, activeToolCalls, statusMessage]); + + const loadModels = async () => { + if (!accessToken) return; + setIsLoadingModels(true); + try { + const fetchedModels = await modelHubCall(accessToken); + if (fetchedModels?.data?.length > 0) { + const models = fetchedModels.data + .map((item: any) => item.model_group as string) + .sort(); + setAvailableModels(models); + } + } catch (error) { + console.error("Failed to load models:", error); + } finally { + setIsLoadingModels(false); + } + }; + + const handleSend = async () => { + if (!accessToken || !inputText.trim() || isLoading) return; + + const userMessage: ChatMessage = { role: "user", content: inputText.trim() }; + const updatedMessages = [...messages, userMessage]; + setMessages(updatedMessages); + setInputText(""); + setIsLoading(true); + setStreamingContent(""); + setStatusMessage(null); + setActiveToolCalls([]); + + const abortController = new AbortController(); + abortControllerRef.current = abortController; + + let accumulated = ""; + const toolCalls: ToolCallStep[] = []; + + try { + await usageAiChatStream( + accessToken, + updatedMessages.slice(-20).map((m) => ({ role: m.role, content: m.content })), + selectedModel || "", + (content: string) => { + setStatusMessage(null); + accumulated += content; + setStreamingContent(accumulated); + }, + () => { + setStatusMessage(null); + setActiveToolCalls([]); + setMessages((prev) => [ + ...prev, + { role: "assistant", content: accumulated, toolCalls: toolCalls.length > 0 ? [...toolCalls] : undefined }, + ]); + setStreamingContent(""); + }, + (errorMsg: string) => { + setStatusMessage(null); + setActiveToolCalls([]); + setMessages((prev) => [ + ...prev, + { role: "assistant", content: `Error: ${errorMsg}` }, + ]); + setStreamingContent(""); + }, + (status: string) => { + setStatusMessage(status); + }, + (event: UsageAiToolCallEvent) => { + const idx = toolCalls.findIndex((tc) => tc.tool_name === event.tool_name); + if (idx >= 0) { + toolCalls[idx] = { ...event }; + } else { + toolCalls.push({ ...event }); + } + setActiveToolCalls([...toolCalls]); + }, + abortController.signal, + ); + } catch (error: any) { + if (error?.name === "AbortError" || abortController.signal.aborted) { + return; + } + const errorMsg = error?.message || "Failed to get response. Please try again."; + setMessages((prev) => [ + ...prev, + { role: "assistant", content: `Error: ${errorMsg}` }, + ]); + setStreamingContent(""); + } finally { + setIsLoading(false); + abortControllerRef.current = null; + } + }; + + const handleKeyDown = (e: React.KeyboardEvent) => { + if (e.key === "Enter" && !e.shiftKey) { + e.preventDefault(); + handleSend(); + } + }; + + const handleClose = () => { + if (abortControllerRef.current) { + abortControllerRef.current.abort(); + } + onClose(); + }; + + const handleClear = () => { + setMessages([]); + setStreamingContent(""); + setActiveToolCalls([]); + setStatusMessage(null); + }; + + return ( +
    + {/* Header */} +
    +
    +
    + + + +

    Ask AI

    +
    + +
    +

    + Ask about your spend, models, keys, and trends +

    +
    + + {/* Model selector */} +
    +