Back to Litellm

Braintrust Prompt Wrapper for LiteLLM

cookbook/litellm_proxy_server/braintrust_prompt_wrapper_README.md

1.84.0-dev.26.5 KB
Original Source

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:

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.