diff --git a/cookbook/litellm_pydantic_ai.ipynb b/cookbook/litellm_pydantic_ai.ipynb new file mode 100644 index 00000000000..75cf06f5e14 --- /dev/null +++ b/cookbook/litellm_pydantic_ai.ipynb @@ -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 +} diff --git a/cookbook/litellm_pydantic_ai.py b/cookbook/litellm_pydantic_ai.py new file mode 100644 index 00000000000..2bb31310201 --- /dev/null +++ b/cookbook/litellm_pydantic_ai.py @@ -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()