diff --git a/docs/my-website/docs/observability/datadog.md b/docs/my-website/docs/observability/datadog.md index 7cf91ced34..6f785be101 100644 --- a/docs/my-website/docs/observability/datadog.md +++ b/docs/my-website/docs/observability/datadog.md @@ -7,6 +7,7 @@ import TabItem from '@theme/TabItem'; LiteLLM Supports logging to the following Datdog Integrations: - `datadog` [Datadog Logs](https://docs.datadoghq.com/logs/) - `datadog_llm_observability` [Datadog LLM Observability](https://www.datadoghq.com/product/llm-observability/) +- `datadog_cost_management` [Datadog Cloud Cost Management](#datadog-cloud-cost-management) - `ddtrace-run` [Datadog Tracing](#datadog-tracing) ## Datadog Logs @@ -73,7 +74,7 @@ Send logs through a local DataDog agent (useful for containerized environments): ```shell LITELLM_DD_AGENT_HOST="localhost" # hostname or IP of DataDog agent LITELLM_DD_AGENT_PORT="10518" # [OPTIONAL] port of DataDog agent (default: 10518) -DD_API_KEY="5f2d0f310***********" # [OPTIONAL] your datadog API Key (agent handles auth) +DD_API_KEY="5f2d0f310***********" # [OPTIONAL] your datadog API Key (Agent handles auth for Logs. REQUIRED for LLM Observability) DD_SOURCE="litellm_dev" # [OPTIONAL] your datadog source ``` @@ -84,6 +85,9 @@ When `LITELLM_DD_AGENT_HOST` is set, logs are sent to the agent instead of direc **Note:** We use `LITELLM_DD_AGENT_HOST` instead of `DD_AGENT_HOST` to avoid conflicts with `ddtrace` which automatically sets `DD_AGENT_HOST` for APM tracing. +> [!IMPORTANT] +> **Datadog LLM Observability**: `DD_API_KEY` is **REQUIRED** even when using the Datadog Agent (`LITELLM_DD_AGENT_HOST`). The agent acts as a proxy but the API key header is mandatory for the LLM Observability endpoint. + **Step 3**: Start the proxy, make a test request Start proxy @@ -161,6 +165,50 @@ On the Datadog LLM Observability page, you should see that both input messages a + + + +## Datadog Cloud Cost Management + +| Feature | Details | +|---------|---------| +| **What is logged** | Aggregated LLM Costs (FOCUS format) | +| **Events** | Periodic Uploads of Aggregated Cost Data | +| **Product Link** | [Datadog Cloud Cost Management](https://docs.datadoghq.com/cost_management/) | + +We will use the `--config` to set `litellm.callbacks = ["datadog_cost_management"]`. This will periodically upload aggregated LLM cost data to Datadog. + +**Step 1**: Create a `config.yaml` file and set `litellm_settings`: `success_callback` + +```yaml +model_list: + - model_name: gpt-3.5-turbo + litellm_params: + model: gpt-3.5-turbo +litellm_settings: + callbacks: ["datadog_cost_management"] +``` + +**Step 2**: Set Required env variables + +```shell +DD_API_KEY="your-api-key" +DD_APP_KEY="your-app-key" # REQUIRED for Cost Management +DD_SITE="us5.datadoghq.com" +``` + +**Step 3**: Start the proxy + +```shell +litellm --config config.yaml +``` + +**How it works** +* LiteLLM aggregates costs in-memory by Provider, Model, Date, and Tags. +* Requires `DD_APP_KEY` for the Custom Costs API. +* Costs are uploaded periodically (flushed). + + ### Datadog Tracing Use `ddtrace-run` to enable [Datadog Tracing](https://ddtrace.readthedocs.io/en/stable/installation_quickstart.html) on litellm proxy @@ -203,5 +251,5 @@ LiteLLM supports customizing the following Datadog environment variables | `POD_NAME` | Pod name tag (useful for Kubernetes deployments) | "unknown" | ❌ No | \* **Required when using Direct API** (default): `DD_API_KEY` and `DD_SITE` are required -\* **Optional when using DataDog Agent**: Set `LITELLM_DD_AGENT_HOST` to use agent mode; `DD_API_KEY` and `DD_SITE` are not required +\* **Optional when using DataDog Agent**: Set `LITELLM_DD_AGENT_HOST` to use agent mode; `DD_API_KEY` and `DD_SITE` are not required for **Datadog Logs**. (**Note: `DD_API_KEY` IS REQUIRED for Datadog LLM Observability**) diff --git a/docs/my-website/docs/providers/vercel_ai_gateway.md b/docs/my-website/docs/providers/vercel_ai_gateway.md index 91f0a18ea1..3ff007171e 100644 --- a/docs/my-website/docs/providers/vercel_ai_gateway.md +++ b/docs/my-website/docs/providers/vercel_ai_gateway.md @@ -11,7 +11,7 @@ import TabItem from '@theme/TabItem'; | Provider Route on LiteLLM | `vercel_ai_gateway/` | | Link to Provider Doc | [Vercel AI Gateway Documentation ↗](https://vercel.com/docs/ai-gateway) | | Base URL | `https://ai-gateway.vercel.sh/v1` | -| Supported Operations | `/chat/completions`, `/models` | +| Supported Operations | `/chat/completions`, `/embeddings`, `/models` |

@@ -73,7 +73,7 @@ messages = [{"content": "Hello, how are you?", "role": "user"}] # Vercel AI Gateway call with streaming response = completion( - model="vercel_ai_gateway/openai/gpt-4o", + model="vercel_ai_gateway/openai/gpt-4o", messages=messages, stream=True ) @@ -82,6 +82,33 @@ for chunk in response: print(chunk) ``` +### Embeddings + +```python showLineNumbers title="Vercel AI Gateway Embeddings" +import os +from litellm import embedding + +os.environ["VERCEL_AI_GATEWAY_API_KEY"] = "your-api-key" + +# Vercel AI Gateway embedding call +response = embedding( + model="vercel_ai_gateway/openai/text-embedding-3-small", + input="Hello world" +) + +print(response.data[0]["embedding"][:5]) # Print first 5 dimensions +``` + +You can also specify the `dimensions` parameter: + +```python showLineNumbers title="Vercel AI Gateway Embeddings with Dimensions" +response = embedding( + model="vercel_ai_gateway/openai/text-embedding-3-small", + input=["Hello world", "Goodbye world"], + dimensions=768 +) +``` + ## Usage - LiteLLM Proxy Add the following to your LiteLLM Proxy configuration file: @@ -97,6 +124,11 @@ model_list: litellm_params: model: vercel_ai_gateway/anthropic/claude-4-sonnet api_key: os.environ/VERCEL_AI_GATEWAY_API_KEY + + - model_name: text-embedding-3-small-gateway + litellm_params: + model: vercel_ai_gateway/openai/text-embedding-3-small + api_key: os.environ/VERCEL_AI_GATEWAY_API_KEY ``` Start your LiteLLM Proxy server: diff --git a/litellm/integrations/callback_configs.json b/litellm/integrations/callback_configs.json index 6b30b6b736..6a003b8c49 100644 --- a/litellm/integrations/callback_configs.json +++ b/litellm/integrations/callback_configs.json @@ -83,6 +83,33 @@ }, "description": "Datadog Logging Integration" }, + { + "id": "datadog_cost_management", + "displayName": "Datadog Cost Management", + "logo": "datadog.png", + "supports_key_team_logging": false, + "dynamic_params": { + "dd_api_key": { + "type": "password", + "ui_name": "API Key", + "description": "Datadog API key for authentication", + "required": true + }, + "dd_app_key": { + "type": "password", + "ui_name": "App Key", + "description": "Datadog Application Key for Cloud Cost Management", + "required": true + }, + "dd_site": { + "type": "text", + "ui_name": "Site", + "description": "Datadog site URL (e.g., us5.datadoghq.com)", + "required": true + } + }, + "description": "Datadog Cloud Cost Management Integration" + }, { "id": "lago", "displayName": "Lago", @@ -407,4 +434,4 @@ }, "description": "SQS Queue (AWS) Logging Integration" } -] +] \ No newline at end of file diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 6a76b57e7f..a5bb530fc5 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -516,7 +516,9 @@ class CustomGuardrail(CustomLogger): from litellm.types.utils import GuardrailMode # Use event_type if provided, otherwise fall back to self.event_hook - guardrail_mode: Union[GuardrailEventHooks, GuardrailMode, List[GuardrailEventHooks]] + guardrail_mode: Union[ + GuardrailEventHooks, GuardrailMode, List[GuardrailEventHooks] + ] if event_type is not None: guardrail_mode = event_type elif isinstance(self.event_hook, Mode): @@ -524,11 +526,21 @@ class CustomGuardrail(CustomLogger): else: guardrail_mode = self.event_hook # type: ignore[assignment] + from litellm.litellm_core_utils.core_helpers import ( + filter_exceptions_from_params, + ) + + # Sanitize the response to ensure it's JSON serializable and free of circular refs + # This prevents RecursionErrors in downstream loggers (Langfuse, Datadog, etc.) + clean_guardrail_response = filter_exceptions_from_params( + guardrail_json_response + ) + slg = StandardLoggingGuardrailInformation( guardrail_name=self.guardrail_name, guardrail_provider=guardrail_provider, guardrail_mode=guardrail_mode, - guardrail_response=guardrail_json_response, + guardrail_response=clean_guardrail_response, guardrail_status=guardrail_status, start_time=start_time, end_time=end_time, diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 503e8d8c87..735d1005d2 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -32,6 +32,7 @@ from litellm.integrations.datadog.datadog_handler import ( get_datadog_service, get_datadog_source, get_datadog_tags, + get_datadog_base_url_from_env, ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.llms.custom_httpx.http_handler import ( @@ -100,7 +101,9 @@ class DataDogLogger( self._configure_dd_direct_api() # Optional override for testing - self._apply_dd_base_url_override() + dd_base_url = get_datadog_base_url_from_env() + if dd_base_url: + self.intake_url = f"{dd_base_url}/api/v2/logs" self.sync_client = _get_httpx_client() asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() @@ -159,18 +162,6 @@ class DataDogLogger( self.DD_API_KEY = os.getenv("DD_API_KEY") self.intake_url = f"https://http-intake.logs.{os.getenv('DD_SITE')}/api/v2/logs" - def _apply_dd_base_url_override(self) -> None: - """ - Apply base URL override for testing purposes - """ - dd_base_url: Optional[str] = ( - os.getenv("_DATADOG_BASE_URL") - or os.getenv("DATADOG_BASE_URL") - or os.getenv("DD_BASE_URL") - ) - if dd_base_url is not None: - self.intake_url = f"{dd_base_url}/api/v2/logs" - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ Async Log success events to Datadog diff --git a/litellm/integrations/datadog/datadog_cost_management.py b/litellm/integrations/datadog/datadog_cost_management.py new file mode 100644 index 0000000000..2eb94b59dd --- /dev/null +++ b/litellm/integrations/datadog/datadog_cost_management.py @@ -0,0 +1,204 @@ +import asyncio +import os +import time +from datetime import datetime +from typing import Dict, List, Optional, Tuple + +from litellm._logging import verbose_logger +from litellm.integrations.custom_batch_logger import CustomBatchLogger +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.integrations.datadog_cost_management import ( + DatadogFOCUSCostEntry, +) +from litellm.types.utils import StandardLoggingPayload + + +class DatadogCostManagementLogger(CustomBatchLogger): + def __init__(self, **kwargs): + self.dd_api_key = os.getenv("DD_API_KEY") + self.dd_app_key = os.getenv("DD_APP_KEY") + self.dd_site = os.getenv("DD_SITE", "datadoghq.com") + + if not self.dd_api_key or not self.dd_app_key: + verbose_logger.warning( + "Datadog Cost Management: DD_API_KEY and DD_APP_KEY are required. Integration will not work." + ) + + self.upload_url = f"https://api.{self.dd_site}/api/v2/cost/custom_costs" + + self.async_client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback + ) + + # Initialize lock and start periodic flush task + self.flush_lock = asyncio.Lock() + asyncio.create_task(self.periodic_flush()) + + # Check if flush_lock is already in kwargs to avoid double passing (unlikely but safe) + if "flush_lock" not in kwargs: + kwargs["flush_lock"] = self.flush_lock + + super().__init__(**kwargs) + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + try: + standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object", None + ) + + if standard_logging_object is None: + return + + # Only log if there is a cost associated + if standard_logging_object.get("response_cost", 0) > 0: + self.log_queue.append(standard_logging_object) + + if len(self.log_queue) >= self.batch_size: + await self.async_send_batch() + + except Exception as e: + verbose_logger.exception( + f"Datadog Cost Management: Error in async_log_success_event: {str(e)}" + ) + + async def async_send_batch(self): + if not self.log_queue: + return + + try: + # Aggregate costs from the batch + aggregated_entries = self._aggregate_costs(self.log_queue) + + if not aggregated_entries: + return + + # Send to Datadog + await self._upload_to_datadog(aggregated_entries) + + # Clear queue only on success (or if we decide to drop on failure) + # CustomBatchLogger clears queue in flush_queue, so we just process here + + except Exception as e: + verbose_logger.exception( + f"Datadog Cost Management: Error in async_send_batch: {str(e)}" + ) + + def _aggregate_costs( + self, logs: List[StandardLoggingPayload] + ) -> List[DatadogFOCUSCostEntry]: + """ + Aggregates costs by Provider, Model, and Date. + Returns a list of DatadogFOCUSCostEntry. + """ + aggregator: Dict[Tuple[str, str, str, Tuple[Tuple[str, str], ...]], DatadogFOCUSCostEntry] = {} + + for log in logs: + try: + # Extract keys for aggregation + provider = log.get("custom_llm_provider") or "unknown" + model = log.get("model") or "unknown" + cost = log.get("response_cost", 0) + + if cost == 0: + continue + + # Get date strings (FOCUS format requires specific keys, but for aggregation we group by Day) + # UTC date + # We interpret "ChargePeriod" as the day of the request. + ts = log.get("startTime") or time.time() + dt = datetime.fromtimestamp(ts) + date_str = dt.strftime("%Y-%m-%d") + + # ChargePeriodStart and End + # If we want daily granularity, end date is usually same day or next day? + # Datadog Custom Costs usually expects periods. + # "ChargePeriodStart": "2023-01-01", "ChargePeriodEnd": "2023-12-31" in example. + # If we send daily, we can say Start=Date, End=Date. + + # Grouping Key: Provider + Model + Date + Tags? + # For simplicity, let's aggregate by Provider + Model + Date first. + # If we handle tags, we need to include them in the key. + + tags = self._extract_tags(log) + tags_key = tuple(sorted(tags.items())) if tags else () + + key = (provider, model, date_str, tags_key) + + if key not in aggregator: + aggregator[key] = { + "ProviderName": provider, + "ChargeDescription": f"LLM Usage for {model}", + "ChargePeriodStart": date_str, + "ChargePeriodEnd": date_str, + "BilledCost": 0.0, + "BillingCurrency": "USD", + "Tags": tags if tags else None, + } + + aggregator[key]["BilledCost"] += cost + + except Exception as e: + verbose_logger.warning( + f"Error processing log for cost aggregation: {e}" + ) + continue + + return list(aggregator.values()) + + def _extract_tags(self, log: StandardLoggingPayload) -> Dict[str, str]: + from litellm.integrations.datadog.datadog_handler import ( + get_datadog_env, + get_datadog_hostname, + get_datadog_pod_name, + get_datadog_service, + ) + + tags = { + "env": get_datadog_env(), + "service": get_datadog_service(), + "host": get_datadog_hostname(), + "pod_name": get_datadog_pod_name(), + } + + # Add metadata as tags + metadata = log.get("metadata", {}) + if metadata: + # Add user info + if "user_api_key_alias" in metadata: + tags["user"] = str(metadata["user_api_key_alias"]) + if "user_api_key_team_alias" in metadata: + tags["team"] = str(metadata["user_api_key_team_alias"]) + # model_group is not in StandardLoggingMetadata TypedDict, so we need to access it via dict.get() + model_group = metadata.get("model_group") # type: ignore[misc] + if model_group: + tags["model_group"] = str(model_group) + + return tags + + async def _upload_to_datadog(self, payload: List[Dict]): + if not self.dd_api_key or not self.dd_app_key: + return + + headers = { + "Content-Type": "application/json", + "DD-API-KEY": self.dd_api_key, + "DD-APPLICATION-KEY": self.dd_app_key, + } + + # The API endpoint expects a list of objects directly in the body (file content behavior) + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + + data_json = safe_dumps(payload) + + response = await self.async_client.put( + self.upload_url, content=data_json, headers=headers + ) + + response.raise_for_status() + + verbose_logger.debug( + f"Datadog Cost Management: Uploaded {len(payload)} cost entries. Status: {response.status_code}" + ) diff --git a/litellm/integrations/datadog/datadog_handler.py b/litellm/integrations/datadog/datadog_handler.py index 26fab77759..e2f30f2f61 100644 --- a/litellm/integrations/datadog/datadog_handler.py +++ b/litellm/integrations/datadog/datadog_handler.py @@ -20,6 +20,14 @@ def get_datadog_hostname() -> str: return os.getenv("HOSTNAME", "") +def get_datadog_base_url_from_env() -> Optional[str]: + """ + Get base URL override from common DD_BASE_URL env var. + This is useful for testing or custom endpoints. + """ + return os.getenv("DD_BASE_URL") + + def get_datadog_env() -> str: return os.getenv("DD_ENV", "unknown") diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 6ffdbc0a00..4f6a5b339a 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -21,6 +21,7 @@ from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.integrations.datadog.datadog_handler import ( get_datadog_service, get_datadog_tags, + get_datadog_base_url_from_env, ) from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -43,24 +44,22 @@ class DataDogLLMObsLogger(CustomBatchLogger): def __init__(self, **kwargs): try: verbose_logger.debug("DataDogLLMObs: Initializing logger") - if os.getenv("DD_API_KEY", None) is None: - raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>'") - if os.getenv("DD_SITE", None) is None: - raise Exception( - "DD_SITE is not set, set 'DD_SITE=<>', example sit = `us5.datadoghq.com`" - ) + # Configure DataDog endpoint (Agent or Direct API) + # Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST + dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST") self.async_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback ) self.DD_API_KEY = os.getenv("DD_API_KEY") - self.DD_SITE = os.getenv("DD_SITE") - self.intake_url = ( - f"https://api.{self.DD_SITE}/api/intake/llm-obs/v1/trace/spans" - ) - # testing base url - dd_base_url = os.getenv("DD_BASE_URL") + if dd_agent_host: + self._configure_dd_agent(dd_agent_host=dd_agent_host) + else: + self._configure_dd_direct_api() + + # Optional override for testing + dd_base_url = get_datadog_base_url_from_env() if dd_base_url: self.intake_url = f"{dd_base_url}/api/intake/llm-obs/v1/trace/spans" @@ -78,6 +77,38 @@ class DataDogLLMObsLogger(CustomBatchLogger): verbose_logger.exception(f"DataDogLLMObs: Error initializing - {str(e)}") raise e + def _configure_dd_agent(self, dd_agent_host: str): + """ + Configure the Datadog logger to send traces to the Agent. + """ + # When using the Agent, LLM Observability Intake does NOT require the API Key + # Reference: https://docs.datadoghq.com/llm_observability/setup/sdk/#agent-setup + + # Use specific port for LLM Obs (Trace Agent) to avoid conflict with Logs Agent (10518) + agent_port = os.getenv("LITELLM_DD_LLM_OBS_PORT", "8126") + self.DD_SITE = "localhost" # Not used for URL construction in agent mode + self.intake_url = ( + f"http://{dd_agent_host}:{agent_port}/api/intake/llm-obs/v1/trace/spans" + ) + verbose_logger.debug(f"DataDogLLMObs: Using DD Agent at {self.intake_url}") + + def _configure_dd_direct_api(self): + """ + Configure the Datadog logger to send traces directly to the Datadog API. + """ + if not self.DD_API_KEY: + raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>'") + + self.DD_SITE = os.getenv("DD_SITE") + if not self.DD_SITE: + raise Exception( + "DD_SITE is not set, set 'DD_SITE=<>', example site = `us5.datadoghq.com`" + ) + + self.intake_url = ( + f"https://api.{self.DD_SITE}/api/intake/llm-obs/v1/trace/spans" + ) + def _get_datadog_llm_obs_params(self) -> Dict: """ Get the datadog_llm_observability_params from litellm.datadog_llm_observability_params @@ -164,13 +195,14 @@ class DataDogLLMObsLogger(CustomBatchLogger): json_payload = safe_dumps(payload) + headers = {"Content-Type": "application/json"} + if self.DD_API_KEY: + headers["DD-API-KEY"] = self.DD_API_KEY + response = await self.async_client.post( url=self.intake_url, content=json_payload, - headers={ - "DD-API-KEY": self.DD_API_KEY, - "Content-Type": "application/json", - }, + headers=headers, ) if response.status_code != 202: diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 46ada3c393..7bf97665fd 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -23,6 +23,7 @@ from litellm.constants import MAX_LANGFUSE_INITIALIZED_CLIENTS from litellm.litellm_core_utils.core_helpers import ( safe_deep_copy, reconstruct_model_name, + filter_exceptions_from_params, ) from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info from litellm.integrations.langfuse.langfuse_mock_client import ( @@ -75,9 +76,8 @@ def _extract_cache_read_input_tokens(usage_obj) -> int: # Check prompt_tokens_details.cached_tokens (used by Gemini and other providers) if hasattr(usage_obj, "prompt_tokens_details"): prompt_tokens_details = getattr(usage_obj, "prompt_tokens_details", None) - if ( - prompt_tokens_details is not None - and hasattr(prompt_tokens_details, "cached_tokens") + if prompt_tokens_details is not None and hasattr( + prompt_tokens_details, "cached_tokens" ): cached_tokens = getattr(prompt_tokens_details, "cached_tokens", None) if ( @@ -540,7 +540,6 @@ class LangFuseLogger: verbose_logger.debug("Langfuse Layer Logging - logging to langfuse v2") try: - metadata = metadata or {} standard_logging_object: Optional[StandardLoggingPayload] = cast( Optional[StandardLoggingPayload], kwargs.get("standard_logging_object", None), @@ -706,9 +705,10 @@ class LangFuseLogger: clean_metadata["litellm_response_cost"] = cost if standard_logging_object is not None: - clean_metadata["hidden_params"] = standard_logging_object[ - "hidden_params" - ] + hidden_params = standard_logging_object.get("hidden_params", {}) + clean_metadata["hidden_params"] = filter_exceptions_from_params( + hidden_params + ) if ( litellm.langfuse_default_tags is not None diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index fadeeffa9c..361f0ade74 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3299,6 +3299,7 @@ def _get_masked_values( "token", "key", "secret", + "vertex_credentials", ] return { k: ( diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index bbe28e3ec2..25ad0a570c 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -21,11 +21,13 @@ from litellm.types.utils import ( ChatCompletionMessageToolCall, ChatCompletionRedactedThinkingBlock, Choices, + CompletionTokensDetailsWrapper, Delta, EmbeddingResponse, Function, HiddenParams, ImageResponse, + PromptTokensDetailsWrapper, ) from litellm.types.utils import Logprobs as TextCompletionLogprobs from litellm.types.utils import ( @@ -304,6 +306,22 @@ class LiteLLMResponseObjectHandler: "text_tokens": 0, } + # Map Responses API naming to Chat Completions API naming for cost calculator + if usage.get("prompt_tokens") is None: + usage["prompt_tokens"] = usage.get("input_tokens", 0) + if usage.get("completion_tokens") is None: + usage["completion_tokens"] = usage.get("output_tokens", 0) + + # Convert dicts to wrapper objects so getattr() works in cost calculation + if isinstance(usage.get("input_tokens_details"), dict): + usage["prompt_tokens_details"] = PromptTokensDetailsWrapper( + **usage["input_tokens_details"] + ) + if isinstance(usage.get("output_tokens_details"), dict): + usage["completion_tokens_details"] = CompletionTokensDetailsWrapper( + **usage["output_tokens_details"] + ) + if model_response_object is None: model_response_object = ImageResponse(**response_object) return model_response_object diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 30263543fc..1d1e38c09d 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -4408,7 +4408,7 @@ def _bedrock_tools_pt(tools: List) -> List[BedrockToolBlock]: ] """ """ - Bedrock toolConfig looks like: + Bedrock toolConfig looks like: "tools": [ { "toolSpec": { @@ -4436,6 +4436,7 @@ def _bedrock_tools_pt(tools: List) -> List[BedrockToolBlock]: tool_block_list: List[BedrockToolBlock] = [] for tool in tools: + # Handle regular function tools parameters = tool.get("function", {}).get( "parameters", {"type": "object", "properties": {}} ) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index cb26a22edf..ec66514207 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -298,6 +298,39 @@ class AmazonConverseConfig(BaseConfig): # Check if the model is specifically Nova Lite 2 return "nova-2-lite" in model_without_region + def _map_web_search_options( + self, + web_search_options: dict, + model: str + ) -> Optional[BedrockToolBlock]: + """ + Map web_search_options to Nova grounding systemTool. + + Nova grounding (web search) is only supported on Amazon Nova models. + Returns None for non-Nova models. + + Args: + web_search_options: The web_search_options dict from the request + model: The model identifier string + + Returns: + BedrockToolBlock with systemTool for Nova models, None otherwise + + Reference: https://docs.aws.amazon.com/nova/latest/userguide/grounding.html + """ + # Only Nova models support nova_grounding + # Model strings can be like: "amazon.nova-pro-v1:0", "us.amazon.nova-pro-v1:0", etc. + if "nova" not in model.lower(): + verbose_logger.debug( + f"web_search_options passed but model {model} is not a Nova model. " + "Nova grounding is only supported on Amazon Nova models." + ) + return None + + # Nova doesn't support search_context_size or user_location params + # (unlike Anthropic), so we just enable grounding with no options + return BedrockToolBlock(systemTool={"name": "nova_grounding"}) + def _transform_reasoning_effort_to_reasoning_config( self, reasoning_effort: str ) -> dict: @@ -438,6 +471,10 @@ class AmazonConverseConfig(BaseConfig): ): supported_params.append("tools") + # Nova models support web_search_options (mapped to nova_grounding systemTool) + if base_model.startswith("amazon.nova"): + supported_params.append("web_search_options") + if litellm.utils.supports_tool_choice( model=model, custom_llm_provider=self.custom_llm_provider ) or litellm.utils.supports_tool_choice( @@ -730,6 +767,13 @@ class AmazonConverseConfig(BaseConfig): if bedrock_tier in ("default", "flex", "priority"): optional_params["serviceTier"] = {"type": bedrock_tier} + if param == "web_search_options" and value and isinstance(value, dict): + grounding_tool = self._map_web_search_options(value, model) + if grounding_tool is not None: + optional_params = self._add_tools_to_optional_params( + optional_params=optional_params, tools=[grounding_tool] + ) + # Only update thinking tokens for non-GPT-OSS models and non-Nova-Lite-2 models # Nova Lite 2 handles token budgeting differently through reasoningConfig if "gpt-oss" not in model and not self._is_nova_lite_2_model(model): @@ -1388,20 +1432,23 @@ class AmazonConverseConfig(BaseConfig): str, List[ChatCompletionToolCallChunk], Optional[List[BedrockConverseReasoningContentBlock]], + Optional[List[CitationsContentBlock]], ]: """ - Translate the message content to a string and a list of tool calls and reasoning content blocks + Translate the message content to a string and a list of tool calls, reasoning content blocks, and citations. Returns: content_str: str tools: List[ChatCompletionToolCallChunk] reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] + citationsContentBlocks: Optional[List[CitationsContentBlock]] - Citations from Nova grounding """ content_str = "" tools: List[ChatCompletionToolCallChunk] = [] reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = ( None ) + citationsContentBlocks: Optional[List[CitationsContentBlock]] = None for idx, content in enumerate(content_blocks): """ - Content is either a tool response or text @@ -1446,10 +1493,15 @@ class AmazonConverseConfig(BaseConfig): if reasoningContentBlocks is None: reasoningContentBlocks = [] reasoningContentBlocks.append(content["reasoningContent"]) + # Handle Nova grounding citations content + if "citationsContent" in content: + if citationsContentBlocks is None: + citationsContentBlocks = [] + citationsContentBlocks.append(content["citationsContent"]) - return content_str, tools, reasoningContentBlocks + return content_str, tools, reasoningContentBlocks, citationsContentBlocks - def _transform_response( + def _transform_response( # noqa: PLR0915 self, model: str, response: httpx.Response, @@ -1525,18 +1577,27 @@ class AmazonConverseConfig(BaseConfig): reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = ( None ) + citationsContentBlocks: Optional[List[CitationsContentBlock]] = None if message is not None: ( content_str, tools, reasoningContentBlocks, + citationsContentBlocks, ) = self._translate_message_content(message["content"]) + # Initialize provider_specific_fields if we have any special content blocks + provider_specific_fields: dict = {} + if reasoningContentBlocks is not None: + provider_specific_fields["reasoningContentBlocks"] = reasoningContentBlocks + if citationsContentBlocks is not None: + provider_specific_fields["citationsContent"] = citationsContentBlocks + + if provider_specific_fields: + chat_completion_message["provider_specific_fields"] = provider_specific_fields + if reasoningContentBlocks is not None: - chat_completion_message["provider_specific_fields"] = { - "reasoningContentBlocks": reasoningContentBlocks, - } chat_completion_message["reasoning_content"] = ( self._transform_reasoning_content(reasoningContentBlocks) ) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index 17474fa022..1c58a11eeb 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -1476,6 +1476,11 @@ class AWSEventStreamDecoder: reasoning_content = ( "" # set to non-empty string to ensure consistency with Anthropic ) + elif "citationsContent" in delta_obj: + # Handle Nova grounding citations in streaming responses + provider_specific_fields = { + "citationsContent": delta_obj["citationsContent"], + } return ( text, tool_use, diff --git a/litellm/llms/openai/image_generation/cost_calculator.py b/litellm/llms/openai/image_generation/cost_calculator.py index 35caaf6e9b..988d562613 100644 --- a/litellm/llms/openai/image_generation/cost_calculator.py +++ b/litellm/llms/openai/image_generation/cost_calculator.py @@ -8,8 +8,7 @@ from typing import Optional from litellm import verbose_logger from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token -from litellm.responses.utils import ResponseAPILoggingUtils -from litellm.types.utils import ImageResponse +from litellm.types.utils import ImageResponse, Usage def cost_calculator( @@ -39,11 +38,18 @@ def cost_calculator( ) return 0.0 - # Transform ImageUsage to Usage using the existing helper - # ImageUsage has the same format as ResponseAPIUsage - chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( - usage - ) + # If usage is already a Usage object with completion_tokens_details set, + # use it directly (it was already transformed in convert_to_image_response) + if isinstance(usage, Usage) and usage.completion_tokens_details is not None: + chat_usage = usage + else: + # Transform ImageUsage to Usage using the existing helper + # ImageUsage has the same format as ResponseAPIUsage + from litellm.responses.utils import ResponseAPILoggingUtils + + chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + usage + ) # Use generic_cost_per_token for cost calculation prompt_cost, completion_cost = generic_cost_per_token( diff --git a/litellm/llms/vercel_ai_gateway/embedding/__init__.py b/litellm/llms/vercel_ai_gateway/embedding/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/litellm/llms/vercel_ai_gateway/embedding/transformation.py b/litellm/llms/vercel_ai_gateway/embedding/transformation.py new file mode 100644 index 0000000000..7238b05f10 --- /dev/null +++ b/litellm/llms/vercel_ai_gateway/embedding/transformation.py @@ -0,0 +1,176 @@ +""" +Vercel AI Gateway Embedding API Configuration. + +This module provides the configuration for Vercel AI Gateway's Embedding API. +Vercel AI Gateway is OpenAI-compatible and supports embeddings via the /v1/embeddings endpoint. + +Docs: https://vercel.com/docs/ai-gateway/openai-compat/embeddings +""" + +from typing import TYPE_CHECKING, Any, Optional + +import httpx + +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllEmbeddingInputValues +from litellm.types.utils import EmbeddingResponse +from litellm.utils import convert_to_model_response_object + +from ..common_utils import VercelAIGatewayException + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class VercelAIGatewayEmbeddingConfig(BaseEmbeddingConfig): + """ + Configuration for Vercel AI Gateway's Embedding API. + + Reference: https://vercel.com/docs/ai-gateway/openai-compat/embeddings + """ + + def validate_environment( + self, + headers: dict, + model: str, + messages: list, + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """ + Validate environment and set up headers for Vercel AI Gateway API. + + Vercel AI Gateway requires: + - Authorization header with Bearer token (API key or OIDC token) + """ + vercel_headers = { + "Content-Type": "application/json", + } + + # Add Authorization header if api_key is provided + if api_key: + vercel_headers["Authorization"] = f"Bearer {api_key}" + + # Merge with existing headers (user's extra_headers take priority) + merged_headers = {**vercel_headers, **headers} + + return merged_headers + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + """ + Get the complete URL for Vercel AI Gateway Embedding API endpoint. + """ + if api_base: + api_base = api_base.rstrip("/") + else: + api_base = ( + get_secret_str("VERCEL_AI_GATEWAY_API_BASE") + or "https://ai-gateway.vercel.sh/v1" + ) + + return f"{api_base}/embeddings" + + def transform_embedding_request( + self, + model: str, + input: AllEmbeddingInputValues, + optional_params: dict, + headers: dict, + ) -> dict: + """ + Transform embedding request to Vercel AI Gateway format (OpenAI-compatible). + """ + # Ensure input is a list + if isinstance(input, str): + input = [input] + + # Strip 'vercel_ai_gateway/' prefix if present + if model.startswith("vercel_ai_gateway/"): + model = model.replace("vercel_ai_gateway/", "", 1) + + return { + "model": model, + "input": input, + **optional_params, + } + + def transform_embedding_response( + self, + model: str, + raw_response: httpx.Response, + model_response: EmbeddingResponse, + logging_obj: LiteLLMLoggingObj, + api_key: Optional[str], + request_data: dict, + optional_params: dict, + litellm_params: dict, + ) -> EmbeddingResponse: + """ + Transform embedding response from Vercel AI Gateway format (OpenAI-compatible). + """ + logging_obj.post_call(original_response=raw_response.text) + + # Vercel AI Gateway returns standard OpenAI-compatible embedding response + response_json = raw_response.json() + + return convert_to_model_response_object( + response_object=response_json, + model_response_object=model_response, + response_type="embedding", + ) + + def get_supported_openai_params(self, model: str) -> list: + """ + Get list of supported OpenAI parameters for Vercel AI Gateway embeddings. + + Vercel AI Gateway supports the standard OpenAI embeddings parameters + and auto-maps 'dimensions' to each provider's expected field. + """ + return [ + "timeout", + "dimensions", + "encoding_format", + "user", + ] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + """ + Map OpenAI parameters to Vercel AI Gateway format. + """ + for param, value in non_default_params.items(): + if param in self.get_supported_openai_params(model): + optional_params[param] = value + return optional_params + + def get_error_class( + self, error_message: str, status_code: int, headers: Any + ) -> Any: + """ + Get the error class for Vercel AI Gateway errors. + """ + return VercelAIGatewayException( + message=error_message, + status_code=status_code, + headers=headers, + ) diff --git a/litellm/main.py b/litellm/main.py index ce84c8988e..3d1eb907a1 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -4866,6 +4866,36 @@ def embedding( # noqa: PLR0915 headers = openrouter_headers + response = base_llm_http_handler.embedding( + model=model, + input=input, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, + logging_obj=logging, + timeout=timeout, + model_response=EmbeddingResponse(), + optional_params=optional_params, + client=client, + aembedding=aembedding, + litellm_params=litellm_params_dict, + headers=headers, + ) + elif custom_llm_provider == "vercel_ai_gateway": + api_base = ( + api_base + or litellm.api_base + or get_secret_str("VERCEL_AI_GATEWAY_API_BASE") + or "https://ai-gateway.vercel.sh/v1" + ) + + api_key = ( + api_key + or litellm.api_key + or get_secret_str("VERCEL_AI_GATEWAY_API_KEY") + or get_secret_str("VERCEL_OIDC_TOKEN") + ) + response = base_llm_http_handler.embedding( model=model, input=input, diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index c0cff84bac..faeca9b2ae 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -274,11 +274,20 @@ def initialize_callbacks_on_proxy( # noqa: PLR0915 WebSearchInterceptionLogger, ) - websearch_interception_obj = WebSearchInterceptionLogger.initialize_from_proxy_config( - litellm_settings=litellm_settings, - callback_specific_params=callback_specific_params, + websearch_interception_obj = ( + WebSearchInterceptionLogger.initialize_from_proxy_config( + litellm_settings=litellm_settings, + callback_specific_params=callback_specific_params, + ) ) imported_list.append(websearch_interception_obj) + elif isinstance(callback, str) and callback == "datadog_cost_management": + from litellm.integrations.datadog.datadog_cost_management import ( + DatadogCostManagementLogger, + ) + + datadog_cost_management_obj = DatadogCostManagementLogger() + imported_list.append(datadog_cost_management_obj) elif isinstance(callback, CustomLogger): imported_list.append(callback) else: @@ -353,17 +362,17 @@ def get_remaining_tokens_and_requests_from_request_data(data: Dict) -> Dict[str, remaining_requests_variable_name = f"litellm-key-remaining-requests-{model_group}" remaining_requests = _metadata.get(remaining_requests_variable_name, None) if remaining_requests: - headers[f"x-litellm-key-remaining-requests-{h11_model_group_name}"] = ( - remaining_requests - ) + headers[ + f"x-litellm-key-remaining-requests-{h11_model_group_name}" + ] = remaining_requests # Remaining Tokens remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}" remaining_tokens = _metadata.get(remaining_tokens_variable_name, None) if remaining_tokens: - headers[f"x-litellm-key-remaining-tokens-{h11_model_group_name}"] = ( - remaining_tokens - ) + headers[ + f"x-litellm-key-remaining-tokens-{h11_model_group_name}" + ] = remaining_tokens return headers @@ -438,9 +447,9 @@ def add_guardrail_response_to_standard_logging_object( ): if litellm_logging_obj is None: return - standard_logging_object: Optional[StandardLoggingPayload] = ( - litellm_logging_obj.model_call_details.get("standard_logging_object") - ) + standard_logging_object: Optional[ + StandardLoggingPayload + ] = litellm_logging_obj.model_call_details.get("standard_logging_object") if standard_logging_object is None: return guardrail_information = standard_logging_object.get("guardrail_information", []) @@ -469,7 +478,9 @@ def get_metadata_variable_name_from_kwargs( return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" -def process_callback(_callback: str, callback_type: str, environment_variables: dict) -> dict: +def process_callback( + _callback: str, callback_type: str, environment_variables: dict +) -> dict: """Process a single callback and return its data with environment variables""" env_vars = CustomLogger.get_callback_env_vars(_callback) @@ -481,11 +492,9 @@ def process_callback(_callback: str, callback_type: str, environment_variables: else: env_vars_dict[_var] = env_variable - return { - "name": _callback, - "variables": env_vars_dict, - "type": callback_type - } + return {"name": _callback, "variables": env_vars_dict, "type": callback_type} + + def normalize_callback_names(callbacks: Iterable[Any]) -> List[Any]: if callbacks is None: return [] diff --git a/litellm/router.py b/litellm/router.py index 54650b120a..9fe37efa3a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1707,8 +1707,11 @@ class Router: litellm_params = deployment.get("litellm_params", {}) dep_num_retries = litellm_params.get("num_retries") - if dep_num_retries is not None and isinstance(dep_num_retries, int): - exception.num_retries = dep_num_retries # type: ignore + if dep_num_retries is not None: + try: + exception.num_retries = int(dep_num_retries) # type: ignore # Handle both int and str + except (ValueError, TypeError): + pass # Skip if value can't be converted to int def _update_kwargs_with_default_litellm_params( self, kwargs: dict, metadata_variable_name: Optional[str] = "metadata" diff --git a/litellm/types/integrations/datadog_cost_management.py b/litellm/types/integrations/datadog_cost_management.py new file mode 100644 index 0000000000..fe04f43ea0 --- /dev/null +++ b/litellm/types/integrations/datadog_cost_management.py @@ -0,0 +1,27 @@ +from typing import Dict, Optional, TypedDict + + +from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams + + +class DatadogCostManagementInitParams(StandardCustomLoggerInitParams): + """ + Init params for Datadog Cost Management + """ + + datadog_cost_management_params: Optional[Dict] = None + + +class DatadogFOCUSCostEntry(TypedDict): + """ + Represents a single cost line item in the FOCUS format. + Ref: https://focus.finops.org/#specification + """ + + ProviderName: str + ChargeDescription: str + ChargePeriodStart: str + ChargePeriodEnd: str + BilledCost: float + BillingCurrency: str + Tags: Optional[Dict[str, str]] diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index ef2f1ba4d5..a85aaafe23 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -93,6 +93,67 @@ class GuardrailConverseContentBlock(TypedDict, total=False): text: GuardrailConverseTextBlock +class CitationWebLocationBlock(TypedDict, total=False): + """ + Web location block for Nova grounding citations. + Contains the URL and domain from web search results. + + Reference: https://docs.aws.amazon.com/nova/latest/userguide/grounding.html + """ + + url: str + domain: str + + +class CitationLocationBlock(TypedDict, total=False): + """ + Location block containing the web location for a citation. + """ + + web: CitationWebLocationBlock + + +class CitationReferenceBlock(TypedDict, total=False): + """ + Citation reference block containing a single citation with its location. + + Each citation contains: + - location.web.url: The URL of the source + - location.web.domain: The domain of the source + """ + + location: CitationLocationBlock + + +class CitationsContentBlock(TypedDict, total=False): + """ + Citations content block returned by Nova grounding (web search) tool. + + When Nova grounding is enabled via systemTool, the model may return + citationsContent blocks containing web search citation references. + + Reference: https://docs.aws.amazon.com/nova/latest/userguide/grounding.html + + Example response structure: + { + "citationsContent": { + "citations": [ + { + "location": { + "web": { + "url": "https://example.com/article", + "domain": "example.com" + } + } + } + ] + } + } + """ + + citations: List[CitationReferenceBlock] + + class ContentBlock(TypedDict, total=False): text: str image: ImageBlock @@ -103,6 +164,7 @@ class ContentBlock(TypedDict, total=False): cachePoint: CachePointBlock reasoningContent: BedrockConverseReasoningContentBlock guardContent: GuardrailConverseContentBlock + citationsContent: CitationsContentBlock class MessageBlock(TypedDict): @@ -159,8 +221,24 @@ class ToolSpecBlock(TypedDict, total=False): description: str +class SystemToolBlock(TypedDict, total=False): + """ + System tool block for Nova grounding and other built-in tools. + + Example: + { + "systemTool": { + "name": "nova_grounding" + } + } + """ + + name: Required[str] + + class ToolBlock(TypedDict, total=False): toolSpec: Optional[ToolSpecBlock] + systemTool: Optional[SystemToolBlock] cachePoint: Optional[CachePointBlock] @@ -210,11 +288,13 @@ class ContentBlockStartEvent(TypedDict, total=False): class ContentBlockDeltaEvent(TypedDict, total=False): """ Either 'text' or 'toolUse' will be specified for Converse API streaming response. + May also include 'citationsContent' when Nova grounding is enabled. """ text: str toolUse: ToolBlockDeltaEvent reasoningContent: BedrockConverseReasoningContentBlockDelta + citationsContent: CitationsContentBlock class PerformanceConfigBlock(TypedDict): @@ -879,3 +959,8 @@ class BedrockGetBatchResponse(TypedDict, total=False): outputDataConfig: BedrockOutputDataConfig timeoutDurationInHours: Optional[int] clientRequestToken: Optional[str] + +class BedrockToolBlock(TypedDict, total=False): + toolSpec: Optional[ToolSpecBlock] + systemTool: Optional[SystemToolBlock] # For Nova grounding + cachePoint: Optional[CachePointBlock] diff --git a/litellm/utils.py b/litellm/utils.py index 584ab8805a..bb95be05b5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8053,6 +8053,12 @@ class ProviderConfigManager: ) return OpenrouterEmbeddingConfig() + elif litellm.LlmProviders.VERCEL_AI_GATEWAY == provider: + from litellm.llms.vercel_ai_gateway.embedding.transformation import ( + VercelAIGatewayEmbeddingConfig, + ) + + return VercelAIGatewayEmbeddingConfig() elif litellm.LlmProviders.GIGACHAT == provider: return litellm.GigaChatEmbeddingConfig() elif litellm.LlmProviders.SAGEMAKER == provider: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d958ea4503..20462db12c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -10232,6 +10232,48 @@ "mode": "completion", "output_cost_per_token": 5e-07 }, + "deepseek-v3-2-251201": { + "input_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "max_input_tokens": 98304, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "glm-4-7-251222": { + "input_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "max_input_tokens": 204800, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 0.0, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "kimi-k2-thinking-251104": { + "input_cost_per_token": 0.0, + "litellm_provider": "volcengine", + "max_input_tokens": 229376, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 0.0, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "doubao-embedding": { "input_cost_per_token": 0.0, "litellm_provider": "volcengine", diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index 7c0db41d13..d48dd1bfd9 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -3954,3 +3954,288 @@ def test_bedrock_openai_error_handling(): assert exc_info.value.status_code == 422 print("✓ Error handling works correctly") + +# ============================================================================ +# Nova Grounding (web_search_options) Unit Tests (Mocked) +# ============================================================================ + +def test_bedrock_nova_grounding_web_search_options_non_streaming(): + """ + Unit test for Nova grounding using web_search_options parameter (non-streaming). + + This test mocks the HTTP call to verify: + 1. web_search_options is correctly mapped to systemTool for Nova models + 2. The request structure is correct + + Related: https://docs.aws.amazon.com/nova/latest/userguide/grounding.html + """ + from unittest.mock import patch, MagicMock + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + client = HTTPHandler() + + messages = [ + { + "role": "user", + "content": "What is the current population of Tokyo, Japan?", + } + ] + + with patch.object(client, "post") as mock_post: + try: + completion( + model="us.amazon.nova-pro-v1:0", # No bedrock/ prefix when using api_base + messages=messages, + web_search_options={}, # Enables Nova grounding + max_tokens=500, + client=client, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com", + ) + except Exception: + pass # Expected - we're just checking the request structure + + # Verify the request was made correctly + if mock_post.called: + request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}")) + print(f"Request body: {json.dumps(request_body, indent=2)}") + + # Verify toolConfig is present with systemTool + assert "toolConfig" in request_body, "toolConfig should be in request" + tool_config = request_body["toolConfig"] + assert "tools" in tool_config, "tools should be in toolConfig" + + # Find the systemTool for nova_grounding + system_tool_found = False + for tool in tool_config["tools"]: + if "systemTool" in tool: + assert tool["systemTool"]["name"] == "nova_grounding" + system_tool_found = True + break + + assert system_tool_found, "systemTool with nova_grounding should be present" + print(f"✓ web_search_options correctly transformed to systemTool (non-streaming)") + + +def test_bedrock_nova_grounding_with_function_tools(): + """ + Unit test for Nova grounding combined with regular function tools. + + This tests the scenario where users want both web grounding AND + custom function calling capabilities. + """ + from unittest.mock import patch + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + client = HTTPHandler() + + # Regular function tool + tools = [ + { + "type": "function", + "function": { + "name": "get_stock_price", + "description": "Get the current stock price for a given ticker symbol", + "parameters": { + "type": "object", + "properties": { + "ticker": { + "type": "string", + "description": "The stock ticker symbol, e.g. AAPL, GOOGL", + } + }, + "required": ["ticker"], + }, + }, + } + ] + + messages = [ + { + "role": "user", + "content": "What is the current market cap of Apple Inc?", + } + ] + + with patch.object(client, "post") as mock_post: + try: + completion( + model="us.amazon.nova-pro-v1:0", # No bedrock/ prefix when using api_base + messages=messages, + tools=tools, + web_search_options={}, # Also enable web grounding + max_tokens=500, + client=client, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com", + ) + except Exception: + pass # Expected - we're just checking the request structure + + # Verify the request was made correctly + if mock_post.called: + request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}")) + print(f"Request body: {json.dumps(request_body, indent=2)}") + + # Verify toolConfig has both function tool and systemTool + assert "toolConfig" in request_body, "toolConfig should be in request" + tool_config = request_body["toolConfig"] + assert "tools" in tool_config, "tools should be in toolConfig" + + tools_in_request = tool_config["tools"] + + # Should have both the function tool and the systemTool + function_tool_found = False + system_tool_found = False + + for tool in tools_in_request: + if "toolSpec" in tool: + assert tool["toolSpec"]["name"] == "get_stock_price" + function_tool_found = True + if "systemTool" in tool: + assert tool["systemTool"]["name"] == "nova_grounding" + system_tool_found = True + + assert function_tool_found, "Function tool (get_stock_price) should be present" + assert system_tool_found, "systemTool (nova_grounding) should be present" + print(f"✓ Both function tools and web_search_options correctly combined") + + +@pytest.mark.asyncio +async def test_bedrock_nova_grounding_async(): + """ + Async unit test for Nova grounding via web_search_options. + + This test verifies the request transformation for async calls. + """ + from unittest.mock import patch, AsyncMock + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler + + client = AsyncHTTPHandler() + + messages = [ + { + "role": "user", + "content": "What is the weather forecast for New York City today?", + } + ] + + with patch.object(client, "post", new=AsyncMock()) as mock_post: + try: + await litellm.acompletion( + model="us.amazon.nova-pro-v1:0", # No bedrock/ prefix when using api_base + messages=messages, + web_search_options={}, + max_tokens=500, + client=client, + api_base="https://bedrock-runtime.us-east-1.amazonaws.com", + ) + except Exception: + pass # Expected - we're just checking the request structure + + # Verify the request was made correctly + if mock_post.called: + request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}")) + print(f"Request body: {json.dumps(request_body, indent=2)}") + + # Verify toolConfig is present with systemTool + assert "toolConfig" in request_body, "toolConfig should be in request" + tool_config = request_body["toolConfig"] + assert "tools" in tool_config, "tools should be in toolConfig" + + # Find the systemTool for nova_grounding + system_tool_found = False + for tool in tool_config["tools"]: + if "systemTool" in tool: + assert tool["systemTool"]["name"] == "nova_grounding" + system_tool_found = True + break + + assert system_tool_found, "systemTool with nova_grounding should be present" + print(f"✓ Async web_search_options correctly transformed to systemTool") + + +def test_bedrock_nova_web_search_options_ignored_for_non_nova(): + """ + Test that web_search_options is ignored for non-Nova Bedrock models. + + Nova grounding is only supported on Nova models. For other models, + the parameter should be silently ignored. + """ + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + config = AmazonConverseConfig() + + # Should return None for non-Nova models + result = config._map_web_search_options({}, "anthropic.claude-3-sonnet-v1") + assert result is None + + result = config._map_web_search_options({}, "amazon.titan-text-express-v1") + assert result is None + + # Should return systemTool for Nova models + result = config._map_web_search_options({}, "amazon.nova-pro-v1:0") + assert result is not None + system_tool = result.get("systemTool") + assert system_tool is not None + assert system_tool["name"] == "nova_grounding" + + result2 = config._map_web_search_options({}, "us.amazon.nova-premier-v1:0") + assert result2 is not None + system_tool2 = result2.get("systemTool") + assert system_tool2 is not None + assert system_tool2["name"] == "nova_grounding" + + +def test_bedrock_nova_grounding_request_transformation(): + """ + Unit test to verify that web_search_options transforms to systemTool in the request. + """ + from unittest.mock import patch, MagicMock + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + client = HTTPHandler() + + messages = [{"role": "user", "content": "What is the population of Tokyo?"}] + + with patch.object(client, "post") as mock_post: + mock_post.return_value = MagicMock( + status_code=200, + json=lambda: { + "output": {"message": {"role": "assistant", "content": [{"text": "Test"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 10, "outputTokens": 5} + } + ) + + try: + response = completion( + model="bedrock/us.amazon.nova-pro-v1:0", + messages=messages, + web_search_options={}, + max_tokens=100, + client=client, + ) + except Exception: + pass # Expected - we're just checking the request + + if mock_post.called: + request_body = json.loads(mock_post.call_args.kwargs.get("data", "{}")) + print(f"Request body: {json.dumps(request_body, indent=2)}") + + # Verify toolConfig is present with systemTool + assert "toolConfig" in request_body, "toolConfig should be in request" + + tool_config = request_body["toolConfig"] + assert "tools" in tool_config, "tools should be in toolConfig" + + tools_in_request = tool_config["tools"] + + # Find the systemTool + system_tool_found = False + for tool in tools_in_request: + if "systemTool" in tool: + assert tool["systemTool"]["name"] == "nova_grounding" + system_tool_found = True + break + + assert system_tool_found, "systemTool with nova_grounding should be present" + print("✓ web_search_options correctly transformed to systemTool") diff --git a/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py b/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py new file mode 100644 index 0000000000..be2084969a --- /dev/null +++ b/tests/test_litellm/integrations/datadog/test_datadog_cost_management.py @@ -0,0 +1,169 @@ +import os +import time +from unittest.mock import AsyncMock + +import pytest +from httpx import Response + +from litellm.integrations.datadog.datadog_cost_management import ( + DatadogCostManagementLogger, +) +from litellm.types.utils import StandardLoggingPayload + + +@pytest.fixture +def clean_env(): + # Save original env + original_api_key = os.environ.get("DD_API_KEY") + original_app_key = os.environ.get("DD_APP_KEY") + original_site = os.environ.get("DD_SITE") + + # Set test env + os.environ["DD_API_KEY"] = "test_api_key" + os.environ["DD_APP_KEY"] = "test_app_key" + os.environ["DD_SITE"] = "test.datadoghq.com" + + yield + + # Restore original env + if original_api_key: + os.environ["DD_API_KEY"] = original_api_key + else: + del os.environ["DD_API_KEY"] + + if original_app_key: + os.environ["DD_APP_KEY"] = original_app_key + else: + del os.environ["DD_APP_KEY"] + + if original_site: + os.environ["DD_SITE"] = original_site + else: + del os.environ["DD_SITE"] + + +@pytest.mark.asyncio +async def test_init(clean_env): + """ + Test initialization sets up clients and url correctly + """ + logger = DatadogCostManagementLogger() + assert logger.dd_api_key == "test_api_key" + assert logger.dd_app_key == "test_app_key" + assert ( + logger.upload_url == "https://api.test.datadoghq.com/api/v2/cost/custom_costs" + ) + + +@pytest.mark.asyncio +async def test_aggregate_costs(clean_env): + """ + Test that costs are correctly aggregated by provider, model, and date + """ + logger = DatadogCostManagementLogger() + + # Mock some log payloads + now = time.time() + day_str = time.strftime("%Y-%m-%d", time.localtime(now)) + + logs = [ + StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.01, + startTime=now, + metadata={"user_api_key_team_alias": "team-a"}, + ), + StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.02, + startTime=now, + metadata={"user_api_key_team_alias": "team-a"}, + ), + StandardLoggingPayload( + custom_llm_provider="anthropic", + model="claude-3", + response_cost=0.05, + startTime=now, + ), + ] + + aggregated = logger._aggregate_costs(logs) + + assert len(aggregated) == 2 + + # Check OpenAI entry + openai_entry = next(e for e in aggregated if e["ProviderName"] == "openai") + assert openai_entry["BilledCost"] == 0.03 + assert openai_entry["ChargeDescription"] == "LLM Usage for gpt-4" + assert openai_entry["ChargePeriodStart"] == day_str + assert openai_entry["Tags"]["team"] == "team-a" + assert "env" in openai_entry["Tags"] + assert "service" in openai_entry["Tags"] + + # Check Anthropic entry + anthropic_entry = next(e for e in aggregated if e["ProviderName"] == "anthropic") + assert anthropic_entry["BilledCost"] == 0.05 + + +@pytest.mark.asyncio +async def test_async_log_success_event(clean_env): + """ + Test that logs are added to queue + """ + logger = DatadogCostManagementLogger(batch_size=10) + + await logger.async_log_success_event( + kwargs={"standard_logging_object": {"response_cost": 0.01}}, + response_obj={}, + start_time=time.time(), + end_time=time.time(), + ) + + assert len(logger.log_queue) == 1 + assert logger.log_queue[0]["response_cost"] == 0.01 + + # Test zero cost ignored + await logger.async_log_success_event( + kwargs={"standard_logging_object": {"response_cost": 0.0}}, + response_obj={}, + start_time=time.time(), + end_time=time.time(), + ) + + assert len(logger.log_queue) == 1 + + +@pytest.mark.asyncio +async def test_async_send_batch(clean_env): + """ + Test that batch is aggregated and uploaded + """ + logger = DatadogCostManagementLogger() + logger.async_client = AsyncMock() + logger.async_client.put.return_value = Response(202, json={"status": "ok"}) + + # Add logs directly to queue + logger.log_queue = [ + StandardLoggingPayload( + custom_llm_provider="openai", + model="gpt-4", + response_cost=0.01, + startTime=time.time(), + ) + ] + + await logger.async_send_batch() + + # Verify API called + assert logger.async_client.put.called + call_args = logger.async_client.put.call_args + assert call_args[0][0] == "https://api.test.datadoghq.com/api/v2/cost/custom_costs" + + import json + + # Use call_args.kwargs['content'] + content = json.loads(call_args[1]["content"]) + assert content[0]["ProviderName"] == "openai" + assert content[0]["BilledCost"] == 0.01 diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs_agent.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs_agent.py new file mode 100644 index 0000000000..2bb51e1e1b --- /dev/null +++ b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs_agent.py @@ -0,0 +1,62 @@ +import os +from unittest.mock import patch +from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger + + +def test_datadog_llm_obs_agent_configuration(): + """ + Test that DataDog LLM Obs logger correctly configures agent endpoint. + """ + test_env = { + "LITELLM_DD_AGENT_HOST": "localhost", + "LITELLM_DD_LLM_OBS_PORT": "10518", + "DD_API_KEY": "test-api-key", # Optional, but checking if it's preserved + } + + # Ensure DD_SITE is NOT set to verify we don't need it in agent mode + + with patch.dict(os.environ, test_env, clear=True): + with patch("asyncio.create_task"): # Prevent periodic flush task from running + dd_logger = DataDogLLMObsLogger() + + expected_url = "http://localhost:10518/api/intake/llm-obs/v1/trace/spans" + assert dd_logger.intake_url == expected_url + assert dd_logger.DD_API_KEY == "test-api-key" + + +def test_datadog_llm_obs_agent_no_api_key_ok(): + """ + Test that agent mode works WITHOUT DD_API_KEY (agent handles auth). + """ + test_env = { + "LITELLM_DD_AGENT_HOST": "localhost", + # No DD_API_KEY + } + + with patch.dict(os.environ, test_env, clear=True): + with patch("asyncio.create_task"): + # Should NOT raise exception anymore + dd_logger = DataDogLLMObsLogger() + + assert dd_logger.DD_API_KEY is None + # Default port is 8126 if not set + expected_url = "http://localhost:8126/api/intake/llm-obs/v1/trace/spans" + assert dd_logger.intake_url == expected_url + + +def test_datadog_llm_obs_direct_api_configuration(): + """ + Test that direct API configuration still works as expected. + """ + test_env = { + "DD_API_KEY": "direct-api-key", + "DD_SITE": "us5.datadoghq.com", + } + + with patch.dict(os.environ, test_env, clear=True): + with patch("asyncio.create_task"): + dd_logger = DataDogLLMObsLogger() + + expected_url = "https://api.us5.datadoghq.com/api/intake/llm-obs/v1/trace/spans" + assert dd_logger.intake_url == expected_url + assert dd_logger.DD_API_KEY == "direct-api-key" diff --git a/tests/test_litellm/integrations/test_custom_guardrail_recursion.py b/tests/test_litellm/integrations/test_custom_guardrail_recursion.py new file mode 100644 index 0000000000..f05b5848bd --- /dev/null +++ b/tests/test_litellm/integrations/test_custom_guardrail_recursion.py @@ -0,0 +1,73 @@ +import pytest +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.types.guardrails import GuardrailEventHooks +import json + + +class TestCustomGuardrailRecursion: + """ + Specific tests for the circular reference / RecursionError fix in logging. + """ + + def test_log_guardrail_information_handles_circular_references(self): + """ + Test that add_standard_logging method sanitizes input data containing circular references + instead of crashing. + + This reproduces the Langfuse crash scenario: + Request -> Metadata -> GuardrailResponse -> DebugContext -> Request + """ + guardrail = CustomGuardrail( + guardrail_name="recursion_test_guardrail", + event_hook=GuardrailEventHooks.pre_call, + ) + + # 1. Setup Circular Data + request_data = {"user_id": "test_recursive_user"} + metadata = {"session_id": "123"} + request_data["metadata"] = metadata + + # Create the danger: Guardrail Response holding a reference back to request_data + dirty_response = { + "flagged": False, + "debug_context": request_data, # <--- ACCESS TO ROOT (Circular Ref) + } + + # 2. Invoke the logging method + # If the fix is working, this will NOT raise RecursionError + try: + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=dirty_response, + request_data=request_data, + guardrail_status="success", + start_time=1.0, + end_time=2.0, + duration=1.0, + masked_entity_count={}, + event_type=GuardrailEventHooks.pre_call, + ) + except RecursionError: + pytest.fail( + "RecursionError raised! The cyclic reference sanitization failed." + ) + + # 3. Verify the data stored is safe + stored_info = request_data["metadata"][ + "standard_logging_guardrail_information" + ][0] + stored_response = stored_info["guardrail_response"] + + # Check that we can dump it to JSON without crashing (Ultimate proof) + try: + json.dumps(stored_response) + except Exception as e: + pytest.fail(f"Stored data is not JSON serializable: {e}") + + # Check content - keys should be preserved but recursion broken + assert "debug_context" in stored_response + debug_context = stored_response["debug_context"] + + # In a sanitized copy, the nested metadata should be a copy, not the original live dict + assert debug_context["user_id"] == "test_recursive_user" + # The 'metadata' inside 'debug_context' would be where recursion stops or is filtered + assert "metadata" in debug_context diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index a22fe13798..e87233a52a 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -1138,6 +1138,73 @@ def test_bedrock_create_bedrock_block_different_document_formats(): assert block["document"]["name"].endswith(f"_{format_type}") assert block["document"]["format"] == format_type +def test_bedrock_nova_web_search_options_mapping(): + """ + Test that web_search_options is correctly mapped to Nova grounding. + + This follows the LiteLLM pattern for web search where: + - Vertex AI maps web_search_options to {"googleSearch": {}} + - Anthropic maps web_search_options to {"type": "web_search_20250305", ...} + - Nova should map web_search_options to {"systemTool": {"name": "nova_grounding"}} + """ + from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + config = AmazonConverseConfig() + + # Test basic mapping for Nova model + result = config._map_web_search_options({}, "amazon.nova-pro-v1:0") + + assert result is not None + system_tool = result.get("systemTool") + assert system_tool is not None + assert system_tool["name"] == "nova_grounding" + + # Test with search_context_size (should be ignored for Nova) + result2 = config._map_web_search_options( + {"search_context_size": "high"}, + "us.amazon.nova-premier-v1:0" + ) + + assert result2 is not None + system_tool2 = result2.get("systemTool") + assert system_tool2 is not None + assert system_tool2["name"] == "nova_grounding" + # Nova doesn't support search_context_size, so it's just ignored + +def test_bedrock_tools_pt_does_not_handle_system_tool(): + """ + Verify that _bedrock_tools_pt does NOT handle system_tool format. + + System tools (nova_grounding) should be added via web_search_options, + not via the tools parameter directly. + """ + + from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_tools_pt + + # Regular function tools should still work + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"} + }, + "required": ["location"] + } + } + } + ] + + result = _bedrock_tools_pt(tools=tools) + + assert len(result) == 1 + tool_spec = result[0].get("toolSpec") + assert tool_spec is not None + assert tool_spec["name"] == "get_weather" def test_convert_to_anthropic_tool_result_image_with_cache_control(): """ @@ -1305,12 +1372,12 @@ def test_convert_to_anthropic_tool_result_image_url_as_http(): assert result["content"][0]["cache_control"]["type"] == "ephemeral" def test_anthropic_messages_pt_server_tool_use_passthrough(): """ - Test that anthropic_messages_pt passes through server_tool_use and + Test that anthropic_messages_pt passes through server_tool_use and tool_search_tool_result blocks in assistant message content. - + These are Anthropic-native content types used for tool search functionality that need to be preserved when reconstructing multi-turn conversations. - + Fixes: https://github.com/BerriAI/litellm/issues/XXXXX """ from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt @@ -1359,15 +1426,15 @@ def test_anthropic_messages_pt_server_tool_use_passthrough(): # Verify we have 3 messages (user, assistant, user) assert len(result) == 3 - + # Verify the assistant message content assistant_msg = result[1] assert assistant_msg["role"] == "assistant" assert isinstance(assistant_msg["content"], list) - + # Find the different content block types content_types = [block.get("type") for block in assistant_msg["content"]] - + # Verify server_tool_use block is preserved assert "server_tool_use" in content_types server_tool_use_block = next( @@ -1376,7 +1443,7 @@ def test_anthropic_messages_pt_server_tool_use_passthrough(): assert server_tool_use_block["id"] == "srvtoolu_01ABC123" assert server_tool_use_block["name"] == "tool_search_tool_regex" assert server_tool_use_block["input"] == {"query": ".*time.*"} - + # Verify tool_search_tool_result block is preserved assert "tool_search_tool_result" in content_types tool_result_block = next( @@ -1385,7 +1452,7 @@ def test_anthropic_messages_pt_server_tool_use_passthrough(): assert tool_result_block["tool_use_id"] == "srvtoolu_01ABC123" assert tool_result_block["content"]["type"] == "tool_search_tool_search_result" assert tool_result_block["content"]["tool_references"][0]["tool_name"] == "get_time" - + # Verify text block is also preserved assert "text" in content_types text_block = next( diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index e035e193fe..1f3f558a49 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -787,11 +787,13 @@ def test_get_masked_values(): "presidio_ad_hoc_recognizers": None, "aws_bedrock_runtime_endpoint": None, "presidio_anonymizer_api_base": None, + "vertex_credentials": "{sensitive_api_key}", } masked_values = _get_masked_values( sensitive_object, unmasked_length=4, number_of_asterisks=4 ) assert masked_values["presidio_anonymizer_api_base"] is None + assert masked_values["vertex_credentials"] == "{s****y}" @pytest.mark.asyncio diff --git a/tests/test_litellm/llms/vercel_ai_gateway/embedding/__init__.py b/tests/test_litellm/llms/vercel_ai_gateway/embedding/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py b/tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py new file mode 100644 index 0000000000..af1e1df92f --- /dev/null +++ b/tests/test_litellm/llms/vercel_ai_gateway/embedding/test_vercel_ai_gateway_embedding.py @@ -0,0 +1,218 @@ +import os +import sys +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +from litellm.llms.vercel_ai_gateway.embedding.transformation import ( + VercelAIGatewayEmbeddingConfig, +) +from litellm.llms.vercel_ai_gateway.common_utils import VercelAIGatewayException +from litellm.types.utils import EmbeddingResponse + + +def test_vercel_ai_gateway_embedding_get_complete_url(): + """Test URL generation for embeddings endpoint""" + config = VercelAIGatewayEmbeddingConfig() + + # Test with default API base + url = config.get_complete_url( + api_base=None, + api_key=None, + model="openai/text-embedding-3-small", + optional_params={}, + litellm_params={}, + ) + assert url == "https://ai-gateway.vercel.sh/v1/embeddings" + + # Test with custom API base + url = config.get_complete_url( + api_base="https://custom.vercel.sh/v1", + api_key=None, + model="openai/text-embedding-3-small", + optional_params={}, + litellm_params={}, + ) + assert url == "https://custom.vercel.sh/v1/embeddings" + + # Test with trailing slash + url = config.get_complete_url( + api_base="https://custom.vercel.sh/v1/", + api_key=None, + model="openai/text-embedding-3-small", + optional_params={}, + litellm_params={}, + ) + assert url == "https://custom.vercel.sh/v1/embeddings" + + +def test_vercel_ai_gateway_embedding_transform_request(): + """Test request transformation for embeddings""" + config = VercelAIGatewayEmbeddingConfig() + + # Test with string input + request = config.transform_embedding_request( + model="openai/text-embedding-3-small", + input="Hello world", + optional_params={}, + headers={}, + ) + assert request["model"] == "openai/text-embedding-3-small" + assert request["input"] == ["Hello world"] + + # Test with list input + request = config.transform_embedding_request( + model="openai/text-embedding-3-small", + input=["Hello", "World"], + optional_params={}, + headers={}, + ) + assert request["model"] == "openai/text-embedding-3-small" + assert request["input"] == ["Hello", "World"] + + # Test stripping vercel_ai_gateway/ prefix + request = config.transform_embedding_request( + model="vercel_ai_gateway/openai/text-embedding-3-small", + input="Hello", + optional_params={}, + headers={}, + ) + assert request["model"] == "openai/text-embedding-3-small" + + +def test_vercel_ai_gateway_embedding_transform_request_with_dimensions(): + """Test request transformation with dimensions parameter""" + config = VercelAIGatewayEmbeddingConfig() + + request = config.transform_embedding_request( + model="openai/text-embedding-3-small", + input="Hello world", + optional_params={"dimensions": 768}, + headers={}, + ) + assert request["model"] == "openai/text-embedding-3-small" + assert request["input"] == ["Hello world"] + assert request["dimensions"] == 768 + + +def test_vercel_ai_gateway_embedding_validate_environment(): + """Test header validation and setup""" + config = VercelAIGatewayEmbeddingConfig() + + headers = config.validate_environment( + headers={}, + model="openai/text-embedding-3-small", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test_key", + ) + assert headers["Content-Type"] == "application/json" + assert headers["Authorization"] == "Bearer test_key" + + # Test with existing headers (should merge) + headers = config.validate_environment( + headers={"X-Custom": "value"}, + model="openai/text-embedding-3-small", + messages=[], + optional_params={}, + litellm_params={}, + api_key="test_key", + ) + assert headers["X-Custom"] == "value" + assert headers["Authorization"] == "Bearer test_key" + + +def test_vercel_ai_gateway_embedding_get_supported_params(): + """Test supported OpenAI parameters""" + config = VercelAIGatewayEmbeddingConfig() + supported = config.get_supported_openai_params("openai/text-embedding-3-small") + + assert "dimensions" in supported + assert "encoding_format" in supported + assert "timeout" in supported + assert "user" in supported + + +def test_vercel_ai_gateway_embedding_map_openai_params(): + """Test OpenAI parameter mapping""" + config = VercelAIGatewayEmbeddingConfig() + + optional_params = config.map_openai_params( + non_default_params={"dimensions": 768, "encoding_format": "float"}, + optional_params={}, + model="openai/text-embedding-3-small", + drop_params=False, + ) + assert optional_params["dimensions"] == 768 + assert optional_params["encoding_format"] == "float" + + +def test_vercel_ai_gateway_embedding_error_class(): + """Test error class creation""" + config = VercelAIGatewayEmbeddingConfig() + + error = config.get_error_class( + error_message="Test error", + status_code=400, + headers={"Content-Type": "application/json"}, + ) + + assert isinstance(error, VercelAIGatewayException) + assert error.message == "Test error" + assert error.status_code == 400 + + +def test_vercel_ai_gateway_embedding_transform_response(): + """Test response transformation""" + config = VercelAIGatewayEmbeddingConfig() + + mock_response = MagicMock(spec=httpx.Response) + mock_response.text = '{"object":"list","data":[{"object":"embedding","index":0,"embedding":[0.1,0.2,0.3]}],"model":"openai/text-embedding-3-small","usage":{"prompt_tokens":2,"total_tokens":2}}' + mock_response.json.return_value = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "model": "openai/text-embedding-3-small", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + } + + mock_logging = MagicMock() + + response = config.transform_embedding_response( + model="openai/text-embedding-3-small", + raw_response=mock_response, + model_response=EmbeddingResponse(), + logging_obj=mock_logging, + api_key="test_key", + request_data={}, + optional_params={}, + litellm_params={}, + ) + + assert response is not None + mock_logging.post_call.assert_called_once() + + +def test_vercel_ai_gateway_embedding_env_vars(): + """Test environment variable handling""" + config = VercelAIGatewayEmbeddingConfig() + + with patch.dict( + os.environ, + { + "VERCEL_AI_GATEWAY_API_BASE": "https://env.vercel.sh/v1", + }, + ): + url = config.get_complete_url( + api_base=None, + api_key=None, + model="openai/text-embedding-3-small", + optional_params={}, + litellm_params={}, + ) + assert url == "https://env.vercel.sh/v1/embeddings" diff --git a/tests/test_litellm/test_gpt_image_cost_calculator.py b/tests/test_litellm/test_gpt_image_cost_calculator.py index 0a2a62b6c9..620c073498 100644 --- a/tests/test_litellm/test_gpt_image_cost_calculator.py +++ b/tests/test_litellm/test_gpt_image_cost_calculator.py @@ -19,10 +19,13 @@ import pytest import litellm from litellm.types.utils import ( + CompletionTokensDetailsWrapper, ImageResponse, ImageObject, ImageUsage, ImageUsageInputTokensDetails, + PromptTokensDetailsWrapper, + Usage, ) @@ -202,6 +205,71 @@ class TestGPTImageCostRouting: assert cost >= 0 +class TestGPTImage15OutputImageTokens: + """ + Test for GitHub issue #19508: + Image usage calculation does not include image tokens in gpt-image-1.5 + + gpt-image-1.5 returns output_tokens_details with separate image_tokens and text_tokens, + and these must be correctly included in cost calculation. + """ + + def test_gpt_image_15_output_image_tokens_cost(self): + """ + Test that output image tokens are correctly included in cost calculation. + + This tests the fix for issue #19508 where output_tokens_details.image_tokens + were not being included in the cost calculation, causing costs to be + underreported (e.g., $0.046 instead of $0.14). + """ + # Simulate gpt-image-1.5 response with output_tokens_details + # This is what the API returns and what convert_to_image_response transforms + usage = Usage( + prompt_tokens=169, + completion_tokens=4599, + total_tokens=4768, + prompt_tokens_details=PromptTokensDetailsWrapper( + text_tokens=169, + image_tokens=0, + ), + completion_tokens_details=CompletionTokensDetailsWrapper( + text_tokens=439, + image_tokens=4160, + ), + ) + + image_response = ImageResponse( + created=1234567890, + data=[ImageObject(b64_json="test")], + ) + image_response.usage = usage + image_response._hidden_params = {"custom_llm_provider": "openai"} + + cost = litellm.completion_cost( + completion_response=image_response, + model="gpt-image-1.5", + call_type="image_generation", + custom_llm_provider="openai", + ) + + # gpt-image-1.5 pricing: + # - input_cost_per_token: 5e-06 ($5/1M for text input) + # - output_cost_per_token: 1e-05 ($10/1M for text output) + # - output_cost_per_image_token: 3.2e-05 ($32/1M for image output) + # + # Expected cost: + # Input text: 169 * $5/1M = $0.000845 + # Output text: 439 * $10/1M = $0.00439 + # Output image: 4160 * $32/1M = $0.13312 + # Total: $0.138355 + expected_cost = 169 * 5e-06 + 439 * 1e-05 + 4160 * 3.2e-05 + + assert abs(cost - expected_cost) < 1e-6, ( + f"Expected {expected_cost}, got {cost}. " + f"Image tokens may not be included in cost calculation." + ) + + class TestCompletionCostIntegration: """Test the full completion_cost integration for gpt-image-1""" diff --git a/tests/test_litellm/test_router_per_deployment_num_retries.py b/tests/test_litellm/test_router_per_deployment_num_retries.py index 4021ca2807..154ba579e4 100644 --- a/tests/test_litellm/test_router_per_deployment_num_retries.py +++ b/tests/test_litellm/test_router_per_deployment_num_retries.py @@ -32,17 +32,17 @@ class TestPerDeploymentNumRetries: ) deployment = router.model_list[0] - + # Create a mock exception without num_retries class MockException(Exception): pass - + exc = MockException("test error") assert not hasattr(exc, "num_retries") or exc.num_retries is None - + # Call the helper router._set_deployment_num_retries_on_exception(exc, deployment) - + # Verify num_retries was set from deployment assert exc.num_retries == 5 @@ -66,16 +66,16 @@ class TestPerDeploymentNumRetries: ) deployment = router.model_list[0] - + # Create an exception that already has num_retries class MockException(Exception): num_retries = 10 # Already set - + exc = MockException("test error") - + # Call the helper router._set_deployment_num_retries_on_exception(exc, deployment) - + # Verify num_retries was NOT overridden assert exc.num_retries == 10 @@ -99,15 +99,15 @@ class TestPerDeploymentNumRetries: ) deployment = router.model_list[0] - + class MockException(Exception): pass - + exc = MockException("test error") - + # Call the helper router._set_deployment_num_retries_on_exception(exc, deployment) - + # Verify num_retries was not set (deployment has no num_retries) assert not hasattr(exc, "num_retries") or exc.num_retries is None @@ -155,3 +155,36 @@ class TestPerDeploymentNumRetries: kwargs = {} router._update_kwargs_before_fallbacks(model="test-model", kwargs=kwargs) assert kwargs["num_retries"] == 7 # Uses global + + def test_set_deployment_num_retries_with_string_value(self): + """ + Test that _set_deployment_num_retries_on_exception handles string values + from environment variables correctly. + GitHub Issue: #19481 + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "test-key", + "num_retries": "6", # String value (as from env var) + }, + }, + ], + num_retries=0, # Global setting + ) + + deployment = router.model_list[0] + + class MockException(Exception): + pass + + exc = MockException("test error") + + # Call the helper + router._set_deployment_num_retries_on_exception(exc, deployment) + + # Verify num_retries was converted from string to int + assert exc.num_retries == 6