diff --git a/cookbook/litellm_pydantic_ai.ipynb b/cookbook/litellm_pydantic_ai.ipynb new file mode 100644 index 00000000000..05e0a524083 --- /dev/null +++ b/cookbook/litellm_pydantic_ai.ipynb @@ -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 +} \ No newline at end of file diff --git a/cookbook/litellm_pydantic_ai.py b/cookbook/litellm_pydantic_ai.py new file mode 100644 index 00000000000..b7b75bfad46 --- /dev/null +++ b/cookbook/litellm_pydantic_ai.py @@ -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() \ No newline at end of file