expose just a json util

Author: aabmass

instrumentation-genai/opentelemetry-instrumentation-vertexai/src/opentelemetry/instrumentation/vertexai/patch.py

@@ -14,7 +14,6 @@
 
 from __future__ import annotations
 
-import json
 from contextlib import contextmanager
 from dataclasses import asdict
 from typing import (
@@ -54,7 +53,7 @@
     InputMessage,
     OutputMessage,
 )
-from opentelemetry.util.genai.utils import Base64JsonEncoder
+from opentelemetry.util.genai.utils import gen_ai_json_dumps
 
 if TYPE_CHECKING:
     from google.cloud.aiplatform_v1.services.prediction_service import client
@@ -224,11 +223,7 @@ def handle_response(
                     ):
                         span.set_attributes(
                             {
-                                k: json.dumps(
-                                    v,
-                                    cls=Base64JsonEncoder,
-                                    separators=(",", ":"),
-                                )
+                                k: gen_ai_json_dumps(v)
                                 for k, v in content_attributes.items()
                             }
                         )

util/opentelemetry-util-genai/src/opentelemetry/util/genai/_upload/completion_hook.py

@@ -15,7 +15,6 @@
 
 from __future__ import annotations
 
-import json
 import logging
 import posixpath
 import threading
@@ -38,7 +37,7 @@
 from opentelemetry.util.genai.environment_variables import (
     OTEL_INSTRUMENTATION_GENAI_UPLOAD_FORMAT,
 )
-from opentelemetry.util.genai.utils import Base64JsonEncoder
+from opentelemetry.util.genai.utils import gen_ai_json_dump
 
 GEN_AI_INPUT_MESSAGES_REF: Final = (
     gen_ai_attributes.GEN_AI_INPUT_MESSAGES + "_ref"
@@ -192,12 +191,7 @@ def _do_upload(
 
         with self._fs.open(path, "w", content_type=content_type) as file:
             for message in message_lines:
-                json.dump(
-                    message,
-                    file,
-                    separators=(",", ":"),
-                    cls=Base64JsonEncoder,
-                )
+                gen_ai_json_dump(message, file)
                 file.write("\n")
 
     def on_completion(

util/opentelemetry-util-genai/src/opentelemetry/util/genai/utils.py

@@ -16,6 +16,7 @@
 import logging
 import os
 from base64 import b64encode
+from functools import partial
 from typing import Any
 
 from opentelemetry.instrumentation._semconv import (
@@ -59,10 +60,24 @@ def get_content_capturing_mode() -> ContentCapturingMode:
         return ContentCapturingMode.NO_CONTENT
 
 
-class Base64JsonEncoder(json.JSONEncoder):
-    """Should be used to serialize python objects to json that may contain bytes."""
+class _GenAiJsonEncoder(json.JSONEncoder):
+    """Should be used by GenAI instrumentations when serializing objects that may contain
+    bytes, datetimes, etc. for GenAI observability."""
 
     def default(self, o: Any) -> Any:
         if isinstance(o, bytes):
             return b64encode(o).decode()
         return super().default(o)
+
+
+gen_ai_json_dump = partial(
+    json.dump, separators=(",", ":"), cls=_GenAiJsonEncoder
+)
+"""Should be used by GenAI instrumentations when serializing objects that may contain
+bytes, datetimes, etc. for GenAI observability."""
+
+gen_ai_json_dumps = partial(
+    json.dumps, separators=(",", ":"), cls=_GenAiJsonEncoder
+)
+"""Should be used by GenAI instrumentations when serializing objects that may contain
+bytes, datetimes, etc. for GenAI observability."""