Handle corner case where flatten function returns compound output.

Author: michaelsafyan

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

@@ -81,8 +81,12 @@ def _flatten_compound_value(
             rename_keys=rename_keys,
             flatten_functions=flatten_functions,
         )
-        if func_output is not None:
+        if func_output is None:
+            return {}
+        elif _is_primitive(func_output) or _is_homogenous_primitive_list(func_output):
             return {key: func_output}
+        else:
+            value = func_output
     if isinstance(value, dict):
         return _flatten_dict(
             value,

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

@@ -19,8 +19,8 @@
 
 class PydanticModel(BaseModel):
     """Used to verify handling of pydantic models in the flattener."""
-    str_value: str
-    int_value: int
+    str_value: str = ""
+    int_value: int = 0
 
 
 class ModelDumpableNotPydantic:
@@ -275,3 +275,41 @@ def test_flatten_list_of_compound_types():
         "list_value[3][0]": "abc",
         "list_value[3][1]": 123,
     }
+
+
+def test_handles_simple_output_from_flatten_func():
+    def f(*args, **kwargs):
+        return "baz"
+
+    input_dict = {
+        "foo": PydanticModel(),
+    }
+
+    output = dict_util.flatten_dict(
+        input_dict,
+        flatten_functions={"foo": f})
+
+    assert output == {
+        "foo": "baz",
+    }
+
+
+def test_handles_compound_output_from_flatten_func():
+    def f(*args, **kwargs):
+        return {
+            "baz": 123,
+            "qux": 456
+        }
+
+    input_dict = {
+        "foo": PydanticModel(),
+    }
+
+    output = dict_util.flatten_dict(
+        input_dict,
+        flatten_functions={"foo": f})
+
+    assert output == {
+        "foo.baz": 123,
+        "foo.qux": 456,
+    }