mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-05 00:25:35 +00:00
Arize Phoenix OSS - Prompt Management Integration (#17750)
* docs(prompt_management.md): document how to onboard prompts to litellm * feat(arize_phoenix_prompt_manager.py): support new prompt management integration allows users to connect arize phoenix prompt manager to litellm * fix(proxy/utils.py): remove prompt variables to avoid re-processing prompt * docs(arize_phoenix_prompts.md): document new prompt management integration
This commit is contained in:
Krish Dholakia
@@ -0,0 +1,134 @@
|
||||
# Arize Phoenix Prompt Management
|
||||
|
||||
Use prompt versions from [Arize Phoenix](https://phoenix.arize.com/) with LiteLLM SDK and Proxy.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### SDK
|
||||
|
||||
```python
|
||||
import litellm
|
||||
|
||||
response = litellm.completion(
|
||||
model="gpt-4o",
|
||||
prompt_id="UHJvbXB0VmVyc2lvbjox",
|
||||
prompt_integration="arize_phoenix",
|
||||
api_key="your-arize-phoenix-token",
|
||||
api_base="https://app.phoenix.arize.com/s/your-workspace",
|
||||
prompt_variables={"question": "What is AI?"},
|
||||
)
|
||||
```
|
||||
|
||||
### Proxy
|
||||
|
||||
**1. Add prompt to config**
|
||||
|
||||
```yaml
|
||||
prompts:
|
||||
- prompt_id: "simple_prompt"
|
||||
litellm_params:
|
||||
prompt_id: "UHJvbXB0VmVyc2lvbjox"
|
||||
prompt_integration: "arize_phoenix"
|
||||
api_base: https://app.phoenix.arize.com/s/your-workspace
|
||||
api_key: os.environ/PHOENIX_API_KEY
|
||||
ignore_prompt_manager_model: true # optional: use model from config instead
|
||||
ignore_prompt_manager_optional_params: true # optional: ignore temp, max_tokens from prompt
|
||||
```
|
||||
|
||||
**2. Make request**
|
||||
|
||||
```bash
|
||||
curl -X POST 'http://0.0.0.0:4000/chat/completions' \
|
||||
-H 'Content-Type: application/json' \
|
||||
-H 'Authorization: Bearer sk-1234' \
|
||||
-d '{
|
||||
"model": "gpt-3.5-turbo",
|
||||
"prompt_id": "simple_prompt",
|
||||
"prompt_variables": {
|
||||
"question": "Explain quantum computing"
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
### Get Arize Phoenix Credentials
|
||||
|
||||
1. **API Token**: Get from [Arize Phoenix Settings](https://app.phoenix.arize.com/)
|
||||
2. **Workspace URL**: `https://app.phoenix.arize.com/s/{your-workspace}`
|
||||
3. **Prompt ID**: Found in prompt version URL
|
||||
|
||||
**Set environment variable**:
|
||||
```bash
|
||||
export PHOENIX_API_KEY="your-token"
|
||||
```
|
||||
|
||||
### SDK + PROXY Options
|
||||
|
||||
| Parameter | Required | Description |
|
||||
|-----------|----------|-------------|
|
||||
| `prompt_id` | Yes | Arize Phoenix prompt version ID |
|
||||
| `prompt_integration` | Yes | Set to `"arize_phoenix"` |
|
||||
| `api_base` | Yes | Workspace URL |
|
||||
| `api_key` | Yes | Access token |
|
||||
| `prompt_variables` | No | Variables for template |
|
||||
|
||||
### Proxy-only Options
|
||||
|
||||
| Parameter | Description |
|
||||
|-----------|-------------|
|
||||
| `ignore_prompt_manager_model` | Use config model instead of prompt's model |
|
||||
| `ignore_prompt_manager_optional_params` | Ignore temperature, max_tokens from prompt |
|
||||
|
||||
## Variable Templates
|
||||
|
||||
Arize Phoenix uses Mustache/Handlebars syntax:
|
||||
|
||||
```python
|
||||
# Template: "Hello {{name}}, question: {{question}}"
|
||||
prompt_variables = {
|
||||
"name": "Alice",
|
||||
"question": "What is ML?"
|
||||
}
|
||||
# Result: "Hello Alice, question: What is ML?"
|
||||
```
|
||||
|
||||
|
||||
## Combine with Additional Messages
|
||||
|
||||
```python
|
||||
response = litellm.completion(
|
||||
model="gpt-4o",
|
||||
prompt_id="UHJvbXB0VmVyc2lvbjox",
|
||||
prompt_integration="arize_phoenix",
|
||||
api_base="https://app.phoenix.arize.com/s/your-workspace",
|
||||
prompt_variables={"question": "Explain AI"},
|
||||
messages=[
|
||||
{"role": "user", "content": "Keep it under 50 words"}
|
||||
]
|
||||
)
|
||||
```
|
||||
|
||||
|
||||
## Error Handling
|
||||
|
||||
```python
|
||||
try:
|
||||
response = litellm.completion(
|
||||
model="gpt-4o",
|
||||
prompt_id="invalid-id",
|
||||
prompt_integration="arize_phoenix",
|
||||
api_base="https://app.phoenix.arize.com/s/workspace"
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
# 404: Prompt not found
|
||||
# 401: Invalid credentials
|
||||
# 403: Access denied
|
||||
```
|
||||
|
||||
## Support
|
||||
|
||||
- [LiteLLM GitHub Issues](https://github.com/BerriAI/litellm/issues)
|
||||
- [Arize Phoenix Docs](https://docs.arize.com/phoenix)
|
||||
|
||||
@@ -12,6 +12,292 @@ Run experiments or change the specific model (e.g. from gpt-4o to gpt4o-mini fin
|
||||
| Langfuse | [Get Started](https://langfuse.com/docs/prompts/get-started) |
|
||||
| Humanloop | [Get Started](../observability/humanloop) |
|
||||
|
||||
## Onboarding Prompts via config.yaml
|
||||
|
||||
You can onboard and initialize prompts directly in your `config.yaml` file. This allows you to:
|
||||
- Load prompts at proxy startup
|
||||
- Manage prompts as code alongside your proxy configuration
|
||||
- Use any supported prompt integration (dotprompt, Langfuse, BitBucket, GitLab, custom)
|
||||
|
||||
### Basic Structure
|
||||
|
||||
Add a `prompts` field to your config.yaml:
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gpt-4
|
||||
litellm_params:
|
||||
model: openai/gpt-4
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
prompts:
|
||||
- prompt_id: "my_prompt_id"
|
||||
litellm_params:
|
||||
prompt_id: "my_prompt_id"
|
||||
prompt_integration: "dotprompt" # or langfuse, bitbucket, gitlab, custom
|
||||
# integration-specific parameters below
|
||||
```
|
||||
|
||||
### Understanding `prompt_integration`
|
||||
|
||||
The `prompt_integration` field determines where and how prompts are loaded:
|
||||
|
||||
- **`dotprompt`**: Load from local `.prompt` files or inline content
|
||||
- **`langfuse`**: Fetch prompts from Langfuse prompt management
|
||||
- **`bitbucket`**: Load from BitBucket repository `.prompt` files (team-based access control)
|
||||
- **`gitlab`**: Load from GitLab repository `.prompt` files (team-based access control)
|
||||
- **`custom`**: Use your own custom prompt management implementation
|
||||
|
||||
Each integration has its own configuration parameters and access control mechanisms.
|
||||
|
||||
### Supported Integrations
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="dotprompt" label="DotPrompt (File-based)">
|
||||
|
||||
**Option 1: Using a prompt directory**
|
||||
|
||||
```yaml
|
||||
prompts:
|
||||
- prompt_id: "hello"
|
||||
litellm_params:
|
||||
prompt_id: "hello"
|
||||
prompt_integration: "dotprompt"
|
||||
prompt_directory: "./prompts" # Directory containing .prompt files
|
||||
|
||||
litellm_settings:
|
||||
global_prompt_directory: "./prompts" # Global setting for all dotprompt integrations
|
||||
```
|
||||
|
||||
**Option 2: Using inline prompt data**
|
||||
|
||||
```yaml
|
||||
prompts:
|
||||
- prompt_id: "my_inline_prompt"
|
||||
litellm_params:
|
||||
prompt_id: "my_inline_prompt"
|
||||
prompt_integration: "dotprompt"
|
||||
prompt_data:
|
||||
my_inline_prompt:
|
||||
content: "Hello {{name}}! How can I help you with {{topic}}?"
|
||||
metadata:
|
||||
model: "gpt-4"
|
||||
temperature: 0.7
|
||||
max_tokens: 150
|
||||
```
|
||||
|
||||
**Option 3: Using dotprompt_content for single prompts**
|
||||
|
||||
```yaml
|
||||
prompts:
|
||||
- prompt_id: "simple_prompt"
|
||||
litellm_params:
|
||||
prompt_id: "simple_prompt"
|
||||
prompt_integration: "dotprompt"
|
||||
dotprompt_content: |
|
||||
---
|
||||
model: gpt-4
|
||||
temperature: 0.7
|
||||
---
|
||||
System: You are a helpful assistant.
|
||||
|
||||
User: {{user_message}}
|
||||
```
|
||||
|
||||
Create `.prompt` files in your prompt directory:
|
||||
|
||||
```yaml
|
||||
# prompts/hello.prompt
|
||||
---
|
||||
model: gpt-4
|
||||
temperature: 0.7
|
||||
---
|
||||
System: You are a helpful assistant.
|
||||
|
||||
User: {{user_message}}
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="langfuse" label="Langfuse">
|
||||
|
||||
```yaml
|
||||
prompts:
|
||||
- prompt_id: "my_langfuse_prompt"
|
||||
litellm_params:
|
||||
prompt_id: "my_langfuse_prompt"
|
||||
prompt_integration: "langfuse"
|
||||
langfuse_public_key: "os.environ/LANGFUSE_PUBLIC_KEY"
|
||||
langfuse_secret_key: "os.environ/LANGFUSE_SECRET_KEY"
|
||||
langfuse_host: "https://cloud.langfuse.com" # optional
|
||||
|
||||
litellm_settings:
|
||||
langfuse_public_key: "os.environ/LANGFUSE_PUBLIC_KEY" # Global setting
|
||||
langfuse_secret_key: "os.environ/LANGFUSE_SECRET_KEY" # Global setting
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="bitbucket" label="BitBucket">
|
||||
|
||||
```yaml
|
||||
prompts:
|
||||
- prompt_id: "my_bitbucket_prompt"
|
||||
litellm_params:
|
||||
prompt_id: "my_bitbucket_prompt"
|
||||
prompt_integration: "bitbucket"
|
||||
bitbucket_workspace: "your-workspace"
|
||||
bitbucket_repository: "your-repo"
|
||||
bitbucket_access_token: "os.environ/BITBUCKET_ACCESS_TOKEN"
|
||||
bitbucket_branch: "main" # optional, defaults to main
|
||||
|
||||
litellm_settings:
|
||||
global_bitbucket_config:
|
||||
workspace: "your-workspace"
|
||||
repository: "your-repo"
|
||||
access_token: "os.environ/BITBUCKET_ACCESS_TOKEN"
|
||||
branch: "main"
|
||||
```
|
||||
|
||||
Your BitBucket repository should contain `.prompt` files:
|
||||
|
||||
```yaml
|
||||
# prompts/my_bitbucket_prompt.prompt
|
||||
---
|
||||
model: gpt-4
|
||||
temperature: 0.7
|
||||
---
|
||||
System: You are a helpful assistant.
|
||||
|
||||
User: {{user_message}}
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="gitlab" label="GitLab">
|
||||
|
||||
```yaml
|
||||
prompts:
|
||||
- prompt_id: "my_gitlab_prompt"
|
||||
litellm_params:
|
||||
prompt_id: "my_gitlab_prompt"
|
||||
prompt_integration: "gitlab"
|
||||
gitlab_project: "group/sub/repo"
|
||||
gitlab_access_token: "os.environ/GITLAB_ACCESS_TOKEN"
|
||||
gitlab_branch: "main" # optional
|
||||
gitlab_prompts_path: "prompts" # optional, defaults to root
|
||||
|
||||
litellm_settings:
|
||||
global_gitlab_config:
|
||||
project: "group/sub/repo"
|
||||
access_token: "os.environ/GITLAB_ACCESS_TOKEN"
|
||||
branch: "main"
|
||||
```
|
||||
|
||||
Your GitLab repository should contain `.prompt` files:
|
||||
|
||||
```yaml
|
||||
# prompts/my_gitlab_prompt.prompt
|
||||
---
|
||||
model: gpt-4
|
||||
temperature: 0.7
|
||||
---
|
||||
System: You are a helpful assistant.
|
||||
|
||||
User: {{user_message}}
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
### Complete Example
|
||||
|
||||
Here's a complete example showing multiple prompts with different integrations:
|
||||
|
||||
```yaml
|
||||
model_list:
|
||||
- model_name: gpt-4
|
||||
litellm_params:
|
||||
model: openai/gpt-4
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
prompts:
|
||||
# File-based dotprompt
|
||||
- prompt_id: "coding_assistant"
|
||||
litellm_params:
|
||||
prompt_id: "coding_assistant"
|
||||
prompt_integration: "dotprompt"
|
||||
prompt_directory: "./prompts"
|
||||
|
||||
# Inline dotprompt
|
||||
- prompt_id: "simple_chat"
|
||||
litellm_params:
|
||||
prompt_id: "simple_chat"
|
||||
prompt_integration: "dotprompt"
|
||||
prompt_data:
|
||||
simple_chat:
|
||||
content: "You are a {{personality}} assistant. User: {{message}}"
|
||||
metadata:
|
||||
model: "gpt-4"
|
||||
temperature: 0.8
|
||||
|
||||
# Langfuse prompt
|
||||
- prompt_id: "langfuse_chat"
|
||||
litellm_params:
|
||||
prompt_id: "langfuse_chat"
|
||||
prompt_integration: "langfuse"
|
||||
langfuse_public_key: "os.environ/LANGFUSE_PUBLIC_KEY"
|
||||
langfuse_secret_key: "os.environ/LANGFUSE_SECRET_KEY"
|
||||
|
||||
litellm_settings:
|
||||
global_prompt_directory: "./prompts"
|
||||
```
|
||||
|
||||
### How It Works
|
||||
|
||||
1. **At Startup**: When the proxy starts, it reads the `prompts` field from `config.yaml`
|
||||
2. **Initialization**: Each prompt is initialized based on its `prompt_integration` type
|
||||
3. **In-Memory Storage**: Prompts are stored in the `IN_MEMORY_PROMPT_REGISTRY`
|
||||
4. **Access**: Use these prompts via the `/v1/chat/completions` endpoint with `prompt_id` in the request
|
||||
|
||||
### Using Config-Loaded Prompts
|
||||
|
||||
After loading prompts via config.yaml, use them in your API requests:
|
||||
|
||||
```bash
|
||||
curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \
|
||||
-H 'Content-Type: application/json' \
|
||||
-H 'Authorization: Bearer sk-1234' \
|
||||
-d '{
|
||||
"model": "gpt-4",
|
||||
"prompt_id": "coding_assistant",
|
||||
"prompt_variables": {
|
||||
"language": "python",
|
||||
"task": "create a web scraper"
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
### Prompt Schema Reference
|
||||
|
||||
Each prompt in the `prompts` list requires:
|
||||
|
||||
- **`prompt_id`** (string, required): Unique identifier for the prompt
|
||||
- **`litellm_params`** (object, required): Configuration for the prompt
|
||||
- **`prompt_id`** (string, required): Must match the top-level prompt_id
|
||||
- **`prompt_integration`** (string, required): One of: `dotprompt`, `langfuse`, `bitbucket`, `gitlab`, `custom`
|
||||
- Additional integration-specific parameters (see tabs above)
|
||||
- **`prompt_info`** (object, optional): Metadata about the prompt
|
||||
- **`prompt_type`** (string): Defaults to `"config"` for config-loaded prompts
|
||||
|
||||
### Notes
|
||||
|
||||
- Config-loaded prompts have `prompt_type: "config"` and **cannot be updated** via the API
|
||||
- To update config prompts, modify your `config.yaml` and restart the proxy
|
||||
- For dynamic prompts that can be updated via API, use the `/prompts` endpoints instead
|
||||
- All supported integrations work with config-loaded prompts
|
||||
|
||||
|
||||
## Quick Start
|
||||
|
||||
|
||||
|
||||
@@ -98,7 +98,8 @@ const sidebars = {
|
||||
"proxy/litellm_prompt_management",
|
||||
"proxy/custom_prompt_management",
|
||||
"proxy/native_litellm_prompt",
|
||||
"proxy/prompt_management"
|
||||
"proxy/prompt_management",
|
||||
"proxy/arize_phoenix_prompts"
|
||||
]
|
||||
},
|
||||
{
|
||||
|
||||
@@ -31,6 +31,8 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
ignore_prompt_manager_optional_params: Optional[bool] = False,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
"""
|
||||
Apply cache control directives based on specified injection points.
|
||||
|
||||
@@ -0,0 +1,210 @@
|
||||
# Arize Phoenix Prompt Management Integration
|
||||
|
||||
This integration enables using prompt versions from Arize Phoenix with LiteLLM's completion function.
|
||||
|
||||
## Features
|
||||
|
||||
- Fetch prompt versions from Arize Phoenix API
|
||||
- Workspace-based access control through Arize Phoenix permissions
|
||||
- Mustache/Handlebars-style variable templating (`{{variable}}`)
|
||||
- Support for multi-message chat templates
|
||||
- Automatic model and parameter configuration from prompt metadata
|
||||
- OpenAI and Anthropic provider parameter support
|
||||
|
||||
## Configuration
|
||||
|
||||
Configure Arize Phoenix access in your application:
|
||||
|
||||
```python
|
||||
import litellm
|
||||
|
||||
# Configure Arize Phoenix access
|
||||
# api_base should include your workspace, e.g., "https://app.phoenix.arize.com/s/your-workspace/v1"
|
||||
api_key = "your-arize-phoenix-token"
|
||||
api_base = "https://app.phoenix.arize.com/s/krrishdholakia/v1"
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Basic Usage
|
||||
|
||||
```python
|
||||
import litellm
|
||||
|
||||
# Use with completion
|
||||
response = litellm.completion(
|
||||
model="arize/gpt-4o",
|
||||
prompt_id="UHJvbXB0VmVyc2lvbjox", # Your prompt version ID
|
||||
prompt_variables={"question": "What is artificial intelligence?"},
|
||||
api_key="your-arize-phoenix-token",
|
||||
api_base="https://app.phoenix.arize.com/s/krrishdholakia/v1",
|
||||
)
|
||||
|
||||
print(response.choices[0].message.content)
|
||||
```
|
||||
|
||||
### With Additional Messages
|
||||
|
||||
You can also combine prompt templates with additional messages:
|
||||
|
||||
```python
|
||||
response = litellm.completion(
|
||||
model="arize/gpt-4o",
|
||||
prompt_id="UHJvbXB0VmVyc2lvbjox",
|
||||
prompt_variables={"question": "Explain quantum computing"},
|
||||
api_key="your-arize-phoenix-token",
|
||||
api_base="https://app.phoenix.arize.com/s/krrishdholakia/v1",
|
||||
messages=[
|
||||
{"role": "user", "content": "Please keep your response under 100 words."}
|
||||
],
|
||||
)
|
||||
```
|
||||
|
||||
### Direct Manager Usage
|
||||
|
||||
You can also use the prompt manager directly:
|
||||
|
||||
```python
|
||||
from litellm.integrations.arize.arize_phoenix_prompt_manager import ArizePhoenixPromptManager
|
||||
|
||||
# Initialize the manager
|
||||
manager = ArizePhoenixPromptManager(
|
||||
api_key="your-arize-phoenix-token",
|
||||
api_base="https://app.phoenix.arize.com/s/krrishdholakia/v1",
|
||||
prompt_id="UHJvbXB0VmVyc2lvbjox",
|
||||
)
|
||||
|
||||
# Get rendered messages
|
||||
messages, metadata = manager.get_prompt_template(
|
||||
prompt_id="UHJvbXB0VmVyc2lvbjox",
|
||||
prompt_variables={"question": "What is machine learning?"}
|
||||
)
|
||||
|
||||
print("Rendered messages:", messages)
|
||||
print("Metadata:", metadata)
|
||||
```
|
||||
|
||||
## Prompt Format
|
||||
|
||||
Arize Phoenix prompts support the following structure:
|
||||
|
||||
```json
|
||||
{
|
||||
"data": {
|
||||
"description": "A chatbot prompt",
|
||||
"model_provider": "OPENAI",
|
||||
"model_name": "gpt-4o",
|
||||
"template": {
|
||||
"type": "chat",
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "You are a chatbot"
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "{{question}}"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
},
|
||||
"template_type": "CHAT",
|
||||
"template_format": "MUSTACHE",
|
||||
"invocation_parameters": {
|
||||
"type": "openai",
|
||||
"openai": {
|
||||
"temperature": 1.0
|
||||
}
|
||||
},
|
||||
"id": "UHJvbXB0VmVyc2lvbjox"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Variable Substitution
|
||||
|
||||
Variables in your prompt templates use Mustache/Handlebars syntax:
|
||||
- `{{variable_name}}` - Simple variable substitution
|
||||
|
||||
Example:
|
||||
```
|
||||
Template: "Hello {{name}}, your order {{order_id}} is ready!"
|
||||
Variables: {"name": "Alice", "order_id": "12345"}
|
||||
Result: "Hello Alice, your order 12345 is ready!"
|
||||
```
|
||||
|
||||
## API Reference
|
||||
|
||||
### ArizePhoenixPromptManager
|
||||
|
||||
Main class for managing Arize Phoenix prompts.
|
||||
|
||||
**Methods:**
|
||||
- `get_prompt_template(prompt_id, prompt_variables)` - Get and render a prompt template
|
||||
- `get_available_prompts()` - List available prompt IDs
|
||||
- `reload_prompts()` - Reload prompts from Arize Phoenix
|
||||
|
||||
### ArizePhoenixClient
|
||||
|
||||
Low-level client for Arize Phoenix API.
|
||||
|
||||
**Methods:**
|
||||
- `get_prompt_version(prompt_version_id)` - Fetch a prompt version
|
||||
- `test_connection()` - Test API connection
|
||||
|
||||
## Error Handling
|
||||
|
||||
The integration provides detailed error messages:
|
||||
|
||||
- **404**: Prompt version not found
|
||||
- **401**: Authentication failed (check your access token)
|
||||
- **403**: Access denied (check workspace permissions)
|
||||
|
||||
Example:
|
||||
```python
|
||||
try:
|
||||
response = litellm.completion(
|
||||
model="arize/gpt-4o",
|
||||
prompt_id="invalid-id",
|
||||
arize_config=arize_config,
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
```
|
||||
|
||||
## Getting Your Prompt Version ID and API Base
|
||||
|
||||
1. Log in to Arize Phoenix
|
||||
2. Navigate to your workspace
|
||||
3. Go to Prompts section
|
||||
4. Select a prompt version
|
||||
5. The ID will be in the URL: `/s/{workspace}/v1/prompt_versions/{PROMPT_VERSION_ID}`
|
||||
|
||||
Your `api_base` should be: `https://app.phoenix.arize.com/s/{workspace}/v1`
|
||||
|
||||
For example:
|
||||
- Workspace: `krrishdholakia`
|
||||
- API Base: `https://app.phoenix.arize.com/s/krrishdholakia/v1`
|
||||
- Prompt Version ID: `UHJvbXB0VmVyc2lvbjox`
|
||||
|
||||
You can also fetch it via API:
|
||||
```bash
|
||||
curl -L -X GET 'https://app.phoenix.arize.com/s/krrishdholakia/v1/prompt_versions/UHJvbXB0VmVyc2lvbjox' \
|
||||
-H 'Authorization: Bearer YOUR_TOKEN'
|
||||
```
|
||||
|
||||
## Support
|
||||
|
||||
For issues or questions:
|
||||
- LiteLLM Issues: https://github.com/BerriAI/litellm/issues
|
||||
- Arize Phoenix Docs: https://docs.arize.com/phoenix
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec
|
||||
from litellm.integrations.custom_prompt_management import CustomPromptManagement
|
||||
|
||||
from litellm.types.prompts.init_prompts import SupportedPromptIntegrations
|
||||
|
||||
from .arize_phoenix_prompt_manager import ArizePhoenixPromptManager
|
||||
|
||||
# Global instances
|
||||
global_arize_config: Optional[dict] = None
|
||||
|
||||
|
||||
def prompt_initializer(
|
||||
litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec"
|
||||
) -> "CustomPromptManagement":
|
||||
"""
|
||||
Initialize a prompt from Arize Phoenix.
|
||||
"""
|
||||
api_key = getattr(litellm_params, "api_key", None) or os.environ.get(
|
||||
"PHOENIX_API_KEY"
|
||||
)
|
||||
api_base = getattr(litellm_params, "api_base", None)
|
||||
prompt_id = getattr(litellm_params, "prompt_id", None)
|
||||
|
||||
if not api_key or not api_base:
|
||||
raise ValueError(
|
||||
"api_key and api_base are required for Arize Phoenix prompt integration"
|
||||
)
|
||||
|
||||
try:
|
||||
arize_prompt_manager = ArizePhoenixPromptManager(
|
||||
**{
|
||||
"api_key": api_key,
|
||||
"api_base": api_base,
|
||||
"prompt_id": prompt_id,
|
||||
**litellm_params.model_dump(
|
||||
exclude={"api_key", "api_base", "prompt_id"}
|
||||
),
|
||||
},
|
||||
)
|
||||
|
||||
return arize_prompt_manager
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
||||
prompt_initializer_registry = {
|
||||
SupportedPromptIntegrations.ARIZE_PHOENIX.value: prompt_initializer,
|
||||
}
|
||||
@@ -0,0 +1,108 @@
|
||||
"""
|
||||
Arize Phoenix API client for fetching prompt versions from Arize Phoenix.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
|
||||
class ArizePhoenixClient:
|
||||
"""
|
||||
Client for interacting with Arize Phoenix API to fetch prompt versions.
|
||||
|
||||
Supports:
|
||||
- Authentication with Bearer tokens
|
||||
- Fetching prompt versions
|
||||
- Direct API base URL configuration
|
||||
"""
|
||||
|
||||
def __init__(self, api_key: Optional[str] = None, api_base: Optional[str] = None):
|
||||
"""
|
||||
Initialize the Arize Phoenix client.
|
||||
|
||||
Args:
|
||||
api_key: Arize Phoenix API token
|
||||
api_base: Base URL for the Arize Phoenix API (e.g., 'https://app.phoenix.arize.com/s/workspace/v1')
|
||||
"""
|
||||
self.api_key = api_key
|
||||
self.api_base = api_base
|
||||
|
||||
if not self.api_key:
|
||||
raise ValueError("api_key is required")
|
||||
|
||||
if not self.api_base:
|
||||
raise ValueError("api_base is required")
|
||||
|
||||
# Set up authentication headers
|
||||
self.headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Accept": "application/json",
|
||||
}
|
||||
|
||||
# Initialize HTTPHandler
|
||||
self.http_handler = HTTPHandler(disable_default_headers=True)
|
||||
|
||||
def get_prompt_version(self, prompt_version_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Fetch a prompt version from Arize Phoenix.
|
||||
|
||||
Args:
|
||||
prompt_version_id: The ID of the prompt version to fetch
|
||||
|
||||
Returns:
|
||||
Dictionary containing prompt version data, or None if not found
|
||||
"""
|
||||
url = f"{self.api_base}/v1/prompt_versions/{prompt_version_id}"
|
||||
|
||||
try:
|
||||
# Use the underlying httpx client directly to avoid query param extraction
|
||||
response = self.http_handler.get(url, headers=self.headers)
|
||||
response.raise_for_status()
|
||||
|
||||
data = response.json()
|
||||
return data.get("data")
|
||||
|
||||
except Exception as e:
|
||||
# Check if it's an HTTP error
|
||||
response = getattr(e, "response", None)
|
||||
if response is not None and hasattr(response, "status_code"):
|
||||
if response.status_code == 404:
|
||||
return None
|
||||
elif response.status_code == 403:
|
||||
raise Exception(
|
||||
f"Access denied to prompt version '{prompt_version_id}'. Check your Arize Phoenix permissions."
|
||||
)
|
||||
elif response.status_code == 401:
|
||||
raise Exception(
|
||||
"Authentication failed. Check your Arize Phoenix API key and permissions."
|
||||
)
|
||||
else:
|
||||
raise Exception(
|
||||
f"Failed to fetch prompt version '{prompt_version_id}': {e}"
|
||||
)
|
||||
else:
|
||||
raise Exception(
|
||||
f"Error fetching prompt version '{prompt_version_id}': {e}"
|
||||
)
|
||||
|
||||
def test_connection(self) -> bool:
|
||||
"""
|
||||
Test the connection to the Arize Phoenix API.
|
||||
|
||||
Returns:
|
||||
True if connection is successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
# Try to access the prompt_versions endpoint to test connection
|
||||
url = f"{self.api_base}/prompt_versions"
|
||||
response = self.http_handler.client.get(url, headers=self.headers)
|
||||
response.raise_for_status()
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
def close(self):
|
||||
"""Close the HTTP handler to free resources."""
|
||||
if hasattr(self, "http_handler"):
|
||||
self.http_handler.close()
|
||||
@@ -0,0 +1,457 @@
|
||||
"""
|
||||
Arize Phoenix prompt manager that integrates with LiteLLM's prompt management system.
|
||||
Fetches prompt versions from Arize Phoenix and provides workspace-based access control.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
from jinja2 import DictLoader, Environment, select_autoescape
|
||||
|
||||
from litellm.integrations.custom_prompt_management import CustomPromptManagement
|
||||
from litellm.integrations.prompt_management_base import (
|
||||
PromptManagementBase,
|
||||
PromptManagementClient,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
from .arize_phoenix_client import ArizePhoenixClient
|
||||
|
||||
|
||||
class ArizePhoenixPromptTemplate:
|
||||
"""
|
||||
Represents a prompt template loaded from Arize Phoenix.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
template_id: str,
|
||||
messages: List[Dict[str, Any]],
|
||||
metadata: Dict[str, Any],
|
||||
model: Optional[str] = None,
|
||||
):
|
||||
self.template_id = template_id
|
||||
self.messages = messages
|
||||
self.metadata = metadata
|
||||
self.model = model or metadata.get("model_name")
|
||||
self.model_provider = metadata.get("model_provider")
|
||||
self.temperature = metadata.get("temperature")
|
||||
self.max_tokens = metadata.get("max_tokens")
|
||||
self.invocation_parameters = metadata.get("invocation_parameters", {})
|
||||
self.description = metadata.get("description", "")
|
||||
self.template_format = metadata.get("template_format", "MUSTACHE")
|
||||
|
||||
def __repr__(self):
|
||||
return (
|
||||
f"ArizePhoenixPromptTemplate(id='{self.template_id}', model='{self.model}')"
|
||||
)
|
||||
|
||||
|
||||
class ArizePhoenixTemplateManager:
|
||||
"""
|
||||
Manager for loading and rendering prompt templates from Arize Phoenix.
|
||||
|
||||
Supports:
|
||||
- Fetching prompt versions from Arize Phoenix API
|
||||
- Workspace-based access control through Arize Phoenix permissions
|
||||
- Mustache/Handlebars-style templating (using Jinja2)
|
||||
- Model configuration and invocation parameters
|
||||
- Multi-message chat templates
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
prompt_id: Optional[str] = None,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.api_base = api_base
|
||||
self.prompt_id = prompt_id
|
||||
self.prompts: Dict[str, ArizePhoenixPromptTemplate] = {}
|
||||
self.arize_client = ArizePhoenixClient(
|
||||
api_key=self.api_key, api_base=self.api_base
|
||||
)
|
||||
|
||||
self.jinja_env = Environment(
|
||||
loader=DictLoader({}),
|
||||
autoescape=select_autoescape(["html", "xml"]),
|
||||
# Use Mustache/Handlebars-style delimiters
|
||||
variable_start_string="{{",
|
||||
variable_end_string="}}",
|
||||
block_start_string="{%",
|
||||
block_end_string="%}",
|
||||
comment_start_string="{#",
|
||||
comment_end_string="#}",
|
||||
)
|
||||
|
||||
# Load prompt from Arize Phoenix if prompt_id is provided
|
||||
if self.prompt_id:
|
||||
self._load_prompt_from_arize(self.prompt_id)
|
||||
|
||||
def _load_prompt_from_arize(self, prompt_version_id: str) -> None:
|
||||
"""Load a specific prompt version from Arize Phoenix."""
|
||||
try:
|
||||
# Fetch the prompt version from Arize Phoenix
|
||||
prompt_data = self.arize_client.get_prompt_version(prompt_version_id)
|
||||
|
||||
if prompt_data:
|
||||
template = self._parse_prompt_data(prompt_data, prompt_version_id)
|
||||
self.prompts[prompt_version_id] = template
|
||||
else:
|
||||
raise ValueError(f"Prompt version '{prompt_version_id}' not found")
|
||||
except Exception as e:
|
||||
raise Exception(
|
||||
f"Failed to load prompt version '{prompt_version_id}' from Arize Phoenix: {e}"
|
||||
)
|
||||
|
||||
def _parse_prompt_data(
|
||||
self, data: Dict[str, Any], prompt_version_id: str
|
||||
) -> ArizePhoenixPromptTemplate:
|
||||
"""Parse Arize Phoenix prompt data and extract messages and metadata."""
|
||||
template_data = data.get("template", {})
|
||||
messages = template_data.get("messages", [])
|
||||
|
||||
# Extract invocation parameters
|
||||
invocation_params = data.get("invocation_parameters", {})
|
||||
provider_params = {}
|
||||
|
||||
# Extract provider-specific parameters
|
||||
if "openai" in invocation_params:
|
||||
provider_params = invocation_params["openai"]
|
||||
elif "anthropic" in invocation_params:
|
||||
provider_params = invocation_params["anthropic"]
|
||||
else:
|
||||
# Try to find any nested provider params
|
||||
for key, value in invocation_params.items():
|
||||
if isinstance(value, dict):
|
||||
provider_params = value
|
||||
break
|
||||
|
||||
# Build metadata dictionary
|
||||
metadata = {
|
||||
"model_name": data.get("model_name"),
|
||||
"model_provider": data.get("model_provider"),
|
||||
"description": data.get("description", ""),
|
||||
"template_type": data.get("template_type"),
|
||||
"template_format": data.get("template_format", "MUSTACHE"),
|
||||
"invocation_parameters": invocation_params,
|
||||
"temperature": provider_params.get("temperature"),
|
||||
"max_tokens": provider_params.get("max_tokens"),
|
||||
}
|
||||
|
||||
return ArizePhoenixPromptTemplate(
|
||||
template_id=prompt_version_id,
|
||||
messages=messages,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
def render_template(
|
||||
self, template_id: str, variables: Optional[Dict[str, Any]] = None
|
||||
) -> List[AllMessageValues]:
|
||||
"""Render a template with the given variables and return formatted messages."""
|
||||
if template_id not in self.prompts:
|
||||
raise ValueError(f"Template '{template_id}' not found")
|
||||
|
||||
template = self.prompts[template_id]
|
||||
rendered_messages: List[AllMessageValues] = []
|
||||
|
||||
for message in template.messages:
|
||||
role = message.get("role", "user")
|
||||
content_parts = message.get("content", [])
|
||||
|
||||
# Render each content part
|
||||
rendered_content_parts = []
|
||||
for part in content_parts:
|
||||
if part.get("type") == "text":
|
||||
text = part.get("text", "")
|
||||
# Render the text with Jinja2 (Mustache-style)
|
||||
jinja_template = self.jinja_env.from_string(text)
|
||||
rendered_text = jinja_template.render(**(variables or {}))
|
||||
rendered_content_parts.append(rendered_text)
|
||||
else:
|
||||
# Handle other content types if needed
|
||||
rendered_content_parts.append(part)
|
||||
|
||||
# Combine rendered content
|
||||
final_content = " ".join(rendered_content_parts)
|
||||
|
||||
rendered_messages.append(
|
||||
{"role": role, "content": final_content} # type: ignore
|
||||
)
|
||||
|
||||
return rendered_messages
|
||||
|
||||
def get_template(self, template_id: str) -> Optional[ArizePhoenixPromptTemplate]:
|
||||
"""Get a template by ID."""
|
||||
return self.prompts.get(template_id)
|
||||
|
||||
def list_templates(self) -> List[str]:
|
||||
"""List all available template IDs."""
|
||||
return list(self.prompts.keys())
|
||||
|
||||
|
||||
class ArizePhoenixPromptManager(CustomPromptManagement):
|
||||
"""
|
||||
Arize Phoenix prompt manager that integrates with LiteLLM's prompt management system.
|
||||
|
||||
This class enables using prompt versions from Arize Phoenix with the
|
||||
litellm completion() function by implementing the PromptManagementBase interface.
|
||||
|
||||
Usage:
|
||||
# Configure Arize Phoenix access
|
||||
arize_config = {
|
||||
"workspace": "your-workspace",
|
||||
"access_token": "your-token",
|
||||
}
|
||||
|
||||
# Use with completion
|
||||
response = litellm.completion(
|
||||
model="arize/gpt-4o",
|
||||
prompt_id="UHJvbXB0VmVyc2lvbjox",
|
||||
prompt_variables={"question": "What is AI?"},
|
||||
arize_config=arize_config,
|
||||
messages=[{"role": "user", "content": "This will be combined with the prompt"}]
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
prompt_id: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
self.api_key = api_key
|
||||
self.api_base = api_base
|
||||
self.prompt_id = prompt_id
|
||||
self._prompt_manager: Optional[ArizePhoenixTemplateManager] = None
|
||||
|
||||
@property
|
||||
def integration_name(self) -> str:
|
||||
"""Integration name used in model names like 'arize/gpt-4o'."""
|
||||
return "arize"
|
||||
|
||||
@property
|
||||
def prompt_manager(self) -> ArizePhoenixTemplateManager:
|
||||
"""Get or create the prompt manager instance."""
|
||||
if self._prompt_manager is None:
|
||||
self._prompt_manager = ArizePhoenixTemplateManager(
|
||||
api_key=self.api_key,
|
||||
api_base=self.api_base,
|
||||
prompt_id=self.prompt_id,
|
||||
)
|
||||
return self._prompt_manager
|
||||
|
||||
def get_prompt_template(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_variables: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[List[AllMessageValues], Dict[str, Any]]:
|
||||
"""
|
||||
Get a prompt template and render it with variables.
|
||||
|
||||
Args:
|
||||
prompt_id: The ID of the prompt version
|
||||
prompt_variables: Variables to substitute in the template
|
||||
|
||||
Returns:
|
||||
Tuple of (rendered_messages, metadata)
|
||||
"""
|
||||
template = self.prompt_manager.get_template(prompt_id)
|
||||
if not template:
|
||||
raise ValueError(f"Prompt template '{prompt_id}' not found")
|
||||
|
||||
# Render the template
|
||||
rendered_messages = self.prompt_manager.render_template(
|
||||
prompt_id, prompt_variables or {}
|
||||
)
|
||||
|
||||
# Extract metadata
|
||||
metadata = {
|
||||
"model": template.model,
|
||||
"temperature": template.temperature,
|
||||
"max_tokens": template.max_tokens,
|
||||
}
|
||||
|
||||
# Add additional invocation parameters
|
||||
invocation_params = template.invocation_parameters
|
||||
provider_params = {}
|
||||
|
||||
if "openai" in invocation_params:
|
||||
provider_params = invocation_params["openai"]
|
||||
elif "anthropic" in invocation_params:
|
||||
provider_params = invocation_params["anthropic"]
|
||||
|
||||
# Add any additional parameters
|
||||
for key, value in provider_params.items():
|
||||
if key not in metadata:
|
||||
metadata[key] = value
|
||||
|
||||
return rendered_messages, metadata
|
||||
|
||||
def pre_call_hook(
|
||||
self,
|
||||
user_id: Optional[str],
|
||||
messages: List[AllMessageValues],
|
||||
function_call: Optional[Union[Dict[str, Any], str]] = None,
|
||||
litellm_params: Optional[Dict[str, Any]] = None,
|
||||
prompt_id: Optional[str] = None,
|
||||
prompt_variables: Optional[Dict[str, Any]] = None,
|
||||
**kwargs,
|
||||
) -> Tuple[List[AllMessageValues], Optional[Dict[str, Any]]]:
|
||||
"""
|
||||
Pre-call hook that processes the prompt template before making the LLM call.
|
||||
"""
|
||||
if not prompt_id:
|
||||
return messages, litellm_params
|
||||
|
||||
try:
|
||||
# Get the rendered messages and metadata
|
||||
rendered_messages, prompt_metadata = self.get_prompt_template(
|
||||
prompt_id, prompt_variables
|
||||
)
|
||||
|
||||
# Merge rendered messages with existing messages
|
||||
if rendered_messages:
|
||||
# Prepend rendered messages to existing messages
|
||||
final_messages = rendered_messages + messages
|
||||
else:
|
||||
final_messages = messages
|
||||
|
||||
# Update litellm_params with prompt metadata
|
||||
if litellm_params is None:
|
||||
litellm_params = {}
|
||||
|
||||
# Apply model and parameters from prompt metadata
|
||||
if prompt_metadata.get("model") and not self.ignore_prompt_manager_model:
|
||||
litellm_params["model"] = prompt_metadata["model"]
|
||||
|
||||
if not self.ignore_prompt_manager_optional_params:
|
||||
for param in [
|
||||
"temperature",
|
||||
"max_tokens",
|
||||
"top_p",
|
||||
"frequency_penalty",
|
||||
"presence_penalty",
|
||||
]:
|
||||
if param in prompt_metadata:
|
||||
litellm_params[param] = prompt_metadata[param]
|
||||
|
||||
return final_messages, litellm_params
|
||||
|
||||
except Exception as e:
|
||||
# Log error but don't fail the call
|
||||
import litellm
|
||||
|
||||
litellm._logging.verbose_proxy_logger.error(
|
||||
f"Error in Arize Phoenix prompt pre_call_hook: {e}"
|
||||
)
|
||||
return messages, litellm_params
|
||||
|
||||
def get_available_prompts(self) -> List[str]:
|
||||
"""Get list of available prompt IDs."""
|
||||
return self.prompt_manager.list_templates()
|
||||
|
||||
def reload_prompts(self) -> None:
|
||||
"""Reload prompts from Arize Phoenix."""
|
||||
if self.prompt_id:
|
||||
self._prompt_manager = None # Reset to force reload
|
||||
self.prompt_manager # This will trigger reload
|
||||
|
||||
def should_run_prompt_management(
|
||||
self,
|
||||
prompt_id: str,
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
"""
|
||||
Determine if prompt management should run based on the prompt_id.
|
||||
|
||||
For Arize Phoenix, we always return True and handle the prompt loading
|
||||
in the _compile_prompt_helper method.
|
||||
"""
|
||||
return True
|
||||
|
||||
def _compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> PromptManagementClient:
|
||||
"""
|
||||
Compile an Arize Phoenix prompt template into a PromptManagementClient structure.
|
||||
|
||||
This method:
|
||||
1. Loads the prompt version from Arize Phoenix
|
||||
2. Renders it with the provided variables
|
||||
3. Returns formatted chat messages
|
||||
4. Extracts model and optional parameters from metadata
|
||||
"""
|
||||
try:
|
||||
# Load the prompt from Arize Phoenix if not already loaded
|
||||
if prompt_id not in self.prompt_manager.prompts:
|
||||
self.prompt_manager._load_prompt_from_arize(prompt_id)
|
||||
|
||||
# Get the rendered messages and metadata
|
||||
rendered_messages, prompt_metadata = self.get_prompt_template(
|
||||
prompt_id, prompt_variables
|
||||
)
|
||||
|
||||
# Extract model from metadata (if specified)
|
||||
template_model = prompt_metadata.get("model")
|
||||
|
||||
# Extract optional parameters from metadata
|
||||
optional_params = {}
|
||||
for param in [
|
||||
"temperature",
|
||||
"max_tokens",
|
||||
"top_p",
|
||||
"frequency_penalty",
|
||||
"presence_penalty",
|
||||
]:
|
||||
if param in prompt_metadata:
|
||||
optional_params[param] = prompt_metadata[param]
|
||||
|
||||
return PromptManagementClient(
|
||||
prompt_id=prompt_id,
|
||||
prompt_template=rendered_messages,
|
||||
prompt_template_model=template_model,
|
||||
prompt_template_optional_params=optional_params,
|
||||
completed_messages=None,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise ValueError(f"Error compiling prompt '{prompt_id}': {e}")
|
||||
|
||||
def get_chat_completion_prompt(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
non_default_params: dict,
|
||||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
ignore_prompt_manager_optional_params: Optional[bool] = False,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
"""
|
||||
Get chat completion prompt from Arize Phoenix and return processed model, messages, and parameters.
|
||||
"""
|
||||
return PromptManagementBase.get_chat_completion_prompt(
|
||||
self,
|
||||
model,
|
||||
messages,
|
||||
non_default_params,
|
||||
prompt_id,
|
||||
prompt_variables,
|
||||
dynamic_callback_params,
|
||||
prompt_label,
|
||||
prompt_version,
|
||||
self.ignore_prompt_manager_model,
|
||||
self.ignore_prompt_manager_optional_params,
|
||||
)
|
||||
@@ -491,6 +491,8 @@ class BitBucketPromptManager(CustomPromptManagement):
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
ignore_prompt_manager_optional_params: Optional[bool] = False,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
"""
|
||||
Get chat completion prompt from BitBucket and return processed model, messages, and parameters.
|
||||
|
||||
@@ -180,6 +180,8 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
ignore_prompt_manager_optional_params: Optional[bool] = False,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
"""
|
||||
Returns:
|
||||
@@ -552,8 +554,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
||||
from copy import copy
|
||||
|
||||
from litellm import Choices, Message, ModelResponse
|
||||
turn_off_message_logging: bool = getattr(self, "turn_off_message_logging", False)
|
||||
|
||||
|
||||
turn_off_message_logging: bool = getattr(
|
||||
self, "turn_off_message_logging", False
|
||||
)
|
||||
|
||||
if turn_off_message_logging is False:
|
||||
return model_call_details
|
||||
|
||||
@@ -579,6 +584,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
||||
if isinstance(response, dict) and "output" in response:
|
||||
# Make a copy to avoid modifying the original
|
||||
from copy import deepcopy
|
||||
|
||||
response_copy = deepcopy(response)
|
||||
# Redact content in output array
|
||||
if isinstance(response_copy.get("output"), list):
|
||||
@@ -587,7 +593,10 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
||||
if isinstance(output_item["content"], list):
|
||||
# Redact text in content items
|
||||
for content_item in output_item["content"]:
|
||||
if isinstance(content_item, dict) and "text" in content_item:
|
||||
if (
|
||||
isinstance(content_item, dict)
|
||||
and "text" in content_item
|
||||
):
|
||||
content_item["text"] = redacted_str
|
||||
standard_logging_object_copy["response"] = response_copy
|
||||
else:
|
||||
@@ -615,29 +624,34 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
||||
def handle_callback_failure(self, callback_name: str):
|
||||
"""
|
||||
Handle callback logging failures by incrementing Prometheus metrics.
|
||||
|
||||
|
||||
Call this method in exception handlers within your callback when logging fails.
|
||||
"""
|
||||
try:
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
||||
|
||||
all_callbacks = litellm.logging_callback_manager._get_all_callbacks()
|
||||
|
||||
|
||||
for callback_obj in all_callbacks:
|
||||
if hasattr(callback_obj, 'increment_callback_logging_failure'):
|
||||
verbose_logger.debug(f"Incrementing callback failure metric for {callback_name}")
|
||||
if hasattr(callback_obj, "increment_callback_logging_failure"):
|
||||
verbose_logger.debug(
|
||||
f"Incrementing callback failure metric for {callback_name}"
|
||||
)
|
||||
callback_obj.increment_callback_logging_failure(callback_name=callback_name) # type: ignore
|
||||
return
|
||||
|
||||
|
||||
verbose_logger.debug(
|
||||
f"No callback with increment_callback_logging_failure method found for {callback_name}. "
|
||||
"Ensure 'prometheus' is in your callbacks config."
|
||||
)
|
||||
|
||||
|
||||
except Exception as e:
|
||||
from litellm._logging import verbose_logger
|
||||
verbose_logger.debug(f"Error in handle_callback_failure for {callback_name}: {str(e)}")
|
||||
|
||||
verbose_logger.debug(
|
||||
f"Error in handle_callback_failure for {callback_name}: {str(e)}"
|
||||
)
|
||||
|
||||
async def _strip_base64_from_messages(
|
||||
self,
|
||||
@@ -656,10 +670,14 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
||||
"""
|
||||
raw_messages: Any = payload.get("messages", [])
|
||||
messages: List[Any] = raw_messages if isinstance(raw_messages, list) else []
|
||||
verbose_logger.debug(f"[CustomLogger] Stripping base64 from {len(messages)} messages")
|
||||
verbose_logger.debug(
|
||||
f"[CustomLogger] Stripping base64 from {len(messages)} messages"
|
||||
)
|
||||
|
||||
if messages:
|
||||
payload["messages"] = self._process_messages(messages=messages, max_depth=max_depth)
|
||||
payload["messages"] = self._process_messages(
|
||||
messages=messages, max_depth=max_depth
|
||||
)
|
||||
|
||||
total_items = 0
|
||||
for m in payload.get("messages", []) or []:
|
||||
@@ -674,7 +692,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
||||
return payload
|
||||
|
||||
def _strip_base64_from_messages_sync(
|
||||
self, payload: "StandardLoggingPayload", max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER
|
||||
self,
|
||||
payload: "StandardLoggingPayload",
|
||||
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER,
|
||||
) -> "StandardLoggingPayload":
|
||||
"""
|
||||
Removes or redacts base64-encoded file data (e.g., PDFs, images, audio)
|
||||
@@ -688,7 +708,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
||||
"""
|
||||
raw_messages: Any = payload.get("messages", [])
|
||||
messages: List[Any] = raw_messages if isinstance(raw_messages, list) else []
|
||||
verbose_logger.debug(f"[CustomLogger] Stripping base64 from {len(messages)} messages")
|
||||
verbose_logger.debug(
|
||||
f"[CustomLogger] Stripping base64 from {len(messages)} messages"
|
||||
)
|
||||
|
||||
if messages:
|
||||
payload["messages"] = self._process_messages(
|
||||
@@ -751,7 +773,11 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
||||
ctype = content.get("type")
|
||||
return not (isinstance(ctype, str) and ctype != "text")
|
||||
|
||||
def _process_messages(self, messages: List[Any], max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER) -> List[Dict[str, Any]]:
|
||||
def _process_messages(
|
||||
self,
|
||||
messages: List[Any],
|
||||
max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER,
|
||||
) -> List[Dict[str, Any]]:
|
||||
filtered_messages: List[Dict[str, Any]] = []
|
||||
for msg in messages:
|
||||
if not isinstance(msg, dict):
|
||||
|
||||
@@ -10,6 +10,17 @@ from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
|
||||
class CustomPromptManagement(CustomLogger, PromptManagementBase):
|
||||
def __init__(
|
||||
self,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
ignore_prompt_manager_optional_params: Optional[bool] = False,
|
||||
**kwargs,
|
||||
):
|
||||
self.ignore_prompt_manager_model = ignore_prompt_manager_model
|
||||
self.ignore_prompt_manager_optional_params = (
|
||||
ignore_prompt_manager_optional_params
|
||||
)
|
||||
|
||||
def get_chat_completion_prompt(
|
||||
self,
|
||||
model: str,
|
||||
@@ -20,6 +31,8 @@ class CustomPromptManagement(CustomLogger, PromptManagementBase):
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
ignore_prompt_manager_optional_params: Optional[bool] = False,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
"""
|
||||
Returns:
|
||||
|
||||
@@ -163,6 +163,8 @@ class DotpromptManager(CustomPromptManagement):
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
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
|
||||
|
||||
@@ -3,40 +3,42 @@ GitLab prompt manager with configurable prompts folder.
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
from jinja2 import DictLoader, Environment, select_autoescape
|
||||
|
||||
from litellm.integrations.custom_prompt_management import CustomPromptManagement
|
||||
from litellm.integrations.gitlab.gitlab_client import GitLabClient
|
||||
from litellm.integrations.prompt_management_base import (
|
||||
PromptManagementBase,
|
||||
PromptManagementClient,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
from litellm.integrations.gitlab.gitlab_client import GitLabClient
|
||||
|
||||
|
||||
GITLAB_PREFIX = "gitlab::"
|
||||
|
||||
|
||||
def encode_prompt_id(raw_id: str) -> str:
|
||||
"""Convert GitLab path IDs like 'invoice/extract' → 'gitlab::invoice::extract'"""
|
||||
if raw_id.startswith(GITLAB_PREFIX):
|
||||
return raw_id # already encoded
|
||||
return f"{GITLAB_PREFIX}{raw_id.replace('/', '::')}"
|
||||
|
||||
|
||||
def decode_prompt_id(encoded_id: str) -> str:
|
||||
"""Convert 'gitlab::invoice::extract' → 'invoice/extract'"""
|
||||
if not encoded_id.startswith(GITLAB_PREFIX):
|
||||
return encoded_id
|
||||
return encoded_id[len(GITLAB_PREFIX):].replace("::", "/")
|
||||
return encoded_id[len(GITLAB_PREFIX) :].replace("::", "/")
|
||||
|
||||
|
||||
class GitLabPromptTemplate:
|
||||
def __init__(
|
||||
self,
|
||||
template_id: str,
|
||||
content: str,
|
||||
metadata: Dict[str, Any],
|
||||
model: Optional[str] = None,
|
||||
self,
|
||||
template_id: str,
|
||||
content: str,
|
||||
metadata: Dict[str, Any],
|
||||
model: Optional[str] = None,
|
||||
):
|
||||
self.template_id = template_id
|
||||
self.content = content
|
||||
@@ -60,13 +62,12 @@ class GitLabTemplateManager:
|
||||
New: supports `prompts_path` (or `folder`) in gitlab_config to scope where prompts live.
|
||||
"""
|
||||
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
gitlab_config: Dict[str, Any],
|
||||
prompt_id: Optional[str] = None,
|
||||
ref: Optional[str] = None,
|
||||
gitlab_client: Optional[GitLabClient] = None
|
||||
self,
|
||||
gitlab_config: Dict[str, Any],
|
||||
prompt_id: Optional[str] = None,
|
||||
ref: Optional[str] = None,
|
||||
gitlab_client: Optional[GitLabClient] = None,
|
||||
):
|
||||
self.gitlab_config = dict(gitlab_config)
|
||||
self.prompt_id = prompt_id
|
||||
@@ -78,9 +79,9 @@ class GitLabTemplateManager:
|
||||
|
||||
# Folder inside repo to look for prompts (e.g., "prompts" or "prompts/chat")
|
||||
self.prompts_path: str = (
|
||||
self.gitlab_config.get("prompts_path")
|
||||
or self.gitlab_config.get("folder")
|
||||
or ""
|
||||
self.gitlab_config.get("prompts_path")
|
||||
or self.gitlab_config.get("folder")
|
||||
or ""
|
||||
).strip("/")
|
||||
|
||||
self.jinja_env = Environment(
|
||||
@@ -120,7 +121,9 @@ class GitLabTemplateManager:
|
||||
|
||||
# ---------- loading ----------
|
||||
|
||||
def _load_prompt_from_gitlab(self, prompt_id: str, *, ref: Optional[str] = None) -> None:
|
||||
def _load_prompt_from_gitlab(
|
||||
self, prompt_id: str, *, ref: Optional[str] = None
|
||||
) -> None:
|
||||
"""Load a specific .prompt file from GitLab (scoped under prompts_path if set)."""
|
||||
try:
|
||||
# prompt_id = decode_prompt_id(prompt_id)
|
||||
@@ -130,7 +133,9 @@ class GitLabTemplateManager:
|
||||
template = self._parse_prompt_file(prompt_content, prompt_id)
|
||||
self.prompts[prompt_id] = template
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to load prompt '{encode_prompt_id(prompt_id)}' from GitLab: {e}")
|
||||
raise Exception(
|
||||
f"Failed to load prompt '{encode_prompt_id(prompt_id)}' from GitLab: {e}"
|
||||
)
|
||||
|
||||
def load_all_prompts(self, *, recursive: bool = True) -> List[str]:
|
||||
"""
|
||||
@@ -146,9 +151,7 @@ class GitLabTemplateManager:
|
||||
|
||||
# ---------- parsing & rendering ----------
|
||||
|
||||
def _parse_prompt_file(
|
||||
self, content: str, prompt_id: str
|
||||
) -> GitLabPromptTemplate:
|
||||
def _parse_prompt_file(self, content: str, prompt_id: str) -> GitLabPromptTemplate:
|
||||
if content.startswith("---"):
|
||||
parts = content.split("---", 2)
|
||||
if len(parts) >= 3:
|
||||
@@ -165,6 +168,7 @@ class GitLabTemplateManager:
|
||||
if frontmatter_str:
|
||||
try:
|
||||
import yaml
|
||||
|
||||
metadata = yaml.safe_load(frontmatter_str) or {}
|
||||
except ImportError:
|
||||
metadata = self._parse_yaml_basic(frontmatter_str)
|
||||
@@ -199,7 +203,7 @@ class GitLabTemplateManager:
|
||||
return result
|
||||
|
||||
def render_template(
|
||||
self, template_id: str, variables: Optional[Dict[str, Any]] = None
|
||||
self, template_id: str, variables: Optional[Dict[str, Any]] = None
|
||||
) -> str:
|
||||
if template_id not in self.prompts:
|
||||
raise ValueError(f"Template '{template_id}' not found")
|
||||
@@ -244,9 +248,14 @@ class GitLabTemplateManager:
|
||||
)
|
||||
# Classic returns GitLab tree entries; filter *.prompt blobs
|
||||
files = []
|
||||
for f in (raw or []):
|
||||
if isinstance(f, dict) and f.get("type") == "blob" and str(f.get("path", "")).endswith(".prompt") and 'path' in f:
|
||||
files.append(f['path'])
|
||||
for f in raw or []:
|
||||
if (
|
||||
isinstance(f, dict)
|
||||
and f.get("type") == "blob"
|
||||
and str(f.get("path", "")).endswith(".prompt")
|
||||
and "path" in f
|
||||
):
|
||||
files.append(f["path"]) # type: ignore
|
||||
|
||||
return [self._repo_path_to_id(p) for p in files]
|
||||
|
||||
@@ -266,11 +275,11 @@ class GitLabPromptManager(CustomPromptManagement):
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
gitlab_config: Dict[str, Any],
|
||||
prompt_id: Optional[str] = None,
|
||||
ref: Optional[str] = None, # tag/branch/SHA override
|
||||
gitlab_client: Optional[GitLabClient] = None
|
||||
self,
|
||||
gitlab_config: Dict[str, Any],
|
||||
prompt_id: Optional[str] = None,
|
||||
ref: Optional[str] = None, # tag/branch/SHA override
|
||||
gitlab_client: Optional[GitLabClient] = None,
|
||||
):
|
||||
self.gitlab_config = gitlab_config
|
||||
self.prompt_id = prompt_id
|
||||
@@ -295,16 +304,16 @@ class GitLabPromptManager(CustomPromptManagement):
|
||||
gitlab_config=self.gitlab_config,
|
||||
prompt_id=self.prompt_id,
|
||||
ref=self._ref_override,
|
||||
gitlab_client=self._injected_gitlab_client
|
||||
gitlab_client=self._injected_gitlab_client,
|
||||
)
|
||||
return self._prompt_manager
|
||||
|
||||
def get_prompt_template(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_variables: Optional[Dict[str, Any]] = None,
|
||||
*,
|
||||
ref: Optional[str] = None,
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_variables: Optional[Dict[str, Any]] = None,
|
||||
*,
|
||||
ref: Optional[str] = None,
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
if prompt_id not in self.prompt_manager.prompts:
|
||||
self.prompt_manager._load_prompt_from_gitlab(prompt_id, ref=ref)
|
||||
@@ -326,15 +335,15 @@ class GitLabPromptManager(CustomPromptManagement):
|
||||
return rendered_prompt, metadata
|
||||
|
||||
def pre_call_hook(
|
||||
self,
|
||||
user_id: Optional[str],
|
||||
messages: List[AllMessageValues],
|
||||
function_call: Optional[Union[Dict[str, Any], str]] = None,
|
||||
litellm_params: Optional[Dict[str, Any]] = None,
|
||||
prompt_id: Optional[str] = None,
|
||||
prompt_variables: Optional[Dict[str, Any]] = None,
|
||||
prompt_version: Optional[str] = None,
|
||||
**kwargs,
|
||||
self,
|
||||
user_id: Optional[str],
|
||||
messages: List[AllMessageValues],
|
||||
function_call: Optional[Union[Dict[str, Any], str]] = None,
|
||||
litellm_params: Optional[Dict[str, Any]] = None,
|
||||
prompt_id: Optional[str] = None,
|
||||
prompt_variables: Optional[Dict[str, Any]] = None,
|
||||
prompt_version: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> Tuple[List[AllMessageValues], Optional[Dict[str, Any]]]:
|
||||
if not prompt_id:
|
||||
return messages, litellm_params
|
||||
@@ -358,16 +367,24 @@ class GitLabPromptManager(CustomPromptManagement):
|
||||
if prompt_metadata.get("model"):
|
||||
litellm_params["model"] = prompt_metadata["model"]
|
||||
|
||||
for param in ["temperature", "max_tokens", "top_p", "frequency_penalty", "presence_penalty"]:
|
||||
for param in [
|
||||
"temperature",
|
||||
"max_tokens",
|
||||
"top_p",
|
||||
"frequency_penalty",
|
||||
"presence_penalty",
|
||||
]:
|
||||
if param in prompt_metadata:
|
||||
litellm_params[param] = prompt_metadata[param]
|
||||
|
||||
return final_messages, litellm_params
|
||||
except Exception as e:
|
||||
import litellm
|
||||
litellm._logging.verbose_proxy_logger.error(f"Error in GitLab prompt pre_call_hook: {e}")
|
||||
return messages, litellm_params
|
||||
|
||||
litellm._logging.verbose_proxy_logger.error(
|
||||
f"Error in GitLab prompt pre_call_hook: {e}"
|
||||
)
|
||||
return messages, litellm_params
|
||||
|
||||
def _parse_prompt_to_messages(self, prompt_content: str) -> List[AllMessageValues]:
|
||||
messages: List[AllMessageValues] = []
|
||||
@@ -405,15 +422,15 @@ class GitLabPromptManager(CustomPromptManagement):
|
||||
return messages
|
||||
|
||||
def post_call_hook(
|
||||
self,
|
||||
user_id: Optional[str],
|
||||
response: Any,
|
||||
input_messages: List[AllMessageValues],
|
||||
function_call: Optional[Union[Dict[str, Any], str]] = None,
|
||||
litellm_params: Optional[Dict[str, Any]] = None,
|
||||
prompt_id: Optional[str] = None,
|
||||
prompt_variables: Optional[Dict[str, Any]] = None,
|
||||
**kwargs,
|
||||
self,
|
||||
user_id: Optional[str],
|
||||
response: Any,
|
||||
input_messages: List[AllMessageValues],
|
||||
function_call: Optional[Union[Dict[str, Any], str]] = None,
|
||||
litellm_params: Optional[Dict[str, Any]] = None,
|
||||
prompt_id: Optional[str] = None,
|
||||
prompt_variables: Optional[Dict[str, Any]] = None,
|
||||
**kwargs,
|
||||
) -> Any:
|
||||
return response
|
||||
|
||||
@@ -436,27 +453,30 @@ class GitLabPromptManager(CustomPromptManagement):
|
||||
_ = self.prompt_manager # trigger re-init/load
|
||||
|
||||
def should_run_prompt_management(
|
||||
self,
|
||||
prompt_id: str,
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
self,
|
||||
prompt_id: str,
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
return True
|
||||
|
||||
def _compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> PromptManagementClient:
|
||||
try:
|
||||
decoded_id = decode_prompt_id(prompt_id)
|
||||
if decoded_id not in self.prompt_manager.prompts:
|
||||
git_ref = getattr(dynamic_callback_params, "extra", {}).get("git_ref") if hasattr(dynamic_callback_params, "extra") else None
|
||||
git_ref = (
|
||||
getattr(dynamic_callback_params, "extra", {}).get("git_ref")
|
||||
if hasattr(dynamic_callback_params, "extra")
|
||||
else None
|
||||
)
|
||||
self.prompt_manager._load_prompt_from_gitlab(decoded_id, ref=git_ref)
|
||||
|
||||
|
||||
rendered_prompt, prompt_metadata = self.get_prompt_template(
|
||||
prompt_id, prompt_variables
|
||||
)
|
||||
@@ -465,7 +485,13 @@ class GitLabPromptManager(CustomPromptManagement):
|
||||
template_model = prompt_metadata.get("model")
|
||||
|
||||
optional_params: Dict[str, Any] = {}
|
||||
for param in ["temperature", "max_tokens", "top_p", "frequency_penalty", "presence_penalty"]:
|
||||
for param in [
|
||||
"temperature",
|
||||
"max_tokens",
|
||||
"top_p",
|
||||
"frequency_penalty",
|
||||
"presence_penalty",
|
||||
]:
|
||||
if param in prompt_metadata:
|
||||
optional_params[param] = prompt_metadata[param]
|
||||
|
||||
@@ -480,15 +506,17 @@ class GitLabPromptManager(CustomPromptManagement):
|
||||
raise ValueError(f"Error compiling prompt '{prompt_id}': {e}")
|
||||
|
||||
def get_chat_completion_prompt(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
non_default_params: dict,
|
||||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
self,
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
non_default_params: dict,
|
||||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
ignore_prompt_manager_optional_params: Optional[bool] = False,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
return PromptManagementBase.get_chat_completion_prompt(
|
||||
self,
|
||||
@@ -537,11 +565,11 @@ class GitLabPromptCache:
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
gitlab_config: Dict[str, Any],
|
||||
*,
|
||||
ref: Optional[str] = None,
|
||||
gitlab_client: Optional[GitLabClient] = None,
|
||||
self,
|
||||
gitlab_config: Dict[str, Any],
|
||||
*,
|
||||
ref: Optional[str] = None,
|
||||
gitlab_client: Optional[GitLabClient] = None,
|
||||
) -> None:
|
||||
# Build a PromptManager (which internally builds TemplateManager + Client)
|
||||
self.prompt_manager = GitLabPromptManager(
|
||||
@@ -550,7 +578,9 @@ class GitLabPromptCache:
|
||||
ref=ref,
|
||||
gitlab_client=gitlab_client,
|
||||
)
|
||||
self.template_manager: GitLabTemplateManager = self.prompt_manager.prompt_manager
|
||||
self.template_manager: GitLabTemplateManager = (
|
||||
self.prompt_manager.prompt_manager
|
||||
)
|
||||
|
||||
# In-memory stores
|
||||
self._by_file: Dict[str, Dict[str, Any]] = {}
|
||||
@@ -565,7 +595,9 @@ class GitLabPromptCache:
|
||||
Scan GitLab for all .prompt files under prompts_path, load and parse each,
|
||||
and return the mapping of repo file path -> JSON-like dict.
|
||||
"""
|
||||
ids = self.template_manager.list_templates(recursive=recursive) # IDs relative to prompts_path
|
||||
ids = self.template_manager.list_templates(
|
||||
recursive=recursive
|
||||
) # IDs relative to prompts_path
|
||||
for pid in ids:
|
||||
# Ensure template is loaded into TemplateManager
|
||||
if pid not in self.template_manager.prompts:
|
||||
@@ -579,7 +611,9 @@ class GitLabPromptCache:
|
||||
if tmpl is None:
|
||||
continue
|
||||
|
||||
file_path = self.template_manager._id_to_repo_path(pid) # "prompts/chat/..../file.prompt"
|
||||
file_path = self.template_manager._id_to_repo_path(
|
||||
pid
|
||||
) # "prompts/chat/..../file.prompt"
|
||||
entry = self._template_to_json(pid, tmpl)
|
||||
|
||||
self._by_file[file_path] = entry
|
||||
@@ -623,7 +657,9 @@ class GitLabPromptCache:
|
||||
# Internals
|
||||
# -------------------------
|
||||
|
||||
def _template_to_json(self, prompt_id: str, tmpl: GitLabPromptTemplate) -> Dict[str, Any]:
|
||||
def _template_to_json(
|
||||
self, prompt_id: str, tmpl: GitLabPromptTemplate
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Normalize a GitLabPromptTemplate into a JSON-like dict that is easy to serialize.
|
||||
"""
|
||||
@@ -637,12 +673,14 @@ class GitLabPromptCache:
|
||||
optional_params = dict(tmpl.optional_params or {})
|
||||
|
||||
return {
|
||||
"id": prompt_id, # e.g. "greet/hi"
|
||||
"path": self.template_manager._id_to_repo_path(prompt_id), # e.g. "prompts/chat/greet/hi.prompt"
|
||||
"content": tmpl.content, # rendered content (without frontmatter)
|
||||
"metadata": md, # parsed frontmatter
|
||||
"id": prompt_id, # e.g. "greet/hi"
|
||||
"path": self.template_manager._id_to_repo_path(
|
||||
prompt_id
|
||||
), # e.g. "prompts/chat/greet/hi.prompt"
|
||||
"content": tmpl.content, # rendered content (without frontmatter)
|
||||
"metadata": md, # parsed frontmatter
|
||||
"model": model,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
"optional_params": optional_params,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -158,6 +158,8 @@ class HumanloopLogger(CustomLogger):
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
ignore_prompt_manager_optional_params: Optional[bool] = False,
|
||||
) -> Tuple[
|
||||
str,
|
||||
List[AllMessageValues],
|
||||
|
||||
@@ -93,6 +93,8 @@ class PromptManagementBase(ABC):
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
ignore_prompt_manager_optional_params: Optional[bool] = False,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
|
||||
if prompt_id is None:
|
||||
@@ -117,13 +119,20 @@ class PromptManagementBase(ABC):
|
||||
prompt_template["prompt_template_optional_params"] or {}
|
||||
)
|
||||
|
||||
updated_non_default_params = {
|
||||
**non_default_params,
|
||||
**prompt_template_optional_params,
|
||||
}
|
||||
if not ignore_prompt_manager_optional_params:
|
||||
updated_non_default_params = {
|
||||
**non_default_params,
|
||||
**prompt_template_optional_params,
|
||||
}
|
||||
else:
|
||||
updated_non_default_params = non_default_params
|
||||
|
||||
if not ignore_prompt_manager_model:
|
||||
model = self._get_model_from_prompt(
|
||||
prompt_management_client=prompt_template, model=model
|
||||
)
|
||||
else:
|
||||
model = model
|
||||
|
||||
model = self._get_model_from_prompt(
|
||||
prompt_management_client=prompt_template, model=model
|
||||
)
|
||||
|
||||
return model, completed_messages, updated_non_default_params
|
||||
|
||||
@@ -3,7 +3,17 @@ import os
|
||||
import ssl
|
||||
import sys
|
||||
import time
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional, Tuple, Union
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Callable,
|
||||
Dict,
|
||||
List,
|
||||
Mapping,
|
||||
Optional,
|
||||
Tuple,
|
||||
Union,
|
||||
)
|
||||
|
||||
import certifi
|
||||
import httpx
|
||||
@@ -54,28 +64,28 @@ def _prepare_request_data_and_content(
|
||||
) -> Tuple[Optional[Union[dict, Mapping]], Any]:
|
||||
"""
|
||||
Helper function to route data/content parameters correctly for httpx requests
|
||||
|
||||
|
||||
This prevents httpx DeprecationWarnings that cause memory leaks.
|
||||
|
||||
|
||||
Background:
|
||||
- httpx shows a DeprecationWarning when you pass bytes/str to `data=`
|
||||
- It wants you to use `content=` instead for bytes/str
|
||||
- The warning itself leaks memory when triggered repeatedly
|
||||
|
||||
|
||||
Solution:
|
||||
- Move bytes/str from `data=` to `content=` before calling build_request
|
||||
- Keep dicts in `data=` (that's still the correct parameter for dicts)
|
||||
|
||||
|
||||
Args:
|
||||
data: Request data (can be dict, str, or bytes)
|
||||
content: Request content (raw bytes/str)
|
||||
|
||||
|
||||
Returns:
|
||||
Tuple of (request_data, request_content) properly routed for httpx
|
||||
"""
|
||||
request_data = None
|
||||
request_content = content
|
||||
|
||||
|
||||
if data is not None:
|
||||
if isinstance(data, (bytes, str)):
|
||||
# Bytes/strings belong in content= (only if not already provided)
|
||||
@@ -84,14 +94,16 @@ def _prepare_request_data_and_content(
|
||||
else:
|
||||
# dict/Mapping stays in data= parameter
|
||||
request_data = data
|
||||
|
||||
|
||||
return request_data, request_content
|
||||
|
||||
|
||||
# Cache for SSL contexts to avoid creating duplicate contexts with the same configuration
|
||||
# Key: tuple of (cafile, ssl_security_level, ssl_ecdh_curve)
|
||||
# Value: ssl.SSLContext
|
||||
_ssl_context_cache: Dict[Tuple[Optional[str], Optional[str], Optional[str]], ssl.SSLContext] = {}
|
||||
_ssl_context_cache: Dict[
|
||||
Tuple[Optional[str], Optional[str], Optional[str]], ssl.SSLContext
|
||||
] = {}
|
||||
|
||||
|
||||
def _create_ssl_context(
|
||||
@@ -201,7 +213,7 @@ def get_ssl_configuration(
|
||||
if ssl_verify is not False:
|
||||
# Create cache key from configuration parameters
|
||||
cache_key = (cafile, ssl_security_level, ssl_ecdh_curve)
|
||||
|
||||
|
||||
# Check if we have a cached SSL context for this configuration
|
||||
if cache_key not in _ssl_context_cache:
|
||||
_ssl_context_cache[cache_key] = _create_ssl_context(
|
||||
@@ -209,7 +221,7 @@ def get_ssl_configuration(
|
||||
ssl_security_level=ssl_security_level,
|
||||
ssl_ecdh_curve=ssl_ecdh_curve,
|
||||
)
|
||||
|
||||
|
||||
# Return the cached SSL context
|
||||
return _ssl_context_cache[cache_key]
|
||||
|
||||
@@ -389,8 +401,10 @@ class AsyncHTTPHandler:
|
||||
timeout = self.timeout
|
||||
|
||||
# Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix)
|
||||
request_data, request_content = _prepare_request_data_and_content(data, content)
|
||||
|
||||
request_data, request_content = _prepare_request_data_and_content(
|
||||
data, content
|
||||
)
|
||||
|
||||
req = self.client.build_request(
|
||||
"POST",
|
||||
url,
|
||||
@@ -401,7 +415,7 @@ class AsyncHTTPHandler:
|
||||
timeout=timeout,
|
||||
files=files,
|
||||
content=request_content,
|
||||
)
|
||||
)
|
||||
response = await self.client.send(req, stream=stream)
|
||||
response.raise_for_status()
|
||||
return response
|
||||
@@ -467,7 +481,9 @@ class AsyncHTTPHandler:
|
||||
timeout = self.timeout
|
||||
|
||||
# Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix)
|
||||
request_data, request_content = _prepare_request_data_and_content(data, content)
|
||||
request_data, request_content = _prepare_request_data_and_content(
|
||||
data, content
|
||||
)
|
||||
|
||||
req = self.client.build_request(
|
||||
"PUT", url, data=request_data, json=json, params=params, headers=headers, timeout=timeout, content=request_content # type: ignore
|
||||
@@ -531,7 +547,9 @@ class AsyncHTTPHandler:
|
||||
timeout = self.timeout
|
||||
|
||||
# Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix)
|
||||
request_data, request_content = _prepare_request_data_and_content(data, content)
|
||||
request_data, request_content = _prepare_request_data_and_content(
|
||||
data, content
|
||||
)
|
||||
|
||||
req = self.client.build_request(
|
||||
"PATCH", url, data=request_data, json=json, params=params, headers=headers, timeout=timeout, content=request_content # type: ignore
|
||||
@@ -593,10 +611,12 @@ class AsyncHTTPHandler:
|
||||
try:
|
||||
if timeout is None:
|
||||
timeout = self.timeout
|
||||
|
||||
|
||||
# Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix)
|
||||
request_data, request_content = _prepare_request_data_and_content(data, content)
|
||||
|
||||
request_data, request_content = _prepare_request_data_and_content(
|
||||
data, content
|
||||
)
|
||||
|
||||
req = self.client.build_request(
|
||||
"DELETE", url, data=request_data, json=json, params=params, headers=headers, timeout=timeout, content=request_content # type: ignore
|
||||
)
|
||||
@@ -648,7 +668,7 @@ class AsyncHTTPHandler:
|
||||
"""
|
||||
# Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix)
|
||||
request_data, request_content = _prepare_request_data_and_content(data, content)
|
||||
|
||||
|
||||
req = client.build_request(
|
||||
"POST", url, data=request_data, json=json, params=params, headers=headers, content=request_content # type: ignore
|
||||
)
|
||||
@@ -802,8 +822,10 @@ class AsyncHTTPHandler:
|
||||
if AIOHTTP_CONNECTOR_LIMIT > 0:
|
||||
transport_connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT
|
||||
if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0:
|
||||
transport_connector_kwargs["limit_per_host"] = AIOHTTP_CONNECTOR_LIMIT_PER_HOST
|
||||
|
||||
transport_connector_kwargs["limit_per_host"] = (
|
||||
AIOHTTP_CONNECTOR_LIMIT_PER_HOST
|
||||
)
|
||||
|
||||
return LiteLLMAiohttpTransport(
|
||||
client=lambda: ClientSession(
|
||||
connector=TCPConnector(**transport_connector_kwargs),
|
||||
@@ -832,6 +854,9 @@ class HTTPHandler:
|
||||
concurrent_limit=None, # Kept for backward compatibility, but ignored (no limits)
|
||||
client: Optional[httpx.Client] = None,
|
||||
ssl_verify: Optional[Union[bool, str]] = None,
|
||||
disable_default_headers: Optional[
|
||||
bool
|
||||
] = False, # arize phoenix returns different API responses when user agent header in request
|
||||
):
|
||||
if timeout is None:
|
||||
timeout = _DEFAULT_TIMEOUT
|
||||
@@ -852,7 +877,7 @@ class HTTPHandler:
|
||||
timeout=timeout,
|
||||
verify=ssl_config,
|
||||
cert=cert,
|
||||
headers=headers,
|
||||
headers=headers if not disable_default_headers else None,
|
||||
follow_redirects=True,
|
||||
)
|
||||
else:
|
||||
@@ -877,7 +902,9 @@ class HTTPHandler:
|
||||
params.update(self.extract_query_params(url))
|
||||
|
||||
response = self.client.get(
|
||||
url, params=params, headers=headers, follow_redirects=_follow_redirects # type: ignore
|
||||
url,
|
||||
params=params,
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
return response
|
||||
@@ -910,8 +937,10 @@ class HTTPHandler:
|
||||
):
|
||||
try:
|
||||
# Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix)
|
||||
request_data, request_content = _prepare_request_data_and_content(data, content)
|
||||
|
||||
request_data, request_content = _prepare_request_data_and_content(
|
||||
data, content
|
||||
)
|
||||
|
||||
if timeout is not None:
|
||||
req = self.client.build_request(
|
||||
"POST",
|
||||
@@ -964,8 +993,10 @@ class HTTPHandler:
|
||||
):
|
||||
try:
|
||||
# Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix)
|
||||
request_data, request_content = _prepare_request_data_and_content(data, content)
|
||||
|
||||
request_data, request_content = _prepare_request_data_and_content(
|
||||
data, content
|
||||
)
|
||||
|
||||
if timeout is not None:
|
||||
req = self.client.build_request(
|
||||
"PATCH", url, data=request_data, json=json, params=params, headers=headers, timeout=timeout, content=request_content # type: ignore
|
||||
@@ -1011,8 +1042,10 @@ class HTTPHandler:
|
||||
):
|
||||
try:
|
||||
# Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix)
|
||||
request_data, request_content = _prepare_request_data_and_content(data, content)
|
||||
|
||||
request_data, request_content = _prepare_request_data_and_content(
|
||||
data, content
|
||||
)
|
||||
|
||||
if timeout is not None:
|
||||
req = self.client.build_request(
|
||||
"PUT", url, data=request_data, json=json, params=params, headers=headers, timeout=timeout, content=request_content # type: ignore
|
||||
@@ -1045,8 +1078,10 @@ class HTTPHandler:
|
||||
):
|
||||
try:
|
||||
# Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix)
|
||||
request_data, request_content = _prepare_request_data_and_content(data, content)
|
||||
|
||||
request_data, request_content = _prepare_request_data_and_content(
|
||||
data, content
|
||||
)
|
||||
|
||||
if timeout is not None:
|
||||
req = self.client.build_request(
|
||||
"DELETE", url, data=request_data, json=json, params=params, headers=headers, timeout=timeout, content=request_content # type: ignore
|
||||
|
||||
@@ -10,13 +10,11 @@ model_list:
|
||||
litellm_params:
|
||||
model: openai/gpt-4.1-mini
|
||||
|
||||
|
||||
guardrails:
|
||||
- guardrail_name: generic-guardrail
|
||||
prompts:
|
||||
- prompt_id: "simple_prompt"
|
||||
litellm_params:
|
||||
guardrail: generic_guardrail_api
|
||||
mode: ["pre_call"]
|
||||
headers:
|
||||
Authorization: Bearer mock-bedrock-token-12345
|
||||
api_base: http://localhost:8080
|
||||
default_on: true
|
||||
prompt_id: "UHJvbXB0VmVyc2lvbjox"
|
||||
prompt_integration: "arize_phoenix"
|
||||
api_base: https://app.phoenix.arize.com/s/krrishdholakia
|
||||
ignore_prompt_manager_model: true # ignores model from prompt manager
|
||||
ignore_prompt_manager_optional_params: true # ignores optional params from prompt manager - e.g. temperature, max_tokens, etc.
|
||||
@@ -17,6 +17,8 @@ class X42PromptManagement(CustomPromptManagement):
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
ignore_prompt_manager_optional_params: Optional[bool] = False,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
"""
|
||||
Returns:
|
||||
|
||||
+30
-18
@@ -66,7 +66,6 @@ from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_a
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.llms.custom_httpx.httpx_handler import HTTPHandler
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.proxy._types import (
|
||||
AlertType,
|
||||
@@ -835,12 +834,12 @@ class ProxyLogging:
|
||||
call_type: CallTypesLiteral,
|
||||
) -> None:
|
||||
"""Process prompt template if applicable."""
|
||||
from litellm.utils import get_non_default_completion_params
|
||||
from litellm.proxy.prompts.prompt_endpoints import (
|
||||
construct_versioned_prompt_id,
|
||||
get_latest_version_prompt_id,
|
||||
)
|
||||
from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY
|
||||
from litellm.utils import get_non_default_completion_params
|
||||
|
||||
if prompt_version is None:
|
||||
lookup_prompt_id = get_latest_version_prompt_id(
|
||||
@@ -879,6 +878,11 @@ class ProxyLogging:
|
||||
data.update(optional_params)
|
||||
data["model"] = model
|
||||
data["messages"] = messages
|
||||
# prevent re-processing the prompt template
|
||||
data.pop("prompt_id", None)
|
||||
data.pop("prompt_variables", None)
|
||||
data.pop("prompt_label", None)
|
||||
data.pop("prompt_version", None)
|
||||
|
||||
def _process_guardrail_metadata(self, data: dict) -> None:
|
||||
"""Process guardrails from metadata and add to applied_guardrails."""
|
||||
@@ -967,11 +971,13 @@ class ProxyLogging:
|
||||
prompt_version = data.get("prompt_version", None)
|
||||
|
||||
## PROMPT TEMPLATE CHECK ##
|
||||
|
||||
if (
|
||||
litellm_logging_obj is not None
|
||||
and prompt_id is not None
|
||||
and (call_type == "completion" or call_type == "acompletion")
|
||||
):
|
||||
|
||||
self._process_prompt_template(
|
||||
data=data,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
@@ -997,7 +1003,7 @@ class ProxyLogging:
|
||||
):
|
||||
result = await self._process_guardrail_callback(
|
||||
callback=_callback,
|
||||
data=data, # type: ignore
|
||||
data=data, # type: ignore
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_type=call_type,
|
||||
)
|
||||
@@ -3477,17 +3483,17 @@ async def update_spend_logs_job(
|
||||
):
|
||||
"""
|
||||
Job to process spend_log_transactions queue.
|
||||
|
||||
|
||||
This job is triggered based on queue size rather than time.
|
||||
Processes spend log transactions when the queue reaches a threshold.
|
||||
"""
|
||||
n_retry_times = 3
|
||||
|
||||
|
||||
queue_size = len(prisma_client.spend_log_transactions)
|
||||
|
||||
|
||||
if queue_size == 0:
|
||||
return
|
||||
|
||||
|
||||
await ProxyUpdateSpend.update_spend_logs(
|
||||
n_retry_times=n_retry_times,
|
||||
prisma_client=prisma_client,
|
||||
@@ -3504,28 +3510,31 @@ async def _monitor_spend_logs_queue(
|
||||
"""
|
||||
Background task that monitors the spend_log_transactions queue size
|
||||
and triggers processing when the threshold is reached.
|
||||
|
||||
|
||||
Args:
|
||||
prisma_client: Prisma client instance
|
||||
db_writer_client: Optional HTTP handler for external spend logs endpoint
|
||||
proxy_logging_obj: Proxy logging object
|
||||
"""
|
||||
from litellm.constants import SPEND_LOG_QUEUE_SIZE_THRESHOLD, SPEND_LOG_QUEUE_POLL_INTERVAL
|
||||
|
||||
from litellm.constants import (
|
||||
SPEND_LOG_QUEUE_POLL_INTERVAL,
|
||||
SPEND_LOG_QUEUE_SIZE_THRESHOLD,
|
||||
)
|
||||
|
||||
threshold = SPEND_LOG_QUEUE_SIZE_THRESHOLD
|
||||
base_interval = SPEND_LOG_QUEUE_POLL_INTERVAL
|
||||
max_backoff = 30.0 # Maximum backoff interval in seconds
|
||||
backoff_multiplier = 1.5 # Exponential backoff multiplier
|
||||
current_interval = base_interval
|
||||
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
f"Starting spend logs queue monitor (threshold: {threshold}, poll_interval: {base_interval}s)"
|
||||
)
|
||||
|
||||
|
||||
while True:
|
||||
try:
|
||||
queue_size = len(prisma_client.spend_log_transactions)
|
||||
|
||||
|
||||
if queue_size > 0:
|
||||
if queue_size >= threshold:
|
||||
verbose_proxy_logger.debug(
|
||||
@@ -3538,8 +3547,10 @@ async def _monitor_spend_logs_queue(
|
||||
f"Spend logs queue size ({queue_size}) below threshold ({threshold}), processing with backoff"
|
||||
)
|
||||
# Exponential backoff when below threshold but still processing
|
||||
current_interval = min(current_interval * backoff_multiplier, max_backoff)
|
||||
|
||||
current_interval = min(
|
||||
current_interval * backoff_multiplier, max_backoff
|
||||
)
|
||||
|
||||
await update_spend_logs_job(
|
||||
prisma_client=prisma_client,
|
||||
db_writer_client=db_writer_client,
|
||||
@@ -3547,8 +3558,10 @@ async def _monitor_spend_logs_queue(
|
||||
)
|
||||
else:
|
||||
# Exponential backoff when no logs to process
|
||||
current_interval = min(current_interval * backoff_multiplier, max_backoff)
|
||||
|
||||
current_interval = min(
|
||||
current_interval * backoff_multiplier, max_backoff
|
||||
)
|
||||
|
||||
await asyncio.sleep(current_interval)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
@@ -3559,7 +3572,6 @@ async def _monitor_spend_logs_queue(
|
||||
await asyncio.sleep(current_interval)
|
||||
|
||||
|
||||
|
||||
def _raise_failed_update_spend_exception(
|
||||
e: Exception, start_time: float, proxy_logging_obj: ProxyLogging
|
||||
):
|
||||
|
||||
@@ -11,6 +11,7 @@ class SupportedPromptIntegrations(str, Enum):
|
||||
CUSTOM = "custom"
|
||||
BITBUCKET = "bitbucket"
|
||||
GITLAB = "gitlab"
|
||||
ARIZE_PHOENIX = "arize_phoenix"
|
||||
|
||||
|
||||
class PromptInfo(BaseModel):
|
||||
@@ -22,6 +23,8 @@ class PromptInfo(BaseModel):
|
||||
class PromptLiteLLMParams(BaseModel):
|
||||
prompt_id: str
|
||||
prompt_integration: str
|
||||
api_key: Optional[str] = None
|
||||
api_base: Optional[str] = None
|
||||
|
||||
dotprompt_content: Optional[str] = None
|
||||
"""
|
||||
|
||||
@@ -35,6 +35,8 @@ class TestCustomPromptManagement(CustomPromptManagement):
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
ignore_prompt_manager_optional_params: Optional[bool] = False,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
print(
|
||||
"TestCustomPromptManagement: running get_chat_completion_prompt for prompt_id: ",
|
||||
|
||||
Reference in New Issue
Block a user