This commit is contained in:
Gautam Kishore 2026-11-06 15:56:31 -05:00 • committed by GitHub
commit 726376b77a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 515 additions and 0 deletions

342
cookbook/litellm_pydantic_ai.ipynb vendored Normal file
View file

@ -0,0 +1,342 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "header"
},
"source": [
"# Using LiteLLM with Pydantic AI\n",
"\n",
"This cookbook demonstrates how to use LiteLLM's proxy and Router with [Pydantic AI](https://pydantic-ai.readthedocs.io/) agents.\n",
"\n",
"LiteLLM exposes an OpenAI-compatible endpoint at `/v1/...` which Pydantic AI's OpenAI model provider can use directly. This allows you to:\n",
"- Route requests across multiple LLM providers (OpenAI, Anthropic, Google, etc.)\n",
"- Add load balancing, fallbacks, and rate limiting\n",
"- Use a single API key and endpoint for all models\n",
"\n",
"## Prerequisites\n",
"\n",
"Install the required packages:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "install"
},
"source": [
"!pip install litellm \"pydantic-ai-slim[openai]\""
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "setup_proxy"
},
"source": [
"## 1. Start the LiteLLM Proxy\n",
"\n",
"First, create a `config.yaml` for the LiteLLM proxy. This configures multiple providers behind a single OpenAI-compatible endpoint:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "config_yaml"
},
"source": [
"%%writefile litellm_config.yaml\n",
"model_list:\n",
" - model_name: gpt-4o\n",
" litellm_params:\n",
" model: openai/gpt-4o\n",
" api_key: os.environ/OPENAI_API_KEY\n",
" - model_name: claude-sonnet\n",
" litellm_params:\n",
" model: anthropic/claude-3-5-sonnet-20241022\n",
" api_key: os.environ/ANTHROPIC_API_KEY\n",
" - model_name: gemini-pro\n",
" litellm_params:\n",
" model: gemini/gemini-2.0-flash\n",
" api_key: os.environ/GEMINI_API_KEY\n",
"general_settings:\n",
" master_key: os.environ/LITELLM_MASTER_KEY # set a strong key in your shell\n"
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "start_proxy_cmd"
},
"source": [
"Start the proxy server bound to localhost only (in a separate terminal).\n",
"Generate a strong key first — the proxy defaults to listening on all\n",
"interfaces, so don't expose this example publicly:\n",
"\n",
"```bash\n",
"export LITELLM_MASTER_KEY=\"$(openssl rand -hex 32)\"\n",
"litellm --config litellm_config.yaml --port 4000 --host 127.0.0.1\n",
"```\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "pydantic_ai_basic"
},
"source": [
"import os\n",
"from pydantic_ai import Agent\n",
"from pydantic_ai.models.openai import OpenAIChatModel\n",
"from pydantic_ai.providers.openai import OpenAIProvider\n",
"\n",
"# Point Pydantic AI at the LiteLLM proxy via an OpenAI-compatible provider.\n",
"# Use the same key you started the proxy with (generate a strong one, e.g.\n",
"# `export LITELLM_MASTER_KEY=\"$(openssl rand -hex 32)\"`).\n",
"model = OpenAIChatModel(\n",
" 'gpt-4o',\n",
" provider=OpenAIProvider(\n",
" base_url='http://localhost:4000/v1',\n",
" api_key=os.environ.get('LITELLM_MASTER_KEY', 'sk-litellm-test'),\n",
" ),\n",
")\n",
"agent = Agent(model)\n",
"\n",
"result = agent.run_sync('What is the capital of France?')\n",
"print(result.data)\n"
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "router_python"
},
"source": [
"## 2b. (Alternative) Use LiteLLM Router Directly in Python\n",
"\n",
"If you don't want to run a proxy server, you can use LiteLLM's `Router` class\n",
"directly in your Python code for native failover and load balancing across\n",
"providers. This path uses LiteLLM only — no Pydantic AI `Agent` is involved:\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "router_direct"
},
"source": [
"from litellm import Router\n",
"\n",
"# Configure a Router with multiple providers\n",
"model_list = [\n",
" {\n",
" \"model_name\": \"gpt-4o\",\n",
" \"litellm_params\": {\n",
" \"model\": \"openai/gpt-4o\",\n",
" \"api_key\": os.environ.get(\"OPENAI_API_KEY\"),\n",
" },\n",
" },\n",
" {\n",
" \"model_name\": \"claude-sonnet\",\n",
" \"litellm_params\": {\n",
" \"model\": \"anthropic/claude-3-5-sonnet-20241022\",\n",
" \"api_key\": os.environ.get(\"ANTHROPIC_API_KEY\"),\n",
" },\n",
" },\n",
"]\n",
"\n",
"router = Router(model_list=model_list)\n",
"\n",
"# Use the router directly with LiteLLM's completion\n",
"response = router.completion(\n",
" model=\"gpt-4o\",\n",
" messages=[{\"role\": \"user\", \"content\": \"Say hello!\"}],\n",
")\n",
"print(response['choices'][0]['message']['content'])"
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "structured_output"
},
"source": [
"## 3. Structured Output with Pydantic AI + LiteLLM Proxy\n",
"\n",
"Pydantic AI's structured result types work seamlessly through the LiteLLM proxy:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "structured"
},
"source": [
"from pydantic import BaseModel\n",
"from pydantic_ai import Agent\n",
"from pydantic_ai.models.openai import OpenAIChatModel\n",
"from pydantic_ai.providers.openai import OpenAIProvider\n",
"import os\n",
"\n",
"\n",
"class City(BaseModel):\n",
" name: str\n",
" country: str\n",
" population: int\n",
"\n",
"\n",
"model = OpenAIChatModel(\n",
" 'gpt-4o',\n",
" provider=OpenAIProvider(\n",
" base_url='http://localhost:4000/v1',\n",
" api_key=os.environ.get('LITELLM_MASTER_KEY', 'sk-litellm-test'),\n",
" ),\n",
")\n",
"agent = Agent(\n",
" model,\n",
" result_type=list[City],\n",
" system_prompt='List the 3 largest cities in Europe with their countries and populations.',\n",
")\n",
"\n",
"result = agent.run_sync('Generate the list')\n",
"for city in result.data:\n",
" print(f\" {city.name}, {city.country} - Population: {city.population:,}\")\n"
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "tool_use"
},
"source": [
"## 4. Tool Use with Pydantic AI through LiteLLM\n",
"\n",
"Pydantic AI's function tool support works through the proxy:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "tools"
},
"source": [
"from pydantic_ai import Agent, RunContext\n",
"from pydantic_ai.models.openai import OpenAIChatModel\n",
"from pydantic_ai.providers.openai import OpenAIProvider\n",
"import os\n",
"\n",
"\n",
"def get_weather(ctx: RunContext, city: str) -> str:\n",
" return f\"The weather in {city} is sunny, 72\\u00b0F.\"\n",
"\n",
"\n",
"model = OpenAIChatModel(\n",
" 'gpt-4o',\n",
" provider=OpenAIProvider(\n",
" base_url='http://localhost:4000/v1',\n",
" api_key=os.environ.get('LITELLM_MASTER_KEY', 'sk-litellm-test'),\n",
" ),\n",
")\n",
"agent = Agent(\n",
" model,\n",
" tools=[get_weather],\n",
")\n",
"\n",
"result = agent.run_sync('What is the weather in Paris?')\n",
"print(result.data)\n"
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "model_selection"
},
"source": [
"## 5. Switching Between Models\n",
"\n",
"Since the proxy exposes all configured models behind the same endpoint, you can switch models easily:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "model_switch"
},
"source": [
"from pydantic_ai import Agent\n",
"from pydantic_ai.models.openai import OpenAIChatModel\n",
"from pydantic_ai.providers.openai import OpenAIProvider\n",
"import os\n",
"\n",
"\n",
"def proxy_model(name: str) -> OpenAIChatModel:\n",
" return OpenAIChatModel(\n",
" name,\n",
" provider=OpenAIProvider(\n",
" base_url='http://localhost:4000/v1',\n",
" api_key=os.environ.get('LITELLM_MASTER_KEY', 'sk-litellm-test'),\n",
" ),\n",
" )\n",
"\n",
"\n",
"# Use Claude via LiteLLM proxy\n",
"claude_agent = Agent(proxy_model('claude-sonnet')) # model_name from proxy config\n",
"\n",
"result = claude_agent.run_sync('Explain quantum computing in one sentence.')\n",
"print(f\"[Claude]: {result.data}\")\n",
"\n",
"# Use Gemini via LiteLLM proxy\n",
"gemini_agent = Agent(proxy_model('gemini-pro'))\n",
"\n",
"result = gemini_agent.run_sync('Explain quantum computing in one sentence.')\n",
"print(f\"[Gemini]: {result.data}\")\n"
],
"outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "summary"
},
"source": [
"## Summary\n",
"\n",
"- LiteLLM's proxy exposes an OpenAI-compatible `/v1/...` endpoint\n",
"- Pydantic AI consumes it through `OpenAIProvider` + `OpenAIChatModel` pointed at the proxy\n",
"- All Pydantic AI features (structured outputs, tools, multi-model) work through the proxy\n",
"- The LiteLLM Router is also available for direct Python integration without a proxy server\n",
"\n",
"For more details, see:\n",
"- [LiteLLM Proxy Docs](https://docs.litellm.ai/docs/proxy)\n",
"- [Pydantic AI Docs](https://pydantic-ai.readthedocs.io/)\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.10.0"
}
},
"nbformat": 4,
"nbformat_minor": 4
}

View file

@ -0,0 +1,173 @@
"""
Using LiteLLM with Pydantic AI
This script demonstrates how to use LiteLLM's proxy and Router with Pydantic AI agents.
Prerequisites:
pip install litellm "pydantic-ai-slim[openai]"
Steps:
1. (Option A) Start the LiteLLM proxy server and point Pydantic AI at it
2. (Option B) Use LiteLLM Router directly (native LiteLLM failover, no proxy server)
For Option A, first run:
export LITELLM_MASTER_KEY="$(openssl rand -hex 32)"
litellm --config litellm_config.yaml --port 4000 --host 127.0.0.1
"""
import os
from pydantic import BaseModel
from pydantic_ai import Agent, RunContext
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.openai import OpenAIProvider
# ============================================================
# Option A: LiteLLM Proxy Server
# ============================================================
# The proxy exposes an OpenAI-compatible endpoint at http://localhost:4000/v1
# which Pydantic AI consumes through OpenAIProvider/OpenAIChatModel.
#
# Save this as litellm_config.yaml:
#
# model_list:
# - model_name: gpt-4o
# litellm_params:
# model: openai/gpt-4o
# api_key: os.environ/OPENAI_API_KEY
# - model_name: claude-sonnet
# litellm_params:
# model: anthropic/claude-3-5-sonnet-20241022
# api_key: os.environ/ANTHROPIC_API_KEY
# - model_name: gemini-pro
# litellm_params:
# model: gemini/gemini-2.0-flash
# api_key: os.environ/GEMINI_API_KEY
# general_settings:
# master_key: os.environ/LITELLM_MASTER_KEY
#
# Start the proxy bound to localhost only (the default 0.0.0.0 listens on all
# interfaces — don't expose this example proxy publicly, anyone with the key
# could send requests through it):
# export LITELLM_MASTER_KEY="$(openssl rand -hex 32)"
# litellm --config litellm_config.yaml --port 4000 --host 127.0.0.1
PROXY_BASE_URL = "http://localhost:4000/v1"
# Same key you started the proxy with. Generate a strong one per the comment
# above instead of hardcoding a secret here.
PROXY_API_KEY = os.environ.get("LITELLM_MASTER_KEY", "sk-litellm-test")
def proxy_model(model_name: str) -> OpenAIChatModel:
"""Build a Pydantic AI model routed through the local LiteLLM proxy."""
return OpenAIChatModel(
model_name,
provider=OpenAIProvider(
base_url=PROXY_BASE_URL,
api_key=PROXY_API_KEY,
),
)
def basic_usage():
agent = Agent(proxy_model("gpt-4o"))
result = agent.run_sync("What is the capital of France?")
print("Basic usage:", result.data)
def structured_output():
class City(BaseModel):
name: str
country: str
population: int
agent = Agent(
proxy_model("gpt-4o"),
result_type=list[City],
system_prompt="List the 3 largest cities in Europe with their countries and populations.",
)
result = agent.run_sync("Generate the list")
for city in result.data:
print(f" {city.name}, {city.country} - Population: {city.population:,}")
def tool_usage():
def get_weather(ctx: RunContext, city: str) -> str:
return f"The weather in {city} is sunny, 72 degrees F."
agent = Agent(
proxy_model("gpt-4o"),
tools=[get_weather],
)
result = agent.run_sync("What is the weather in Paris?")
print("Tool usage:", result.data)
def switch_models():
claude_agent = Agent(proxy_model("claude-sonnet"))
result = claude_agent.run_sync("Explain quantum computing in one sentence.")
print(f"[Claude]: {result.data}")
gemini_agent = Agent(proxy_model("gemini-pro"))
result = gemini_agent.run_sync("Explain quantum computing in one sentence.")
print(f"[Gemini]: {result.data}")
# ============================================================
# Option B: LiteLLM Router (Direct Python, no proxy server)
# ============================================================
# This path uses LiteLLM's Router natively (failover / load balancing across
# providers) without running a proxy server and without a Pydantic AI Agent.
def router_usage():
from litellm import Router
model_list = [
{
"model_name": "gpt-4o",
"litellm_params": {
"model": "openai/gpt-4o",
"api_key": os.environ.get("OPENAI_API_KEY"),
},
},
{
"model_name": "claude-sonnet",
"litellm_params": {
"model": "anthropic/claude-3-5-sonnet-20241022",
"api_key": os.environ.get("ANTHROPIC_API_KEY"),
},
},
]
router = Router(model_list=model_list)
response = router.completion(
model="gpt-4o",
messages=[{"role": "user", "content": "Say hello!"}],
)
print("Router usage:", response["choices"][0]["message"]["content"])
if __name__ == "__main__":
print("=== LiteLLM + Pydantic AI Cookbook ===")
print()
print("--- Basic Usage ---")
basic_usage()
print()
print("--- Structured Output ---")
structured_output()
print()
print("--- Tool Usage ---")
tool_usage()
print()
print("--- Switch Models ---")
switch_models()
print()
print("--- Router (Direct Python) ---")
router_usage()