Merge pull request #23440 from BerriAI/litellm_oss_staging_03_11_2026

Litellm oss staging 03 11 2026
This commit is contained in:
Cesar Garcia
2026-03-12 16:24:14 -03:00
committed by GitHub
955 changed files with 23392 additions and 20598 deletions
+2 -2
View File
@@ -33,10 +33,10 @@ jobs:
poetry lock
poetry install --with dev
- name: Run Black formatting
- name: Check Black formatting
run: |
cd litellm
poetry run black .
poetry run black --check --exclude '/enterprise/' .
cd ..
- name: Debug - Check file state
@@ -20,6 +20,9 @@ spec:
selector:
matchLabels:
{{- include "litellm.selectorLabels" . | nindent 6 }}
{{- if .Values.deploymentMinReadySeconds }}
minReadySeconds: {{ .Values.deploymentMinReadySeconds }}
{{- end }}
template:
metadata:
annotations:
@@ -306,3 +306,16 @@ tests:
- equal:
path: spec.template.spec.containers[0].resources
value: {}
- it: should be able to set minReadySeconds
template: deployment.yaml
set:
deploymentMinReadySeconds: 5
asserts:
- equal:
path: spec.minReadySeconds
value: 5
- it: should have minReadySeconds absent when deploymentMinReadySeconds is not set
template: deployment.yaml
asserts:
- notExists:
path: spec.minReadySeconds
+2
View File
@@ -31,6 +31,8 @@ serviceAccount:
# annotations for litellm deployment
deploymentAnnotations: {}
deploymentLabels: {}
deploymentMinReadySeconds: 0
# annotations for litellm pods
podAnnotations: {}
podLabels: {}
+6
View File
@@ -326,4 +326,10 @@ print("file content=", content.text)
### [Bedrock](./providers/bedrock_batches#4-retrieve-batch-results)
### [Anthropic](./providers/anthropic#files-api)
:::note
Anthropic Files API has a different purpose than OpenAI's. It's **not** for Batches or Fine-tuning—it's for uploading files once and referencing them by `file_id` in multiple messages, avoiding re-uploads. File API operations are free — file content used in Messages requests is priced as input tokens.
:::
## [Swagger API Reference](https://litellm-api.up.railway.app/#/files)
@@ -1965,6 +1965,98 @@ curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
</TabItem>
</Tabs>
## Files API
Upload files once and reference them by `file_id` in multiple requests—no need to re-upload content each time.
:::info
The `file_id` obtained from Anthropic only works with Anthropic Claude models. You cannot use it with other providers (OpenAI, Bedrock, etc.).
:::
- **Max file size:** 500 MB | **Total storage:** 100 GB per org
- **Pricing:** File API operations are free. File content used in Messages requests is priced as input tokens.
**Supported models by file type:**
- **Images:** All Claude 3+ models
- **PDFs:** All Claude 3.5+ models
- **Other file types** (for code execution): Claude 3.5 Haiku + all Claude 3.7+ models
### Quick Start
```python
import litellm
import os
os.environ["ANTHROPIC_API_KEY"] = "sk-ant-..."
# 1. Upload a file once
file = litellm.create_file(
file=open("document.pdf", "rb"),
purpose="messages",
custom_llm_provider="anthropic",
)
# 2. Use file_id in messages (no re-upload needed)
response = litellm.completion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[{
"role": "user",
"content": [
{"type": "text", "text": "Summarize this document"},
{"type": "file", "file": {"file_id": file.id, "format": "application/pdf"}}
]
}]
)
```
### File Operations
| Operation | Function |
|-----------|----------|
| Upload | `litellm.create_file(file, purpose="messages", custom_llm_provider="anthropic")` |
| List | `litellm.file_list(custom_llm_provider="anthropic")` |
| Retrieve | `litellm.file_retrieve(file_id, custom_llm_provider="anthropic")` |
| Delete | `litellm.file_delete(file_id, custom_llm_provider="anthropic")` |
| Download | `litellm.file_content(file_id, custom_llm_provider="anthropic")` |
:::note
Download only works for files created by the [code execution tool](https://docs.anthropic.com/en/docs/agents-and-tools/tool-use/code-execution-tool), not uploaded files.
:::
### Supported Formats
| File Type | Format Value |
|-----------|-------------|
| PDF | `application/pdf` |
| Plain text | `text/plain` |
| JPEG | `image/jpeg` |
| PNG | `image/png` |
| GIF | `image/gif` |
| WebP | `image/webp` |
### Using Images
```python
# Upload image
image = litellm.create_file(
file=open("photo.jpg", "rb"),
purpose="messages",
custom_llm_provider="anthropic",
)
# Use in message
response = litellm.completion(
model="anthropic/claude-sonnet-4-5-20250929",
messages=[{
"role": "user",
"content": [
{"type": "text", "text": "What's in this image?"},
{"type": "file", "file": {"file_id": image.id, "format": "image/jpeg"}}
]
}]
)
```
## Usage - passing 'user_id' to Anthropic
LiteLLM translates the OpenAI `user` param to Anthropic's `metadata[user_id]` param.
+3 -1
View File
@@ -638,7 +638,9 @@ This is useful when you want to use [Responses API](https://platform.openai.com/
:::tip gpt-5.4 + reasoning_effort + function tools
OpenAI does not support `reasoning_effort` with function tools for `gpt-5.4` in `/v1/chat/completions`. Use the responses bridge instead:
LiteLLM drops `reasoning_effort` from `gpt-5.4` requests to `litellm.completion()` that include tools, since that combination is supported in the Responses API.
If you need reasoning **and** tools together, use the responses bridge instead:
```python
response = litellm.completion(
@@ -11,6 +11,7 @@ import TabItem from '@theme/TabItem';
|----------|---------------|---------------|
| Anthropic (Claude) | `vertex_ai/claude-*` | [Vertex AI - Anthropic Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude) |
| DeepSeek | `vertex_ai/deepseek-ai/{MODEL}` | [Vertex AI - DeepSeek Models](https://cloud.google.com/vertex-ai/generative-ai/docs/maas/deepseek) |
| ZAI (GLM) | `vertex_ai/zai-org/{MODEL}` | [Vertex AI - GLM Models](https://cloud.google.com/vertex-ai/generative-ai/docs/maas/zaiorg/glm-47) |
| Meta/Llama | `vertex_ai/meta/{MODEL}` | [Vertex AI - Meta Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/llama) |
| Mistral | `vertex_ai/mistral-*` | [Vertex AI - Mistral Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/mistral) |
| AI21 (Jamba) | `vertex_ai/jamba-*` | [Vertex AI - AI21 Models](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/ai21) |
@@ -226,6 +227,79 @@ ModelResponse(
|------------------|------------------------------|
| vertex_ai/deepseek-ai/deepseek-r1-0528-maas | `completion('vertex_ai/deepseek-ai/deepseek-r1-0528-maas', messages)` |
## VertexAI ZAI (GLM)
| Property | Details |
|----------|---------|
| Provider Route | `vertex_ai/zai-org/{MODEL}` |
| Vertex Documentation | [Vertex AI - GLM Models](https://cloud.google.com/vertex-ai/generative-ai/docs/maas/zaiorg/glm-47) |
**LiteLLM Supports all Vertex AI GLM Models.** Ensure you use the `vertex_ai/zai-org/` prefix for all Vertex AI GLM models.
| Model Name | Usage |
|------------|-------|
| vertex_ai/zai-org/glm-4.7-maas | `completion('vertex_ai/zai-org/glm-4.7-maas', messages)` |
#### Usage
<Tabs>
<TabItem value="sdk" label="SDK">
```python
from litellm import completion
import os
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = ""
response = completion(
model="vertex_ai/zai-org/glm-4.7-maas",
messages=[{"role": "user", "content": "hi"}],
vertex_project="your-vertex-project",
# vertex_location routes to "global"
)
print("\nModel Response", response)
```
</TabItem>
<TabItem value="proxy" label="Proxy">
**1. Add to config**
```yaml
model_list:
- model_name: glm-4.7
litellm_params:
model: vertex_ai/zai-org/glm-4.7-maas
vertex_project: "my-project"
# vertex_location routes to "global"
```
**2. Start proxy**
```bash
litellm --config /path/to/config.yaml
# RUNNING at http://0.0.0.0:4000
```
**3. Test it!**
```bash
curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"model": "glm-4.7",
"messages": [
{
"role": "user",
"content": "what llm are you"
}
],
}'
```
</TabItem>
</Tabs>
## VertexAI Meta/Llama API
+3 -1
View File
@@ -594,7 +594,9 @@ Expected Response
:::tip gpt-5.4: reasoning_effort + function tools
OpenAI does not support `reasoning_effort` with function tools for `gpt-5.4` in `/v1/chat/completions`. Use `openai/responses/gpt-5.4` to route through the Responses API instead. See [Responses API Bridge](/docs/providers/openai#openai-chat-completion-to-responses-api-bridge) for details.
LiteLLM drops `reasoning_effort` from `gpt-5.4` requests to `litellm.completion()` that include tools, since that combination is supported in the Responses API.
If you need reasoning **and** tools together, use `openai/responses/gpt-5.4` to route through the Responses API instead. See [Responses API Bridge](/docs/providers/openai#openai-chat-completion-to-responses-api-bridge) for details.
:::
+4 -4
View File
@@ -6,7 +6,7 @@
"": {
"dependencies": {
"@hono/node-server": "^1.10.1",
"hono": "^4.10.3"
"hono": "^4.12.7"
},
"devDependencies": {
"@types/node": "^20.11.17",
@@ -548,9 +548,9 @@
}
},
"node_modules/hono": {
"version": "4.10.6",
"resolved": "https://registry.npmjs.org/hono/-/hono-4.10.6.tgz",
"integrity": "sha512-BIdolzGpDO9MQ4nu3AUuDwHZZ+KViNm+EZ75Ae55eMXMqLVhDFqEMXxtUe9Qh8hjL+pIna/frs2j6Y2yD5Ua/g==",
"version": "4.12.7",
"resolved": "https://registry.npmjs.org/hono/-/hono-4.12.7.tgz",
"integrity": "sha512-jq9l1DM0zVIvsm3lv9Nw9nlJnMNPOcAtsbsgiUhWcFzPE99Gvo6yRTlszSLLYacMeQ6quHD6hMfId8crVHvexw==",
"license": "MIT",
"engines": {
"node": ">=16.9.0"
+1 -1
View File
@@ -4,7 +4,7 @@
},
"dependencies": {
"@hono/node-server": "^1.10.1",
"hono": "^4.10.3"
"hono": "^4.12.7"
},
"devDependencies": {
"@types/node": "^20.11.17",
+561 -218
View File
File diff suppressed because it is too large Load Diff
+51 -41
View File
@@ -55,7 +55,7 @@ from ._lazy_imports_registry import (
def _get_litellm_globals() -> dict:
"""
Get the globals dictionary of the litellm module.
This is where we cache imported attributes so we don't import them twice.
When you do `litellm.some_function`, it gets stored in this dictionary.
"""
@@ -65,12 +65,13 @@ def _get_litellm_globals() -> dict:
def _get_utils_globals() -> dict:
"""
Get the globals dictionary of the utils module.
This is where we cache imported attributes so we don't import them twice.
When you do `litellm.utils.some_function`, it gets stored in this dictionary.
"""
return sys.modules["litellm.utils"].__dict__
# These are special lazy loaders for things that are used internally
# They're separate from the main lazy import system because they have specific use cases
@@ -81,10 +82,10 @@ _default_encoding: Optional[Any] = None
def _get_default_encoding() -> Any:
"""
Lazily load and cache the default OpenAI encoding.
This avoids importing `litellm.litellm_core_utils.default_encoding` (and thus tiktoken)
at `litellm` import time. The encoding is cached after the first import.
This is used internally by utils.py functions that need the encoding but shouldn't
trigger its import during module load.
"""
@@ -103,10 +104,10 @@ _get_modified_max_tokens_func: Optional[Any] = None
def _get_modified_max_tokens() -> Any:
"""
Lazily load and cache the get_modified_max_tokens function.
This avoids importing `litellm.litellm_core_utils.token_counter` at `litellm` import time.
The function is cached after the first import.
This is used internally by utils.py functions that need the token counter but shouldn't
trigger its import during module load.
"""
@@ -127,10 +128,10 @@ _token_counter_new_func: Optional[Any] = None
def _get_token_counter_new() -> Any:
"""
Lazily load and cache the token_counter function (aliased as token_counter_new).
This avoids importing `litellm.litellm_core_utils.token_counter` at `litellm` import time.
The function is cached after the first import.
This is used internally by utils.py functions that need the token counter but shouldn't
trigger its import during module load.
"""
@@ -157,10 +158,10 @@ _LAZY_IMPORT_REGISTRY: Optional[dict[str, Callable[[str], Any]]] = None
def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]:
"""
Build the registry that maps attribute names to their handler functions.
This is called once, the first time someone accesses a lazy-loaded attribute.
After that, we just look up the handler function in this dictionary.
Returns:
Dictionary like {"ModelResponse": _lazy_import_utils, ...}
"""
@@ -199,17 +200,19 @@ def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]:
_LAZY_IMPORT_REGISTRY[name] = _lazy_import_llm_provider_logic
for name in UTILS_MODULE_NAMES:
_LAZY_IMPORT_REGISTRY[name] = _lazy_import_utils_module
return _LAZY_IMPORT_REGISTRY
def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> Any:
def _generic_lazy_import(
name: str, import_map: dict[str, tuple[str, str]], category: str
) -> Any:
"""
Generic function that handles lazy importing for most attributes.
This is the workhorse function - it does the actual importing and caching.
Most handler functions just call this with their specific import map.
Steps:
1. Check if the name exists in the import map (if not, raise error)
2. Check if we've already imported it (if yes, return cached value)
@@ -218,7 +221,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
5. Get the attribute from the module
6. Cache it in _globals so we don't import again
7. Return it
Args:
name: The attribute name someone is trying to access (e.g., "ModelResponse")
import_map: Dictionary telling us where to find each attribute
@@ -228,19 +231,19 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
# Step 1: Make sure this attribute exists in our map
if name not in import_map:
raise AttributeError(f"{category} lazy import: unknown attribute {name!r}")
# Step 2: Get the cache (where we store imported things)
_globals = _get_litellm_globals()
# Step 3: If we've already imported it, just return the cached version
if name in _globals:
return _globals[name]
# Step 4: Look up where to find this attribute
# The map tells us: (module_path, attribute_name)
# Example: (".utils", "ModelResponse") means "look in .utils module, get ModelResponse"
module_path, attr_name = import_map[name]
# Step 5: Import the module
# Python automatically caches modules in sys.modules, so calling this twice is fast
# If module_path starts with ".", it's a relative import (needs package="litellm")
@@ -249,14 +252,14 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
module = importlib.import_module(module_path, package="litellm")
else:
module = importlib.import_module(module_path)
# Step 6: Get the actual attribute from the module
# Example: getattr(utils_module, "ModelResponse") returns the ModelResponse class
value = getattr(module, attr_name)
# Step 7: Cache it so we don't have to import again next time
_globals[name] = value
# Step 8: Return it
return value
@@ -268,6 +271,7 @@ def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], cate
# Most of them just call _generic_lazy_import with their specific import map.
# The registry (above) maps attribute names to these handler functions.
def _lazy_import_utils(name: str) -> Any:
"""Handler for utils module attributes (ModelResponse, token_counter, etc.)"""
return _generic_lazy_import(name, _UTILS_IMPORT_MAP, "Utils")
@@ -297,6 +301,7 @@ def _lazy_import_caching(name: str) -> Any:
"""Handler for caching classes (Cache, DualCache, RedisCache, etc.)"""
return _generic_lazy_import(name, _CACHING_IMPORT_MAP, "Caching")
def _lazy_import_dotprompt(name: str) -> Any:
"""Handler for dotprompt integration globals"""
return _generic_lazy_import(name, _DOTPROMPT_IMPORT_MAP, "Dotprompt")
@@ -311,6 +316,7 @@ def _lazy_import_llm_configs(name: str) -> Any:
"""Handler for LLM config classes (AnthropicConfig, OpenAILikeChatConfig, etc.)"""
return _generic_lazy_import(name, _LLM_CONFIGS_IMPORT_MAP, "LLM config")
def _lazy_import_litellm_logging(name: str) -> Any:
"""Handler for litellm_logging module (Logging, modify_integration)"""
return _generic_lazy_import(name, _LITELLM_LOGGING_IMPORT_MAP, "Litellm logging")
@@ -318,87 +324,91 @@ def _lazy_import_litellm_logging(name: str) -> Any:
def _lazy_import_llm_provider_logic(name: str) -> Any:
"""Handler for LLM provider logic functions (get_llm_provider, etc.)"""
return _generic_lazy_import(name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic")
return _generic_lazy_import(
name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic"
)
def _lazy_import_utils_module(name: str) -> Any:
"""
Handler for utils module lazy imports.
This uses a custom implementation because utils module needs to use
_get_utils_globals() instead of _get_litellm_globals() for caching.
"""
# Check if this attribute exists in our map
if name not in _UTILS_MODULE_IMPORT_MAP:
raise AttributeError(f"Utils module lazy import: unknown attribute {name!r}")
# Get the cache (where we store imported things) - use utils globals
_globals = _get_utils_globals()
# If we've already imported it, just return the cached version
if name in _globals:
return _globals[name]
# Look up where to find this attribute
module_path, attr_name = _UTILS_MODULE_IMPORT_MAP[name]
# Import the module
if module_path.startswith("."):
module = importlib.import_module(module_path, package="litellm")
else:
module = importlib.import_module(module_path)
# Get the actual attribute from the module
value = getattr(module, attr_name)
# Cache it so we don't have to import again next time
_globals[name] = value
# Return it
return value
# ============================================================================
# SPECIAL HANDLERS
# ============================================================================
# These handlers have custom logic that doesn't fit the generic pattern
def _lazy_import_llm_client_cache(name: str) -> Any:
"""
Handler for LLM client cache - has special logic for singleton instance.
This one is different because:
- "LLMClientCache" is the class itself
- "in_memory_llm_clients_cache" is a singleton instance of that class
So we need custom logic to handle both cases.
"""
_globals = _get_litellm_globals()
# If already cached, return it
if name in _globals:
return _globals[name]
# Import the class
module = importlib.import_module("litellm.caching.llm_caching_handler")
LLMClientCache = getattr(module, "LLMClientCache")
# If they want the class itself, return it
if name == "LLMClientCache":
_globals["LLMClientCache"] = LLMClientCache
return LLMClientCache
# If they want the singleton instance, create it (only once)
if name == "in_memory_llm_clients_cache":
instance = LLMClientCache()
_globals["in_memory_llm_clients_cache"] = instance
return instance
raise AttributeError(f"LLM client cache lazy import: unknown attribute {name!r}")
def _lazy_import_http_handlers(name: str) -> Any:
"""
Handler for HTTP clients - has special logic for creating client instances.
This one is different because:
- These aren't just imports, they're actual client instances that need to be created
- They need configuration (timeout, etc.) from the module globals
@@ -413,14 +423,14 @@ def _lazy_import_http_handlers(name: str) -> Any:
# Get timeout from module config (if set)
timeout = _globals.get("request_timeout")
params = {"timeout": timeout, "client_alias": "module level aclient"}
# Create the client instance
provider_id = cast(Any, "litellm_module_level_client")
async_client = get_async_httpx_client(
llm_provider=provider_id,
params=params,
)
# Cache it so we don't create it again
_globals["module_level_aclient"] = async_client
return async_client
@@ -431,7 +441,7 @@ def _lazy_import_http_handlers(name: str) -> Any:
timeout = _globals.get("request_timeout")
sync_client = HTTPHandler(timeout=timeout)
# Cache it
_globals["module_level_client"] = sync_client
return sync_client
+8 -2
View File
@@ -677,7 +677,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
"FireworksAIRerankConfig",
),
"VoyageRerankConfig": (".llms.voyage.rerank.transformation", "VoyageRerankConfig"),
"IBMWatsonXRerankConfig": (".llms.watsonx.rerank.transformation", "IBMWatsonXRerankConfig"),
"IBMWatsonXRerankConfig": (
".llms.watsonx.rerank.transformation",
"IBMWatsonXRerankConfig",
),
"ClarifaiConfig": (".llms.clarifai.chat.transformation", "ClarifaiConfig"),
"AI21ChatConfig": (".llms.ai21.chat.transformation", "AI21ChatConfig"),
"LlamaAPIConfig": (".llms.meta_llama.chat.transformation", "LlamaAPIConfig"),
@@ -859,7 +862,10 @@ _LLM_CONFIGS_IMPORT_MAP = {
"OpenAITextCompletionConfig",
),
"GroqChatConfig": (".llms.groq.chat.transformation", "GroqChatConfig"),
"BedrockMantleChatConfig": (".llms.bedrock_mantle.chat.transformation", "BedrockMantleChatConfig"),
"BedrockMantleChatConfig": (
".llms.bedrock_mantle.chat.transformation",
"BedrockMantleChatConfig",
),
"A2AConfig": (".llms.a2a.chat.transformation", "A2AConfig"),
"GenAIHubOrchestrationConfig": (
".llms.sap.chat.transformation",
+71 -44
View File
@@ -34,7 +34,12 @@ def _get_redis_kwargs():
"retry",
}
include_args = ["url", "redis_connect_func", "gcp_service_account", "gcp_ssl_ca_certs"]
include_args = [
"url",
"redis_connect_func",
"gcp_service_account",
"gcp_ssl_ca_certs",
]
available_args = [x for x in arg_spec.args if x not in exclude_args] + include_args
@@ -75,7 +80,9 @@ def _get_redis_cluster_kwargs(client=None):
available_args.append("ssl_cert_reqs")
available_args.append("ssl_check_hostname")
available_args.append("ssl_ca_certs")
available_args.append("redis_connect_func") # Needed for sync clusters and IAM detection
available_args.append(
"redis_connect_func"
) # Needed for sync clusters and IAM detection
available_args.append("gcp_service_account")
available_args.append("gcp_ssl_ca_certs")
available_args.append("max_connections")
@@ -103,10 +110,10 @@ def _redis_kwargs_from_environment():
def _generate_gcp_iam_access_token(service_account: str) -> str:
"""
Generate GCP IAM access token for Redis authentication.
Args:
service_account: GCP service account in format 'projects/-/serviceAccounts/name@project.iam.gserviceaccount.com'
Returns:
Access token string for GCP IAM authentication
"""
@@ -117,11 +124,11 @@ def _generate_gcp_iam_access_token(service_account: str) -> str:
"google-cloud-iam is required for GCP IAM Redis authentication. "
"Install it with: pip install google-cloud-iam"
)
client = iam_credentials_v1.IAMCredentialsClient()
request = iam_credentials_v1.GenerateAccessTokenRequest(
name=service_account,
scope=['https://www.googleapis.com/auth/cloud-platform'],
scope=["https://www.googleapis.com/auth/cloud-platform"],
)
response = client.generate_access_token(request=request)
return str(response.access_token)
@@ -133,14 +140,15 @@ def create_gcp_iam_redis_connect_func(
) -> Callable:
"""
Creates a custom Redis connection function for GCP IAM authentication.
Args:
service_account: GCP service account in format 'projects/-/serviceAccounts/name@project.iam.gserviceaccount.com'
ssl_ca_certs: Path to SSL CA certificate file for secure connections
Returns:
A connection function that can be used with Redis clients
"""
def iam_connect(self):
"""Initialize the connection and authenticate using GCP IAM"""
from redis.exceptions import (
@@ -148,25 +156,25 @@ def create_gcp_iam_redis_connect_func(
AuthenticationWrongNumberOfArgsError,
)
from redis.utils import str_if_bytes
self._parser.on_connect(self)
auth_args = (_generate_gcp_iam_access_token(service_account),)
self.send_command("AUTH", *auth_args, check_health=False)
try:
auth_response = self.read_response()
except AuthenticationWrongNumberOfArgsError:
# Fallback to password auth if IAM fails
if hasattr(self, 'password') and self.password:
if hasattr(self, "password") and self.password:
self.send_command("AUTH", self.password, check_health=False)
auth_response = self.read_response()
else:
raise
if str_if_bytes(auth_response) != "OK":
raise AuthenticationError("GCP IAM authentication failed")
return iam_connect
@@ -178,22 +186,20 @@ def get_redis_url_from_environment():
raise ValueError(
"Either 'REDIS_URL' or both 'REDIS_HOST' and 'REDIS_PORT' must be specified for Redis."
)
if "REDIS_SSL" in os.environ and os.environ["REDIS_SSL"].lower() == "true":
redis_protocol = "rediss"
else:
redis_protocol = "redis"
# Build authentication part of URL
auth_part = ""
if "REDIS_USERNAME" in os.environ and "REDIS_PASSWORD" in os.environ:
auth_part = f"{os.environ['REDIS_USERNAME']}:{os.environ['REDIS_PASSWORD']}@"
elif "REDIS_PASSWORD" in os.environ:
auth_part = f"{os.environ['REDIS_PASSWORD']}@"
return (
f"{redis_protocol}://{auth_part}{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}"
)
return f"{redis_protocol}://{auth_part}{os.environ['REDIS_HOST']}:{os.environ['REDIS_PORT']}"
def _get_redis_client_logic(**env_overrides):
@@ -241,22 +247,27 @@ def _get_redis_client_logic(**env_overrides):
redis_kwargs["service_name"] = _service_name
# Handle GCP IAM authentication
_gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT")
_gcp_ssl_ca_certs = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str("REDIS_GCP_SSL_CA_CERTS")
_gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str(
"REDIS_GCP_SERVICE_ACCOUNT"
)
_gcp_ssl_ca_certs = redis_kwargs.get("gcp_ssl_ca_certs") or get_secret_str(
"REDIS_GCP_SSL_CA_CERTS"
)
if _gcp_service_account is not None:
verbose_logger.debug("Setting up GCP IAM authentication for Redis with service account.")
verbose_logger.debug(
"Setting up GCP IAM authentication for Redis with service account."
)
redis_kwargs["redis_connect_func"] = create_gcp_iam_redis_connect_func(
service_account=_gcp_service_account,
ssl_ca_certs=_gcp_ssl_ca_certs
service_account=_gcp_service_account, ssl_ca_certs=_gcp_ssl_ca_certs
)
# Store GCP service account in redis_connect_func for async cluster access
redis_kwargs["redis_connect_func"]._gcp_service_account = _gcp_service_account
# Remove GCP-specific kwargs that shouldn't be passed to Redis client
redis_kwargs.pop("gcp_service_account", None)
redis_kwargs.pop("gcp_ssl_ca_certs", None)
# Only enable SSL if explicitly requested AND SSL CA certs are provided
if _gcp_ssl_ca_certs and redis_kwargs.get("ssl", False):
redis_kwargs["ssl_ca_certs"] = _gcp_ssl_ca_certs
@@ -377,7 +388,8 @@ def get_redis_client(**env_overrides):
def get_redis_async_client(
connection_pool: Optional[async_redis.BlockingConnectionPool] = None, **env_overrides,
connection_pool: Optional[async_redis.BlockingConnectionPool] = None,
**env_overrides,
) -> Union[async_redis.Redis, async_redis.RedisCluster]:
redis_kwargs = _get_redis_client_logic(**env_overrides)
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
@@ -411,39 +423,50 @@ def get_redis_async_client(
# Get GCP service account - first try from redis_connect_func, then from environment
gcp_service_account = None
if redis_connect_func and hasattr(redis_connect_func, '_gcp_service_account'):
if redis_connect_func and hasattr(redis_connect_func, "_gcp_service_account"):
gcp_service_account = redis_connect_func._gcp_service_account
else:
gcp_service_account = redis_kwargs.get("gcp_service_account") or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT")
verbose_logger.debug(f"DEBUG: Redis cluster kwargs: redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}")
gcp_service_account = redis_kwargs.get(
"gcp_service_account"
) or get_secret_str("REDIS_GCP_SERVICE_ACCOUNT")
verbose_logger.debug(
f"DEBUG: Redis cluster kwargs: redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}"
)
# If GCP IAM is configured (indicated by redis_connect_func), generate access token and use as password
if redis_connect_func and gcp_service_account:
verbose_logger.debug("DEBUG: Generating IAM token for service account (value not logged for security reasons)")
verbose_logger.debug(
"DEBUG: Generating IAM token for service account (value not logged for security reasons)"
)
try:
# Generate IAM access token using the helper function
access_token = _generate_gcp_iam_access_token(gcp_service_account)
cluster_kwargs["password"] = access_token
verbose_logger.debug("DEBUG: Successfully generated GCP IAM access token for async Redis cluster")
verbose_logger.debug(
"DEBUG: Successfully generated GCP IAM access token for async Redis cluster"
)
except Exception as e:
verbose_logger.error(f"Failed to generate GCP IAM access token: {e}")
from redis.exceptions import AuthenticationError
raise AuthenticationError("Failed to generate GCP IAM access token")
else:
verbose_logger.debug(f"DEBUG: Not using GCP IAM auth - redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}")
verbose_logger.debug(
f"DEBUG: Not using GCP IAM auth - redis_connect_func={redis_connect_func is not None}, gcp_service_account_provided={gcp_service_account is not None}"
)
new_startup_nodes: List[ClusterNode] = []
for item in redis_kwargs["startup_nodes"]:
new_startup_nodes.append(ClusterNode(**item))
cluster_kwargs.pop("startup_nodes", None)
# Create async RedisCluster with IAM token as password if available
cluster_client = async_redis.RedisCluster(
startup_nodes=new_startup_nodes, **cluster_kwargs # type: ignore
)
return cluster_client
# Check for Redis Sentinel
@@ -463,7 +486,10 @@ def get_redis_connection_pool(**env_overrides):
redis_kwargs = _get_redis_client_logic(**env_overrides)
verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs)
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
pool_kwargs = {"timeout": REDIS_CONNECTION_POOL_TIMEOUT, "url": redis_kwargs["url"]}
pool_kwargs = {
"timeout": REDIS_CONNECTION_POOL_TIMEOUT,
"url": redis_kwargs["url"],
}
if "max_connections" in redis_kwargs:
try:
pool_kwargs["max_connections"] = int(redis_kwargs["max_connections"])
@@ -483,6 +509,7 @@ def get_redis_connection_pool(**env_overrides):
timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs
)
def _pretty_print_redis_config(redis_kwargs: dict) -> None:
"""Pretty print the Redis configuration using rich with sensitive data masking"""
try:
@@ -492,6 +519,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
from rich.panel import Panel
from rich.table import Table
from rich.text import Text
if not verbose_logger.isEnabledFor(logging.DEBUG):
return
@@ -499,7 +527,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
# Initialize the sensitive data masker
masker = SensitiveDataMasker()
# Mask sensitive data in redis_kwargs
masked_redis_kwargs = masker.mask_dict(redis_kwargs)
@@ -531,7 +559,7 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
value_str = str(value)
else:
value_str = str(value)
config_table.add_row(key, value_str)
# Determine connection type
@@ -568,4 +596,3 @@ def _pretty_print_redis_config(redis_kwargs: dict) -> None:
verbose_logger.info(f"Redis configuration: {masked_redis_kwargs}")
except Exception as e:
verbose_logger.error(f"Error pretty printing Redis configuration: {e}")
+1 -1
View File
@@ -317,7 +317,7 @@ class ServiceLogging(CustomLogger):
await self.async_service_success_hook(
service=ServiceTypes.LITELLM,
duration=_duration,
call_type=kwargs.get("call_type", "unknown")
call_type=kwargs.get("call_type", "unknown"),
)
except Exception as e:
raise e
+3 -1
View File
@@ -103,5 +103,7 @@ class A2AClient:
from litellm.a2a_protocol.main import asend_message_streaming
a2a_client = await self._get_client()
async for chunk in asend_message_streaming(a2a_client=a2a_client, request=request):
async for chunk in asend_message_streaming(
a2a_client=a2a_client, request=request
):
yield chunk
+6 -2
View File
@@ -97,7 +97,11 @@ class A2ACostCalculator:
completion_tokens = getattr(usage, "completion_tokens", 0) or 0
# Calculate costs
input_cost = prompt_tokens * (float(input_cost_per_token) if input_cost_per_token else 0.0)
output_cost = completion_tokens * (float(output_cost_per_token) if output_cost_per_token else 0.0)
input_cost = prompt_tokens * (
float(input_cost_per_token) if input_cost_per_token else 0.0
)
output_cost = completion_tokens * (
float(output_cost_per_token) if output_cost_per_token else 0.0
)
return input_cost + output_cost
@@ -50,30 +50,28 @@ class A2ACompletionBridgeHandler:
a2a_provider_config = A2AProviderConfigManager.get_provider_config(
custom_llm_provider=custom_llm_provider
)
# If provider config exists, use it
if a2a_provider_config is not None:
if api_base is None:
raise ValueError(f"api_base is required for {custom_llm_provider}")
verbose_logger.info(
f"A2A: Using provider config for {custom_llm_provider}"
)
verbose_logger.info(f"A2A: Using provider config for {custom_llm_provider}")
response_data = await a2a_provider_config.handle_non_streaming(
request_id=request_id,
params=params,
api_base=api_base,
)
return response_data
# Extract message from params
message = params.get("message", {})
# Transform A2A message to OpenAI format
openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(
message
openai_messages = (
A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
)
# Get completion params
@@ -100,7 +98,8 @@ class A2ACompletionBridgeHandler:
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v for k, v in litellm_params.items()
k: v
for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider")
}
completion_params.update(litellm_params_to_add)
@@ -109,9 +108,11 @@ class A2ACompletionBridgeHandler:
response = await litellm.acompletion(**completion_params)
# Transform response to A2A format
a2a_response = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
response=response,
request_id=request_id,
a2a_response = (
A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
response=response,
request_id=request_id,
)
)
verbose_logger.info(f"A2A completion bridge completed: request_id={request_id}")
@@ -148,25 +149,25 @@ class A2ACompletionBridgeHandler:
a2a_provider_config = A2AProviderConfigManager.get_provider_config(
custom_llm_provider=custom_llm_provider
)
# If provider config exists, use it
if a2a_provider_config is not None:
if api_base is None:
raise ValueError(f"api_base is required for {custom_llm_provider}")
verbose_logger.info(
f"A2A: Using provider config for {custom_llm_provider} (streaming)"
)
async for chunk in a2a_provider_config.handle_streaming(
request_id=request_id,
params=params,
api_base=api_base,
):
yield chunk
return
# Extract message from params
message = params.get("message", {})
@@ -177,8 +178,8 @@ class A2ACompletionBridgeHandler:
)
# Transform A2A message to OpenAI format
openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(
message
openai_messages = (
A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
)
# Get completion params
@@ -205,7 +206,8 @@ class A2ACompletionBridgeHandler:
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v for k, v in litellm_params.items()
k: v
for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider")
}
completion_params.update(litellm_params_to_add)
@@ -244,9 +246,11 @@ class A2ACompletionBridgeHandler:
# Emit artifact update with accumulated content
if accumulated_text:
artifact_event = A2ACompletionBridgeTransformation.create_artifact_update_event(
ctx=ctx,
text=accumulated_text,
artifact_event = (
A2ACompletionBridgeTransformation.create_artifact_update_event(
ctx=ctx,
text=accumulated_text,
)
)
yield artifact_event
@@ -124,9 +124,7 @@ class A2ACompletionBridgeTransformation:
},
}
verbose_logger.debug(
f"OpenAI -> A2A transform: content_length={len(content)}"
)
verbose_logger.debug(f"OpenAI -> A2A transform: content_length={len(content)}")
return a2a_response
+4 -6
View File
@@ -120,9 +120,9 @@ def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str:
litellm_logging_obj.model = model
litellm_logging_obj.custom_llm_provider = custom_llm_provider
litellm_logging_obj.model_call_details["model"] = model
litellm_logging_obj.model_call_details["custom_llm_provider"] = (
custom_llm_provider
)
litellm_logging_obj.model_call_details[
"custom_llm_provider"
] = custom_llm_provider
return agent_name
@@ -664,9 +664,7 @@ async def create_a2a_client(
if extra_headers:
# Encode headers into a cache-key-only param so each unique header
# set produces a distinct cache key.
_client_params["disable_aiohttp_transport"] = str(
sorted(extra_headers.items())
)
_client_params["disable_aiohttp_transport"] = str(sorted(extra_headers.items()))
_async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.A2AProvider,
params=_client_params,
@@ -8,4 +8,3 @@ from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager
__all__ = ["BaseA2AProviderConfig", "A2AProviderConfigManager"]
+1 -2
View File
@@ -9,7 +9,7 @@ from typing import Any, AsyncIterator, Dict
class BaseA2AProviderConfig(ABC):
"""
Base configuration class for A2A protocol providers.
Each provider should implement this interface to define how to handle
A2A requests for their specific agent type.
"""
@@ -60,4 +60,3 @@ class BaseA2AProviderConfig(ABC):
# The yield is here to make this a generator function
if False: # pragma: no cover
yield {}
@@ -12,7 +12,7 @@ from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig
class A2AProviderConfigManager:
"""
Manager for A2A provider configurations.
Similar to ProviderConfigManager in litellm.utils but specifically for A2A providers.
"""
@@ -31,7 +31,7 @@ class A2AProviderConfigManager:
"""
if custom_llm_provider is None:
return None
if custom_llm_provider == "pydantic_ai_agents":
from litellm.a2a_protocol.providers.pydantic_ai_agents.config import (
PydanticAIProviderConfig,
@@ -45,4 +45,3 @@ class A2AProviderConfigManager:
# return AnotherProviderConfig()
return None
@@ -3,4 +3,3 @@ LiteLLM Completion bridge provider for A2A protocol.
Routes A2A requests through litellm.acompletion based on custom_llm_provider.
"""
@@ -52,26 +52,26 @@ class A2ACompletionBridgeHandler:
if custom_llm_provider == "pydantic_ai_agents":
if api_base is None:
raise ValueError("api_base is required for Pydantic AI agents")
verbose_logger.info(
f"Pydantic AI: Routing to Pydantic AI agent at {api_base}"
)
# Send request directly to Pydantic AI agent
response_data = await PydanticAITransformation.send_non_streaming_request(
api_base=api_base,
request_id=request_id,
params=params,
)
return response_data
# Extract message from params
message = params.get("message", {})
# Transform A2A message to OpenAI format
openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(
message
openai_messages = (
A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
)
# Get completion params
@@ -98,7 +98,8 @@ class A2ACompletionBridgeHandler:
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v for k, v in litellm_params.items()
k: v
for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider")
}
completion_params.update(litellm_params_to_add)
@@ -107,9 +108,11 @@ class A2ACompletionBridgeHandler:
response = await litellm.acompletion(**completion_params)
# Transform response to A2A format
a2a_response = A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
response=response,
request_id=request_id,
a2a_response = (
A2ACompletionBridgeTransformation.openai_response_to_a2a_response(
response=response,
request_id=request_id,
)
)
verbose_logger.info(f"A2A completion bridge completed: request_id={request_id}")
@@ -146,27 +149,27 @@ class A2ACompletionBridgeHandler:
if custom_llm_provider == "pydantic_ai_agents":
if api_base is None:
raise ValueError("api_base is required for Pydantic AI agents")
verbose_logger.info(
f"Pydantic AI: Faking streaming for Pydantic AI agent at {api_base}"
)
# Get non-streaming response first
response_data = await PydanticAITransformation.send_non_streaming_request(
api_base=api_base,
request_id=request_id,
params=params,
)
# Convert to fake streaming
async for chunk in PydanticAITransformation.fake_streaming_from_response(
response_data=response_data,
request_id=request_id,
):
yield chunk
return
# Extract message from params
message = params.get("message", {})
@@ -177,8 +180,8 @@ class A2ACompletionBridgeHandler:
)
# Transform A2A message to OpenAI format
openai_messages = A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(
message
openai_messages = (
A2ACompletionBridgeTransformation.a2a_message_to_openai_messages(message)
)
# Get completion params
@@ -205,7 +208,8 @@ class A2ACompletionBridgeHandler:
}
# Add litellm_params (contains api_key, client_id, client_secret, tenant_id, etc.)
litellm_params_to_add = {
k: v for k, v in litellm_params.items()
k: v
for k, v in litellm_params.items()
if k not in ("model", "custom_llm_provider")
}
completion_params.update(litellm_params_to_add)
@@ -244,9 +248,11 @@ class A2ACompletionBridgeHandler:
# Emit artifact update with accumulated content
if accumulated_text:
artifact_event = A2ACompletionBridgeTransformation.create_artifact_update_event(
ctx=ctx,
text=accumulated_text,
artifact_event = (
A2ACompletionBridgeTransformation.create_artifact_update_event(
ctx=ctx,
text=accumulated_text,
)
)
yield artifact_event
@@ -124,9 +124,7 @@ class A2ACompletionBridgeTransformation:
},
}
verbose_logger.debug(
f"OpenAI -> A2A transform: content_length={len(content)}"
)
verbose_logger.debug(f"OpenAI -> A2A transform: content_length={len(content)}")
return a2a_response
@@ -14,4 +14,3 @@ from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import (
)
__all__ = ["PydanticAIHandler", "PydanticAITransformation", "PydanticAIProviderConfig"]
@@ -11,7 +11,7 @@ from litellm.a2a_protocol.providers.pydantic_ai_agents.handler import PydanticAI
class PydanticAIProviderConfig(BaseA2AProviderConfig):
"""
Provider configuration for Pydantic AI agents.
Pydantic AI agents follow A2A protocol but don't support streaming natively.
This config provides fake streaming by converting non-streaming responses into streaming chunks.
"""
@@ -48,4 +48,3 @@ class PydanticAIProviderConfig(BaseA2AProviderConfig):
delay_ms=kwargs.get("delay_ms", 10),
):
yield chunk
@@ -16,7 +16,7 @@ from litellm.a2a_protocol.providers.pydantic_ai_agents.transformation import (
class PydanticAIHandler:
"""
Handler for Pydantic AI agent requests.
Provides:
- Direct non-streaming requests to Pydantic AI agents
- Fake streaming by converting non-streaming responses into streaming chunks
@@ -41,9 +41,7 @@ class PydanticAIHandler:
Returns:
A2A SendMessageResponse dict
"""
verbose_logger.info(
f"Pydantic AI: Routing to Pydantic AI agent at {api_base}"
)
verbose_logger.info(f"Pydantic AI: Routing to Pydantic AI agent at {api_base}")
# Send request directly to Pydantic AI agent
response_data = await PydanticAITransformation.send_non_streaming_request(
@@ -102,5 +100,3 @@ class PydanticAIHandler:
delay_ms=delay_ms,
):
yield chunk
@@ -10,13 +10,16 @@ from typing import Any, AsyncIterator, Dict, cast
from uuid import uuid4
from litellm._logging import verbose_logger
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client,
)
class PydanticAITransformation:
"""
Transformation layer for Pydantic AI agents.
Handles:
- Direct A2A requests to Pydantic AI endpoints
- Polling for task completion (since Pydantic AI doesn't support streaming)
@@ -27,13 +30,13 @@ class PydanticAITransformation:
def _remove_none_values(obj: Any) -> Any:
"""
Recursively remove None values from a dict/list structure.
FastA2A/Pydantic AI servers don't accept None values for optional fields -
they expect those fields to be omitted entirely.
Args:
obj: Dict, list, or other value to clean
Returns:
Cleaned object with None values removed
"""
@@ -56,10 +59,10 @@ class PydanticAITransformation:
def _params_to_dict(params: Any) -> Dict[str, Any]:
"""
Convert params to a dict, handling Pydantic models.
Args:
params: Dict or Pydantic model
Returns:
Dict representation of params
"""
@@ -86,7 +89,7 @@ class PydanticAITransformation:
) -> Dict[str, Any]:
"""
Poll for task completion using tasks/get method.
Args:
client: HTTPX async client
endpoint: API endpoint URL
@@ -94,7 +97,7 @@ class PydanticAITransformation:
request_id: JSON-RPC request ID
max_attempts: Maximum polling attempts
poll_interval: Seconds between poll attempts
Returns:
Completed task response
"""
@@ -105,7 +108,7 @@ class PydanticAITransformation:
"method": "tasks/get",
"params": {"id": task_id},
}
response = await client.post(
endpoint,
json=poll_request,
@@ -113,23 +116,25 @@ class PydanticAITransformation:
)
response.raise_for_status()
poll_data = response.json()
result = poll_data.get("result", {})
status = result.get("status", {})
state = status.get("state", "")
verbose_logger.debug(
f"Pydantic AI: Poll attempt {attempt + 1}/{max_attempts}, state={state}"
)
if state == "completed":
return poll_data
elif state in ("failed", "canceled"):
raise Exception(f"Task {task_id} ended with state: {state}")
await asyncio.sleep(poll_interval)
raise TimeoutError(f"Task {task_id} did not complete within {max_attempts * poll_interval} seconds")
raise TimeoutError(
f"Task {task_id} did not complete within {max_attempts * poll_interval} seconds"
)
@staticmethod
async def _send_and_poll_raw(
@@ -140,7 +145,7 @@ class PydanticAITransformation:
) -> Dict[str, Any]:
"""
Send a request to Pydantic AI agent and return the raw task response.
This is an internal method used by both non-streaming and streaming handlers.
Returns the raw Pydantic AI task format with history/artifacts.
@@ -155,10 +160,10 @@ class PydanticAITransformation:
"""
# Convert params to dict if it's a Pydantic model
params_dict = PydanticAITransformation._params_to_dict(params)
# Remove None values - FastA2A doesn't accept null for optional fields
params_dict = PydanticAITransformation._remove_none_values(params_dict)
# Ensure the message has 'kind': 'message' as required by FastA2A/Pydantic AI
if "message" in params_dict:
params_dict["message"]["kind"] = "message"
@@ -174,9 +179,7 @@ class PydanticAITransformation:
# FastA2A uses root endpoint (/) not /messages
endpoint = api_base.rstrip("/")
verbose_logger.info(
f"Pydantic AI: Sending non-streaming request to {endpoint}"
)
verbose_logger.info(f"Pydantic AI: Sending non-streaming request to {endpoint}")
# Send request to Pydantic AI agent using shared async HTTP client
client = get_async_httpx_client(
@@ -190,12 +193,12 @@ class PydanticAITransformation:
)
response.raise_for_status()
response_data = response.json()
# Check if task is already completed
result = response_data.get("result", {})
status = result.get("status", {})
state = status.get("state", "")
if state != "completed":
# Need to poll for completion
task_id = result.get("id")
@@ -210,7 +213,9 @@ class PydanticAITransformation:
request_id=request_id,
)
verbose_logger.info(f"Pydantic AI: Received completed response for request_id={request_id}")
verbose_logger.info(
f"Pydantic AI: Received completed response for request_id={request_id}"
)
return response_data
@@ -256,7 +261,7 @@ class PydanticAITransformation:
) -> Dict[str, Any]:
"""
Send a request to Pydantic AI agent and return the raw task response.
Used by streaming handler to get raw response for fake streaming.
Args:
@@ -282,7 +287,7 @@ class PydanticAITransformation:
) -> Dict[str, Any]:
"""
Transform Pydantic AI task response to standard A2A non-streaming format.
Pydantic AI returns a task with history/artifacts, but the standard A2A
non-streaming format expects:
{
@@ -296,11 +301,11 @@ class PydanticAITransformation:
}
}
}
Args:
response_data: Pydantic AI task response
request_id: Original request ID
Returns:
Standard A2A non-streaming response format
"""
@@ -308,14 +313,14 @@ class PydanticAITransformation:
full_text, message_id, parts = PydanticAITransformation._extract_response_text(
response_data
)
# Build standard A2A message
a2a_message = {
"role": "agent",
"parts": parts if parts else [{"kind": "text", "text": full_text}],
"messageId": message_id,
}
# Return standard A2A non-streaming format
return {
"jsonrpc": "2.0",
@@ -329,19 +334,19 @@ class PydanticAITransformation:
def _extract_response_text(response_data: Dict[str, Any]) -> tuple[str, str, list]:
"""
Extract response text from completed task response.
Pydantic AI returns completed tasks with:
- history: list of messages (user and agent)
- artifacts: list of result artifacts
Args:
response_data: Completed task response
Returns:
Tuple of (full_text, message_id, parts)
"""
result = response_data.get("result", {})
# Try to extract from artifacts first (preferred for results)
artifacts = result.get("artifacts", [])
if artifacts:
@@ -352,7 +357,7 @@ class PydanticAITransformation:
text = part.get("text", "")
if text:
return text, str(uuid4()), parts
# Fall back to history - get the last agent message
history = result.get("history", [])
for msg in reversed(history):
@@ -365,7 +370,7 @@ class PydanticAITransformation:
full_text += part.get("text", "")
if full_text:
return full_text, message_id, parts
# Fall back to message field (original format)
message = result.get("message", {})
if message:
@@ -376,7 +381,7 @@ class PydanticAITransformation:
if part.get("kind") == "text":
full_text += part.get("text", "")
return full_text, message_id, parts
return "", str(uuid4()), []
@staticmethod
@@ -408,7 +413,7 @@ class PydanticAITransformation:
full_text, message_id, parts = PydanticAITransformation._extract_response_text(
response_data
)
# Extract input message from raw response for history
result = response_data.get("result", {})
history = result.get("history", [])
@@ -436,7 +441,9 @@ class PydanticAITransformation:
"contextId": context_id,
"kind": "message",
"messageId": input_message_id,
"parts": input_message.get("parts", [{"kind": "text", "text": ""}]),
"parts": input_message.get(
"parts", [{"kind": "text", "text": ""}]
),
"role": "user",
"taskId": task_id,
}
@@ -475,7 +482,7 @@ class PydanticAITransformation:
if full_text:
# Split text into chunks
for i in range(0, len(full_text), chunk_size):
chunk_text = full_text[i:i + chunk_size]
chunk_text = full_text[i : i + chunk_size]
is_last_chunk = (i + chunk_size) >= len(full_text)
artifact_event = {
@@ -521,5 +528,3 @@ class PydanticAITransformation:
verbose_logger.info(
f"Pydantic AI: Fake streaming completed for request_id={request_id}"
)
+16 -5
View File
@@ -71,7 +71,11 @@ class A2AStreamingIterator:
def _collect_text_from_chunk(self, chunk: Any) -> None:
"""Extract text from a streaming chunk and add to collected parts."""
try:
chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {}
chunk_dict = (
chunk.model_dump(mode="json", exclude_none=True)
if hasattr(chunk, "model_dump")
else {}
)
text = A2ARequestUtils.extract_text_from_response(chunk_dict)
if text:
self.collected_text_parts.append(text)
@@ -81,7 +85,11 @@ class A2AStreamingIterator:
def _is_completed_chunk(self, chunk: Any) -> bool:
"""Check if chunk indicates stream completion."""
try:
chunk_dict = chunk.model_dump(mode="json", exclude_none=True) if hasattr(chunk, "model_dump") else {}
chunk_dict = (
chunk.model_dump(mode="json", exclude_none=True)
if hasattr(chunk, "model_dump")
else {}
)
result = chunk_dict.get("result", {})
if isinstance(result, dict):
status = result.get("status", {})
@@ -102,7 +110,9 @@ class A2AStreamingIterator:
prompt_tokens = A2ARequestUtils.count_tokens(input_text)
# Use the last (most complete) text from chunks
output_text = self.collected_text_parts[-1] if self.collected_text_parts else ""
output_text = (
self.collected_text_parts[-1] if self.collected_text_parts else ""
)
completion_tokens = A2ARequestUtils.count_tokens(output_text)
total_tokens = prompt_tokens + completion_tokens
@@ -158,7 +168,9 @@ class A2AStreamingIterator:
result: Dict[str, Any] = {
"id": getattr(self.request, "id", "unknown"),
"jsonrpc": "2.0",
"usage": usage.model_dump() if hasattr(usage, "model_dump") else dict(usage),
"usage": usage.model_dump()
if hasattr(usage, "model_dump")
else dict(usage),
}
# Add final chunk result if available
@@ -170,4 +182,3 @@ class A2AStreamingIterator:
pass
return result
+57 -49
View File
@@ -38,7 +38,7 @@ _BETA_HEADERS_CONFIG: Optional[Dict] = None
class GetAnthropicBetaHeadersConfig:
"""
Handles fetching, validating, and loading the Anthropic beta headers configuration.
Similar to GetModelCostMap, this class manages the lifecycle of the beta headers
configuration with support for remote fetching and local fallback.
"""
@@ -62,7 +62,7 @@ class GetAnthropicBetaHeadersConfig:
"bedrock": {},
"bedrock_converse": {},
"vertex_ai": {},
"provider_aliases": {}
"provider_aliases": {},
}
@staticmethod
@@ -84,9 +84,15 @@ class GetAnthropicBetaHeadersConfig:
return False
# Check for at least one provider key
provider_keys = ["anthropic", "azure_ai", "bedrock", "bedrock_converse", "vertex_ai"]
provider_keys = [
"anthropic",
"azure_ai",
"bedrock",
"bedrock_converse",
"vertex_ai",
]
has_provider = any(key in fetched_config for key in provider_keys)
if not has_provider:
verbose_logger.warning(
"LiteLLM: Fetched beta headers config missing provider keys. "
@@ -100,7 +106,7 @@ class GetAnthropicBetaHeadersConfig:
def validate_beta_headers_config(cls, fetched_config: dict) -> bool:
"""
Validate the integrity of a fetched beta headers config.
Returns True if all checks pass, False otherwise.
"""
return cls._check_is_valid_dict(fetched_config)
@@ -109,7 +115,7 @@ class GetAnthropicBetaHeadersConfig:
def fetch_remote_beta_headers_config(url: str, timeout: int = 5) -> dict:
"""
Fetch the beta headers config from a remote URL.
Returns the parsed JSON dict. Raises on network/parse errors
(caller is expected to handle).
"""
@@ -121,14 +127,14 @@ class GetAnthropicBetaHeadersConfig:
def get_beta_headers_config(url: str) -> dict:
"""
Public entry point returns the beta headers config dict.
1. If ``LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS`` is set, uses the local backup only.
2. Otherwise fetches from ``url``, validates integrity, and falls back
to the local backup on any failure.
Args:
url: URL to fetch the remote beta headers configuration from
Returns:
Dict containing the beta headers configuration
"""
@@ -149,7 +155,9 @@ def get_beta_headers_config(url: str) -> dict:
return GetAnthropicBetaHeadersConfig.load_local_beta_headers_config()
# Validate the fetched config
if not GetAnthropicBetaHeadersConfig.validate_beta_headers_config(fetched_config=content):
if not GetAnthropicBetaHeadersConfig.validate_beta_headers_config(
fetched_config=content
):
verbose_logger.warning(
"LiteLLM: Fetched beta headers config failed integrity check. "
"Using local backup instead. url=%s",
@@ -164,23 +172,23 @@ def _load_beta_headers_config() -> Dict:
"""
Load the beta headers configuration.
Uses caching to avoid repeated fetches/file reads.
This function is called by all public API functions and manages the global cache.
Returns:
Dict containing the beta headers configuration
"""
global _BETA_HEADERS_CONFIG
if _BETA_HEADERS_CONFIG is not None:
return _BETA_HEADERS_CONFIG
# Get the URL from environment or use default
from litellm import anthropic_beta_headers_url
_BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url)
verbose_logger.debug("Loaded and cached beta headers config")
return _BETA_HEADERS_CONFIG
@@ -188,7 +196,7 @@ def reload_beta_headers_config() -> Dict:
"""
Force reload the beta headers configuration from source (remote or local).
Clears the cache and fetches fresh configuration.
Returns:
Dict containing the newly loaded beta headers configuration
"""
@@ -201,10 +209,10 @@ def reload_beta_headers_config() -> Dict:
def get_provider_name(provider: str) -> str:
"""
Resolve provider aliases to canonical provider names.
Args:
provider: Provider name (may be an alias)
Returns:
Canonical provider name
"""
@@ -219,53 +227,53 @@ def filter_and_transform_beta_headers(
) -> List[str]:
"""
Filter and transform beta headers based on provider's mapping configuration.
This function:
1. Only allows headers that are present in the provider's mapping keys
2. Filters out headers with null values (unsupported)
3. Maps headers to provider-specific names (e.g., advanced-tool-use -> tool-search-tool)
Args:
beta_headers: List of Anthropic beta header values
provider: Provider name (e.g., "anthropic", "bedrock", "vertex_ai")
Returns:
List of filtered and transformed beta headers for the provider
"""
if not beta_headers:
return []
config = _load_beta_headers_config()
provider = get_provider_name(provider)
# Get the header mapping for this provider
provider_mapping = config.get(provider, {})
filtered_headers: Set[str] = set()
for header in beta_headers:
header = header.strip()
# Check if header is in the mapping
if header not in provider_mapping:
verbose_logger.debug(
f"Dropping unknown beta header '{header}' for provider '{provider}' (not in mapping)"
)
continue
# Get the mapped header value
mapped_header = provider_mapping[header]
# Skip if header is unsupported (null value)
if mapped_header is None:
verbose_logger.debug(
f"Dropping unsupported beta header '{header}' for provider '{provider}'"
)
continue
# Add the mapped header
filtered_headers.add(mapped_header)
return sorted(list(filtered_headers))
@@ -275,18 +283,18 @@ def is_beta_header_supported(
) -> bool:
"""
Check if a specific beta header is supported by a provider.
Args:
beta_header: The Anthropic beta header value
provider: Provider name
Returns:
True if the header is in the mapping with a non-null value, False otherwise
"""
config = _load_beta_headers_config()
provider = get_provider_name(provider)
provider_mapping = config.get(provider, {})
# Header is supported if it's in the mapping and has a non-null value
return beta_header in provider_mapping and provider_mapping[beta_header] is not None
@@ -297,26 +305,26 @@ def get_provider_beta_header(
) -> Optional[str]:
"""
Get the provider-specific beta header name for a given Anthropic beta header.
This function handles header transformations/mappings (e.g., advanced-tool-use -> tool-search-tool).
Args:
anthropic_beta_header: The Anthropic beta header value
provider: Provider name
Returns:
The provider-specific header name if supported, or None if unsupported/unknown
"""
config = _load_beta_headers_config()
provider = get_provider_name(provider)
# Get the header mapping for this provider
provider_mapping = config.get(provider, {})
# Check if header is in the mapping
if anthropic_beta_header not in provider_mapping:
return None
# Return the mapped value (could be None if unsupported)
return provider_mapping[anthropic_beta_header]
@@ -328,50 +336,50 @@ def update_headers_with_filtered_beta(
"""
Update headers dict by filtering and transforming anthropic-beta header values.
Modifies the headers dict in place and returns it.
Args:
headers: Request headers dict (will be modified in place)
provider: Provider name
Returns:
Updated headers dict
"""
existing_beta = headers.get("anthropic-beta")
if not existing_beta:
return headers
# Parse existing beta headers
beta_values = [b.strip() for b in existing_beta.split(",") if b.strip()]
# Filter and transform based on provider
filtered_beta_values = filter_and_transform_beta_headers(
beta_headers=beta_values,
provider=provider,
)
# Update or remove the header
if filtered_beta_values:
headers["anthropic-beta"] = ",".join(filtered_beta_values)
else:
# Remove the header if no values remain
headers.pop("anthropic-beta", None)
return headers
def get_unsupported_headers(provider: str) -> List[str]:
"""
Get all beta headers that are unsupported by a provider (have null values in mapping).
Args:
provider: Provider name
Returns:
List of unsupported Anthropic beta header names
"""
config = _load_beta_headers_config()
provider = get_provider_name(provider)
provider_mapping = config.get(provider, {})
# Return headers with null values
return [header for header, value in provider_mapping.items() if value is None]
@@ -149,7 +149,9 @@ class AnthropicExceptionMapping:
parsed = None
# If parsed and already in Anthropic format - passthrough
if parsed is not None and AnthropicExceptionMapping._is_anthropic_error_dict(parsed):
if parsed is not None and AnthropicExceptionMapping._is_anthropic_error_dict(
parsed
):
# Optionally add request_id if provided and not present
if request_id and "request_id" not in parsed:
parsed["request_id"] = request_id
@@ -157,7 +159,9 @@ class AnthropicExceptionMapping:
# Extract message - use parsed dict if available, otherwise raw string
if parsed is not None:
message = AnthropicExceptionMapping._extract_message_from_dict(parsed, raw_message)
message = AnthropicExceptionMapping._extract_message_from_dict(
parsed, raw_message
)
else:
message = raw_message
+75 -36
View File
@@ -10,7 +10,9 @@ from litellm.utils import token_counter
async def calculate_batch_cost_and_usage(
file_content_dictionary: List[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"
],
model_name: Optional[str] = None,
model_info: Optional[ModelInfo] = None,
) -> Tuple[float, Usage, List[str]]:
@@ -34,19 +36,23 @@ async def calculate_batch_cost_and_usage(
custom_llm_provider=custom_llm_provider,
model_name=model_name,
)
batch_models = _get_batch_models_from_file_content(file_content_dictionary, model_name)
batch_models = _get_batch_models_from_file_content(
file_content_dictionary, model_name
)
return batch_cost, batch_usage, batch_models
async def _handle_completed_batch(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"
],
model_name: Optional[str] = None,
litellm_params: Optional[dict] = None,
) -> Tuple[float, Usage, List[str]]:
"""Helper function to process a completed batch and handle logging
Args:
batch: The batch object
custom_llm_provider: The LLM provider
@@ -70,7 +76,9 @@ async def _handle_completed_batch(
model_name=model_name,
)
batch_models = _get_batch_models_from_file_content(file_content_dictionary, model_name)
batch_models = _get_batch_models_from_file_content(
file_content_dictionary, model_name
)
return batch_cost, batch_usage, batch_models
@@ -96,7 +104,9 @@ def _get_batch_models_from_file_content(
def _batch_cost_calculator(
file_content_dictionary: List[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"
] = "openai",
model_name: Optional[str] = None,
model_info: Optional[ModelInfo] = None,
) -> float:
@@ -105,10 +115,12 @@ def _batch_cost_calculator(
"""
# Handle Vertex AI with specialized method
if custom_llm_provider == "vertex_ai" and model_name:
batch_cost, _ = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
batch_cost, _ = calculate_vertex_ai_batch_cost_and_usage(
file_content_dictionary, model_name
)
verbose_logger.debug("vertex_ai_total_cost=%s", batch_cost)
return batch_cost
# For other providers, use the existing logic
total_cost = _get_batch_job_cost_from_file_content(
file_content_dictionary=file_content_dictionary,
@@ -173,7 +185,10 @@ def calculate_vertex_ai_batch_cost_and_usage(
verbose_logger.info(
"vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d",
total_cost, prompt_tokens, completion_tokens, total_tokens,
total_cost,
prompt_tokens,
completion_tokens,
total_tokens,
)
return total_cost, Usage(
@@ -185,12 +200,14 @@ def calculate_vertex_ai_batch_cost_and_usage(
async def _get_batch_output_file_content_as_dictionary(
batch: Batch,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"
] = "openai",
litellm_params: Optional[dict] = None,
) -> List[dict]:
"""
Get the batch output file content as a list of dictionaries
Args:
batch: The batch object
custom_llm_provider: The LLM provider
@@ -198,8 +215,9 @@ async def _get_batch_output_file_content_as_dictionary(
Required for Azure and other providers that need authentication
"""
from litellm.files.main import afile_content
from litellm.proxy.openai_files_endpoints.common_utils import \
_is_base64_encoded_unified_file_id
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
if custom_llm_provider == "vertex_ai":
raise ValueError("Vertex AI does not support file content retrieval")
@@ -211,21 +229,27 @@ async def _get_batch_output_file_content_as_dictionary(
is_base64_unified_file_id = _is_base64_encoded_unified_file_id(file_id)
if is_base64_unified_file_id:
try:
file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(";")[0]
verbose_logger.debug(f"Extracted LLM output file ID from unified file ID: {file_id}")
file_id = is_base64_unified_file_id.split("llm_output_file_id,")[1].split(
";"
)[0]
verbose_logger.debug(
f"Extracted LLM output file ID from unified file ID: {file_id}"
)
except (IndexError, AttributeError) as e:
verbose_logger.error(f"Failed to extract LLM output file ID from unified file ID: {batch.output_file_id}, error: {e}")
verbose_logger.error(
f"Failed to extract LLM output file ID from unified file ID: {batch.output_file_id}, error: {e}"
)
# Build kwargs for afile_content with credentials from litellm_params
file_content_kwargs = {
"file_id": file_id,
"custom_llm_provider": custom_llm_provider,
}
# Extract and add credentials for file access
credentials = _extract_file_access_credentials(litellm_params)
file_content_kwargs.update(credentials)
_file_content = await afile_content(**file_content_kwargs) # type: ignore[reportArgumentType]
return _get_file_content_as_dictionary(_file_content.content)
@@ -233,30 +257,37 @@ async def _get_batch_output_file_content_as_dictionary(
def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict:
"""
Extract credentials from litellm_params for file access operations.
This method extracts relevant authentication and configuration parameters
needed for accessing files across different providers (Azure, Vertex AI, etc.).
Args:
litellm_params: Dictionary containing litellm parameters with credentials
Returns:
Dictionary containing only the credentials needed for file access
"""
credentials = {}
if litellm_params:
# List of credential keys that should be passed to file operations
credential_keys = [
"api_key", "api_base", "api_version", "organization",
"azure_ad_token", "azure_ad_token_provider",
"vertex_project", "vertex_location", "vertex_credentials",
"timeout", "max_retries"
"api_key",
"api_base",
"api_version",
"organization",
"azure_ad_token",
"azure_ad_token_provider",
"vertex_project",
"vertex_location",
"vertex_credentials",
"timeout",
"max_retries",
]
for key in credential_keys:
if key in litellm_params:
credentials[key] = litellm_params[key]
return credentials
@@ -279,7 +310,9 @@ def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]:
def _get_batch_job_cost_from_file_content(
file_content_dictionary: List[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"
] = "openai",
model_info: Optional[ModelInfo] = None,
) -> float:
"""
@@ -321,7 +354,9 @@ def _get_batch_job_cost_from_file_content(
def _get_batch_job_total_usage_from_file_content(
file_content_dictionary: List[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"
] = "openai",
model_name: Optional[str] = None,
) -> Usage:
"""
@@ -329,9 +364,11 @@ def _get_batch_job_total_usage_from_file_content(
"""
# Handle Vertex AI with specialized method
if custom_llm_provider == "vertex_ai" and model_name:
_, batch_usage = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
_, batch_usage = calculate_vertex_ai_batch_cost_and_usage(
file_content_dictionary, model_name
)
return batch_usage
# For other providers, use the existing logic
total_tokens: int = 0
prompt_tokens: int = 0
@@ -349,6 +386,7 @@ def _get_batch_job_total_usage_from_file_content(
completion_tokens=completion_tokens,
)
def _get_batch_job_input_file_usage(
file_content_dictionary: List[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
@@ -358,25 +396,26 @@ def _get_batch_job_input_file_usage(
Count the number of tokens in the input file
Used for batch rate limiting to count the number of tokens in the input file
"""
"""
prompt_tokens: int = 0
completion_tokens: int = 0
for _item in file_content_dictionary:
body = _item.get("body", {})
model = body.get("model", model_name or "")
messages = body.get("messages", [])
if messages:
item_tokens = token_counter(model=model, messages=messages)
prompt_tokens += item_tokens
return Usage(
total_tokens=prompt_tokens + completion_tokens,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
)
def _get_batch_job_usage_from_response_body(response_body: dict) -> Usage:
"""
Get the tokens of a batch job from the response body
@@ -400,4 +439,4 @@ def _batch_response_was_successful(batch_job_output_file: dict) -> bool:
Check if the batch job response status == 200
"""
_response: dict = batch_job_output_file.get("response", None) or {}
return _response.get("status_code", None) == 200
return _response.get("status_code", None) == 200
+38 -12
View File
@@ -109,7 +109,9 @@ async def acreate_batch(
completion_window: Literal["24h"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
input_file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai",
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"
] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@@ -159,7 +161,9 @@ def create_batch( # noqa: PLR0915
completion_window: Literal["24h"],
endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"],
input_file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai",
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"
] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@@ -220,7 +224,9 @@ def create_batch( # noqa: PLR0915
extra_body=extra_body,
)
if output_expires_after is not None:
_create_batch_request["output_expires_after"] = cast(FileExpiresAfter, output_expires_after)
_create_batch_request["output_expires_after"] = cast(
FileExpiresAfter, output_expires_after
)
if model is not None:
provider_config = ProviderConfigManager.get_provider_batches_config(
model=model,
@@ -364,7 +370,9 @@ def create_batch( # noqa: PLR0915
@client
async def aretrieve_batch(
batch_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"
] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@@ -410,7 +418,9 @@ def _handle_retrieve_batch_providers_without_provider_config(
litellm_params: dict,
_retrieve_batch_request: RetrieveBatchRequest,
_is_async: bool,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"
] = "openai",
logging_obj: Optional[Any] = None,
):
api_base: Optional[str] = None
@@ -549,7 +559,9 @@ def _handle_retrieve_batch_providers_without_provider_config(
@client
def retrieve_batch(
batch_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"] = "openai",
custom_llm_provider: Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic"
] = "openai",
metadata: Optional[Dict[str, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@@ -929,7 +941,6 @@ def cancel_batch(
LiteLLM Equivalent of POST https://api.openai.com/v1/batches/{batch_id}/cancel
"""
try:
try:
if model is not None:
_, custom_llm_provider, _, _ = get_llm_provider(
@@ -1097,25 +1108,40 @@ def _handle_async_invoke_status(
"inprogress": "in_progress",
"in_progress": "in_progress",
}
normalized_status: BatchJobStatus = status_mapping.get(aws_status_lower, "failed") # Default to "failed" if unknown status
normalized_status: BatchJobStatus = status_mapping.get(
aws_status_lower, "failed"
) # Default to "failed" if unknown status
# Get output S3 URI safely
output_s3_uri = ""
try:
output_s3_uri = status_response["outputDataConfig"]["s3OutputDataConfig"]["s3Uri"]
output_s3_uri = status_response["outputDataConfig"]["s3OutputDataConfig"][
"s3Uri"
]
except (KeyError, TypeError):
pass
# Use BedrockBatchesConfig's timestamp parsing method (expects raw AWS status string)
import time
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
created_at, in_progress_at, completed_at, failed_at, _, _ = BedrockBatchesConfig()._parse_timestamps_and_status(status_response, aws_status_raw)
(
created_at,
in_progress_at,
completed_at,
failed_at,
_,
_,
) = BedrockBatchesConfig()._parse_timestamps_and_status(
status_response, aws_status_raw
)
result = LiteLLMBatch(
id=status_response["invocationArn"],
object="batch",
status=normalized_status,
created_at=created_at or int(time.time()), # Provide default timestamp if None
created_at=created_at
or int(time.time()), # Provide default timestamp if None
in_progress_at=in_progress_at,
completed_at=completed_at,
failed_at=failed_at,
+8 -4
View File
@@ -22,7 +22,9 @@ class AzureBlobCache(BaseCache):
from azure.storage.blob import BlobServiceClient
from azure.core.exceptions import ResourceExistsError
from azure.identity import DefaultAzureCredential
from azure.identity.aio import DefaultAzureCredential as AsyncDefaultAzureCredential
from azure.identity.aio import (
DefaultAzureCredential as AsyncDefaultAzureCredential,
)
from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient
self.container_client = BlobServiceClient(
@@ -50,14 +52,16 @@ class AzureBlobCache(BaseCache):
print_verbose(f"LiteLLM SET Cache - Azure Blob. Key={key}. Value={value}")
serialized_value = json.dumps(value)
try:
await self.async_container_client.upload_blob(key, serialized_value, overwrite=True)
await self.async_container_client.upload_blob(
key, serialized_value, overwrite=True
)
except Exception as e:
# NON blocking - notify users Azure Blob is throwing an exception
print_verbose(f"LiteLLM set_cache() - Got exception from Azure Blob: {e}")
def get_cache(self, key, **kwargs):
from azure.core.exceptions import ResourceNotFoundError
try:
print_verbose(f"Get Azure Blob Cache: key: {key}")
as_bytes = self.container_client.download_blob(key).readall()
@@ -74,7 +78,7 @@ class AzureBlobCache(BaseCache):
async def async_get_cache(self, key, **kwargs):
from azure.core.exceptions import ResourceNotFoundError
try:
print_verbose(f"Get Azure Blob Cache: key: {key}")
blob = await self.async_container_client.download_blob(key)
+3 -3
View File
@@ -53,12 +53,12 @@ class BaseCache(ABC):
async def disconnect(self):
raise NotImplementedError
async def test_connection(self) -> dict:
"""
Test the cache connection.
Returns:
dict: {"status": "success" | "failed", "message": str, "error": Optional[str]}
"""
raise NotImplementedError
raise NotImplementedError
+23 -20
View File
@@ -78,9 +78,7 @@ class CachingHandlerResponse(BaseModel):
cached_result: Optional[Any] = None
final_embedding_cached_response: Optional[EmbeddingResponse] = None
embedding_all_elements_cache_hit: bool = (
False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call
)
embedding_all_elements_cache_hit: bool = False # this is set to True when all elements in the list have a cache hit in the embedding cache, if true return the final_embedding_cached_response no need to make an API call
in_memory_cache_obj = InMemoryCache()
@@ -159,7 +157,7 @@ class LLMCachingHandler:
#########################################################
parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs)
kwargs["parent_otel_span"] = parent_otel_span
if litellm.cache is not None and self._is_call_type_supported_by_cache(
original_function=original_function
):
@@ -181,7 +179,9 @@ class LLMCachingHandler:
api_base=kwargs.get("api_base", None),
api_key=kwargs.get("api_key", None),
)
cache_duration_ms = (cache_check_end_time - cache_check_start_time) * 1000
cache_duration_ms = (
cache_check_end_time - cache_check_start_time
) * 1000
self._update_litellm_logging_obj_environment(
logging_obj=logging_obj,
model=model,
@@ -194,7 +194,6 @@ class LLMCachingHandler:
call_type = original_function.__name__
cached_result = self._convert_cached_result_to_model_response(
cached_result=cached_result,
call_type=call_type,
@@ -244,7 +243,7 @@ class LLMCachingHandler:
final_embedding_cached_response=final_embedding_cached_response,
embedding_all_elements_cache_hit=embedding_all_elements_cache_hit,
)
verbose_logger.debug(f"CACHE RESULT: {cached_result}")
return CachingHandlerResponse(
cached_result=cached_result,
@@ -265,9 +264,8 @@ class LLMCachingHandler:
) -> CachingHandlerResponse:
from litellm.utils import CustomStreamWrapper
cached_result: Optional[Any] = None
# Check if caching should be performed BEFORE doing expensive kwargs copy
if litellm.cache is not None and self._is_call_type_supported_by_cache(
original_function=original_function
@@ -325,7 +323,7 @@ class LLMCachingHandler:
result=cached_result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit
cache_hit=cache_hit,
)
cache_key = litellm.cache.get_cache_key(**kwargs)
if (
@@ -554,12 +552,18 @@ class LLMCachingHandler:
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
async_coroutine=logging_obj.async_success_handler(
result=cached_result, start_time=start_time, end_time=end_time, cache_hit=cache_hit
result=cached_result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
)
)
logging_obj.handle_sync_success_callbacks_for_async_calls(
result=cached_result, start_time=start_time, end_time=end_time, cache_hit=cache_hit
result=cached_result,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
)
async def _retrieve_from_cache(
@@ -728,10 +732,9 @@ class LLMCachingHandler:
response_type="audio_transcription",
hidden_params=hidden_params,
)
elif (
call_type == "aresponses"
or call_type == "responses"
) and isinstance(cached_result, dict):
elif (call_type == "aresponses" or call_type == "responses") and isinstance(
cached_result, dict
):
# Convert cached dict back to ResponsesAPIResponse object
cached_result = ResponsesAPIResponse(**cached_result)
@@ -741,7 +744,7 @@ class LLMCachingHandler:
and isinstance(cached_result._hidden_params, dict)
):
cached_result._hidden_params["cache_hit"] = True
#########################################################
# Add final timing metrics to the cached result
#########################################################
@@ -1011,9 +1014,9 @@ class LLMCachingHandler:
}
if litellm.cache is not None:
litellm_params["preset_cache_key"] = (
litellm.cache._get_preset_cache_key_from_kwargs(**kwargs)
)
litellm_params[
"preset_cache_key"
] = litellm.cache._get_preset_cache_key_from_kwargs(**kwargs)
else:
litellm_params["preset_cache_key"] = None
+10 -4
View File
@@ -319,18 +319,20 @@ class DualCache(BaseCache):
previous_access_times
)
raise
# Short-circuit if redis_result is None or contains only None values
if redis_result is None or all(v is None for v in redis_result.values()):
if redis_result is None or all(
v is None for v in redis_result.values()
):
return result
# Pre-compute key-to-index mapping for O(1) lookup
key_to_index = {key: i for i, key in enumerate(keys)}
# Update both result and in-memory cache in a single loop
for key, value in redis_result.items():
result[key_to_index[key]] = value
if value is not None and self.in_memory_cache is not None:
await self.in_memory_cache.async_set_cache(
key, value, **kwargs
@@ -346,6 +348,8 @@ class DualCache(BaseCache):
)
try:
if self.in_memory_cache is not None:
if "ttl" not in kwargs and self.default_in_memory_ttl is not None:
kwargs["ttl"] = self.default_in_memory_ttl
await self.in_memory_cache.async_set_cache(key, value, **kwargs)
if self.redis_cache is not None and local_only is False:
@@ -367,6 +371,8 @@ class DualCache(BaseCache):
)
try:
if self.in_memory_cache is not None:
if "ttl" not in kwargs and self.default_in_memory_ttl is not None:
kwargs["ttl"] = self.default_in_memory_ttl
await self.in_memory_cache.async_set_cache_pipeline(
cache_list=cache_list, **kwargs
)
+22 -6
View File
@@ -16,13 +16,23 @@ from .base_cache import BaseCache
class GCSCache(BaseCache):
def __init__(self, bucket_name: Optional[str] = None, path_service_account: Optional[str] = None, gcs_path: Optional[str] = None) -> None:
def __init__(
self,
bucket_name: Optional[str] = None,
path_service_account: Optional[str] = None,
gcs_path: Optional[str] = None,
) -> None:
super().__init__()
self.bucket_name = bucket_name or GCSBucketBase(bucket_name=None).BUCKET_NAME
self.path_service_account = path_service_account or GCSBucketBase(bucket_name=None).path_service_account_json
self.path_service_account = (
path_service_account
or GCSBucketBase(bucket_name=None).path_service_account_json
)
self.key_prefix = gcs_path.rstrip("/") + "/" if gcs_path else ""
# create httpx clients
self.async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
self.async_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
self.sync_client = _get_httpx_client()
def _construct_headers(self) -> dict:
@@ -52,7 +62,9 @@ class GCSCache(BaseCache):
data = json.dumps(value)
await self.async_client.post(url=url, data=data, headers=headers)
except Exception as e:
print_verbose(f"GCS Caching: async_set_cache() - Got exception from GCS: {e}")
print_verbose(
f"GCS Caching: async_set_cache() - Got exception from GCS: {e}"
)
def get_cache(self, key, **kwargs):
try:
@@ -69,7 +81,9 @@ class GCSCache(BaseCache):
return cached_response
return None
except Exception as e:
verbose_logger.error(f"GCS Caching: get_cache() - Got exception from GCS: {e}")
verbose_logger.error(
f"GCS Caching: get_cache() - Got exception from GCS: {e}"
)
async def async_get_cache(self, key, **kwargs):
try:
@@ -82,7 +96,9 @@ class GCSCache(BaseCache):
return json.loads(response.text)
return None
except Exception as e:
verbose_logger.error(f"GCS Caching: async_get_cache() - Got exception from GCS: {e}")
verbose_logger.error(
f"GCS Caching: async_get_cache() - Got exception from GCS: {e}"
)
def flush_cache(self):
pass
+3 -1
View File
@@ -54,7 +54,9 @@ class QdrantSemanticCache(BaseCache):
raise Exception("similarity_threshold must be provided, passed None")
self.similarity_threshold = similarity_threshold
self.embedding_model = embedding_model
self.vector_size = vector_size if vector_size is not None else QDRANT_VECTOR_SIZE
self.vector_size = (
vector_size if vector_size is not None else QDRANT_VECTOR_SIZE
)
headers = {}
# check if defined as os.environ/ variable
+11 -14
View File
@@ -268,19 +268,19 @@ class RedisCache(BaseCache):
def _parse_redis_major_version(self) -> int:
"""
Parse Redis version to extract the major version number.
Handles multiple version formats:
- Strings: "7.0.0", "6", "7.0.0-rc1", " 7.0.0 "
- Floats: 7.0 (e.g., from AWS ElastiCache Valkey)
- Integers: 7
- Malformed: "latest", "", "Unknown" (defaults to DEFAULT_REDIS_MAJOR_VERSION)
Returns:
int: The major version number (defaults to DEFAULT_REDIS_MAJOR_VERSION if unparseable)
"""
if self.redis_version == "Unknown":
return DEFAULT_REDIS_MAJOR_VERSION
try:
version_str = str(self.redis_version).strip()
# Handle cases where there's no dot (e.g., "7" or 7)
@@ -1113,14 +1113,14 @@ class RedisCache(BaseCache):
self.redis_client.close()
except Exception as e:
verbose_logger.debug("Error closing sync Redis client: %s", e)
async def test_connection(self) -> dict:
"""
Test the Redis connection by creating a new client and pinging it.
This creates a fresh connection without using cached clients or connection pools
to ensure the credentials are actually valid.
Returns:
dict: {"status": "success" | "failed", "message": str, "error": Optional[str]}
"""
@@ -1129,29 +1129,26 @@ class RedisCache(BaseCache):
# Create a fresh Redis client with current settings
redis_client = redis_async.Redis(**self.redis_kwargs)
# Test the connection
ping_result = await redis_client.ping() # type: ignore[misc]
# Close the connection
await redis_client.aclose() # type: ignore[attr-defined]
if ping_result:
return {
"status": "success",
"message": "Redis connection test successful"
"message": "Redis connection test successful",
}
else:
return {
"status": "failed",
"message": "Redis ping returned False"
}
return {"status": "failed", "message": "Redis ping returned False"}
except Exception as e:
verbose_logger.error(f"Redis connection test failed: {str(e)}")
return {
"status": "failed",
"message": f"Redis connection failed: {str(e)}",
"error": str(e)
"error": str(e),
}
async def async_delete_cache(self, key: str):
+11 -10
View File
@@ -57,11 +57,11 @@ class RedisClusterCache(RedisCache):
"""
async_redis_cluster_client = self.init_async_client()
return await async_redis_cluster_client.mget_nonatomic(keys=keys) # type: ignore
async def test_connection(self) -> dict:
"""
Test the Redis Cluster connection.
Returns:
dict: {"status": "success" | "failed", "message": str, "error": Optional[str]}
"""
@@ -72,37 +72,38 @@ class RedisClusterCache(RedisCache):
# Create ClusterNode objects from startup_nodes
cluster_kwargs = self.redis_kwargs.copy()
startup_nodes = cluster_kwargs.pop("startup_nodes", [])
new_startup_nodes: List[ClusterNode] = []
for item in startup_nodes:
new_startup_nodes.append(ClusterNode(**item))
# Create a fresh Redis Cluster client with current settings
redis_client = redis_async.RedisCluster(
startup_nodes=new_startup_nodes, **cluster_kwargs # type: ignore
)
# Test the connection
ping_result = await redis_client.ping() # type: ignore[attr-defined, misc]
# Close the connection
await redis_client.aclose() # type: ignore[attr-defined]
if ping_result:
return {
"status": "success",
"message": "Redis Cluster connection test successful"
"message": "Redis Cluster connection test successful",
}
else:
return {
"status": "failed",
"message": "Redis Cluster ping returned False"
"message": "Redis Cluster ping returned False",
}
except Exception as e:
from litellm._logging import verbose_logger
verbose_logger.error(f"Redis Cluster connection test failed: {str(e)}")
return {
"status": "failed",
"message": f"Redis Cluster connection failed: {str(e)}",
"error": str(e)
}
"error": str(e),
}
+4 -2
View File
@@ -110,7 +110,9 @@ class S3Cache(BaseCache):
func = partial(self.set_cache, key, value, **kwargs)
await loop.run_in_executor(None, func)
except Exception as e:
verbose_logger.error(f"S3 Caching: async_set_cache() - Got exception from S3: {e}")
verbose_logger.error(
f"S3 Caching: async_set_cache() - Got exception from S3: {e}"
)
def get_cache(self, key, **kwargs):
import botocore
@@ -126,7 +128,7 @@ class S3Cache(BaseCache):
if cached_response is not None:
if "Expires" in cached_response:
expires_time = cached_response['Expires']
expires_time = cached_response["Expires"]
current_time = datetime.now(expires_time.tzinfo)
if current_time > expires_time:
@@ -61,9 +61,7 @@ class ResponsesToCompletionBridgeHandler:
existing.setdefault(key, value)
return response
def _collect_response_from_stream(
self, stream_iter: Any
) -> "ResponsesAPIResponse":
def _collect_response_from_stream(self, stream_iter: Any) -> "ResponsesAPIResponse":
for _ in stream_iter:
pass
@@ -144,7 +142,9 @@ class ResponsesToCompletionBridgeHandler:
custom_llm_provider=custom_llm_provider,
)
def completion(self, *args, **kwargs) -> Union[
def completion(
self, *args, **kwargs
) -> Union[
Coroutine[Any, Any, Union["ModelResponse", "CustomStreamWrapper"]],
"ModelResponse",
"CustomStreamWrapper",
@@ -63,7 +63,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
def __init__(self):
pass
def _handle_raw_dict_response_item(self, item: Dict[str, Any], index: int) -> Tuple[Optional[Any], int]:
def _handle_raw_dict_response_item(
self, item: Dict[str, Any], index: int
) -> Tuple[Optional[Any], int]:
"""
Handle raw dict response items from Responses API (e.g., GPT-5 Codex format).
@@ -106,9 +108,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if item_type == "function_call":
# Extract provider_specific_fields if present and pass through as-is
provider_specific_fields = item.get("provider_specific_fields")
if provider_specific_fields and not isinstance(provider_specific_fields, dict):
if provider_specific_fields and not isinstance(
provider_specific_fields, dict
):
provider_specific_fields = (
dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {}
dict(provider_specific_fields)
if hasattr(provider_specific_fields, "__dict__")
else {}
)
tool_call_dict = {
@@ -124,7 +130,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if provider_specific_fields:
tool_call_dict["provider_specific_fields"] = provider_specific_fields
# Also add to function's provider_specific_fields for consistency
tool_call_dict["function"]["provider_specific_fields"] = provider_specific_fields
tool_call_dict["function"][
"provider_specific_fields"
] = provider_specific_fields
msg = Message(
content=None,
@@ -232,10 +240,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if key in ("max_tokens", "max_completion_tokens"):
responses_api_request["max_output_tokens"] = value
elif key == "tools" and value is not None:
responses_api_request["tools"] = (
self._convert_tools_to_responses_format(
cast(List[Dict[str, Any]], value)
)
responses_api_request[
"tools"
] = self._convert_tools_to_responses_format(
cast(List[Dict[str, Any]], value)
)
elif key == "response_format":
text_format = self._transform_response_format_to_text_format(value)
@@ -250,9 +258,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
elif key == "web_search_options":
self._add_web_search_tool(responses_api_request, value)
def _build_sanitized_litellm_params(
self, litellm_params: dict
) -> Dict[str, Any]:
def _build_sanitized_litellm_params(self, litellm_params: dict) -> Dict[str, Any]:
"""Build sanitized litellm_params with merged metadata."""
responses_optional_param_keys = set(
ResponsesAPIOptionalRequestParams.__annotations__.keys()
@@ -337,7 +343,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
previous_response_id = optional_params.get("previous_response_id")
if previous_response_id:
# Use the existing session handler for responses API
verbose_logger.debug(f"Chat provider: Warning ignoring previous response ID: {previous_response_id}")
verbose_logger.debug(
f"Chat provider: Warning ignoring previous response ID: {previous_response_id}"
)
# Convert back to responses API format for the actual request
@@ -347,9 +355,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
setattr(litellm_logging_obj, "call_type", CallTypes.responses.value)
sanitized_litellm_params = self._build_sanitized_litellm_params(
litellm_params
)
sanitized_litellm_params = self._build_sanitized_litellm_params(litellm_params)
request_data = {
"model": api_model,
@@ -359,7 +365,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
"client": client,
}
verbose_logger.debug(f"Chat provider: Final request model={api_model}, input_items={len(input_items)}")
verbose_logger.debug(
f"Chat provider: Final request model={api_model}, input_items={len(input_items)}"
)
self._merge_responses_api_request_into_request_data(
request_data, responses_api_request, instructions
@@ -390,6 +398,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
ResponseOutputMessage,
ResponseReasoningItem,
)
from openai.types.responses.response_output_item import (
ResponseApplyPatchToolCall,
)
from litellm.types.utils import Choices, Message
@@ -439,11 +450,21 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
LiteLLMCompletionResponsesConfig,
)
tool_call_dict = (
LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call(
tool_call_item=item,
index=tool_call_index,
)
tool_call_dict = LiteLLMCompletionResponsesConfig.convert_response_function_tool_call_to_chat_completion_tool_call(
tool_call_item=item,
index=tool_call_index,
)
accumulated_tool_calls.append(tool_call_dict)
tool_call_index += 1
elif isinstance(item, ResponseApplyPatchToolCall):
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
tool_call_dict = LiteLLMCompletionResponsesConfig.convert_apply_patch_tool_call_to_chat_completion_tool_call(
tool_call_item=item,
index=tool_call_index,
)
accumulated_tool_calls.append(tool_call_dict)
tool_call_index += 1
@@ -463,7 +484,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
tool_calls=accumulated_tool_calls,
reasoning_content=reasoning_content,
)
choices.append(Choices(message=msg, finish_reason="tool_calls", index=index))
choices.append(
Choices(message=msg, finish_reason="tool_calls", index=index)
)
reasoning_content = None
return choices
@@ -499,10 +522,17 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
)
if len(choices) == 0:
if raw_response.incomplete_details is not None and raw_response.incomplete_details.reason is not None:
raise ValueError(f"{model} unable to complete request: {raw_response.incomplete_details.reason}")
if (
raw_response.incomplete_details is not None
and raw_response.incomplete_details.reason is not None
):
raise ValueError(
f"{model} unable to complete request: {raw_response.incomplete_details.reason}"
)
else:
raise ValueError(f"Unknown items in responses API response: {raw_response.output}")
raise ValueError(
f"Unknown items in responses API response: {raw_response.output}"
)
setattr(model_response, "choices", choices)
@@ -511,21 +541,28 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
setattr(
model_response,
"usage",
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(raw_response.usage),
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
raw_response.usage
),
)
# Preserve hidden params from the ResponsesAPIResponse, especially the headers
# which contain important provider information like x-request-id
raw_response_hidden_params = getattr(raw_response, "_hidden_params", {})
if raw_response_hidden_params:
if not hasattr(model_response, "_hidden_params") or model_response._hidden_params is None:
if (
not hasattr(model_response, "_hidden_params")
or model_response._hidden_params is None
):
model_response._hidden_params = {}
# Merge the raw_response hidden params with model_response hidden params
# Preserve existing keys in model_response but add/override with raw_response params
for key, value in raw_response_hidden_params.items():
if key == "additional_headers" and key in model_response._hidden_params:
# Merge additional_headers to preserve both sets
existing_additional_headers = model_response._hidden_params.get("additional_headers", {})
existing_additional_headers = model_response._hidden_params.get(
"additional_headers", {}
)
merged_headers = {**value, **existing_additional_headers}
model_response._hidden_params[key] = merged_headers
else:
@@ -535,13 +572,19 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
def get_model_response_iterator(
self,
streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel"],
streaming_response: Union[
Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel"
],
sync_stream: bool,
json_mode: Optional[bool] = False,
) -> BaseModelResponseIterator:
return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode)
return OpenAiResponsesToChatCompletionStreamIterator(
streaming_response, sync_stream, json_mode
)
def _convert_content_str_to_input_text(self, content: str, role: str) -> Dict[str, Any]:
def _convert_content_str_to_input_text(
self, content: str, role: str
) -> Dict[str, Any]:
if role == "user" or role == "system" or role == "tool":
return {"type": "input_text", "text": content}
else:
@@ -568,7 +611,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if actual_image_url is None:
raise ValueError(f"Invalid image URL: {content_image_url}")
image_param = ResponseInputImageParam(image_url=actual_image_url, detail="auto", type="input_image")
image_param = ResponseInputImageParam(
image_url=actual_image_url, detail="auto", type="input_image"
)
if detail:
image_param["detail"] = detail
@@ -581,7 +626,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
Union[
str,
List[Any],
Iterable[Union["OpenAIMessageContentListBlock", "ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]],
Iterable[
Union[
"OpenAIMessageContentListBlock",
"ChatCompletionThinkingBlock",
"ChatCompletionRedactedThinkingBlock",
]
],
]
],
role: str,
@@ -589,7 +640,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
"""Convert chat completion content to responses API format"""
from litellm.types.llms.openai import ChatCompletionImageObject
verbose_logger.debug(f"Chat provider: Converting content to responses format - input type: {type(content)}")
verbose_logger.debug(
f"Chat provider: Converting content to responses format - input type: {type(content)}"
)
if content is None:
return [self._convert_content_str_to_input_text("", role)]
@@ -600,7 +653,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
elif isinstance(content, list):
result = []
for i, item in enumerate(content):
verbose_logger.debug(f"Chat provider: Processing content item {i}: {type(item)} = {item}")
verbose_logger.debug(
f"Chat provider: Processing content item {i}: {type(item)} = {item}"
)
if isinstance(item, str):
converted = self._convert_content_str_to_input_text(item, role)
result.append(converted)
@@ -609,7 +664,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
# Handle multimodal content
original_type = item.get("type")
if original_type == "text":
converted = self._convert_content_str_to_input_text(item.get("text", ""), role)
converted = self._convert_content_str_to_input_text(
item.get("text", ""), role
)
result.append(converted)
verbose_logger.debug(f"Chat provider: text -> {converted}")
elif original_type == "image_url":
@@ -621,14 +678,18 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
),
)
result.append(converted)
verbose_logger.debug(f"Chat provider: image_url -> {converted}")
verbose_logger.debug(
f"Chat provider: image_url -> {converted}"
)
else:
# Try to map other types to responses API format
item_type = original_type or "input_text"
if item_type == "image":
converted = {"type": "input_image", **item}
result.append(converted)
verbose_logger.debug(f"Chat provider: image -> {converted}")
verbose_logger.debug(
f"Chat provider: image -> {converted}"
)
elif item_type in [
"input_text",
"input_image",
@@ -640,12 +701,18 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
]:
# Already in responses API format
result.append(item)
verbose_logger.debug(f"Chat provider: passthrough -> {item}")
verbose_logger.debug(
f"Chat provider: passthrough -> {item}"
)
else:
# Default to input_text for unknown types
converted = self._convert_content_str_to_input_text(str(item.get("text", item)), role)
converted = self._convert_content_str_to_input_text(
str(item.get("text", item)), role
)
result.append(converted)
verbose_logger.debug(f"Chat provider: unknown({original_type}) -> {converted}")
verbose_logger.debug(
f"Chat provider: unknown({original_type}) -> {converted}"
)
verbose_logger.debug(f"Chat provider: Final converted content: {result}")
return result
else:
@@ -653,13 +720,17 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
verbose_logger.debug(f"Chat provider: Other content type -> {result}")
return result
def _convert_tools_to_responses_format(self, tools: List[Dict[str, Any]]) -> List["ALL_RESPONSES_API_TOOL_PARAMS"]:
def _convert_tools_to_responses_format(
self, tools: List[Dict[str, Any]]
) -> List["ALL_RESPONSES_API_TOOL_PARAMS"]:
"""Convert chat completion tools to responses API tools format"""
responses_tools: List["ALL_RESPONSES_API_TOOL_PARAMS"] = []
for tool in tools:
# convert function tool from chat completion to responses API format
if tool.get("type") == "function":
function_tool = cast(ChatCompletionToolParamFunctionChunk, tool.get("function"))
function_tool = cast(
ChatCompletionToolParamFunctionChunk, tool.get("function")
)
responses_tools.append(
FunctionToolParam(
name=function_tool["name"],
@@ -685,7 +756,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if not extra_body:
return optional_params
supported_responses_api_params = set(ResponsesAPIOptionalRequestParams.__annotations__.keys())
supported_responses_api_params = set(
ResponsesAPIOptionalRequestParams.__annotations__.keys()
)
# Also include params we handle specially
supported_responses_api_params.update(
{
@@ -703,7 +776,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
return optional_params
def _map_reasoning_effort(self, reasoning_effort: Union[str, Dict[str, Any]]) -> Optional[Reasoning]:
def _map_reasoning_effort(
self, reasoning_effort: Union[str, Dict[str, Any]]
) -> Optional[Reasoning]:
# If dict is passed, convert it directly to Reasoning object
if isinstance(reasoning_effort, dict):
return Reasoning(**reasoning_effort) # type: ignore[typeddict-item]
@@ -711,25 +786,38 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
# Check if auto-summary is enabled via flag or environment variable
# Priority: litellm.reasoning_auto_summary flag > LITELLM_REASONING_AUTO_SUMMARY env var
auto_summary_enabled = (
litellm.reasoning_auto_summary or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true"
litellm.reasoning_auto_summary
or os.getenv("LITELLM_REASONING_AUTO_SUMMARY", "false").lower() == "true"
)
# If string is passed, map with optional summary based on flag/env var
if reasoning_effort == "none":
return Reasoning(effort="none", summary="detailed") if auto_summary_enabled else Reasoning(effort="none") # type: ignore
elif reasoning_effort == "high":
return Reasoning(effort="high", summary="detailed") if auto_summary_enabled else Reasoning(effort="high")
return (
Reasoning(effort="high", summary="detailed")
if auto_summary_enabled
else Reasoning(effort="high")
)
elif reasoning_effort == "xhigh":
return Reasoning(effort="xhigh", summary="detailed") if auto_summary_enabled else Reasoning(effort="xhigh") # type: ignore[typeddict-item]
elif reasoning_effort == "medium":
return (
Reasoning(effort="medium", summary="detailed") if auto_summary_enabled else Reasoning(effort="medium")
Reasoning(effort="medium", summary="detailed")
if auto_summary_enabled
else Reasoning(effort="medium")
)
elif reasoning_effort == "low":
return Reasoning(effort="low", summary="detailed") if auto_summary_enabled else Reasoning(effort="low")
return (
Reasoning(effort="low", summary="detailed")
if auto_summary_enabled
else Reasoning(effort="low")
)
elif reasoning_effort == "minimal":
return (
Reasoning(effort="minimal", summary="detailed") if auto_summary_enabled else Reasoning(effort="minimal")
Reasoning(effort="minimal", summary="detailed")
if auto_summary_enabled
else Reasoning(effort="minimal")
)
return None
@@ -745,7 +833,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
responses_api_request: The responses API request dict to modify
web_search_options: Web search configuration (dict or other value)
"""
if "tools" not in responses_api_request or responses_api_request["tools"] is None:
if (
"tools" not in responses_api_request
or responses_api_request["tools"] is None
):
responses_api_request["tools"] = []
# Get the tools list with proper type narrowing
@@ -835,13 +926,17 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
annotation_dict = annotation
else:
# Skip unsupported annotation types
verbose_logger.debug(f"Skipping unsupported annotation type: {type(annotation)}")
verbose_logger.debug(
f"Skipping unsupported annotation type: {type(annotation)}"
)
continue
result.append(annotation_dict) # type: ignore
except Exception as e:
# Skip malformed annotations
verbose_logger.debug(f"Skipping malformed annotation: {annotation}, error: {e}")
verbose_logger.debug(
f"Skipping malformed annotation: {annotation}, error: {e}"
)
continue
return result if result else None
@@ -862,7 +957,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
def __init__(self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False):
def __init__(
self, streaming_response, sync_stream: bool, json_mode: Optional[bool] = False
):
super().__init__(streaming_response, sync_stream, json_mode)
def _handle_string_chunk(
@@ -875,7 +972,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
if not str_line or str_line.startswith("event:"):
# ignore.
return GenericStreamingChunk(text="", tool_use=None, is_finished=False, finish_reason="", usage=None)
return GenericStreamingChunk(
text="", tool_use=None, is_finished=False, finish_reason="", usage=None
)
index = str_line.find("data:")
if index != -1:
str_line = str_line[index + 5 :]
@@ -938,9 +1037,13 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
if output_item.get("type") == "function_call":
# Extract provider_specific_fields if present
provider_specific_fields = output_item.get("provider_specific_fields")
if provider_specific_fields and not isinstance(provider_specific_fields, dict):
if provider_specific_fields and not isinstance(
provider_specific_fields, dict
):
provider_specific_fields = (
dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {}
dict(provider_specific_fields)
if hasattr(provider_specific_fields, "__dict__")
else {}
)
function_chunk = ChatCompletionToolCallFunctionChunk(
@@ -949,7 +1052,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
)
if provider_specific_fields:
function_chunk["provider_specific_fields"] = provider_specific_fields
function_chunk[
"provider_specific_fields"
] = provider_specific_fields
tool_call_index = parsed_chunk.get("output_index", 0)
tool_call_chunk = ChatCompletionToolCallChunk(
@@ -986,7 +1091,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
id=None,
index=tool_call_index,
type="function",
function=ChatCompletionToolCallFunctionChunk(name=None, arguments=content_part),
function=ChatCompletionToolCallFunctionChunk(
name=None, arguments=content_part
),
)
]
),
@@ -995,16 +1102,22 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
]
)
else:
raise ValueError(f"Chat provider: Invalid function argument delta {parsed_chunk}")
raise ValueError(
f"Chat provider: Invalid function argument delta {parsed_chunk}"
)
elif event_type == "response.output_item.done":
# New output item added
output_item = parsed_chunk.get("item", {})
if output_item.get("type") == "function_call":
# Extract provider_specific_fields if present
provider_specific_fields = output_item.get("provider_specific_fields")
if provider_specific_fields and not isinstance(provider_specific_fields, dict):
if provider_specific_fields and not isinstance(
provider_specific_fields, dict
):
provider_specific_fields = (
dict(provider_specific_fields) if hasattr(provider_specific_fields, "__dict__") else {}
dict(provider_specific_fields)
if hasattr(provider_specific_fields, "__dict__")
else {}
)
function_chunk = ChatCompletionToolCallFunctionChunk(
@@ -1014,7 +1127,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
# Add provider_specific_fields to function if present
if provider_specific_fields:
function_chunk["provider_specific_fields"] = provider_specific_fields
function_chunk[
"provider_specific_fields"
] = provider_specific_fields
tool_call_index = parsed_chunk.get("output_index", 0)
tool_call_chunk = ChatCompletionToolCallChunk(
@@ -1090,11 +1205,22 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
output_items = response_data.get("output", []) if response_data else []
has_function_calls = any(
item.get("type") == "function_call" for item in output_items if isinstance(item, dict)
item.get("type") == "function_call"
for item in output_items
if isinstance(item, dict)
)
finish_reason = "tool_calls" if has_function_calls else "stop"
usage = None
if response_data.get("usage"):
from litellm.responses.utils import ResponseAPILoggingUtils
usage = (
ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage(
response_data.get("usage")
)
)
return ModelResponseStream(
choices=[
StreamingChoices(
@@ -1102,12 +1228,15 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
delta=Delta(content=""),
finish_reason=finish_reason,
)
]
],
usage=usage,
)
else:
pass
# For any unhandled event types, create a minimal valid chunk or skip
verbose_logger.debug(f"Chat provider: Unhandled event type '{event_type}', creating empty chunk")
verbose_logger.debug(
f"Chat provider: Unhandled event type '{event_type}', creating empty chunk"
)
# Return a minimal valid chunk for unknown events
return ModelResponseStream(
@@ -1130,5 +1259,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator):
Returns:
ModelResponseStream: OpenAI-formatted streaming chunk
"""
verbose_logger.debug(f"Chat provider: transform_streaming_response called with chunk: {chunk}")
return OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(chunk)
verbose_logger.debug(
f"Chat provider: transform_streaming_response called with chunk: {chunk}"
)
return OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(
chunk
)
+2 -6
View File
@@ -60,9 +60,7 @@ LITELLM_MAX_STREAMING_DURATION_SECONDS = (
# Maximum number of base64 characters to keep in logging payloads.
# Data URIs exceeding this are replaced with a size placeholder.
# Set to 0 to disable truncation.
MAX_BASE64_LENGTH_FOR_LOGGING = int(
os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64)
)
MAX_BASE64_LENGTH_FOR_LOGGING = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64))
# When true, adds detailed per-phase timing breakdown headers to responses.
# Headers: x-litellm-timing-{pre-processing,llm-api,post-processing,message-copy}-ms
@@ -1421,9 +1419,7 @@ SPECIAL_LITELLM_AUTH_TOKEN = ["ui-token"]
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL = int(
os.getenv("DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL", 60)
)
DEFAULT_ACCESS_GROUP_CACHE_TTL = int(
os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600)
)
DEFAULT_ACCESS_GROUP_CACHE_TTL = int(os.getenv("DEFAULT_ACCESS_GROUP_CACHE_TTL", 600))
# Sentry Scrubbing Configuration
SENTRY_DENYLIST = [
-1
View File
@@ -42,4 +42,3 @@ __all__ = [
"retrieve_container_file",
"retrieve_container_file_content",
]
+20 -14
View File
@@ -43,13 +43,13 @@ def _load_endpoints_config() -> Dict:
def create_sync_endpoint_function(endpoint_config: Dict) -> Callable:
"""
Create a sync SDK function from endpoint config.
Uses the generic container handler instead of individual handler methods.
"""
endpoint_name = endpoint_config["name"]
response_type = RESPONSE_TYPES.get(endpoint_config["response_type"])
path_params = endpoint_config.get("path_params", [])
@client
def endpoint_func(
timeout: int = 600,
@@ -76,14 +76,16 @@ def create_sync_endpoint_function(endpoint_config: Dict) -> Callable:
# Get provider config
litellm_params = GenericLiteLLMParams(**kwargs)
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
container_provider_config: Optional[
BaseContainerConfig
] = ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
if container_provider_config is None:
raise ValueError(f"Container provider config not found for: {custom_llm_provider}")
raise ValueError(
f"Container provider config not found for: {custom_llm_provider}"
)
# Build optional params for logging
optional_params = {k: kwargs.get(k) for k in path_params if k in kwargs}
@@ -126,7 +128,7 @@ def create_async_endpoint_function(
endpoint_config: Dict,
) -> Callable:
"""Create an async SDK function that wraps the sync function."""
@client
async def async_endpoint_func(
timeout: int = 600,
@@ -176,21 +178,21 @@ def create_async_endpoint_function(
def generate_container_endpoints() -> Dict[str, Callable]:
"""
Generate all container endpoint functions from the JSON config.
Returns a dict mapping function names to their implementations.
"""
config = _load_endpoints_config()
endpoints = {}
for endpoint_config in config["endpoints"]:
# Create sync function
sync_func = create_sync_endpoint_function(endpoint_config)
endpoints[endpoint_config["name"]] = sync_func
# Create async function
async_func = create_async_endpoint_function(sync_func, endpoint_config)
endpoints[endpoint_config["async_name"]] = async_func
return endpoints
@@ -222,5 +224,9 @@ retrieve_container_file = _generated_endpoints.get("retrieve_container_file")
aretrieve_container_file = _generated_endpoints.get("aretrieve_container_file")
delete_container_file = _generated_endpoints.get("delete_container_file")
adelete_container_file = _generated_endpoints.get("adelete_container_file")
retrieve_container_file_content = _generated_endpoints.get("retrieve_container_file_content")
aretrieve_container_file_content = _generated_endpoints.get("aretrieve_container_file_content")
retrieve_container_file_content = _generated_endpoints.get(
"retrieve_container_file_content"
)
aretrieve_container_file_content = _generated_endpoints.get(
"aretrieve_container_file_content"
)
+62 -58
View File
@@ -39,6 +39,7 @@ __all__ = [
"upload_container_file",
]
##### Container Create #######################
@client
async def acreate_container(
@@ -164,10 +165,7 @@ def create_container(
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Union[
ContainerObject,
Coroutine[Any, Any, ContainerObject],
]:
) -> Union[ContainerObject, Coroutine[Any, Any, ContainerObject],]:
"""Create a container using the OpenAI Container API.
Currently supports OpenAI
@@ -175,7 +173,7 @@ def create_container(
Example:
```python
import litellm
response = litellm.create_container(
name="My Container",
custom_llm_provider="openai",
@@ -207,19 +205,23 @@ def create_container(
**kwargs,
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
container_provider_config: Optional[
BaseContainerConfig
] = ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
if container_provider_config is None:
raise ValueError(f"container operations are not supported for {custom_llm_provider}")
raise ValueError(
f"container operations are not supported for {custom_llm_provider}"
)
local_vars.update(kwargs)
# Get ContainerCreateOptionalRequestParams with only valid parameters
container_create_optional_params: ContainerCreateOptionalRequestParams = (
ContainerRequestUtils.get_requested_container_create_optional_param(local_vars)
ContainerRequestUtils.get_requested_container_create_optional_param(
local_vars
)
)
# Get optional parameters for the container API
@@ -388,10 +390,7 @@ def list_containers(
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Union[
ContainerListResponse,
Coroutine[Any, Any, ContainerListResponse],
]:
) -> Union[ContainerListResponse, Coroutine[Any, Any, ContainerListResponse],]:
"""List containers using the OpenAI Container API.
Currently supports OpenAI
@@ -420,18 +419,22 @@ def list_containers(
**kwargs,
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
container_provider_config: Optional[
BaseContainerConfig
] = ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
if container_provider_config is None:
raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}")
raise ValueError(
f"Container provider config not found for provider: {custom_llm_provider}"
)
# Get container list request parameters
container_list_optional_params: ContainerListOptionalRequestParams = (
ContainerRequestUtils.get_requested_container_list_optional_param(local_vars)
ContainerRequestUtils.get_requested_container_list_optional_param(
local_vars
)
)
# Pre Call logging
@@ -582,10 +585,7 @@ def retrieve_container(
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Union[
ContainerObject,
Coroutine[Any, Any, ContainerObject],
]:
) -> Union[ContainerObject, Coroutine[Any, Any, ContainerObject],]:
"""Retrieve a container using the OpenAI Container API.
Currently supports OpenAI
@@ -614,14 +614,16 @@ def retrieve_container(
**kwargs,
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
container_provider_config: Optional[
BaseContainerConfig
] = ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
if container_provider_config is None:
raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}")
raise ValueError(
f"Container provider config not found for provider: {custom_llm_provider}"
)
# Pre Call logging
litellm_logging_obj.update_environment_variables(
@@ -768,10 +770,7 @@ def delete_container(
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Union[
DeleteContainerResult,
Coroutine[Any, Any, DeleteContainerResult],
]:
) -> Union[DeleteContainerResult, Coroutine[Any, Any, DeleteContainerResult],]:
"""Delete a container using the OpenAI Container API.
Currently supports OpenAI
@@ -800,14 +799,16 @@ def delete_container(
**kwargs,
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
container_provider_config: Optional[
BaseContainerConfig
] = ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
if container_provider_config is None:
raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}")
raise ValueError(
f"Container provider config not found for provider: {custom_llm_provider}"
)
# Pre Call logging
litellm_logging_obj.update_environment_variables(
@@ -968,10 +969,7 @@ def list_container_files(
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Union[
ContainerFileListResponse,
Coroutine[Any, Any, ContainerFileListResponse],
]:
) -> Union[ContainerFileListResponse, Coroutine[Any, Any, ContainerFileListResponse],]:
"""List files in a container using the OpenAI Container API.
Currently supports OpenAI
@@ -1000,19 +998,26 @@ def list_container_files(
**kwargs,
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
container_provider_config: Optional[
BaseContainerConfig
] = ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
if container_provider_config is None:
raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}")
raise ValueError(
f"Container provider config not found for provider: {custom_llm_provider}"
)
# Pre Call logging
litellm_logging_obj.update_environment_variables(
model="",
optional_params={"container_id": container_id, "after": after, "limit": limit, "order": order},
optional_params={
"container_id": container_id,
"after": after,
"limit": limit,
"order": order,
},
litellm_params={
"litellm_call_id": litellm_call_id,
},
@@ -1180,10 +1185,7 @@ def upload_container_file(
extra_query: Optional[Dict[str, Any]] = None,
extra_body: Optional[Dict[str, Any]] = None,
**kwargs,
) -> Union[
ContainerFileObject,
Coroutine[Any, Any, ContainerFileObject],
]:
) -> Union[ContainerFileObject, Coroutine[Any, Any, ContainerFileObject],]:
"""Upload a file to a container using the OpenAI Container API.
This endpoint allows uploading files directly to a container session,
@@ -1241,14 +1243,16 @@ def upload_container_file(
**kwargs,
)
# get provider config
container_provider_config: Optional[BaseContainerConfig] = (
ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
container_provider_config: Optional[
BaseContainerConfig
] = ProviderConfigManager.get_provider_container_config(
provider=litellm.LlmProviders(custom_llm_provider),
)
if container_provider_config is None:
raise ValueError(f"Container provider config not found for provider: {custom_llm_provider}")
raise ValueError(
f"Container provider config not found for provider: {custom_llm_provider}"
)
# Pre Call logging
litellm_logging_obj.update_environment_variables(
+4 -1
View File
@@ -1,7 +1,10 @@
from typing import Dict
from litellm.llms.base_llm.containers.transformation import BaseContainerConfig
from litellm.types.containers.main import ContainerCreateOptionalRequestParams, ContainerListOptionalRequestParams
from litellm.types.containers.main import (
ContainerCreateOptionalRequestParams,
ContainerListOptionalRequestParams,
)
class ContainerRequestUtils:
+51 -36
View File
@@ -120,37 +120,49 @@ else:
LitellmLoggingObject = Any
# Pre-resolved CallTypes enum values for fast membership checks
_A2A_CALL_TYPES = frozenset({
CallTypes.asend_message.value,
CallTypes.send_message.value,
})
_A2A_CALL_TYPES = frozenset(
{
CallTypes.asend_message.value,
CallTypes.send_message.value,
}
)
_VIDEO_CALL_TYPES = frozenset({
CallTypes.create_video.value,
CallTypes.acreate_video.value,
CallTypes.video_remix.value,
CallTypes.avideo_remix.value,
})
_VIDEO_CALL_TYPES = frozenset(
{
CallTypes.create_video.value,
CallTypes.acreate_video.value,
CallTypes.video_remix.value,
CallTypes.avideo_remix.value,
}
)
_SPEECH_CALL_TYPES = frozenset({
CallTypes.speech.value,
CallTypes.aspeech.value,
})
_SPEECH_CALL_TYPES = frozenset(
{
CallTypes.speech.value,
CallTypes.aspeech.value,
}
)
_TRANSCRIPTION_CALL_TYPES = frozenset({
CallTypes.atranscription.value,
CallTypes.transcription.value,
})
_TRANSCRIPTION_CALL_TYPES = frozenset(
{
CallTypes.atranscription.value,
CallTypes.transcription.value,
}
)
_RERANK_CALL_TYPES = frozenset({
CallTypes.rerank.value,
CallTypes.arerank.value,
})
_RERANK_CALL_TYPES = frozenset(
{
CallTypes.rerank.value,
CallTypes.arerank.value,
}
)
_SEARCH_CALL_TYPES = frozenset({
CallTypes.search.value,
CallTypes.asearch.value,
})
_SEARCH_CALL_TYPES = frozenset(
{
CallTypes.search.value,
CallTypes.asearch.value,
}
)
_AREALTIME_CALL_TYPE = CallTypes.arealtime.value
_MCP_CALL_TYPE = CallTypes.call_mcp_tool.value
@@ -522,7 +534,10 @@ def cost_per_token( # noqa: PLR0915
return dashscope_cost_per_token(model=model, usage=usage_block)
elif custom_llm_provider == "azure_ai":
return azure_ai_cost_per_token(
model=model, usage=usage_block, response_time_ms=response_time_ms, request_model=request_model
model=model,
usage=usage_block,
response_time_ms=response_time_ms,
request_model=request_model,
)
else:
model_info = _cached_get_model_info_helper(
@@ -1114,9 +1129,9 @@ def completion_cost( # noqa: PLR0915
or isinstance(completion_response, dict)
): # tts returns a custom class
if isinstance(completion_response, dict):
usage_obj: Optional[Union[dict, Usage]] = (
completion_response.get("usage", {})
)
usage_obj: Optional[
Union[dict, Usage]
] = completion_response.get("usage", {})
else:
usage_obj = getattr(completion_response, "usage", {})
if isinstance(usage_obj, BaseModel) and not _is_known_usage_objects(
@@ -1501,7 +1516,6 @@ def completion_cost( # noqa: PLR0915
else:
additional_costs = None
_final_cost = (
prompt_tokens_cost_usd_dollar + completion_tokens_cost_usd_dollar
)
@@ -1519,7 +1533,11 @@ def completion_cost( # noqa: PLR0915
# Apply discount from module-level config if configured
original_cost = _final_cost
if litellm.cost_discount_config:
_final_cost, discount_percent, discount_amount = _apply_cost_discount(
(
_final_cost,
discount_percent,
discount_amount,
) = _apply_cost_discount(
base_cost=_final_cost,
custom_llm_provider=custom_llm_provider,
)
@@ -1976,9 +1994,7 @@ def default_video_cost_calculator(
cost_info = litellm.model_cost[prefixed_model]
if cost_info is None:
raise Exception(
f"Model not found in cost map for model={model}"
)
raise Exception(f"Model not found in cost map for model={model}")
# Check for video-specific cost per second first
video_cost_per_second = cost_info.get("output_cost_per_video_per_second")
@@ -2250,4 +2266,3 @@ def handle_realtime_stream_cost_calculation(
total_cost = input_cost_per_token + output_cost_per_token
return total_cost
+100 -69
View File
@@ -152,16 +152,14 @@ def create_eval(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = (
ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
evals_api_provider_config: Optional[
BaseEvalsAPIConfig
] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
if evals_api_provider_config is None:
raise ValueError(
f"CREATE eval is not supported for {custom_llm_provider}"
)
raise ValueError(f"CREATE eval is not supported for {custom_llm_provider}")
# Build create request
create_request: CreateEvalRequest = {
@@ -344,10 +342,10 @@ def list_evals(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = (
ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
evals_api_provider_config: Optional[
BaseEvalsAPIConfig
] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
if evals_api_provider_config is None:
@@ -513,10 +511,10 @@ def get_eval(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = (
ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
evals_api_provider_config: Optional[
BaseEvalsAPIConfig
] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
if evals_api_provider_config is None:
@@ -681,16 +679,14 @@ def update_eval(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = (
ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
evals_api_provider_config: Optional[
BaseEvalsAPIConfig
] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
if evals_api_provider_config is None:
raise ValueError(
f"UPDATE eval is not supported for {custom_llm_provider}"
)
raise ValueError(f"UPDATE eval is not supported for {custom_llm_provider}")
# Build update request
update_request: UpdateEvalRequest = {}
@@ -701,20 +697,41 @@ def update_eval(
if metadata is not None:
# List of internal LiteLLM metadata keys that should NOT be sent to OpenAI
internal_keys = {
"headers", "requester_metadata", "user_api_key_hash", "user_api_key_alias",
"user_api_key_spend", "user_api_key_max_budget", "user_api_key_team_id",
"user_api_key_user_id", "user_api_key_org_id", "user_api_key_team_alias",
"user_api_key_end_user_id", "user_api_key_user_email", "user_api_key_request_route",
"user_api_key_budget_reset_at", "user_api_key_auth_metadata", "user_api_key",
"user_api_end_user_max_budget", "user_api_key_auth", "litellm_api_version",
"global_max_parallel_requests", "user_api_key_team_max_budget",
"user_api_key_team_spend", "user_api_key_model_max_budget",
"user_api_key_user_spend", "user_api_key_user_max_budget",
"user_api_key_metadata", "endpoint", "litellm_parent_otel_span",
"requester_ip_address", "user_agent",
"headers",
"requester_metadata",
"user_api_key_hash",
"user_api_key_alias",
"user_api_key_spend",
"user_api_key_max_budget",
"user_api_key_team_id",
"user_api_key_user_id",
"user_api_key_org_id",
"user_api_key_team_alias",
"user_api_key_end_user_id",
"user_api_key_user_email",
"user_api_key_request_route",
"user_api_key_budget_reset_at",
"user_api_key_auth_metadata",
"user_api_key",
"user_api_end_user_max_budget",
"user_api_key_auth",
"litellm_api_version",
"global_max_parallel_requests",
"user_api_key_team_max_budget",
"user_api_key_team_spend",
"user_api_key_model_max_budget",
"user_api_key_user_spend",
"user_api_key_user_max_budget",
"user_api_key_metadata",
"endpoint",
"litellm_parent_otel_span",
"requester_ip_address",
"user_agent",
}
# Only include user-provided metadata keys
filtered_metadata = {k: v for k, v in metadata.items() if k not in internal_keys}
filtered_metadata = {
k: v for k, v in metadata.items() if k not in internal_keys
}
if filtered_metadata: # Only add if there's user metadata
update_request["metadata"] = filtered_metadata
@@ -730,7 +747,11 @@ def update_eval(
# Transform request
api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE
url, headers, request_body = evals_api_provider_config.transform_update_eval_request(
(
url,
headers,
request_body,
) = evals_api_provider_config.transform_update_eval_request(
eval_id=eval_id,
update_request=update_request,
api_base=api_base,
@@ -868,10 +889,10 @@ def delete_eval(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = (
ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
evals_api_provider_config: Optional[
BaseEvalsAPIConfig
] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
if evals_api_provider_config is None:
@@ -1021,10 +1042,10 @@ def cancel_eval(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = (
ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
evals_api_provider_config: Optional[
BaseEvalsAPIConfig
] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
if evals_api_provider_config is None:
@@ -1038,7 +1059,11 @@ def cancel_eval(
# Transform request
api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE
url, headers, request_body = evals_api_provider_config.transform_cancel_eval_request(
(
url,
headers,
request_body,
) = evals_api_provider_config.transform_cancel_eval_request(
eval_id=eval_id,
api_base=api_base,
litellm_params=litellm_params,
@@ -1199,16 +1224,14 @@ def create_run(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = (
ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
evals_api_provider_config: Optional[
BaseEvalsAPIConfig
] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
if evals_api_provider_config is None:
raise ValueError(
f"CREATE run is not supported for {custom_llm_provider}"
)
raise ValueError(f"CREATE run is not supported for {custom_llm_provider}")
# Build create request
create_request: CreateRunRequest = {
@@ -1388,10 +1411,10 @@ def list_runs(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = (
ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
evals_api_provider_config: Optional[
BaseEvalsAPIConfig
] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
if evals_api_provider_config is None:
@@ -1561,10 +1584,10 @@ def get_run(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = (
ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
evals_api_provider_config: Optional[
BaseEvalsAPIConfig
] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
if evals_api_provider_config is None:
@@ -1720,10 +1743,10 @@ def cancel_run(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = (
ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
evals_api_provider_config: Optional[
BaseEvalsAPIConfig
] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
if evals_api_provider_config is None:
@@ -1737,7 +1760,11 @@ def cancel_run(
# Transform request
api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE
url, headers, request_body = evals_api_provider_config.transform_cancel_run_request(
(
url,
headers,
request_body,
) = evals_api_provider_config.transform_cancel_run_request(
eval_id=eval_id,
run_id=run_id,
api_base=api_base,
@@ -1884,10 +1911,10 @@ def delete_run(
custom_llm_provider = "openai"
# Get provider config
evals_api_provider_config: Optional[BaseEvalsAPIConfig] = (
ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
evals_api_provider_config: Optional[
BaseEvalsAPIConfig
] = ProviderConfigManager.get_provider_evals_api_config( # type: ignore
provider=litellm.LlmProviders(custom_llm_provider),
)
if evals_api_provider_config is None:
@@ -1901,7 +1928,11 @@ def delete_run(
# Transform request
api_base = litellm_params.api_base or DEFAULT_OPENAI_API_BASE
url, headers, request_body = evals_api_provider_config.transform_delete_run_request(
(
url,
headers,
request_body,
) = evals_api_provider_config.transform_delete_run_request(
eval_id=eval_id,
run_id=run_id,
api_base=api_base,
+2 -4
View File
@@ -25,9 +25,7 @@ def _get_minimal_error_response() -> httpx.Response:
if _MINIMAL_ERROR_RESPONSE is None:
_MINIMAL_ERROR_RESPONSE = httpx.Response(
status_code=400,
request=httpx.Request(
method="GET", url="https://litellm.ai"
),
request=httpx.Request(method="GET", url="https://litellm.ai"),
)
return _MINIMAL_ERROR_RESPONSE
@@ -996,7 +994,7 @@ class MidStreamFallbackError(ServiceUnavailableError): # type: ignore
max_retries=self.max_retries,
num_retries=self.num_retries,
)
# Restore the propagated status and original response/request objects
self.status_code = int(original_status) if original_status is not None else 503
self.response = _saved_response
+33 -16
View File
@@ -4,7 +4,18 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers.
import asyncio
import base64
from typing import Any, Awaitable, Callable, Dict, Generator, List, Optional, Tuple, TypeVar, Union
from typing import (
Any,
Awaitable,
Callable,
Dict,
Generator,
List,
Optional,
Tuple,
TypeVar,
Union,
)
import httpx
from mcp import ClientSession, ReadResourceResult, Resource, StdioServerParameters
@@ -14,7 +25,10 @@ from mcp.client.stdio import stdio_client
streamable_http_client: Optional[Any] = None
try:
import mcp.client.streamable_http as streamable_http_module # type: ignore
streamable_http_client = getattr(streamable_http_module, "streamable_http_client", None)
streamable_http_client = getattr(
streamable_http_module, "streamable_http_client", None
)
except ImportError:
pass
from mcp.types import CallToolRequestParams as MCPCallToolRequestParams
@@ -188,12 +202,15 @@ class MCPClient:
if self.transport_type == MCPTransport.sse:
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
return sse_client(
url=self.server_url,
timeout=self.timeout,
headers=headers,
httpx_client_factory=httpx_client_factory,
), None
return (
sse_client(
url=self.server_url,
timeout=self.timeout,
headers=headers,
httpx_client_factory=httpx_client_factory,
),
None,
)
# HTTP transport (default)
if streamable_http_client is None:
@@ -201,12 +218,10 @@ class MCPClient:
"streamable_http_client is not available. "
"Please install mcp with HTTP support."
)
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
verbose_logger.debug(
"litellm headers for streamable_http_client: %s", headers
)
verbose_logger.debug("litellm headers for streamable_http_client: %s", headers)
http_client = httpx_client_factory(
headers=headers,
timeout=httpx.Timeout(self.timeout),
@@ -392,7 +407,7 @@ class MCPClient:
async def call_tool(
self,
call_tool_request_params: MCPCallToolRequestParams,
host_progress_callback: Optional[Callable] = None
host_progress_callback: Optional[Callable] = None,
) -> MCPCallToolResult:
"""
Call an MCP Tool.
@@ -401,13 +416,15 @@ class MCPClient:
f"MCP client calling tool '{call_tool_request_params.name}' with arguments: {call_tool_request_params.arguments}"
)
async def on_progress(progress: float, total: float | None, message: str | None):
async def on_progress(
progress: float, total: float | None, message: str | None
):
percentage = (progress / total * 100) if total else 0
verbose_logger.info(
f"MCP Tool '{call_tool_request_params.name}' progress: "
f"{progress}/{total} ({percentage:.0f}%) - {message or ''}"
)
# Forward to Host if callback provided
if host_progress_callback:
try:
@@ -421,8 +438,8 @@ class MCPClient:
name=call_tool_request_params.name,
arguments=call_tool_request_params.arguments,
progress_callback=on_progress,
)
try:
tool_result = await self.run_with_session(_call_tool_operation)
verbose_logger.info(
+13 -14
View File
@@ -18,7 +18,7 @@ from litellm.types.utils import ChatCompletionMessageToolCall
def transform_mcp_tool_to_openai_tool(mcp_tool: MCPTool) -> ChatCompletionToolParam:
"""Convert an MCP tool to an OpenAI tool."""
normalized_parameters = _normalize_mcp_input_schema(mcp_tool.inputSchema)
return ChatCompletionToolParam(
type="function",
function=FunctionDefinition(
@@ -33,41 +33,39 @@ def transform_mcp_tool_to_openai_tool(mcp_tool: MCPTool) -> ChatCompletionToolPa
def _normalize_mcp_input_schema(input_schema: dict) -> dict:
"""
Normalize MCP input schema to ensure it's valid for OpenAI function calling.
OpenAI requires that function parameters have:
- type: 'object'
- properties: dict (can be empty)
- additionalProperties: false (recommended)
"""
if not input_schema:
return {
"type": "object",
"properties": {},
"additionalProperties": False
}
return {"type": "object", "properties": {}, "additionalProperties": False}
# Make a copy to avoid modifying the original
normalized_schema = dict(input_schema)
# Ensure type is 'object'
if "type" not in normalized_schema:
normalized_schema["type"] = "object"
# Ensure properties exists (can be empty)
if "properties" not in normalized_schema:
normalized_schema["properties"] = {}
# Add additionalProperties if not present (recommended by OpenAI)
if "additionalProperties" not in normalized_schema:
normalized_schema["additionalProperties"] = False
return normalized_schema
def transform_mcp_tool_to_openai_responses_api_tool(mcp_tool: MCPTool) -> FunctionToolParam:
def transform_mcp_tool_to_openai_responses_api_tool(
mcp_tool: MCPTool,
) -> FunctionToolParam:
"""Convert an MCP tool to an OpenAI Responses API tool."""
normalized_parameters = _normalize_mcp_input_schema(mcp_tool.inputSchema)
return FunctionToolParam(
name=mcp_tool.name,
parameters=normalized_parameters,
@@ -76,6 +74,7 @@ def transform_mcp_tool_to_openai_responses_api_tool(mcp_tool: MCPTool) -> Functi
description=mcp_tool.description or "",
)
async def load_mcp_tools(
session: ClientSession, format: Literal["mcp", "openai"] = "mcp"
) -> Union[List[MCPTool], List[ChatCompletionToolParam]]:
+91 -38
View File
@@ -14,11 +14,30 @@ from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast
import httpx
# Type aliases for provider parameters
FileCreateProvider = Literal[
"openai",
"azure",
"gemini",
"vertex_ai",
"bedrock",
"hosted_vllm",
"manus",
"anthropic",
]
FileRetrieveProvider = Literal[
"openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus", "anthropic"
]
FileDeleteProvider = Literal["openai", "azure", "gemini", "manus", "anthropic"]
FileListProvider = Literal["openai", "azure", "manus", "anthropic"]
FileContentProvider = Literal[
"openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"
]
import litellm
from litellm import get_secret_str
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.anthropic.files.handler import AnthropicFilesHandler
from litellm.llms.azure.common_utils import get_azure_credentials
from litellm.llms.azure.files.handler import AzureOpenAIFilesAPI
from litellm.llms.bedrock.files.handler import BedrockFilesHandler
@@ -54,16 +73,15 @@ openai_files_instance = OpenAIFilesAPI()
azure_files_instance = AzureOpenAIFilesAPI()
vertex_ai_files_instance = VertexAIFilesHandler()
bedrock_files_instance = BedrockFilesHandler()
anthropic_files_instance = AnthropicFilesHandler()
#################################################
@client
async def acreate_file(
file: FileTypes,
purpose: Literal["assistants", "batch", "fine-tune"],
purpose: Literal["assistants", "batch", "fine-tune", "messages"],
expires_after: Optional[FileExpiresAfter] = None,
custom_llm_provider: Literal["openai", "azure", "gemini", "vertex_ai", "bedrock", "hosted_vllm", "manus"] = "openai",
custom_llm_provider: FileCreateProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -106,9 +124,9 @@ async def acreate_file(
@client
def create_file(
file: FileTypes,
purpose: Literal["assistants", "batch", "fine-tune"],
purpose: Literal["assistants", "batch", "fine-tune", "messages"],
expires_after: Optional[FileExpiresAfter] = None,
custom_llm_provider: Optional[Literal["openai", "azure", "gemini", "vertex_ai", "bedrock", "hosted_vllm", "manus"]] = None,
custom_llm_provider: Optional[FileCreateProvider] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -218,7 +236,7 @@ def create_file(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'create_file'. Only ['openai', 'azure', 'vertex_ai', 'manus'] are supported.".format(
message="LiteLLM doesn't support {} for 'create_file'. Only ['openai', 'azure', 'vertex_ai', 'manus', 'anthropic'] are supported.".format(
custom_llm_provider
),
model="n/a",
@@ -237,7 +255,7 @@ def create_file(
@client
async def afile_retrieve(
file_id: str,
custom_llm_provider: Literal["openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus"] = "openai",
custom_llm_provider: FileRetrieveProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -278,7 +296,7 @@ async def afile_retrieve(
@client
def file_retrieve(
file_id: str,
custom_llm_provider: Literal["openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus"] = "openai",
custom_llm_provider: FileRetrieveProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -348,22 +366,25 @@ def file_retrieve(
litellm_params_dict = get_litellm_params(**kwargs)
litellm_params_dict["api_key"] = optional_params.api_key
litellm_params_dict["api_base"] = optional_params.api_base
logging_obj = kwargs.get("litellm_logging_obj")
if logging_obj is None:
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
logging_obj = LiteLLMLoggingObj(
model="",
messages=[],
stream=False,
call_type="afile_retrieve" if _is_async else "file_retrieve",
start_time=time.time(),
litellm_call_id=kwargs.get("litellm_call_id", str(uuid_module.uuid4())),
litellm_call_id=kwargs.get(
"litellm_call_id", str(uuid_module.uuid4())
),
function_id=str(kwargs.get("id") or ""),
)
client = kwargs.get("client")
response = base_llm_http_handler.retrieve_file(
file_id=file_id,
@@ -382,7 +403,7 @@ def file_retrieve(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'file_retrieve'. Only 'openai', 'azure', and 'manus' are supported.".format(
message="LiteLLM doesn't support {} for 'file_retrieve'. Only 'openai', 'azure', 'manus', and 'anthropic' are supported.".format(
custom_llm_provider
),
model="n/a",
@@ -403,7 +424,7 @@ def file_retrieve(
@client
async def afile_delete(
file_id: str,
custom_llm_provider: Literal["openai", "azure", "gemini", "manus"] = "openai",
custom_llm_provider: FileDeleteProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -447,7 +468,7 @@ async def afile_delete(
def file_delete(
file_id: str,
model: Optional[str] = None,
custom_llm_provider: Union[Literal["openai", "azure", "gemini", "manus"], str] = "openai",
custom_llm_provider: Union[FileDeleteProvider, str] = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -525,22 +546,25 @@ def file_delete(
if provider_config is not None:
litellm_params_dict["api_key"] = optional_params.api_key
litellm_params_dict["api_base"] = optional_params.api_base
logging_obj = kwargs.get("litellm_logging_obj")
if logging_obj is None:
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
logging_obj = LiteLLMLoggingObj(
model="",
messages=[],
stream=False,
call_type="afile_delete" if _is_async else "file_delete",
start_time=time.time(),
litellm_call_id=kwargs.get("litellm_call_id", str(uuid_module.uuid4())),
litellm_call_id=kwargs.get(
"litellm_call_id", str(uuid_module.uuid4())
),
function_id=str(kwargs.get("id") or ""),
)
response = base_llm_http_handler.delete_file(
file_id=file_id,
provider_config=provider_config,
@@ -558,7 +582,7 @@ def file_delete(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'file_delete'. Only 'openai', 'azure', 'gemini', and 'manus' are supported.".format(
message="LiteLLM doesn't support {} for 'file_delete'. Only 'openai', 'azure', 'gemini', 'manus', and 'anthropic' are supported.".format(
custom_llm_provider
),
model="n/a",
@@ -577,7 +601,7 @@ def file_delete(
# List files
@client
async def afile_list(
custom_llm_provider: Literal["openai", "azure", "manus"] = "openai",
custom_llm_provider: FileListProvider = "openai",
purpose: Optional[str] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@@ -618,7 +642,7 @@ async def afile_list(
@client
def file_list(
custom_llm_provider: Literal["openai", "azure", "manus"] = "openai",
custom_llm_provider: FileListProvider = "openai",
purpose: Optional[str] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
@@ -648,7 +672,7 @@ def file_list(
timeout = 600.0
_is_async = kwargs.pop("is_async", False) is True
# Check if provider has a custom files config (e.g., Manus, Bedrock, Vertex AI)
provider_config = ProviderConfigManager.get_provider_files_config(
model="",
@@ -658,22 +682,25 @@ def file_list(
litellm_params_dict = get_litellm_params(**kwargs)
litellm_params_dict["api_key"] = optional_params.api_key
litellm_params_dict["api_base"] = optional_params.api_base
logging_obj = kwargs.get("litellm_logging_obj")
if logging_obj is None:
from litellm.litellm_core_utils.litellm_logging import (
Logging as LiteLLMLoggingObj,
)
logging_obj = LiteLLMLoggingObj(
model="",
messages=[],
stream=False,
call_type="afile_list" if _is_async else "file_list",
start_time=time.time(),
litellm_call_id=kwargs.get("litellm_call_id", str(uuid_module.uuid4())),
litellm_call_id=kwargs.get(
"litellm_call_id", str(uuid_module.uuid4())
),
function_id=str(kwargs.get("id", "")),
)
client = kwargs.get("client")
response = base_llm_http_handler.list_files(
purpose=purpose,
@@ -723,7 +750,7 @@ def file_list(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'file_list'. Only 'openai', 'azure', and 'manus' are supported.".format(
message="LiteLLM doesn't support {} for 'file_list'. Only 'openai', 'azure', 'manus', and 'anthropic' are supported.".format(
custom_llm_provider
),
model="n/a",
@@ -742,7 +769,7 @@ def file_list(
@client
async def afile_content(
file_id: str,
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"] = "openai",
custom_llm_provider: FileContentProvider = "openai",
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -786,9 +813,7 @@ async def afile_content(
def file_content(
file_id: str,
model: Optional[str] = None,
custom_llm_provider: Optional[
Union[Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm", "anthropic", "manus"], str]
] = None,
custom_llm_provider: Optional[Union[FileContentProvider, str]] = None,
extra_headers: Optional[Dict[str, str]] = None,
extra_body: Optional[Dict[str, str]] = None,
**kwargs,
@@ -834,15 +859,43 @@ def file_content(
_is_async = kwargs.pop("afile_content", False) is True
# Check if this is an Anthropic batch results request
if custom_llm_provider == "anthropic":
response = anthropic_files_instance.file_content(
_is_async=_is_async,
# Check if provider has a custom files config (e.g., Anthropic, Manus)
provider_config = ProviderConfigManager.get_provider_files_config(
model="",
provider=LlmProviders(custom_llm_provider),
)
if provider_config is not None:
litellm_params_dict["api_key"] = optional_params.api_key
litellm_params_dict["api_base"] = optional_params.api_base
logging_obj = kwargs.get("litellm_logging_obj")
if logging_obj is None:
logging_obj = LiteLLMLoggingObj(
model="",
messages=[],
stream=False,
call_type="afile_content" if _is_async else "file_content",
start_time=time.time(),
litellm_call_id=kwargs.get(
"litellm_call_id", str(uuid_module.uuid4())
),
function_id=str(kwargs.get("id") or ""),
)
response = base_llm_http_handler.retrieve_file_content(
file_content_request=_file_content_request,
api_base=optional_params.api_base,
api_key=optional_params.api_key,
provider_config=provider_config,
litellm_params=litellm_params_dict,
headers=extra_headers or {},
logging_obj=logging_obj,
_is_async=_is_async,
client=(
client
if client is not None
and isinstance(client, (HTTPHandler, AsyncHTTPHandler))
else None
),
timeout=timeout,
max_retries=optional_params.max_retries,
)
return response
@@ -915,7 +968,7 @@ def file_content(
)
else:
raise litellm.exceptions.BadRequestError(
message="LiteLLM doesn't support {} for 'file_content'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock', 'manus'.".format(
message="LiteLLM doesn't support {} for 'file_content'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock', 'manus', 'anthropic'.".format(
custom_llm_provider
),
model="n/a",
+8 -3
View File
@@ -8,17 +8,22 @@ class FilesAPIUtils:
"""
Utils for files API interface on litellm
"""
@staticmethod
def is_batch_jsonl_file(create_file_data: CreateFileRequest, extracted_file_data: ExtractedFileData) -> bool:
def is_batch_jsonl_file(
create_file_data: CreateFileRequest, extracted_file_data: ExtractedFileData
) -> bool:
"""
Check if the file is a batch jsonl file
"""
return (
create_file_data.get("purpose") == "batch"
and FilesAPIUtils.valid_content_type(extracted_file_data.get("content_type"))
and FilesAPIUtils.valid_content_type(
extracted_file_data.get("content_type")
)
and extracted_file_data.get("content") is not None
)
@staticmethod
def valid_content_type(content_type: Optional[str]) -> bool:
"""
+38 -16
View File
@@ -41,34 +41,34 @@ def _prepare_azure_extra_body(
) -> Dict[str, Any]:
"""
Prepare extra_body for Azure fine-tuning API by combining Azure-specific parameters.
Azure fine-tuning API accepts additional parameters beyond the standard OpenAI spec:
- trainingType: Type of training (e.g., 1 for supervised fine-tuning)
- prompt_loss_weight: Weight for prompt loss in training
These parameters must be passed in the extra_body field when calling the Azure OpenAI SDK.
Args:
extra_body: Optional existing extra_body dict
kwargs: Request kwargs that may contain Azure-specific parameters
azure_specific_hyperparams: Dict of Azure-specific hyperparameters already extracted
Returns:
Dict containing all Azure-specific parameters to be passed in extra_body
"""
if extra_body is None:
extra_body = {}
# Azure-specific root-level parameters
azure_specific_params = ["trainingType"]
for param in azure_specific_params:
if param in kwargs:
extra_body[param] = kwargs[param]
# Add Azure-specific hyperparameters
if azure_specific_hyperparams:
extra_body.update(azure_specific_hyperparams)
return extra_body
@@ -126,7 +126,9 @@ async def acreate_fine_tuning_job(
raise e
def _build_fine_tuning_job_data(model, training_file, hyperparameters, suffix, validation_file, integrations, seed):
def _build_fine_tuning_job_data(
model, training_file, hyperparameters, suffix, validation_file, integrations, seed
):
return FineTuningJobCreate(
model=model,
training_file=training_file,
@@ -177,7 +179,7 @@ def create_fine_tuning_job(
# handle hyperparameters
hyperparameters = hyperparameters or {} # original hyperparameters
# For Azure, extract Azure-specific hyperparameters before creating OpenAI-spec hyperparameters
azure_specific_hyperparams = {}
if custom_llm_provider == "azure":
@@ -185,7 +187,7 @@ def create_fine_tuning_job(
for key in azure_hyperparameter_keys:
if key in hyperparameters:
azure_specific_hyperparams[key] = hyperparameters.pop(key)
_oai_hyperparameters: Hyperparameters = Hyperparameters(
**hyperparameters
) # Typed Hyperparameters for OpenAI Spec
@@ -219,7 +221,13 @@ def create_fine_tuning_job(
)
create_fine_tuning_job_data_dict = _build_fine_tuning_job_data(
model, training_file, _oai_hyperparameters, suffix, validation_file, integrations, seed,
model,
training_file,
_oai_hyperparameters,
suffix,
validation_file,
integrations,
seed,
).model_dump(exclude_none=True)
response = openai_fine_tuning_apis_instance.create_fine_tuning_job(
@@ -258,12 +266,20 @@ def create_fine_tuning_job(
extra_body.pop("azure_ad_token", None)
else:
get_secret_str("AZURE_AD_TOKEN") # type: ignore
# Prepare Azure-specific parameters for extra_body
extra_body = _prepare_azure_extra_body(extra_body, kwargs, azure_specific_hyperparams)
extra_body = _prepare_azure_extra_body(
extra_body, kwargs, azure_specific_hyperparams
)
create_fine_tuning_job_data_dict = _build_fine_tuning_job_data(
model, training_file, _oai_hyperparameters, suffix, validation_file, integrations, seed,
model,
training_file,
_oai_hyperparameters,
suffix,
validation_file,
integrations,
seed,
).model_dump(exclude_none=True)
# Add extra_body if it has Azure-specific parameters
@@ -298,7 +314,13 @@ def create_fine_tuning_job(
response = vertex_fine_tuning_apis_instance.create_fine_tuning_job(
_is_async=_is_async,
create_fine_tuning_job_data=_build_fine_tuning_job_data(
model, training_file, _oai_hyperparameters, suffix, validation_file, integrations, seed,
model,
training_file,
_oai_hyperparameters,
suffix,
validation_file,
integrations,
seed,
),
vertex_credentials=vertex_credentials,
vertex_project=vertex_ai_project,
+2 -2
View File
@@ -13,7 +13,7 @@ from .main import (
__all__ = [
"generate_content",
"agenerate_content",
"agenerate_content",
"generate_content_stream",
"agenerate_content_stream",
]
]
+3 -3
View File
@@ -13,7 +13,7 @@ from .handler import GenerateContentToCompletionHandler
from .transformation import GoogleGenAIAdapter, GoogleGenAIStreamWrapper
__all__ = [
"GoogleGenAIAdapter",
"GoogleGenAIAdapter",
"GoogleGenAIStreamWrapper",
"GenerateContentToCompletionHandler"
]
"GenerateContentToCompletionHandler",
]
+9 -3
View File
@@ -168,7 +168,9 @@ class GenerateContentHelper:
)
)
# Extract systemInstruction from kwargs to pass to transform
system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction")
system_instruction = kwargs.get("systemInstruction") or kwargs.get(
"system_instruction"
)
request_body = (
generate_content_provider_config.transform_generate_content_request(
model=model,
@@ -318,7 +320,9 @@ def generate_content(
)
# Extract systemInstruction from kwargs to pass to handler
system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction")
system_instruction = kwargs.get("systemInstruction") or kwargs.get(
"system_instruction"
)
# Check if we should use the adapter (when provider config is None)
if setup_result.generate_content_provider_config is None:
@@ -407,7 +411,9 @@ async def agenerate_content_stream(
)
# Extract systemInstruction from kwargs to pass to handler
system_instruction = kwargs.get("systemInstruction") or kwargs.get("system_instruction")
system_instruction = kwargs.get("systemInstruction") or kwargs.get(
"system_instruction"
)
# Check if we should use the adapter (when provider config is None)
if setup_result.generate_content_provider_config is None:
+12 -4
View File
@@ -17,6 +17,7 @@ else:
GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ = PassThroughEndpointLogging()
class BaseGoogleGenAIGenerateContentStreamingIterator:
"""
Base class for Google GenAI Generate Content streaming iterators that provides common logic
@@ -42,6 +43,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
from litellm.proxy.pass_through_endpoints.streaming_handler import (
PassThroughStreamingHandler,
)
end_time = datetime.now()
asyncio.create_task(
PassThroughStreamingHandler._route_streaming_logging_to_handler(
@@ -58,7 +60,9 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
)
class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContentStreamingIterator):
class GoogleGenAIGenerateContentStreamingIterator(
BaseGoogleGenAIGenerateContentStreamingIterator
):
"""
Streaming iterator specifically for Google GenAI generate content API.
"""
@@ -105,10 +109,14 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent
async def __anext__(self):
# This should not be used for sync responses
# If you need async iteration, use AsyncGoogleGenAIGenerateContentStreamingIterator
raise NotImplementedError("Use AsyncGoogleGenAIGenerateContentStreamingIterator for async iteration")
raise NotImplementedError(
"Use AsyncGoogleGenAIGenerateContentStreamingIterator for async iteration"
)
class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContentStreamingIterator):
class AsyncGoogleGenAIGenerateContentStreamingIterator(
BaseGoogleGenAIGenerateContentStreamingIterator
):
"""
Async streaming iterator specifically for Google GenAI generate content API.
"""
@@ -148,4 +156,4 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo
return chunk
except StopAsyncIteration:
await self._handle_async_streaming_logging()
raise StopAsyncIteration
raise StopAsyncIteration
+52 -49
View File
@@ -86,7 +86,6 @@ def _get_ImageEditRequestUtils() -> "ImageEditRequestUtils":
return _ImageEditRequestUtils_cache
##### Image Generation #######################
@client
async def aimage_generation(*args, **kwargs) -> ImageResponse:
@@ -212,10 +211,7 @@ def image_generation( # noqa: PLR0915
api_version: Optional[str] = None,
custom_llm_provider=None,
**kwargs,
) -> Union[
ImageResponse,
Coroutine[Any, Any, ImageResponse],
]:
) -> Union[ImageResponse, Coroutine[Any, Any, ImageResponse],]:
"""
Maps the https://api.openai.com/v1/images/generations endpoint.
@@ -346,7 +342,7 @@ def image_generation( # noqa: PLR0915
azure_ad_token = optional_params.pop(
"azure_ad_token", None
) or get_secret_str("AZURE_AD_TOKEN")
# Create azure_ad_token_provider from tenant_id, client_id, client_secret if not already provided
if azure_ad_token_provider is None:
from litellm.llms.azure.common_utils import (
@@ -357,8 +353,11 @@ def image_generation( # noqa: PLR0915
tenant_id = litellm_params_dict.get("tenant_id")
client_id = litellm_params_dict.get("client_id")
client_secret = litellm_params_dict.get("client_secret")
azure_scope = litellm_params_dict.get("azure_scope") or "https://cognitiveservices.azure.com/.default"
azure_scope = (
litellm_params_dict.get("azure_scope")
or "https://cognitiveservices.azure.com/.default"
)
# Create token provider if credentials are available
if tenant_id and client_id and client_secret:
azure_ad_token_provider = get_azure_ad_token_from_entra_id(
@@ -375,7 +374,7 @@ def image_generation( # noqa: PLR0915
# Azure AD authentication will use Authorization header instead
if api_key is not None:
default_headers["api-key"] = api_key
for k, v in default_headers.items():
if k not in headers:
headers[k] = v
@@ -462,7 +461,7 @@ def image_generation( # noqa: PLR0915
# Azure AD authentication will use Authorization header instead
if api_key is not None:
default_headers["api-key"] = api_key
for k, v in default_headers.items():
if k not in headers:
headers[k] = v
@@ -738,7 +737,7 @@ def image_variation(
@client
def image_edit( # noqa: PLR0915
image: Optional[Union[FileTypes, List[FileTypes]]] = None,
prompt: Optional[str]= None,
prompt: Optional[str] = None,
model: Optional[str] = None,
mask: Optional[str] = None,
n: Optional[int] = None,
@@ -762,23 +761,23 @@ def image_edit( # noqa: PLR0915
local_vars = locals()
try:
openai_params = [
"user",
"request_timeout",
"api_base",
"api_version",
"api_key",
"deployment_id",
"organization",
"base_url",
"default_headers",
"timeout",
"max_retries",
"n",
"quality",
"size",
"style",
"async_call",
]
"user",
"request_timeout",
"api_base",
"api_version",
"api_key",
"deployment_id",
"organization",
"base_url",
"default_headers",
"timeout",
"max_retries",
"n",
"quality",
"size",
"style",
"async_call",
]
litellm_params_list = all_litellm_params
default_params = openai_params + litellm_params_list
non_default_params = {
@@ -791,7 +790,9 @@ def image_edit( # noqa: PLR0915
_is_async = kwargs.pop("async_call", False) is True
# add images / or return a single image
images = image if isinstance(image, list) else ([image] if image is not None else [])
images = (
image if isinstance(image, list) else ([image] if image is not None else [])
)
headers_from_kwargs = kwargs.get("headers")
merged_extra_headers: Dict[str, Any] = {}
@@ -864,11 +865,11 @@ def image_edit( # noqa: PLR0915
)
# get provider config
image_edit_provider_config: Optional[BaseImageEditConfig] = (
ProviderConfigManager.get_provider_image_edit_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
)
image_edit_provider_config: Optional[
BaseImageEditConfig
] = ProviderConfigManager.get_provider_image_edit_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
)
if image_edit_provider_config is None:
@@ -877,7 +878,9 @@ def image_edit( # noqa: PLR0915
local_vars.update(kwargs)
# Get ImageEditOptionalRequestParams with only valid parameters
image_edit_optional_params: ImageEditOptionalRequestParams = (
_get_ImageEditRequestUtils().get_requested_image_edit_optional_param(local_vars)
_get_ImageEditRequestUtils().get_requested_image_edit_optional_param(
local_vars
)
)
# Get optional parameters for the responses API
image_edit_request_params: Dict = (
@@ -926,20 +929,20 @@ def image_edit( # noqa: PLR0915
elif custom_llm_provider == "stability":
image_edit_request_params.update(non_default_params)
return base_llm_http_handler.image_edit_handler(
model=model,
image=images,
prompt=prompt,
image_edit_provider_config=image_edit_provider_config,
image_edit_optional_request_params=image_edit_request_params,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
extra_headers=extra_headers,
extra_body=extra_body,
timeout=timeout or DEFAULT_REQUEST_TIMEOUT,
_is_async=_is_async,
client=kwargs.get("client"),
)
model=model,
image=images,
prompt=prompt,
image_edit_provider_config=image_edit_provider_config,
image_edit_optional_request_params=image_edit_request_params,
custom_llm_provider=custom_llm_provider,
litellm_params=litellm_params,
logging_obj=litellm_logging_obj,
extra_headers=extra_headers,
extra_body=extra_body,
timeout=timeout or DEFAULT_REQUEST_TIMEOUT,
_is_async=_is_async,
client=kwargs.get("client"),
)
elif custom_llm_provider == "black_forest_labs":
# Route to BFL-specific handler (polling required)
if model is None:
+1 -3
View File
@@ -40,9 +40,7 @@ class ImageEditRequestUtils:
filtered_optional_params.pop(param, None)
unsupported_params = [
param
for param in filtered_optional_params
if param not in supported_params
param for param in filtered_optional_params if param not in supported_params
]
if unsupported_params:
@@ -102,10 +102,10 @@ class AlertingHangingRequestCheck:
)
for request_id in hanging_requests:
hanging_request_data: Optional[HangingRequestData] = (
await self.hanging_request_cache.async_get_cache(
key=request_id,
)
hanging_request_data: Optional[
HangingRequestData
] = await self.hanging_request_cache.async_get_cache(
key=request_id,
)
if hanging_request_data is None:
@@ -96,7 +96,9 @@ class SlackAlerting(CustomBatchLogger):
self.alert_type_config: Dict[str, AlertTypeConfig] = {}
if alert_type_config:
for key, val in alert_type_config.items():
self.alert_type_config[key] = AlertTypeConfig(**val) if isinstance(val, dict) else val
self.alert_type_config[key] = (
AlertTypeConfig(**val) if isinstance(val, dict) else val
)
self.digest_buckets: Dict[str, DigestEntry] = {}
self.digest_lock = asyncio.Lock()
super().__init__(**kwargs, flush_lock=self.flush_lock)
@@ -126,7 +128,9 @@ class SlackAlerting(CustomBatchLogger):
self.periodic_started = True
if alert_type_config is not None:
for key, val in alert_type_config.items():
self.alert_type_config[key] = AlertTypeConfig(**val) if isinstance(val, dict) else val
self.alert_type_config[key] = (
AlertTypeConfig(**val) if isinstance(val, dict) else val
)
if alert_to_webhook_url is not None:
# update the dict
@@ -1367,7 +1371,7 @@ Model Info:
return False
async def send_alert( # noqa: PLR0915
async def send_alert( # noqa: PLR0915
self,
message: str,
level: Literal["Low", "Medium", "High"],
@@ -1439,7 +1443,9 @@ Model Info:
self.alert_to_webhook_url is not None
and alert_type in self.alert_to_webhook_url
):
_digest_webhook: Optional[Union[str, List[str]]] = self.alert_to_webhook_url[alert_type]
_digest_webhook: Optional[
Union[str, List[str]]
] = self.alert_to_webhook_url[alert_type]
elif self.default_webhook_url is not None:
_digest_webhook = self.default_webhook_url
else:
@@ -1588,11 +1594,21 @@ Model Info:
if isinstance(webhook_url, list):
for url in webhook_url:
self.log_queue.append(
{"url": url, "headers": headers, "payload": payload, "alert_type": alert_type_name}
{
"url": url,
"headers": headers,
"payload": payload,
"alert_type": alert_type_name,
}
)
else:
self.log_queue.append(
{"url": webhook_url, "headers": headers, "payload": payload, "alert_type": alert_type_name}
{
"url": webhook_url,
"headers": headers,
"payload": payload,
"alert_type": alert_type_name,
}
)
flushed_keys.append(key)
+13 -5
View File
@@ -73,11 +73,15 @@ class SpanAttributes:
"""
Number of tokens in the prompt.
"""
LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE = "llm.token_count.prompt_details.cache_write"
LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_WRITE = (
"llm.token_count.prompt_details.cache_write"
)
"""
Number of tokens in the prompt that were written to cache.
"""
LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ = "llm.token_count.prompt_details.cache_read"
LLM_TOKEN_COUNT_PROMPT_DETAILS_CACHE_READ = (
"llm.token_count.prompt_details.cache_read"
)
"""
Number of tokens in the prompt that were read from cache.
"""
@@ -89,11 +93,15 @@ class SpanAttributes:
"""
Number of tokens in the completion.
"""
LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING = "llm.token_count.completion_details.reasoning"
LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING = (
"llm.token_count.completion_details.reasoning"
)
"""
Number of tokens used for reasoning steps in the completion.
"""
LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO = "llm.token_count.completion_details.audio"
LLM_TOKEN_COUNT_COMPLETION_DETAILS_AUDIO = (
"llm.token_count.completion_details.audio"
)
"""
The number of audio input tokens generated by the model
"""
@@ -396,7 +404,7 @@ class OpenInferenceLLMProviderValues(Enum):
class ErrorAttributes:
"""
Attributes for error information in spans.
These attributes follow OpenTelemetry semantic conventions for exceptions
and are used to record error information from StandardLoggingPayloadErrorInformation.
"""
+1 -1
View File
@@ -1,3 +1,3 @@
from .agentops import AgentOps
__all__ = ["AgentOps"]
__all__ = ["AgentOps"]
+17 -19
View File
@@ -7,6 +7,7 @@ from typing import Optional, Dict, Any
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
@dataclass
class AgentOpsConfig:
endpoint: str = "https://otlp.agentops.cloud/v1/traces"
@@ -22,9 +23,10 @@ class AgentOpsConfig:
api_key=os.getenv("AGENTOPS_API_KEY"),
service_name=os.getenv("AGENTOPS_SERVICE_NAME", "agentops"),
deployment_environment=os.getenv("AGENTOPS_ENVIRONMENT", "production"),
auth_endpoint="https://api.agentops.ai/v3/auth/token"
auth_endpoint="https://api.agentops.ai/v3/auth/token",
)
class AgentOps(OpenTelemetry):
"""
AgentOps integration - built on top of OpenTelemetry
@@ -32,7 +34,7 @@ class AgentOps(OpenTelemetry):
Example usage:
```python
import litellm
litellm.success_callback = ["agentops"]
response = litellm.completion(
@@ -41,6 +43,7 @@ class AgentOps(OpenTelemetry):
)
```
"""
def __init__(
self,
config: Optional[AgentOpsConfig] = None,
@@ -60,18 +63,13 @@ class AgentOps(OpenTelemetry):
pass
headers = f"Authorization=Bearer {jwt_token}" if jwt_token else None
otel_config = OpenTelemetryConfig(
exporter="otlp_http",
endpoint=config.endpoint,
headers=headers
exporter="otlp_http", endpoint=config.endpoint, headers=headers
)
# Initialize OpenTelemetry with our config
super().__init__(
config=otel_config,
callback_name="agentops"
)
super().__init__(config=otel_config, callback_name="agentops")
# Set AgentOps-specific resource attributes
resource_attrs = {
@@ -79,20 +77,20 @@ class AgentOps(OpenTelemetry):
"deployment.environment": config.deployment_environment or "production",
"telemetry.sdk.name": "agentops",
}
if project_id:
resource_attrs["project.id"] = project_id
self.resource_attributes = resource_attrs
def _fetch_auth_token(self, api_key: str, auth_endpoint: str) -> Dict[str, Any]:
"""
Fetch JWT authentication token from AgentOps API
Args:
api_key: AgentOps API key
auth_endpoint: Authentication endpoint
Returns:
Dict containing JWT token and project ID
"""
@@ -100,19 +98,19 @@ class AgentOps(OpenTelemetry):
"Content-Type": "application/json",
"Connection": "keep-alive",
}
client = _get_httpx_client()
try:
response = client.post(
url=auth_endpoint,
headers=headers,
json={"api_key": api_key},
timeout=10
timeout=10,
)
if response.status_code != 200:
raise Exception(f"Failed to fetch auth token: {response.text}")
return response.json()
finally:
client.close()
client.close()
@@ -99,10 +99,10 @@ class AnthropicCacheControlHook(CustomPromptManagement):
targetted_index += len(messages)
if 0 <= targetted_index < len(messages):
messages[targetted_index] = (
AnthropicCacheControlHook._safe_insert_cache_control_in_message(
messages[targetted_index], control
)
messages[
targetted_index
] = AnthropicCacheControlHook._safe_insert_cache_control_in_message(
messages[targetted_index], control
)
else:
verbose_logger.warning(
+77 -25
View File
@@ -14,12 +14,12 @@ from litellm.types.utils import StandardLoggingPayload
if TYPE_CHECKING:
from opentelemetry.trace import Span
from litellm.integrations._types.open_inference import (
MessageAttributes,
ImageAttributes,
SpanAttributes,
AudioAttributes,
EmbeddingAttributes,
OpenInferenceSpanKindValues
MessageAttributes,
ImageAttributes,
SpanAttributes,
AudioAttributes,
EmbeddingAttributes,
OpenInferenceSpanKindValues,
)
@@ -158,7 +158,9 @@ def _set_audio_outputs(span: "Span", response_obj, audio_attrs, span_attrs):
audio_transcript = audio_item.get("transcript")
if audio_transcript:
safe_set_attribute(span, f"{audio_attrs.AUDIO_TRANSCRIPT}.{i}", audio_transcript)
safe_set_attribute(
span, f"{audio_attrs.AUDIO_TRANSCRIPT}.{i}", audio_transcript
)
def _set_embedding_outputs(span: "Span", response_obj, embedding_attrs, span_attrs):
@@ -212,7 +214,9 @@ def _set_structured_outputs(span: "Span", response_obj, msg_attrs, span_attrs):
message_content = getattr(first_content, "text", "")
message_role = getattr(item, "role", "assistant")
safe_set_attribute(span, span_attrs.OUTPUT_VALUE, message_content)
safe_set_attribute(span, f"{prefix}.{msg_attrs.MESSAGE_CONTENT}", message_content)
safe_set_attribute(
span, f"{prefix}.{msg_attrs.MESSAGE_CONTENT}", message_content
)
safe_set_attribute(span, f"{prefix}.{msg_attrs.MESSAGE_ROLE}", message_role)
@@ -221,16 +225,24 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs):
if not usage:
return
safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_TOTAL, usage.get("total_tokens"))
safe_set_attribute(
span, span_attrs.LLM_TOKEN_COUNT_TOTAL, usage.get("total_tokens")
)
completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens")
if completion_tokens:
safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_COMPLETION, completion_tokens)
safe_set_attribute(
span, span_attrs.LLM_TOKEN_COUNT_COMPLETION, completion_tokens
)
prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens")
if prompt_tokens:
safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_PROMPT, prompt_tokens)
reasoning_tokens = usage.get("output_tokens_details", {}).get("reasoning_tokens")
if reasoning_tokens:
safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, reasoning_tokens)
safe_set_attribute(
span,
span_attrs.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING,
reasoning_tokens,
)
def _infer_open_inference_span_kind(call_type: Optional[str]) -> str:
@@ -281,11 +293,15 @@ def _infer_open_inference_span_kind(call_type: Optional[str]) -> str:
):
return OpenInferenceSpanKindValues.LLM.value
if any(keyword in lowered for keyword in ("file", "batch", "container", "fine_tuning_job")):
if any(
keyword in lowered
for keyword in ("file", "batch", "container", "fine_tuning_job")
):
return OpenInferenceSpanKindValues.CHAIN.value
return OpenInferenceSpanKindValues.UNKNOWN.value
def _set_tool_attributes(
span: "Span", optional_tools: Optional[list], metadata_tools: Optional[list]
):
@@ -294,18 +310,30 @@ def _set_tool_attributes(
for idx, tool in enumerate(optional_tools):
if not isinstance(tool, dict):
continue
function = tool.get("function") if isinstance(tool.get("function"), dict) else None
function = (
tool.get("function") if isinstance(tool.get("function"), dict) else None
)
if not function:
continue
tool_name = function.get("name")
if tool_name:
safe_set_attribute(span, f"{SpanAttributes.LLM_TOOLS}.{idx}.name", tool_name)
safe_set_attribute(
span, f"{SpanAttributes.LLM_TOOLS}.{idx}.name", tool_name
)
tool_description = function.get("description")
if tool_description:
safe_set_attribute(span, f"{SpanAttributes.LLM_TOOLS}.{idx}.description", tool_description)
safe_set_attribute(
span,
f"{SpanAttributes.LLM_TOOLS}.{idx}.description",
tool_description,
)
params = function.get("parameters")
if params is not None:
safe_set_attribute(span, f"{SpanAttributes.LLM_TOOLS}.{idx}.parameters", json.dumps(params))
safe_set_attribute(
span,
f"{SpanAttributes.LLM_TOOLS}.{idx}.parameters",
json.dumps(params),
)
if metadata_tools and isinstance(metadata_tools, list):
for idx, tool in enumerate(metadata_tools):
@@ -343,7 +371,11 @@ def set_attributes(
if standard_logging_payload is None:
raise ValueError("standard_logging_object not found in kwargs")
metadata = standard_logging_payload.get("metadata") if standard_logging_payload else None
metadata = (
standard_logging_payload.get("metadata")
if standard_logging_payload
else None
)
_set_metadata_attributes(span, metadata, SpanAttributes)
metadata_tools = _extract_metadata_tools(metadata)
@@ -362,13 +394,19 @@ def set_attributes(
span_kind = _infer_open_inference_span_kind(call_type=call_type)
_set_tool_attributes(span, optional_tools, metadata_tools)
if (optional_tools or metadata_tools) and span_kind != OpenInferenceSpanKindValues.TOOL.value:
if (
optional_tools or metadata_tools
) and span_kind != OpenInferenceSpanKindValues.TOOL.value:
span_kind = OpenInferenceSpanKindValues.TOOL.value
safe_set_attribute(span, SpanAttributes.OPENINFERENCE_SPAN_KIND, span_kind)
attributes.set_messages(span, kwargs)
model_params = standard_logging_payload.get("model_parameters") if standard_logging_payload else None
model_params = (
standard_logging_payload.get("model_parameters")
if standard_logging_payload
else None
)
_set_model_params(span, model_params, SpanAttributes)
_set_response_attributes(span=span, response_obj=response_obj)
@@ -418,17 +456,29 @@ def _set_request_attributes(
if kwargs.get("model"):
safe_set_attribute(span, span_attrs.LLM_MODEL_NAME, kwargs.get("model"))
safe_set_attribute(span, "llm.request.type", standard_logging_payload.get("call_type"))
safe_set_attribute(span, span_attrs.LLM_PROVIDER, litellm_params.get("custom_llm_provider", "Unknown"))
safe_set_attribute(
span, "llm.request.type", standard_logging_payload.get("call_type")
)
safe_set_attribute(
span,
span_attrs.LLM_PROVIDER,
litellm_params.get("custom_llm_provider", "Unknown"),
)
if optional_params.get("max_tokens"):
safe_set_attribute(span, "llm.request.max_tokens", optional_params.get("max_tokens"))
safe_set_attribute(
span, "llm.request.max_tokens", optional_params.get("max_tokens")
)
if optional_params.get("temperature"):
safe_set_attribute(span, "llm.request.temperature", optional_params.get("temperature"))
safe_set_attribute(
span, "llm.request.temperature", optional_params.get("temperature")
)
if optional_params.get("top_p"):
safe_set_attribute(span, "llm.request.top_p", optional_params.get("top_p"))
safe_set_attribute(span, "llm.is_streaming", str(optional_params.get("stream", False)))
safe_set_attribute(
span, "llm.is_streaming", str(optional_params.get("stream", False))
)
if optional_params.get("user"):
safe_set_attribute(span, "llm.user", optional_params.get("user"))
@@ -443,7 +493,9 @@ def _set_model_params(span: "Span", model_params: Optional[dict], span_attrs) ->
if not model_params:
return
safe_set_attribute(span, span_attrs.LLM_INVOCATION_PARAMETERS, safe_dumps(model_params))
safe_set_attribute(
span, span_attrs.LLM_INVOCATION_PARAMETERS, safe_dumps(model_params)
)
if model_params.get("user"):
user_id = model_params.get("user")
if user_id is not None:
+18 -11
View File
@@ -12,7 +12,9 @@ if TYPE_CHECKING:
from opentelemetry.trace import SpanKind
from litellm.integrations.opentelemetry import OpenTelemetry as _OpenTelemetry
from litellm.integrations.opentelemetry import OpenTelemetryConfig as _OpenTelemetryConfig
from litellm.integrations.opentelemetry import (
OpenTelemetryConfig as _OpenTelemetryConfig,
)
from litellm.types.integrations.arize import Protocol as _Protocol
Protocol = _Protocol
@@ -91,7 +93,9 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
@staticmethod
def set_arize_phoenix_attributes(span: Span, kwargs, response_obj):
from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import safe_set_attribute
from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import (
safe_set_attribute,
)
_utils.set_attributes(span, kwargs, response_obj, ArizeOTELAttributes)
@@ -103,7 +107,9 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
# Fall back to static config from env var
config = ArizePhoenixLogger.get_arize_phoenix_config()
if config.project_name:
safe_set_attribute(span, "openinference.project.name", config.project_name)
safe_set_attribute(
span, "openinference.project.name", config.project_name
)
return
@@ -172,7 +178,9 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
start_time_val = kwargs.get("start_time", kwargs.get("api_call_start_time"))
parent_span = self.tracer.start_span(
name="litellm_proxy_request",
start_time=self._to_ns(start_time_val) if start_time_val is not None else None,
start_time=self._to_ns(start_time_val)
if start_time_val is not None
else None,
context=traceparent_ctx,
kind=self.span_kind.SERVER,
)
@@ -212,9 +220,7 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
# Raw-request sub-span (if enabled) — must be created before
# ending the parent span so the hierarchy is valid.
self._maybe_log_raw_request(
kwargs, response_obj, start_time, end_time, span
)
self._maybe_log_raw_request(kwargs, response_obj, start_time, end_time, span)
span.end(end_time=self._to_ns(end_time))
# Guardrail span
@@ -290,7 +296,9 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
if collector_endpoint:
# Parse the endpoint to determine protocol
if collector_endpoint.startswith("grpc://") or (":4317" in collector_endpoint and "/v1/traces" not in collector_endpoint):
if collector_endpoint.startswith("grpc://") or (
":4317" in collector_endpoint and "/v1/traces" not in collector_endpoint
):
endpoint = collector_endpoint
protocol = "otlp_grpc"
else:
@@ -334,11 +342,10 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
endpoint=endpoint,
project_name=project_name,
)
## cannot suppress additional proxy server spans, removed previous methods.
async def async_health_check(self):
config = self.get_arize_phoenix_config()
if not config.otlp_auth_headers:
@@ -350,4 +357,4 @@ class ArizePhoenixLogger(OpenTelemetry): # type: ignore
return {
"status": "healthy",
"message": "Arize-Phoenix credentials are configured properly",
}
}
@@ -1,4 +1,3 @@
from litellm.integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger
__all__ = ["AzureSentinelLogger"]
@@ -62,18 +62,22 @@ class AzureSentinelLogger(CustomBatchLogger):
llm_provider=httpxSpecialProvider.LoggingCallback
)
self.dcr_immutable_id = (
dcr_immutable_id or os.getenv("AZURE_SENTINEL_DCR_IMMUTABLE_ID")
self.dcr_immutable_id = dcr_immutable_id or os.getenv(
"AZURE_SENTINEL_DCR_IMMUTABLE_ID"
)
self.stream_name = stream_name or os.getenv(
"AZURE_SENTINEL_STREAM_NAME", "Custom-LiteLLM"
)
self.endpoint = endpoint or os.getenv("AZURE_SENTINEL_ENDPOINT")
self.tenant_id = tenant_id or os.getenv("AZURE_SENTINEL_TENANT_ID") or os.getenv(
"AZURE_TENANT_ID"
self.tenant_id = (
tenant_id
or os.getenv("AZURE_SENTINEL_TENANT_ID")
or os.getenv("AZURE_TENANT_ID")
)
self.client_id = client_id or os.getenv("AZURE_SENTINEL_CLIENT_ID") or os.getenv(
"AZURE_CLIENT_ID"
self.client_id = (
client_id
or os.getenv("AZURE_SENTINEL_CLIENT_ID")
or os.getenv("AZURE_CLIENT_ID")
)
self.client_secret = (
client_secret
@@ -103,9 +107,7 @@ class AzureSentinelLogger(CustomBatchLogger):
)
# Build API endpoint: {Endpoint}/dataCollectionRules/{DCR Immutable ID}/streams/{Stream Name}?api-version=2023-01-01
self.api_endpoint = (
f"{self.endpoint.rstrip('/')}/dataCollectionRules/{self.dcr_immutable_id}/streams/{self.stream_name}?api-version=2023-01-01"
)
self.api_endpoint = f"{self.endpoint.rstrip('/')}/dataCollectionRules/{self.dcr_immutable_id}/streams/{self.stream_name}?api-version=2023-01-01"
# OAuth2 scope for Azure Monitor
self.oauth_scope = "https://monitor.azure.com/.default"
@@ -139,7 +141,9 @@ class AzureSentinelLogger(CustomBatchLogger):
assert self.client_id is not None, "client_id is required"
assert self.client_secret is not None, "client_secret is required"
token_url = f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token"
token_url = (
f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token"
)
token_data = {
"client_id": self.client_id,
@@ -173,9 +177,7 @@ class AzureSentinelLogger(CustomBatchLogger):
return self.oauth_token
async def async_log_success_event(
self, kwargs, response_obj, start_time, end_time
):
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
"""
Async Log success events to Azure Sentinel
@@ -209,9 +211,7 @@ class AzureSentinelLogger(CustomBatchLogger):
)
pass
async def async_log_failure_event(
self, kwargs, response_obj, start_time, end_time
):
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
"""
Async Log failure events to Azure Sentinel
@@ -54,12 +54,12 @@ class AzureBlobStorageLogger(CustomBatchLogger):
self._service_client_timeout: Optional[float] = None
# Internal variables used for Token based authentication
self.azure_auth_token: Optional[str] = (
None # the Azure AD token to use for Azure Storage API requests
)
self.token_expiry: Optional[datetime] = (
None # the expiry time of the currentAzure AD token
)
self.azure_auth_token: Optional[
str
] = None # the Azure AD token to use for Azure Storage API requests
self.token_expiry: Optional[
datetime
] = None # the expiry time of the currentAzure AD token
asyncio.create_task(self.periodic_flush())
self.flush_lock = asyncio.Lock()
+9 -7
View File
@@ -41,7 +41,9 @@ class BraintrustLogger(CustomLogger):
self.is_mock_mode = should_use_braintrust_mock()
if self.is_mock_mode:
create_mock_braintrust_client()
verbose_logger.info("[BRAINTRUST MOCK] Braintrust logger initialized in mock mode")
verbose_logger.info(
"[BRAINTRUST MOCK] Braintrust logger initialized in mock mode"
)
self.validate_environment(api_key=api_key)
self.api_base = api_base or os.getenv("BRAINTRUST_API_BASE") or API_BASE
self.default_project_id = None
@@ -50,9 +52,9 @@ class BraintrustLogger(CustomLogger):
"Authorization": "Bearer " + self.api_key,
"Content-Type": "application/json",
}
self._project_id_cache: Dict[str, str] = (
{}
) # Cache mapping project names to IDs
self._project_id_cache: Dict[
str, str
] = {} # Cache mapping project names to IDs
self.global_braintrust_http_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
@@ -214,7 +216,7 @@ class BraintrustLogger(CustomLogger):
# Allow metadata override for span name
span_name = dynamic_metadata.get("span_name", "Chat Completion")
# Span parents is a special case
span_parents = dynamic_metadata.get("span_parents")
@@ -236,7 +238,7 @@ class BraintrustLogger(CustomLogger):
"span_attributes": {"name": span_name, "type": "llm"},
}
# Braintrust cannot specify 'tags' for non-root spans
# Braintrust cannot specify 'tags' for non-root spans
if dynamic_metadata.get("root_span_id") is None:
request_data["tags"] = tags
@@ -386,7 +388,7 @@ class BraintrustLogger(CustomLogger):
"span_attributes": {"name": span_name, "type": "llm"},
}
# Braintrust cannot specify 'tags' for non-root spans
# Braintrust cannot specify 'tags' for non-root spans
if dynamic_metadata.get("root_span_id") is None:
request_data["tags"] = tags
+52 -17
View File
@@ -13,7 +13,11 @@ import time
from urllib.parse import urlparse
from litellm._logging import verbose_logger
from litellm.integrations.mock_client_factory import MockClientConfig, MockResponse, create_mock_client_factory
from litellm.integrations.mock_client_factory import (
MockClientConfig,
MockResponse,
create_mock_client_factory,
)
# Use factory for should_use_mock and MockResponse
# Braintrust uses both HTTPHandler (sync) and AsyncHTTPHandler (async)
@@ -37,7 +41,10 @@ _config = MockClientConfig(
# Get should_use_mock and create_mock_client from factory
# We need to call the factory's create_mock_client to patch AsyncHTTPHandler.post
create_mock_braintrust_factory_client, should_use_braintrust_mock = create_mock_client_factory(_config)
(
create_mock_braintrust_factory_client,
should_use_braintrust_mock,
) = create_mock_client_factory(_config)
# Store original HTTPHandler.post method (Braintrust-specific for sync calls with custom logic)
_original_http_handler_post = None
@@ -66,7 +73,19 @@ def _is_braintrust_url(url: str) -> bool:
)
def _mock_http_handler_post(self, url, data=None, json=None, params=None, headers=None, timeout=None, stream=False, files=None, content=None, logging_obj=None):
def _mock_http_handler_post(
self,
url,
data=None,
json=None,
params=None,
headers=None,
timeout=None,
stream=False,
files=None,
content=None,
logging_obj=None,
):
"""Monkey-patched HTTPHandler.post that intercepts Braintrust calls with endpoint-specific responses."""
# Only mock Braintrust API calls
if isinstance(url, str) and _is_braintrust_url(url):
@@ -86,46 +105,62 @@ def _mock_http_handler_post(self, url, data=None, json=None, params=None, header
status_code=_config.default_status_code,
json_data=mock_data,
url=url,
elapsed_seconds=_MOCK_LATENCY_SECONDS
elapsed_seconds=_MOCK_LATENCY_SECONDS,
)
if _original_http_handler_post is not None:
return _original_http_handler_post(self, url=url, data=data, json=json, params=params, headers=headers, timeout=timeout, stream=stream, files=files, content=content, logging_obj=logging_obj)
return _original_http_handler_post(
self,
url=url,
data=data,
json=json,
params=params,
headers=headers,
timeout=timeout,
stream=stream,
files=files,
content=content,
logging_obj=logging_obj,
)
raise RuntimeError("Original HTTPHandler.post not available")
def create_mock_braintrust_client():
"""
Monkey-patch HTTPHandler.post to intercept Braintrust sync calls.
Braintrust uses HTTPHandler for sync calls and AsyncHTTPHandler for async calls.
HTTPHandler.post uses self.client.send(), not self.client.post(), so we need
custom patching for sync (similar to Helicone).
AsyncHTTPHandler.post is patched by the factory.
We use custom patching instead of factory's patch_http_handler because we need
endpoint-specific responses (different for /project vs /project_logs).
This function is idempotent - it only initializes mocks once, even if called multiple times.
"""
global _original_http_handler_post, _mocks_initialized
if _mocks_initialized:
return
verbose_logger.debug("[BRAINTRUST MOCK] Initializing Braintrust mock client...")
from litellm.llms.custom_httpx.http_handler import HTTPHandler
if _original_http_handler_post is None:
_original_http_handler_post = HTTPHandler.post
HTTPHandler.post = _mock_http_handler_post # type: ignore
verbose_logger.debug("[BRAINTRUST MOCK] Patched HTTPHandler.post")
# CRITICAL: Call the factory's initialization function to patch AsyncHTTPHandler.post
# This is required for async calls to be mocked
create_mock_braintrust_factory_client()
verbose_logger.debug(f"[BRAINTRUST MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms")
verbose_logger.debug("[BRAINTRUST MOCK] Braintrust mock client initialization complete")
verbose_logger.debug(
f"[BRAINTRUST MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms"
)
verbose_logger.debug(
"[BRAINTRUST MOCK] Braintrust mock client initialization complete"
)
_mocks_initialized = True
@@ -30,7 +30,9 @@ class CZEntityType(str, Enum):
class CZRNGenerator:
"""Generate CloudZero Resource Names (CZRNs) for LiteLLM resources."""
CZRN_REGEX = re.compile(r'^czrn:([a-z0-9-]+):([a-zA-Z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):(.+)$')
CZRN_REGEX = re.compile(
r"^czrn:([a-z0-9-]+):([a-zA-Z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):([a-z0-9-]+):(.+)$"
)
def __init__(self):
"""Initialize CZRN generator."""
@@ -38,9 +40,9 @@ class CZRNGenerator:
def create_from_litellm_data(self, row: dict[str, Any]) -> str:
"""Create a CZRN from LiteLLM daily spend data.
CZRN format: czrn:<service-type>:<provider>:<region>:<owner-account-id>:<resource-type>:<cloud-local-id>
For LiteLLM resources, we map:
- service-type: 'litellm' (the service managing the LLM calls)
- provider: The custom_llm_provider (e.g., 'openai', 'anthropic', 'azure')
@@ -49,18 +51,18 @@ class CZRNGenerator:
- resource-type: 'llm-usage' (represents LLM usage/inference)
- cloud-local-id: model
"""
service_type = 'litellm'
provider = self._normalize_provider(row.get('custom_llm_provider', 'unknown'))
region = 'cross-region'
service_type = "litellm"
provider = self._normalize_provider(row.get("custom_llm_provider", "unknown"))
region = "cross-region"
# Use the actual entity_id (team_id or user_id) as the owner account
team_id = row.get('team_id', 'unknown')
team_id = row.get("team_id", "unknown")
owner_account_id = self._normalize_component(team_id)
resource_type = 'llm-usage'
resource_type = "llm-usage"
# Create a unique identifier with just the model (entity info already in owner_account_id)
model = row.get('model', 'unknown')
model = row.get("model", "unknown")
cloud_local_id = model
@@ -70,7 +72,7 @@ class CZRNGenerator:
region=region,
owner_account_id=owner_account_id,
resource_type=resource_type,
cloud_local_id=cloud_local_id
cloud_local_id=cloud_local_id,
)
def create_from_components(
@@ -80,7 +82,7 @@ class CZRNGenerator:
region: str,
owner_account_id: str,
resource_type: str,
cloud_local_id: str
cloud_local_id: str,
) -> str:
"""Create a CZRN from individual components."""
# Normalize components to ensure they meet CZRN requirements
@@ -104,7 +106,7 @@ class CZRNGenerator:
def extract_components(self, czrn: str) -> tuple[str, str, str, str, str, str]:
"""Extract all components from a CZRN.
Returns: (service_type, provider, region, owner_account_id, resource_type, cloud_local_id)
"""
match = self.CZRN_REGEX.match(czrn)
@@ -117,42 +119,43 @@ class CZRNGenerator:
"""Normalize provider names to standard CZRN format."""
# Map common provider names to CZRN standards
provider_map = {
litellm.LlmProviders.AZURE.value: 'azure',
litellm.LlmProviders.AZURE_AI.value: 'azure',
litellm.LlmProviders.ANTHROPIC.value: 'anthropic',
litellm.LlmProviders.BEDROCK.value: 'aws',
litellm.LlmProviders.VERTEX_AI.value: 'gcp',
litellm.LlmProviders.GEMINI.value: 'google',
litellm.LlmProviders.COHERE.value: 'cohere',
litellm.LlmProviders.HUGGINGFACE.value: 'huggingface',
litellm.LlmProviders.REPLICATE.value: 'replicate',
litellm.LlmProviders.TOGETHER_AI.value: 'together-ai',
litellm.LlmProviders.AZURE.value: "azure",
litellm.LlmProviders.AZURE_AI.value: "azure",
litellm.LlmProviders.ANTHROPIC.value: "anthropic",
litellm.LlmProviders.BEDROCK.value: "aws",
litellm.LlmProviders.VERTEX_AI.value: "gcp",
litellm.LlmProviders.GEMINI.value: "google",
litellm.LlmProviders.COHERE.value: "cohere",
litellm.LlmProviders.HUGGINGFACE.value: "huggingface",
litellm.LlmProviders.REPLICATE.value: "replicate",
litellm.LlmProviders.TOGETHER_AI.value: "together-ai",
}
normalized = provider.lower().replace('_', '-')
normalized = provider.lower().replace("_", "-")
# use litellm custom llm provider if not in provider_map
if normalized not in provider_map:
return normalized
return provider_map.get(normalized, normalized)
def _normalize_component(self, component: str, allow_uppercase: bool = False) -> str:
def _normalize_component(
self, component: str, allow_uppercase: bool = False
) -> str:
"""Normalize a CZRN component to meet format requirements."""
if not component:
return 'unknown'
return "unknown"
# Convert to lowercase unless uppercase is allowed
if not allow_uppercase:
component = component.lower()
# Replace invalid characters with hyphens
component = re.sub(r'[^a-zA-Z0-9-]', '-', component)
component = re.sub(r"[^a-zA-Z0-9-]", "-", component)
# Remove consecutive hyphens
component = re.sub(r'-+', '-', component)
component = re.sub(r"-+", "-", component)
# Remove leading/trailing hyphens
component = component.strip('-')
return component or 'unknown'
component = component.strip("-")
return component or "unknown"
+95 -44
View File
@@ -30,7 +30,9 @@ from rich.console import Console
class CloudZeroStreamer:
"""Stream CBF data to CloudZero AnyCost API with proper batching and timezone handling."""
def __init__(self, api_key: str, connection_id: str, user_timezone: Optional[str] = None):
def __init__(
self, api_key: str, connection_id: str, user_timezone: Optional[str] = None
):
"""Initialize CloudZero streamer with credentials."""
self.api_key = api_key
self.connection_id = connection_id
@@ -43,12 +45,16 @@ class CloudZeroStreamer:
try:
self.user_timezone = zoneinfo.ZoneInfo(user_timezone)
except zoneinfo.ZoneInfoNotFoundError:
self.console.print(f"[yellow]Warning: Unknown timezone '{user_timezone}', using UTC[/yellow]")
self.console.print(
f"[yellow]Warning: Unknown timezone '{user_timezone}', using UTC[/yellow]"
)
self.user_timezone = timezone.utc
else:
self.user_timezone = timezone.utc
def send_batched(self, data: pl.DataFrame, operation: str = "replace_hourly") -> None:
def send_batched(
self, data: pl.DataFrame, operation: str = "replace_hourly"
) -> None:
"""Send CBF data in daily batches to CloudZero AnyCost API."""
if data.is_empty():
self.console.print("[yellow]No data to send to CloudZero[/yellow]")
@@ -61,7 +67,9 @@ class CloudZeroStreamer:
self.console.print("[yellow]No valid daily batches to send[/yellow]")
return
self.console.print(f"[blue]Sending {len(daily_batches)} daily batch(es) with operation '{operation}'[/blue]")
self.console.print(
f"[blue]Sending {len(daily_batches)} daily batch(es) with operation '{operation}'[/blue]"
)
for batch_date, batch_data in daily_batches.items():
self._send_daily_batch(batch_date, batch_data, operation)
@@ -71,21 +79,23 @@ class CloudZeroStreamer:
daily_batches: dict[str, list[dict[str, Any]]] = {}
# Ensure we have the required columns
if 'time/usage_start' not in data.columns:
self.console.print("[red]Error: Missing 'time/usage_start' column for date grouping[/red]")
if "time/usage_start" not in data.columns:
self.console.print(
"[red]Error: Missing 'time/usage_start' column for date grouping[/red]"
)
return {}
timestamp_str: Optional[str] = None
for row in data.iter_rows(named=True):
try:
# Parse the timestamp and convert to UTC
timestamp_str = row.get('time/usage_start')
timestamp_str = row.get("time/usage_start")
if not timestamp_str:
continue
# Parse timestamp and handle timezone conversion
dt = self._parse_and_convert_timestamp(timestamp_str)
batch_date = dt.strftime('%Y-%m-%d')
batch_date = dt.strftime("%Y-%m-%d")
if batch_date not in daily_batches:
daily_batches[batch_date] = []
@@ -93,25 +103,54 @@ class CloudZeroStreamer:
daily_batches[batch_date].append(row)
except Exception as e:
self.console.print(f"[yellow]Warning: Could not process timestamp '{timestamp_str}': {e}[/yellow]")
self.console.print(
f"[yellow]Warning: Could not process timestamp '{timestamp_str}': {e}[/yellow]"
)
continue
# Convert lists back to DataFrames
return {date_key: pl.DataFrame(records) for date_key, records in daily_batches.items() if records}
return {
date_key: pl.DataFrame(records)
for date_key, records in daily_batches.items()
if records
}
def _parse_and_convert_timestamp(self, timestamp_str: str) -> datetime:
"""Parse timestamp string and convert to UTC."""
# Try to parse the timestamp string
try:
# Handle various ISO 8601 formats
if timestamp_str.endswith('Z'):
dt = datetime.fromisoformat(timestamp_str.replace('Z', '+00:00'))
elif '+' in timestamp_str or timestamp_str.endswith(('-00:00', '-01:00', '-02:00', '-03:00',
'-04:00', '-05:00', '-06:00', '-07:00',
'-08:00', '-09:00', '-10:00', '-11:00',
'-12:00', '+01:00', '+02:00', '+03:00',
'+04:00', '+05:00', '+06:00', '+07:00',
'+08:00', '+09:00', '+10:00', '+11:00', '+12:00')):
if timestamp_str.endswith("Z"):
dt = datetime.fromisoformat(timestamp_str.replace("Z", "+00:00"))
elif "+" in timestamp_str or timestamp_str.endswith(
(
"-00:00",
"-01:00",
"-02:00",
"-03:00",
"-04:00",
"-05:00",
"-06:00",
"-07:00",
"-08:00",
"-09:00",
"-10:00",
"-11:00",
"-12:00",
"+01:00",
"+02:00",
"+03:00",
"+04:00",
"+05:00",
"+06:00",
"+07:00",
"+08:00",
"+09:00",
"+10:00",
"+11:00",
"+12:00",
)
):
dt = datetime.fromisoformat(timestamp_str)
else:
# Assume user timezone if no timezone info
@@ -125,14 +164,16 @@ class CloudZeroStreamer:
except ValueError as e:
raise ValueError(f"Could not parse timestamp '{timestamp_str}': {e}")
def _send_daily_batch(self, batch_date: str, batch_data: pl.DataFrame, operation: str) -> None:
def _send_daily_batch(
self, batch_date: str, batch_data: pl.DataFrame, operation: str
) -> None:
"""Send a single daily batch to CloudZero API."""
if batch_data.is_empty():
return
headers = {
'Authorization': f'Bearer {self.api_key}',
'Content-Type': 'application/json'
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json",
}
# Use the correct API endpoint format from documentation
@@ -143,29 +184,39 @@ class CloudZeroStreamer:
try:
with httpx.Client(timeout=30.0) as client:
self.console.print(f"[blue]Sending batch for {batch_date} ({len(batch_data)} records)[/blue]")
self.console.print(
f"[blue]Sending batch for {batch_date} ({len(batch_data)} records)[/blue]"
)
response = client.post(url, headers=headers, json=payload)
response.raise_for_status()
self.console.print(f"[green]✓ Successfully sent batch for {batch_date} ({len(batch_data)} records)[/green]")
self.console.print(
f"[green]✓ Successfully sent batch for {batch_date} ({len(batch_data)} records)[/green]"
)
except httpx.RequestError as e:
self.console.print(f"[red]✗ Network error sending batch for {batch_date}: {e}[/red]")
self.console.print(
f"[red]✗ Network error sending batch for {batch_date}: {e}[/red]"
)
raise
except httpx.HTTPStatusError as e:
self.console.print(f"[red]✗ HTTP error sending batch for {batch_date}: {e.response.status_code} {e.response.text}[/red]")
self.console.print(
f"[red]✗ HTTP error sending batch for {batch_date}: {e.response.status_code} {e.response.text}[/red]"
)
raise
def _prepare_batch_payload(self, batch_date: str, batch_data: pl.DataFrame, operation: str) -> dict[str, Any]:
def _prepare_batch_payload(
self, batch_date: str, batch_data: pl.DataFrame, operation: str
) -> dict[str, Any]:
"""Prepare batch payload according to CloudZero AnyCost API format."""
# Convert batch_date to month for the API (YYYY-MM format)
try:
date_obj = datetime.strptime(batch_date, '%Y-%m-%d')
month_str = date_obj.strftime('%Y-%m')
date_obj = datetime.strptime(batch_date, "%Y-%m-%d")
month_str = date_obj.strftime("%Y-%m")
except ValueError:
# Fallback to current month
month_str = datetime.now().strftime('%Y-%m')
month_str = datetime.now().strftime("%Y-%m")
# Convert DataFrame rows to API format
data_records = []
@@ -174,15 +225,13 @@ class CloudZeroStreamer:
if record:
data_records.append(record)
payload = {
'month': month_str,
'operation': operation,
'data': data_records
}
payload = {"month": month_str, "operation": operation, "data": data_records}
return payload
def _convert_cbf_to_api_format(self, row: dict[str, Any]) -> Optional[dict[str, Any]]:
def _convert_cbf_to_api_format(
self, row: dict[str, Any]
) -> Optional[dict[str, Any]]:
"""Convert CBF row to CloudZero API format - keeping CBF field names as CloudZero expects them."""
try:
# CloudZero expects CBF format field names directly, not converted names
@@ -196,20 +245,24 @@ class CloudZeroStreamer:
# Format floats to avoid scientific notation
if isinstance(value, float):
# Use a reasonable precision that avoids scientific notation
api_record[key] = f"{value:.10f}".rstrip('0').rstrip('.')
api_record[key] = f"{value:.10f}".rstrip("0").rstrip(".")
else:
api_record[key] = str(value)
else:
api_record[key] = value
# Ensure timestamp is in UTC format
if 'time/usage_start' in api_record:
api_record['time/usage_start'] = self._ensure_utc_timestamp(api_record['time/usage_start'])
if "time/usage_start" in api_record:
api_record["time/usage_start"] = self._ensure_utc_timestamp(
api_record["time/usage_start"]
)
return api_record
except Exception as e:
self.console.print(f"[yellow]Warning: Could not convert record to API format: {e}[/yellow]")
self.console.print(
f"[yellow]Warning: Could not convert record to API format: {e}[/yellow]"
)
return None
def _ensure_utc_timestamp(self, timestamp_str: str) -> str:
@@ -219,9 +272,7 @@ class CloudZeroStreamer:
try:
dt = self._parse_and_convert_timestamp(timestamp_str)
return dt.isoformat().replace('+00:00', 'Z')
return dt.isoformat().replace("+00:00", "Z")
except Exception:
# Fallback to current time in UTC
return datetime.now(timezone.utc).isoformat().replace('+00:00', 'Z')
return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
+86 -63
View File
@@ -41,8 +41,8 @@ class CBFTransformer:
# Filter out records with zero successful_requests first
original_count = len(data)
if 'successful_requests' in data.columns:
filtered_data = data.filter(pl.col('successful_requests') > 0)
if "successful_requests" in data.columns:
filtered_data = data.filter(pl.col("successful_requests") > 0)
zero_requests_dropped = original_count - len(filtered_data)
else:
filtered_data = data
@@ -64,16 +64,23 @@ class CBFTransformer:
# Print summary of dropped records if any
from rich.console import Console
console = Console()
if zero_requests_dropped > 0:
console.print(f"[yellow]⚠️ Dropped {zero_requests_dropped:,} of {original_count:,} records with zero successful_requests[/yellow]")
console.print(
f"[yellow]⚠️ Dropped {zero_requests_dropped:,} of {original_count:,} records with zero successful_requests[/yellow]"
)
if czrn_dropped_count > 0:
console.print(f"[yellow]⚠️ Dropped {czrn_dropped_count:,} of {filtered_count:,} filtered records due to invalid CZRNs[/yellow]")
console.print(
f"[yellow]⚠️ Dropped {czrn_dropped_count:,} of {filtered_count:,} filtered records due to invalid CZRNs[/yellow]"
)
if len(cbf_data) > 0:
console.print(f"[green]✓ Successfully transformed {len(cbf_data):,} records[/green]")
console.print(
f"[green]✓ Successfully transformed {len(cbf_data):,} records[/green]"
)
return pl.DataFrame(cbf_data)
@@ -81,99 +88,116 @@ class CBFTransformer:
"""Create a single CBF record from LiteLLM daily spend row."""
# Parse date (daily spend tables use date strings like '2025-04-19')
usage_date = self._parse_date(row.get('date'))
usage_date = self._parse_date(row.get("date"))
# Calculate total tokens
prompt_tokens = int(row.get('prompt_tokens', 0))
completion_tokens = int(row.get('completion_tokens', 0))
prompt_tokens = int(row.get("prompt_tokens", 0))
completion_tokens = int(row.get("completion_tokens", 0))
total_tokens = prompt_tokens + completion_tokens
# Create CloudZero Resource Name (CZRN) as resource_id
resource_id = self.czrn_generator.create_from_litellm_data(row)
# Build dimensions for CloudZero
model = str(row.get('model', ''))
api_key_hash = str(row.get('api_key', ''))[:8] # First 8 chars for identification
model = str(row.get("model", ""))
api_key_hash = str(row.get("api_key", ""))[
:8
] # First 8 chars for identification
# Handle team information with fallbacks
team_id = row.get('team_id')
team_alias = row.get('team_alias')
user_email = row.get('user_email')
team_id = row.get("team_id")
team_alias = row.get("team_alias")
user_email = row.get("user_email")
# Use team_alias if available, otherwise team_id, otherwise fallback to 'unknown'
entity_id = str(team_alias) if team_alias else (str(team_id) if team_id else 'unknown')
entity_id = (
str(team_alias) if team_alias else (str(team_id) if team_id else "unknown")
)
# Get alias fields if they exist
api_key_alias = row.get('api_key_alias')
organization_alias = row.get('organization_alias')
project_alias = row.get('project_alias')
user_alias = row.get('user_alias')
api_key_alias = row.get("api_key_alias")
organization_alias = row.get("organization_alias")
project_alias = row.get("project_alias")
user_alias = row.get("user_alias")
dimensions = {
'entity_type': CZEntityType.TEAM.value,
'entity_id': entity_id,
'team_alias': str(team_alias) if team_alias else 'unknown',
'model': model,
'model_group': str(row.get('model_group', '')),
'provider': str(row.get('custom_llm_provider', '')),
'api_key_prefix': api_key_hash,
'api_key_alias': str(row.get('api_key_alias', '')),
'user_email': str(user_email) if user_email else '',
'api_requests': str(row.get('api_requests', 0)),
'successful_requests': str(row.get('successful_requests', 0)),
'failed_requests': str(row.get('failed_requests', 0)),
'cache_creation_tokens': str(row.get('cache_creation_input_tokens', 0)),
'cache_read_tokens': str(row.get('cache_read_input_tokens', 0)),
'organization_alias': str(organization_alias) if organization_alias else '',
'project_alias': str(project_alias) if project_alias else '',
'user_alias': str(user_alias) if user_alias else '',
"entity_type": CZEntityType.TEAM.value,
"entity_id": entity_id,
"team_alias": str(team_alias) if team_alias else "unknown",
"model": model,
"model_group": str(row.get("model_group", "")),
"provider": str(row.get("custom_llm_provider", "")),
"api_key_prefix": api_key_hash,
"api_key_alias": str(row.get("api_key_alias", "")),
"user_email": str(user_email) if user_email else "",
"api_requests": str(row.get("api_requests", 0)),
"successful_requests": str(row.get("successful_requests", 0)),
"failed_requests": str(row.get("failed_requests", 0)),
"cache_creation_tokens": str(row.get("cache_creation_input_tokens", 0)),
"cache_read_tokens": str(row.get("cache_read_input_tokens", 0)),
"organization_alias": str(organization_alias) if organization_alias else "",
"project_alias": str(project_alias) if project_alias else "",
"user_alias": str(user_alias) if user_alias else "",
}
# Extract CZRN components to populate corresponding CBF columns
czrn_components = self.czrn_generator.extract_components(resource_id)
service_type, provider, region, owner_account_id, resource_type, cloud_local_id = czrn_components
(
service_type,
provider,
region,
owner_account_id,
resource_type,
cloud_local_id,
) = czrn_components
# Build resource/account as concat of api_key_alias and api_key_prefix
resource_account = f"{api_key_alias}|{api_key_hash}" if api_key_alias else api_key_hash
resource_account = (
f"{api_key_alias}|{api_key_hash}" if api_key_alias else api_key_hash
)
# CloudZero CBF format with proper column names
cbf_record = {
# Required CBF fields
'time/usage_start': usage_date.isoformat() if usage_date else None, # Required: ISO-formatted UTC datetime
'cost/cost': float(row.get('spend', 0.0)), # Required: billed cost
'resource/id': resource_id, # CZRN (CloudZero Resource Name)
"time/usage_start": usage_date.isoformat()
if usage_date
else None, # Required: ISO-formatted UTC datetime
"cost/cost": float(row.get("spend", 0.0)), # Required: billed cost
"resource/id": resource_id, # CZRN (CloudZero Resource Name)
# Usage metrics for token consumption
'usage/amount': total_tokens, # Numeric value of tokens consumed
'usage/units': 'tokens', # Description of token units
"usage/amount": total_tokens, # Numeric value of tokens consumed
"usage/units": "tokens", # Description of token units
# CBF fields - updated per LIT-1907
'resource/service': str(row.get('model_group', '')), # Send model_group
'resource/account': resource_account, # Send api_key_alias|api_key_prefix
'resource/region': region, # Maps to CZRN region (cross-region)
'resource/usage_family': str(row.get('custom_llm_provider', '')), # Send provider
"resource/service": str(row.get("model_group", "")), # Send model_group
"resource/account": resource_account, # Send api_key_alias|api_key_prefix
"resource/region": region, # Maps to CZRN region (cross-region)
"resource/usage_family": str(
row.get("custom_llm_provider", "")
), # Send provider
# Action field
'action/operation': str(team_id) if team_id else '', # Send team_id
"action/operation": str(team_id) if team_id else "", # Send team_id
# Line item details
'lineitem/type': 'Usage', # Standard usage line item
"lineitem/type": "Usage", # Standard usage line item
}
# Add CZRN components that don't have direct CBF column mappings as resource tags
cbf_record['resource/tag:provider'] = provider # CZRN provider component
cbf_record['resource/tag:model'] = cloud_local_id # CZRN cloud-local-id component (model)
cbf_record["resource/tag:provider"] = provider # CZRN provider component
cbf_record[
"resource/tag:model"
] = cloud_local_id # CZRN cloud-local-id component (model)
# Add resource tags for all dimensions (using resource/tag:<key> format)
for key, value in dimensions.items():
if value and value != 'N/A' and value != 'unknown': # Only add meaningful tags
cbf_record[f'resource/tag:{key}'] = str(value)
if (
value and value != "N/A" and value != "unknown"
): # Only add meaningful tags
cbf_record[f"resource/tag:{key}"] = str(value)
# Add token breakdown as resource tags for analysis (excluding total_tokens per LIT-1907)
if prompt_tokens > 0:
cbf_record['resource/tag:prompt_tokens'] = str(prompt_tokens)
cbf_record["resource/tag:prompt_tokens"] = str(prompt_tokens)
if completion_tokens > 0:
cbf_record['resource/tag:completion_tokens'] = str(completion_tokens)
cbf_record["resource/tag:completion_tokens"] = str(completion_tokens)
return CBFRecord(cbf_record)
@@ -197,4 +221,3 @@ class CBFTransformer:
return None
return None
+4 -4
View File
@@ -670,7 +670,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
return final_response
"""
pass
async def async_should_run_chat_completion_agentic_loop(
self,
response: Any,
@@ -869,9 +869,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
model_response_dict = model_response.model_dump()
standard_logging_object_copy["response"] = model_response_dict
model_call_details_copy["standard_logging_object"] = (
standard_logging_object_copy
)
model_call_details_copy[
"standard_logging_object"
] = standard_logging_object_copy
return model_call_details_copy
async def get_proxy_server_request_from_cold_storage_with_object_key(
@@ -100,9 +100,7 @@ class CustomSecretManager(BaseSecretManager):
"""
super().__init__()
self.secret_manager_name = secret_manager_name or "custom_secret_manager"
verbose_logger.info(
"Initialized custom secret manager"
)
verbose_logger.info("Initialized custom secret manager")
@abstractmethod
async def async_read_secret(
+2 -1
View File
@@ -13,6 +13,7 @@ class CustomSSOLoginHandler(CustomLogger):
Useful when you have an OAuth proxy in front of LiteLLM
and you want to use the headers from the proxy to sign in the user
"""
async def handle_custom_ui_sso_sign_in(
self,
request: Request,
@@ -26,4 +27,4 @@ class CustomSSOLoginHandler(CustomLogger):
display_name="Test",
picture="https://test.com/test.png",
provider="test",
)
)
@@ -48,13 +48,15 @@ class DataDogLLMObsLogger(CustomBatchLogger):
def __init__(self, **kwargs):
try:
verbose_logger.debug("DataDogLLMObs: Initializing logger")
self.is_mock_mode = should_use_datadog_mock()
if self.is_mock_mode:
create_mock_datadog_client()
verbose_logger.debug("[DATADOG MOCK] DataDogLLMObs logger initialized in mock mode")
verbose_logger.debug(
"[DATADOG MOCK] DataDogLLMObs logger initialized in mock mode"
)
# Configure DataDog endpoint (Agent or Direct API)
# Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST
# Check for agent mode FIRST - agent mode doesn't require DD_API_KEY or DD_SITE
@@ -189,9 +191,11 @@ class DataDogLLMObsLogger(CustomBatchLogger):
verbose_logger.debug(
f"DataDogLLMObs: Flushing {len(self.log_queue)} events"
)
if self.is_mock_mode:
verbose_logger.debug("[DATADOG MOCK] Mock mode enabled - API calls will be intercepted")
verbose_logger.debug(
"[DATADOG MOCK] Mock mode enabled - API calls will be intercepted"
)
# Prepare the payload
payload = {
@@ -8,7 +8,10 @@ Usage:
Set DATADOG_MOCK=true in environment variables or config to enable mock mode.
"""
from litellm.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory
from litellm.integrations.mock_client_factory import (
MockClientConfig,
create_mock_client_factory,
)
# Create mock client using factory
_config = MockClientConfig(
@@ -25,4 +28,6 @@ _config = MockClientConfig(
patch_sync_client=True,
)
create_mock_datadog_client, should_use_datadog_mock = create_mock_client_factory(_config)
create_mock_datadog_client, should_use_datadog_mock = create_mock_client_factory(
_config
)
+5 -7
View File
@@ -25,6 +25,7 @@ def set_global_prompt_directory(directory: str) -> None:
litellm.global_prompt_directory = directory # type: ignore
def _get_prompt_data_from_dotprompt_content(dotprompt_content: str) -> dict:
"""
Get the prompt data from the dotprompt content.
@@ -36,12 +37,10 @@ def _get_prompt_data_from_dotprompt_content(dotprompt_content: str) -> dict:
# Parse the dotprompt content to extract frontmatter and content
temp_manager = PromptManager()
metadata, content = temp_manager._parse_frontmatter(dotprompt_content)
# Convert to prompt_data format
return {
"content": content.strip(),
"metadata": metadata
}
return {"content": content.strip(), "metadata": metadata}
def prompt_initializer(
litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec"
@@ -58,7 +57,7 @@ def prompt_initializer(
)
prompt_file = getattr(litellm_params, "prompt_file", None)
# Handle dotprompt_content from database
dotprompt_content = getattr(litellm_params, "dotprompt_content", None)
if dotprompt_content and not prompt_data and not prompt_file:
@@ -74,7 +73,6 @@ def prompt_initializer(
return dot_prompt_manager
except Exception as e:
raise e
@@ -128,7 +128,6 @@ class DotpromptManager(CustomPromptManagement):
raise ValueError("prompt_id is required for dotprompt manager")
try:
# Get the prompt template (versioned or base)
template = self.prompt_manager.get_prompt(
prompt_id=prompt_id, version=prompt_version
@@ -205,7 +204,6 @@ class DotpromptManager(CustomPromptManagement):
ignore_prompt_manager_model: Optional[bool] = False,
ignore_prompt_manager_optional_params: Optional[bool] = False,
) -> Tuple[str, List[AllMessageValues], dict]:
from litellm.integrations.prompt_management_base import PromptManagementBase
return PromptManagementBase.get_chat_completion_prompt(
@@ -205,7 +205,7 @@ class PromptManager:
"""
# Get the template (versioned or base)
template = self.get_prompt(prompt_id=prompt_id, version=version)
if template is None:
available_prompts = list(self.prompts.keys())
version_str = f" (version {version})" if version else ""
@@ -266,11 +266,11 @@ class PromptManager:
) -> Optional[PromptTemplate]:
"""
Get a prompt template by ID and optional version.
Args:
prompt_id: The base prompt ID
version: Optional version number. If provided, looks for {prompt_id}.v{version}
Returns:
The prompt template if found, None otherwise
"""
@@ -279,7 +279,7 @@ class PromptManager:
versioned_id = f"{prompt_id}.v{version}"
if versioned_id in self.prompts:
return self.prompts[versioned_id]
# Fall back to base prompt_id
return self.prompts.get(prompt_id)
@@ -222,4 +222,3 @@ response = client.chat.completions.create(<br>
</body>
</html>
"""
@@ -131,4 +131,4 @@ MAX_BUDGET_ALERT_EMAIL_TEMPLATE = """
Best, <br />
The LiteLLM team <br />
"""
"""
+4 -1
View File
@@ -56,7 +56,9 @@ class FocusLogger(CustomLogger):
self.interval_seconds = int(raw_interval) if raw_interval is not None else None
env_prefix = os.getenv("FOCUS_PREFIX")
self.prefix: str = (
prefix if prefix is not None else (env_prefix if env_prefix else "focus_exports")
prefix
if prefix is not None
else (env_prefix if env_prefix else "focus_exports")
)
self._destination_config = destination_config
@@ -208,4 +210,5 @@ class FocusLogger(CustomLogger):
frequency=self.frequency,
)
__all__ = ["FocusLogger"]
+46 -27
View File
@@ -34,7 +34,10 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
os.getenv("GCS_FLUSH_INTERVAL", GCS_DEFAULT_FLUSH_INTERVAL_SECONDS)
)
self.use_batched_logging = (
os.getenv("GCS_USE_BATCHED_LOGGING", str(GCS_DEFAULT_USE_BATCHED_LOGGING).lower()).lower() == "true"
os.getenv(
"GCS_USE_BATCHED_LOGGING", str(GCS_DEFAULT_USE_BATCHED_LOGGING).lower()
).lower()
== "true"
)
self.flush_lock = asyncio.Lock()
super().__init__(
@@ -112,9 +115,9 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
def _drain_queue_batch(self) -> List[GCSLogQueueItem]:
"""
Drain items from the queue (non-blocking), respecting batch_size limit.
This prevents unbounded queue growth when processing is slower than log accumulation.
Returns:
List of items to process, up to batch_size items
"""
@@ -137,33 +140,45 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
"""
Extract a synchronous grouping key from kwargs to group items by GCS config.
This allows us to batch items with the same bucket/credentials together.
Returns a string key that uniquely identifies the GCS config combination.
This key may contain sensitive information (bucket names, paths) - use _sanitize_config_key()
for logging purposes.
"""
standard_callback_dynamic_params = kwargs.get("standard_callback_dynamic_params", None) or {}
bucket_name = standard_callback_dynamic_params.get("gcs_bucket_name", None) or self.BUCKET_NAME or "default"
path_service_account = standard_callback_dynamic_params.get("gcs_path_service_account", None) or self.path_service_account_json or "default"
standard_callback_dynamic_params = (
kwargs.get("standard_callback_dynamic_params", None) or {}
)
bucket_name = (
standard_callback_dynamic_params.get("gcs_bucket_name", None)
or self.BUCKET_NAME
or "default"
)
path_service_account = (
standard_callback_dynamic_params.get("gcs_path_service_account", None)
or self.path_service_account_json
or "default"
)
return f"{bucket_name}|{path_service_account}"
def _sanitize_config_key(self, config_key: str) -> str:
"""
Create a sanitized version of the config key for logging.
Uses a hash to avoid exposing sensitive bucket names or service account paths.
Returns a short hash prefix for safe logging.
"""
hash_obj = hashlib.sha256(config_key.encode('utf-8'))
hash_obj = hashlib.sha256(config_key.encode("utf-8"))
return f"config-{hash_obj.hexdigest()[:8]}"
def _group_items_by_config(self, items: List[GCSLogQueueItem]) -> Dict[str, List[GCSLogQueueItem]]:
def _group_items_by_config(
self, items: List[GCSLogQueueItem]
) -> Dict[str, List[GCSLogQueueItem]]:
"""
Group items by their GCS config (bucket + credentials).
This ensures items with different configs are processed separately.
Returns a dict mapping config_key -> list of items with that config.
"""
grouped: Dict[str, List[GCSLogQueueItem]] = {}
@@ -186,18 +201,20 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
lines.append(json_line)
return "\n".join(lines)
async def _send_grouped_batch(self, items: List[GCSLogQueueItem], config_key: str) -> Tuple[int, int]:
async def _send_grouped_batch(
self, items: List[GCSLogQueueItem], config_key: str
) -> Tuple[int, int]:
"""
Send a batch of items that share the same GCS config.
Returns:
(success_count, error_count)
"""
if not items:
return (0, 0)
first_kwargs = items[0]["kwargs"]
try:
gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config(
first_kwargs
@@ -208,23 +225,25 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
service_account_json=gcs_logging_config["path_service_account"],
)
bucket_name = gcs_logging_config["bucket_name"]
current_date = self._get_object_date_from_datetime(datetime.now(timezone.utc))
current_date = self._get_object_date_from_datetime(
datetime.now(timezone.utc)
)
batch_id = f"{int(time.time() * 1000)}-{uuid.uuid4().hex[:8]}"
object_name = self._generate_batch_object_name(current_date, batch_id)
combined_payload = self._combine_payloads_to_ndjson(items)
await self._log_json_data_on_gcs(
headers=headers,
bucket_name=bucket_name,
object_name=object_name,
logging_payload=combined_payload,
)
success_count = len(items)
error_count = 0
return (success_count, error_count)
except Exception as e:
success_count = 0
error_count = len(items)
@@ -255,13 +274,13 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
service_account_json=gcs_logging_config["path_service_account"],
)
bucket_name = gcs_logging_config["bucket_name"]
object_name = self._get_object_name(
kwargs=item["kwargs"],
logging_payload=item["payload"],
response_obj=item["response_obj"],
)
await self._log_json_data_on_gcs(
headers=headers,
bucket_name=bucket_name,
@@ -289,7 +308,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils):
if self.use_batched_logging:
grouped_items = self._group_items_by_config(items_to_process)
for config_key, group_items in grouped_items.items():
await self._send_grouped_batch(group_items, config_key)
else:
@@ -7,7 +7,7 @@ from litellm.integrations.gcs_bucket.gcs_bucket_mock_client import (
create_mock_gcs_client,
mock_vertex_auth_methods,
)
from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
@@ -28,11 +28,11 @@ IAM_AUTH_KEY = "IAM_AUTH"
class GCSBucketBase(CustomBatchLogger):
def __init__(self, bucket_name: Optional[str] = None, **kwargs) -> None:
self.is_mock_mode = should_use_gcs_mock()
if self.is_mock_mode:
mock_vertex_auth_methods()
create_mock_gcs_client()
self.async_httpx_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)
@@ -85,10 +85,10 @@ class GCSBucketBase(CustomBatchLogger):
from litellm import vertex_chat_completion
# Get project_id from environment if available, otherwise None
# This helps support use of this library to auth to pull secrets
# This helps support use of this library to auth to pull secrets
# from Secret Manager.
project_id = os.getenv("GOOGLE_SECRET_MANAGER_PROJECT_ID")
_auth_header, vertex_project = vertex_chat_completion._ensure_access_token(
credentials=self.path_service_account_json,
project_id=project_id,
@@ -11,7 +11,11 @@ Usage:
import asyncio
from litellm._logging import verbose_logger
from litellm.integrations.mock_client_factory import MockClientConfig, create_mock_client_factory, MockResponse
from litellm.integrations.mock_client_factory import (
MockClientConfig,
create_mock_client_factory,
MockResponse,
)
# Use factory for POST handler
_config = MockClientConfig(
@@ -34,10 +38,14 @@ _mocks_initialized = False
# Default mock latency in seconds (simulates network round-trip)
# Typical GCS API calls take 100-300ms for uploads, 50-150ms for GET/DELETE
_MOCK_LATENCY_SECONDS = float(__import__("os").getenv("GCS_MOCK_LATENCY_MS", "150")) / 1000.0
_MOCK_LATENCY_SECONDS = (
float(__import__("os").getenv("GCS_MOCK_LATENCY_MS", "150")) / 1000.0
)
async def _mock_async_handler_get(self, url, params=None, headers=None, follow_redirects=None):
async def _mock_async_handler_get(
self, url, params=None, headers=None, follow_redirects=None
):
"""Monkey-patched AsyncHTTPHandler.get that intercepts GCS calls."""
# Only mock GCS API calls
if isinstance(url, str) and "storage.googleapis.com" in url:
@@ -86,14 +94,30 @@ async def _mock_async_handler_get(self, url, params=None, headers=None, follow_r
status_code=200,
json_data=mock_payload,
url=url,
elapsed_seconds=_MOCK_LATENCY_SECONDS
elapsed_seconds=_MOCK_LATENCY_SECONDS,
)
if _original_async_handler_get is not None:
return await _original_async_handler_get(self, url=url, params=params, headers=headers, follow_redirects=follow_redirects)
return await _original_async_handler_get(
self,
url=url,
params=params,
headers=headers,
follow_redirects=follow_redirects,
)
raise RuntimeError("Original AsyncHTTPHandler.get not available")
async def _mock_async_handler_delete(self, url, data=None, json=None, params=None, headers=None, timeout=None, stream=False, content=None):
async def _mock_async_handler_delete(
self,
url,
data=None,
json=None,
params=None,
headers=None,
timeout=None,
stream=False,
content=None,
):
"""Monkey-patched AsyncHTTPHandler.delete that intercepts GCS calls."""
# Only mock GCS API calls
if isinstance(url, str) and "storage.googleapis.com" in url:
@@ -104,49 +128,61 @@ async def _mock_async_handler_delete(self, url, data=None, json=None, params=Non
status_code=204,
json_data=None, # Empty body for DELETE
url=url,
elapsed_seconds=_MOCK_LATENCY_SECONDS
elapsed_seconds=_MOCK_LATENCY_SECONDS,
)
if _original_async_handler_delete is not None:
return await _original_async_handler_delete(self, url=url, data=data, json=json, params=params, headers=headers, timeout=timeout, stream=stream, content=content)
return await _original_async_handler_delete(
self,
url=url,
data=data,
json=json,
params=params,
headers=headers,
timeout=timeout,
stream=stream,
content=content,
)
raise RuntimeError("Original AsyncHTTPHandler.delete not available")
def create_mock_gcs_client():
"""
Monkey-patch AsyncHTTPHandler methods to intercept GCS calls.
AsyncHTTPHandler is used by LiteLLM's get_async_httpx_client() which is what
GCSBucketBase uses for making API calls.
This function is idempotent - it only initializes mocks once, even if called multiple times.
"""
global _original_async_handler_get, _original_async_handler_delete, _mocks_initialized
# Use factory for POST handler
_create_mock_gcs_post()
# If already initialized, skip GET/DELETE patching
if _mocks_initialized:
return
verbose_logger.debug("[GCS MOCK] Initializing GCS GET/DELETE handlers...")
# Patch GET and DELETE handlers (GCS-specific)
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
if _original_async_handler_get is None:
_original_async_handler_get = AsyncHTTPHandler.get
AsyncHTTPHandler.get = _mock_async_handler_get # type: ignore
verbose_logger.debug("[GCS MOCK] Patched AsyncHTTPHandler.get")
if _original_async_handler_delete is None:
_original_async_handler_delete = AsyncHTTPHandler.delete
AsyncHTTPHandler.delete = _mock_async_handler_delete # type: ignore
verbose_logger.debug("[GCS MOCK] Patched AsyncHTTPHandler.delete")
verbose_logger.debug(f"[GCS MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms")
verbose_logger.debug(
f"[GCS MOCK] Mock latency set to {_MOCK_LATENCY_SECONDS*1000:.0f}ms"
)
verbose_logger.debug("[GCS MOCK] GCS mock client initialization complete")
_mocks_initialized = True
@@ -154,38 +190,64 @@ def mock_vertex_auth_methods():
"""
Monkey-patch Vertex AI auth methods to return fake tokens.
This prevents auth failures when GCS_MOCK is enabled.
This function is idempotent - it only patches once, even if called multiple times.
"""
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
# Store original methods if not already stored
if not hasattr(VertexBase, '_original_ensure_access_token_async'):
setattr(VertexBase, '_original_ensure_access_token_async', VertexBase._ensure_access_token_async)
setattr(VertexBase, '_original_ensure_access_token', VertexBase._ensure_access_token)
setattr(VertexBase, '_original_get_token_and_url', VertexBase._get_token_and_url)
async def _mock_ensure_access_token_async(self, credentials, project_id, custom_llm_provider):
if not hasattr(VertexBase, "_original_ensure_access_token_async"):
setattr(
VertexBase,
"_original_ensure_access_token_async",
VertexBase._ensure_access_token_async,
)
setattr(
VertexBase, "_original_ensure_access_token", VertexBase._ensure_access_token
)
setattr(
VertexBase, "_original_get_token_and_url", VertexBase._get_token_and_url
)
async def _mock_ensure_access_token_async(
self, credentials, project_id, custom_llm_provider
):
"""Mock async auth method - returns fake token."""
verbose_logger.debug("[GCS MOCK] Vertex AI auth: _ensure_access_token_async called")
verbose_logger.debug(
"[GCS MOCK] Vertex AI auth: _ensure_access_token_async called"
)
return ("mock-gcs-token", "mock-project-id")
def _mock_ensure_access_token(self, credentials, project_id, custom_llm_provider):
def _mock_ensure_access_token(
self, credentials, project_id, custom_llm_provider
):
"""Mock sync auth method - returns fake token."""
verbose_logger.debug("[GCS MOCK] Vertex AI auth: _ensure_access_token called")
verbose_logger.debug(
"[GCS MOCK] Vertex AI auth: _ensure_access_token called"
)
return ("mock-gcs-token", "mock-project-id")
def _mock_get_token_and_url(self, model, auth_header, vertex_credentials, vertex_project,
vertex_location, gemini_api_key, stream, custom_llm_provider, api_base):
def _mock_get_token_and_url(
self,
model,
auth_header,
vertex_credentials,
vertex_project,
vertex_location,
gemini_api_key,
stream,
custom_llm_provider,
api_base,
):
"""Mock get_token_and_url - returns fake token."""
verbose_logger.debug("[GCS MOCK] Vertex AI auth: _get_token_and_url called")
return ("mock-gcs-token", "https://storage.googleapis.com")
# Patch the methods
VertexBase._ensure_access_token_async = _mock_ensure_access_token_async # type: ignore
VertexBase._ensure_access_token = _mock_ensure_access_token # type: ignore
VertexBase._get_token_and_url = _mock_get_token_and_url # type: ignore
verbose_logger.debug("[GCS MOCK] Patched Vertex AI auth methods")
@@ -164,7 +164,11 @@ class GenericAPILogger(CustomBatchLogger):
self.callback_name: Optional[str] = callback_name
# Validate and store log_format
if log_format is not None and log_format not in ["json_array", "ndjson", "single"]:
if log_format is not None and log_format not in [
"json_array",
"ndjson",
"single",
]:
raise ValueError(
f"Invalid log_format: {log_format}. Must be one of: 'json_array', 'ndjson', 'single'"
)
@@ -120,7 +120,6 @@ class GenericPromptManager(CustomPromptManagement):
http_client = _get_httpx_client()
try:
response = http_client.get(
url,
params=params,
@@ -325,7 +324,6 @@ class GenericPromptManager(CustomPromptManagement):
prompt_label: Optional[str] = None,
prompt_version: Optional[int] = None,
) -> PromptManagementClient:
# Check cache first
cached_prompt = self._common_caching_logic(
prompt_id=prompt_id,
+4 -6
View File
@@ -39,11 +39,8 @@ def prompt_initializer(
gitlab_config = getattr(litellm_params, "gitlab_config", None)
prompt_id = getattr(litellm_params, "prompt_id", None)
if not gitlab_config:
raise ValueError(
"gitlab_config is required for gitlab prompt integration"
)
raise ValueError("gitlab_config is required for gitlab prompt integration")
try:
gitlab_prompt_manager = GitLabPromptManager(
@@ -55,9 +52,10 @@ def prompt_initializer(
except Exception as e:
raise e
def _gitlab_prompt_initializer(
litellm_params: PromptLiteLLMParams,
prompt: PromptSpec,
litellm_params: PromptLiteLLMParams,
prompt: PromptSpec,
) -> CustomPromptManagement:
"""
Build a GitLab-backed prompt manager for this prompt.
+40 -16
View File
@@ -45,7 +45,7 @@ class GitLabClient:
self.auth_method = config.get("auth_method", "token") # 'token' or 'oauth'
self.branch = config.get("branch", None)
if not self.branch:
self.branch = 'main'
self.branch = "main"
self.tag = config.get("tag")
self.base_url = config.get("base_url", "https://gitlab.com/api/v4")
@@ -86,7 +86,13 @@ class GitLabClient:
ref_q = quote(ref or self.ref, safe="")
return f"{self.base_url}/projects/{self._project_enc}/repository/files/{file_enc}?ref={ref_q}"
def _tree_url(self, directory_path: str = "", recursive: bool = False, *, ref: Optional[str] = None) -> str:
def _tree_url(
self,
directory_path: str = "",
recursive: bool = False,
*,
ref: Optional[str] = None,
) -> str:
path_q = f"&path={quote(directory_path, safe='')}" if directory_path else ""
rec_q = "&recursive=true" if recursive else ""
ref_q = quote(ref or self.ref, safe="")
@@ -102,7 +108,9 @@ class GitLabClient:
raise ValueError("ref must be a non-empty string")
self.ref = ref
def get_file_content(self, file_path: str, *, ref: Optional[str] = None) -> Optional[str]:
def get_file_content(
self, file_path: str, *, ref: Optional[str] = None
) -> Optional[str]:
"""
Fetch the content of a file from the GitLab repository at the given ref
(tag, branch, or commit SHA). If `ref` is None, uses self.ref.
@@ -124,7 +132,11 @@ class GitLabClient:
resp.raise_for_status()
ctype = (resp.headers.get("content-type") or "").lower()
if ctype.startswith("text/") or "charset=" in ctype or ctype.startswith("application/json"):
if (
ctype.startswith("text/")
or "charset=" in ctype
or ctype.startswith("application/json")
):
return resp.text
try:
return resp.content.decode("utf-8")
@@ -140,10 +152,14 @@ class GitLabClient:
f"Access denied to file '{file_path}'. Check your GitLab permissions for project '{self.project}'."
)
if status == 401:
raise Exception("Authentication failed. Check your GitLab token and auth_method.")
raise Exception(
"Authentication failed. Check your GitLab token and auth_method."
)
raise Exception(f"Failed to fetch file '{file_path}': {e}")
def _get_file_content_via_json(self, file_path: str, *, ref: Optional[str] = None) -> Optional[str]:
def _get_file_content_via_json(
self, file_path: str, *, ref: Optional[str] = None
) -> Optional[str]:
"""
Fallback for get_file_content(): use the JSON file API which returns base64 content.
"""
@@ -171,16 +187,20 @@ class GitLabClient:
f"Access denied to file '{file_path}'. Check your GitLab permissions for project '{self.project}'."
)
if status == 401:
raise Exception("Authentication failed. Check your GitLab token and auth_method.")
raise Exception(f"Failed to fetch file '{file_path}' via JSON endpoint: {e}")
raise Exception(
"Authentication failed. Check your GitLab token and auth_method."
)
raise Exception(
f"Failed to fetch file '{file_path}' via JSON endpoint: {e}"
)
def list_files(
self,
directory_path: str = "",
file_extension: str = ".prompt",
recursive: bool = False,
*,
ref: Optional[str] = None,
self,
directory_path: str = "",
file_extension: str = ".prompt",
recursive: bool = False,
*,
ref: Optional[str] = None,
) -> List[str]:
"""
List files in a directory with a specific extension using the repository tree API.
@@ -220,7 +240,9 @@ class GitLabClient:
f"Access denied to directory '{directory_path}'. Check your GitLab permissions for project '{self.project}'."
)
if status == 401:
raise Exception("Authentication failed. Check your GitLab token and auth_method.")
raise Exception(
"Authentication failed. Check your GitLab token and auth_method."
)
raise Exception(f"Failed to list files in '{directory_path}': {e}")
def get_repository_info(self) -> Dict[str, Any]:
@@ -252,7 +274,9 @@ class GitLabClient:
except Exception as e:
raise Exception(f"Failed to get branches: {e}")
def get_file_metadata(self, file_path: str, *, ref: Optional[str] = None) -> Optional[Dict[str, Any]]:
def get_file_metadata(
self, file_path: str, *, ref: Optional[str] = None
) -> Optional[Dict[str, Any]]:
"""
Get minimal metadata about a file via RAW endpoint headers at a given ref.

Some files were not shown because too many files have changed in this diff Show More