Skip to content

Commit ffd0d01

Browse files
cursoragentviniciusdsmello
cursoragent
and
viniciusdsmello
committed
feat: Add Google Gemini tracing support
Integrates tracing for Google Gemini LLM API calls. Includes a new example notebook. Co-authored-by: vinicius <[email protected]>
1 parent 90b940c commit ffd0d01

File tree

4 files changed

+953
-0
lines changed

4 files changed

+953
-0
lines changed

‎examples/tracing/google-gemini/gemini_tracing.ipynb‎

Lines changed: 231 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,231 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "markdown",
5+
"id": "2722b419",
6+
"metadata": {},
7+
"source": [
8+
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/openlayer-ai/openlayer-python/blob/main/examples/tracing/google-gemini/gemini_tracing.ipynb)\n",
9+
"\n",
10+
"\n",
11+
"# <a id=\"top\">Google Gemini API tracing</a>\n",
12+
"\n",
13+
"This notebook illustrates how to get started tracing Google Gemini API calls with Openlayer."
14+
]
15+
},
16+
{
17+
"cell_type": "code",
18+
"execution_count": null,
19+
"id": "020c8f6a",
20+
"metadata": {},
21+
"outputs": [],
22+
"source": [
23+
"!pip install google-generativeai openlayer"
24+
]
25+
},
26+
{
27+
"cell_type": "markdown",
28+
"id": "75c2a473",
29+
"metadata": {},
30+
"source": [
31+
"## 1. Set the environment variables"
32+
]
33+
},
34+
{
35+
"cell_type": "code",
36+
"execution_count": null,
37+
"id": "f3f4fa13",
38+
"metadata": {},
39+
"outputs": [],
40+
"source": [
41+
"import os\n",
42+
"\n",
43+
"import google.generativeai as genai\n",
44+
"\n",
45+
"# Gemini API key\n",
46+
"os.environ[\"GOOGLE_API_KEY\"] = \"YOUR_GOOGLE_API_KEY_HERE\"\n",
47+
"\n",
48+
"# Openlayer env variables\n",
49+
"os.environ[\"OPENLAYER_API_KEY\"] = \"YOUR_OPENLAYER_API_KEY_HERE\"\n",
50+
"os.environ[\"OPENLAYER_INFERENCE_PIPELINE_ID\"] = \"YOUR_OPENLAYER_INFERENCE_PIPELINE_ID_HERE\""
51+
]
52+
},
53+
{
54+
"cell_type": "markdown",
55+
"id": "9758533f",
56+
"metadata": {},
57+
"source": [
58+
"## 2. Configure Gemini and create a traced model"
59+
]
60+
},
61+
{
62+
"cell_type": "code",
63+
"execution_count": null,
64+
"id": "c35d9860-dc41-4f7c-8d69-cc2ac7e5e485",
65+
"metadata": {},
66+
"outputs": [],
67+
"source": [
68+
"from openlayer.lib import trace_gemini\n",
69+
"\n",
70+
"genai.configure(api_key=os.environ[\"GOOGLE_API_KEY\"])\n",
71+
"\n",
72+
"model = genai.GenerativeModel(\"gemini-pro\")\n",
73+
"traced_model = trace_gemini(model)"
74+
]
75+
},
76+
{
77+
"cell_type": "markdown",
78+
"id": "72a6b954",
79+
"metadata": {},
80+
"source": [
81+
"## 3. Use the traced Gemini model normally"
82+
]
83+
},
84+
{
85+
"cell_type": "markdown",
86+
"id": "76a350b4",
87+
"metadata": {},
88+
"source": [
89+
"That's it! Now you can continue using the traced Gemini model normally. The data is automatically published to Openlayer and you can start creating tests around it!"
90+
]
91+
},
92+
{
93+
"cell_type": "markdown",
94+
"id": "fb5ebdad",
95+
"metadata": {},
96+
"source": [
97+
"### 3.1 Non-streaming generation"
98+
]
99+
},
100+
{
101+
"cell_type": "code",
102+
"execution_count": null,
103+
"id": "e00c1c79",
104+
"metadata": {},
105+
"outputs": [],
106+
"source": [
107+
"response = traced_model.generate_content(\"What is the meaning of life?\")"
108+
]
109+
},
110+
{
111+
"cell_type": "code",
112+
"execution_count": null,
113+
"id": "b5e8c9f0",
114+
"metadata": {},
115+
"outputs": [],
116+
"source": [
117+
"response.text"
118+
]
119+
},
120+
{
121+
"cell_type": "markdown",
122+
"id": "09d39983",
123+
"metadata": {},
124+
"source": [
125+
"### 3.2 Streaming generation"
126+
]
127+
},
128+
{
129+
"cell_type": "code",
130+
"execution_count": null,
131+
"id": "9a86642c",
132+
"metadata": {},
133+
"outputs": [],
134+
"source": [
135+
"response = traced_model.generate_content(\"Tell me a short story.\", stream=True)\n",
136+
"\n",
137+
"for chunk in response:\n",
138+
" if hasattr(chunk, 'text'):\n",
139+
" continue # Process chunks as needed"
140+
]
141+
},
142+
{
143+
"cell_type": "markdown",
144+
"id": "4e6fb396",
145+
"metadata": {},
146+
"source": [
147+
"### 3.3 Multi-turn conversation"
148+
]
149+
},
150+
{
151+
"cell_type": "code",
152+
"execution_count": null,
153+
"id": "21369c42",
154+
"metadata": {},
155+
"outputs": [],
156+
"source": [
157+
"chat = traced_model.start_chat(history=[])\n",
158+
"\n",
159+
"response1 = chat.send_message(\"Hello, I'm learning about AI.\")\n",
160+
"response2 = chat.send_message(\"Can you explain neural networks?\")"
161+
]
162+
},
163+
{
164+
"cell_type": "code",
165+
"execution_count": null,
166+
"id": "22369c43",
167+
"metadata": {},
168+
"outputs": [],
169+
"source": [
170+
"response2.text"
171+
]
172+
},
173+
{
174+
"cell_type": "markdown",
175+
"id": "5e6fb397",
176+
"metadata": {},
177+
"source": [
178+
"### 3.4 With generation configuration"
179+
]
180+
},
181+
{
182+
"cell_type": "code",
183+
"execution_count": null,
184+
"id": "31369c44",
185+
"metadata": {},
186+
"outputs": [],
187+
"source": [
188+
"response = traced_model.generate_content(\n",
189+
" \"Write a haiku about technology.\",\n",
190+
" generation_config=genai.types.GenerationConfig(\n",
191+
" temperature=0.7,\n",
192+
" top_p=0.9,\n",
193+
" top_k=40,\n",
194+
" max_output_tokens=100,\n",
195+
" ),\n",
196+
")"
197+
]
198+
},
199+
{
200+
"cell_type": "code",
201+
"execution_count": null,
202+
"id": "41369c45",
203+
"metadata": {},
204+
"outputs": [],
205+
"source": [
206+
"response.text"
207+
]
208+
}
209+
],
210+
"metadata": {
211+
"kernelspec": {
212+
"display_name": "Python 3 (ipykernel)",
213+
"language": "python",
214+
"name": "python3"
215+
},
216+
"language_info": {
217+
"codemirror_mode": {
218+
"name": "ipython",
219+
"version": 3
220+
},
221+
"file_extension": ".py",
222+
"mimetype": "text/x-python",
223+
"name": "python",
224+
"nbconvert_exporter": "python",
225+
"pygments_lexer": "ipython3",
226+
"version": "3.9.18"
227+
}
228+
},
229+
"nbformat": 4,
230+
"nbformat_minor": 5
231+
}

‎src/openlayer/lib/__init__.py‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
"trace_litellm",
1717
"trace_google_adk",
1818
"unpatch_google_adk",
19+
"trace_gemini",
1920
"update_current_trace",
2021
"update_current_step",
2122
# Offline buffer management functions
@@ -235,3 +236,47 @@ def unpatch_google_adk():
235236
from .integrations import google_adk_tracer
236237

237238
return google_adk_tracer.unpatch_google_adk()
239+
240+
241+
# -------------------------------- Google Gemini --------------------------------- #
242+
def trace_gemini(client):
243+
"""Trace Google Gemini chat completions.
244+
245+
This function patches a Google GenerativeModel client to automatically trace
246+
all generate_content calls, including both streaming and non-streaming requests.
247+
248+
Parameters
249+
----------
250+
client : google.generativeai.GenerativeModel
251+
The Gemini GenerativeModel to patch.
252+
253+
Returns
254+
-------
255+
google.generativeai.GenerativeModel
256+
The patched Gemini client.
257+
258+
Example
259+
-------
260+
>>> import google.generativeai as genai
261+
>>> from openlayer.lib import trace_gemini
262+
>>> # Configure Gemini
263+
>>> genai.configure(api_key="your-api-key")
264+
>>> model = genai.GenerativeModel("gemini-pro")
265+
>>> # Enable tracing
266+
>>> traced_model = trace_gemini(model)
267+
>>> # Use normally - tracing happens automatically
268+
>>> response = traced_model.generate_content("Hello!")
269+
"""
270+
# pylint: disable=import-outside-toplevel
271+
try:
272+
import google.generativeai as genai
273+
except ImportError:
274+
raise ImportError(
275+
"google-generativeai is required for Gemini tracing. Install with: pip install google-generativeai"
276+
)
277+
278+
from .integrations import gemini_tracer
279+
280+
if not isinstance(client, genai.GenerativeModel):
281+
raise ValueError("Invalid client. Please provide a google.generativeai.GenerativeModel instance.")
282+
return gemini_tracer.trace_gemini(client)

‎src/openlayer/lib/integrations/__init__.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,3 +31,10 @@
3131
__all__.extend(["trace_google_adk", "unpatch_google_adk"])
3232
except ImportError:
3333
pass
34+
35+
try:
36+
from .gemini_tracer import trace_gemini
37+
38+
__all__.extend(["trace_gemini"])
39+
except ImportError:
40+
pass

0 commit comments

Comments
 (0)