mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-05 10:24:03 +00:00
Merge pull request #23440 from BerriAI/litellm_oss_staging_03_11_2026
Litellm oss staging 03 11 2026
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -31,6 +31,8 @@ serviceAccount:
|
||||
# annotations for litellm deployment
|
||||
deploymentAnnotations: {}
|
||||
deploymentLabels: {}
|
||||
deploymentMinReadySeconds: 0
|
||||
|
||||
# annotations for litellm pods
|
||||
podAnnotations: {}
|
||||
podLabels: {}
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
:::
|
||||
|
||||
|
||||
Generated
+4
-4
@@ -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"
|
||||
|
||||
@@ -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
File diff suppressed because it is too large
Load Diff
+51
-41
@@ -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
|
||||
|
||||
@@ -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
@@ -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}")
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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"]
|
||||
|
||||
|
||||
@@ -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}"
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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),
|
||||
}
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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 = [
|
||||
|
||||
@@ -42,4 +42,3 @@ __all__ = [
|
||||
"retrieve_container_file",
|
||||
"retrieve_container_file_content",
|
||||
]
|
||||
|
||||
|
||||
@@ -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
@@ -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(
|
||||
|
||||
@@ -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
@@ -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
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
@@ -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,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
@@ -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,
|
||||
|
||||
@@ -13,7 +13,7 @@ from .main import (
|
||||
|
||||
__all__ = [
|
||||
"generate_content",
|
||||
"agenerate_content",
|
||||
"agenerate_content",
|
||||
"generate_content_stream",
|
||||
"agenerate_content_stream",
|
||||
]
|
||||
]
|
||||
|
||||
@@ -13,7 +13,7 @@ from .handler import GenerateContentToCompletionHandler
|
||||
from .transformation import GoogleGenAIAdapter, GoogleGenAIStreamWrapper
|
||||
|
||||
__all__ = [
|
||||
"GoogleGenAIAdapter",
|
||||
"GoogleGenAIAdapter",
|
||||
"GoogleGenAIStreamWrapper",
|
||||
"GenerateContentToCompletionHandler"
|
||||
]
|
||||
"GenerateContentToCompletionHandler",
|
||||
]
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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,3 +1,3 @@
|
||||
from .agentops import AgentOps
|
||||
|
||||
__all__ = ["AgentOps"]
|
||||
__all__ = ["AgentOps"]
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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 />
|
||||
"""
|
||||
"""
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
Reference in New Issue
Block a user