Improve data model and add some rudimentary type checking.

Author: michaelsafyan

instrumentation-genai/opentelemetry-instrumentation-google-genai/Makefile

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-.PHONY: all test install clean lint
+.PHONY: all test install clean lint typecheck
 
-all: build test lint
+all: build test lint typecheck
 
 test:
 	./tools/test.sh
@@ -27,6 +27,8 @@ build: ./dist/opentelemetry_instrumentation_google_genai-0.0.1-py3-none-any.whl
 install: ./dist/opentelemetry_instrumentation_google_genai-0.0.1-py3-none-any.whl
 	pip install ./dist/opentelemetry_instrumentation_google_genai-0.0.1-py3-none-any.whl
 
+typecheck:
+	./tools/typecheck.sh
 
 lint:
 	./tools/lint.sh

instrumentation-genai/opentelemetry-instrumentation-google-genai/pyproject.toml

@@ -65,3 +65,15 @@ include = [
 
 [tool.hatch.build.targets.wheel]
 packages = ["src/opentelemetry"]
+
+[tool.pyright]
+include = [
+  "src",
+]
+exclude = [
+  "**/__pycache__",
+]
+stubPath = "types"
+reportMissingImports = "error"
+reportMissingTypeStubs = false
+pythonVersion = "3.9"

instrumentation-genai/opentelemetry-instrumentation-google-genai/src/opentelemetry/instrumentation/google_genai/generate_content.py

@@ -12,11 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import asyncio
 import functools
 import logging
 import os
 import time
-from typing import AsyncIterator, Awaitable, Iterator, Optional, Union
+from typing import Any, AsyncIterator, Awaitable, Iterator, Optional, Union
 
 from google.genai.models import AsyncModels, Models
 from google.genai.types import (
@@ -29,6 +30,7 @@
 
 from opentelemetry import trace
 from opentelemetry.semconv._incubating.attributes import gen_ai_attributes
+from opentelemetry.semconv._incubating.attributes import code_attributes
 from opentelemetry.semconv.attributes import error_attributes
 
 from .flags import is_content_recording_enabled
@@ -120,11 +122,11 @@ def _determine_genai_system(models_object: Union[Models, AsyncModels]):
 
 def _get_config_property(
     config: Optional[GenerateContentConfigOrDict],
-    path: str):
+    path: str) -> Any:
     if config is None:
         return None
     path_segments = path.split(".")
-    current_context = config
+    current_context: Any = config
     for path_segment in path_segments:
         if current_context is None:
             return None
@@ -186,11 +188,12 @@ def __init__(
         self._output_tokens = 0
         self._content_recording_enabled = is_content_recording_enabled()
 
-    def start_span_as_current_span(self, name):
+    def start_span_as_current_span(self, model_name, function_name):
         return self._otel_wrapper.start_as_current_span(
-            name,
+            f'generate_content [{model_name}]',
             start_time=self._start_time,
             attributes={
+                code_attributes.CODE_FUNCTION_NAME: function_name,
                 gen_ai_attributes.GEN_AI_SYSTEM: self._genai_system,
                 gen_ai_attributes.GEN_AI_REQUEST_MODEL: self._genai_request_model,
                 gen_ai_attributes.GEN_AI_OPERATION_NAME: "GenerateContent",
@@ -230,9 +233,9 @@ def finalize_processing(self):
     def _maybe_update_token_counts(self, response: GenerateContentResponse):
         input_tokens = _get_response_property(response, "usage_metadata.prompt_token_count")
         output_tokens = _get_response_property(response, "usage_metadata.candidates_token_count")
-        if input_tokens:
+        if input_tokens and isinstance(input_tokens, int):
             self._input_tokens += input_tokens
-        if output_tokens:
+        if output_tokens and isinstance(output_tokens, int):
             self._output_tokens += output_tokens
 
     def _maybe_update_error_type(self, response: GenerateContentResponse):
@@ -331,7 +334,7 @@ def instrumented_generate_content(
         contents: Union[ContentListUnion, ContentListUnionDict],
         config: Optional[GenerateContentConfigOrDict] = None) -> GenerateContentResponse:
         helper = _GenerateContentInstrumentationHelper(self, otel_wrapper, model)
-        with helper.start_span_as_current_span("google.genai.Models.generate_content"):
+        with helper.start_span_as_current_span(model, "google.genai.Models.generate_content"):
             helper.process_request(contents, config)
             try:
                 response = wrapped_func(self, model=model, contents=contents, config=config)
@@ -360,7 +363,7 @@ def instrumented_generate_content_stream(
         contents: Union[ContentListUnion, ContentListUnionDict],
         config: Optional[GenerateContentConfigOrDict] = None) -> Iterator[GenerateContentResponse]:
         helper = _GenerateContentInstrumentationHelper(self, otel_wrapper, model)
-        with helper.start_span_as_current_span("google.genai.Models.generate_content_stream"):
+        with helper.start_span_as_current_span(model, "google.genai.Models.generate_content_stream"):
             helper.process_request(contents, config)
             try:
                 for response in wrapped_func(self, model=model, contents=contents, config=config):
@@ -389,7 +392,7 @@ async def instrumented_generate_content(
         contents: Union[ContentListUnion, ContentListUnionDict],
         config: Optional[GenerateContentConfigOrDict] = None) -> GenerateContentResponse:
         helper = _GenerateContentInstrumentationHelper(self, otel_wrapper, model)
-        with helper.start_span_as_current_span("google.genai.AsyncModels.generate_content"):
+        with helper.start_span_as_current_span(model, "google.genai.AsyncModels.generate_content"):
             helper.process_request(contents, config)
             try:
                 response = await wrapped_func(self, model=model, contents=contents, config=config)
@@ -403,7 +406,8 @@ async def instrumented_generate_content(
     return instrumented_generate_content
 
 
-def _create_instrumented_async_generate_content_stream(
+# Disabling type checking because this is not yet implemented and tested fully.
+def _create_instrumented_async_generate_content_stream(  # pyright: ignore
     snapshot: _MethodsSnapshot,
     otel_wrapper: OTelWrapper):
     wrapped_func = snapshot.async_generate_content_stream
@@ -416,14 +420,14 @@ async def instrumented_generate_content_stream(
         *,
         model: str,
         contents: Union[ContentListUnion, ContentListUnionDict],
-        config: Optional[GenerateContentConfigOrDict] = None) -> Awaitable[AsyncIterator[GenerateContentResponse]]:
+        config: Optional[GenerateContentConfigOrDict] = None) -> Awaitable[AsyncIterator[GenerateContentResponse]]:  # pyright: ignore
         helper = _GenerateContentInstrumentationHelper(self, otel_wrapper, model)
-        with helper.start_span_as_current_span("google.genai.AsyncModels.generate_content_stream"):
+        with helper.start_span_as_current_span(model, "google.genai.AsyncModels.generate_content_stream"):
             helper.process_request(contents, config)
             try:
-                async for response in wrapped_func(self, model=model, contents=contents, config=config):
+                async for response in await wrapped_func(self, model=model, contents=contents, config=config):  # pyright: ignore
                     helper.process_response(response)
-                    yield response
+                    yield response  # pyright: ignore
             except Exception as error:
                 helper.process_error(error)
                 raise

instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/common/otel_mocker.py

@@ -168,7 +168,7 @@ def get_span_named(self, name):
 
     def assert_has_span_named(self, name):
         span = self.get_span_named(name)
-        finished_spans = self.get_finished_spans()
+        finished_spans = [span.name for span in self.get_finished_spans()]
         assert span is not None, f'Could not find span named "{name}"; finished spans: {finished_spans}'
 
     def get_event_named(self, event_name):

instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/generate_content/nonstreaming_base.py

@@ -56,8 +56,8 @@ def setUp(self):  # pylint: disable=invalid-name
     def generate_content(self, *args, **kwargs):
         raise NotImplementedError("Must implement 'generate_content'.")
 
-    def expected_span_name(self):
-        raise NotImplementedError("Must implement 'expected_span_name'.")
+    def expected_function_name(self):
+        raise NotImplementedError("Must implement 'expected_function_name'.")
 
     def configure_valid_response(self, response_text="The model_response", input_tokens=10, output_tokens=20):
         self.requests.add_response(create_valid_response(
@@ -78,26 +78,43 @@ def test_generates_span(self):
             model="gemini-2.0-flash",
             contents="Does this work?")
         self.assertEqual(response.text, "Yep, it works!")
-        self.otel.assert_has_span_named(self.expected_span_name)
+        self.otel.assert_has_span_named("generate_content [gemini-2.0-flash]")
+
+    def test_model_reflected_into_span_name(self):
+        self.configure_valid_response(response_text="Yep, it works!")
+        response = self.generate_content(
+            model="gemini-1.5-flash",
+            contents="Does this work?")
+        self.assertEqual(response.text, "Yep, it works!")
+        self.otel.assert_has_span_named("generate_content [gemini-1.5-flash]")
 
     def test_generated_span_has_minimal_genai_attributes(self):
         self.configure_valid_response(response_text="Yep, it works!")
         self.generate_content(
             model="gemini-2.0-flash",
             contents="Does this work?")
-        self.otel.assert_has_span_named(self.expected_span_name)
-        span = self.otel.get_span_named(self.expected_span_name)
+        self.otel.assert_has_span_named("generate_content [gemini-2.0-flash]")
+        span = self.otel.get_span_named("generate_content [gemini-2.0-flash]")
         self.assertEqual(span.attributes["gen_ai.system"], "gemini")
         self.assertEqual(span.attributes["gen_ai.operation.name"], "GenerateContent")
 
+    def test_generated_span_has_correct_function_name(self):
+        self.configure_valid_response(response_text="Yep, it works!")
+        self.generate_content(
+            model="gemini-2.0-flash",
+            contents="Does this work?")
+        self.otel.assert_has_span_named("generate_content [gemini-2.0-flash]")
+        span = self.otel.get_span_named("generate_content [gemini-2.0-flash]")
+        self.assertEqual(span.attributes["code.function.name"], self.expected_function_name)
+
     def test_generated_span_has_vertex_ai_system_when_configured(self):
         self.set_use_vertex(True)
         self.configure_valid_response(response_text="Yep, it works!")
         self.generate_content(
             model="gemini-2.0-flash",
             contents="Does this work?")
-        self.otel.assert_has_span_named(self.expected_span_name)
-        span = self.otel.get_span_named(self.expected_span_name)
+        self.otel.assert_has_span_named("generate_content [gemini-2.0-flash]")
+        span = self.otel.get_span_named("generate_content [gemini-2.0-flash]")
         self.assertEqual(span.attributes["gen_ai.system"], "vertex_ai")
         self.assertEqual(span.attributes["gen_ai.operation.name"], "GenerateContent")
 
@@ -106,8 +123,8 @@ def test_generated_span_counts_tokens(self):
         self.generate_content(
             model="gemini-2.0-flash",
             contents="Some input")
-        self.otel.assert_has_span_named(self.expected_span_name)
-        span = self.otel.get_span_named(self.expected_span_name)
+        self.otel.assert_has_span_named("generate_content [gemini-2.0-flash]")
+        span = self.otel.get_span_named("generate_content [gemini-2.0-flash]")
         self.assertEqual(span.attributes["gen_ai.usage.input_tokens"], 123)
         self.assertEqual(span.attributes["gen_ai.usage.output_tokens"], 456)
 

instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/generate_content/test_sync_nonstreaming.py

@@ -25,7 +25,7 @@ def generate_content(self, *args, **kwargs):
         return self.client.models.generate_content(*args, **kwargs)
 
     @property
-    def expected_span_name(self):
+    def expected_function_name(self):
         return "google.genai.Models.generate_content"
 
 

instrumentation-genai/opentelemetry-instrumentation-google-genai/tools/typecheck.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+SCRIPT_DIR=$(cd $(dirname "${BASH_SOURCE:-$0}"); pwd)
+PROJECT_DIR=$(readlink -f "${SCRIPT_DIR}/..")
+TESTS_DIR="${PROJECT_DIR}/tests"
+REQUIREMENTS_FILE="${TESTS_DIR}/requirements.txt"
+TYPECHECK_ENV="${PROJECT_DIR}/.test/.typecheck-venv"
+
+function main() {
+    if [ ! -d "${TYPECHECK_ENV}" ] ; then
+      mkdir -p "${TYPECHECK_ENV}" || exit 1
+    fi
+    if [ ! -e "${TYPECHECK_ENV}/bin/activate" ] ; then
+      python3 -m venv "${TYPECHECK_ENV}" || exit 1
+    fi
+    source "${TYPECHECK_ENV}/bin/activate" || exit 1
+    pip install pyright || exit 1
+    pip install -r "${REQUIREMENTS_FILE}" || exit 1
+    cd "${PROJECT_DIR}" || exit 1
+    pyright --venvpath "${TYPECHECK_ENV}" || exit $?
+}
+
+main