mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-10 09:04:10 +00:00
Merge branch 'BerriAI:main' into LangfuseUsageDetails
This commit is contained in:
Fabrício Ceschin
+2
-1
@@ -95,4 +95,5 @@ test.py
|
||||
litellm_config.yaml
|
||||
.cursor
|
||||
.vscode/launch.json
|
||||
litellm/proxy/to_delete_loadtest_work/*
|
||||
litellm/proxy/to_delete_loadtest_work/*
|
||||
update_model_cost_map.py
|
||||
|
||||
@@ -1,3 +1,6 @@
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
|
||||
# Helicone - OSS LLM Observability Platform
|
||||
|
||||
:::tip
|
||||
@@ -9,9 +12,68 @@ https://github.com/BerriAI/litellm
|
||||
|
||||
[Helicone](https://helicone.ai/) is an open source observability platform that proxies your LLM requests and provides key insights into your usage, spend, latency and more.
|
||||
|
||||
## Using Helicone with LiteLLM
|
||||
## Quick Start
|
||||
|
||||
LiteLLM provides `success_callbacks` and `failure_callbacks`, allowing you to easily log data to Helicone based on the status of your responses.
|
||||
<Tabs>
|
||||
<TabItem value="sdk" label="Python SDK">
|
||||
|
||||
Use just 1 line of code to instantly log your responses **across all providers** with Helicone:
|
||||
|
||||
```python
|
||||
import os
|
||||
from litellm import completion
|
||||
|
||||
## Set env variables
|
||||
os.environ["HELICONE_API_KEY"] = "your-helicone-key"
|
||||
os.environ["OPENAI_API_KEY"] = "your-openai-key"
|
||||
|
||||
# Set callbacks
|
||||
litellm.success_callback = ["helicone"]
|
||||
|
||||
# OpenAI call
|
||||
response = completion(
|
||||
model="gpt-4o",
|
||||
messages=[{"role": "user", "content": "Hi 👋 - I'm OpenAI"}],
|
||||
)
|
||||
|
||||
print(response)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="proxy" label="LiteLLM Proxy">
|
||||
|
||||
Add Helicone to your LiteLLM proxy configuration:
|
||||
|
||||
```yaml title="config.yaml"
|
||||
model_list:
|
||||
- model_name: gpt-4
|
||||
litellm_params:
|
||||
model: gpt-4
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
# Add Helicone callback
|
||||
litellm_settings:
|
||||
success_callback: ["helicone"]
|
||||
|
||||
# Set Helicone API key
|
||||
environment_variables:
|
||||
HELICONE_API_KEY: "your-helicone-key"
|
||||
```
|
||||
|
||||
Start the proxy:
|
||||
```bash
|
||||
litellm --config config.yaml
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Integration Methods
|
||||
|
||||
There are two main approaches to integrate Helicone with LiteLLM:
|
||||
|
||||
1. **Callbacks**: Log to Helicone while using any provider
|
||||
2. **Proxy Mode**: Use Helicone as a proxy for advanced features
|
||||
|
||||
### Supported LLM Providers
|
||||
|
||||
@@ -26,27 +88,16 @@ Helicone can log requests across [various LLM providers](https://docs.helicone.a
|
||||
- Replicate
|
||||
- And more
|
||||
|
||||
### Integration Methods
|
||||
## Method 1: Using Callbacks
|
||||
|
||||
There are two main approaches to integrate Helicone with LiteLLM:
|
||||
Log requests to Helicone while using any LLM provider directly.
|
||||
|
||||
1. Using callbacks
|
||||
2. Using Helicone as a proxy
|
||||
|
||||
Let's explore each method in detail.
|
||||
|
||||
### Approach 1: Use Callbacks
|
||||
|
||||
Use just 1 line of code to instantly log your responses **across all providers** with Helicone:
|
||||
|
||||
```python
|
||||
litellm.success_callback = ["helicone"]
|
||||
```
|
||||
|
||||
Complete Code
|
||||
<Tabs>
|
||||
<TabItem value="sdk" label="Python SDK">
|
||||
|
||||
```python
|
||||
import os
|
||||
import litellm
|
||||
from litellm import completion
|
||||
|
||||
## Set env variables
|
||||
@@ -66,28 +117,78 @@ response = completion(
|
||||
print(response)
|
||||
```
|
||||
|
||||
### Approach 2: Use Helicone as a proxy
|
||||
</TabItem>
|
||||
<TabItem value="proxy" label="LiteLLM Proxy">
|
||||
|
||||
```yaml title="config.yaml"
|
||||
model_list:
|
||||
- model_name: gpt-4
|
||||
litellm_params:
|
||||
model: gpt-4
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
- model_name: claude-3
|
||||
litellm_params:
|
||||
model: anthropic/claude-3-sonnet-20240229
|
||||
api_key: os.environ/ANTHROPIC_API_KEY
|
||||
|
||||
# Add Helicone logging
|
||||
litellm_settings:
|
||||
success_callback: ["helicone"]
|
||||
|
||||
# Environment variables
|
||||
environment_variables:
|
||||
HELICONE_API_KEY: "your-helicone-key"
|
||||
OPENAI_API_KEY: "your-openai-key"
|
||||
ANTHROPIC_API_KEY: "your-anthropic-key"
|
||||
```
|
||||
|
||||
Start the proxy:
|
||||
```bash
|
||||
litellm --config config.yaml
|
||||
```
|
||||
|
||||
Make requests to your proxy:
|
||||
```python
|
||||
import openai
|
||||
|
||||
client = openai.OpenAI(
|
||||
api_key="anything", # proxy doesn't require real API key
|
||||
base_url="http://localhost:4000"
|
||||
)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="gpt-4", # This gets logged to Helicone
|
||||
messages=[{"role": "user", "content": "Hello!"}]
|
||||
)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Method 2: Using Helicone as a Proxy
|
||||
|
||||
Helicone's proxy provides [advanced functionality](https://docs.helicone.ai/getting-started/proxy-vs-async) like caching, rate limiting, LLM security through [PromptArmor](https://promptarmor.com/) and more.
|
||||
|
||||
To use Helicone as a proxy for your LLM requests:
|
||||
<Tabs>
|
||||
<TabItem value="sdk" label="Python SDK">
|
||||
|
||||
1. Set Helicone as your base URL via: litellm.api_base
|
||||
2. Pass in Helicone request headers via: litellm.metadata
|
||||
|
||||
Complete Code:
|
||||
Set Helicone as your base URL and pass authentication headers:
|
||||
|
||||
```python
|
||||
import os
|
||||
import litellm
|
||||
from litellm import completion
|
||||
|
||||
# Configure LiteLLM to use Helicone proxy
|
||||
litellm.api_base = "https://oai.hconeai.com/v1"
|
||||
litellm.headers = {
|
||||
"Helicone-Auth": f"Bearer {os.getenv('HELICONE_API_KEY')}", # Authenticate to send requests to Helicone API
|
||||
"Helicone-Auth": f"Bearer {os.getenv('HELICONE_API_KEY')}",
|
||||
}
|
||||
|
||||
response = litellm.completion(
|
||||
# Set your OpenAI API key
|
||||
os.environ["OPENAI_API_KEY"] = "your-openai-key"
|
||||
|
||||
response = completion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "How does a court case get to the Supreme Court?"}]
|
||||
)
|
||||
@@ -140,32 +241,112 @@ litellm.metadata = {
|
||||
|
||||
Track multi-step and agentic LLM interactions using session IDs and paths:
|
||||
|
||||
```python
|
||||
litellm.metadata = {
|
||||
"Helicone-Auth": f"Bearer {os.getenv('HELICONE_API_KEY')}", # Authenticate to send requests to Helicone API
|
||||
"Helicone-Session-Id": "session-abc-123", # The session ID you want to track
|
||||
"Helicone-Session-Path": "parent-trace/child-trace", # The path of the session
|
||||
}
|
||||
```
|
||||
|
||||
- `Helicone-Session-Id`: Use this to specify the unique identifier for the session you want to track. This allows you to group related requests together.
|
||||
- `Helicone-Session-Path`: This header defines the path of the session, allowing you to represent parent and child traces. For example, "parent/child" represents a child trace of a parent trace.
|
||||
|
||||
By using these two headers, you can effectively group and visualize multi-step LLM interactions, gaining insights into complex AI workflows.
|
||||
|
||||
### Retry and Fallback Mechanisms
|
||||
|
||||
Set up retry mechanisms and fallback options:
|
||||
<Tabs>
|
||||
<TabItem value="sdk" label="Python SDK">
|
||||
|
||||
```python
|
||||
import litellm
|
||||
|
||||
litellm.api_base = "https://oai.hconeai.com/v1"
|
||||
litellm.metadata = {
|
||||
"Helicone-Auth": f"Bearer {os.getenv('HELICONE_API_KEY')}", # Authenticate to send requests to Helicone API
|
||||
"Helicone-Retry-Enabled": "true", # Enable retry mechanism
|
||||
"helicone-retry-num": "3", # Set number of retries
|
||||
"helicone-retry-factor": "2", # Set exponential backoff factor
|
||||
"Helicone-Fallbacks": '["gpt-3.5-turbo", "gpt-4"]', # Set fallback models
|
||||
"Helicone-Auth": f"Bearer {os.getenv('HELICONE_API_KEY')}",
|
||||
"Helicone-Session-Id": "session-abc-123",
|
||||
"Helicone-Session-Path": "parent-trace/child-trace",
|
||||
}
|
||||
|
||||
response = litellm.completion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "Start a conversation"}]
|
||||
)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="proxy" label="LiteLLM Proxy">
|
||||
|
||||
```python
|
||||
import openai
|
||||
|
||||
client = openai.OpenAI(
|
||||
api_key="anything",
|
||||
base_url="http://localhost:4000"
|
||||
)
|
||||
|
||||
# First request in session
|
||||
response1 = client.chat.completions.create(
|
||||
model="gpt-4",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
extra_headers={
|
||||
"Helicone-Session-Id": "session-abc-123",
|
||||
"Helicone-Session-Path": "conversation/greeting"
|
||||
}
|
||||
)
|
||||
|
||||
# Follow-up request in same session
|
||||
response2 = client.chat.completions.create(
|
||||
model="gpt-4",
|
||||
messages=[{"role": "user", "content": "Tell me more"}],
|
||||
extra_headers={
|
||||
"Helicone-Session-Id": "session-abc-123",
|
||||
"Helicone-Session-Path": "conversation/follow-up"
|
||||
}
|
||||
)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
- `Helicone-Session-Id`: Unique identifier for the session to group related requests
|
||||
- `Helicone-Session-Path`: Hierarchical path to represent parent/child traces (e.g., "parent/child")
|
||||
|
||||
## Retry and Fallback Mechanisms
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="sdk" label="Python SDK">
|
||||
|
||||
```python
|
||||
import litellm
|
||||
|
||||
litellm.api_base = "https://oai.hconeai.com/v1"
|
||||
litellm.metadata = {
|
||||
"Helicone-Auth": f"Bearer {os.getenv('HELICONE_API_KEY')}",
|
||||
"Helicone-Retry-Enabled": "true",
|
||||
"helicone-retry-num": "3",
|
||||
"helicone-retry-factor": "2", # Exponential backoff
|
||||
"Helicone-Fallbacks": '["gpt-3.5-turbo", "gpt-4"]',
|
||||
}
|
||||
|
||||
response = litellm.completion(
|
||||
model="gpt-4",
|
||||
messages=[{"role": "user", "content": "Hello"}]
|
||||
)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="proxy" label="LiteLLM Proxy">
|
||||
|
||||
```yaml title="config.yaml"
|
||||
model_list:
|
||||
- model_name: gpt-4
|
||||
litellm_params:
|
||||
model: gpt-4
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
api_base: "https://oai.hconeai.com/v1"
|
||||
|
||||
default_litellm_params:
|
||||
headers:
|
||||
Helicone-Auth: "Bearer ${HELICONE_API_KEY}"
|
||||
Helicone-Retry-Enabled: "true"
|
||||
helicone-retry-num: "3"
|
||||
helicone-retry-factor: "2"
|
||||
Helicone-Fallbacks: '["gpt-3.5-turbo", "gpt-4"]'
|
||||
|
||||
environment_variables:
|
||||
HELICONE_API_KEY: "your-helicone-key"
|
||||
OPENAI_API_KEY: "your-openai-key"
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
> **Supported Headers** - For a full list of supported Helicone headers and their descriptions, please refer to the [Helicone documentation](https://docs.helicone.ai/getting-started/quick-start).
|
||||
> By utilizing these headers and metadata options, you can gain deeper insights into your LLM usage, optimize performance, and better manage your AI workflows with Helicone and LiteLLM.
|
||||
|
||||
@@ -889,6 +889,19 @@ curl http://0.0.0.0:4000/v1/chat/completions \
|
||||
|
||||
Example of using [Bedrock Guardrails with LiteLLM](https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails-use-converse-api.html)
|
||||
|
||||
### Selective Content Moderation with `guarded_text`
|
||||
|
||||
LiteLLM supports selective content moderation using the `guarded_text` content type. This allows you to wrap only specific content that should be moderated by Bedrock Guardrails, rather than evaluating the entire conversation.
|
||||
|
||||
**How it works:**
|
||||
- Content with `type: "guarded_text"` gets automatically wrapped in `guardrailConverseContent` blocks
|
||||
- Only the wrapped content is evaluated by Bedrock Guardrails
|
||||
- Regular content with `type: "text"` bypasses guardrail evaluation
|
||||
|
||||
:::note
|
||||
If `guarded_text` is not used, the entire conversation history will be sent to the guardrail for evaluation, which can increase latency and costs.
|
||||
:::
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="sdk" label="LiteLLM SDK">
|
||||
|
||||
@@ -915,6 +928,24 @@ response = completion(
|
||||
"trace": "disabled", # The trace behavior for the guardrail. Can either be "disabled" or "enabled"
|
||||
},
|
||||
)
|
||||
|
||||
# Selective guardrail usage with guarded_text - only specific content is evaluated
|
||||
response_guard = completion(
|
||||
model="anthropic.claude-v2",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is the main topic of this legal document?"},
|
||||
{"type": "guarded_text", "text": "This document contains sensitive legal information that should be moderated by guardrails."}
|
||||
]
|
||||
}
|
||||
],
|
||||
guardrailConfig={
|
||||
"guardrailIdentifier": "gr-abc123",
|
||||
"guardrailVersion": "DRAFT"
|
||||
}
|
||||
)
|
||||
```
|
||||
</TabItem>
|
||||
<TabItem value="proxy" label="Proxy on request">
|
||||
@@ -993,7 +1024,20 @@ response = client.chat.completions.create(model="bedrock-claude-v1", messages =
|
||||
temperature=0.7
|
||||
)
|
||||
|
||||
print(response)
|
||||
# For adding selective guardrail usage with guarded_text
|
||||
response_guard = client.chat.completions.create(model="bedrock-claude-v1", messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is the main topic of this legal document?"},
|
||||
{"type": "guarded_text", "text": "This document contains sensitive legal information that should be moderated by guardrails."}
|
||||
]
|
||||
}
|
||||
],
|
||||
temperature=0.7
|
||||
)
|
||||
|
||||
print(response_guard)
|
||||
```
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
@@ -4,7 +4,7 @@ import TabItem from '@theme/TabItem';
|
||||
# CompactifAI
|
||||
https://docs.compactif.ai/
|
||||
|
||||
CompactifAI offers highly compressed versions of leading language models, delivering up to **70% lower inference costs**, **4x throughput gains**, and **low-latency inference** with minimal quality loss (<5%). CompactifAI's OpenAI-compatible API makes integration straightforward, enabling developers to build ultra-efficient, scalable AI applications with superior concurrency and resource efficiency.
|
||||
CompactifAI offers highly compressed versions of leading language models, delivering up to **70% lower inference costs**, **4x throughput gains**, and **low-latency inference** with minimal quality loss (under 5%). CompactifAI's OpenAI-compatible API makes integration straightforward, enabling developers to build ultra-efficient, scalable AI applications with superior concurrency and resource efficiency.
|
||||
|
||||
| Property | Details |
|
||||
|-------|-------|
|
||||
@@ -192,7 +192,7 @@ Common model formats:
|
||||
## Benefits
|
||||
|
||||
- **Cost Efficient**: Up to 70% lower inference costs compared to standard models
|
||||
- **High Performance**: 4x throughput gains with minimal quality loss (<5%)
|
||||
- **High Performance**: 4x throughput gains with minimal quality loss (under 5%)
|
||||
- **Low Latency**: Optimized for fast response times
|
||||
- **Drop-in Replacement**: Full OpenAI API compatibility
|
||||
- **Scalable**: Superior concurrency and resource efficiency
|
||||
|
||||
@@ -2758,6 +2758,44 @@ curl http://localhost:4000/v1/fine_tuning/jobs \
|
||||
</Tabs>
|
||||
|
||||
|
||||
## Labels
|
||||
|
||||
|
||||
Google enables you to add custom metadata to its `generateContent` and `streamGenerateContent` calls.
|
||||
This mechanism is useful in Vertex AI because it allows costs and usage tracking over multiple
|
||||
different applications or users.
|
||||
|
||||
|
||||
### Usage
|
||||
|
||||
You can use that feature through LiteLLM by sending `labels` or `metadata` field in your requests.
|
||||
|
||||
If the client sets the `labels` field in the request to the LiteLLM,
|
||||
the LiteLLM will pass the `labels` field to the Vertex AI backend.
|
||||
|
||||
If the client sets the `metadata` field in the request to the LiteLLM and the `labels` field is not set,
|
||||
the LiteLLM will create the `labels` field filled with `metadata` key/value pairs for all string values and
|
||||
pass it to the Vertex AI backend.
|
||||
|
||||
|
||||
Here is an example JSON request demonstrating the labels usage:
|
||||
|
||||
```json
|
||||
{
|
||||
"model": "gemini-2.0-flash-lite",
|
||||
"messages": [
|
||||
{ "role": "user", "content": "respond in 20 words. who are you?" }
|
||||
],
|
||||
"labels": {
|
||||
"client_app": "acme_comp_financial_app",
|
||||
"department": "finance",
|
||||
"project": "acme_ai"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
|
||||
## Extra
|
||||
|
||||
### Using `GOOGLE_APPLICATION_CREDENTIALS`
|
||||
|
||||
@@ -93,6 +93,8 @@ callback_settings:
|
||||
|
||||
general_settings:
|
||||
completion_model: string
|
||||
store_prompts_in_spend_logs: boolean
|
||||
forward_client_headers_to_llm_api: boolean
|
||||
disable_spend_logs: boolean # turn off writing each transaction to the db
|
||||
disable_master_key_return: boolean # turn off returning master key on UI (checked on '/user/info' endpoint)
|
||||
disable_retry_on_max_parallel_request_limit_error: boolean # turn off retries when max parallel request limit is reached
|
||||
@@ -121,6 +123,35 @@ general_settings:
|
||||
alerting: ["slack", "email"]
|
||||
alerting_threshold: 0
|
||||
use_client_credentials_pass_through_routes: boolean # use client credentials for all pass through routes like "/vertex-ai", /bedrock/. When this is True Virtual Key auth will not be applied on these endpoints
|
||||
|
||||
router_settings:
|
||||
routing_strategy: simple-shuffle # Literal["simple-shuffle", "least-busy", "usage-based-routing","latency-based-routing"], default="simple-shuffle" - RECOMMENDED for best performance
|
||||
redis_host: <your-redis-host> # string
|
||||
redis_password: <your-redis-password> # string
|
||||
redis_port: <your-redis-port> # string
|
||||
enable_pre_call_checks: true # bool - Before call is made check if a call is within model context window
|
||||
allowed_fails: 3 # cooldown model if it fails > 1 call in a minute.
|
||||
cooldown_time: 30 # (in seconds) how long to cooldown model if fails/min > allowed_fails
|
||||
disable_cooldowns: True # bool - Disable cooldowns for all models
|
||||
enable_tag_filtering: True # bool - Use tag based routing for requests
|
||||
retry_policy: { # Dict[str, int]: retry policy for different types of exceptions
|
||||
"AuthenticationErrorRetries": 3,
|
||||
"TimeoutErrorRetries": 3,
|
||||
"RateLimitErrorRetries": 3,
|
||||
"ContentPolicyViolationErrorRetries": 4,
|
||||
"InternalServerErrorRetries": 4
|
||||
}
|
||||
allowed_fails_policy: {
|
||||
"BadRequestErrorAllowedFails": 1000, # Allow 1000 BadRequestErrors before cooling down a deployment
|
||||
"AuthenticationErrorAllowedFails": 10, # int
|
||||
"TimeoutErrorAllowedFails": 12, # int
|
||||
"RateLimitErrorAllowedFails": 10000, # int
|
||||
"ContentPolicyViolationErrorAllowedFails": 15, # int
|
||||
"InternalServerErrorAllowedFails": 20, # int
|
||||
}
|
||||
content_policy_fallbacks=[{"claude-2": ["my-fallback-model"]}] # List[Dict[str, List[str]]]: Fallback model for content policy violations
|
||||
fallbacks=[{"claude-2": ["my-fallback-model"]}] # List[Dict[str, List[str]]]: Fallback model for all errors
|
||||
|
||||
```
|
||||
|
||||
### litellm_settings - Reference
|
||||
|
||||
@@ -61,6 +61,11 @@ Inherits from `StandardLoggingUserAPIKeyMetadata` and adds:
|
||||
| `requester_metadata` | `Optional[dict]` | Additional requester metadata |
|
||||
| `vector_store_request_metadata` | `Optional[List[StandardLoggingVectorStoreRequest]]` | Vector store request metadata |
|
||||
| `requester_custom_headers` | Dict[str, str] | Any custom (`x-`) headers sent by the client to the proxy. |
|
||||
| `prompt_management_metadata` | `Optional[StandardLoggingPromptManagementMetadata]` | Prompt management and versioning metadata |
|
||||
| `mcp_tool_call_metadata` | `Optional[StandardLoggingMCPToolCall]` | MCP (Model Context Protocol) tool call information and cost tracking |
|
||||
| `applied_guardrails` | `Optional[List[str]]` | List of applied guardrail names |
|
||||
| `usage_object` | `Optional[dict]` | Raw usage object from the LLM provider |
|
||||
| `cold_storage_object_key` | `Optional[str]` | S3/GCS object key for cold storage retrieval |
|
||||
| `guardrail_information` | `Optional[StandardLoggingGuardrailInformation]` | Guardrail information |
|
||||
|
||||
|
||||
@@ -145,4 +150,82 @@ A literal type with two possible values:
|
||||
| `duration` | `Optional[float]` | Duration of the guardrail in seconds |
|
||||
| `masked_entity_count` | `Optional[Dict[str, int]]` | Count of masked entities |
|
||||
|
||||
## StandardLoggingPromptManagementMetadata
|
||||
|
||||
Used for tracking prompt versioning and management information.
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `prompt_id` | `str` | **Required**. Unique identifier for the prompt template or version |
|
||||
| `prompt_variables` | `Optional[dict]` | Variables/parameters used in the prompt template (e.g., `{"user_name": "John", "context": "support"}`) |
|
||||
| `prompt_integration` | `str` | **Required**. Integration or system managing the prompt (e.g., `"langfuse"`, `"promptlayer"`, `"custom"`) |
|
||||
|
||||
## StandardLoggingMCPToolCall
|
||||
|
||||
Used to track Model Context Protocol (MCP) tool calls within LiteLLM requests. This provides detailed logging for external tool integrations.
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `name` | `str` | **Required**. The name of the tool being called (e.g., `"get_weather"`, `"search_database"`) |
|
||||
| `arguments` | `dict` | **Required**. Arguments passed to the tool as key-value pairs |
|
||||
| `result` | `Optional[dict]` | The response/result returned by the tool execution (populated by custom logging hooks) |
|
||||
| `mcp_server_name` | `Optional[str]` | Name of the MCP server that handled the tool call (e.g., `"weather-service"`, `"database-connector"`) |
|
||||
| `mcp_server_logo_url` | `Optional[str]` | URL for the MCP server's logo (used for UI display in LiteLLM dashboard) |
|
||||
| `namespaced_tool_name` | `Optional[str]` | Fully qualified tool name including server prefix (e.g., `"deepwiki-mcp/get_page_content"`, `"github-mcp/create_issue"`) |
|
||||
| `mcp_server_cost_info` | `Optional[MCPServerCostInfo]` | Cost tracking information for the tool call |
|
||||
|
||||
### MCPServerCostInfo
|
||||
|
||||
Cost tracking structure for MCP server tool calls:
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `default_cost_per_query` | `Optional[float]` | Default cost in USD for any tool call to this MCP server |
|
||||
| `tool_name_to_cost_per_query` | `Optional[Dict[str, float]]` | Per-tool cost mapping for granular pricing (e.g., `{"search": 0.01, "create": 0.05}`) |
|
||||
|
||||
### Usage
|
||||
|
||||
```python
|
||||
# Basic MCP tool call metadata
|
||||
mcp_tool_call = {
|
||||
"name": "search_documents",
|
||||
"arguments": {
|
||||
"query": "machine learning tutorials",
|
||||
"limit": 10,
|
||||
"filter": "type:pdf"
|
||||
},
|
||||
"mcp_server_name": "document-search-service",
|
||||
"namespaced_tool_name": "docs-mcp/search_documents",
|
||||
"mcp_server_cost_info": {
|
||||
"default_cost_per_query": 0.02,
|
||||
"tool_name_to_cost_per_query": {
|
||||
"search_documents": 0.02,
|
||||
"get_document": 0.01
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# optional result field (via custom logging hooks)
|
||||
mcp_tool_call_with_result = {
|
||||
"name": "search_documents",
|
||||
"arguments": {
|
||||
"query": "machine learning tutorials",
|
||||
"limit": 10,
|
||||
"filter": "type:pdf"
|
||||
},
|
||||
"result": {
|
||||
"documents": [...],
|
||||
"total_found": 42,
|
||||
"search_time_ms": 150
|
||||
},
|
||||
"mcp_server_name": "document-search-service",
|
||||
"namespaced_tool_name": "docs-mcp/search_documents",
|
||||
"mcp_server_cost_info": {
|
||||
"default_cost_per_query": 0.02,
|
||||
"tool_name_to_cost_per_query": {
|
||||
"search_documents": 0.02,
|
||||
"get_document": 0.01
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
@@ -2,7 +2,7 @@
|
||||
[Schedule Demo 👋](https://calendly.com/d/4mp-gd3-k5k/berriai-1-1-onboarding-litellm-hosted-version)
|
||||
|
||||
[Community Discord 💭](https://discord.gg/wuPM9dRgDw)
|
||||
[Community Slack 💭](https://join.slack.com/share/enQtOTE0ODczMzk2Nzk4NC01YjUxNjY2YjBlYTFmNDRiZTM3NDFiYTM3MzVkODFiMDVjOGRjMmNmZTZkZTMzOWQzZGQyZWIwYjQ0MWExYmE3)
|
||||
[Community Slack 💭](https://litellmossslack.slack.com/)
|
||||
|
||||
Our numbers 📞 +1 (770) 8783-106 / +1 (412) 618-6238
|
||||
|
||||
|
||||
+169
-107
@@ -340,7 +340,7 @@ def create_batch(
|
||||
@client
|
||||
async def aretrieve_batch(
|
||||
batch_id: str,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai",
|
||||
metadata: Optional[Dict[str, str]] = None,
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
extra_body: Optional[Dict[str, str]] = None,
|
||||
@@ -378,11 +378,129 @@ async def aretrieve_batch(
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
def _handle_retrieve_batch_providers_without_provider_config(
|
||||
batch_id: str,
|
||||
optional_params: GenericLiteLLMParams,
|
||||
timeout: Union[float, httpx.Timeout],
|
||||
litellm_params: dict,
|
||||
_retrieve_batch_request: RetrieveBatchRequest,
|
||||
_is_async: bool,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai",
|
||||
):
|
||||
api_base: Optional[str] = None
|
||||
if custom_llm_provider == "openai":
|
||||
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
|
||||
api_base = (
|
||||
optional_params.api_base
|
||||
or litellm.api_base
|
||||
or os.getenv("OPENAI_BASE_URL")
|
||||
or os.getenv("OPENAI_API_BASE")
|
||||
or "https://api.openai.com/v1"
|
||||
)
|
||||
organization = (
|
||||
optional_params.organization
|
||||
or litellm.organization
|
||||
or os.getenv("OPENAI_ORGANIZATION", None)
|
||||
or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
|
||||
)
|
||||
# set API KEY
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
|
||||
or litellm.openai_key
|
||||
or os.getenv("OPENAI_API_KEY")
|
||||
)
|
||||
|
||||
response = openai_batches_instance.retrieve_batch(
|
||||
_is_async=_is_async,
|
||||
retrieve_batch_data=_retrieve_batch_request,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
organization=organization,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
)
|
||||
elif custom_llm_provider == "azure":
|
||||
api_base = (
|
||||
optional_params.api_base
|
||||
or litellm.api_base
|
||||
or get_secret_str("AZURE_API_BASE")
|
||||
)
|
||||
api_version = (
|
||||
optional_params.api_version
|
||||
or litellm.api_version
|
||||
or get_secret_str("AZURE_API_VERSION")
|
||||
)
|
||||
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
or litellm.api_key
|
||||
or litellm.azure_key
|
||||
or get_secret_str("AZURE_OPENAI_API_KEY")
|
||||
or get_secret_str("AZURE_API_KEY")
|
||||
)
|
||||
|
||||
extra_body = optional_params.get("extra_body", {})
|
||||
if extra_body is not None:
|
||||
extra_body.pop("azure_ad_token", None)
|
||||
else:
|
||||
get_secret_str("AZURE_AD_TOKEN") # type: ignore
|
||||
|
||||
response = azure_batches_instance.retrieve_batch(
|
||||
_is_async=_is_async,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
api_version=api_version,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
retrieve_batch_data=_retrieve_batch_request,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
elif custom_llm_provider == "vertex_ai":
|
||||
api_base = optional_params.api_base or ""
|
||||
vertex_ai_project = (
|
||||
optional_params.vertex_project
|
||||
or litellm.vertex_project
|
||||
or get_secret_str("VERTEXAI_PROJECT")
|
||||
)
|
||||
vertex_ai_location = (
|
||||
optional_params.vertex_location
|
||||
or litellm.vertex_location
|
||||
or get_secret_str("VERTEXAI_LOCATION")
|
||||
)
|
||||
vertex_credentials = optional_params.vertex_credentials or get_secret_str(
|
||||
"VERTEXAI_CREDENTIALS"
|
||||
)
|
||||
|
||||
response = vertex_ai_batches_instance.retrieve_batch(
|
||||
_is_async=_is_async,
|
||||
batch_id=batch_id,
|
||||
api_base=api_base,
|
||||
vertex_project=vertex_ai_project,
|
||||
vertex_location=vertex_ai_location,
|
||||
vertex_credentials=vertex_credentials,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
)
|
||||
else:
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format(
|
||||
custom_llm_provider
|
||||
),
|
||||
model="n/a",
|
||||
llm_provider=custom_llm_provider,
|
||||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
),
|
||||
)
|
||||
return response
|
||||
|
||||
@client
|
||||
def retrieve_batch(
|
||||
batch_id: str,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai",
|
||||
metadata: Optional[Dict[str, str]] = None,
|
||||
extra_headers: Optional[Dict[str, str]] = None,
|
||||
extra_body: Optional[Dict[str, str]] = None,
|
||||
@@ -430,115 +548,59 @@ def retrieve_batch(
|
||||
)
|
||||
|
||||
_is_async = kwargs.pop("aretrieve_batch", False) is True
|
||||
api_base: Optional[str] = None
|
||||
if custom_llm_provider == "openai":
|
||||
# for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there
|
||||
api_base = (
|
||||
optional_params.api_base
|
||||
or litellm.api_base
|
||||
or os.getenv("OPENAI_BASE_URL")
|
||||
or os.getenv("OPENAI_API_BASE")
|
||||
or "https://api.openai.com/v1"
|
||||
)
|
||||
organization = (
|
||||
optional_params.organization
|
||||
or litellm.organization
|
||||
or os.getenv("OPENAI_ORGANIZATION", None)
|
||||
or None # default - https://github.com/openai/openai-python/blob/284c1799070c723c6a553337134148a7ab088dd8/openai/util.py#L105
|
||||
)
|
||||
# set API KEY
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
or litellm.api_key # for deepinfra/perplexity/anyscale we check in get_llm_provider and pass in the api key from there
|
||||
or litellm.openai_key
|
||||
or os.getenv("OPENAI_API_KEY")
|
||||
)
|
||||
|
||||
response = openai_batches_instance.retrieve_batch(
|
||||
_is_async=_is_async,
|
||||
retrieve_batch_data=_retrieve_batch_request,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
organization=organization,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
)
|
||||
elif custom_llm_provider == "azure":
|
||||
api_base = (
|
||||
optional_params.api_base
|
||||
or litellm.api_base
|
||||
or get_secret_str("AZURE_API_BASE")
|
||||
)
|
||||
api_version = (
|
||||
optional_params.api_version
|
||||
or litellm.api_version
|
||||
or get_secret_str("AZURE_API_VERSION")
|
||||
)
|
||||
|
||||
api_key = (
|
||||
optional_params.api_key
|
||||
or litellm.api_key
|
||||
or litellm.azure_key
|
||||
or get_secret_str("AZURE_OPENAI_API_KEY")
|
||||
or get_secret_str("AZURE_API_KEY")
|
||||
)
|
||||
|
||||
extra_body = optional_params.get("extra_body", {})
|
||||
if extra_body is not None:
|
||||
extra_body.pop("azure_ad_token", None)
|
||||
else:
|
||||
get_secret_str("AZURE_AD_TOKEN") # type: ignore
|
||||
|
||||
response = azure_batches_instance.retrieve_batch(
|
||||
_is_async=_is_async,
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
api_version=api_version,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
retrieve_batch_data=_retrieve_batch_request,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
elif custom_llm_provider == "vertex_ai":
|
||||
api_base = optional_params.api_base or ""
|
||||
vertex_ai_project = (
|
||||
optional_params.vertex_project
|
||||
or litellm.vertex_project
|
||||
or get_secret_str("VERTEXAI_PROJECT")
|
||||
)
|
||||
vertex_ai_location = (
|
||||
optional_params.vertex_location
|
||||
or litellm.vertex_location
|
||||
or get_secret_str("VERTEXAI_LOCATION")
|
||||
)
|
||||
vertex_credentials = optional_params.vertex_credentials or get_secret_str(
|
||||
"VERTEXAI_CREDENTIALS"
|
||||
)
|
||||
|
||||
response = vertex_ai_batches_instance.retrieve_batch(
|
||||
_is_async=_is_async,
|
||||
batch_id=batch_id,
|
||||
api_base=api_base,
|
||||
vertex_project=vertex_ai_project,
|
||||
vertex_location=vertex_ai_location,
|
||||
vertex_credentials=vertex_credentials,
|
||||
timeout=timeout,
|
||||
max_retries=optional_params.max_retries,
|
||||
client = kwargs.get("client", None)
|
||||
|
||||
# Try to use provider config first (for providers like bedrock)
|
||||
model: Optional[str] = kwargs.get("model", None)
|
||||
if model is not None:
|
||||
provider_config = ProviderConfigManager.get_provider_batches_config(
|
||||
model=model,
|
||||
provider=LlmProviders(custom_llm_provider),
|
||||
)
|
||||
else:
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format(
|
||||
custom_llm_provider
|
||||
),
|
||||
model="n/a",
|
||||
llm_provider=custom_llm_provider,
|
||||
response=httpx.Response(
|
||||
status_code=400,
|
||||
content="Unsupported provider",
|
||||
request=httpx.Request(method="create_thread", url="https://github.com/BerriAI/litellm"), # type: ignore
|
||||
provider_config = None
|
||||
|
||||
if provider_config is not None:
|
||||
response = base_llm_http_handler.retrieve_batch(
|
||||
batch_id=batch_id,
|
||||
provider_config=provider_config,
|
||||
litellm_params=litellm_params,
|
||||
headers=extra_headers or {},
|
||||
api_base=optional_params.api_base,
|
||||
api_key=optional_params.api_key,
|
||||
logging_obj=litellm_logging_obj or LiteLLMLoggingObj(
|
||||
model=model or "bedrock/unknown",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="batch_retrieve",
|
||||
start_time=None,
|
||||
litellm_call_id="batch_retrieve_" + batch_id,
|
||||
function_id="batch_retrieve",
|
||||
),
|
||||
_is_async=_is_async,
|
||||
client=client
|
||||
if client is not None
|
||||
and isinstance(client, (HTTPHandler, AsyncHTTPHandler))
|
||||
else None,
|
||||
timeout=timeout,
|
||||
model=model,
|
||||
)
|
||||
return response
|
||||
return response
|
||||
|
||||
|
||||
#########################################################
|
||||
# Handle providers without provider config
|
||||
#########################################################
|
||||
return _handle_retrieve_batch_providers_without_provider_config(
|
||||
batch_id=batch_id,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
_retrieve_batch_request=_retrieve_batch_request,
|
||||
_is_async=_is_async,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
||||
@@ -148,10 +148,22 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger):
|
||||
),
|
||||
),
|
||||
}
|
||||
verbose_logger.debug("payload %s", json.dumps(payload, indent=4))
|
||||
|
||||
# serialize datetime objects - for budget reset time in spend metrics
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
|
||||
try:
|
||||
verbose_logger.debug("payload %s", safe_dumps(payload))
|
||||
except Exception as debug_error:
|
||||
verbose_logger.debug(
|
||||
"payload serialization failed: %s", str(debug_error)
|
||||
)
|
||||
|
||||
json_payload = safe_dumps(payload)
|
||||
|
||||
response = await self.async_client.post(
|
||||
url=self.intake_url,
|
||||
json=payload,
|
||||
content=json_payload,
|
||||
headers={
|
||||
"DD-API-KEY": self.DD_API_KEY,
|
||||
"Content-Type": "application/json",
|
||||
@@ -494,6 +506,12 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger):
|
||||
latency_metrics = self._get_latency_metrics(standard_logging_payload)
|
||||
_metadata.update({"latency_metrics": dict(latency_metrics)})
|
||||
|
||||
#########################################################
|
||||
# Add spend metrics to metadata
|
||||
#########################################################
|
||||
spend_metrics = self._get_spend_metrics(standard_logging_payload)
|
||||
_metadata.update({"spend_metrics": dict(spend_metrics)})
|
||||
|
||||
## extract tool calls and add to metadata
|
||||
tool_call_metadata = self._extract_tool_call_metadata(standard_logging_payload)
|
||||
_metadata.update(tool_call_metadata)
|
||||
@@ -543,6 +561,71 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger):
|
||||
|
||||
return latency_metrics
|
||||
|
||||
def _get_spend_metrics(
|
||||
self, standard_logging_payload: StandardLoggingPayload
|
||||
) -> DDLLMObsSpendMetrics:
|
||||
"""
|
||||
Get the spend metrics from the standard logging payload
|
||||
"""
|
||||
spend_metrics: DDLLMObsSpendMetrics = DDLLMObsSpendMetrics()
|
||||
|
||||
# send response cost
|
||||
spend_metrics["response_cost"] = standard_logging_payload.get(
|
||||
"response_cost", 0.0
|
||||
)
|
||||
|
||||
# Get budget information from metadata
|
||||
metadata = standard_logging_payload.get("metadata", {})
|
||||
|
||||
# API key max budget
|
||||
user_api_key_max_budget = metadata.get("user_api_key_max_budget")
|
||||
if user_api_key_max_budget is not None:
|
||||
spend_metrics["user_api_key_max_budget"] = float(user_api_key_max_budget)
|
||||
|
||||
# API key spend
|
||||
user_api_key_spend = metadata.get("user_api_key_spend")
|
||||
if user_api_key_spend is not None:
|
||||
try:
|
||||
spend_metrics["user_api_key_spend"] = float(user_api_key_spend)
|
||||
except (ValueError, TypeError):
|
||||
verbose_logger.debug(
|
||||
f"Invalid user_api_key_spend value: {user_api_key_spend}"
|
||||
)
|
||||
|
||||
# API key budget reset datetime
|
||||
user_api_key_budget_reset_at = metadata.get("user_api_key_budget_reset_at")
|
||||
if user_api_key_budget_reset_at is not None:
|
||||
try:
|
||||
from datetime import datetime, timezone
|
||||
|
||||
budget_reset_at = None
|
||||
if isinstance(user_api_key_budget_reset_at, str):
|
||||
# Handle ISO format strings that might have 'Z' suffix
|
||||
iso_string = user_api_key_budget_reset_at.replace("Z", "+00:00")
|
||||
budget_reset_at = datetime.fromisoformat(iso_string)
|
||||
elif isinstance(user_api_key_budget_reset_at, datetime):
|
||||
budget_reset_at = user_api_key_budget_reset_at
|
||||
|
||||
if budget_reset_at is not None:
|
||||
# Preserve timezone info if already present
|
||||
if budget_reset_at.tzinfo is None:
|
||||
budget_reset_at = budget_reset_at.replace(tzinfo=timezone.utc)
|
||||
|
||||
# Convert to ISO string format for JSON serialization
|
||||
# This prevents circular reference issues and ensures proper timezone representation
|
||||
iso_string = budget_reset_at.isoformat()
|
||||
spend_metrics["user_api_key_budget_reset_at"] = iso_string
|
||||
|
||||
# Debug logging to verify the conversion
|
||||
verbose_logger.debug(
|
||||
f"Converted budget_reset_at to ISO format: {iso_string}"
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.debug(f"Error processing budget reset datetime: {e}")
|
||||
verbose_logger.debug(f"Original value: {user_api_key_budget_reset_at}")
|
||||
|
||||
return spend_metrics
|
||||
|
||||
def _process_input_messages_preserving_tool_calls(
|
||||
self, messages: List[Any]
|
||||
) -> List[Dict[str, Any]]:
|
||||
|
||||
@@ -15,7 +15,7 @@ from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info
|
||||
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.integrations.langfuse import *
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse
|
||||
from litellm.types.utils import (
|
||||
EmbeddingResponse,
|
||||
ImageResponse,
|
||||
@@ -196,6 +196,7 @@ class LangFuseLogger:
|
||||
TranscriptionResponse,
|
||||
RerankResponse,
|
||||
HttpxBinaryResponseContent,
|
||||
ResponsesAPIResponse,
|
||||
],
|
||||
start_time: Optional[datetime] = None,
|
||||
end_time: Optional[datetime] = None,
|
||||
@@ -305,6 +306,7 @@ class LangFuseLogger:
|
||||
TranscriptionResponse,
|
||||
RerankResponse,
|
||||
HttpxBinaryResponseContent,
|
||||
ResponsesAPIResponse,
|
||||
],
|
||||
prompt: dict,
|
||||
level: str,
|
||||
@@ -369,6 +371,11 @@ class LangFuseLogger:
|
||||
):
|
||||
input = prompt
|
||||
output = response_obj.results
|
||||
elif response_obj is not None and isinstance(
|
||||
response_obj, litellm.ResponsesAPIResponse
|
||||
):
|
||||
input = prompt
|
||||
output = self._get_responses_api_content_for_langfuse(response_obj)
|
||||
elif (
|
||||
kwargs.get("call_type") is not None
|
||||
and kwargs.get("call_type") == "_arealtime"
|
||||
@@ -775,6 +782,19 @@ class LangFuseLogger:
|
||||
else:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_responses_api_content_for_langfuse(
|
||||
response_obj: ResponsesAPIResponse,
|
||||
):
|
||||
"""
|
||||
Get the responses API content for Langfuse logging
|
||||
"""
|
||||
if hasattr(response_obj, 'output') and response_obj.output:
|
||||
# ResponsesAPIResponse.output is a list of strings
|
||||
return response_obj.output
|
||||
else:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_langfuse_tags(
|
||||
standard_logging_object: Optional[StandardLoggingPayload],
|
||||
|
||||
@@ -228,9 +228,11 @@ def safe_deep_copy(data):
|
||||
"""
|
||||
Safe Deep Copy
|
||||
|
||||
The LiteLLM Request has some object that can-not be pickled / deep copied
|
||||
|
||||
Use this function to safely deep copy the LiteLLM Request
|
||||
The LiteLLM request may contain objects that cannot be pickled/deep-copied
|
||||
(e.g., tracing spans, locks, clients).
|
||||
|
||||
This helper deep-copies each top-level key independently; on failure keeps
|
||||
original ref
|
||||
"""
|
||||
import copy
|
||||
|
||||
@@ -255,9 +257,22 @@ def safe_deep_copy(data):
|
||||
"litellm_parent_otel_span"
|
||||
)
|
||||
data["litellm_metadata"]["litellm_parent_otel_span"] = "placeholder"
|
||||
new_data = copy.deepcopy(data)
|
||||
|
||||
# Step 2: re-add the litellm_parent_otel_span after doing a deep copy
|
||||
# Step 2: Per-key deepcopy with fallback
|
||||
if isinstance(data, dict):
|
||||
new_data = {}
|
||||
for k, v in data.items():
|
||||
try:
|
||||
new_data[k] = copy.deepcopy(v)
|
||||
except Exception:
|
||||
new_data[k] = v
|
||||
else:
|
||||
try:
|
||||
new_data = copy.deepcopy(data)
|
||||
except Exception:
|
||||
new_data = data
|
||||
|
||||
# Step 3: re-add the litellm_parent_otel_span after doing a deep copy
|
||||
if isinstance(data, dict) and litellm_parent_otel_span is not None:
|
||||
if "metadata" in data and "litellm_parent_otel_span" in data["metadata"]:
|
||||
data["metadata"]["litellm_parent_otel_span"] = litellm_parent_otel_span
|
||||
@@ -268,4 +283,4 @@ def safe_deep_copy(data):
|
||||
data["litellm_metadata"][
|
||||
"litellm_parent_otel_span"
|
||||
] = litellm_parent_otel_span
|
||||
return new_data
|
||||
return new_data
|
||||
@@ -300,9 +300,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
self.litellm_trace_id: str = litellm_trace_id or str(uuid.uuid4())
|
||||
self.function_id = function_id
|
||||
self.streaming_chunks: List[Any] = [] # for generating complete stream response
|
||||
self.sync_streaming_chunks: List[Any] = (
|
||||
[]
|
||||
) # for generating complete stream response
|
||||
self.sync_streaming_chunks: List[
|
||||
Any
|
||||
] = [] # for generating complete stream response
|
||||
self.log_raw_request_response = log_raw_request_response
|
||||
|
||||
# Initialize dynamic callbacks
|
||||
@@ -672,24 +672,23 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook(
|
||||
non_default_params
|
||||
):
|
||||
self.model_call_details["prompt_integration"] = (
|
||||
anthropic_cache_control_logger.__class__.__name__
|
||||
)
|
||||
self.model_call_details[
|
||||
"prompt_integration"
|
||||
] = anthropic_cache_control_logger.__class__.__name__
|
||||
return anthropic_cache_control_logger
|
||||
|
||||
#########################################################
|
||||
# Vector Store / Knowledge Base hooks
|
||||
#########################################################
|
||||
if litellm.vector_store_registry is not None:
|
||||
|
||||
vector_store_custom_logger = _init_custom_logger_compatible_class(
|
||||
logging_integration="vector_store_pre_call_hook",
|
||||
internal_usage_cache=None,
|
||||
llm_router=None,
|
||||
)
|
||||
self.model_call_details["prompt_integration"] = (
|
||||
vector_store_custom_logger.__class__.__name__
|
||||
)
|
||||
self.model_call_details[
|
||||
"prompt_integration"
|
||||
] = vector_store_custom_logger.__class__.__name__
|
||||
return vector_store_custom_logger
|
||||
|
||||
return None
|
||||
@@ -741,9 +740,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
model
|
||||
): # if model name was changes pre-call, overwrite the initial model call name with the new one
|
||||
self.model_call_details["model"] = model
|
||||
self.model_call_details["litellm_params"]["api_base"] = (
|
||||
self._get_masked_api_base(additional_args.get("api_base", ""))
|
||||
)
|
||||
self.model_call_details["litellm_params"][
|
||||
"api_base"
|
||||
] = self._get_masked_api_base(additional_args.get("api_base", ""))
|
||||
|
||||
def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR0915
|
||||
# Log the exact input to the LLM API
|
||||
@@ -772,10 +771,10 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
try:
|
||||
# [Non-blocking Extra Debug Information in metadata]
|
||||
if turn_off_message_logging is True:
|
||||
_metadata["raw_request"] = (
|
||||
"redacted by litellm. \
|
||||
_metadata[
|
||||
"raw_request"
|
||||
] = "redacted by litellm. \
|
||||
'litellm.turn_off_message_logging=True'"
|
||||
)
|
||||
else:
|
||||
curl_command = self._get_request_curl_command(
|
||||
api_base=additional_args.get("api_base", ""),
|
||||
@@ -786,32 +785,32 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
|
||||
_metadata["raw_request"] = str(curl_command)
|
||||
# split up, so it's easier to parse in the UI
|
||||
self.model_call_details["raw_request_typed_dict"] = (
|
||||
RawRequestTypedDict(
|
||||
raw_request_api_base=str(
|
||||
additional_args.get("api_base") or ""
|
||||
),
|
||||
raw_request_body=self._get_raw_request_body(
|
||||
additional_args.get("complete_input_dict", {})
|
||||
),
|
||||
raw_request_headers=self._get_masked_headers(
|
||||
additional_args.get("headers", {}) or {},
|
||||
ignore_sensitive_headers=True,
|
||||
),
|
||||
error=None,
|
||||
)
|
||||
self.model_call_details[
|
||||
"raw_request_typed_dict"
|
||||
] = RawRequestTypedDict(
|
||||
raw_request_api_base=str(
|
||||
additional_args.get("api_base") or ""
|
||||
),
|
||||
raw_request_body=self._get_raw_request_body(
|
||||
additional_args.get("complete_input_dict", {})
|
||||
),
|
||||
raw_request_headers=self._get_masked_headers(
|
||||
additional_args.get("headers", {}) or {},
|
||||
ignore_sensitive_headers=True,
|
||||
),
|
||||
error=None,
|
||||
)
|
||||
except Exception as e:
|
||||
self.model_call_details["raw_request_typed_dict"] = (
|
||||
RawRequestTypedDict(
|
||||
error=str(e),
|
||||
)
|
||||
self.model_call_details[
|
||||
"raw_request_typed_dict"
|
||||
] = RawRequestTypedDict(
|
||||
error=str(e),
|
||||
)
|
||||
_metadata["raw_request"] = (
|
||||
"Unable to Log \
|
||||
_metadata[
|
||||
"raw_request"
|
||||
] = "Unable to Log \
|
||||
raw request: {}".format(
|
||||
str(e)
|
||||
)
|
||||
str(e)
|
||||
)
|
||||
if getattr(self, "logger_fn", None) and callable(self.logger_fn):
|
||||
try:
|
||||
@@ -1112,13 +1111,13 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
for callback in callbacks:
|
||||
try:
|
||||
if isinstance(callback, CustomLogger):
|
||||
response: Optional[MCPPostCallResponseObject] = (
|
||||
await callback.async_post_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
response_obj=post_mcp_tool_call_response_obj,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
response: Optional[
|
||||
MCPPostCallResponseObject
|
||||
] = await callback.async_post_mcp_tool_call_hook(
|
||||
kwargs=kwargs,
|
||||
response_obj=post_mcp_tool_call_response_obj,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
)
|
||||
######################################################################
|
||||
# if any of the callbacks modify the response, use the modified response
|
||||
@@ -1238,9 +1237,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
verbose_logger.debug(
|
||||
f"response_cost_failure_debug_information: {debug_info}"
|
||||
)
|
||||
self.model_call_details["response_cost_failure_debug_information"] = (
|
||||
debug_info
|
||||
)
|
||||
self.model_call_details[
|
||||
"response_cost_failure_debug_information"
|
||||
] = debug_info
|
||||
return None
|
||||
|
||||
try:
|
||||
@@ -1265,9 +1264,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
verbose_logger.debug(
|
||||
f"response_cost_failure_debug_information: {debug_info}"
|
||||
)
|
||||
self.model_call_details["response_cost_failure_debug_information"] = (
|
||||
debug_info
|
||||
)
|
||||
self.model_call_details[
|
||||
"response_cost_failure_debug_information"
|
||||
] = debug_info
|
||||
|
||||
return None
|
||||
|
||||
@@ -1411,9 +1410,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
end_time = datetime.datetime.now()
|
||||
if self.completion_start_time is None:
|
||||
self.completion_start_time = end_time
|
||||
self.model_call_details["completion_start_time"] = (
|
||||
self.completion_start_time
|
||||
)
|
||||
self.model_call_details[
|
||||
"completion_start_time"
|
||||
] = self.completion_start_time
|
||||
self.model_call_details["log_event_type"] = "successful_api_call"
|
||||
self.model_call_details["end_time"] = end_time
|
||||
self.model_call_details["cache_hit"] = cache_hit
|
||||
@@ -1466,39 +1465,39 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
"response_cost"
|
||||
]
|
||||
else:
|
||||
self.model_call_details["response_cost"] = (
|
||||
self._response_cost_calculator(result=logging_result)
|
||||
)
|
||||
self.model_call_details[
|
||||
"response_cost"
|
||||
] = self._response_cost_calculator(result=logging_result)
|
||||
## STANDARDIZED LOGGING PAYLOAD
|
||||
|
||||
self.model_call_details["standard_logging_object"] = (
|
||||
get_standard_logging_object_payload(
|
||||
kwargs=self.model_call_details,
|
||||
init_response_obj=logging_result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=self,
|
||||
status="success",
|
||||
standard_built_in_tools_params=self.standard_built_in_tools_params,
|
||||
)
|
||||
self.model_call_details[
|
||||
"standard_logging_object"
|
||||
] = get_standard_logging_object_payload(
|
||||
kwargs=self.model_call_details,
|
||||
init_response_obj=logging_result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=self,
|
||||
status="success",
|
||||
standard_built_in_tools_params=self.standard_built_in_tools_params,
|
||||
)
|
||||
elif isinstance(result, dict) or isinstance(result, list):
|
||||
## STANDARDIZED LOGGING PAYLOAD
|
||||
self.model_call_details["standard_logging_object"] = (
|
||||
get_standard_logging_object_payload(
|
||||
kwargs=self.model_call_details,
|
||||
init_response_obj=result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=self,
|
||||
status="success",
|
||||
standard_built_in_tools_params=self.standard_built_in_tools_params,
|
||||
)
|
||||
self.model_call_details[
|
||||
"standard_logging_object"
|
||||
] = get_standard_logging_object_payload(
|
||||
kwargs=self.model_call_details,
|
||||
init_response_obj=result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=self,
|
||||
status="success",
|
||||
standard_built_in_tools_params=self.standard_built_in_tools_params,
|
||||
)
|
||||
elif standard_logging_object is not None:
|
||||
self.model_call_details["standard_logging_object"] = (
|
||||
standard_logging_object
|
||||
)
|
||||
self.model_call_details[
|
||||
"standard_logging_object"
|
||||
] = standard_logging_object
|
||||
else: # streaming chunks + image gen.
|
||||
self.model_call_details["response_cost"] = None
|
||||
|
||||
@@ -1597,7 +1596,6 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
)
|
||||
|
||||
if complete_streaming_response is not None:
|
||||
|
||||
self.success_handler(result=complete_streaming_response)
|
||||
return
|
||||
|
||||
@@ -1650,23 +1648,23 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
verbose_logger.debug(
|
||||
"Logging Details LiteLLM-Success Call streaming complete"
|
||||
)
|
||||
self.model_call_details["complete_streaming_response"] = (
|
||||
complete_streaming_response
|
||||
)
|
||||
self.model_call_details["response_cost"] = (
|
||||
self._response_cost_calculator(result=complete_streaming_response)
|
||||
)
|
||||
self.model_call_details[
|
||||
"complete_streaming_response"
|
||||
] = complete_streaming_response
|
||||
self.model_call_details[
|
||||
"response_cost"
|
||||
] = self._response_cost_calculator(result=complete_streaming_response)
|
||||
## STANDARDIZED LOGGING PAYLOAD
|
||||
self.model_call_details["standard_logging_object"] = (
|
||||
get_standard_logging_object_payload(
|
||||
kwargs=self.model_call_details,
|
||||
init_response_obj=complete_streaming_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=self,
|
||||
status="success",
|
||||
standard_built_in_tools_params=self.standard_built_in_tools_params,
|
||||
)
|
||||
self.model_call_details[
|
||||
"standard_logging_object"
|
||||
] = get_standard_logging_object_payload(
|
||||
kwargs=self.model_call_details,
|
||||
init_response_obj=complete_streaming_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=self,
|
||||
status="success",
|
||||
standard_built_in_tools_params=self.standard_built_in_tools_params,
|
||||
)
|
||||
callbacks = self.get_combined_callback_list(
|
||||
dynamic_success_callbacks=self.dynamic_success_callbacks,
|
||||
@@ -1994,10 +1992,10 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
)
|
||||
else:
|
||||
if self.stream and complete_streaming_response:
|
||||
self.model_call_details["complete_response"] = (
|
||||
self.model_call_details.get(
|
||||
"complete_streaming_response", {}
|
||||
)
|
||||
self.model_call_details[
|
||||
"complete_response"
|
||||
] = self.model_call_details.get(
|
||||
"complete_streaming_response", {}
|
||||
)
|
||||
result = self.model_call_details["complete_response"]
|
||||
openMeterLogger.log_success_event(
|
||||
@@ -2036,10 +2034,10 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
)
|
||||
else:
|
||||
if self.stream and complete_streaming_response:
|
||||
self.model_call_details["complete_response"] = (
|
||||
self.model_call_details.get(
|
||||
"complete_streaming_response", {}
|
||||
)
|
||||
self.model_call_details[
|
||||
"complete_response"
|
||||
] = self.model_call_details.get(
|
||||
"complete_streaming_response", {}
|
||||
)
|
||||
result = self.model_call_details["complete_response"]
|
||||
|
||||
@@ -2141,10 +2139,12 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
result.usage = batch_usage
|
||||
|
||||
elif not is_base64_unified_file_id: # only run for non-unified file ids
|
||||
response_cost, batch_usage, batch_models = (
|
||||
await _handle_completed_batch(
|
||||
batch=result, custom_llm_provider=self.custom_llm_provider
|
||||
)
|
||||
(
|
||||
response_cost,
|
||||
batch_usage,
|
||||
batch_models,
|
||||
) = await _handle_completed_batch(
|
||||
batch=result, custom_llm_provider=self.custom_llm_provider
|
||||
)
|
||||
|
||||
result._hidden_params["response_cost"] = response_cost
|
||||
@@ -2175,9 +2175,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
if complete_streaming_response is not None:
|
||||
print_verbose("Async success callbacks: Got a complete streaming response")
|
||||
|
||||
self.model_call_details["async_complete_streaming_response"] = (
|
||||
complete_streaming_response
|
||||
)
|
||||
self.model_call_details[
|
||||
"async_complete_streaming_response"
|
||||
] = complete_streaming_response
|
||||
|
||||
try:
|
||||
if self.model_call_details.get("cache_hit", False) is True:
|
||||
@@ -2188,10 +2188,10 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
model_call_details=self.model_call_details
|
||||
)
|
||||
# base_model defaults to None if not set on model_info
|
||||
self.model_call_details["response_cost"] = (
|
||||
self._response_cost_calculator(
|
||||
result=complete_streaming_response
|
||||
)
|
||||
self.model_call_details[
|
||||
"response_cost"
|
||||
] = self._response_cost_calculator(
|
||||
result=complete_streaming_response
|
||||
)
|
||||
|
||||
verbose_logger.debug(
|
||||
@@ -2204,16 +2204,16 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
self.model_call_details["response_cost"] = None
|
||||
|
||||
## STANDARDIZED LOGGING PAYLOAD
|
||||
self.model_call_details["standard_logging_object"] = (
|
||||
get_standard_logging_object_payload(
|
||||
kwargs=self.model_call_details,
|
||||
init_response_obj=complete_streaming_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=self,
|
||||
status="success",
|
||||
standard_built_in_tools_params=self.standard_built_in_tools_params,
|
||||
)
|
||||
self.model_call_details[
|
||||
"standard_logging_object"
|
||||
] = get_standard_logging_object_payload(
|
||||
kwargs=self.model_call_details,
|
||||
init_response_obj=complete_streaming_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=self,
|
||||
status="success",
|
||||
standard_built_in_tools_params=self.standard_built_in_tools_params,
|
||||
)
|
||||
callbacks = self.get_combined_callback_list(
|
||||
dynamic_success_callbacks=self.dynamic_async_success_callbacks,
|
||||
@@ -2426,18 +2426,18 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||
|
||||
## STANDARDIZED LOGGING PAYLOAD
|
||||
|
||||
self.model_call_details["standard_logging_object"] = (
|
||||
get_standard_logging_object_payload(
|
||||
kwargs=self.model_call_details,
|
||||
init_response_obj={},
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=self,
|
||||
status="failure",
|
||||
error_str=str(exception),
|
||||
original_exception=exception,
|
||||
standard_built_in_tools_params=self.standard_built_in_tools_params,
|
||||
)
|
||||
self.model_call_details[
|
||||
"standard_logging_object"
|
||||
] = get_standard_logging_object_payload(
|
||||
kwargs=self.model_call_details,
|
||||
init_response_obj={},
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=self,
|
||||
status="failure",
|
||||
error_str=str(exception),
|
||||
original_exception=exception,
|
||||
standard_built_in_tools_params=self.standard_built_in_tools_params,
|
||||
)
|
||||
return start_time, end_time
|
||||
|
||||
@@ -3326,9 +3326,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
||||
endpoint=arize_config.endpoint,
|
||||
)
|
||||
|
||||
os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
|
||||
f"space_id={arize_config.space_key},api_key={arize_config.api_key}"
|
||||
)
|
||||
os.environ[
|
||||
"OTEL_EXPORTER_OTLP_TRACES_HEADERS"
|
||||
] = f"space_id={arize_config.space_key},api_key={arize_config.api_key}"
|
||||
for callback in _in_memory_loggers:
|
||||
if (
|
||||
isinstance(callback, ArizeLogger)
|
||||
@@ -3352,9 +3352,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
||||
|
||||
# auth can be disabled on local deployments of arize phoenix
|
||||
if arize_phoenix_config.otlp_auth_headers is not None:
|
||||
os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
|
||||
arize_phoenix_config.otlp_auth_headers
|
||||
)
|
||||
os.environ[
|
||||
"OTEL_EXPORTER_OTLP_TRACES_HEADERS"
|
||||
] = arize_phoenix_config.otlp_auth_headers
|
||||
|
||||
for callback in _in_memory_loggers:
|
||||
if (
|
||||
@@ -3462,9 +3462,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
||||
exporter="otlp_http",
|
||||
endpoint="https://langtrace.ai/api/trace",
|
||||
)
|
||||
os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
|
||||
f"api_key={os.getenv('LANGTRACE_API_KEY')}"
|
||||
)
|
||||
os.environ[
|
||||
"OTEL_EXPORTER_OTLP_TRACES_HEADERS"
|
||||
] = f"api_key={os.getenv('LANGTRACE_API_KEY')}"
|
||||
for callback in _in_memory_loggers:
|
||||
if (
|
||||
isinstance(callback, OpenTelemetry)
|
||||
@@ -4114,10 +4114,10 @@ class StandardLoggingPayloadSetup:
|
||||
for key in StandardLoggingHiddenParams.__annotations__.keys():
|
||||
if key in hidden_params:
|
||||
if key == "additional_headers":
|
||||
clean_hidden_params["additional_headers"] = (
|
||||
StandardLoggingPayloadSetup.get_additional_headers(
|
||||
hidden_params[key]
|
||||
)
|
||||
clean_hidden_params[
|
||||
"additional_headers"
|
||||
] = StandardLoggingPayloadSetup.get_additional_headers(
|
||||
hidden_params[key]
|
||||
)
|
||||
else:
|
||||
clean_hidden_params[key] = hidden_params[key] # type: ignore
|
||||
@@ -4150,15 +4150,28 @@ class StandardLoggingPayloadSetup:
|
||||
from litellm.integrations.s3 import get_s3_object_key
|
||||
|
||||
# Only generate object key if cold storage is configured
|
||||
if litellm.configured_cold_storage_logger is None:
|
||||
configured_cold_storage_logger = litellm.configured_cold_storage_logger
|
||||
if configured_cold_storage_logger is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
# Generate file name in same format as litellm.utils.get_logging_id
|
||||
s3_file_name = f"time-{start_time.strftime('%H-%M-%S-%f')}_{response_id}"
|
||||
|
||||
# Get the actual s3_path from the configured cold storage logger instance
|
||||
s3_path = "" # default value
|
||||
|
||||
# Try to get the actual logger instance from the logger name
|
||||
try:
|
||||
custom_logger = litellm.logging_callback_manager.get_active_custom_logger_for_callback_name(configured_cold_storage_logger)
|
||||
if custom_logger and hasattr(custom_logger, 's3_path') and custom_logger.s3_path:
|
||||
s3_path = custom_logger.s3_path
|
||||
except Exception:
|
||||
# If any error occurs in getting the logger instance, use default empty s3_path
|
||||
pass
|
||||
|
||||
s3_object_key = get_s3_object_key(
|
||||
s3_path="", # Use empty path as default
|
||||
s3_path=s3_path, # Use actual s3_path from logger configuration
|
||||
team_alias_prefix="", # Don't split by team alias for cold storage
|
||||
start_time=start_time,
|
||||
s3_file_name=s3_file_name,
|
||||
@@ -4553,6 +4566,9 @@ def get_standard_logging_metadata(
|
||||
clean_metadata = StandardLoggingMetadata(
|
||||
user_api_key_hash=None,
|
||||
user_api_key_alias=None,
|
||||
user_api_key_spend=None,
|
||||
user_api_key_max_budget=None,
|
||||
user_api_key_budget_reset_at=None,
|
||||
user_api_key_team_id=None,
|
||||
user_api_key_org_id=None,
|
||||
user_api_key_user_id=None,
|
||||
@@ -4602,9 +4618,9 @@ def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]):
|
||||
):
|
||||
for k, v in metadata["user_api_key_metadata"].items():
|
||||
if k == "logging": # prevent logging user logging keys
|
||||
cleaned_user_api_key_metadata[k] = (
|
||||
"scrubbed_by_litellm_for_sensitive_keys"
|
||||
)
|
||||
cleaned_user_api_key_metadata[
|
||||
k
|
||||
] = "scrubbed_by_litellm_for_sensitive_keys"
|
||||
else:
|
||||
cleaned_user_api_key_metadata[k] = v
|
||||
|
||||
|
||||
@@ -16,8 +16,8 @@ from litellm import verbose_logger
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client
|
||||
from litellm.types.files import get_file_extension_from_mime_type
|
||||
from litellm.types.llms.anthropic import *
|
||||
from litellm.types.llms.bedrock import MessageBlock as BedrockMessageBlock
|
||||
from litellm.types.llms.bedrock import CachePointBlock
|
||||
from litellm.types.llms.bedrock import MessageBlock as BedrockMessageBlock
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
from litellm.types.llms.ollama import OllamaVisionModelObject
|
||||
from litellm.types.llms.openai import (
|
||||
@@ -1067,10 +1067,10 @@ def convert_to_gemini_tool_call_invoke(
|
||||
if tool_calls is not None:
|
||||
for tool in tool_calls:
|
||||
if "function" in tool:
|
||||
gemini_function_call: Optional[VertexFunctionCall] = (
|
||||
_gemini_tool_call_invoke_helper(
|
||||
function_call_params=tool["function"]
|
||||
)
|
||||
gemini_function_call: Optional[
|
||||
VertexFunctionCall
|
||||
] = _gemini_tool_call_invoke_helper(
|
||||
function_call_params=tool["function"]
|
||||
)
|
||||
if gemini_function_call is not None:
|
||||
_parts_list.append(
|
||||
@@ -1589,9 +1589,9 @@ def anthropic_messages_pt( # noqa: PLR0915
|
||||
)
|
||||
|
||||
if "cache_control" in _content_element:
|
||||
_anthropic_content_element["cache_control"] = (
|
||||
_content_element["cache_control"]
|
||||
)
|
||||
_anthropic_content_element[
|
||||
"cache_control"
|
||||
] = _content_element["cache_control"]
|
||||
user_content.append(_anthropic_content_element)
|
||||
elif m.get("type", "") == "text":
|
||||
m = cast(ChatCompletionTextObject, m)
|
||||
@@ -1629,9 +1629,9 @@ def anthropic_messages_pt( # noqa: PLR0915
|
||||
)
|
||||
|
||||
if "cache_control" in _content_element:
|
||||
_anthropic_content_text_element["cache_control"] = (
|
||||
_content_element["cache_control"]
|
||||
)
|
||||
_anthropic_content_text_element[
|
||||
"cache_control"
|
||||
] = _content_element["cache_control"]
|
||||
|
||||
user_content.append(_anthropic_content_text_element)
|
||||
|
||||
@@ -2482,8 +2482,7 @@ class BedrockImageProcessor:
|
||||
|
||||
if is_document:
|
||||
return BedrockImageProcessor._get_document_format(
|
||||
mime_type=mime_type,
|
||||
supported_doc_formats=supported_doc_formats
|
||||
mime_type=mime_type, supported_doc_formats=supported_doc_formats
|
||||
)
|
||||
|
||||
else:
|
||||
@@ -2495,12 +2494,9 @@ class BedrockImageProcessor:
|
||||
f"Unsupported image format: {image_format}. Supported formats: {supported_image_and_video_formats}"
|
||||
)
|
||||
return image_format
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _get_document_format(
|
||||
mime_type: str,
|
||||
supported_doc_formats: List[str]
|
||||
) -> str:
|
||||
def _get_document_format(mime_type: str, supported_doc_formats: List[str]) -> str:
|
||||
"""
|
||||
Get the document format from the mime type
|
||||
|
||||
@@ -2519,13 +2515,9 @@ class BedrockImageProcessor:
|
||||
The document format
|
||||
"""
|
||||
valid_extensions: Optional[List[str]] = None
|
||||
potential_extensions = mimetypes.guess_all_extensions(
|
||||
mime_type, strict=False
|
||||
)
|
||||
potential_extensions = mimetypes.guess_all_extensions(mime_type, strict=False)
|
||||
valid_extensions = [
|
||||
ext[1:]
|
||||
for ext in potential_extensions
|
||||
if ext[1:] in supported_doc_formats
|
||||
ext[1:] for ext in potential_extensions if ext[1:] in supported_doc_formats
|
||||
]
|
||||
|
||||
# Fallback to types/files.py if mimetypes doesn't return valid extensions
|
||||
@@ -2689,10 +2681,12 @@ def _convert_to_bedrock_tool_call_invoke(
|
||||
)
|
||||
bedrock_content_block = BedrockContentBlock(toolUse=bedrock_tool)
|
||||
_parts_list.append(bedrock_content_block)
|
||||
|
||||
|
||||
# Check for cache_control and add a separate cachePoint block
|
||||
if tool.get("cache_control", None) is not None:
|
||||
cache_point_block = BedrockContentBlock(cachePoint=CachePointBlock(type="default"))
|
||||
cache_point_block = BedrockContentBlock(
|
||||
cachePoint=CachePointBlock(type="default")
|
||||
)
|
||||
_parts_list.append(cache_point_block)
|
||||
return _parts_list
|
||||
except Exception as e:
|
||||
@@ -2754,7 +2748,7 @@ def _convert_to_bedrock_tool_call_result(
|
||||
for content in content_list:
|
||||
if content["type"] == "text":
|
||||
content_str += content["text"]
|
||||
|
||||
|
||||
message.get("name", "")
|
||||
id = str(message.get("tool_call_id", str(uuid.uuid4())))
|
||||
|
||||
@@ -2763,7 +2757,7 @@ def _convert_to_bedrock_tool_call_result(
|
||||
content=[tool_result_content_block],
|
||||
toolUseId=id,
|
||||
)
|
||||
|
||||
|
||||
content_block = BedrockContentBlock(toolResult=tool_result)
|
||||
|
||||
return content_block
|
||||
@@ -3085,6 +3079,7 @@ class BedrockConverseMessagesProcessor:
|
||||
messages.append(DEFAULT_USER_CONTINUE_MESSAGE)
|
||||
return messages
|
||||
|
||||
|
||||
@staticmethod
|
||||
async def _bedrock_converse_messages_pt_async( # noqa: PLR0915
|
||||
messages: List,
|
||||
@@ -3128,6 +3123,12 @@ class BedrockConverseMessagesProcessor:
|
||||
if element["type"] == "text":
|
||||
_part = BedrockContentBlock(text=element["text"])
|
||||
_parts.append(_part)
|
||||
elif element["type"] == "guarded_text":
|
||||
# Wrap guarded_text in guardrailConverseContent block
|
||||
_part = BedrockContentBlock(
|
||||
guardrailConverseContent={"text": element["text"]}
|
||||
)
|
||||
_parts.append(_part)
|
||||
elif element["type"] == "image_url":
|
||||
format: Optional[str] = None
|
||||
if isinstance(element["image_url"], dict):
|
||||
@@ -3170,6 +3171,7 @@ class BedrockConverseMessagesProcessor:
|
||||
|
||||
msg_i += 1
|
||||
if user_content:
|
||||
|
||||
if len(contents) > 0 and contents[-1]["role"] == "user":
|
||||
if (
|
||||
assistant_continue_message is not None
|
||||
@@ -3199,26 +3201,29 @@ class BedrockConverseMessagesProcessor:
|
||||
current_message = messages[msg_i]
|
||||
tool_call_result = _convert_to_bedrock_tool_call_result(current_message)
|
||||
tool_content.append(tool_call_result)
|
||||
|
||||
|
||||
# Check if we need to add a separate cachePoint block
|
||||
has_cache_control = False
|
||||
|
||||
|
||||
# Check for message-level cache_control
|
||||
if current_message.get("cache_control", None) is not None:
|
||||
has_cache_control = True
|
||||
# Check for content-level cache_control in list content
|
||||
elif isinstance(current_message.get("content"), list):
|
||||
for content_element in current_message["content"]:
|
||||
if (isinstance(content_element, dict) and
|
||||
content_element.get("cache_control", None) is not None):
|
||||
if (
|
||||
isinstance(content_element, dict)
|
||||
and content_element.get("cache_control", None) is not None
|
||||
):
|
||||
has_cache_control = True
|
||||
break
|
||||
|
||||
|
||||
# Add a separate cachePoint block if cache_control is present
|
||||
if has_cache_control:
|
||||
cache_point_block = BedrockContentBlock(cachePoint=CachePointBlock(type="default"))
|
||||
cache_point_block = BedrockContentBlock(
|
||||
cachePoint=CachePointBlock(type="default")
|
||||
)
|
||||
tool_content.append(cache_point_block)
|
||||
|
||||
|
||||
msg_i += 1
|
||||
if tool_content:
|
||||
@@ -3299,7 +3304,7 @@ class BedrockConverseMessagesProcessor:
|
||||
image_url=image_url
|
||||
)
|
||||
assistants_parts.append(assistants_part)
|
||||
# Add cache point block for assistant content elements
|
||||
# Add cache point block for assistant content elements
|
||||
_cache_point_block = (
|
||||
litellm.AmazonConverseConfig()._get_cache_point_block(
|
||||
message_block=cast(
|
||||
@@ -3311,8 +3316,12 @@ class BedrockConverseMessagesProcessor:
|
||||
if _cache_point_block is not None:
|
||||
assistants_parts.append(_cache_point_block)
|
||||
assistant_content.extend(assistants_parts)
|
||||
elif _assistant_content is not None and isinstance(_assistant_content, str):
|
||||
assistant_content.append(BedrockContentBlock(text=_assistant_content))
|
||||
elif _assistant_content is not None and isinstance(
|
||||
_assistant_content, str
|
||||
):
|
||||
assistant_content.append(
|
||||
BedrockContentBlock(text=_assistant_content)
|
||||
)
|
||||
# Add cache point block for assistant string content
|
||||
_cache_point_block = (
|
||||
litellm.AmazonConverseConfig()._get_cache_point_block(
|
||||
@@ -3496,6 +3505,12 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915
|
||||
if element["type"] == "text":
|
||||
_part = BedrockContentBlock(text=element["text"])
|
||||
_parts.append(_part)
|
||||
elif element["type"] == "guarded_text":
|
||||
# Wrap guarded_text in guardrailConverseContent block
|
||||
_part = BedrockContentBlock(
|
||||
guardrailConverseContent={"text": element["text"]}
|
||||
)
|
||||
_parts.append(_part)
|
||||
elif element["type"] == "image_url":
|
||||
format: Optional[str] = None
|
||||
if isinstance(element["image_url"], dict):
|
||||
@@ -3539,6 +3554,7 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915
|
||||
|
||||
msg_i += 1
|
||||
if user_content:
|
||||
|
||||
if len(contents) > 0 and contents[-1]["role"] == "user":
|
||||
if (
|
||||
assistant_continue_message is not None
|
||||
@@ -3565,29 +3581,33 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915
|
||||
while msg_i < len(messages) and messages[msg_i]["role"] == "tool":
|
||||
tool_call_result = _convert_to_bedrock_tool_call_result(messages[msg_i])
|
||||
current_message = messages[msg_i]
|
||||
|
||||
|
||||
# Add the tool result first
|
||||
tool_content.append(tool_call_result)
|
||||
|
||||
|
||||
# Check if we need to add a separate cachePoint block
|
||||
has_cache_control = False
|
||||
|
||||
|
||||
# Check for message-level cache_control
|
||||
if current_message.get("cache_control", None) is not None:
|
||||
has_cache_control = True
|
||||
# Check for content-level cache_control in list content
|
||||
elif isinstance(current_message.get("content"), list):
|
||||
for content_element in current_message["content"]:
|
||||
if (isinstance(content_element, dict) and
|
||||
content_element.get("cache_control", None) is not None):
|
||||
if (
|
||||
isinstance(content_element, dict)
|
||||
and content_element.get("cache_control", None) is not None
|
||||
):
|
||||
has_cache_control = True
|
||||
break
|
||||
|
||||
|
||||
# Add a separate cachePoint block if cache_control is present
|
||||
if has_cache_control:
|
||||
cache_point_block = BedrockContentBlock(cachePoint=CachePointBlock(type="default"))
|
||||
cache_point_block = BedrockContentBlock(
|
||||
cachePoint=CachePointBlock(type="default")
|
||||
)
|
||||
tool_content.append(cache_point_block)
|
||||
|
||||
|
||||
msg_i += 1
|
||||
if tool_content:
|
||||
# if last message was a 'user' message, then add a blank assistant message (bedrock requires alternating roles)
|
||||
@@ -3852,10 +3872,9 @@ def function_call_prompt(messages: list, functions: list):
|
||||
if isinstance(message["content"], str):
|
||||
message["content"] += f""" {function_prompt}"""
|
||||
else:
|
||||
message["content"].append({
|
||||
"type": "text",
|
||||
"text": f""" {function_prompt}"""
|
||||
})
|
||||
message["content"].append(
|
||||
{"type": "text", "text": f""" {function_prompt}"""}
|
||||
)
|
||||
function_added_to_prompt = True
|
||||
|
||||
if function_added_to_prompt is False:
|
||||
|
||||
@@ -200,8 +200,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
|
||||
)
|
||||
|
||||
_allowed_properties = set(AnthropicInputSchema.__annotations__.keys())
|
||||
input_schema_filtered = {k: v for k, v in _input_schema.items() if k in _allowed_properties}
|
||||
input_anthropic_schema: AnthropicInputSchema = AnthropicInputSchema(**input_schema_filtered)
|
||||
input_schema_filtered = {
|
||||
k: v for k, v in _input_schema.items() if k in _allowed_properties
|
||||
}
|
||||
input_anthropic_schema: AnthropicInputSchema = AnthropicInputSchema(
|
||||
**input_schema_filtered
|
||||
)
|
||||
|
||||
_tool = AnthropicMessagesTool(
|
||||
name=tool["function"]["name"],
|
||||
|
||||
@@ -158,6 +158,48 @@ class BaseBatchesConfig(ABC):
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def transform_retrieve_batch_request(
|
||||
self,
|
||||
batch_id: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> Union[bytes, str, Dict[str, Any]]:
|
||||
"""
|
||||
Transform the batch retrieval request to provider-specific format.
|
||||
|
||||
Args:
|
||||
batch_id: Batch ID to retrieve
|
||||
optional_params: Optional parameters
|
||||
litellm_params: LiteLLM parameters
|
||||
|
||||
Returns:
|
||||
Transformed request data
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def transform_retrieve_batch_response(
|
||||
self,
|
||||
model: Optional[str],
|
||||
raw_response: httpx.Response,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
) -> LiteLLMBatch:
|
||||
"""
|
||||
Transform provider-specific batch retrieval response to LiteLLM format.
|
||||
|
||||
Args:
|
||||
model: Model name
|
||||
raw_response: Raw HTTP response
|
||||
logging_obj: Logging object
|
||||
litellm_params: LiteLLM parameters
|
||||
|
||||
Returns:
|
||||
LiteLLM batch object
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_error_class(
|
||||
self, error_message: str, status_code: int, headers: Union[Dict, Headers]
|
||||
|
||||
@@ -7,7 +7,6 @@ from httpx import Headers, Response
|
||||
from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.types.llms.bedrock import (
|
||||
BedrockBatchJobStatus,
|
||||
BedrockCreateBatchRequest,
|
||||
BedrockCreateBatchResponse,
|
||||
BedrockInputDataConfig,
|
||||
@@ -200,19 +199,23 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
||||
|
||||
# Extract information from typed Bedrock response
|
||||
job_arn = response_data.get("jobArn", "")
|
||||
status: BedrockBatchJobStatus = response_data.get("status", "Submitted")
|
||||
status_str: str = str(response_data.get("status", "Submitted"))
|
||||
|
||||
# Map Bedrock status to OpenAI-compatible status
|
||||
status_mapping: Dict[BedrockBatchJobStatus, str] = {
|
||||
status_mapping: Dict[str, str] = {
|
||||
"Submitted": "validating",
|
||||
"Validating": "validating",
|
||||
"Scheduled": "in_progress",
|
||||
"InProgress": "in_progress",
|
||||
"PartiallyCompleted": "completed",
|
||||
"Completed": "completed",
|
||||
"Failed": "failed",
|
||||
"Stopping": "cancelling",
|
||||
"Stopped": "cancelled"
|
||||
"Stopped": "cancelled",
|
||||
"Expired": "expired",
|
||||
}
|
||||
|
||||
openai_status = cast(Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"], status_mapping.get(status, "validating"))
|
||||
openai_status = cast(Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"], status_mapping.get(status_str, "validating"))
|
||||
|
||||
# Get original request data from litellm_params if available
|
||||
original_request = litellm_params.get("original_batch_request", {})
|
||||
@@ -229,7 +232,7 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
||||
output_file_id=None, # Will be populated when job completes
|
||||
error_file_id=None,
|
||||
created_at=int(time.time()),
|
||||
in_progress_at=int(time.time()) if status == "InProgress" else None,
|
||||
in_progress_at=int(time.time()) if status_str == "InProgress" else None,
|
||||
expires_at=None,
|
||||
finalizing_at=None,
|
||||
completed_at=None,
|
||||
@@ -241,6 +244,203 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig):
|
||||
metadata=original_request.get("metadata", {}),
|
||||
)
|
||||
|
||||
def transform_retrieve_batch_request(
|
||||
self,
|
||||
batch_id: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Transform batch retrieval request for Bedrock.
|
||||
|
||||
Args:
|
||||
batch_id: Bedrock job ARN
|
||||
optional_params: Optional parameters
|
||||
litellm_params: LiteLLM parameters
|
||||
|
||||
Returns:
|
||||
Transformed request data for Bedrock GetModelInvocationJob API
|
||||
"""
|
||||
# For Bedrock, batch_id should be the full job ARN
|
||||
# The GetModelInvocationJob API expects the full ARN as the identifier
|
||||
if not batch_id.startswith("arn:aws:bedrock:"):
|
||||
raise ValueError(f"Invalid batch_id format. Expected ARN, got: {batch_id}")
|
||||
|
||||
# Extract the job identifier from the ARN - use the full ARN path part
|
||||
# ARN format: arn:aws:bedrock:region:account:model-invocation-job/job-name
|
||||
arn_parts = batch_id.split(":")
|
||||
if len(arn_parts) < 6:
|
||||
raise ValueError(f"Invalid ARN format: {batch_id}")
|
||||
|
||||
region = arn_parts[3]
|
||||
# arn_parts[5] contains "model-invocation-job/{jobId}"
|
||||
|
||||
# Build the endpoint URL for GetModelInvocationJob
|
||||
# AWS API format: GET /model-invocation-job/{jobIdentifier}
|
||||
# Use the FULL ARN as jobIdentifier and URL-encode it (includes ':' and '/')
|
||||
import urllib.parse as _ul
|
||||
encoded_arn = _ul.quote(batch_id, safe="")
|
||||
endpoint_url = f"https://bedrock.{region}.amazonaws.com/model-invocation-job/{encoded_arn}"
|
||||
|
||||
# Use common utility for AWS signing
|
||||
signed_headers, _ = self.common_utils.sign_aws_request(
|
||||
service_name="bedrock",
|
||||
data={}, # GET request has no body
|
||||
endpoint_url=endpoint_url,
|
||||
optional_params=optional_params,
|
||||
method="GET"
|
||||
)
|
||||
|
||||
# Return pre-signed request format
|
||||
return {
|
||||
"method": "GET",
|
||||
"url": endpoint_url,
|
||||
"headers": signed_headers,
|
||||
"data": None
|
||||
}
|
||||
|
||||
def _parse_timestamps_and_status(self, response_data, status_str: str):
|
||||
"""Helper to parse timestamps based on status."""
|
||||
import datetime
|
||||
def parse_timestamp(ts_str: Optional[str]) -> Optional[int]:
|
||||
if not ts_str:
|
||||
return None
|
||||
try:
|
||||
dt = datetime.datetime.fromisoformat(ts_str.replace('Z', '+00:00'))
|
||||
return int(dt.timestamp())
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
created_at = parse_timestamp(str(response_data.get("submitTime")) if response_data.get("submitTime") is not None else None)
|
||||
in_progress_states = {"InProgress", "Validating", "Scheduled"}
|
||||
in_progress_at = (
|
||||
parse_timestamp(str(response_data.get("lastModifiedTime")) if response_data.get("lastModifiedTime") is not None else None)
|
||||
if status_str in in_progress_states
|
||||
else None
|
||||
)
|
||||
completed_at = parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str in {"Completed", "PartiallyCompleted"} else None
|
||||
failed_at = parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str == "Failed" else None
|
||||
cancelled_at = parse_timestamp(str(response_data.get("endTime")) if response_data.get("endTime") is not None else None) if status_str == "Stopped" else None
|
||||
expires_at = parse_timestamp(str(response_data.get("jobExpirationTime")) if response_data.get("jobExpirationTime") is not None else None)
|
||||
|
||||
return created_at, in_progress_at, completed_at, failed_at, cancelled_at, expires_at
|
||||
|
||||
def _extract_file_configs(self, response_data):
|
||||
"""Helper to extract input and output file configurations."""
|
||||
# Extract input file ID
|
||||
input_file_id = ""
|
||||
input_data_config = response_data.get("inputDataConfig", {})
|
||||
if isinstance(input_data_config, dict):
|
||||
s3_input_config = input_data_config.get("s3InputDataConfig", {})
|
||||
if isinstance(s3_input_config, dict):
|
||||
input_file_id = s3_input_config.get("s3Uri", "")
|
||||
|
||||
# Extract output file ID
|
||||
output_file_id = None
|
||||
output_data_config = response_data.get("outputDataConfig", {})
|
||||
if isinstance(output_data_config, dict):
|
||||
s3_output_config = output_data_config.get("s3OutputDataConfig", {})
|
||||
if isinstance(s3_output_config, dict):
|
||||
output_file_id = s3_output_config.get("s3Uri", "")
|
||||
|
||||
return input_file_id, output_file_id
|
||||
|
||||
def _extract_errors_and_metadata(self, response_data, raw_response):
|
||||
"""Helper to extract errors and enriched metadata."""
|
||||
# Extract errors
|
||||
message = response_data.get("message")
|
||||
errors = None
|
||||
if message:
|
||||
from openai.types.batch import Errors
|
||||
from openai.types.batch_error import BatchError
|
||||
errors = Errors(
|
||||
data=[BatchError(message=message, code=str(raw_response.status_code))],
|
||||
object="list"
|
||||
)
|
||||
|
||||
# Enrich metadata with useful Bedrock fields
|
||||
enriched_metadata_raw: Dict[str, Any] = {
|
||||
"jobName": response_data.get("jobName"),
|
||||
"clientRequestToken": response_data.get("clientRequestToken"),
|
||||
"modelId": response_data.get("modelId"),
|
||||
"roleArn": response_data.get("roleArn"),
|
||||
"timeoutDurationInHours": response_data.get("timeoutDurationInHours"),
|
||||
"vpcConfig": response_data.get("vpcConfig"),
|
||||
}
|
||||
import json as _json
|
||||
enriched_metadata: Dict[str, str] = {}
|
||||
for _k, _v in enriched_metadata_raw.items():
|
||||
if _v is None:
|
||||
continue
|
||||
if isinstance(_v, (dict, list)):
|
||||
try:
|
||||
enriched_metadata[_k] = _json.dumps(_v)
|
||||
except Exception:
|
||||
enriched_metadata[_k] = str(_v)
|
||||
else:
|
||||
enriched_metadata[_k] = str(_v)
|
||||
|
||||
return errors, enriched_metadata
|
||||
|
||||
def transform_retrieve_batch_response(
|
||||
self,
|
||||
model: Optional[str],
|
||||
raw_response: Response,
|
||||
logging_obj: Any,
|
||||
litellm_params: dict,
|
||||
) -> LiteLLMBatch:
|
||||
"""
|
||||
Transform Bedrock batch retrieval response to LiteLLM format.
|
||||
"""
|
||||
from litellm.types.llms.bedrock import BedrockGetBatchResponse
|
||||
try:
|
||||
response_data: BedrockGetBatchResponse = raw_response.json()
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to parse Bedrock batch response: {e}")
|
||||
|
||||
job_arn = response_data.get("jobArn", "")
|
||||
status_str: str = str(response_data.get("status", "Submitted"))
|
||||
|
||||
# Map Bedrock status to OpenAI-compatible status
|
||||
status_mapping: Dict[str, str] = {
|
||||
"Submitted": "validating", "Validating": "validating", "Scheduled": "in_progress",
|
||||
"InProgress": "in_progress", "PartiallyCompleted": "completed", "Completed": "completed",
|
||||
"Failed": "failed", "Stopping": "cancelling", "Stopped": "cancelled", "Expired": "expired"
|
||||
}
|
||||
openai_status = cast(Literal["validating", "failed", "in_progress", "finalizing", "completed", "expired", "cancelling", "cancelled"], status_mapping.get(status_str, "validating"))
|
||||
|
||||
# Parse timestamps
|
||||
created_at, in_progress_at, completed_at, failed_at, cancelled_at, expires_at = self._parse_timestamps_and_status(response_data, status_str)
|
||||
|
||||
# Extract file configurations
|
||||
input_file_id, output_file_id = self._extract_file_configs(response_data)
|
||||
|
||||
# Extract errors and metadata
|
||||
errors, enriched_metadata = self._extract_errors_and_metadata(response_data, raw_response)
|
||||
|
||||
return LiteLLMBatch(
|
||||
id=job_arn,
|
||||
object="batch",
|
||||
endpoint="/v1/chat/completions",
|
||||
errors=errors,
|
||||
input_file_id=input_file_id,
|
||||
completion_window="24h",
|
||||
status=openai_status,
|
||||
output_file_id=output_file_id,
|
||||
error_file_id=None,
|
||||
created_at=created_at or int(time.time()),
|
||||
in_progress_at=in_progress_at,
|
||||
expires_at=expires_at,
|
||||
finalizing_at=None,
|
||||
completed_at=completed_at,
|
||||
failed_at=failed_at,
|
||||
expired_at=None,
|
||||
cancelling_at=None,
|
||||
cancelled_at=cancelled_at,
|
||||
request_counts=None,
|
||||
metadata=enriched_metadata,
|
||||
)
|
||||
|
||||
def get_error_class(
|
||||
self, error_message: str, status_code: int, headers: Union[Dict, Headers]
|
||||
) -> BaseLLMException:
|
||||
|
||||
@@ -501,7 +501,6 @@ class AmazonConverseConfig(BaseConfig):
|
||||
)
|
||||
and not is_thinking_enabled
|
||||
):
|
||||
|
||||
optional_params["tool_choice"] = ToolChoiceValuesBlock(
|
||||
tool=SpecificToolChoiceBlock(name=RESPONSE_FORMAT_TOOL_NAME)
|
||||
)
|
||||
@@ -995,7 +994,9 @@ class AmazonConverseConfig(BaseConfig):
|
||||
|
||||
return message, returned_finish_reason
|
||||
|
||||
def _translate_message_content(self, content_blocks: List[ContentBlock]) -> Tuple[
|
||||
def _translate_message_content(
|
||||
self, content_blocks: List[ContentBlock]
|
||||
) -> Tuple[
|
||||
str,
|
||||
List[ChatCompletionToolCallChunk],
|
||||
Optional[List[BedrockConverseReasoningContentBlock]],
|
||||
@@ -1010,9 +1011,9 @@ class AmazonConverseConfig(BaseConfig):
|
||||
"""
|
||||
content_str = ""
|
||||
tools: List[ChatCompletionToolCallChunk] = []
|
||||
reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = (
|
||||
None
|
||||
)
|
||||
reasoningContentBlocks: Optional[
|
||||
List[BedrockConverseReasoningContentBlock]
|
||||
] = None
|
||||
for idx, content in enumerate(content_blocks):
|
||||
"""
|
||||
- Content is either a tool response or text
|
||||
@@ -1133,9 +1134,9 @@ class AmazonConverseConfig(BaseConfig):
|
||||
chat_completion_message: ChatCompletionResponseMessage = {"role": "assistant"}
|
||||
content_str = ""
|
||||
tools: List[ChatCompletionToolCallChunk] = []
|
||||
reasoningContentBlocks: Optional[List[BedrockConverseReasoningContentBlock]] = (
|
||||
None
|
||||
)
|
||||
reasoningContentBlocks: Optional[
|
||||
List[BedrockConverseReasoningContentBlock]
|
||||
] = None
|
||||
|
||||
if message is not None:
|
||||
(
|
||||
@@ -1148,12 +1149,12 @@ class AmazonConverseConfig(BaseConfig):
|
||||
chat_completion_message["provider_specific_fields"] = {
|
||||
"reasoningContentBlocks": reasoningContentBlocks,
|
||||
}
|
||||
chat_completion_message["reasoning_content"] = (
|
||||
self._transform_reasoning_content(reasoningContentBlocks)
|
||||
)
|
||||
chat_completion_message["thinking_blocks"] = (
|
||||
self._transform_thinking_blocks(reasoningContentBlocks)
|
||||
)
|
||||
chat_completion_message[
|
||||
"reasoning_content"
|
||||
] = self._transform_reasoning_content(reasoningContentBlocks)
|
||||
chat_completion_message[
|
||||
"thinking_blocks"
|
||||
] = self._transform_thinking_blocks(reasoningContentBlocks)
|
||||
chat_completion_message["content"] = content_str
|
||||
if (
|
||||
json_mode is True
|
||||
@@ -1171,7 +1172,6 @@ class AmazonConverseConfig(BaseConfig):
|
||||
# Bedrock returns the response wrapped in a "properties" object
|
||||
# We need to extract the actual content from this wrapper
|
||||
try:
|
||||
|
||||
response_data = json.loads(json_mode_content_str)
|
||||
|
||||
# If Bedrock wrapped the response in "properties", extract the content
|
||||
|
||||
@@ -738,19 +738,24 @@ class CommonBatchFilesUtils:
|
||||
)
|
||||
|
||||
# Prepare the request data
|
||||
if isinstance(data, dict):
|
||||
import json
|
||||
request_data = json.dumps(data)
|
||||
method_upper = method.upper()
|
||||
if method_upper == "GET":
|
||||
# GET requests should be signed with an empty payload
|
||||
request_data = ""
|
||||
headers = {}
|
||||
else:
|
||||
request_data = data
|
||||
|
||||
# Prepare headers
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if isinstance(data, dict):
|
||||
import json
|
||||
request_data = json.dumps(data)
|
||||
else:
|
||||
request_data = data
|
||||
# Prepare headers for non-GET requests
|
||||
headers = {"Content-Type": "application/json"}
|
||||
|
||||
# Create AWS request and sign it
|
||||
sigv4 = SigV4Auth(credentials, service_name, aws_region_name)
|
||||
request = AWSRequest(
|
||||
method=method.upper(), url=endpoint_url, data=request_data, headers=headers
|
||||
method=method_upper, url=endpoint_url, data=request_data, headers=headers
|
||||
)
|
||||
sigv4.add_auth(request)
|
||||
prepped = request.prepare()
|
||||
|
||||
@@ -2520,6 +2520,95 @@ class BaseLLMHTTPHandler:
|
||||
litellm_params=litellm_params_with_request,
|
||||
)
|
||||
|
||||
def retrieve_batch(
|
||||
self,
|
||||
batch_id: str,
|
||||
litellm_params: dict,
|
||||
provider_config: "BaseBatchesConfig",
|
||||
headers: dict,
|
||||
api_base: Optional[str],
|
||||
api_key: Optional[str],
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
_is_async: bool = False,
|
||||
client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
model: Optional[str] = None,
|
||||
) -> Union["LiteLLMBatch", Coroutine[Any, Any, "LiteLLMBatch"]]:
|
||||
"""
|
||||
Retrieve a batch using provider-specific configuration.
|
||||
"""
|
||||
# Transform the request using provider config
|
||||
transformed_request = provider_config.transform_retrieve_batch_request(
|
||||
batch_id=batch_id,
|
||||
optional_params=litellm_params,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
if _is_async:
|
||||
return self.async_retrieve_batch(
|
||||
transformed_request=transformed_request,
|
||||
litellm_params=litellm_params,
|
||||
provider_config=provider_config,
|
||||
headers=headers,
|
||||
api_base=api_base,
|
||||
logging_obj=logging_obj,
|
||||
client=client,
|
||||
timeout=timeout,
|
||||
batch_id=batch_id,
|
||||
model=model,
|
||||
)
|
||||
|
||||
if client is None or not isinstance(client, HTTPHandler):
|
||||
sync_httpx_client = _get_httpx_client()
|
||||
else:
|
||||
sync_httpx_client = client
|
||||
|
||||
try:
|
||||
if (
|
||||
isinstance(transformed_request, dict)
|
||||
and "method" in transformed_request
|
||||
):
|
||||
# Handle pre-signed requests (e.g., from Bedrock with AWS auth)
|
||||
method = transformed_request["method"].lower()
|
||||
request_kwargs = {
|
||||
"url": transformed_request["url"],
|
||||
"headers": transformed_request["headers"],
|
||||
}
|
||||
|
||||
# Only add data for non-GET requests
|
||||
if method != "get" and transformed_request.get("data") is not None:
|
||||
request_kwargs["data"] = transformed_request["data"]
|
||||
|
||||
batch_response = getattr(sync_httpx_client, method)(**request_kwargs)
|
||||
elif isinstance(transformed_request, dict) and api_base:
|
||||
# For other providers that use JSON requests
|
||||
batch_response = sync_httpx_client.get(
|
||||
url=api_base,
|
||||
headers={**headers, "Content-Type": "application/json"},
|
||||
params=transformed_request,
|
||||
)
|
||||
else:
|
||||
# Handle other request types if needed
|
||||
if not api_base:
|
||||
raise ValueError("api_base is required for non-pre-signed requests")
|
||||
batch_response = sync_httpx_client.get(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"Error retrieving batch: {e}")
|
||||
raise self._handle_error(
|
||||
e=e,
|
||||
provider_config=provider_config,
|
||||
)
|
||||
|
||||
return provider_config.transform_retrieve_batch_response(
|
||||
model=model,
|
||||
raw_response=batch_response,
|
||||
logging_obj=logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
async def async_create_batch(
|
||||
self,
|
||||
transformed_request: Union[bytes, str, dict],
|
||||
@@ -2606,6 +2695,89 @@ class BaseLLMHTTPHandler:
|
||||
litellm_params=litellm_params_with_request,
|
||||
)
|
||||
|
||||
async def async_retrieve_batch(
|
||||
self,
|
||||
transformed_request: Union[bytes, str, dict],
|
||||
litellm_params: dict,
|
||||
provider_config: "BaseBatchesConfig",
|
||||
headers: dict,
|
||||
api_base: Optional[str],
|
||||
logging_obj: "LiteLLMLoggingObj",
|
||||
client: Optional[Union["HTTPHandler", "AsyncHTTPHandler"]] = None,
|
||||
timeout: Optional[Union[float, httpx.Timeout]] = None,
|
||||
batch_id: Optional[str] = None,
|
||||
model: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Async version of retrieve_batch
|
||||
"""
|
||||
if client is None or not isinstance(client, AsyncHTTPHandler):
|
||||
async_httpx_client = get_async_httpx_client(
|
||||
llm_provider=provider_config.custom_llm_provider
|
||||
)
|
||||
else:
|
||||
async_httpx_client = client
|
||||
|
||||
#########################################################
|
||||
# Debug Logging
|
||||
#########################################################
|
||||
logging_obj.pre_call(
|
||||
input="",
|
||||
api_key="",
|
||||
additional_args={
|
||||
"complete_input_dict": transformed_request,
|
||||
"api_base": api_base,
|
||||
"headers": headers,
|
||||
"batch_id": batch_id,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
if (
|
||||
isinstance(transformed_request, dict)
|
||||
and "method" in transformed_request
|
||||
):
|
||||
# Handle pre-signed requests (e.g., from Bedrock with AWS auth)
|
||||
method = transformed_request["method"].lower()
|
||||
request_kwargs = {
|
||||
"url": transformed_request["url"],
|
||||
"headers": transformed_request["headers"],
|
||||
}
|
||||
|
||||
# Only add data for non-GET requests
|
||||
if method != "get" and transformed_request.get("data") is not None:
|
||||
request_kwargs["data"] = transformed_request["data"]
|
||||
|
||||
batch_response = await getattr(async_httpx_client, method)(**request_kwargs)
|
||||
elif isinstance(transformed_request, dict) and api_base:
|
||||
# For other providers that use JSON requests
|
||||
batch_response = await async_httpx_client.get(
|
||||
url=api_base,
|
||||
headers={**headers, "Content-Type": "application/json"},
|
||||
params=transformed_request,
|
||||
)
|
||||
else:
|
||||
# Handle other request types if needed
|
||||
if not api_base:
|
||||
raise ValueError("api_base is required for non-pre-signed requests")
|
||||
batch_response = await async_httpx_client.get(
|
||||
url=api_base,
|
||||
headers=headers,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"Error retrieving batch: {e}")
|
||||
raise self._handle_error(
|
||||
e=e,
|
||||
provider_config=provider_config,
|
||||
)
|
||||
|
||||
return provider_config.transform_retrieve_batch_response(
|
||||
model=model,
|
||||
raw_response=batch_response,
|
||||
logging_obj=logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
|
||||
def cancel_response_api_handler(
|
||||
self,
|
||||
response_id: str,
|
||||
|
||||
@@ -378,7 +378,7 @@ class OCIChatConfig(BaseConfig):
|
||||
or not oci_compartment_id
|
||||
):
|
||||
raise Exception(
|
||||
"Missing required parameters: oci_user, oci_fingerprint, oci_tenancy, "
|
||||
"Missing required parameters: oci_user, oci_fingerprint, oci_tenancy, oci_compartment_id "
|
||||
"and at least one of oci_key or oci_key_file."
|
||||
)
|
||||
|
||||
|
||||
@@ -239,6 +239,7 @@ class VertexBase:
|
||||
stream=stream,
|
||||
auth_header=None,
|
||||
url=default_api_base,
|
||||
model=model,
|
||||
)
|
||||
return api_base
|
||||
|
||||
@@ -292,6 +293,7 @@ class VertexBase:
|
||||
stream: Optional[bool],
|
||||
auth_header: Optional[str],
|
||||
url: str,
|
||||
model: Optional[str] = None,
|
||||
) -> Tuple[Optional[str], str]:
|
||||
"""
|
||||
for cloudflare ai gateway - https://github.com/BerriAI/litellm/issues/4317
|
||||
@@ -301,7 +303,12 @@ class VertexBase:
|
||||
"""
|
||||
if api_base:
|
||||
if custom_llm_provider == "gemini":
|
||||
url = "{}:{}".format(api_base, endpoint)
|
||||
# For Gemini (Google AI Studio), construct the full path like other providers
|
||||
if model is None:
|
||||
raise ValueError(
|
||||
"Model parameter is required for Gemini custom API base URLs"
|
||||
)
|
||||
url = "{}/models/{}:{}".format(api_base, model, endpoint)
|
||||
if gemini_api_key is None:
|
||||
raise ValueError(
|
||||
"Missing gemini_api_key, please set `GEMINI_API_KEY`"
|
||||
@@ -373,6 +380,7 @@ class VertexBase:
|
||||
endpoint=endpoint,
|
||||
stream=stream,
|
||||
url=url,
|
||||
model=model,
|
||||
)
|
||||
|
||||
def _handle_reauthentication(
|
||||
@@ -384,19 +392,19 @@ class VertexBase:
|
||||
) -> Tuple[str, str]:
|
||||
"""
|
||||
Handle reauthentication when credentials refresh fails.
|
||||
|
||||
|
||||
This method clears the cached credentials and attempts to reload them once.
|
||||
It should only be called when "Reauthentication is needed" error occurs.
|
||||
|
||||
|
||||
Args:
|
||||
credentials: The original credentials
|
||||
project_id: The project ID
|
||||
credential_cache_key: The cache key to clear
|
||||
error: The original error that triggered reauthentication
|
||||
|
||||
|
||||
Returns:
|
||||
Tuple of (access_token, project_id)
|
||||
|
||||
|
||||
Raises:
|
||||
The original error if reauthentication fails
|
||||
"""
|
||||
@@ -404,11 +412,11 @@ class VertexBase:
|
||||
f"Handling reauthentication for project_id: {project_id}. "
|
||||
f"Clearing cache and retrying once."
|
||||
)
|
||||
|
||||
|
||||
# Clear the cached credentials
|
||||
if credential_cache_key in self._credentials_project_mapping:
|
||||
del self._credentials_project_mapping[credential_cache_key]
|
||||
|
||||
|
||||
# Retry once with _retry_reauth=True to prevent infinite recursion
|
||||
try:
|
||||
return self.get_access_token(
|
||||
@@ -438,12 +446,12 @@ class VertexBase:
|
||||
3. Check if loaded credentials have expired
|
||||
4. If expired, refresh credentials
|
||||
5. Return access token and project id
|
||||
|
||||
|
||||
Args:
|
||||
credentials: The credentials to use for authentication
|
||||
project_id: The Google Cloud project ID
|
||||
_retry_reauth: Internal flag to prevent infinite recursion during reauthentication
|
||||
|
||||
|
||||
Returns:
|
||||
Tuple of (access_token, project_id)
|
||||
"""
|
||||
|
||||
+3
-3
@@ -116,6 +116,7 @@ from litellm.utils import (
|
||||
|
||||
from ._logging import verbose_logger
|
||||
from .caching.caching import disable_cache, enable_cache, update_cache
|
||||
from .litellm_core_utils.core_helpers import safe_deep_copy
|
||||
from .litellm_core_utils.fallback_utils import (
|
||||
async_completion_with_fallbacks,
|
||||
completion_with_fallbacks,
|
||||
@@ -2847,8 +2848,7 @@ def completion( # type: ignore # noqa: PLR0915
|
||||
)
|
||||
|
||||
api_base = api_base or litellm.api_base or get_secret("GEMINI_API_BASE")
|
||||
|
||||
new_params = deepcopy(optional_params)
|
||||
new_params = safe_deep_copy(optional_params or {})
|
||||
response = vertex_chat_completion.completion( # type: ignore
|
||||
model=model,
|
||||
messages=messages,
|
||||
@@ -2892,7 +2892,7 @@ def completion( # type: ignore # noqa: PLR0915
|
||||
|
||||
api_base = api_base or litellm.api_base or get_secret("VERTEXAI_API_BASE")
|
||||
|
||||
new_params = deepcopy(optional_params)
|
||||
new_params = safe_deep_copy(optional_params or {})
|
||||
if vertex_partner_models_chat_completion.is_vertex_partner_model(model):
|
||||
model_response = vertex_partner_models_chat_completion.completion(
|
||||
model=model,
|
||||
|
||||
+20210
-19930
File diff suppressed because it is too large
Load Diff
@@ -469,16 +469,13 @@ async def get_end_user_object(
|
||||
# check if in cache
|
||||
cached_user_obj = await user_api_key_cache.async_get_cache(key=_key)
|
||||
if cached_user_obj is not None:
|
||||
if isinstance(cached_user_obj, dict):
|
||||
return_obj = LiteLLM_EndUserTable(**cached_user_obj)
|
||||
check_in_budget(end_user_obj=return_obj)
|
||||
return return_obj
|
||||
elif isinstance(cached_user_obj, LiteLLM_EndUserTable):
|
||||
return_obj = cached_user_obj
|
||||
check_in_budget(end_user_obj=return_obj)
|
||||
return return_obj
|
||||
# Convert cached dict to LiteLLM_EndUserTable instance
|
||||
return_obj = LiteLLM_EndUserTable(**cached_user_obj)
|
||||
check_in_budget(end_user_obj=return_obj)
|
||||
return return_obj
|
||||
|
||||
# else, check db
|
||||
try:
|
||||
try:
|
||||
response = await prisma_client.db.litellm_endusertable.find_unique(
|
||||
where={"user_id": end_user_id},
|
||||
include={"litellm_budget_table": True},
|
||||
@@ -487,9 +484,9 @@ async def get_end_user_object(
|
||||
if response is None:
|
||||
raise Exception
|
||||
|
||||
# save the end-user object to cache
|
||||
# save the end-user object to cache (always store as dict for consistency)
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key="end_user_id:{}".format(end_user_id), value=response
|
||||
key="end_user_id:{}".format(end_user_id), value=response.dict()
|
||||
)
|
||||
|
||||
_response = LiteLLM_EndUserTable(**response.dict())
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
"""
|
||||
Performance utilities for LiteLLM proxy server.
|
||||
|
||||
This module provides performance monitoring and profiling functionality for endpoint
|
||||
performance analysis using cProfile with configurable sampling rates.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import cProfile
|
||||
import functools
|
||||
import threading
|
||||
from pathlib import Path as PathLib
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
# Global profiling state
|
||||
_profile_lock = threading.Lock()
|
||||
_profiler = None
|
||||
_last_profile_file_path = None
|
||||
_sample_counter = 0
|
||||
_sample_counter_lock = threading.Lock()
|
||||
|
||||
|
||||
def _should_sample(profile_sampling_rate: float) -> bool:
|
||||
"""Determine if current request should be sampled based on sampling rate."""
|
||||
if profile_sampling_rate >= 1.0:
|
||||
return True # Always sample
|
||||
elif profile_sampling_rate <= 0.0:
|
||||
return False # Never sample
|
||||
|
||||
# Use deterministic sampling based on counter for consistent rate
|
||||
global _sample_counter
|
||||
with _sample_counter_lock:
|
||||
_sample_counter += 1
|
||||
# Sample based on rate (e.g., 0.1 means sample every 10th request)
|
||||
should_sample = (_sample_counter % int(1.0 / profile_sampling_rate)) == 0
|
||||
return should_sample
|
||||
|
||||
|
||||
def _start_profiling(profile_sampling_rate: float) -> None:
|
||||
"""Start cProfile profiling once globally."""
|
||||
global _profiler
|
||||
with _profile_lock:
|
||||
if _profiler is None:
|
||||
_profiler = cProfile.Profile()
|
||||
_profiler.enable()
|
||||
verbose_proxy_logger.info(f"Profiling started with sampling rate: {profile_sampling_rate}")
|
||||
|
||||
|
||||
def _start_profiling_for_request(profile_sampling_rate: float) -> bool:
|
||||
"""Start profiling for a specific request (if sampling allows)."""
|
||||
if _should_sample(profile_sampling_rate):
|
||||
_start_profiling(profile_sampling_rate)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _save_stats(profile_file: PathLib) -> None:
|
||||
"""Save current stats directly to file."""
|
||||
with _profile_lock:
|
||||
if _profiler is None:
|
||||
return
|
||||
try:
|
||||
# Disable profiler temporarily to dump stats
|
||||
_profiler.disable()
|
||||
_profiler.dump_stats(str(profile_file))
|
||||
# Re-enable profiler to continue profiling
|
||||
_profiler.enable()
|
||||
verbose_proxy_logger.debug(f"Profiling stats saved to {profile_file}")
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"Error saving profiling stats: {e}")
|
||||
# Make sure profiler is re-enabled even if there's an error
|
||||
try:
|
||||
_profiler.enable()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def profile_endpoint(sampling_rate: float = 1.0):
|
||||
"""Decorator to sample endpoint hits and save to a profile file.
|
||||
|
||||
Args:
|
||||
sampling_rate: Rate of requests to profile (0.0 to 1.0)
|
||||
- 1.0: Profile all requests (100%)
|
||||
- 0.1: Profile 1 in 10 requests (10%)
|
||||
- 0.0: Profile no requests (0%)
|
||||
"""
|
||||
def decorator(func):
|
||||
def set_last_profile_path(path: PathLib) -> None:
|
||||
global _last_profile_file_path
|
||||
_last_profile_file_path = path
|
||||
|
||||
if asyncio.iscoroutinefunction(func):
|
||||
@functools.wraps(func)
|
||||
async def async_wrapper(*args, **kwargs):
|
||||
is_sampling = _start_profiling_for_request(sampling_rate)
|
||||
file_path_obj = PathLib("endpoint_profile.pstat")
|
||||
set_last_profile_path(file_path_obj)
|
||||
try:
|
||||
result = await func(*args, **kwargs)
|
||||
if is_sampling:
|
||||
_save_stats(file_path_obj)
|
||||
return result
|
||||
except Exception:
|
||||
if is_sampling:
|
||||
_save_stats(file_path_obj)
|
||||
raise
|
||||
return async_wrapper
|
||||
else:
|
||||
@functools.wraps(func)
|
||||
def sync_wrapper(*args, **kwargs):
|
||||
is_sampling = _start_profiling_for_request(sampling_rate)
|
||||
file_path_obj = PathLib("endpoint_profile.pstat")
|
||||
set_last_profile_path(file_path_obj)
|
||||
try:
|
||||
result = func(*args, **kwargs)
|
||||
if is_sampling:
|
||||
_save_stats(file_path_obj)
|
||||
return result
|
||||
except Exception:
|
||||
if is_sampling:
|
||||
_save_stats(file_path_obj)
|
||||
raise
|
||||
return sync_wrapper
|
||||
return decorator
|
||||
@@ -832,7 +832,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
||||
litellm_parent_otel_span: Union[Span, None] = (
|
||||
_get_parent_otel_span_from_kwargs(kwargs)
|
||||
)
|
||||
user_api_key = kwargs["litellm_params"]["metadata"].get("user_api_key")
|
||||
litellm_metadata = kwargs["litellm_params"]["metadata"]
|
||||
user_api_key = litellm_metadata.get("user_api_key") if litellm_metadata else None
|
||||
pipeline_operations: List[RedisPipelineIncrementOperation] = []
|
||||
|
||||
if user_api_key:
|
||||
|
||||
@@ -169,12 +169,12 @@ def _get_dynamic_logging_metadata(
|
||||
user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig
|
||||
) -> Optional[TeamCallbackMetadata]:
|
||||
callback_settings_obj: Optional[TeamCallbackMetadata] = None
|
||||
key_dynamic_logging_settings: Optional[dict] = (
|
||||
KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)
|
||||
)
|
||||
team_dynamic_logging_settings: Optional[dict] = (
|
||||
KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict)
|
||||
)
|
||||
key_dynamic_logging_settings: Optional[
|
||||
dict
|
||||
] = KeyAndTeamLoggingSettings.get_key_dynamic_logging_settings(user_api_key_dict)
|
||||
team_dynamic_logging_settings: Optional[
|
||||
dict
|
||||
] = KeyAndTeamLoggingSettings.get_team_dynamic_logging_settings(user_api_key_dict)
|
||||
#########################################################################################
|
||||
# Key-based callbacks
|
||||
#########################################################################################
|
||||
@@ -562,6 +562,8 @@ class LiteLLMProxyRequestSetup:
|
||||
user_api_key_logged_metadata = StandardLoggingUserAPIKeyMetadata(
|
||||
user_api_key_hash=user_api_key_dict.api_key, # just the hashed token
|
||||
user_api_key_alias=user_api_key_dict.key_alias,
|
||||
user_api_key_spend=user_api_key_dict.spend,
|
||||
user_api_key_max_budget=user_api_key_dict.max_budget,
|
||||
user_api_key_team_id=user_api_key_dict.team_id,
|
||||
user_api_key_user_id=user_api_key_dict.user_id,
|
||||
user_api_key_org_id=user_api_key_dict.org_id,
|
||||
@@ -569,6 +571,7 @@ class LiteLLMProxyRequestSetup:
|
||||
user_api_key_end_user_id=user_api_key_dict.end_user_id,
|
||||
user_api_key_user_email=user_api_key_dict.user_email,
|
||||
user_api_key_request_route=user_api_key_dict.request_route,
|
||||
user_api_key_budget_reset_at=user_api_key_dict.budget_reset_at,
|
||||
)
|
||||
return user_api_key_logged_metadata
|
||||
|
||||
@@ -611,11 +614,11 @@ class LiteLLMProxyRequestSetup:
|
||||
|
||||
## KEY-LEVEL SPEND LOGS / TAGS
|
||||
if "tags" in key_metadata and key_metadata["tags"] is not None:
|
||||
data[_metadata_variable_name]["tags"] = (
|
||||
LiteLLMProxyRequestSetup._merge_tags(
|
||||
request_tags=data[_metadata_variable_name].get("tags"),
|
||||
tags_to_add=key_metadata["tags"],
|
||||
)
|
||||
data[_metadata_variable_name][
|
||||
"tags"
|
||||
] = LiteLLMProxyRequestSetup._merge_tags(
|
||||
request_tags=data[_metadata_variable_name].get("tags"),
|
||||
tags_to_add=key_metadata["tags"],
|
||||
)
|
||||
if "spend_logs_metadata" in key_metadata and isinstance(
|
||||
key_metadata["spend_logs_metadata"], dict
|
||||
@@ -844,9 +847,9 @@ async def add_litellm_data_to_request( # noqa: PLR0915
|
||||
data[_metadata_variable_name]["litellm_api_version"] = version
|
||||
|
||||
if general_settings is not None:
|
||||
data[_metadata_variable_name]["global_max_parallel_requests"] = (
|
||||
general_settings.get("global_max_parallel_requests", None)
|
||||
)
|
||||
data[_metadata_variable_name][
|
||||
"global_max_parallel_requests"
|
||||
] = general_settings.get("global_max_parallel_requests", None)
|
||||
|
||||
### KEY-LEVEL Controls
|
||||
key_metadata = user_api_key_dict.metadata
|
||||
|
||||
@@ -474,6 +474,9 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
|
||||
user_api_key_team_alias=user_api_key_dict.team_alias,
|
||||
user_api_key_end_user_id=user_api_key_dict.end_user_id,
|
||||
user_api_key_request_route=user_api_key_dict.request_route,
|
||||
user_api_key_spend=user_api_key_dict.spend,
|
||||
user_api_key_max_budget=user_api_key_dict.max_budget,
|
||||
user_api_key_budget_reset_at=user_api_key_dict.budget_reset_at,
|
||||
)
|
||||
)
|
||||
|
||||
@@ -1003,7 +1006,7 @@ class InitPassThroughEndpointHelpers:
|
||||
):
|
||||
"""Add exact path route for pass-through endpoint"""
|
||||
route_key = f"{endpoint_id}:exact:{path}"
|
||||
|
||||
|
||||
# Check if this exact route is already registered
|
||||
if route_key in _registered_pass_through_routes:
|
||||
verbose_proxy_logger.debug(
|
||||
@@ -1011,7 +1014,7 @@ class InitPassThroughEndpointHelpers:
|
||||
path,
|
||||
)
|
||||
return
|
||||
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"adding exact pass through endpoint: %s, dependencies: %s",
|
||||
path,
|
||||
@@ -1032,12 +1035,12 @@ class InitPassThroughEndpointHelpers:
|
||||
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
|
||||
dependencies=dependencies,
|
||||
)
|
||||
|
||||
|
||||
# Register the route to prevent duplicates
|
||||
_registered_pass_through_routes[route_key] = {
|
||||
"endpoint_id": endpoint_id,
|
||||
"path": path,
|
||||
"type": "exact"
|
||||
"type": "exact",
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
@@ -1055,7 +1058,7 @@ class InitPassThroughEndpointHelpers:
|
||||
"""Add wildcard route for sub-paths"""
|
||||
wildcard_path = f"{path}/{{subpath:path}}"
|
||||
route_key = f"{endpoint_id}:subpath:{path}"
|
||||
|
||||
|
||||
# Check if this subpath route is already registered
|
||||
if route_key in _registered_pass_through_routes:
|
||||
verbose_proxy_logger.debug(
|
||||
@@ -1063,7 +1066,7 @@ class InitPassThroughEndpointHelpers:
|
||||
wildcard_path,
|
||||
)
|
||||
return
|
||||
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"adding wildcard pass through endpoint: %s, dependencies: %s",
|
||||
wildcard_path,
|
||||
@@ -1085,19 +1088,20 @@ class InitPassThroughEndpointHelpers:
|
||||
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
|
||||
dependencies=dependencies,
|
||||
)
|
||||
|
||||
|
||||
# Register the route to prevent duplicates
|
||||
_registered_pass_through_routes[route_key] = {
|
||||
"endpoint_id": endpoint_id,
|
||||
"path": path,
|
||||
"type": "subpath"
|
||||
"type": "subpath",
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def remove_endpoint_routes(endpoint_id: str):
|
||||
"""Remove all routes for a specific endpoint ID from the registry"""
|
||||
keys_to_remove = [
|
||||
key for key, value in _registered_pass_through_routes.items()
|
||||
key
|
||||
for key, value in _registered_pass_through_routes.items()
|
||||
if value["endpoint_id"] == endpoint_id
|
||||
]
|
||||
for key in keys_to_remove:
|
||||
@@ -1480,7 +1484,7 @@ async def delete_pass_through_endpoints(
|
||||
pass_through_endpoint_data.pop(endpoint_index)
|
||||
response_obj = found_endpoint
|
||||
|
||||
# Remove routes from registry
|
||||
# Remove routes from registry
|
||||
InitPassThroughEndpointHelpers.remove_endpoint_routes(endpoint_id)
|
||||
|
||||
## Update db
|
||||
|
||||
@@ -141,7 +141,7 @@ class ResponsesSessionHandler:
|
||||
# Add Output messages for this Spend Log
|
||||
############################################################
|
||||
_response_output = spend_log.get("response", "{}")
|
||||
if isinstance(_response_output, dict):
|
||||
if isinstance(_response_output, dict) and _response_output and _response_output != {}:
|
||||
# transform `ChatCompletion Response` to `ResponsesAPIResponse`
|
||||
model_response = ModelResponse(**_response_output)
|
||||
for choice in model_response.choices:
|
||||
|
||||
@@ -83,3 +83,10 @@ class DDLLMObsLatencyMetrics(TypedDict, total=False):
|
||||
time_to_first_token_ms: float
|
||||
litellm_overhead_time_ms: float
|
||||
guardrail_overhead_time_ms: float
|
||||
|
||||
|
||||
class DDLLMObsSpendMetrics(TypedDict, total=False):
|
||||
response_cost: float
|
||||
user_api_key_spend: float
|
||||
user_api_key_max_budget: float
|
||||
user_api_key_budget_reset_at: str
|
||||
|
||||
@@ -3,14 +3,9 @@ from typing import Any, List, Literal, Optional, Union
|
||||
|
||||
from typing_extensions import (
|
||||
TYPE_CHECKING,
|
||||
Protocol,
|
||||
Required,
|
||||
Self,
|
||||
TypedDict,
|
||||
TypeGuard,
|
||||
get_origin,
|
||||
override,
|
||||
runtime_checkable,
|
||||
)
|
||||
|
||||
from .openai import ChatCompletionToolCallChunk
|
||||
@@ -93,6 +88,12 @@ class BedrockConverseReasoningContentBlockDelta(TypedDict, total=False):
|
||||
text: str
|
||||
|
||||
|
||||
class GuardrailConverseContentBlock(TypedDict, total=False):
|
||||
"""Content block for selective guardrail evaluation in Bedrock Converse API"""
|
||||
|
||||
text: str
|
||||
|
||||
|
||||
class ContentBlock(TypedDict, total=False):
|
||||
text: str
|
||||
image: ImageBlock
|
||||
@@ -102,6 +103,7 @@ class ContentBlock(TypedDict, total=False):
|
||||
toolUse: ToolUseBlock
|
||||
cachePoint: CachePointBlock
|
||||
reasoningContent: BedrockConverseReasoningContentBlock
|
||||
guardrailConverseContent: GuardrailConverseContentBlock
|
||||
|
||||
|
||||
class MessageBlock(TypedDict):
|
||||
@@ -581,30 +583,35 @@ class AmazonDeepSeekR1StreamingResponse(TypedDict):
|
||||
|
||||
class BedrockS3InputDataConfig(TypedDict):
|
||||
"""S3 input data configuration for Bedrock batch jobs."""
|
||||
|
||||
s3Uri: str
|
||||
|
||||
|
||||
class BedrockInputDataConfig(TypedDict):
|
||||
"""Input data configuration for Bedrock batch jobs."""
|
||||
|
||||
s3InputDataConfig: BedrockS3InputDataConfig
|
||||
|
||||
|
||||
class BedrockS3OutputDataConfig(TypedDict):
|
||||
"""S3 output data configuration for Bedrock batch jobs."""
|
||||
|
||||
s3Uri: str
|
||||
|
||||
|
||||
class BedrockOutputDataConfig(TypedDict):
|
||||
"""Output data configuration for Bedrock batch jobs."""
|
||||
|
||||
s3OutputDataConfig: BedrockS3OutputDataConfig
|
||||
|
||||
|
||||
class BedrockCreateBatchRequest(TypedDict, total=False):
|
||||
"""
|
||||
Request structure for creating a Bedrock batch inference job.
|
||||
|
||||
|
||||
Reference: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_CreateModelInvocationJob.html
|
||||
"""
|
||||
|
||||
jobName: str
|
||||
roleArn: str
|
||||
modelId: str
|
||||
@@ -616,21 +623,17 @@ class BedrockCreateBatchRequest(TypedDict, total=False):
|
||||
|
||||
|
||||
BedrockBatchJobStatus = Literal[
|
||||
"Submitted",
|
||||
"InProgress",
|
||||
"Completed",
|
||||
"Failed",
|
||||
"Stopping",
|
||||
"Stopped"
|
||||
"Submitted", "InProgress", "Completed", "Failed", "Stopping", "Stopped"
|
||||
]
|
||||
|
||||
|
||||
class BedrockCreateBatchResponse(TypedDict):
|
||||
"""
|
||||
Response structure from creating a Bedrock batch inference job.
|
||||
|
||||
|
||||
Reference: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_CreateModelInvocationJob.html
|
||||
"""
|
||||
|
||||
jobArn: str
|
||||
jobName: str
|
||||
status: BedrockBatchJobStatus
|
||||
@@ -639,9 +642,10 @@ class BedrockCreateBatchResponse(TypedDict):
|
||||
class BedrockGetBatchResponse(TypedDict, total=False):
|
||||
"""
|
||||
Response structure from getting a Bedrock batch inference job.
|
||||
|
||||
|
||||
Reference: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_GetModelInvocationJob.html
|
||||
"""
|
||||
|
||||
jobArn: str
|
||||
jobName: str
|
||||
modelId: str
|
||||
|
||||
@@ -723,6 +723,7 @@ ValidUserMessageContentTypes = [
|
||||
"input_audio",
|
||||
"audio_url",
|
||||
"document",
|
||||
"guarded_text",
|
||||
"video_url",
|
||||
"file",
|
||||
] # used for validating user messages. Prevent users from accidentally sending anthropic messages.
|
||||
|
||||
@@ -123,6 +123,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
||||
max_output_tokens: Required[Optional[int]]
|
||||
input_cost_per_token: Required[float]
|
||||
cache_creation_input_token_cost: Optional[float]
|
||||
cache_creation_input_token_cost_above_1hr: Optional[float]
|
||||
cache_read_input_token_cost: Optional[float]
|
||||
input_cost_per_character: Optional[float] # only for vertex ai models
|
||||
input_cost_per_audio_token: Optional[float]
|
||||
@@ -162,7 +163,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
|
||||
SearchContextCostPerQuery
|
||||
] # Cost for using web search tool
|
||||
citation_cost_per_token: Optional[float] # Cost per citation token for Perplexity
|
||||
tiered_pricing: Optional[List[Dict[str, Any]]] # Tiered pricing structure for models like Dashscope
|
||||
tiered_pricing: Optional[
|
||||
List[Dict[str, Any]]
|
||||
] # Tiered pricing structure for models like Dashscope
|
||||
litellm_provider: Required[str]
|
||||
mode: Required[
|
||||
Literal[
|
||||
@@ -1807,6 +1810,9 @@ class AdapterCompletionStreamWrapper:
|
||||
class StandardLoggingUserAPIKeyMetadata(TypedDict):
|
||||
user_api_key_hash: Optional[str] # hash of the litellm virtual key used
|
||||
user_api_key_alias: Optional[str]
|
||||
user_api_key_spend: Optional[float]
|
||||
user_api_key_max_budget: Optional[float]
|
||||
user_api_key_budget_reset_at: Optional[str]
|
||||
user_api_key_org_id: Optional[str]
|
||||
user_api_key_team_id: Optional[str]
|
||||
user_api_key_user_id: Optional[str]
|
||||
|
||||
+20210
-19930
File diff suppressed because it is too large
Load Diff
@@ -75,6 +75,19 @@ async def test_async_file_and_batch():
|
||||
)
|
||||
print("CREATED BATCH RESPONSE=", create_batch_response)
|
||||
|
||||
# retrieve batch
|
||||
retrieve_batch_response = await litellm.aretrieve_batch(
|
||||
batch_id=create_batch_response.id,
|
||||
custom_llm_provider="bedrock",
|
||||
model="us.anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
)
|
||||
print("RETRIEVED BATCH RESPONSE=", retrieve_batch_response)
|
||||
|
||||
# Validate the response
|
||||
assert retrieve_batch_response.id == create_batch_response.id
|
||||
assert retrieve_batch_response.object == "batch"
|
||||
assert retrieve_batch_response.status in ["validating", "in_progress", "completed", "failed", "cancelled"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_mock_bedrock_file_url_mapping():
|
||||
@@ -118,3 +131,65 @@ async def test_mock_bedrock_file_url_mapping():
|
||||
expected_s3_uri, _ = bedrock_config._convert_https_url_to_s3_uri(captured_put_url)
|
||||
assert file_obj.id == expected_s3_uri
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_bedrock_retrieve_batch():
|
||||
"""
|
||||
Test bedrock batch retrieval functionality, validating that input and output file IDs
|
||||
are correctly extracted from the Bedrock response and included in the final transformed response.
|
||||
"""
|
||||
print("Testing bedrock batch retrieval")
|
||||
|
||||
# Mock bedrock batch response
|
||||
mock_bedrock_response = {
|
||||
"jobArn": "arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/test-job-123",
|
||||
"jobName": "test-job-123",
|
||||
"modelId": "us.anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
"roleArn": "arn:aws:iam::123456789012:role/service-role/AmazonBedrockExecutionRoleForAgents_TEST",
|
||||
"status": "InProgress",
|
||||
"message": "Job is in progress",
|
||||
"submitTime": "2024-01-01T12:00:00Z",
|
||||
"lastModifiedTime": "2024-01-01T12:30:00Z",
|
||||
"inputDataConfig": {
|
||||
"s3InputDataConfig": {
|
||||
"s3Uri": "s3://test-bucket/input/test-input.jsonl"
|
||||
}
|
||||
},
|
||||
"outputDataConfig": {
|
||||
"s3OutputDataConfig": {
|
||||
"s3Uri": "s3://test-bucket/output/"
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
# Mock the HTTP response
|
||||
mock_response = MagicMock()
|
||||
mock_response.json.return_value = mock_bedrock_response
|
||||
mock_response.status_code = 200
|
||||
|
||||
# Print the mock response to debug
|
||||
print("MOCK RESPONSE DATA:", mock_bedrock_response)
|
||||
|
||||
with patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get") as mock_get:
|
||||
mock_response.raise_for_status.return_value = None
|
||||
mock_get.return_value = mock_response
|
||||
|
||||
# Test retrieve batch
|
||||
batch_response = await litellm.aretrieve_batch(
|
||||
batch_id="arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/test-job-123",
|
||||
custom_llm_provider="bedrock",
|
||||
model="us.anthropic.claude-3-5-sonnet-20240620-v1:0",
|
||||
)
|
||||
|
||||
print("MOCKED BATCH RESPONSE=", batch_response)
|
||||
|
||||
# Validate the response
|
||||
assert batch_response.id == "arn:aws:bedrock:us-west-2:123456789012:model-invocation-job/test-job-123"
|
||||
assert batch_response.object == "batch"
|
||||
assert batch_response.status == "in_progress" # Bedrock "InProgress" maps to "in_progress"
|
||||
assert batch_response.endpoint == "/v1/chat/completions"
|
||||
|
||||
# Validate input and output file IDs in the final transformed response
|
||||
assert batch_response.input_file_id == "s3://test-bucket/input/test-input.jsonl"
|
||||
assert batch_response.output_file_id == "s3://test-bucket/output/"
|
||||
|
||||
|
||||
@@ -1,203 +0,0 @@
|
||||
import os
|
||||
import sys
|
||||
import traceback
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
import io
|
||||
import os
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../..")
|
||||
) # Adds the parent directory to the system path
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import completion
|
||||
from litellm.llms.cohere.completion.transformation import CohereTextConfig
|
||||
|
||||
|
||||
def test_cohere_generate_api_completion():
|
||||
try:
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
client = HTTPHandler()
|
||||
litellm.set_verbose = True
|
||||
messages = [
|
||||
{"role": "system", "content": "You're a good bot"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Hey",
|
||||
},
|
||||
]
|
||||
|
||||
with patch.object(client, "post") as mock_client:
|
||||
try:
|
||||
completion(
|
||||
model="cohere/command",
|
||||
messages=messages,
|
||||
max_tokens=10,
|
||||
client=client,
|
||||
)
|
||||
except Exception as e:
|
||||
print(e)
|
||||
mock_client.assert_called_once()
|
||||
print("mock_client.call_args.kwargs", mock_client.call_args.kwargs)
|
||||
|
||||
assert (
|
||||
mock_client.call_args.kwargs["url"]
|
||||
== "https://api.cohere.ai/v1/generate"
|
||||
)
|
||||
json_data = json.loads(mock_client.call_args.kwargs["data"])
|
||||
assert json_data["model"] == "command"
|
||||
assert json_data["prompt"] == "You're a good bot Hey"
|
||||
assert json_data["max_tokens"] == 10
|
||||
except Exception as e:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cohere_generate_api_stream():
|
||||
try:
|
||||
litellm.set_verbose = True
|
||||
messages = [
|
||||
{"role": "system", "content": "You're a good bot"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Hey",
|
||||
},
|
||||
]
|
||||
response = await litellm.acompletion(
|
||||
model="cohere/command",
|
||||
messages=messages,
|
||||
max_tokens=10,
|
||||
stream=True,
|
||||
)
|
||||
print("async cohere stream response", response)
|
||||
async for chunk in response:
|
||||
print(chunk)
|
||||
except Exception as e:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
def test_completion_cohere_stream_bad_key():
|
||||
try:
|
||||
api_key = "bad-key"
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "how does a court case get to the Supreme Court?",
|
||||
},
|
||||
]
|
||||
completion(
|
||||
model="command",
|
||||
messages=messages,
|
||||
stream=True,
|
||||
max_tokens=50,
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
except litellm.AuthenticationError as e:
|
||||
pass
|
||||
except Exception as e:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
def test_cohere_transform_request():
|
||||
try:
|
||||
config = CohereTextConfig()
|
||||
messages = [
|
||||
{"role": "system", "content": "You're a helpful bot"},
|
||||
{"role": "user", "content": "Hello"},
|
||||
]
|
||||
optional_params = {"max_tokens": 10, "temperature": 0.7}
|
||||
headers = {}
|
||||
|
||||
transformed_request = config.transform_request(
|
||||
model="command",
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
print("transformed_request", json.dumps(transformed_request, indent=4))
|
||||
|
||||
assert transformed_request["model"] == "command"
|
||||
assert transformed_request["prompt"] == "You're a helpful bot Hello"
|
||||
assert transformed_request["max_tokens"] == 10
|
||||
assert transformed_request["temperature"] == 0.7
|
||||
except Exception as e:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
def test_cohere_transform_request_with_tools():
|
||||
try:
|
||||
config = CohereTextConfig()
|
||||
messages = [{"role": "user", "content": "What's the weather?"}]
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather information",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"location": {"type": "string"}},
|
||||
},
|
||||
},
|
||||
}
|
||||
]
|
||||
optional_params = {"tools": tools}
|
||||
|
||||
transformed_request = config.transform_request(
|
||||
model="command",
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
print("transformed_request", json.dumps(transformed_request, indent=4))
|
||||
assert "tools" in transformed_request
|
||||
assert transformed_request["tools"] == {"tools": tools}
|
||||
except Exception as e:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
|
||||
|
||||
def test_cohere_map_openai_params():
|
||||
try:
|
||||
config = CohereTextConfig()
|
||||
openai_params = {
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 100,
|
||||
"n": 2,
|
||||
"top_p": 0.9,
|
||||
"frequency_penalty": 0.5,
|
||||
"presence_penalty": 0.5,
|
||||
"stop": ["END"],
|
||||
"stream": True,
|
||||
}
|
||||
|
||||
mapped_params = config.map_openai_params(
|
||||
non_default_params=openai_params,
|
||||
optional_params={},
|
||||
model="command",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert mapped_params["temperature"] == 0.7
|
||||
assert mapped_params["max_tokens"] == 100
|
||||
assert mapped_params["num_generations"] == 2
|
||||
assert mapped_params["p"] == 0.9
|
||||
assert mapped_params["frequency_penalty"] == 0.5
|
||||
assert mapped_params["presence_penalty"] == 0.5
|
||||
assert mapped_params["stop_sequences"] == ["END"]
|
||||
assert mapped_params["stream"] == True
|
||||
except Exception as e:
|
||||
pytest.fail(f"Error occurred: {e}")
|
||||
@@ -300,6 +300,76 @@ async def test_anthropic_api_prompt_caching_basic():
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_anthropic_api_prompt_caching_basic_with_cache_creation():
|
||||
from uuid import uuid4
|
||||
|
||||
random_id = uuid4()
|
||||
|
||||
litellm.set_verbose = True
|
||||
response = await litellm.acompletion(
|
||||
model="anthropic/claude-3-5-sonnet-20240620",
|
||||
messages=[
|
||||
# System Message
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Here is the full text of a complex legal agreement {}".format(
|
||||
random_id
|
||||
)
|
||||
* 400,
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
},
|
||||
# marked for caching with the cache_control parameter, so that this checkpoint can read from the previous cache.
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What are the key terms and conditions in this agreement?",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"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": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What are the key terms and conditions in this agreement?",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
temperature=0.2,
|
||||
max_tokens=10,
|
||||
extra_headers={
|
||||
"anthropic-version": "2023-06-01",
|
||||
"anthropic-beta": "prompt-caching-2024-07-31",
|
||||
},
|
||||
)
|
||||
|
||||
print("response=", response)
|
||||
|
||||
assert "cache_read_input_tokens" in response.usage
|
||||
assert "cache_creation_input_tokens" in response.usage
|
||||
|
||||
# Assert either a cache entry was created or cache was read - changes depending on the anthropic api ttl
|
||||
assert (response.usage.cache_read_input_tokens > 0) or (
|
||||
response.usage.cache_creation_input_tokens > 0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio()
|
||||
async def test_anthropic_api_prompt_caching_with_content_str():
|
||||
system_message = [
|
||||
|
||||
@@ -145,9 +145,6 @@ def test_cancel_response():
|
||||
|
||||
# verify cancel response structure
|
||||
assert hasattr(cancel_response, "id")
|
||||
# Note: Cancel response returns ResponsesAPIResponse, not DeleteResponseResult
|
||||
# The actual response structure depends on the provider implementation
|
||||
assert isinstance(cancel_response, ResponsesAPIResponse)
|
||||
except Exception as e:
|
||||
if "Cannot cancel a completed response" in str(e):
|
||||
pass
|
||||
@@ -179,9 +176,6 @@ def test_cancel_streaming_response():
|
||||
cancel_response = client.responses.cancel(response_id)
|
||||
print("CANCEL streaming response=", cancel_response)
|
||||
assert hasattr(cancel_response, "id")
|
||||
# Note: Cancel response returns ResponsesAPIResponse, not DeleteResponseResult
|
||||
# The actual response structure depends on the provider implementation
|
||||
assert isinstance(cancel_response, ResponsesAPIResponse)
|
||||
except Exception as e:
|
||||
if "Cannot cancel a completed response" in str(e):
|
||||
pass
|
||||
|
||||
@@ -48,7 +48,7 @@ async def test_get_end_user_object(customer_spend, customer_budget):
|
||||
)
|
||||
_cache = DualCache()
|
||||
_key = "end_user_id:{}".format(end_user_id)
|
||||
_cache.set_cache(key=_key, value=end_user_obj)
|
||||
_cache.set_cache(key=_key, value=end_user_obj.model_dump())
|
||||
try:
|
||||
await get_end_user_object(
|
||||
end_user_id=end_user_id,
|
||||
|
||||
@@ -345,6 +345,130 @@ async def test_generationconfig_to_config_mapping(sample_request_payload):
|
||||
print("✅ generationConfig to config mapping test passed")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_custom_api_base_proxy_integration():
|
||||
"""
|
||||
Test that Gemini models work correctly with custom API base URLs in proxy context.
|
||||
|
||||
This test verifies that when a custom api_base is provided for Gemini models,
|
||||
the URL is correctly constructed using the _check_custom_proxy method.
|
||||
"""
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
|
||||
# Test the _check_custom_proxy method directly
|
||||
vertex_base = VertexBase()
|
||||
|
||||
# Test case 1: Custom API base for Gemini
|
||||
custom_api_base = "https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta"
|
||||
model = "gemini-2.5-flash-lite"
|
||||
endpoint = "generateContent"
|
||||
|
||||
auth_header, result_url = vertex_base._check_custom_proxy(
|
||||
api_base=custom_api_base,
|
||||
custom_llm_provider="gemini",
|
||||
gemini_api_key="test-api-key",
|
||||
endpoint=endpoint,
|
||||
stream=False,
|
||||
auth_header=None,
|
||||
url=f"https://generativelanguage.googleapis.com/v1beta/models/{model}:{endpoint}",
|
||||
model=model,
|
||||
)
|
||||
|
||||
# Verify the URL is correctly constructed
|
||||
expected_url = f"{custom_api_base}/models/{model}:{endpoint}"
|
||||
assert result_url == expected_url, f"Expected {expected_url}, got {result_url}"
|
||||
|
||||
# Verify the auth header is set to the API key
|
||||
assert auth_header == "test-api-key", f"Expected 'test-api-key', got {auth_header}"
|
||||
|
||||
print(f"✅ Custom API base URL construction test passed: {result_url}")
|
||||
|
||||
# Test case 2: Custom API base with streaming
|
||||
auth_header_streaming, result_url_streaming = vertex_base._check_custom_proxy(
|
||||
api_base=custom_api_base,
|
||||
custom_llm_provider="gemini",
|
||||
gemini_api_key="test-api-key",
|
||||
endpoint=endpoint,
|
||||
stream=True,
|
||||
auth_header=None,
|
||||
url=f"https://generativelanguage.googleapis.com/v1beta/models/{model}:{endpoint}",
|
||||
model=model,
|
||||
)
|
||||
|
||||
# Verify streaming URL has ?alt=sse parameter
|
||||
expected_streaming_url = f"{custom_api_base}/models/{model}:{endpoint}?alt=sse"
|
||||
assert result_url_streaming == expected_streaming_url, f"Expected {expected_streaming_url}, got {result_url_streaming}"
|
||||
|
||||
print(f"✅ Custom API base streaming URL test passed: {result_url_streaming}")
|
||||
|
||||
# Test case 3: Error handling - missing API key
|
||||
with pytest.raises(ValueError, match="Missing gemini_api_key"):
|
||||
vertex_base._check_custom_proxy(
|
||||
api_base=custom_api_base,
|
||||
custom_llm_provider="gemini",
|
||||
gemini_api_key=None, # Missing API key
|
||||
endpoint=endpoint,
|
||||
stream=False,
|
||||
auth_header=None,
|
||||
url=f"https://generativelanguage.googleapis.com/v1beta/models/{model}:{endpoint}",
|
||||
model=model,
|
||||
)
|
||||
|
||||
print("✅ Missing API key error handling test passed")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_proxy_config_with_custom_api_base():
|
||||
"""
|
||||
Test that proxy configuration correctly handles custom API base for Gemini models.
|
||||
|
||||
This test simulates the proxy configuration scenario where a model is configured
|
||||
with a custom api_base in the config.yaml file.
|
||||
"""
|
||||
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
|
||||
|
||||
# Simulate proxy configuration
|
||||
model_config = {
|
||||
"model_name": "byok-gemini/*",
|
||||
"litellm_params": {
|
||||
"model": "gemini/*",
|
||||
"api_key": "dummy-key-for-testing",
|
||||
"api_base": "https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta"
|
||||
}
|
||||
}
|
||||
|
||||
vertex_base = VertexBase()
|
||||
|
||||
# Test with different Gemini models
|
||||
test_models = [
|
||||
"gemini-2.5-flash-lite",
|
||||
"gemini-2.5-pro",
|
||||
"gemini-1.5-flash",
|
||||
"gemini-1.5-pro"
|
||||
]
|
||||
|
||||
for model in test_models:
|
||||
# Test generateContent endpoint
|
||||
auth_header, result_url = vertex_base._check_custom_proxy(
|
||||
api_base=model_config["litellm_params"]["api_base"],
|
||||
custom_llm_provider="gemini",
|
||||
gemini_api_key=model_config["litellm_params"]["api_key"],
|
||||
endpoint="generateContent",
|
||||
stream=False,
|
||||
auth_header=None,
|
||||
url=f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent",
|
||||
model=model,
|
||||
)
|
||||
|
||||
expected_url = f"{model_config['litellm_params']['api_base']}/models/{model}:generateContent"
|
||||
assert result_url == expected_url, f"Expected {expected_url}, got {result_url} for model {model}"
|
||||
assert auth_header == model_config["litellm_params"]["api_key"], f"Expected API key, got {auth_header} for model {model}"
|
||||
|
||||
print(f"✅ Model {model} configuration test passed: {result_url}")
|
||||
|
||||
print("✅ Proxy configuration with custom API base test passed")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run the tests
|
||||
pytest.main([__file__, "-v"])
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import asyncio
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timedelta
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import Optional
|
||||
from unittest.mock import Mock, patch, MagicMock
|
||||
|
||||
@@ -661,17 +661,24 @@ def create_standard_logging_payload_with_tool_calls() -> StandardLoggingPayload:
|
||||
"""Create a StandardLoggingPayload object with tool calls for testing"""
|
||||
return {
|
||||
"id": "test-request-id-tool-calls",
|
||||
"trace_id": "test-trace-id-tool-calls",
|
||||
"call_type": "completion",
|
||||
"stream": None,
|
||||
"response_cost": 0.05,
|
||||
"response_cost_failure_debug_info": None,
|
||||
"status": "success",
|
||||
"custom_llm_provider": "openai",
|
||||
"total_tokens": 50,
|
||||
"prompt_tokens": 20,
|
||||
"completion_tokens": 30,
|
||||
"startTime": 1234567890.0,
|
||||
"endTime": 1234567891.0,
|
||||
"completionStartTime": 1234567890.5,
|
||||
"model_map_information": {"model_map_key": "gpt-4", "model_map_value": None},
|
||||
"response_time": 1.0,
|
||||
"model_map_information": {
|
||||
"model_map_key": "gpt-4",
|
||||
"model_map_value": None
|
||||
},
|
||||
"model": "gpt-4",
|
||||
"model_id": "model-123",
|
||||
"model_group": "openai-gpt",
|
||||
@@ -746,6 +753,7 @@ def create_standard_logging_payload_with_tool_calls() -> StandardLoggingPayload:
|
||||
]
|
||||
},
|
||||
"error_str": None,
|
||||
"error_information": None,
|
||||
"model_parameters": {"temperature": 0.7},
|
||||
"hidden_params": {
|
||||
"model_id": "model-123",
|
||||
@@ -758,14 +766,9 @@ def create_standard_logging_payload_with_tool_calls() -> StandardLoggingPayload:
|
||||
"litellm_model_name": None,
|
||||
"usage_object": None,
|
||||
},
|
||||
"stream": None,
|
||||
"response_time": 1.0,
|
||||
"error_information": None,
|
||||
"guardrail_information": None,
|
||||
"standard_built_in_tools_params": None,
|
||||
"trace_id": "test-trace-id-tool-calls",
|
||||
"custom_llm_provider": "openai",
|
||||
}
|
||||
} # type: ignore
|
||||
|
||||
|
||||
class TestDataDogLLMObsLoggerToolCalls:
|
||||
@@ -897,3 +900,204 @@ class TestDataDogLLMObsLoggerToolCalls:
|
||||
assert len(output_tool_calls) == 1
|
||||
output_function_info = output_tool_calls[0].get("function", {})
|
||||
assert output_function_info.get("name") == "format_response"
|
||||
|
||||
def create_standard_logging_payload_with_spend_metrics() -> StandardLoggingPayload:
|
||||
"""Create a StandardLoggingPayload object with spend metrics for testing"""
|
||||
from datetime import datetime, timezone
|
||||
|
||||
# Create a budget reset time 10 days from now (using "10d" format)
|
||||
budget_reset_at = datetime.now(timezone.utc) + timedelta(days=10)
|
||||
|
||||
return {
|
||||
"id": "test-request-id-spend",
|
||||
"trace_id": "test-trace-id-spend",
|
||||
"call_type": "completion",
|
||||
"stream": None,
|
||||
"response_cost": 0.15,
|
||||
"response_cost_failure_debug_info": None,
|
||||
"status": "success",
|
||||
"custom_llm_provider": "openai",
|
||||
"total_tokens": 30,
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 20,
|
||||
"startTime": 1234567890.0,
|
||||
"endTime": 1234567891.0,
|
||||
"completionStartTime": 1234567890.5,
|
||||
"response_time": 1.0,
|
||||
"model_map_information": {
|
||||
"model_map_key": "gpt-4",
|
||||
"model_map_value": None
|
||||
},
|
||||
"model": "gpt-4",
|
||||
"model_id": "model-123",
|
||||
"model_group": "openai-gpt",
|
||||
"api_base": "https://api.openai.com",
|
||||
"metadata": {
|
||||
"user_api_key_hash": "test_hash",
|
||||
"user_api_key_org_id": None,
|
||||
"user_api_key_alias": "test_alias",
|
||||
"user_api_key_team_id": "test_team",
|
||||
"user_api_key_user_id": "test_user",
|
||||
"user_api_key_team_alias": "test_team_alias",
|
||||
"user_api_key_user_email": None,
|
||||
"user_api_key_end_user_id": None,
|
||||
"user_api_key_request_route": None,
|
||||
"user_api_key_spend": 0.67,
|
||||
"user_api_key_max_budget": 10.0, # $10 max budget
|
||||
"user_api_key_budget_reset_at": budget_reset_at.isoformat(), # ISO format: 2025-09-26T...
|
||||
"spend_logs_metadata": None,
|
||||
"requester_ip_address": "127.0.0.1",
|
||||
"requester_metadata": None,
|
||||
"requester_custom_headers": None,
|
||||
"prompt_management_metadata": None,
|
||||
"mcp_tool_call_metadata": None,
|
||||
"vector_store_request_metadata": None,
|
||||
"applied_guardrails": None,
|
||||
"usage_object": None,
|
||||
"cold_storage_object_key": None,
|
||||
},
|
||||
"cache_hit": False,
|
||||
"cache_key": None,
|
||||
"saved_cache_cost": 0.0,
|
||||
"request_tags": [],
|
||||
"end_user": None,
|
||||
"requester_ip_address": "127.0.0.1",
|
||||
"messages": [{"role": "user", "content": "Hello, world!"}],
|
||||
"response": {"choices": [{"message": {"content": "Hi there!"}}]},
|
||||
"error_str": None,
|
||||
"error_information": None,
|
||||
"model_parameters": {"stream": False},
|
||||
"hidden_params": {
|
||||
"model_id": "model-123",
|
||||
"cache_key": None,
|
||||
"api_base": "https://api.openai.com",
|
||||
"response_cost": "0.15",
|
||||
"litellm_overhead_time_ms": None,
|
||||
"additional_headers": None,
|
||||
"batch_models": None,
|
||||
"litellm_model_name": None,
|
||||
"usage_object": None,
|
||||
},
|
||||
"guardrail_information": None,
|
||||
"standard_built_in_tools_params": None,
|
||||
} # type: ignore
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_datadog_llm_obs_spend_metrics(mock_env_vars):
|
||||
"""Test that budget metrics are properly extracted and logged"""
|
||||
datadog_llm_obs_logger = DataDogLLMObsLogger()
|
||||
|
||||
# Create a standard logging payload with spend metrics
|
||||
payload = create_standard_logging_payload_with_spend_metrics()
|
||||
|
||||
# Show the budget reset time in ISO format
|
||||
budget_reset_iso = payload["metadata"]["user_api_key_budget_reset_at"]
|
||||
print(f"Budget reset time (ISO format): {budget_reset_iso}")
|
||||
from datetime import datetime, timezone
|
||||
print(f"Current time: {datetime.now(timezone.utc).isoformat()}")
|
||||
|
||||
# Test the _get_spend_metrics method
|
||||
spend_metrics = datadog_llm_obs_logger._get_spend_metrics(payload)
|
||||
|
||||
# Verify budget metrics are present
|
||||
assert "user_api_key_max_budget" in spend_metrics
|
||||
assert spend_metrics["user_api_key_max_budget"] == 10.0
|
||||
|
||||
assert "user_api_key_budget_reset_at" in spend_metrics
|
||||
# The budget reset should be a datetime string in ISO format
|
||||
budget_reset = spend_metrics["user_api_key_budget_reset_at"]
|
||||
assert isinstance(budget_reset, str)
|
||||
print(f"Budget reset datetime: {budget_reset}")
|
||||
# Should be close to 10 days from now
|
||||
budget_reset_dt = datetime.fromisoformat(budget_reset.replace('Z', '+00:00'))
|
||||
now = datetime.now(timezone.utc)
|
||||
time_diff = (budget_reset_dt - now).total_seconds() / 86400 # days
|
||||
assert 9.5 <= time_diff <= 10.5 # Should be close to 10 days
|
||||
|
||||
print(f"Spend metrics: {spend_metrics}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_datadog_llm_obs_spend_metrics_no_budget(mock_env_vars):
|
||||
"""Test that spend metrics work when no budget is set"""
|
||||
datadog_llm_obs_logger = DataDogLLMObsLogger()
|
||||
|
||||
# Create a standard logging payload without budget metadata
|
||||
payload = create_standard_logging_payload_with_spend_metrics()
|
||||
|
||||
# Remove budget-related metadata to test no-budget scenario
|
||||
payload["metadata"].pop("user_api_key_max_budget", None)
|
||||
payload["metadata"].pop("user_api_key_budget_reset_at", None)
|
||||
|
||||
# Test the _get_spend_metrics method
|
||||
spend_metrics = datadog_llm_obs_logger._get_spend_metrics(payload)
|
||||
|
||||
# Verify only response cost is present
|
||||
assert "response_cost" in spend_metrics
|
||||
assert spend_metrics["response_cost"] == 0.15
|
||||
|
||||
# Budget metrics should not be present
|
||||
assert "user_api_key_max_budget" not in spend_metrics
|
||||
assert "user_api_key_budget_reset_at" not in spend_metrics
|
||||
|
||||
print(f"Spend metrics (no budget): {spend_metrics}")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_spend_metrics_in_datadog_payload(mock_env_vars):
|
||||
"""Test that spend metrics are correctly included in DataDog LLM Observability payloads"""
|
||||
from datetime import datetime
|
||||
datadog_llm_obs_logger = DataDogLLMObsLogger()
|
||||
|
||||
standard_payload = create_standard_logging_payload_with_spend_metrics()
|
||||
|
||||
kwargs = {
|
||||
"standard_logging_object": standard_payload,
|
||||
"litellm_params": {"metadata": {}},
|
||||
}
|
||||
|
||||
start_time = datetime.now()
|
||||
end_time = datetime.now()
|
||||
|
||||
payload = datadog_llm_obs_logger.create_llm_obs_payload(kwargs, start_time, end_time)
|
||||
|
||||
# Verify basic payload structure
|
||||
assert payload.get("name") == "litellm_llm_call"
|
||||
assert payload.get("status") == "ok"
|
||||
|
||||
# Verify spend metrics are included in metadata
|
||||
meta = payload.get("meta", {})
|
||||
assert meta is not None, "Meta section should exist in payload"
|
||||
|
||||
metadata = meta.get("metadata", {})
|
||||
assert metadata is not None, "Metadata section should exist in meta"
|
||||
|
||||
spend_metrics = metadata.get("spend_metrics", {})
|
||||
assert spend_metrics, "Spend metrics should exist in metadata"
|
||||
|
||||
# Check that all metrics are present
|
||||
assert "response_cost" in spend_metrics
|
||||
assert "user_api_key_spend" in spend_metrics
|
||||
assert "user_api_key_max_budget" in spend_metrics
|
||||
assert "user_api_key_budget_reset_at" in spend_metrics
|
||||
|
||||
# Verify the values are correct
|
||||
assert spend_metrics["response_cost"] == 0.15 # response_cost
|
||||
assert spend_metrics["user_api_key_spend"] == 0.67 # lol
|
||||
assert spend_metrics["user_api_key_max_budget"] == 10.0 # max budget
|
||||
|
||||
# Verify budget reset is a datetime string in ISO format
|
||||
budget_reset = spend_metrics["user_api_key_budget_reset_at"]
|
||||
assert isinstance(budget_reset, str)
|
||||
print(f"Budget reset in payload: {budget_reset}") # In StandardLoggingUserAPIKeyMetadata
|
||||
user_api_key_budget_reset_at: Optional[str] = None
|
||||
|
||||
# In DDLLMObsSpendMetrics
|
||||
user_api_key_budget_reset_at: str
|
||||
# Should be close to 10 days from now
|
||||
from datetime import datetime, timezone
|
||||
budget_reset_dt = datetime.fromisoformat(budget_reset.replace('Z', '+00:00'))
|
||||
now = datetime.now(timezone.utc)
|
||||
time_diff = (budget_reset_dt - now).total_seconds() / 86400 # days
|
||||
assert 9.5 <= time_diff <= 10.5 # Should be close to 10 days
|
||||
|
||||
@@ -6,6 +6,8 @@ import pytest
|
||||
|
||||
from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger
|
||||
from litellm.types.integrations.langfuse_otel import LangfuseOtelConfig
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from datetime import datetime
|
||||
|
||||
|
||||
class TestLangfuseOtelIntegration:
|
||||
@@ -241,6 +243,134 @@ class TestLangfuseOtelIntegration:
|
||||
_ = LangfuseOtelLogger.get_langfuse_otel_config()
|
||||
|
||||
assert os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT") == "https://otel-host.com/api/public/otel"
|
||||
|
||||
|
||||
class TestLangfuseOtelResponsesAPI:
|
||||
"""Test suite for Langfuse OTEL integration with ResponsesAPI"""
|
||||
|
||||
def test_langfuse_otel_with_responses_api(self):
|
||||
"""Test that Langfuse OTEL logger works with ResponsesAPI responses and logs metadata."""
|
||||
# Create a mock ResponsesAPIResponse
|
||||
mock_response = ResponsesAPIResponse(
|
||||
id="response-123",
|
||||
created_at=1234567890,
|
||||
output=[
|
||||
{
|
||||
"type": "message",
|
||||
"content": [{"type": "text", "text": "Hello from responses API"}]
|
||||
}
|
||||
],
|
||||
parallel_tool_calls=False,
|
||||
tool_choice="auto",
|
||||
tools=[],
|
||||
top_p=1.0
|
||||
)
|
||||
|
||||
# Create kwargs with metadata that should be logged
|
||||
test_metadata = {
|
||||
"user_id": "test123",
|
||||
"session_id": "abc456",
|
||||
"custom_field": "test_value",
|
||||
"generation_name": "responses_test_generation",
|
||||
"trace_name": "responses_api_trace"
|
||||
}
|
||||
|
||||
kwargs = {
|
||||
"call_type": "responses",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"model": "gpt-4o",
|
||||
"optional_params": {},
|
||||
"litellm_params": {"metadata": test_metadata}
|
||||
}
|
||||
|
||||
mock_span = MagicMock()
|
||||
|
||||
with patch('litellm.integrations.arize._utils.set_attributes') as mock_set_attributes:
|
||||
with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute:
|
||||
logger = LangfuseOtelLogger()
|
||||
logger.set_langfuse_otel_attributes(mock_span, kwargs, mock_response)
|
||||
|
||||
# Verify that set_attributes was called for general attributes
|
||||
mock_set_attributes.assert_called_once_with(mock_span, kwargs, mock_response)
|
||||
|
||||
# Verify that Langfuse-specific attributes were set
|
||||
mock_safe_set_attribute.assert_any_call(
|
||||
mock_span, "langfuse.generation.name", "responses_test_generation"
|
||||
)
|
||||
mock_safe_set_attribute.assert_any_call(
|
||||
mock_span, "langfuse.trace.name", "responses_api_trace"
|
||||
)
|
||||
|
||||
def test_responses_api_metadata_extraction(self):
|
||||
"""Test that metadata is correctly extracted from ResponsesAPI kwargs."""
|
||||
# Clean up any existing module mocks
|
||||
import sys
|
||||
if "litellm.integrations.langfuse.langfuse" in sys.modules:
|
||||
original_module = sys.modules["litellm.integrations.langfuse.langfuse"]
|
||||
|
||||
test_metadata = {
|
||||
"user_id": "responses_user_123",
|
||||
"session_id": "responses_session_456",
|
||||
"custom_metadata": {"key": "value"},
|
||||
"generation_name": "responses_generation",
|
||||
"trace_id": "custom_trace_id"
|
||||
}
|
||||
|
||||
kwargs = {
|
||||
"call_type": "responses",
|
||||
"model": "gpt-4o",
|
||||
"litellm_params": {"metadata": test_metadata}
|
||||
}
|
||||
|
||||
extracted_metadata = LangfuseOtelLogger._extract_langfuse_metadata(kwargs)
|
||||
|
||||
# Verify all expected metadata was extracted (may have additional fields from header enrichment)
|
||||
for key, value in test_metadata.items():
|
||||
assert extracted_metadata[key] == value
|
||||
|
||||
assert extracted_metadata["user_id"] == "responses_user_123"
|
||||
assert extracted_metadata["generation_name"] == "responses_generation"
|
||||
assert extracted_metadata["trace_id"] == "custom_trace_id"
|
||||
|
||||
def test_responses_api_langfuse_specific_attributes(self):
|
||||
"""Test that ResponsesAPI metadata maps correctly to Langfuse OTEL attributes."""
|
||||
metadata = {
|
||||
"generation_name": "responses_gen",
|
||||
"generation_id": "resp_gen_123",
|
||||
"trace_name": "responses_trace",
|
||||
"trace_user_id": "resp_user_456",
|
||||
"session_id": "resp_session_789",
|
||||
"tags": ["responses", "api", "test"],
|
||||
"trace_metadata": {"source": "responses_api", "version": "1.0"}
|
||||
}
|
||||
|
||||
kwargs = {
|
||||
"call_type": "responses",
|
||||
"litellm_params": {"metadata": metadata}
|
||||
}
|
||||
|
||||
mock_span = MagicMock()
|
||||
|
||||
with patch('litellm.integrations.arize._utils.safe_set_attribute') as mock_safe_set_attribute:
|
||||
LangfuseOtelLogger._set_langfuse_specific_attributes(mock_span, kwargs)
|
||||
|
||||
# Verify specific attributes were set
|
||||
from litellm.types.integrations.langfuse_otel import LangfuseSpanAttributes
|
||||
|
||||
expected_calls = [
|
||||
(mock_span, LangfuseSpanAttributes.GENERATION_NAME.value, "responses_gen"),
|
||||
(mock_span, LangfuseSpanAttributes.GENERATION_ID.value, "resp_gen_123"),
|
||||
(mock_span, LangfuseSpanAttributes.TRACE_NAME.value, "responses_trace"),
|
||||
(mock_span, LangfuseSpanAttributes.TRACE_USER_ID.value, "resp_user_456"),
|
||||
(mock_span, LangfuseSpanAttributes.SESSION_ID.value, "resp_session_789"),
|
||||
(mock_span, LangfuseSpanAttributes.TAGS.value, json.dumps(["responses", "api", "test"])),
|
||||
(mock_span, LangfuseSpanAttributes.TRACE_METADATA.value,
|
||||
json.dumps({"source": "responses_api", "version": "1.0"}))
|
||||
]
|
||||
|
||||
for expected_call in expected_calls:
|
||||
mock_safe_set_attribute.assert_any_call(*expected_call)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -9,7 +9,11 @@ sys.path.insert(
|
||||
0, os.path.abspath("../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs, safe_divide
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
get_litellm_metadata_from_kwargs,
|
||||
safe_divide,
|
||||
safe_deep_copy
|
||||
)
|
||||
|
||||
|
||||
def test_get_litellm_metadata_from_kwargs():
|
||||
@@ -127,3 +131,43 @@ def test_safe_divide_weight_scenario():
|
||||
expected_zero = [0, 0, 0]
|
||||
|
||||
assert normalized_zero_weights == expected_zero, f"Expected {expected_zero}, got {normalized_zero_weights}"
|
||||
|
||||
|
||||
def test_safe_deep_copy_with_non_pickleables_and_span():
|
||||
"""
|
||||
Verify safe_deep_copy:
|
||||
- does not crash when non-pickleables are present,
|
||||
- preserves structure/keys,
|
||||
- deep-copies JSON-y payloads (e.g., messages),
|
||||
- keeps non-pickleables by reference,
|
||||
- redacts OTEL span in the copy and restores it in the original.
|
||||
"""
|
||||
import threading
|
||||
rlock = threading.RLock()
|
||||
data = {
|
||||
"metadata": {"litellm_parent_otel_span": rlock, "x": 1},
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"optional_params": {"handle": rlock},
|
||||
"ok": True,
|
||||
}
|
||||
|
||||
copied = safe_deep_copy(data)
|
||||
|
||||
# Structure preserved
|
||||
assert set(copied.keys()) == set(data.keys())
|
||||
|
||||
# Messages are deep-copied (new object, same content)
|
||||
assert copied["messages"] is not data["messages"]
|
||||
assert copied["messages"][0] == data["messages"][0]
|
||||
|
||||
# Non-pickleable subtree kept by reference (no crash)
|
||||
assert copied["optional_params"] is data["optional_params"]
|
||||
assert copied["optional_params"]["handle"] is rlock
|
||||
|
||||
# OTEL span: redacted in the copy, restored in original
|
||||
assert copied["metadata"]["litellm_parent_otel_span"] == "placeholder"
|
||||
assert data["metadata"]["litellm_parent_otel_span"] is rlock
|
||||
|
||||
# Other simple fields unchanged
|
||||
assert copied["ok"] is True
|
||||
assert copied["metadata"]["x"] == 1
|
||||
|
||||
@@ -438,6 +438,97 @@ async def test_e2e_generate_cold_storage_object_key_successful():
|
||||
assert isinstance(result, str)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path():
|
||||
"""
|
||||
Test that _generate_cold_storage_object_key uses s3_path from custom logger instance.
|
||||
"""
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
|
||||
# Create test data
|
||||
start_time = datetime(2025, 1, 15, 10, 30, 45, 123456, timezone.utc)
|
||||
response_id = "chatcmpl-test-12345"
|
||||
|
||||
# Create mock custom logger with s3_path
|
||||
mock_custom_logger = MagicMock()
|
||||
mock_custom_logger.s3_path = "storage"
|
||||
|
||||
with patch("litellm.configured_cold_storage_logger", "s3_v2"), \
|
||||
patch("litellm.logging_callback_manager.get_active_custom_logger_for_callback_name") as mock_get_logger, \
|
||||
patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key:
|
||||
|
||||
# Setup mocks
|
||||
mock_get_logger.return_value = mock_custom_logger
|
||||
mock_get_s3_key.return_value = "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
|
||||
|
||||
# Call the function
|
||||
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
|
||||
start_time=start_time,
|
||||
response_id=response_id
|
||||
)
|
||||
|
||||
# Verify logger was queried correctly
|
||||
mock_get_logger.assert_called_once_with("s3_v2")
|
||||
|
||||
# Verify the S3 function was called with the custom logger's s3_path
|
||||
mock_get_s3_key.assert_called_once_with(
|
||||
s3_path="storage", # Should use custom logger's s3_path
|
||||
team_alias_prefix="",
|
||||
start_time=start_time,
|
||||
s3_file_name="time-10-30-45-123456_chatcmpl-test-12345"
|
||||
)
|
||||
|
||||
# Verify the result
|
||||
assert result == "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path():
|
||||
"""
|
||||
Test that _generate_cold_storage_object_key falls back to empty s3_path when logger has no s3_path.
|
||||
"""
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
|
||||
|
||||
# Create test data
|
||||
start_time = datetime(2025, 1, 15, 10, 30, 45, 123456, timezone.utc)
|
||||
response_id = "chatcmpl-test-12345"
|
||||
|
||||
# Create mock custom logger without s3_path
|
||||
mock_custom_logger = MagicMock()
|
||||
mock_custom_logger.s3_path = None # or could be missing attribute
|
||||
|
||||
with patch("litellm.configured_cold_storage_logger", "s3_v2"), \
|
||||
patch("litellm.logging_callback_manager.get_active_custom_logger_for_callback_name") as mock_get_logger, \
|
||||
patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key:
|
||||
|
||||
# Setup mocks
|
||||
mock_get_logger.return_value = mock_custom_logger
|
||||
mock_get_s3_key.return_value = "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
|
||||
|
||||
# Call the function
|
||||
result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
|
||||
start_time=start_time,
|
||||
response_id=response_id
|
||||
)
|
||||
|
||||
# Verify the S3 function was called with empty s3_path (fallback)
|
||||
mock_get_s3_key.assert_called_once_with(
|
||||
s3_path="", # Should fall back to empty string
|
||||
team_alias_prefix="",
|
||||
start_time=start_time,
|
||||
s3_file_name="time-10-30-45-123456_chatcmpl-test-12345"
|
||||
)
|
||||
|
||||
# Verify the result
|
||||
assert result == "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_e2e_generate_cold_storage_object_key_not_configured():
|
||||
"""
|
||||
|
||||
@@ -1589,4 +1589,282 @@ async def test_no_cache_control_no_cache_point():
|
||||
# Tool message should only have tool result, no cachePoint
|
||||
tool_content = result[2]["content"]
|
||||
assert len(tool_content) == 1
|
||||
assert "toolResult" in tool_content[0]
|
||||
assert "toolResult" in tool_content[0]
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Guarded Text Feature Tests
|
||||
# ============================================================================
|
||||
|
||||
def test_guarded_text_wraps_in_guardrail_converse_content():
|
||||
"""Test that guarded_text content type gets wrapped in guardrailConverseContent blocks."""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Regular text content"},
|
||||
{"type": "guarded_text", "text": "This should be guarded"},
|
||||
{"type": "text", "text": "More regular text"}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
result = _bedrock_converse_messages_pt(
|
||||
messages=messages,
|
||||
model="us.amazon.nova-pro-v1:0",
|
||||
llm_provider="bedrock_converse"
|
||||
)
|
||||
|
||||
# Should have 1 message
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "user"
|
||||
|
||||
# Should have 3 content blocks
|
||||
content = result[0]["content"]
|
||||
assert len(content) == 3
|
||||
|
||||
# First and third should be regular text
|
||||
assert "text" in content[0]
|
||||
assert content[0]["text"] == "Regular text content"
|
||||
assert "text" in content[2]
|
||||
assert content[2]["text"] == "More regular text"
|
||||
|
||||
# Second should be guardrailConverseContent
|
||||
assert "guardrailConverseContent" in content[1]
|
||||
assert content[1]["guardrailConverseContent"]["text"] == "This should be guarded"
|
||||
|
||||
|
||||
def test_guarded_text_with_system_messages():
|
||||
"""Test guarded_text with system messages using the full transformation."""
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What is the main topic of this legal document?"},
|
||||
{"type": "guarded_text", "text": "This is a set of very long instructions that you will follow. Here is a legal document that you will use to answer the user's question."}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
optional_params = {
|
||||
"guardrailConfig": {
|
||||
"guardrailIdentifier": "gr-abc123",
|
||||
"guardrailVersion": "DRAFT"
|
||||
}
|
||||
}
|
||||
|
||||
result = config._transform_request(
|
||||
model="us.amazon.nova-pro-v1:0",
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
|
||||
# Should have system content blocks
|
||||
assert "system" in result
|
||||
assert len(result["system"]) == 1
|
||||
assert result["system"][0]["text"] == "You are a helpful assistant."
|
||||
|
||||
# Should have 1 message (system messages are removed)
|
||||
assert "messages" in result
|
||||
assert len(result["messages"]) == 1
|
||||
|
||||
# User message should have both regular text and guarded text
|
||||
user_message = result["messages"][0]
|
||||
assert user_message["role"] == "user"
|
||||
content = user_message["content"]
|
||||
assert len(content) == 2
|
||||
|
||||
# First should be regular text
|
||||
assert "text" in content[0]
|
||||
assert content[0]["text"] == "What is the main topic of this legal document?"
|
||||
|
||||
# Second should be guardrailConverseContent
|
||||
assert "guardrailConverseContent" in content[1]
|
||||
assert content[1]["guardrailConverseContent"]["text"] == "This is a set of very long instructions that you will follow. Here is a legal document that you will use to answer the user's question."
|
||||
|
||||
|
||||
def test_guarded_text_with_mixed_content_types():
|
||||
"""Test guarded_text with mixed content types including images."""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Look at this image"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,test"}},
|
||||
{"type": "guarded_text", "text": "This sensitive content should be guarded"}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
result = _bedrock_converse_messages_pt(
|
||||
messages=messages,
|
||||
model="us.amazon.nova-pro-v1:0",
|
||||
llm_provider="bedrock_converse"
|
||||
)
|
||||
|
||||
# Should have 1 message
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "user"
|
||||
|
||||
# Should have 3 content blocks
|
||||
content = result[0]["content"]
|
||||
assert len(content) == 3
|
||||
|
||||
# First should be regular text
|
||||
assert "text" in content[0]
|
||||
assert content[0]["text"] == "Look at this image"
|
||||
|
||||
# Second should be image
|
||||
assert "image" in content[1]
|
||||
|
||||
# Third should be guardrailConverseContent
|
||||
assert "guardrailConverseContent" in content[2]
|
||||
assert content[2]["guardrailConverseContent"]["text"] == "This sensitive content should be guarded"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_guarded_text():
|
||||
"""Test async version of guarded_text processing."""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import BedrockConverseMessagesProcessor
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Hello"},
|
||||
{"type": "guarded_text", "text": "This should be guarded"}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async(
|
||||
messages=messages,
|
||||
model="us.amazon.nova-pro-v1:0",
|
||||
llm_provider="bedrock_converse"
|
||||
)
|
||||
|
||||
# Should have 1 message
|
||||
assert len(result) == 1
|
||||
assert result[0]["role"] == "user"
|
||||
|
||||
# Should have 2 content blocks
|
||||
content = result[0]["content"]
|
||||
assert len(content) == 2
|
||||
|
||||
# First should be regular text
|
||||
assert "text" in content[0]
|
||||
assert content[0]["text"] == "Hello"
|
||||
|
||||
# Second should be guardrailConverseContent
|
||||
assert "guardrailConverseContent" in content[1]
|
||||
assert content[1]["guardrailConverseContent"]["text"] == "This should be guarded"
|
||||
|
||||
|
||||
def test_guarded_text_with_tool_calls():
|
||||
"""Test guarded_text with tool calls in the conversation."""
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import _bedrock_converse_messages_pt
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "What's the weather?"},
|
||||
{"type": "guarded_text", "text": "Please be careful with sensitive information"}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_123",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": "{}"}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_123",
|
||||
"content": "It's sunny and 25°C"
|
||||
}
|
||||
]
|
||||
|
||||
result = _bedrock_converse_messages_pt(
|
||||
messages=messages,
|
||||
model="us.amazon.nova-pro-v1:0",
|
||||
llm_provider="bedrock_converse"
|
||||
)
|
||||
|
||||
# Should have 3 messages
|
||||
assert len(result) == 3
|
||||
|
||||
# First message (user) should have both text and guarded_text
|
||||
user_message = result[0]
|
||||
assert user_message["role"] == "user"
|
||||
content = user_message["content"]
|
||||
assert len(content) == 2
|
||||
|
||||
# First should be regular text
|
||||
assert "text" in content[0]
|
||||
assert content[0]["text"] == "What's the weather?"
|
||||
|
||||
# Second should be guardrailConverseContent
|
||||
assert "guardrailConverseContent" in content[1]
|
||||
assert content[1]["guardrailConverseContent"]["text"] == "Please be careful with sensitive information"
|
||||
|
||||
# Other messages should not have guardrailConverseContent
|
||||
for i in range(1, 3):
|
||||
content = result[i]["content"]
|
||||
for block in content:
|
||||
assert "guardrailConverseContent" not in block
|
||||
|
||||
|
||||
def test_guarded_text_guardrail_config_preserved():
|
||||
"""Test that guardrailConfig is preserved when using guarded_text."""
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Hello"},
|
||||
{"type": "guarded_text", "text": "This should be guarded"}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
optional_params = {
|
||||
"guardrailConfig": {
|
||||
"guardrailIdentifier": "gr-abc123",
|
||||
"guardrailVersion": "DRAFT"
|
||||
}
|
||||
}
|
||||
|
||||
result = config._transform_request(
|
||||
model="us.amazon.nova-pro-v1:0",
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers={}
|
||||
)
|
||||
|
||||
# GuardrailConfig should be present at top level
|
||||
assert "guardrailConfig" in result
|
||||
assert result["guardrailConfig"]["guardrailIdentifier"] == "gr-abc123"
|
||||
|
||||
# GuardrailConfig should also be in inferenceConfig
|
||||
assert "inferenceConfig" in result
|
||||
assert "guardrailConfig" in result["inferenceConfig"]
|
||||
assert result["inferenceConfig"]["guardrailConfig"]["guardrailIdentifier"] == "gr-abc123"
|
||||
|
||||
|
||||
|
||||
@@ -704,3 +704,177 @@ class TestVertexBase:
|
||||
vertex_base.get_api_base(api_base=api_base, vertex_location=vertex_location)
|
||||
== expected
|
||||
), f"Expected {expected} with api_base {api_base} and vertex_location {vertex_location}"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"api_base, custom_llm_provider, gemini_api_key, endpoint, stream, auth_header, url, model, expected_auth_header, expected_url",
|
||||
[
|
||||
# Test case 1: Gemini with custom API base
|
||||
(
|
||||
"https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta",
|
||||
"gemini",
|
||||
"test-api-key",
|
||||
"generateContent",
|
||||
False,
|
||||
None,
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent",
|
||||
"gemini-2.5-flash-lite",
|
||||
"test-api-key",
|
||||
"https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent"
|
||||
),
|
||||
# Test case 2: Gemini with custom API base and streaming
|
||||
(
|
||||
"https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta",
|
||||
"gemini",
|
||||
"test-api-key",
|
||||
"generateContent",
|
||||
True,
|
||||
None,
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent",
|
||||
"gemini-2.5-flash-lite",
|
||||
"test-api-key",
|
||||
"https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse"
|
||||
),
|
||||
# Test case 3: Non-Gemini provider with custom API base
|
||||
(
|
||||
"https://custom-vertex-api.com",
|
||||
"vertex_ai",
|
||||
None,
|
||||
"generateContent",
|
||||
False,
|
||||
"Bearer token123",
|
||||
"https://aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-pro:generateContent",
|
||||
"gemini-pro",
|
||||
"Bearer token123",
|
||||
"https://custom-vertex-api.com:generateContent"
|
||||
),
|
||||
# Test case 4: No API base provided (should return original values)
|
||||
(
|
||||
None,
|
||||
"gemini",
|
||||
"test-api-key",
|
||||
"generateContent",
|
||||
False,
|
||||
"Bearer token123",
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent",
|
||||
"gemini-2.5-flash-lite",
|
||||
"Bearer token123",
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent"
|
||||
),
|
||||
# Test case 5: Gemini without API key (should raise ValueError)
|
||||
(
|
||||
"https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta",
|
||||
"gemini",
|
||||
None,
|
||||
"generateContent",
|
||||
False,
|
||||
None,
|
||||
"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent",
|
||||
"gemini-2.5-flash-lite",
|
||||
None, # This should raise an exception
|
||||
None
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_check_custom_proxy(
|
||||
self,
|
||||
api_base,
|
||||
custom_llm_provider,
|
||||
gemini_api_key,
|
||||
endpoint,
|
||||
stream,
|
||||
auth_header,
|
||||
url,
|
||||
model,
|
||||
expected_auth_header,
|
||||
expected_url
|
||||
):
|
||||
"""Test the _check_custom_proxy method for handling custom API base URLs"""
|
||||
vertex_base = VertexBase()
|
||||
|
||||
if custom_llm_provider == "gemini" and api_base and gemini_api_key is None:
|
||||
# Test case 5: Should raise ValueError for Gemini without API key
|
||||
with pytest.raises(ValueError, match="Missing gemini_api_key"):
|
||||
vertex_base._check_custom_proxy(
|
||||
api_base=api_base,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
gemini_api_key=gemini_api_key,
|
||||
endpoint=endpoint,
|
||||
stream=stream,
|
||||
auth_header=auth_header,
|
||||
url=url,
|
||||
model=model,
|
||||
)
|
||||
else:
|
||||
# Test cases 1-4: Should work correctly
|
||||
result_auth_header, result_url = vertex_base._check_custom_proxy(
|
||||
api_base=api_base,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
gemini_api_key=gemini_api_key,
|
||||
endpoint=endpoint,
|
||||
stream=stream,
|
||||
auth_header=auth_header,
|
||||
url=url,
|
||||
model=model,
|
||||
)
|
||||
|
||||
assert result_auth_header == expected_auth_header, f"Expected auth_header {expected_auth_header}, got {result_auth_header}"
|
||||
assert result_url == expected_url, f"Expected URL {expected_url}, got {result_url}"
|
||||
|
||||
def test_check_custom_proxy_gemini_url_construction(self):
|
||||
"""Test that Gemini URLs are constructed correctly with custom API base"""
|
||||
vertex_base = VertexBase()
|
||||
|
||||
# Test various Gemini models with custom API base
|
||||
test_cases = [
|
||||
("gemini-2.5-flash-lite", "generateContent", "https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent"),
|
||||
("gemini-2.5-pro", "generateContent", "https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-pro:generateContent"),
|
||||
("gemini-1.5-flash", "streamGenerateContent", "https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:streamGenerateContent"),
|
||||
]
|
||||
|
||||
for model, endpoint, expected_url in test_cases:
|
||||
_, result_url = vertex_base._check_custom_proxy(
|
||||
api_base="https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta",
|
||||
custom_llm_provider="gemini",
|
||||
gemini_api_key="test-api-key",
|
||||
endpoint=endpoint,
|
||||
stream=False,
|
||||
auth_header=None,
|
||||
url=f"https://generativelanguage.googleapis.com/v1beta/models/{model}:{endpoint}",
|
||||
model=model,
|
||||
)
|
||||
|
||||
assert result_url == expected_url, f"Expected {expected_url}, got {result_url} for model {model}"
|
||||
|
||||
def test_check_custom_proxy_streaming_parameter(self):
|
||||
"""Test that streaming parameter correctly adds ?alt=sse to URLs"""
|
||||
vertex_base = VertexBase()
|
||||
|
||||
# Test with streaming enabled
|
||||
_, result_url_streaming = vertex_base._check_custom_proxy(
|
||||
api_base="https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta",
|
||||
custom_llm_provider="gemini",
|
||||
gemini_api_key="test-api-key",
|
||||
endpoint="generateContent",
|
||||
stream=True,
|
||||
auth_header=None,
|
||||
url="https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent",
|
||||
model="gemini-2.5-flash-lite",
|
||||
)
|
||||
|
||||
expected_streaming_url = "https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent?alt=sse"
|
||||
assert result_url_streaming == expected_streaming_url, f"Expected {expected_streaming_url}, got {result_url_streaming}"
|
||||
|
||||
# Test with streaming disabled
|
||||
_, result_url_no_streaming = vertex_base._check_custom_proxy(
|
||||
api_base="https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta",
|
||||
custom_llm_provider="gemini",
|
||||
gemini_api_key="test-api-key",
|
||||
endpoint="generateContent",
|
||||
stream=False,
|
||||
auth_header=None,
|
||||
url="https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent",
|
||||
model="gemini-2.5-flash-lite",
|
||||
)
|
||||
|
||||
expected_no_streaming_url = "https://proxy.zapier.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent"
|
||||
assert result_url_no_streaming == expected_no_streaming_url, f"Expected {expected_no_streaming_url}, got {result_url_no_streaming}"
|
||||
|
||||
+50
@@ -364,3 +364,53 @@ async def test_should_check_cold_storage_for_full_payload():
|
||||
with patch.object(litellm, 'configured_cold_storage_logger', None):
|
||||
result5 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf)
|
||||
assert result5 == False, "Should return False when cold storage is not configured, even with truncated content"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_chat_completion_message_history_empty_response_dict():
|
||||
"""
|
||||
Test that empty response dict is handled correctly without processing.
|
||||
This tests the fix for response validation to check for empty dict responses.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
# Mock spend logs with empty response dict
|
||||
mock_spend_logs = [
|
||||
{
|
||||
"request_id": "chatcmpl-test-empty-response",
|
||||
"call_type": "aresponses",
|
||||
"api_key": "test_key",
|
||||
"spend": 0.001,
|
||||
"total_tokens": 0,
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 0,
|
||||
"startTime": "2025-01-15T10:30:00.000+00:00",
|
||||
"endTime": "2025-01-15T10:30:01.000+00:00",
|
||||
"model": "gpt-4",
|
||||
"session_id": "test-session",
|
||||
"proxy_server_request": {
|
||||
"input": "test input",
|
||||
"model": "gpt-4"
|
||||
},
|
||||
"response": {} # Empty dict - should not be processed
|
||||
}
|
||||
]
|
||||
|
||||
with patch.object(ResponsesSessionHandler, "get_all_spend_logs_for_previous_response_id") as mock_get_spend_logs:
|
||||
mock_get_spend_logs.return_value = mock_spend_logs
|
||||
|
||||
# Call the function
|
||||
result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
|
||||
"chatcmpl-test-empty-response"
|
||||
)
|
||||
|
||||
# Verify that user message was added but no assistant response
|
||||
# Since response is empty dict, no assistant response should be processed
|
||||
# But user input from proxy_server_request should still be included
|
||||
messages = result["messages"]
|
||||
assert len(messages) == 1 # Only user message, no assistant response
|
||||
assert messages[0]["role"] == "user"
|
||||
assert messages[0]["content"] == "test input"
|
||||
|
||||
# Verify the session was still created correctly
|
||||
assert result["litellm_session_id"] == "test-session"
|
||||
|
||||
@@ -508,6 +508,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
|
||||
"supports_computer_use": {"type": "boolean"},
|
||||
"cache_creation_input_audio_token_cost": {"type": "number"},
|
||||
"cache_creation_input_token_cost": {"type": "number"},
|
||||
"cache_creation_input_token_cost_above_1hr": {"type": "number"},
|
||||
"cache_creation_input_token_cost_above_200k_tokens": {"type": "number"},
|
||||
"cache_read_input_token_cost": {"type": "number"},
|
||||
"cache_read_input_token_cost_above_200k_tokens": {"type": "number"},
|
||||
@@ -661,16 +662,16 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
|
||||
"type": "array",
|
||||
"items": {"type": "number"},
|
||||
"minItems": 2,
|
||||
"maxItems": 2
|
||||
"maxItems": 2,
|
||||
},
|
||||
"input_cost_per_token": {"type": "number"},
|
||||
"output_cost_per_token": {"type": "number"},
|
||||
"cache_read_input_token_cost": {"type": "number"},
|
||||
"output_cost_per_reasoning_token": {"type": "number"}
|
||||
"output_cost_per_reasoning_token": {"type": "number"},
|
||||
},
|
||||
"required": ["range"],
|
||||
"additionalProperties": False
|
||||
}
|
||||
"additionalProperties": False,
|
||||
},
|
||||
},
|
||||
},
|
||||
"additionalProperties": False,
|
||||
@@ -843,7 +844,6 @@ for commitment in BEDROCK_COMMITMENTS:
|
||||
print("block_list", block_list)
|
||||
|
||||
|
||||
|
||||
def test_supports_computer_use_utility():
|
||||
"""
|
||||
Tests the litellm.utils.supports_computer_use utility function.
|
||||
@@ -925,8 +925,7 @@ def test_pre_process_non_default_params(model, custom_llm_provider):
|
||||
from litellm.utils import ProviderConfigManager, pre_process_non_default_params
|
||||
|
||||
provider_config = ProviderConfigManager.get_provider_chat_config(
|
||||
model=model,
|
||||
provider=LlmProviders(custom_llm_provider)
|
||||
model=model, provider=LlmProviders(custom_llm_provider)
|
||||
)
|
||||
|
||||
class ResponseFormat(BaseModel):
|
||||
|
||||
@@ -857,10 +857,6 @@ const CreateKey: React.FC<CreateKeyProps> = ({
|
||||
className="mt-4"
|
||||
help={premiumUser ? "Select existing guardrails or enter new ones" : "Premium feature - Upgrade to set guardrails by key"}
|
||||
>
|
||||
<Tooltip
|
||||
title={!premiumUser ? "Setting guardrails by key is a premium feature" : ""}
|
||||
placement="top"
|
||||
>
|
||||
<Select
|
||||
mode="tags"
|
||||
style={{ width: '100%' }}
|
||||
@@ -872,7 +868,6 @@ const CreateKey: React.FC<CreateKeyProps> = ({
|
||||
}
|
||||
options={guardrailsList.map(name => ({ value: name, label: name }))}
|
||||
/>
|
||||
</Tooltip>
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
@@ -894,10 +889,6 @@ const CreateKey: React.FC<CreateKeyProps> = ({
|
||||
className="mt-4"
|
||||
help={premiumUser ? "Select existing prompts or enter new ones" : "Premium feature - Upgrade to set prompts by key"}
|
||||
>
|
||||
<Tooltip
|
||||
title={!premiumUser ? "Setting prompts by key is a premium feature" : ""}
|
||||
placement="top"
|
||||
>
|
||||
<Select
|
||||
mode="tags"
|
||||
style={{ width: '100%' }}
|
||||
@@ -909,7 +900,6 @@ const CreateKey: React.FC<CreateKeyProps> = ({
|
||||
}
|
||||
options={promptsList.map(name => ({ value: name, label: name }))}
|
||||
/>
|
||||
</Tooltip>
|
||||
</Form.Item>
|
||||
<Form.Item
|
||||
label={
|
||||
|
||||
@@ -11,6 +11,7 @@ import EditLoggingSettings from "../team/EditLoggingSettings"
|
||||
import { extractLoggingSettings, formatMetadataForDisplay } from "../key_info_utils"
|
||||
import { fetchMCPAccessGroups } from "../networking"
|
||||
import { mapInternalToDisplayNames, mapDisplayToInternalNames } from "../callback_info_helpers"
|
||||
import GuardrailSelector from "@/components/guardrails/GuardrailSelector"
|
||||
|
||||
interface KeyEditViewProps {
|
||||
keyData: KeyResponse
|
||||
@@ -220,20 +221,9 @@ export function KeyEditView({
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item label="Guardrails" name="guardrails">
|
||||
<Tooltip title={!premiumUser ? "Setting guardrails by key is a premium feature" : ""} placement="top">
|
||||
<Select
|
||||
mode="tags"
|
||||
style={{ width: "100%" }}
|
||||
disabled={!premiumUser}
|
||||
placeholder={
|
||||
!premiumUser
|
||||
? "Premium feature - Upgrade to set guardrails by key"
|
||||
: Array.isArray(keyData.metadata?.guardrails) && keyData.metadata.guardrails.length > 0
|
||||
? `Current: ${keyData.metadata.guardrails.join(", ")}`
|
||||
: "Select or enter guardrails"
|
||||
}
|
||||
/>
|
||||
</Tooltip>
|
||||
{ accessToken &&
|
||||
<GuardrailSelector onChange={(v) => {form.setFieldValue("guardrails", v)}} accessToken={accessToken} />
|
||||
}
|
||||
</Form.Item>
|
||||
|
||||
<Form.Item label="Prompts" name="prompts">
|
||||
|
||||
Reference in New Issue
Block a user