Update recording mechanism to record more request options.

Author: michaelsafyan

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

@@ -0,0 +1,18 @@
+# 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.
+
+
+# Prefix to use for LLM model request attributes that are unique GCP
+# (or that have not yet been formally defined in the GenAI/LLM SIG).
+CUSTOM_LLM_REQUEST_PREFIX = "gen_ai.gcp.request"

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

@@ -13,20 +13,18 @@
 # limitations under the License.
 
 
-from typing import Dict, Optional, Set, Any, Callable, TypeAlias, Union, List
+from typing import Dict, Optional, Set, Any, Callable, Union, List, Sequence
 
 import json
 
-Primitive: TypeAlias = Union[bool, str, int, float]
-
-BoolList: TypeAlias = list[bool]
-StringList: TypeAlias = list[str]
-IntList: TypeAlias = list[int]
-FloatList: TypeAlias = list[float]
-HomogenousPrimitiveList: TypeAlias = Union[BoolList, StringList, IntList, FloatList]
-
-FlattenedValue: TypeAlias = Union[Primitive, HomogenousPrimitiveList]
-FlattenedDict: TypeAlias = Dict[str, FlattenedValue]
+Primitive = Union[bool, str, int, float]
+BoolList = list[bool]
+StringList = list[str]
+IntList = list[int]
+FloatList = list[float]
+HomogenousPrimitiveList = Union[BoolList, StringList, IntList, FloatList]
+FlattenedValue = Union[Primitive, HomogenousPrimitiveList]
+FlattenedDict = Dict[str, FlattenedValue]
 
 
 def _concat_key(prefix: Optional[str], suffix: str):
@@ -84,7 +82,9 @@ def _flatten_value(
     return {key: value}
   flatten_func = _get_flatten_func(flatten_functions, key_names)
   if flatten_func is not None:
-    return flatten_func(key, value, exclude_keys=exclude_keys, rename_keys=rename_keys, flatten_functions=flatten_functions)
+    func_output = flatten_func(key, value, exclude_keys=exclude_keys, rename_keys=rename_keys, flatten_functions=flatten_functions)
+    if func_output is not None:
+      return {key: func_output}
   if isinstance(value, dict):
     return _flatten_dict(value, key_prefix=key, exclude_keys=exclude_keys, rename_keys=rename_keys, flatten_functions=flatten_functions)
   if isinstance(value, list):
@@ -108,6 +108,8 @@ def _flatten_dict(
     flatten_functions: Dict[str, Callable]) -> FlattenedDict:
  result = {}
  for key, value in d.items():
+    if key in exclude_keys:
+      continue
     full_key = _concat_key(key_prefix, key)
     flattened = _flatten_value(full_key, value, exclude_keys=exclude_keys, rename_keys=rename_keys, flatten_functions=flatten_functions)
     result.update(flattened)
@@ -134,7 +136,7 @@ def _flatten_list(
 def flatten_dict(
     d: Dict[str, Any],
     key_prefix: Optional[str] = None,
-    exclude_keys: Optional[Union[List[str]|Set[str]]] = None,
+    exclude_keys: Optional[Sequence[str]] = None,
     rename_keys: Optional[Dict[str, str]] = None,
     flatten_functions: Optional[Dict[str, Callable]] = None):
   key_prefix = key_prefix or ""
@@ -143,4 +145,5 @@ def flatten_dict(
   elif isinstance(exclude_keys, list):
     exclude_keys = set(exclude_keys)
   rename_keys = rename_keys or {}
+  flatten_functions = flatten_functions or {}
   return _flatten_dict(d, key_prefix=key_prefix, exclude_keys=exclude_keys, rename_keys=rename_keys, flatten_functions=flatten_functions)

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

@@ -41,6 +41,8 @@
 
 from .flags import is_content_recording_enabled
 from .otel_wrapper import OTelWrapper
+from .dict_util import flatten_dict
+from .custom_semconv import CUSTOM_LLM_REQUEST_PREFIX
 
 _logger = logging.getLogger(__name__)
 
@@ -129,21 +131,45 @@ def _determine_genai_system(models_object: Union[Models, AsyncModels]):
     return _get_gemini_system_name()
 
 
-def _get_config_property(
-    config: Optional[GenerateContentConfigOrDict], path: str
-) -> Any:
+def _to_dict(value: object):
+    if isinstance(value, dict):
+        return value
+    if hasattr(value, "model_dump"):
+        return value.model_dump()
+    return json.loads(json.dumps(value))
+
+
+def _add_request_options_to_span(span, config: Optional[GenerateContentConfigOrDict]):
     if config is None:
-        return None
-    path_segments = path.split(".")
-    current_context: Any = config
-    for path_segment in path_segments:
-        if current_context is None:
-            return None
-        if isinstance(current_context, dict):
-            current_context = current_context.get(path_segment)
-        else:
-            current_context = getattr(current_context, path_segment)
-    return current_context
+        return
+    span_context = span.get_span_context()
+    if not span_context.trace_flags.sampled:
+        return
+    attributes = flatten_dict(
+        _to_dict(config),
+        key_prefix=CUSTOM_LLM_REQUEST_PREFIX,
+        exclude_keys=[
+            # System instruction can be overly long for a span attribute.
+            # Additionally, it is recorded as an event (log), instead.
+            "gen_ai.gcp.request.system_instruction",
+            # Headers could include sensitive information, therefore it is
+            # best that we not record these options.
+            "gen_ai.gcp.request.http_options.headers",
+        ],
+        rename_keys={
+            # TODO: add more entries here as more semantic conventions are
+            # generalized to cover more of the available config options.
+            "gen_ai.gcp.request.temperature": "gen_ai.request.temperature",
+            "gen_ai.gcp.request.top_k": "gen_ai.request.top_k",
+            "gen_ai.gcp.request.top_p": "gen_ai.request.top_p",
+            "gen_ai.gcp.request.candidate_count": "gen_ai.request.choice.count",
+            "gen_ai.gcp.request.max_output_tokens": "gen_ai.request.max_tokens",
+            "gen_ai.gcp.request.stop_sequences": "gen_ai.request.stop_sequences",
+            "gen_ai.gcp.request.presence_penalty": "gen_ai.request.presence_penalty",
+            "gen_ai.gcp.request.seed": "gen_ai.request.seed",
+        }
+    )
+    span.set_attributes(attributes)
 
 
 def _get_response_property(response: GenerateContentResponse, path: str):
@@ -159,44 +185,6 @@ def _get_response_property(response: GenerateContentResponse, path: str):
     return current_context
 
 
-def _get_temperature(config: Optional[GenerateContentConfigOrDict]):
-    return _get_config_property(config, "temperature")
-
-
-def _get_top_k(config: Optional[GenerateContentConfigOrDict]):
-    return _get_config_property(config, "top_k")
-
-
-def _get_top_p(config: Optional[GenerateContentConfigOrDict]):
-    return _get_config_property(config, "top_p")
-
-
-# A map from define attributes to the function that can obtain
-# the relevant information from the request object.
-#
-# TODO: expand this to cover a larger set of the available
-# span attributes from GenAI semantic conventions.
-#
-# TODO: define semantic conventions for attributes that
-# are relevant for the Google GenAI SDK which are not
-# currently covered by the existing semantic conventions.
-#
-# See also: TODOS.md
-_SPAN_ATTRIBUTE_TO_CONFIG_EXTRACTOR = {
-    gen_ai_attributes.GEN_AI_REQUEST_TEMPERATURE: _get_temperature,
-    gen_ai_attributes.GEN_AI_REQUEST_TOP_K: _get_top_k,
-    gen_ai_attributes.GEN_AI_REQUEST_TOP_P: _get_top_p,
-}
-
-
-def _to_dict(value: object):
-    if isinstance(value, dict):
-        return value
-    if hasattr(value, "model_dump"):
-        return value.model_dump()
-    return json.loads(json.dumps(value))
-
-
 class _GenerateContentInstrumentationHelper:
     def __init__(
         self,
@@ -237,13 +225,7 @@ def process_request(
         config: Optional[GenerateContentConfigOrDict],
     ):
         span = trace.get_current_span()
-        for (
-            attribute_key,
-            extractor,
-        ) in _SPAN_ATTRIBUTE_TO_CONFIG_EXTRACTOR.items():
-            attribute_value = extractor(config)
-            if attribute_value is not None:
-                span.set_attribute(attribute_key, attribute_value)
+        _add_request_options_to_span(span, config)
         self._maybe_log_system_instruction(config=config)
         self._maybe_log_user_prompt(contents)
 
@@ -317,7 +299,12 @@ def _maybe_update_error_type(self, response: GenerateContentResponse):
     def _maybe_log_system_instruction(
         self, config: Optional[GenerateContentConfigOrDict] = None
     ):
-        system_instruction = _get_config_property(config, "system_instruction")
+        system_instruction = None
+        if config is not None:
+            if isinstance(config, dict):
+                system_instruction = config["system_instruction"]
+            else:
+                system_instruction = config.system_instruction
         if not system_instruction:
             return
         attributes = {

instrumentation-genai/opentelemetry-instrumentation-google-genai/tests/utils/test_dict_util.py

@@ -20,6 +20,7 @@ def test_flatten_empty_dict():
     d = {}
     assert dict_util.flatten_dict(d) == d
 
+
 def test_flatten_simple_dict():
     d = {
         "int_key": 1,
@@ -30,6 +31,31 @@ def test_flatten_simple_dict():
     assert dict_util.flatten_dict(d) == d
 
 
+def test_flatten_nested_dict():
+    d = {
+        "int_key": 1,
+        "string_key": "somevalue",
+        "float_key": 3.14,
+        "bool_key": True,
+        "object_key": {
+            "nested": {
+                "foo": 1,
+                "bar": "baz",
+            },
+            "qux": 54321
+        }
+    }
+    assert dict_util.flatten_dict(d) == {
+        "int_key": 1,
+        "string_key": "somevalue",
+        "float_key": 3.14,
+        "bool_key": True,
+        "object_key.nested.foo": 1,
+        "object_key.nested.bar": "baz",
+        "object_key.qux": 54321,
+    }
+
+
 def test_flatten_with_key_exclusion():
     d = {
         "int_key": 1,
@@ -80,3 +106,30 @@ def test_flatten_with_prefixing():
         "someprefix.float_key": 3.14,
         "someprefix.bool_key": True
     }
+
+
+def test_flatten_with_custom_flatten_func():
+    def summarize_int_list(key, value, **kwargs):
+        total = 0
+        for item in value:
+            total += item
+        avg = total / len(value)
+        return f"{len(value)} items (total: {total}, average: {avg})"
+    flatten_functions = {
+         "some.deeply.nested.key": summarize_int_list
+    }
+    d = {
+        "some": {
+            "deeply": {
+                "nested": {
+                    "key": [1, 2, 3, 4, 5, 6, 7, 8, 9],
+                },
+            },
+        },
+        "other": [1, 2, 3, 4, 5, 6, 7, 8, 9]
+    }
+    output = dict_util.flatten_dict(d, flatten_functions=flatten_functions)
+    assert output == {
+        "some.deeply.nested.key": "9 items (total: 45, average: 5.0)",
+        "other": [1, 2, 3, 4, 5, 6, 7, 8, 9],
+    }