Support google.genai

Author: alizenhom

src/examples/gemini_example/main.py

@@ -20,16 +20,16 @@ async def async_demo():
 
 def basic():
     generate()
-    generate(stream=True, with_tools=True)
+    # generate(stream=True, with_tools=True)
 
     # image_to_text()
     # audio_to_text()
-    asyncio.run(async_demo())
+    # asyncio.run(async_demo())
 
 
 def generate(stream=False, with_tools=False):
     model = genai.GenerativeModel(
-        "gemini-1.5-pro", system_instruction="You are a cat. Your name is Neko."
+        "gemini-2.0-flash-exp", system_instruction="You are a cat. Your name is Neko."
     )
 
     response = model.generate_content(

src/examples/google_genai_example/__init__.py

@@ -0,0 +1,7 @@
+from .main import generate_content, generate_content_streaming
+
+
+class GoogleGenaiRunner:
+    def run(self):
+        # generate_content()
+        generate_content_streaming()

src/examples/google_genai_example/main.py

@@ -0,0 +1,27 @@
+from google import genai
+from dotenv import load_dotenv
+import os
+from langtrace_python_sdk import langtrace
+
+load_dotenv()
+langtrace.init(write_spans_to_console=False)
+
+
+def generate_content():
+    # Only run this block for Google AI API
+    client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
+    response = client.models.generate_content(
+        model="gemini-2.0-flash-exp", contents="What is your name?"
+    )
+
+    print(response.text)
+
+
+def generate_content_streaming():
+    client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))
+    response = client.models.generate_content_stream(
+        model="gemini-2.0-flash-exp", contents="What is your name?"
+    )
+
+    for chunk in response:
+        pass

src/langtrace_python_sdk/instrumentation/__init__.py

@@ -24,6 +24,7 @@
 from .pymongo import PyMongoInstrumentation
 from .cerebras import CerebrasInstrumentation
 from .milvus import MilvusInstrumentation
+from .google_genai import GoogleGenaiInstrumentation
 
 __all__ = [
     "AnthropicInstrumentation",
@@ -52,4 +53,5 @@
     "AWSBedrockInstrumentation",
     "CerebrasInstrumentation",
     "MilvusInstrumentation",
+    "GoogleGenaiInstrumentation",
 ]

src/langtrace_python_sdk/instrumentation/google_genai/__init__.py

@@ -0,0 +1,3 @@
+from .instrumentation import GoogleGenaiInstrumentation
+
+__all__ = ["GoogleGenaiInstrumentation"]

src/langtrace_python_sdk/instrumentation/google_genai/instrumentation.py

@@ -0,0 +1,30 @@
+from typing import Collection
+from importlib_metadata import version as v
+from wrapt import wrap_function_wrapper as _W
+from opentelemetry.instrumentation.instrumentor import BaseInstrumentor
+from opentelemetry.trace import get_tracer
+from .patch import patch_google_genai, patch_google_genai_streaming
+
+
+class GoogleGenaiInstrumentation(BaseInstrumentor):
+    def instrumentation_dependencies(self) -> Collection[str]:
+        return ["google-genai >= 0.1.0", "google-generativeai < 1.0.0"]
+
+    def _instrument(self, **kwargs):
+        trace_provider = kwargs.get("tracer_provider")
+        tracer = get_tracer(__name__, "", trace_provider)
+        version = v("google-genai")
+
+        _W(
+            module="google.genai",
+            name="models.Models.generate_content",
+            wrapper=patch_google_genai(tracer, version),
+        )
+        _W(
+            module="google.genai",
+            name="models.Models.generate_content_stream",
+            wrapper=patch_google_genai_streaming(tracer, version),
+        )
+
+    def _uninstrument(self, **kwargs):
+        pass

src/langtrace_python_sdk/instrumentation/google_genai/patch.py

@@ -0,0 +1,129 @@
+from langtrace_python_sdk.utils.llm import (
+    get_langtrace_attributes,
+    get_llm_request_attributes,
+    set_span_attributes,
+    set_usage_attributes,
+    set_span_attribute,
+    set_event_completion,
+)
+from langtrace_python_sdk.utils import handle_span_error
+
+from opentelemetry.trace import Tracer, SpanKind
+from opentelemetry.sdk.trace import Span
+from langtrace.trace_attributes import SpanAttributes
+from google.genai.types import GenerateContentResponse
+
+from typing import Iterator
+
+
+def patch_google_genai(tracer: Tracer, version: str):
+    def traced_method(wrapped, instance, args, kwargs):
+        prompt = [
+            {
+                "role": "user",
+                "content": kwargs["contents"],
+            }
+        ]
+        span_attributes = {
+            **get_langtrace_attributes(
+                service_provider="google_genai", version=version
+            ),
+            **get_llm_request_attributes(kwargs=kwargs, prompts=prompt),
+        }
+        with tracer.start_as_current_span(
+            name="google.genai.generate_content",
+            kind=SpanKind.CLIENT,
+        ) as span:
+            try:
+                set_span_attributes(span, span_attributes)
+                response = wrapped(*args, **kwargs)
+                set_response_attributes(span, response)
+                return response
+            except Exception as error:
+                handle_span_error(span, error)
+                raise
+
+    return traced_method
+
+
+def patch_google_genai_streaming(tracer: Tracer, version: str):
+    def traced_method(wrapped, instance, args, kwargs):
+        prompt = [
+            {
+                "role": "user",
+                "content": kwargs["contents"],
+            }
+        ]
+        span_attributes = {
+            **get_langtrace_attributes(
+                service_provider="google_genai", version=version
+            ),
+            **get_llm_request_attributes(kwargs=kwargs, prompts=prompt),
+        }
+        with tracer.start_as_current_span(
+            name="google.genai.generate_content_stream",
+            kind=SpanKind.CLIENT,
+        ) as span:
+            set_span_attributes(span, span_attributes)
+            response = wrapped(*args, **kwargs)
+            set_streaming_response_attributes(span, response)
+            return response
+
+    return traced_method
+
+
+def set_streaming_response_attributes(
+    span: Span, response: Iterator[GenerateContentResponse]
+):
+    accum_completion = ""
+    for chunk in response:
+        set_span_attribute(
+            span,
+            SpanAttributes.LLM_RESPONSE_MODEL,
+            chunk.model_version,
+        )
+        candidates = chunk.candidates
+        for candidate in candidates:
+            set_span_attribute(
+                span,
+                SpanAttributes.LLM_RESPONSE_FINISH_REASON,
+                candidate.finish_reason,
+            )
+
+            accum_completion += candidate.content.parts[0].text
+
+        if chunk.usage_metadata:
+            set_usage_attributes(
+                span,
+                {
+                    "input_tokens": chunk.usage_metadata.prompt_token_count,
+                    "output_tokens": chunk.usage_metadata.candidates_token_count,
+                },
+            )
+    set_event_completion(span, [{"role": "assistant", "content": accum_completion}])
+
+
+def set_response_attributes(span: Span, response: GenerateContentResponse):
+    completions = []
+    for candidate in response.candidates:
+        set_span_attribute(
+            span, SpanAttributes.LLM_RESPONSE_FINISH_REASON, candidate.finish_reason
+        )
+        parts = candidate.content.parts
+        role = candidate.content.role
+        completion = {
+            "role": role or "assistant",
+            "content": [part.text for part in parts],
+        }
+        completions.append(completion)
+
+    set_span_attribute(span, SpanAttributes.LLM_RESPONSE_MODEL, response.model_version)
+    set_event_completion(span, completions)
+    if response.usage_metadata:
+        set_usage_attributes(
+            span,
+            {
+                "input_tokens": response.usage_metadata.prompt_token_count,
+                "output_tokens": response.usage_metadata.candidates_token_count,
+            },
+        )

src/langtrace_python_sdk/langtrace.py

@@ -68,6 +68,7 @@
     PyMongoInstrumentation,
     CerebrasInstrumentation,
     MilvusInstrumentation,
+    GoogleGenaiInstrumentation,
 )
 from opentelemetry.util.re import parse_env_headers
 
@@ -301,6 +302,7 @@ def init(
         "vertexai": VertexAIInstrumentation(),
         "google-cloud-aiplatform": VertexAIInstrumentation(),
         "google-generativeai": GeminiInstrumentation(),
+        "google-genai": GoogleGenaiInstrumentation(),
         "mistralai": MistralInstrumentation(),
         "boto3": AWSBedrockInstrumentation(),
         "autogen": AutogenInstrumentation(),

src/langtrace_python_sdk/version.py

@@ -1 +1 @@
-__version__ = "3.3.22"
+__version__ = "3.3.23"

src/run_example.py

@@ -20,8 +20,9 @@
     "vertexai": False,
     "gemini": False,
     "mistral": False,
-    "awsbedrock": True,
+    "awsbedrock": False,
     "cerebras": False,
+    "google_genai": True,
 }
 
 if ENABLED_EXAMPLES["anthropic"]:
@@ -137,3 +138,9 @@
 
     print(Fore.BLUE + "Running Cerebras example" + Fore.RESET)
     CerebrasRunner().run()
+
+if ENABLED_EXAMPLES["google_genai"]:
+    from examples.google_genai_example import GoogleGenaiRunner
+
+    print(Fore.BLUE + "Running Google Genai example" + Fore.RESET)
+    GoogleGenaiRunner().run()