Create a utility to simplify recording request attributes.

Author: michaelsafyan

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

@@ -0,0 +1,146 @@
+# 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.
+
+
+from typing import Dict, Optional, Set, Any, Callable, TypeAlias, Union, List
+
+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]
+
+
+def _concat_key(prefix: Optional[str], suffix: str):
+    if not prefix:
+        return suffix
+    return f"{prefix}.{suffix}"
+
+
+def _is_primitive(v):
+    for t in [str, bool, int, float]:
+        if isinstance(v, t):
+            return True
+    return False
+
+
+def _is_homogenous_primitive_list(v):
+    if not isinstance(v, list):
+        return False
+    if len(v) == 0:
+        return True
+    if not _is_primitive(v[0]):
+        return False
+    t = type(v[0])
+    for entry in v[1:]:
+        if not isinstance(entry, t):
+            return False
+    return True
+
+
+def _get_flatten_func(
+    flatten_functions: Dict[str, Callable],
+    key_names: set[str]):
+    for key in key_names:
+        flatten_func = flatten_functions.get(key)
+        if flatten_func is not None:
+            return flatten_func
+    return None
+
+
+def _flatten_value(
+    key: str,
+    value: Any,
+    exclude_keys: Set[str],
+    rename_keys: Dict[str, str],
+    flatten_functions: Dict[str, Callable],
+    _from_json=False) -> FlattenedDict:
+  key_names = set([key])
+  renamed_key = rename_keys.get(key)
+  if renamed_key is not None:
+    key_names.add(renamed_key)
+    key = renamed_key
+  if key_names & exclude_keys:
+    return {}
+  if _is_primitive(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)
+  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):
+    if _is_homogenous_primitive_list(value):
+        return {key: value}
+    return _flatten_list(value, key_prefix=key, exclude_keys=exclude_keys, rename_keys=rename_keys, flatten_functions=flatten_functions)
+  if hasattr(value, "model_dump"):
+    d = value.model_dump()
+    return _flatten_dict(d, key_prefix=key, exclude_keys=exclude_keys, rename_keys=rename_keys, flatten_functions=flatten_functions)
+  if _from_json:
+    raise ValueError(f"Cannot flatten value with key {key}; value: {value}")
+  json_value = json.loads(json.dumps(value))
+  return _flatten_value(key, json_value, exclude_keys=exclude_keys, rename_keys=rename_keys, flatten_functions=flatten_functions, _from_json=True)
+
+
+def _flatten_dict(
+    d: Dict[str, Any],
+    key_prefix: str,
+    exclude_keys: Set[str],
+    rename_keys: Dict[str, str],
+    flatten_functions: Dict[str, Callable]) -> FlattenedDict:
+ result = {}
+ for key, value in d.items():
+    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)
+ return result
+
+
+
+def _flatten_list(
+    lst: list[Any],
+    key_prefix: str,
+    exclude_keys: Set[str],
+    rename_keys: Dict[str, str],
+    flatten_functions: Dict[str, Callable]) -> FlattenedDict:
+ result = {}
+ result[_concat_key(key_prefix, "length")] = len(lst)
+ for index, value in enumerate(lst):
+    full_key = f"{key_prefix}[{index}]"
+    flattened = _flatten_value(full_key, value, exclude_keys=exclude_keys, rename_keys=rename_keys, flatten_functions=flatten_functions)
+    result.update(flattened)
+ return result
+
+
+
+def flatten_dict(
+    d: Dict[str, Any],
+    key_prefix: Optional[str] = None,
+    exclude_keys: Optional[Union[List[str]|Set[str]]] = None,
+    rename_keys: Optional[Dict[str, str]] = None,
+    flatten_functions: Optional[Dict[str, Callable]] = None):
+  key_prefix = key_prefix or ""
+  if exclude_keys is None:
+    exclude_keys = set()
+  elif isinstance(exclude_keys, list):
+    exclude_keys = set(exclude_keys)
+  rename_keys = rename_keys 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/tests/utils/test_dict_util.py

@@ -0,0 +1,82 @@
+# 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.
+
+
+from opentelemetry.instrumentation.google_genai import dict_util
+
+
+def test_flatten_empty_dict():
+    d = {}
+    assert dict_util.flatten_dict(d) == d
+
+def test_flatten_simple_dict():
+    d = {
+        "int_key": 1,
+        "string_key": "somevalue",
+        "float_key": 3.14,
+        "bool_key": True
+    }
+    assert dict_util.flatten_dict(d) == d
+
+
+def test_flatten_with_key_exclusion():
+    d = {
+        "int_key": 1,
+        "string_key": "somevalue",
+        "float_key": 3.14,
+        "bool_key": True
+    }
+    output = dict_util.flatten_dict(d, exclude_keys=["int_key"])
+    assert "int_key" not in output
+    assert output == {
+        "string_key": "somevalue",
+        "float_key": 3.14,
+        "bool_key": True
+    }
+
+
+def test_flatten_with_renaming():
+    d = {
+        "int_key": 1,
+        "string_key": "somevalue",
+        "float_key": 3.14,
+        "bool_key": True
+    }
+    output = dict_util.flatten_dict(
+        d,
+        rename_keys={"float_key": "math_key"})
+    assert "float_key" not in output
+    assert "math_key" in output
+    assert output == {
+        "int_key": 1,
+        "string_key": "somevalue",
+        "math_key": 3.14,
+        "bool_key": True
+    }
+
+
+def test_flatten_with_prefixing():
+    d = {
+        "int_key": 1,
+        "string_key": "somevalue",
+        "float_key": 3.14,
+        "bool_key": True
+    }
+    output = dict_util.flatten_dict(d, key_prefix="someprefix")
+    assert output == {
+        "someprefix.int_key": 1,
+        "someprefix.string_key": "somevalue",
+        "someprefix.float_key": 3.14,
+        "someprefix.bool_key": True
+    }