Skip to content

Commit 74b9560

Browse files
committed
Add gemini example using litellm
1 parent c29afd0 commit 74b9560

File tree

8 files changed

+616
-0
lines changed

8 files changed

+616
-0
lines changed
Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
# Python
2+
__pycache__/
3+
*.py[cod]
4+
*$py.class
5+
*.so
6+
.Python
7+
build/
8+
develop-eggs/
9+
dist/
10+
downloads/
11+
eggs/
12+
.eggs/
13+
lib/
14+
lib64/
15+
parts/
16+
sdist/
17+
var/
18+
wheels/
19+
*.egg-info/
20+
.installed.cfg
21+
*.egg
22+
23+
# Environments
24+
.env**
25+
.venv
26+
env/
27+
venv/
28+
ENV/
29+
env.bak/
30+
venv.bak/
31+
32+
# IDE
33+
.idea/
34+
.vscode/
35+
*.swp
36+
*.swo
37+
38+
# Git
39+
.git
40+
.gitignore
41+
42+
# Misc
43+
.DS_Store
Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
# syntax=docker/dockerfile:1.3
2+
FROM python:3.12-slim
3+
COPY --from=ghcr.io/astral-sh/uv:0.6.4 /uv /uvx /bin/
4+
5+
# Install system dependencies
6+
RUN apt-get update && apt-get install -y \
7+
htop \
8+
vim \
9+
curl \
10+
tar \
11+
python3-dev \
12+
postgresql-client \
13+
build-essential \
14+
libpq-dev \
15+
gcc \
16+
cmake \
17+
netcat-openbsd \
18+
nodejs \
19+
npm \
20+
&& apt-get clean \
21+
&& rm -rf /var/lib/apt/lists/**
22+
23+
# Install tctl (Temporal CLI)
24+
RUN curl -L https://github.com/temporalio/tctl/releases/download/v1.18.1/tctl_1.18.1_linux_arm64.tar.gz -o /tmp/tctl.tar.gz && \
25+
tar -xzf /tmp/tctl.tar.gz -C /usr/local/bin && \
26+
chmod +x /usr/local/bin/tctl && \
27+
rm /tmp/tctl.tar.gz
28+
29+
RUN uv pip install --system --upgrade pip setuptools wheel
30+
31+
ENV UV_HTTP_TIMEOUT=1000
32+
33+
# Copy pyproject.toml and README.md to install dependencies
34+
COPY 10_async/10_temporal/100_gemini_litellm/pyproject.toml /app/100_gemini_litellm/pyproject.toml
35+
COPY 10_async/10_temporal/100_gemini_litellm/README.md /app/100_gemini_litellm/README.md
36+
37+
WORKDIR /app/100_gemini_litellm
38+
39+
# Copy the project code
40+
COPY 10_async/10_temporal/100_gemini_litellm/project /app/100_gemini_litellm/project
41+
42+
# Install the required Python packages
43+
RUN uv pip install --system .
44+
45+
WORKDIR /app/100_gemini_litellm
46+
47+
ENV PYTHONPATH=/app
48+
ENV AGENT_NAME=at100-gemini-litellm
49+
50+
# Run the ACP server using uvicorn
51+
CMD ["uvicorn", "project.acp:acp", "--host", "0.0.0.0", "--port", "8000"]
52+
53+
# When we deploy the worker, we will replace the CMD with the following
54+
# CMD ["python", "-m", "run_worker"]
Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
# [Temporal] Using Alternative Models with LiteLLM (Gemini)
2+
3+
**Part of the [OpenAI SDK + Temporal integration series](../README.md)**
4+
5+
## What You'll Learn
6+
7+
This tutorial demonstrates how to use Google's Gemini models (or any other LLM provider) with the OpenAI Agents SDK through LiteLLM. The key insight is that LiteLLM provides a unified interface, allowing you to swap models without changing your agent code structure.
8+
9+
**Key insight:** You can use the same OpenAI Agents SDK patterns with any LLM provider supported by LiteLLM - Gemini, Anthropic Claude, Mistral, and many more.
10+
11+
## Prerequisites
12+
- Development environment set up (see [main repo README](https://github.com/scaleapi/scale-agentex))
13+
- Backend services running: `make dev` from repository root (includes Temporal)
14+
- Temporal UI available at http://localhost:8233
15+
- **Google Gemini API key** (see setup below)
16+
- Understanding of OpenAI Agents SDK basics (see [060_open_ai_agents_sdk_hello_world](../060_open_ai_agents_sdk_hello_world/))
17+
18+
## Setup
19+
20+
### 1. Get a Gemini API Key
21+
22+
1. Go to [Google AI Studio](https://aistudio.google.com/apikey)
23+
2. Create a new API key
24+
3. Copy the key for the next step
25+
26+
### 2. Configure the API Key
27+
28+
Add to your environment or `manifest.yaml`:
29+
30+
**Option A: Environment variable**
31+
```bash
32+
export GEMINI_API_KEY="your-gemini-api-key-here"
33+
```
34+
35+
**Option B: In manifest.yaml**
36+
```yaml
37+
agent:
38+
env:
39+
GEMINI_API_KEY: "your-gemini-api-key-here"
40+
```
41+
42+
### 3. Install LiteLLM Dependency
43+
44+
The `pyproject.toml` already includes `litellm>=1.52.0`. When you run the agent, dependencies are installed automatically.
45+
46+
## Quick Start
47+
48+
```bash
49+
cd examples/tutorials/10_async/10_temporal/100_gemini_litellm
50+
uv run agentex agents run --manifest manifest.yaml
51+
```
52+
53+
**Monitor:** Open Temporal UI at http://localhost:8233 to see workflow execution.
54+
55+
## Key Code Changes
56+
57+
The main difference from OpenAI examples is using `LitellmModel`:
58+
59+
```python
60+
from agents.extensions.models.litellm_model import LitellmModel
61+
62+
# Create a LiteLLM model pointing to Gemini
63+
gemini_model = LitellmModel(model="gemini/gemini-2.0-flash")
64+
65+
agent = Agent(
66+
name="Gemini Assistant",
67+
instructions="You are a helpful assistant powered by Gemini.",
68+
model=gemini_model, # Use the LiteLLM model instead of default
69+
)
70+
71+
# Run works exactly the same way
72+
result = await Runner.run(agent, user_messages)
73+
```
74+
75+
## Supported Models
76+
77+
LiteLLM supports many providers. Just change the model string:
78+
79+
| Provider | Model String Example |
80+
|----------|---------------------|
81+
| Google Gemini | `gemini/gemini-2.0-flash`, `gemini/gemini-1.5-pro` |
82+
| Anthropic | `anthropic/claude-3-sonnet-20240229` |
83+
| Mistral | `mistral/mistral-large-latest` |
84+
| Cohere | `cohere/command-r-plus` |
85+
| AWS Bedrock | `bedrock/anthropic.claude-3-sonnet` |
86+
87+
See [LiteLLM Providers](https://docs.litellm.ai/docs/providers) for the full list.
88+
89+
## Why LiteLLM?
90+
91+
**Model Flexibility:** Switch between providers without code changes - just update the model string.
92+
93+
**Unified Interface:** Same OpenAI Agents SDK patterns work with any provider.
94+
95+
**Cost Optimization:** Easily compare costs across providers by switching models.
96+
97+
**Fallback Support:** LiteLLM supports automatic fallbacks if a provider is unavailable.
98+
99+
## Architecture Notes
100+
101+
The Temporal integration remains identical:
102+
- Workflows are durable and survive restarts
103+
- LLM calls are wrapped as activities automatically
104+
- Full observability in Temporal UI
105+
- Automatic retries on failures
106+
107+
The only change is the model provider - everything else works the same.
108+
109+
## When to Use
110+
111+
- Want to use non-OpenAI models with OpenAI Agents SDK
112+
- Need to compare model performance across providers
113+
- Building multi-model systems with fallbacks
114+
- Cost optimization across different providers
115+
- Regulatory requirements for specific model providers
116+
117+
## Troubleshooting
118+
119+
**"GEMINI_API_KEY environment variable is not set"**
120+
- Ensure you've exported the API key or added it to manifest.yaml
121+
122+
**"Model not found" errors**
123+
- Check the model string format matches LiteLLM's expected format
124+
- See [LiteLLM Providers](https://docs.litellm.ai/docs/providers) for correct model names
125+
126+
**Rate limiting errors**
127+
- Gemini has different rate limits than OpenAI
128+
- Consider adding retry logic or using LiteLLM's built-in retry support
129+
130+
**Previous:** [090_claude_agents_sdk_mvp](../090_claude_agents_sdk_mvp/) - Claude SDK integration
Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
# Gemini LiteLLM Tutorial
Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
import os
2+
from datetime import timedelta
3+
4+
from temporalio.contrib.openai_agents import OpenAIAgentsPlugin, ModelActivityParameters
5+
from agents.extensions.models.litellm_provider import LitellmProvider
6+
7+
# === DEBUG SETUP (AgentEx CLI Debug Support) ===
8+
if os.getenv("AGENTEX_DEBUG_ENABLED") == "true":
9+
import debugpy
10+
debug_port = int(os.getenv("AGENTEX_DEBUG_PORT", "5679"))
11+
debugpy.configure(subProcess=False)
12+
debugpy.listen(debug_port)
13+
if os.getenv("AGENTEX_DEBUG_WAIT_FOR_ATTACH", "false").lower() == "true":
14+
debugpy.wait_for_client()
15+
# === END DEBUG SETUP ===
16+
17+
from agentex.lib.types.fastacp import TemporalACPConfig
18+
from agentex.lib.sdk.fastacp.fastacp import FastACP
19+
from agentex.lib.core.temporal.plugins.openai_agents.interceptors.context_interceptor import ContextInterceptor
20+
21+
context_interceptor = ContextInterceptor()
22+
23+
# Create the ACP server
24+
# We use LitellmProvider instead of TemporalStreamingModelProvider
25+
# to enable using Gemini and other models through LiteLLM
26+
acp = FastACP.create(
27+
acp_type="async",
28+
config=TemporalACPConfig(
29+
# When deployed to the cluster, the Temporal address will automatically be set to the cluster address
30+
# For local development, we set the address manually to talk to the local Temporal service set up via docker compose
31+
#
32+
# We use the OpenAI Agents SDK plugin because Temporal has built-in support for it,
33+
# handling serialization and activity wrapping automatically. LitellmProvider lets us
34+
# route to different model providers (like Gemini) while keeping all that infrastructure.
35+
type="temporal",
36+
temporal_address=os.getenv("TEMPORAL_ADDRESS", "localhost:7233"),
37+
plugins=[OpenAIAgentsPlugin(
38+
model_params=ModelActivityParameters(
39+
start_to_close_timeout=timedelta(days=1)
40+
),
41+
model_provider=LitellmProvider(),
42+
)],
43+
interceptors=[context_interceptor]
44+
)
45+
)
46+
47+
48+
# Notice that we don't need to register any handlers when we use type="temporal"
49+
# If you look at the code in agentex.sdk.fastacp.impl.temporal_acp
50+
# You can see that these handlers are automatically registered when the ACP is created
51+
52+
# @acp.on_task_create
53+
# This will be handled by the method in your workflow that is decorated with @workflow.run
54+
55+
# @acp.on_task_event_send
56+
# This will be handled by the method in your workflow that is decorated with @workflow.signal(name=SignalName.RECEIVE_MESSAGE)
57+
58+
# @acp.on_task_cancel
59+
# This does not need to be handled by your workflow.
60+
# It is automatically handled by the temporal client which cancels the workflow directly
Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
import asyncio
2+
from datetime import timedelta
3+
4+
from temporalio.contrib.openai_agents import OpenAIAgentsPlugin, ModelActivityParameters
5+
from agents.extensions.models.litellm_provider import LitellmProvider
6+
7+
from project.workflow import At100GeminiLitellmWorkflow
8+
from agentex.lib.utils.debug import setup_debug_if_enabled
9+
from agentex.lib.utils.logging import make_logger
10+
from agentex.lib.environment_variables import EnvironmentVariables
11+
from agentex.lib.core.temporal.activities import get_all_activities
12+
from agentex.lib.core.temporal.workers.worker import AgentexWorker
13+
from agentex.lib.core.temporal.plugins.openai_agents.interceptors.context_interceptor import ContextInterceptor
14+
15+
environment_variables = EnvironmentVariables.refresh()
16+
17+
logger = make_logger(__name__)
18+
19+
20+
async def main():
21+
# Setup debug mode if enabled
22+
setup_debug_if_enabled()
23+
24+
task_queue_name = environment_variables.WORKFLOW_TASK_QUEUE
25+
if task_queue_name is None:
26+
raise ValueError("WORKFLOW_TASK_QUEUE is not set")
27+
28+
# Add activities to the worker
29+
all_activities = get_all_activities() + [] # add your own activities here
30+
31+
# ============================================================================
32+
# LITELLM SETUP: Interceptor + LitellmProvider
33+
# ============================================================================
34+
# The ContextInterceptor threads task_id through activity headers using
35+
# Temporal's interceptor pattern. This enables runtime context without
36+
# forking the Temporal plugin.
37+
#
38+
# We use LitellmProvider instead of TemporalStreamingModelProvider to
39+
# enable routing to Gemini and other models through LiteLLM.
40+
context_interceptor = ContextInterceptor()
41+
42+
# Create a worker with automatic tracing
43+
# IMPORTANT: We use the STANDARD temporalio.contrib.openai_agents.OpenAIAgentsPlugin
44+
# but with LitellmProvider to handle model routing to Gemini.
45+
worker = AgentexWorker(
46+
task_queue=task_queue_name,
47+
plugins=[OpenAIAgentsPlugin(
48+
model_params=ModelActivityParameters(
49+
start_to_close_timeout=timedelta(days=1)
50+
),
51+
model_provider=LitellmProvider(),
52+
)],
53+
interceptors=[context_interceptor]
54+
)
55+
56+
await worker.run(
57+
activities=all_activities,
58+
workflow=At100GeminiLitellmWorkflow,
59+
)
60+
61+
if __name__ == "__main__":
62+
asyncio.run(main())

0 commit comments

Comments
 (0)