mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
feat(cookbook): add cookbook for using LiteLLM with Pydantic AI
Demonstrates two integration paths: - LiteLLM proxy server (OpenAI-compatible /v1 endpoint consumed by Pydantic AI) - LiteLLM Router (direct Python usage) Covers basic usage, structured output, tool use, and model switching.
This commit is contained in:
Gautam Kishore
parent
60729f733e
commit
7eb7357a97
2 changed files with 477 additions and 0 deletions
309
cookbook/litellm_pydantic_ai.ipynb
vendored
Normal file
309
cookbook/litellm_pydantic_ai.ipynb
vendored
Normal file
|
|
@ -0,0 +1,309 @@
|
|||
{
|
||||
"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"
|
||||
],
|
||||
"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\ngeneral_settings:\n",
|
||||
" master_key: sk-litellm-test # API key for proxy auth"
|
||||
],
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "start_proxy_cmd"
|
||||
},
|
||||
"source": [
|
||||
"Start the proxy server (in a separate terminal):\n",
|
||||
"\n",
|
||||
"```bash\n",
|
||||
"litellm --config litellm_config.yaml --port 4000\n",
|
||||
"```\n",
|
||||
"\n",
|
||||
"## 2. Use Pydantic AI with the LiteLLM Proxy\n",
|
||||
"\n",
|
||||
"Once the proxy is running, configure your Pydantic AI agent to point at the proxy's OpenAI-compatible endpoint:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"metadata": {
|
||||
"id": "pydantic_ai_basic"
|
||||
},
|
||||
"source": [
|
||||
"import os\n",
|
||||
"from pydantic_ai import Agent\n",
|
||||
"\n",
|
||||
"# Point Pydantic AI at the LiteLLM proxy\n",
|
||||
"# The proxy speaks the OpenAI API protocol, so we use the 'openai' provider\n",
|
||||
"agent = Agent(\n",
|
||||
" 'openai:gpt-4o',\n",
|
||||
" base_url='http://localhost:4000/v1',\n",
|
||||
" api_key='sk-litellm-test',\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"result = agent.run_sync('What is the capital of France?')\n",
|
||||
"print(result.data)"
|
||||
],
|
||||
"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 directly in your Python code. Pydantic AI can still connect through a lightweight local proxy:"
|
||||
]
|
||||
},
|
||||
{
|
||||
"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",
|
||||
"\n",
|
||||
"\n",
|
||||
"class City(BaseModel):\n",
|
||||
" name: str\n",
|
||||
" country: str\n",
|
||||
" population: int\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"agent = Agent(\n",
|
||||
" 'openai:gpt-4o',\n",
|
||||
" base_url='http://localhost:4000/v1',\n",
|
||||
" api_key='sk-litellm-test',\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:,}\")"
|
||||
],
|
||||
"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 RunContext\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",
|
||||
"agent = Agent(\n",
|
||||
" 'openai:gpt-4o',\n",
|
||||
" base_url='http://localhost:4000/v1',\n",
|
||||
" api_key='sk-litellm-test',\n",
|
||||
" tools=[get_weather],\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"result = agent.run_sync('What is the weather in Paris?')\n",
|
||||
"print(result.data)"
|
||||
],
|
||||
"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": [
|
||||
"# Use Claude via LiteLLM proxy\n",
|
||||
"claude_agent = Agent(\n",
|
||||
" 'openai:claude-sonnet', # model_name from proxy config\n",
|
||||
" base_url='http://localhost:4000/v1',\n",
|
||||
" api_key='sk-litellm-test',\n",
|
||||
")\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(\n",
|
||||
" 'openai:gemini-pro',\n",
|
||||
" base_url='http://localhost:4000/v1',\n",
|
||||
" api_key='sk-litellm-test',\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"result = gemini_agent.run_sync('Explain quantum computing in one sentence.')\n",
|
||||
"print(f\"[Gemini]: {result.data}\")"
|
||||
],
|
||||
"outputs": []
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "summary"
|
||||
},
|
||||
"source": [
|
||||
"## Summary\n",
|
||||
"\n",
|
||||
"- LiteLLM's proxy exposes an OpenAI-compatible `/v1/...` endpoint\n",
|
||||
"- Pydantic AI's `openai:` provider can consume this endpoint by setting `base_url` and `api_key`\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/)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"language": "python",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python",
|
||||
"version": "3.10.0"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 4
|
||||
}
|
||||
168
cookbook/litellm_pydantic_ai.py
Normal file
168
cookbook/litellm_pydantic_ai.py
Normal file
|
|
@ -0,0 +1,168 @@
|
|||
"""
|
||||
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
|
||||
|
||||
Steps:
|
||||
1. (Option A) Start the LiteLLM proxy server and point Pydantic AI at it
|
||||
2. (Option B) Use LiteLLM Router directly with Pydantic AI's completion
|
||||
|
||||
For Option A, first run:
|
||||
litellm --config litellm_config.yaml --port 4000
|
||||
"""
|
||||
|
||||
import os
|
||||
from pydantic import BaseModel
|
||||
from pydantic_ai import Agent, RunContext
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Option A: LiteLLM Proxy Server
|
||||
# ============================================================
|
||||
# The proxy exposes an OpenAI-compatible endpoint at http://localhost:4000/v1
|
||||
# which Pydantic AI's openai provider can consume directly.
|
||||
#
|
||||
# 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: sk-litellm-test
|
||||
#
|
||||
# Start the proxy:
|
||||
# litellm --config litellm_config.yaml --port 4000
|
||||
|
||||
|
||||
PROXY_BASE_URL = "http://localhost:4000/v1"
|
||||
PROXY_API_KEY = "sk-litellm-test"
|
||||
|
||||
|
||||
def basic_usage():
|
||||
agent = Agent(
|
||||
"openai:gpt-4o",
|
||||
base_url=PROXY_BASE_URL,
|
||||
api_key=PROXY_API_KEY,
|
||||
)
|
||||
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(
|
||||
"openai:gpt-4o",
|
||||
base_url=PROXY_BASE_URL,
|
||||
api_key=PROXY_API_KEY,
|
||||
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(
|
||||
"openai:gpt-4o",
|
||||
base_url=PROXY_BASE_URL,
|
||||
api_key=PROXY_API_KEY,
|
||||
tools=[get_weather],
|
||||
)
|
||||
result = agent.run_sync("What is the weather in Paris?")
|
||||
print("Tool usage:", result.data)
|
||||
|
||||
|
||||
def switch_models():
|
||||
claude_agent = Agent(
|
||||
"openai:claude-sonnet",
|
||||
base_url=PROXY_BASE_URL,
|
||||
api_key=PROXY_API_KEY,
|
||||
)
|
||||
result = claude_agent.run_sync("Explain quantum computing in one sentence.")
|
||||
print(f"[Claude]: {result.data}")
|
||||
|
||||
gemini_agent = Agent(
|
||||
"openai:gemini-pro",
|
||||
base_url=PROXY_BASE_URL,
|
||||
api_key=PROXY_API_KEY,
|
||||
)
|
||||
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)
|
||||
# ============================================================
|
||||
|
||||
|
||||
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()
|
||||
Loading…
Add table
Reference in a new issue