From 03cea4eb615ef048b90dba854a499844b9e01ed8 Mon Sep 17 00:00:00 2001 From: Ali Waleed Date: Mon, 13 Jan 2025 17:52:50 +0200 Subject: [PATCH 1/2] Support `google.genai` --- src/examples/gemini_example/main.py | 6 +- src/examples/google_genai_example/__init__.py | 7 + src/examples/google_genai_example/main.py | 27 ++++ .../instrumentation/__init__.py | 2 + .../instrumentation/google_genai/__init__.py | 3 + .../google_genai/instrumentation.py | 30 ++++ .../instrumentation/google_genai/patch.py | 129 ++++++++++++++++++ src/langtrace_python_sdk/langtrace.py | 2 + src/langtrace_python_sdk/version.py | 2 +- src/run_example.py | 9 +- 10 files changed, 212 insertions(+), 5 deletions(-) create mode 100644 src/examples/google_genai_example/__init__.py create mode 100644 src/examples/google_genai_example/main.py create mode 100644 src/langtrace_python_sdk/instrumentation/google_genai/__init__.py create mode 100644 src/langtrace_python_sdk/instrumentation/google_genai/instrumentation.py create mode 100644 src/langtrace_python_sdk/instrumentation/google_genai/patch.py diff --git a/src/examples/gemini_example/main.py b/src/examples/gemini_example/main.py index 2990d518..2ca59303 100644 --- a/src/examples/gemini_example/main.py +++ b/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( diff --git a/src/examples/google_genai_example/__init__.py b/src/examples/google_genai_example/__init__.py new file mode 100644 index 00000000..bed04e13 --- /dev/null +++ b/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() diff --git a/src/examples/google_genai_example/main.py b/src/examples/google_genai_example/main.py new file mode 100644 index 00000000..8db51d96 --- /dev/null +++ b/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 diff --git a/src/langtrace_python_sdk/instrumentation/__init__.py b/src/langtrace_python_sdk/instrumentation/__init__.py index 952a4e10..bef3fa4f 100644 --- a/src/langtrace_python_sdk/instrumentation/__init__.py +++ b/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", ] diff --git a/src/langtrace_python_sdk/instrumentation/google_genai/__init__.py b/src/langtrace_python_sdk/instrumentation/google_genai/__init__.py new file mode 100644 index 00000000..e90722e8 --- /dev/null +++ b/src/langtrace_python_sdk/instrumentation/google_genai/__init__.py @@ -0,0 +1,3 @@ +from .instrumentation import GoogleGenaiInstrumentation + +__all__ = ["GoogleGenaiInstrumentation"] diff --git a/src/langtrace_python_sdk/instrumentation/google_genai/instrumentation.py b/src/langtrace_python_sdk/instrumentation/google_genai/instrumentation.py new file mode 100644 index 00000000..eeb92f47 --- /dev/null +++ b/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 diff --git a/src/langtrace_python_sdk/instrumentation/google_genai/patch.py b/src/langtrace_python_sdk/instrumentation/google_genai/patch.py new file mode 100644 index 00000000..7d97743a --- /dev/null +++ b/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, + }, + ) diff --git a/src/langtrace_python_sdk/langtrace.py b/src/langtrace_python_sdk/langtrace.py index 90bfbfcc..8077e605 100644 --- a/src/langtrace_python_sdk/langtrace.py +++ b/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(), diff --git a/src/langtrace_python_sdk/version.py b/src/langtrace_python_sdk/version.py index 4e499ef0..086040fa 100644 --- a/src/langtrace_python_sdk/version.py +++ b/src/langtrace_python_sdk/version.py @@ -1 +1 @@ -__version__ = "3.3.22" +__version__ = "3.3.23" diff --git a/src/run_example.py b/src/run_example.py index fcdde43a..80a8cf42 100644 --- a/src/run_example.py +++ b/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() From 82598a228a8df838beb418e08f7cf46c06bb3a78 Mon Sep 17 00:00:00 2001 From: Ali Waleed Date: Mon, 13 Jan 2025 17:58:02 +0200 Subject: [PATCH 2/2] remove type --- .../instrumentation/google_genai/patch.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/src/langtrace_python_sdk/instrumentation/google_genai/patch.py b/src/langtrace_python_sdk/instrumentation/google_genai/patch.py index 7d97743a..02cf9ef5 100644 --- a/src/langtrace_python_sdk/instrumentation/google_genai/patch.py +++ b/src/langtrace_python_sdk/instrumentation/google_genai/patch.py @@ -11,7 +11,6 @@ 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 @@ -72,9 +71,7 @@ def traced_method(wrapped, instance, args, kwargs): return traced_method -def set_streaming_response_attributes( - span: Span, response: Iterator[GenerateContentResponse] -): +def set_streaming_response_attributes(span: Span, response): accum_completion = "" for chunk in response: set_span_attribute( @@ -103,7 +100,7 @@ def set_streaming_response_attributes( set_event_completion(span, [{"role": "assistant", "content": accum_completion}]) -def set_response_attributes(span: Span, response: GenerateContentResponse): +def set_response_attributes(span: Span, response): completions = [] for candidate in response.candidates: set_span_attribute(