mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Prompt Management - new API for integrating providers (#17829)
* Prompt Management API - new API to interact with Prompt Management integrations (no PR required) (#17800) * feat: initial commit adding prompt management api * feat: initial commit adding prompt management api * fix: refactoring to make sure get prompt is async * fix: additional fixes * fix: partially working generic api prompt management
This commit is contained in:
parent
d693596e87
commit
7e58931ec1
32 changed files with 1656 additions and 67 deletions
|
|
@ -0,0 +1,279 @@
|
|||
# Braintrust Prompt Wrapper for LiteLLM
|
||||
|
||||
This directory contains a wrapper server that enables LiteLLM to use prompts from [Braintrust](https://www.braintrust.dev/) through the generic prompt management API.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌─────────────┐ ┌──────────────────────┐ ┌─────────────┐
|
||||
│ LiteLLM │ ──────> │ Wrapper Server │ ──────> │ Braintrust │
|
||||
│ Client │ │ (This Server) │ │ API │
|
||||
└─────────────┘ └──────────────────────┘ └─────────────┘
|
||||
Uses generic Transforms Stores actual
|
||||
prompt manager Braintrust format prompt templates
|
||||
to LiteLLM format
|
||||
```
|
||||
|
||||
## Components
|
||||
|
||||
### 1. Generic Prompt Manager (`litellm/integrations/generic_prompt_management/`)
|
||||
|
||||
A generic client that can work with any API implementing the `/beta/litellm_prompt_management` endpoint.
|
||||
|
||||
**Expected API Response Format:**
|
||||
```json
|
||||
{
|
||||
"prompt_id": "string",
|
||||
"prompt_template": [
|
||||
{"role": "system", "content": "You are a helpful assistant"},
|
||||
{"role": "user", "content": "Hello {name}"}
|
||||
],
|
||||
"prompt_template_model": "gpt-4",
|
||||
"prompt_template_optional_params": {
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 100
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### 2. Braintrust Wrapper Server (`braintrust_prompt_wrapper_server.py`)
|
||||
|
||||
A FastAPI server that:
|
||||
- Implements the `/beta/litellm_prompt_management` endpoint
|
||||
- Fetches prompts from Braintrust API
|
||||
- Transforms Braintrust response format to LiteLLM format
|
||||
|
||||
## Setup
|
||||
|
||||
### Install Dependencies
|
||||
|
||||
```bash
|
||||
pip install fastapi uvicorn httpx litellm
|
||||
```
|
||||
|
||||
### Set Environment Variables
|
||||
|
||||
```bash
|
||||
export BRAINTRUST_API_KEY="your-braintrust-api-key"
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Step 1: Start the Wrapper Server
|
||||
|
||||
```bash
|
||||
python braintrust_prompt_wrapper_server.py
|
||||
```
|
||||
|
||||
The server will start on `http://localhost:8080` by default.
|
||||
|
||||
You can customize the port and host:
|
||||
```bash
|
||||
export PORT=8000
|
||||
export HOST=0.0.0.0
|
||||
python braintrust_prompt_wrapper_server.py
|
||||
```
|
||||
|
||||
### Step 2: Use with LiteLLM
|
||||
|
||||
```python
|
||||
import litellm
|
||||
from litellm.integrations.generic_prompt_management import GenericPromptManager
|
||||
|
||||
# Configure the generic prompt manager to use your wrapper server
|
||||
generic_config = {
|
||||
"api_base": "http://localhost:8080",
|
||||
"api_key": "your-braintrust-api-key", # Will be passed to Braintrust
|
||||
"timeout": 30,
|
||||
}
|
||||
|
||||
# Create the prompt manager
|
||||
prompt_manager = GenericPromptManager(**generic_config)
|
||||
|
||||
# Use with completion
|
||||
response = litellm.completion(
|
||||
model="generic_prompt/gpt-4",
|
||||
prompt_id="your-braintrust-prompt-id",
|
||||
prompt_variables={"name": "World"}, # Variables to substitute
|
||||
messages=[{"role": "user", "content": "Additional message"}]
|
||||
)
|
||||
|
||||
print(response)
|
||||
```
|
||||
|
||||
### Step 3: Direct API Testing
|
||||
|
||||
You can also test the wrapper API directly:
|
||||
|
||||
```bash
|
||||
# Test with curl
|
||||
curl -H "Authorization: Bearer YOUR_BRAINTRUST_TOKEN" \
|
||||
"http://localhost:8080/beta/litellm_prompt_management?prompt_id=YOUR_PROMPT_ID"
|
||||
|
||||
# Health check
|
||||
curl http://localhost:8080/health
|
||||
|
||||
# Service info
|
||||
curl http://localhost:8080/
|
||||
```
|
||||
|
||||
## API Documentation
|
||||
|
||||
Once the server is running, visit:
|
||||
- Swagger UI: `http://localhost:8080/docs`
|
||||
- ReDoc: `http://localhost:8080/redoc`
|
||||
|
||||
## Braintrust Format Transformation
|
||||
|
||||
The wrapper automatically transforms Braintrust's response format:
|
||||
|
||||
**Braintrust API Response:**
|
||||
```json
|
||||
{
|
||||
"id": "prompt-123",
|
||||
"prompt_data": {
|
||||
"prompt": {
|
||||
"type": "chat",
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful assistant"
|
||||
}
|
||||
]
|
||||
},
|
||||
"options": {
|
||||
"model": "gpt-4",
|
||||
"params": {
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 100
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Transformed to LiteLLM Format:**
|
||||
```json
|
||||
{
|
||||
"prompt_id": "prompt-123",
|
||||
"prompt_template": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are a helpful assistant"
|
||||
}
|
||||
],
|
||||
"prompt_template_model": "gpt-4",
|
||||
"prompt_template_optional_params": {
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 100
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Supported Parameters
|
||||
|
||||
The wrapper automatically maps these Braintrust parameters to LiteLLM:
|
||||
|
||||
- `temperature`
|
||||
- `max_tokens` / `max_completion_tokens`
|
||||
- `top_p`
|
||||
- `frequency_penalty`
|
||||
- `presence_penalty`
|
||||
- `n`
|
||||
- `stop`
|
||||
- `response_format`
|
||||
- `tool_choice`
|
||||
- `function_call`
|
||||
- `tools`
|
||||
|
||||
## Variable Substitution
|
||||
|
||||
The generic prompt manager supports simple variable substitution:
|
||||
|
||||
```python
|
||||
# In your Braintrust prompt:
|
||||
# "Hello {name}, welcome to {place}!"
|
||||
|
||||
# In your code:
|
||||
prompt_variables = {
|
||||
"name": "Alice",
|
||||
"place": "Wonderland"
|
||||
}
|
||||
|
||||
# Result:
|
||||
# "Hello Alice, welcome to Wonderland!"
|
||||
```
|
||||
|
||||
Supports both `{variable}` and `{{variable}}` syntax.
|
||||
|
||||
## Error Handling
|
||||
|
||||
The wrapper provides detailed error messages:
|
||||
|
||||
- **401**: Missing or invalid Braintrust API token
|
||||
- **404**: Prompt not found in Braintrust
|
||||
- **502**: Failed to connect to Braintrust API
|
||||
- **500**: Error transforming response
|
||||
|
||||
## Production Deployment
|
||||
|
||||
For production use:
|
||||
|
||||
1. **Use HTTPS**: Deploy behind a reverse proxy with SSL
|
||||
2. **Authentication**: Add authentication to the wrapper endpoint if needed
|
||||
3. **Rate Limiting**: Implement rate limiting to prevent abuse
|
||||
4. **Caching**: Consider caching prompt responses
|
||||
5. **Monitoring**: Add logging and monitoring
|
||||
|
||||
Example with Docker:
|
||||
|
||||
```dockerfile
|
||||
FROM python:3.11-slim
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
RUN pip install fastapi uvicorn httpx
|
||||
|
||||
COPY braintrust_prompt_wrapper_server.py .
|
||||
|
||||
ENV PORT=8080
|
||||
ENV HOST=0.0.0.0
|
||||
|
||||
EXPOSE 8080
|
||||
|
||||
CMD ["python", "braintrust_prompt_wrapper_server.py"]
|
||||
```
|
||||
|
||||
## Extending to Other Providers
|
||||
|
||||
This pattern can be used with any prompt management provider:
|
||||
|
||||
1. Create a wrapper server that implements `/beta/litellm_prompt_management`
|
||||
2. Transform the provider's response to LiteLLM format
|
||||
3. Use the generic prompt manager to connect
|
||||
|
||||
Example providers:
|
||||
- Langsmith
|
||||
- PromptLayer
|
||||
- Humanloop
|
||||
- Custom internal systems
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### "No Braintrust API token provided"
|
||||
- Set `BRAINTRUST_API_KEY` environment variable
|
||||
- Or pass token in `Authorization: Bearer TOKEN` header
|
||||
|
||||
### "Failed to connect to Braintrust API"
|
||||
- Check your internet connection
|
||||
- Verify Braintrust API is accessible
|
||||
- Check firewall settings
|
||||
|
||||
### "Prompt not found"
|
||||
- Verify the prompt ID exists in Braintrust
|
||||
- Check that your API token has access to the prompt
|
||||
|
||||
## License
|
||||
|
||||
This wrapper is part of the LiteLLM project and follows the same license.
|
||||
|
||||
|
|
@ -0,0 +1,274 @@
|
|||
"""
|
||||
Mock server that implements the /beta/litellm_prompt_management endpoint
|
||||
and acts as a wrapper for calling the Braintrust API.
|
||||
|
||||
This server transforms Braintrust's prompt API response into the format
|
||||
expected by LiteLLM's generic prompt management client.
|
||||
|
||||
Usage:
|
||||
python braintrust_prompt_wrapper_server.py
|
||||
|
||||
# Then test with:
|
||||
curl -H "Authorization: Bearer YOUR_BRAINTRUST_TOKEN" \
|
||||
"http://localhost:8080/beta/litellm_prompt_management?prompt_id=YOUR_PROMPT_ID"
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import httpx
|
||||
from fastapi import FastAPI, HTTPException, Header, Query
|
||||
from fastapi.responses import JSONResponse
|
||||
import uvicorn
|
||||
|
||||
|
||||
app = FastAPI(
|
||||
title="Braintrust Prompt Wrapper",
|
||||
description="Wrapper server for Braintrust prompts to work with LiteLLM",
|
||||
version="1.0.0",
|
||||
)
|
||||
|
||||
|
||||
def transform_braintrust_message(message: Dict[str, Any]) -> Dict[str, str]:
|
||||
"""
|
||||
Transform a Braintrust message to LiteLLM format.
|
||||
|
||||
Braintrust message format:
|
||||
{
|
||||
"role": "system",
|
||||
"content": "...",
|
||||
"name": "..." (optional)
|
||||
}
|
||||
|
||||
LiteLLM format:
|
||||
{
|
||||
"role": "system",
|
||||
"content": "..."
|
||||
}
|
||||
"""
|
||||
result = {
|
||||
"role": message.get("role", "user"),
|
||||
"content": message.get("content", ""),
|
||||
}
|
||||
|
||||
# Include name if present
|
||||
if "name" in message:
|
||||
result["name"] = message["name"]
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def transform_braintrust_response(
|
||||
braintrust_response: Dict[str, Any],
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Transform Braintrust API response to LiteLLM prompt management format.
|
||||
|
||||
Braintrust response format:
|
||||
{
|
||||
"objects": [{
|
||||
"id": "prompt_id",
|
||||
"prompt_data": {
|
||||
"prompt": {
|
||||
"type": "chat",
|
||||
"messages": [...],
|
||||
"tools": "..."
|
||||
},
|
||||
"options": {
|
||||
"model": "gpt-4",
|
||||
"params": {
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 100,
|
||||
...
|
||||
}
|
||||
}
|
||||
}
|
||||
}]
|
||||
}
|
||||
|
||||
LiteLLM format:
|
||||
{
|
||||
"prompt_id": "prompt_id",
|
||||
"prompt_template": [...],
|
||||
"prompt_template_model": "gpt-4",
|
||||
"prompt_template_optional_params": {...}
|
||||
}
|
||||
"""
|
||||
# Extract the first object from the objects array if it exists
|
||||
if "objects" in braintrust_response and len(braintrust_response["objects"]) > 0:
|
||||
prompt_object = braintrust_response["objects"][0]
|
||||
else:
|
||||
prompt_object = braintrust_response
|
||||
|
||||
prompt_data = prompt_object.get("prompt_data", {})
|
||||
prompt_info = prompt_data.get("prompt", {})
|
||||
options = prompt_data.get("options", {})
|
||||
|
||||
# Extract messages
|
||||
messages = prompt_info.get("messages", [])
|
||||
transformed_messages = [transform_braintrust_message(msg) for msg in messages]
|
||||
|
||||
# Extract model
|
||||
model = options.get("model")
|
||||
|
||||
# Extract optional parameters
|
||||
params = options.get("params", {})
|
||||
optional_params: Dict[str, Any] = {}
|
||||
|
||||
# Map common parameters
|
||||
param_mapping = {
|
||||
"temperature": "temperature",
|
||||
"max_tokens": "max_tokens",
|
||||
"max_completion_tokens": "max_tokens", # Alternative name
|
||||
"top_p": "top_p",
|
||||
"frequency_penalty": "frequency_penalty",
|
||||
"presence_penalty": "presence_penalty",
|
||||
"n": "n",
|
||||
"stop": "stop",
|
||||
}
|
||||
|
||||
for braintrust_param, litellm_param in param_mapping.items():
|
||||
if braintrust_param in params:
|
||||
value = params[braintrust_param]
|
||||
if value is not None:
|
||||
optional_params[litellm_param] = value
|
||||
|
||||
# Handle response_format
|
||||
if "response_format" in params:
|
||||
optional_params["response_format"] = params["response_format"]
|
||||
|
||||
# Handle tool_choice
|
||||
if "tool_choice" in params:
|
||||
optional_params["tool_choice"] = params["tool_choice"]
|
||||
|
||||
# Handle function_call
|
||||
if "function_call" in params:
|
||||
optional_params["function_call"] = params["function_call"]
|
||||
|
||||
# Add tools if present
|
||||
if "tools" in prompt_info and prompt_info["tools"]:
|
||||
optional_params["tools"] = prompt_info["tools"]
|
||||
|
||||
# Handle tool_functions from prompt_data
|
||||
if "tool_functions" in prompt_data and prompt_data["tool_functions"]:
|
||||
optional_params["tool_functions"] = prompt_data["tool_functions"]
|
||||
|
||||
return {
|
||||
"prompt_id": prompt_object.get("id"),
|
||||
"prompt_template": transformed_messages,
|
||||
"prompt_template_model": model,
|
||||
"prompt_template_optional_params": optional_params if optional_params else None,
|
||||
}
|
||||
|
||||
|
||||
@app.get("/beta/litellm_prompt_management")
|
||||
async def get_prompt(
|
||||
prompt_id: str = Query(..., description="The Braintrust prompt ID to fetch"),
|
||||
authorization: Optional[str] = Header(
|
||||
None, description="Bearer token for Braintrust API"
|
||||
),
|
||||
) -> JSONResponse:
|
||||
"""
|
||||
Fetch a prompt from Braintrust and transform it to LiteLLM format.
|
||||
|
||||
Args:
|
||||
prompt_id: The Braintrust prompt ID
|
||||
authorization: Bearer token for Braintrust API (from header)
|
||||
|
||||
Returns:
|
||||
JSONResponse with the transformed prompt data
|
||||
"""
|
||||
# Extract token from Authorization header or environment
|
||||
braintrust_token = None
|
||||
if authorization and authorization.startswith("Bearer "):
|
||||
braintrust_token = authorization.replace("Bearer ", "")
|
||||
else:
|
||||
braintrust_token = os.getenv("BRAINTRUST_API_KEY")
|
||||
|
||||
if not braintrust_token:
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="No Braintrust API token provided. Pass via Authorization header or set BRAINTRUST_API_KEY environment variable.",
|
||||
)
|
||||
|
||||
# Call Braintrust API
|
||||
braintrust_url = f"https://api.braintrust.dev/v1/prompt/{prompt_id}"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {braintrust_token}",
|
||||
"Accept": "application/json",
|
||||
}
|
||||
print(f"headers: {headers}")
|
||||
print(f"braintrust_url: {braintrust_url}")
|
||||
print(f"braintrust_token: {braintrust_token}")
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.get(braintrust_url, headers=headers)
|
||||
response.raise_for_status()
|
||||
braintrust_data = response.json()
|
||||
except httpx.HTTPStatusError as e:
|
||||
raise HTTPException(
|
||||
status_code=e.response.status_code,
|
||||
detail=f"Braintrust API error: {e.response.text}",
|
||||
)
|
||||
except httpx.RequestError as e:
|
||||
raise HTTPException(
|
||||
status_code=502,
|
||||
detail=f"Failed to connect to Braintrust API: {str(e)}",
|
||||
)
|
||||
except json.JSONDecodeError as e:
|
||||
raise HTTPException(
|
||||
status_code=502,
|
||||
detail=f"Failed to parse Braintrust API response: {str(e)}",
|
||||
)
|
||||
|
||||
print(f"braintrust_data: {braintrust_data}")
|
||||
# Transform the response
|
||||
try:
|
||||
transformed_data = transform_braintrust_response(braintrust_data)
|
||||
print(f"transformed_data: {transformed_data}")
|
||||
return JSONResponse(content=transformed_data)
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to transform Braintrust response: {str(e)}",
|
||||
)
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health_check():
|
||||
"""Health check endpoint."""
|
||||
return {"status": "healthy", "service": "braintrust-prompt-wrapper"}
|
||||
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
"""Root endpoint with service information."""
|
||||
return {
|
||||
"service": "Braintrust Prompt Wrapper for LiteLLM",
|
||||
"version": "1.0.0",
|
||||
"endpoints": {
|
||||
"prompt_management": "/beta/litellm_prompt_management?prompt_id=<id>",
|
||||
"health": "/health",
|
||||
},
|
||||
"documentation": "/docs",
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
"""Run the server."""
|
||||
port = int(os.getenv("PORT", "8080"))
|
||||
host = os.getenv("HOST", "0.0.0.0")
|
||||
|
||||
print(f"🚀 Starting Braintrust Prompt Wrapper Server on {host}:{port}")
|
||||
print(f"📚 API Documentation available at http://{host}:{port}/docs")
|
||||
print(
|
||||
f"🔑 Make sure to set BRAINTRUST_API_KEY environment variable or pass token in Authorization header"
|
||||
)
|
||||
|
||||
uvicorn.run(app, host=host, port=port)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -7,16 +7,18 @@ Users can define
|
|||
"""
|
||||
|
||||
import copy
|
||||
from typing import Dict, List, Optional, Tuple, Union, cast
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union, cast
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.custom_prompt_management import CustomPromptManagement
|
||||
from litellm.integrations.prompt_management_base import PromptManagementClient
|
||||
from litellm.types.integrations.anthropic_cache_control_hook import (
|
||||
CacheControlInjectionPoint,
|
||||
CacheControlMessageInjectionPoint,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionCachedContent
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
|
||||
|
|
@ -29,6 +31,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
|
|
@ -141,6 +144,78 @@ class AnthropicCacheControlHook(CustomPromptManagement):
|
|||
"""Return the integration name for this hook."""
|
||||
return "anthropic_cache_control_hook"
|
||||
|
||||
def should_run_prompt_management(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
"""Always return False since this is not a true prompt management system."""
|
||||
return False
|
||||
|
||||
def _compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> PromptManagementClient:
|
||||
"""Not used - this hook only modifies messages, doesn't fetch prompts."""
|
||||
return PromptManagementClient(
|
||||
prompt_id=prompt_id,
|
||||
prompt_template=[],
|
||||
prompt_template_model=None,
|
||||
prompt_template_optional_params=None,
|
||||
completed_messages=None,
|
||||
)
|
||||
|
||||
async def async_compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> PromptManagementClient:
|
||||
"""Not used - this hook only modifies messages, doesn't fetch prompts."""
|
||||
return self._compile_prompt_helper(
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_variables=prompt_variables,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
async def async_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,
|
||||
litellm_logging_obj: Any,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
tools: Optional[List[Dict]] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
"""Async version - delegates to sync since no async operations needed."""
|
||||
return self.get_chat_completion_prompt(
|
||||
model=model,
|
||||
messages=messages,
|
||||
non_default_params=non_default_params,
|
||||
prompt_id=prompt_id,
|
||||
prompt_variables=prompt_variables,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def should_use_anthropic_cache_control_hook(non_default_params: Dict) -> bool:
|
||||
if non_default_params.get("cache_control_injection_points", None):
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from litellm.integrations.prompt_management_base import (
|
|||
PromptManagementClient,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
from .bitbucket_client import BitBucketClient
|
||||
|
|
@ -414,7 +415,8 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
|
||||
def should_run_prompt_management(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
"""
|
||||
|
|
@ -423,11 +425,12 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
For BitBucket, we always return True and handle the prompt loading
|
||||
in the _compile_prompt_helper method.
|
||||
"""
|
||||
return True
|
||||
return prompt_id is not None
|
||||
|
||||
def _compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
|
|
@ -442,6 +445,9 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
3. Converts the rendered text into chat messages
|
||||
4. Extracts model and optional parameters from metadata
|
||||
"""
|
||||
if prompt_id is None:
|
||||
raise ValueError("prompt_id is required for BitBucket prompt manager")
|
||||
|
||||
try:
|
||||
# Load the prompt from BitBucket if not already loaded
|
||||
if prompt_id not in self.prompt_manager.prompts:
|
||||
|
|
@ -481,6 +487,31 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
except Exception as e:
|
||||
raise ValueError(f"Error compiling prompt '{prompt_id}': {e}")
|
||||
|
||||
async def async_compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> PromptManagementClient:
|
||||
"""
|
||||
Async version of compile prompt helper. Since BitBucket operations use sync client,
|
||||
this simply delegates to the sync version.
|
||||
"""
|
||||
if prompt_id is None:
|
||||
raise ValueError("prompt_id is required for BitBucket prompt manager")
|
||||
|
||||
return self._compile_prompt_helper(
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_variables=prompt_variables,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
def get_chat_completion_prompt(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -489,6 +520,7 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
|
|
@ -505,6 +537,39 @@ class BitBucketPromptManager(CustomPromptManagement):
|
|||
prompt_id,
|
||||
prompt_variables,
|
||||
dynamic_callback_params,
|
||||
prompt_label,
|
||||
prompt_version,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
async def async_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,
|
||||
litellm_logging_obj: Any,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
tools: Optional[List[Dict]] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
"""
|
||||
Async version - delegates to PromptManagementBase async implementation.
|
||||
"""
|
||||
return await PromptManagementBase.async_get_chat_completion_prompt(
|
||||
self,
|
||||
model,
|
||||
messages,
|
||||
non_default_params,
|
||||
prompt_id=prompt_id,
|
||||
prompt_variables=prompt_variables,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_spec=prompt_spec,
|
||||
tools=tools,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from litellm.caching.caching import DualCache
|
|||
from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER
|
||||
from litellm.types.integrations.argilla import ArgillaItem
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.utils import (
|
||||
AdapterCompletionStreamWrapper,
|
||||
CallTypes,
|
||||
|
|
@ -158,9 +159,12 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
tools: Optional[List[Dict]] = None,
|
||||
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:
|
||||
|
|
@ -178,6 +182,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from litellm.integrations.prompt_management_base import (
|
|||
PromptManagementClient,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
|
||||
|
|
@ -29,6 +30,7 @@ class CustomPromptManagement(CustomLogger, PromptManagementBase):
|
|||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
|
|
@ -48,14 +50,16 @@ class CustomPromptManagement(CustomLogger, PromptManagementBase):
|
|||
|
||||
def should_run_prompt_management(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
return True
|
||||
|
||||
def _compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
|
|||
from litellm.integrations.custom_prompt_management import CustomPromptManagement
|
||||
from litellm.integrations.prompt_management_base import PromptManagementClient
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
from .prompt_manager import PromptManager, PromptTemplate
|
||||
|
|
@ -82,7 +83,8 @@ class DotpromptManager(CustomPromptManagement):
|
|||
|
||||
def should_run_prompt_management(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
"""
|
||||
|
|
@ -90,6 +92,8 @@ class DotpromptManager(CustomPromptManagement):
|
|||
|
||||
Returns True if the prompt_id exists in our prompt manager.
|
||||
"""
|
||||
if prompt_id is None:
|
||||
return False
|
||||
try:
|
||||
return prompt_id in self.prompt_manager.list_prompts()
|
||||
except Exception:
|
||||
|
|
@ -98,7 +102,8 @@ class DotpromptManager(CustomPromptManagement):
|
|||
|
||||
def _compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
|
|
@ -114,6 +119,9 @@ class DotpromptManager(CustomPromptManagement):
|
|||
4. Extracts model and optional parameters from metadata
|
||||
"""
|
||||
|
||||
if prompt_id is None:
|
||||
raise ValueError("prompt_id is required for dotprompt manager")
|
||||
|
||||
try:
|
||||
|
||||
# Get the prompt template (versioned or base)
|
||||
|
|
@ -153,6 +161,31 @@ class DotpromptManager(CustomPromptManagement):
|
|||
except Exception as e:
|
||||
raise ValueError(f"Error compiling prompt '{prompt_id}': {e}")
|
||||
|
||||
async def async_compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> PromptManagementClient:
|
||||
"""
|
||||
Async version of compile prompt helper. Since dotprompt operations are synchronous,
|
||||
this simply delegates to the sync version.
|
||||
"""
|
||||
if prompt_id is None:
|
||||
raise ValueError("prompt_id is required for dotprompt manager")
|
||||
|
||||
return self._compile_prompt_helper(
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_variables=prompt_variables,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
def get_chat_completion_prompt(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -161,6 +194,7 @@ class DotpromptManager(CustomPromptManagement):
|
|||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
|
|
@ -177,8 +211,43 @@ class DotpromptManager(CustomPromptManagement):
|
|||
prompt_id,
|
||||
prompt_variables,
|
||||
dynamic_callback_params,
|
||||
prompt_label,
|
||||
prompt_version,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
async def async_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,
|
||||
litellm_logging_obj: Any,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
tools: Optional[List[Dict]] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
"""
|
||||
Async version - delegates to PromptManagementBase async implementation.
|
||||
"""
|
||||
from litellm.integrations.prompt_management_base import PromptManagementBase
|
||||
|
||||
return await PromptManagementBase.async_get_chat_completion_prompt(
|
||||
self,
|
||||
model,
|
||||
messages,
|
||||
non_default_params,
|
||||
prompt_id=prompt_id,
|
||||
prompt_variables=prompt_variables,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_spec=prompt_spec,
|
||||
tools=tools,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
def _convert_to_messages(self, rendered_content: str) -> List[AllMessageValues]:
|
||||
|
|
|
|||
80
litellm/integrations/generic_prompt_management/__init__.py
Normal file
80
litellm/integrations/generic_prompt_management/__init__.py
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
"""Generic prompt management integration for LiteLLM."""
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .generic_prompt_manager import GenericPromptManager
|
||||
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 .generic_prompt_manager import GenericPromptManager
|
||||
|
||||
# Global instances
|
||||
global_generic_prompt_config: Optional[dict] = None
|
||||
|
||||
|
||||
def set_global_generic_prompt_config(config: dict) -> None:
|
||||
"""
|
||||
Set the global generic prompt configuration.
|
||||
|
||||
Args:
|
||||
config: Dictionary containing generic prompt configuration
|
||||
- api_base: Base URL for the API
|
||||
- api_key: Optional API key for authentication
|
||||
- timeout: Request timeout in seconds (default: 30)
|
||||
"""
|
||||
import litellm
|
||||
|
||||
litellm.global_generic_prompt_config = config # type: ignore
|
||||
|
||||
|
||||
def prompt_initializer(
|
||||
litellm_params: "PromptLiteLLMParams", prompt_spec: "PromptSpec"
|
||||
) -> "CustomPromptManagement":
|
||||
"""
|
||||
Initialize a prompt from a generic prompt management API.
|
||||
"""
|
||||
prompt_id = getattr(litellm_params, "prompt_id", None)
|
||||
|
||||
api_base = litellm_params.api_base
|
||||
api_key = litellm_params.api_key
|
||||
if not api_base:
|
||||
raise ValueError("api_base is required in generic_prompt_config")
|
||||
|
||||
provider_specific_query_params = litellm_params.provider_specific_query_params
|
||||
|
||||
try:
|
||||
generic_prompt_manager = GenericPromptManager(
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
prompt_id=prompt_id,
|
||||
additional_provider_specific_query_params=provider_specific_query_params,
|
||||
**litellm_params.model_dump(
|
||||
exclude_none=True,
|
||||
exclude={
|
||||
"prompt_id",
|
||||
"api_key",
|
||||
"provider_specific_query_params",
|
||||
"api_base",
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
return generic_prompt_manager
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
|
||||
prompt_initializer_registry = {
|
||||
SupportedPromptIntegrations.GENERIC_PROMPT_MANAGEMENT.value: prompt_initializer,
|
||||
}
|
||||
|
||||
# Export public API
|
||||
__all__ = [
|
||||
"GenericPromptManager",
|
||||
"set_global_generic_prompt_config",
|
||||
"global_generic_prompt_config",
|
||||
"prompt_initializer_registry",
|
||||
]
|
||||
|
|
@ -0,0 +1,501 @@
|
|||
"""
|
||||
Generic prompt manager that integrates with LiteLLM's prompt management system.
|
||||
Fetches prompts from any API that implements the /beta/litellm_prompt_management endpoint.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.integrations.custom_prompt_management import CustomPromptManagement
|
||||
from litellm.integrations.prompt_management_base import (
|
||||
PromptManagementBase,
|
||||
PromptManagementClient,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
_get_httpx_client,
|
||||
get_async_httpx_client,
|
||||
)
|
||||
from litellm.types.llms.custom_http import httpxSpecialProvider
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
|
||||
class GenericPromptManager(CustomPromptManagement):
|
||||
"""
|
||||
Generic prompt manager that integrates with LiteLLM's prompt management system.
|
||||
|
||||
This class enables using prompts from any API that implements the
|
||||
/beta/litellm_prompt_management endpoint.
|
||||
|
||||
Usage:
|
||||
# Configure API access
|
||||
generic_config = {
|
||||
"api_base": "https://your-api.com",
|
||||
"api_key": "your-api-key", # optional
|
||||
"timeout": 30, # optional, defaults to 30
|
||||
}
|
||||
|
||||
# Use with completion
|
||||
response = litellm.completion(
|
||||
model="generic_prompt/gpt-4",
|
||||
prompt_id="my_prompt_id",
|
||||
prompt_variables={"variable": "value"},
|
||||
generic_prompt_config=generic_config,
|
||||
messages=[{"role": "user", "content": "Additional message"}]
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_base: str,
|
||||
api_key: Optional[str] = None,
|
||||
timeout: int = 30,
|
||||
prompt_id: Optional[str] = None,
|
||||
additional_provider_specific_query_params: Optional[Dict[str, Any]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Initialize the Generic Prompt Manager.
|
||||
|
||||
Args:
|
||||
api_base: Base URL for the API (e.g., "https://your-api.com")
|
||||
api_key: Optional API key for authentication
|
||||
timeout: Request timeout in seconds (default: 30)
|
||||
prompt_id: Optional prompt ID to pre-load
|
||||
"""
|
||||
super().__init__(**kwargs)
|
||||
self.api_base = api_base.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.timeout = timeout
|
||||
self.prompt_id = prompt_id
|
||||
self.additional_provider_specific_query_params = (
|
||||
additional_provider_specific_query_params
|
||||
)
|
||||
self._prompt_cache: Dict[str, PromptManagementClient] = {}
|
||||
|
||||
@property
|
||||
def integration_name(self) -> str:
|
||||
"""Integration name used in model names like 'generic_prompt/gpt-4'."""
|
||||
return "generic_prompt"
|
||||
|
||||
def _get_headers(self) -> Dict[str, str]:
|
||||
"""Get HTTP headers for API requests."""
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "application/json",
|
||||
}
|
||||
if self.api_key:
|
||||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||||
return headers
|
||||
|
||||
def _fetch_prompt_from_api(
|
||||
self, prompt_id: Optional[str], prompt_spec: Optional[PromptSpec]
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Fetch a prompt from the API.
|
||||
|
||||
Args:
|
||||
prompt_id: The ID of the prompt to fetch
|
||||
|
||||
Returns:
|
||||
The prompt data from the API
|
||||
|
||||
Raises:
|
||||
Exception: If the API request fails
|
||||
"""
|
||||
if prompt_id is None and prompt_spec is None:
|
||||
raise ValueError("prompt_id or prompt_spec is required")
|
||||
|
||||
url = f"{self.api_base}/beta/litellm_prompt_management"
|
||||
params = {
|
||||
"prompt_id": prompt_id,
|
||||
**(self.additional_provider_specific_query_params or {}),
|
||||
}
|
||||
http_client = _get_httpx_client()
|
||||
|
||||
try:
|
||||
|
||||
response = http_client.get(
|
||||
url,
|
||||
params=params,
|
||||
headers=self._get_headers(),
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except httpx.HTTPError as e:
|
||||
raise Exception(f"Failed to fetch prompt '{prompt_id}' from API: {e}")
|
||||
except json.JSONDecodeError as e:
|
||||
raise Exception(f"Failed to parse prompt response for '{prompt_id}': {e}")
|
||||
|
||||
async def async_fetch_prompt_from_api(
|
||||
self, prompt_id: Optional[str], prompt_spec: Optional[PromptSpec]
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Fetch a prompt from the API asynchronously.
|
||||
"""
|
||||
if prompt_id is None and prompt_spec is None:
|
||||
raise ValueError("prompt_id or prompt_spec is required")
|
||||
|
||||
url = f"{self.api_base}/beta/litellm_prompt_management"
|
||||
params = {
|
||||
"prompt_id": prompt_id,
|
||||
**(
|
||||
prompt_spec.litellm_params.provider_specific_query_params
|
||||
if prompt_spec
|
||||
and prompt_spec.litellm_params.provider_specific_query_params
|
||||
else {}
|
||||
),
|
||||
}
|
||||
|
||||
http_client = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.PromptManagement,
|
||||
)
|
||||
|
||||
try:
|
||||
response = await http_client.get(
|
||||
url,
|
||||
params=params,
|
||||
headers=self._get_headers(),
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except httpx.HTTPError as e:
|
||||
raise Exception(f"Failed to fetch prompt '{prompt_id}' from API: {e}")
|
||||
except json.JSONDecodeError as e:
|
||||
raise Exception(f"Failed to parse prompt response for '{prompt_id}': {e}")
|
||||
|
||||
def _parse_api_response(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
api_response: Dict[str, Any],
|
||||
) -> PromptManagementClient:
|
||||
"""
|
||||
Parse the API response into a PromptManagementClient structure.
|
||||
|
||||
Expected API response format:
|
||||
{
|
||||
"prompt_id": "string",
|
||||
"prompt_template": [
|
||||
{"role": "system", "content": "..."},
|
||||
{"role": "user", "content": "..."}
|
||||
],
|
||||
"prompt_template_model": "gpt-4", # optional
|
||||
"prompt_template_optional_params": { # optional
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 100
|
||||
}
|
||||
}
|
||||
|
||||
Args:
|
||||
prompt_id: The ID of the prompt
|
||||
api_response: The response from the API
|
||||
|
||||
Returns:
|
||||
PromptManagementClient structure
|
||||
"""
|
||||
return PromptManagementClient(
|
||||
prompt_id=prompt_id,
|
||||
prompt_template=api_response.get("prompt_template", []),
|
||||
prompt_template_model=api_response.get("prompt_template_model"),
|
||||
prompt_template_optional_params=api_response.get(
|
||||
"prompt_template_optional_params"
|
||||
),
|
||||
completed_messages=None,
|
||||
)
|
||||
|
||||
def should_run_prompt_management(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
"""
|
||||
Determine if prompt management should run based on the prompt_id.
|
||||
|
||||
For Generic Prompt Manager, we always return True and handle the prompt loading
|
||||
in the _compile_prompt_helper method.
|
||||
"""
|
||||
if prompt_id is not None or (
|
||||
prompt_spec is not None
|
||||
and prompt_spec.litellm_params.provider_specific_query_params is not None
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _get_cache_key(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> str:
|
||||
return f"{prompt_id}:{prompt_label}:{prompt_version}"
|
||||
|
||||
def _common_caching_logic(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
prompt_variables: Optional[dict] = None,
|
||||
) -> Optional[PromptManagementClient]:
|
||||
"""
|
||||
Common caching logic for the prompt manager.
|
||||
"""
|
||||
# Check cache first
|
||||
cache_key = self._get_cache_key(prompt_id, prompt_label, prompt_version)
|
||||
if cache_key in self._prompt_cache:
|
||||
cached_prompt = self._prompt_cache[cache_key]
|
||||
# Return a copy with variables applied if needed
|
||||
if prompt_variables:
|
||||
return self._apply_variables(cached_prompt, prompt_variables)
|
||||
return cached_prompt
|
||||
return None
|
||||
|
||||
def _compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> PromptManagementClient:
|
||||
"""
|
||||
Compile a prompt template into a PromptManagementClient structure.
|
||||
|
||||
This method:
|
||||
1. Fetches the prompt from the API (with caching)
|
||||
2. Applies any prompt variables (if the API supports it)
|
||||
3. Returns the structured prompt data
|
||||
|
||||
Args:
|
||||
prompt_id: The ID of the prompt
|
||||
prompt_variables: Variables to substitute in the template (optional)
|
||||
dynamic_callback_params: Dynamic callback parameters
|
||||
prompt_label: Optional label for the prompt version
|
||||
prompt_version: Optional specific version number
|
||||
|
||||
Returns:
|
||||
PromptManagementClient structure
|
||||
"""
|
||||
cached_prompt = self._common_caching_logic(
|
||||
prompt_id=prompt_id,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
prompt_variables=prompt_variables,
|
||||
)
|
||||
if cached_prompt:
|
||||
return cached_prompt
|
||||
|
||||
cache_key = self._get_cache_key(prompt_id, prompt_label, prompt_version)
|
||||
try:
|
||||
# Fetch from API
|
||||
api_response = self._fetch_prompt_from_api(prompt_id, prompt_spec)
|
||||
|
||||
# Parse the response
|
||||
prompt_client = self._parse_api_response(
|
||||
prompt_id, prompt_spec, api_response
|
||||
)
|
||||
|
||||
# Cache the result
|
||||
self._prompt_cache[cache_key] = prompt_client
|
||||
|
||||
# Apply variables if provided
|
||||
if prompt_variables:
|
||||
prompt_client = self._apply_variables(prompt_client, prompt_variables)
|
||||
|
||||
return prompt_client
|
||||
|
||||
except Exception as e:
|
||||
raise ValueError(f"Error compiling prompt '{prompt_id}': {e}")
|
||||
|
||||
async def async_compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> PromptManagementClient:
|
||||
|
||||
# Check cache first
|
||||
cached_prompt = self._common_caching_logic(
|
||||
prompt_id=prompt_id,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
prompt_variables=prompt_variables,
|
||||
)
|
||||
if cached_prompt:
|
||||
return cached_prompt
|
||||
|
||||
cache_key = self._get_cache_key(prompt_id, prompt_label, prompt_version)
|
||||
|
||||
try:
|
||||
# Fetch from API
|
||||
|
||||
api_response = await self.async_fetch_prompt_from_api(
|
||||
prompt_id=prompt_id, prompt_spec=prompt_spec
|
||||
)
|
||||
|
||||
# Parse the response
|
||||
prompt_client = self._parse_api_response(
|
||||
prompt_id, prompt_spec, api_response
|
||||
)
|
||||
|
||||
# Cache the result
|
||||
self._prompt_cache[cache_key] = prompt_client
|
||||
|
||||
# Apply variables if provided
|
||||
if prompt_variables:
|
||||
prompt_client = self._apply_variables(prompt_client, prompt_variables)
|
||||
|
||||
return prompt_client
|
||||
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Error compiling prompt '{prompt_id}': {e}, prompt_spec: {prompt_spec}"
|
||||
)
|
||||
|
||||
def _apply_variables(
|
||||
self,
|
||||
prompt_client: PromptManagementClient,
|
||||
variables: Dict[str, Any],
|
||||
) -> PromptManagementClient:
|
||||
"""
|
||||
Apply variables to the prompt template.
|
||||
|
||||
This performs simple string substitution using {variable_name} syntax.
|
||||
|
||||
Args:
|
||||
prompt_client: The prompt client structure
|
||||
variables: Variables to substitute
|
||||
|
||||
Returns:
|
||||
Updated PromptManagementClient with variables applied
|
||||
"""
|
||||
# Create a copy of the prompt template with variables applied
|
||||
updated_messages: List[AllMessageValues] = []
|
||||
for message in prompt_client["prompt_template"]:
|
||||
updated_message = dict(message) # type: ignore
|
||||
if "content" in updated_message and isinstance(
|
||||
updated_message["content"], str
|
||||
):
|
||||
content = updated_message["content"]
|
||||
for key, value in variables.items():
|
||||
content = content.replace(f"{{{key}}}", str(value))
|
||||
content = content.replace(
|
||||
f"{{{{{key}}}}}", str(value)
|
||||
) # Also support {{key}}
|
||||
updated_message["content"] = content
|
||||
updated_messages.append(updated_message) # type: ignore
|
||||
|
||||
return PromptManagementClient(
|
||||
prompt_id=prompt_client["prompt_id"],
|
||||
prompt_template=updated_messages,
|
||||
prompt_template_model=prompt_client["prompt_template_model"],
|
||||
prompt_template_optional_params=prompt_client[
|
||||
"prompt_template_optional_params"
|
||||
],
|
||||
completed_messages=None,
|
||||
)
|
||||
|
||||
async def async_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,
|
||||
litellm_logging_obj: "LiteLLMLoggingObj",
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
tools: Optional[List[Dict]] = None,
|
||||
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 and return processed model, messages, and parameters.
|
||||
"""
|
||||
|
||||
return await PromptManagementBase.async_get_chat_completion_prompt(
|
||||
self,
|
||||
model,
|
||||
messages,
|
||||
non_default_params,
|
||||
prompt_id=prompt_id,
|
||||
prompt_variables=prompt_variables,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_spec=prompt_spec,
|
||||
tools=tools,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
ignore_prompt_manager_model=(
|
||||
ignore_prompt_manager_model
|
||||
or prompt_spec.litellm_params.ignore_prompt_manager_model
|
||||
if prompt_spec
|
||||
else False
|
||||
),
|
||||
ignore_prompt_manager_optional_params=(
|
||||
ignore_prompt_manager_optional_params
|
||||
or prompt_spec.litellm_params.ignore_prompt_manager_optional_params
|
||||
if prompt_spec
|
||||
else False
|
||||
),
|
||||
)
|
||||
|
||||
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_spec: Optional[PromptSpec] = None,
|
||||
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 and return processed model, messages, and parameters.
|
||||
"""
|
||||
return PromptManagementBase.get_chat_completion_prompt(
|
||||
self,
|
||||
model,
|
||||
messages,
|
||||
non_default_params,
|
||||
prompt_id=prompt_id,
|
||||
prompt_variables=prompt_variables,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
ignore_prompt_manager_model=(
|
||||
ignore_prompt_manager_model
|
||||
or prompt_spec.litellm_params.ignore_prompt_manager_model
|
||||
if prompt_spec
|
||||
else False
|
||||
),
|
||||
ignore_prompt_manager_optional_params=(
|
||||
ignore_prompt_manager_optional_params
|
||||
or prompt_spec.litellm_params.ignore_prompt_manager_optional_params
|
||||
if prompt_spec
|
||||
else False
|
||||
),
|
||||
)
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
"""Clear the prompt cache."""
|
||||
self._prompt_cache.clear()
|
||||
|
|
@ -13,6 +13,7 @@ from litellm.integrations.prompt_management_base import (
|
|||
PromptManagementClient,
|
||||
)
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
GITLAB_PREFIX = "gitlab::"
|
||||
|
|
@ -454,19 +455,24 @@ class GitLabPromptManager(CustomPromptManagement):
|
|||
|
||||
def should_run_prompt_management(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
return True
|
||||
return prompt_id is not None
|
||||
|
||||
def _compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> PromptManagementClient:
|
||||
if prompt_id is None:
|
||||
raise ValueError("prompt_id is required for GitLab prompt manager")
|
||||
|
||||
try:
|
||||
decoded_id = decode_prompt_id(prompt_id)
|
||||
if decoded_id not in self.prompt_manager.prompts:
|
||||
|
|
@ -505,6 +511,31 @@ class GitLabPromptManager(CustomPromptManagement):
|
|||
except Exception as e:
|
||||
raise ValueError(f"Error compiling prompt '{prompt_id}': {e}")
|
||||
|
||||
async def async_compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> PromptManagementClient:
|
||||
"""
|
||||
Async version of compile prompt helper. Since GitLab operations use sync client,
|
||||
this simply delegates to the sync version.
|
||||
"""
|
||||
if prompt_id is None:
|
||||
raise ValueError("prompt_id is required for GitLab prompt manager")
|
||||
|
||||
return self._compile_prompt_helper(
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_variables=prompt_variables,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
def get_chat_completion_prompt(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -513,6 +544,7 @@ class GitLabPromptManager(CustomPromptManagement):
|
|||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
|
|
@ -526,8 +558,41 @@ class GitLabPromptManager(CustomPromptManagement):
|
|||
prompt_id,
|
||||
prompt_variables,
|
||||
dynamic_callback_params,
|
||||
prompt_label,
|
||||
prompt_version,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
async def async_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,
|
||||
litellm_logging_obj: Any,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
tools: Optional[List[Dict]] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
"""
|
||||
Async version - delegates to PromptManagementBase async implementation.
|
||||
"""
|
||||
return await PromptManagementBase.async_get_chat_completion_prompt(
|
||||
self,
|
||||
model,
|
||||
messages,
|
||||
non_default_params,
|
||||
prompt_id=prompt_id,
|
||||
prompt_variables=prompt_variables,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_spec=prompt_spec,
|
||||
tools=tools,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
from litellm.integrations.prompt_management_base import PromptManagementClient
|
||||
from litellm.litellm_core_utils.asyncify import run_async_function
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionSystemMessage
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayload
|
||||
|
||||
from ...litellm_core_utils.specialty_caches.dynamic_logging_cache import (
|
||||
|
|
@ -183,6 +184,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge
|
|||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
tools: Optional[List[Dict]] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
|
|
@ -200,9 +202,12 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge
|
|||
|
||||
def should_run_prompt_management(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
if prompt_id is None:
|
||||
return False
|
||||
langfuse_client = langfuse_client_init(
|
||||
langfuse_public_key=dynamic_callback_params.get("langfuse_public_key"),
|
||||
langfuse_secret=dynamic_callback_params.get("langfuse_secret"),
|
||||
|
|
@ -217,12 +222,16 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge
|
|||
|
||||
def _compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> PromptManagementClient:
|
||||
if prompt_id is None:
|
||||
raise ValueError("prompt_id is required for Langfuse prompt management")
|
||||
|
||||
langfuse_client = langfuse_client_init(
|
||||
langfuse_public_key=dynamic_callback_params.get("langfuse_public_key"),
|
||||
langfuse_secret=dynamic_callback_params.get("langfuse_secret"),
|
||||
|
|
@ -257,6 +266,24 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge
|
|||
completed_messages=None,
|
||||
)
|
||||
|
||||
async def async_compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> PromptManagementClient:
|
||||
return self._compile_prompt_helper(
|
||||
prompt_id=prompt_id,
|
||||
prompt_variables=prompt_variables,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
return run_async_function(
|
||||
self.async_log_success_event, kwargs, response_obj, start_time, end_time
|
||||
|
|
|
|||
|
|
@ -1,14 +1,18 @@
|
|||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from typing_extensions import TypedDict
|
||||
from typing_extensions import TYPE_CHECKING, TypedDict
|
||||
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
|
||||
class PromptManagementClient(TypedDict):
|
||||
prompt_id: str
|
||||
prompt_id: Optional[str]
|
||||
prompt_template: List[AllMessageValues]
|
||||
prompt_template_model: Optional[str]
|
||||
prompt_template_optional_params: Optional[Dict[str, Any]]
|
||||
|
|
@ -24,7 +28,8 @@ class PromptManagementBase(ABC):
|
|||
@abstractmethod
|
||||
def should_run_prompt_management(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
pass
|
||||
|
|
@ -32,7 +37,8 @@ class PromptManagementBase(ABC):
|
|||
@abstractmethod
|
||||
def _compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_id: Optional[str],
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
|
|
@ -40,6 +46,18 @@ class PromptManagementBase(ABC):
|
|||
) -> PromptManagementClient:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def async_compile_prompt_helper(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> PromptManagementClient:
|
||||
pass
|
||||
|
||||
def merge_messages(
|
||||
self,
|
||||
prompt_template: List[AllMessageValues],
|
||||
|
|
@ -55,10 +73,41 @@ class PromptManagementBase(ABC):
|
|||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
) -> PromptManagementClient:
|
||||
|
||||
compiled_prompt_client = self._compile_prompt_helper(
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_variables=prompt_variables,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
try:
|
||||
messages = compiled_prompt_client["prompt_template"] + client_messages
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Error compiling prompt: {e}. Prompt id={prompt_id}, prompt_variables={prompt_variables}, client_messages={client_messages}, dynamic_callback_params={dynamic_callback_params}"
|
||||
)
|
||||
|
||||
compiled_prompt_client["completed_messages"] = messages
|
||||
return compiled_prompt_client
|
||||
|
||||
async def async_compile_prompt(
|
||||
self,
|
||||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
client_messages: List[AllMessageValues],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
) -> PromptManagementClient:
|
||||
compiled_prompt_client = await self.async_compile_prompt_helper(
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_variables=prompt_variables,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_label=prompt_label,
|
||||
|
|
@ -83,6 +132,39 @@ class PromptManagementBase(ABC):
|
|||
else:
|
||||
return model.replace("{}/".format(self.integration_name), "")
|
||||
|
||||
def post_compile_prompt_processing(
|
||||
self,
|
||||
prompt_template: PromptManagementClient,
|
||||
messages: List[AllMessageValues],
|
||||
non_default_params: dict,
|
||||
model: str,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
ignore_prompt_manager_optional_params: Optional[bool] = False,
|
||||
):
|
||||
completed_messages = prompt_template["completed_messages"] or messages
|
||||
|
||||
prompt_template_optional_params = (
|
||||
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
|
||||
else {}
|
||||
),
|
||||
}
|
||||
|
||||
if not ignore_prompt_manager_model:
|
||||
model = self._get_model_from_prompt(
|
||||
prompt_management_client=prompt_template, model=model
|
||||
)
|
||||
else:
|
||||
model = model
|
||||
|
||||
return model, completed_messages, updated_non_default_params
|
||||
|
||||
def get_chat_completion_prompt(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -91,6 +173,7 @@ class PromptManagementBase(ABC):
|
|||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
|
|
@ -100,7 +183,9 @@ class PromptManagementBase(ABC):
|
|||
if prompt_id is None:
|
||||
raise ValueError("prompt_id is required for Prompt Management Base class")
|
||||
if not self.should_run_prompt_management(
|
||||
prompt_id=prompt_id, dynamic_callback_params=dynamic_callback_params
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
):
|
||||
return model, messages, non_default_params
|
||||
|
||||
|
|
@ -113,26 +198,53 @@ class PromptManagementBase(ABC):
|
|||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
completed_messages = prompt_template["completed_messages"] or messages
|
||||
|
||||
prompt_template_optional_params = (
|
||||
prompt_template["prompt_template_optional_params"] or {}
|
||||
return self.post_compile_prompt_processing(
|
||||
prompt_template=prompt_template,
|
||||
messages=messages,
|
||||
non_default_params=non_default_params,
|
||||
model=model,
|
||||
ignore_prompt_manager_model=ignore_prompt_manager_model,
|
||||
ignore_prompt_manager_optional_params=ignore_prompt_manager_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
|
||||
async def async_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,
|
||||
litellm_logging_obj: "LiteLLMLoggingObj",
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
tools: Optional[List[Dict]] = None,
|
||||
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 not self.should_run_prompt_management(
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
):
|
||||
return model, messages, 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
|
||||
prompt_template = await self.async_compile_prompt(
|
||||
prompt_id=prompt_id,
|
||||
prompt_variables=prompt_variables,
|
||||
client_messages=messages,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
)
|
||||
|
||||
|
||||
return model, completed_messages, updated_non_default_params
|
||||
return self.post_compile_prompt_processing(
|
||||
prompt_template=prompt_template,
|
||||
messages=messages,
|
||||
non_default_params=non_default_params,
|
||||
model=model,
|
||||
ignore_prompt_manager_model=ignore_prompt_manager_model,
|
||||
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -85,6 +85,7 @@ from litellm.types.llms.openai import (
|
|||
ResponsesAPIResponse,
|
||||
)
|
||||
from litellm.types.mcp import MCPPostCallResponseObject
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
from litellm.types.rerank import RerankResponse
|
||||
from litellm.types.utils import (
|
||||
CachingDetails,
|
||||
|
|
@ -265,6 +266,7 @@ def _get_cached_prometheus_logger():
|
|||
global _PrometheusLogger
|
||||
if _PrometheusLogger is None:
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
|
||||
_PrometheusLogger = PrometheusLogger
|
||||
return _PrometheusLogger
|
||||
|
||||
|
|
@ -601,8 +603,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
non_default_params: Dict,
|
||||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
prompt_id: Optional[str] = None,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_management_logger: Optional[CustomLogger] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
|
|
@ -613,6 +616,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
model=model,
|
||||
non_default_params=non_default_params,
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
dynamic_callback_params=self.standard_callback_dynamic_params,
|
||||
)
|
||||
)
|
||||
|
|
@ -627,6 +631,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
messages=messages,
|
||||
non_default_params=non_default_params or {},
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_variables=prompt_variables,
|
||||
dynamic_callback_params=self.standard_callback_dynamic_params,
|
||||
prompt_label=prompt_label,
|
||||
|
|
@ -640,8 +645,9 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
non_default_params: Dict,
|
||||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
prompt_id: Optional[str] = None,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_management_logger: Optional[CustomLogger] = None,
|
||||
tools: Optional[List[Dict]] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
|
|
@ -654,6 +660,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
tools=tools,
|
||||
non_default_params=non_default_params,
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
dynamic_callback_params=self.standard_callback_dynamic_params,
|
||||
)
|
||||
)
|
||||
|
|
@ -668,6 +675,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
messages=messages,
|
||||
non_default_params=non_default_params or {},
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_variables=prompt_variables,
|
||||
dynamic_callback_params=self.standard_callback_dynamic_params,
|
||||
litellm_logging_obj=self,
|
||||
|
|
@ -681,6 +689,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
def _auto_detect_prompt_management_logger(
|
||||
self,
|
||||
prompt_id: str,
|
||||
prompt_spec: Optional[PromptSpec],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
) -> Optional[CustomLogger]:
|
||||
"""
|
||||
|
|
@ -706,6 +715,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
try:
|
||||
if logger.should_run_prompt_management(
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
):
|
||||
self.model_call_details["prompt_integration"] = (
|
||||
|
|
@ -724,6 +734,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
non_default_params: Dict,
|
||||
tools: Optional[List[Dict]] = None,
|
||||
prompt_id: Optional[str] = None,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
dynamic_callback_params: Optional[StandardCallbackDynamicParams] = None,
|
||||
) -> Optional[CustomLogger]:
|
||||
"""
|
||||
|
|
@ -756,6 +767,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
if prompt_id and dynamic_callback_params is not None:
|
||||
auto_detected_logger = self._auto_detect_prompt_management_logger(
|
||||
prompt_id=prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
)
|
||||
if auto_detected_logger is not None:
|
||||
|
|
@ -3516,7 +3528,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
|||
return _literalai_logger # type: ignore
|
||||
elif logging_integration == "prometheus":
|
||||
PrometheusLogger = _get_cached_prometheus_logger()
|
||||
|
||||
|
||||
for callback in _in_memory_loggers:
|
||||
if isinstance(callback, PrometheusLogger):
|
||||
return callback # type: ignore
|
||||
|
|
@ -3835,9 +3847,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
|||
_in_memory_loggers.append(_otel_logger)
|
||||
return _otel_logger # type: ignore
|
||||
elif logging_integration == "weave_otel":
|
||||
from litellm.integrations.opentelemetry import (
|
||||
OpenTelemetryConfig,
|
||||
)
|
||||
from litellm.integrations.opentelemetry import OpenTelemetryConfig
|
||||
from litellm.integrations.weave.weave_otel import (
|
||||
WeaveOtelLogger,
|
||||
get_weave_otel_config,
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
|
|
@ -10,11 +10,25 @@ model_list:
|
|||
litellm_params:
|
||||
model: openai/gpt-4.1-mini
|
||||
|
||||
|
||||
# guardrails:
|
||||
# - guardrail_name: generic-guardrail
|
||||
# litellm_params:
|
||||
# guardrail: generic_guardrail_api
|
||||
# mode: ["pre_call"]
|
||||
# headers:
|
||||
# Authorization: Bearer mock-bedrock-token-12345
|
||||
# api_base: http://localhost:8080
|
||||
# default_on: true
|
||||
|
||||
prompts:
|
||||
- prompt_id: "simple_prompt"
|
||||
litellm_params:
|
||||
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.
|
||||
prompt_integration: "generic_prompt_management"
|
||||
provider_specific_query_params:
|
||||
project_name: litellm
|
||||
slug: hello-world-prompt-2bac
|
||||
api_base: http://localhost:8080
|
||||
api_key: os.environ/BRAINTRUST_API_KEY
|
||||
ignore_prompt_manager_model: true
|
||||
ignore_prompt_manager_optional_params: true
|
||||
|
|
|
|||
|
|
@ -824,7 +824,7 @@ class ProxyLogging:
|
|||
|
||||
return data
|
||||
|
||||
def _process_prompt_template(
|
||||
async def _process_prompt_template(
|
||||
self,
|
||||
data: dict,
|
||||
litellm_logging_obj: Any,
|
||||
|
|
@ -833,6 +833,7 @@ class ProxyLogging:
|
|||
call_type: CallTypesLiteral,
|
||||
) -> None:
|
||||
"""Process prompt template if applicable."""
|
||||
|
||||
from litellm.proxy.prompts.prompt_endpoints import (
|
||||
construct_versioned_prompt_id,
|
||||
get_latest_version_prompt_id,
|
||||
|
|
@ -857,21 +858,24 @@ class ProxyLogging:
|
|||
litellm_prompt_id: Optional[str] = None
|
||||
if prompt_spec is not None:
|
||||
litellm_prompt_id = prompt_spec.litellm_params.prompt_id
|
||||
data.pop("prompt_id", None)
|
||||
|
||||
if custom_logger and prompt_spec is not None:
|
||||
|
||||
if custom_logger and litellm_prompt_id is not None:
|
||||
(
|
||||
model,
|
||||
messages,
|
||||
optional_params,
|
||||
) = litellm_logging_obj.get_chat_completion_prompt(
|
||||
) = await litellm_logging_obj.async_get_chat_completion_prompt(
|
||||
model=data.get("model", ""),
|
||||
messages=data.get("messages", []),
|
||||
non_default_params=get_non_default_completion_params(kwargs=data),
|
||||
non_default_params=get_non_default_completion_params(kwargs=data) or {},
|
||||
prompt_id=litellm_prompt_id,
|
||||
prompt_spec=prompt_spec,
|
||||
prompt_management_logger=custom_logger,
|
||||
prompt_variables=data.get("prompt_variables", None),
|
||||
prompt_label=data.get("prompt_label", None),
|
||||
prompt_version=data.get("prompt_version", None),
|
||||
prompt_variables=data.pop("prompt_variables", None) or {},
|
||||
prompt_label=data.pop("prompt_label", None) or {},
|
||||
prompt_version=data.pop("prompt_version", None) or {},
|
||||
)
|
||||
|
||||
data.update(optional_params)
|
||||
|
|
@ -976,8 +980,7 @@ class ProxyLogging:
|
|||
and prompt_id is not None
|
||||
and (call_type == "completion" or call_type == "acompletion")
|
||||
):
|
||||
|
||||
self._process_prompt_template(
|
||||
await self._process_prompt_template(
|
||||
data=data,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
prompt_id=prompt_id,
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ class httpxSpecialProvider(str, Enum):
|
|||
MCP = "mcp"
|
||||
RAG = "rag"
|
||||
A2A = "a2a"
|
||||
PromptManagement = "prompt_management"
|
||||
|
||||
|
||||
VerifyTypes = Union[str, bool, ssl.SSLContext]
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ class SupportedPromptIntegrations(str, Enum):
|
|||
CUSTOM = "custom"
|
||||
BITBUCKET = "bitbucket"
|
||||
GITLAB = "gitlab"
|
||||
GENERIC_PROMPT_MANAGEMENT = "generic_prompt_management"
|
||||
ARIZE_PHOENIX = "arize_phoenix"
|
||||
|
||||
|
||||
|
|
@ -21,10 +22,16 @@ class PromptInfo(BaseModel):
|
|||
|
||||
|
||||
class PromptLiteLLMParams(BaseModel):
|
||||
prompt_id: str
|
||||
prompt_id: Optional[str] = None
|
||||
prompt_integration: str
|
||||
api_key: Optional[str] = None
|
||||
|
||||
api_base: Optional[str] = None
|
||||
api_key: Optional[str] = None
|
||||
|
||||
provider_specific_query_params: Optional[Dict[str, Any]] = None
|
||||
|
||||
ignore_prompt_manager_model: Optional[bool] = False
|
||||
ignore_prompt_manager_optional_params: Optional[bool] = False
|
||||
|
||||
dotprompt_content: Optional[str] = None
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue