* 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
6.5 KiB
Braintrust Prompt Wrapper for LiteLLM
This directory contains a wrapper server that enables LiteLLM to use prompts from Braintrust 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:
{
"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_managementendpoint - Fetches prompts from Braintrust API
- Transforms Braintrust response format to LiteLLM format
Setup
Install Dependencies
pip install fastapi uvicorn httpx litellm
Set Environment Variables
export BRAINTRUST_API_KEY="your-braintrust-api-key"
Usage
Step 1: Start the Wrapper Server
python braintrust_prompt_wrapper_server.py
The server will start on http://localhost:8080 by default.
You can customize the port and host:
export PORT=8000
export HOST=0.0.0.0
python braintrust_prompt_wrapper_server.py
Step 2: Use with LiteLLM
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:
# 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:
{
"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:
{
"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:
temperaturemax_tokens/max_completion_tokenstop_pfrequency_penaltypresence_penaltynstopresponse_formattool_choicefunction_calltools
Variable Substitution
The generic prompt manager supports simple variable substitution:
# 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:
- Use HTTPS: Deploy behind a reverse proxy with SSL
- Authentication: Add authentication to the wrapper endpoint if needed
- Rate Limiting: Implement rate limiting to prevent abuse
- Caching: Consider caching prompt responses
- Monitoring: Add logging and monitoring
Example with Docker:
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:
- Create a wrapper server that implements
/beta/litellm_prompt_management - Transform the provider's response to LiteLLM format
- Use the generic prompt manager to connect
Example providers:
- Langsmith
- PromptLayer
- Humanloop
- Custom internal systems
Troubleshooting
"No Braintrust API token provided"
- Set
BRAINTRUST_API_KEYenvironment variable - Or pass token in
Authorization: Bearer TOKENheader
"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.