Uploader class for _ref events

Author: aabmass

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

@@ -37,9 +37,10 @@ classifiers = [
   "Programming Language :: Python :: 3.12"
 ]
 dependencies = [
+  "fsspec>=2025.5.1",
   "opentelemetry-api >=1.31.1, <2",
   "opentelemetry-instrumentation >=0.52b1, <2",
-  "opentelemetry-semantic-conventions >=0.52b1, <2"
+  "opentelemetry-semantic-conventions >=0.52b1, <2",
 ]
 
 [project.optional-dependencies]

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

@@ -17,20 +17,18 @@
 import json
 import logging
 import os
-from re import S
 import time
 from typing import Any, AsyncIterator, Awaitable, Iterator, Optional, Union
+from uuid import uuid4
 
-from google.genai.models import AsyncModels, Models, t as transformers
+from google.genai.models import AsyncModels, Models
+from google.genai.models import t as transformers
 from google.genai.types import (
     BlockedReason,
     Candidate,
     Content,
     ContentListUnion,
-    GenerateContentConfig,
-    ContentUnion,
     ContentListUnionDict,
-    Content,
     ContentUnion,
     ContentUnionDict,
     GenerateContentConfig,
@@ -49,14 +47,13 @@
 from .custom_semconv import GCP_GENAI_OPERATION_CONFIG
 from .dict_util import flatten_dict
 from .flags import is_content_recording_enabled
-from .otel_wrapper import OTelWrapper
-from .tool_call_wrapper import wrapped as wrapped_tool
 from .message import (
-    ContentUnion,
     to_input_messages,
     to_output_message,
     to_system_instruction,
 )
+from .otel_wrapper import OTelWrapper
+from .tool_call_wrapper import wrapped as wrapped_tool
 
 _logger = logging.getLogger(__name__)
 
@@ -330,7 +327,7 @@ def process_completion(
         self._update_finish_reasons(response)
         self._maybe_update_token_counts(response)
         self._maybe_update_error_type(response)
-        self._maybe_log_response(response)
+        # self._maybe_log_response(response)
         self._response_index += 1
 
     def process_error(self, e: Exception):
@@ -441,6 +438,15 @@ def _transform_content(
             attributes=attributes,
         )
 
+        # Forward looking remote storage refs
+        self._otel_wrapper.log_completion_details_refs(
+            system_instructions=system_instruction,
+            input_messages=input_messages,
+            output_messages=output_message,
+            attributes=attributes,
+            response_id=response.response_id or str(uuid4()),
+        )
+
     def _maybe_log_system_instruction(
         self, config: Optional[GenerateContentConfigOrDict] = None
     ):

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

@@ -13,39 +13,28 @@
 # limitations under the License.
 
 import logging
+
+from google.genai.types import (
+    Candidate,
+    Content,
+    Part,
+)
+
 from .message_models import (
     BlobPart,
     ChatMessage,
     Choice,
     FileDataPart,
     InputMessages,
     MessagePart,
-    SystemMessage,
     OutputMessages,
+    SystemMessage,
     TextPart,
     ToolCallPart,
     ToolCallResponsePart,
     UnknownPart,
 )
 
-from google.genai.models import t as transfomers
-
-from google.genai.types import (
-    BlockedReason,
-    Candidate,
-    Content,
-    ContentListUnion,
-    ContentListUnionDict,
-    ContentUnion,
-    GenerateContentConfig,
-    ContentUnionDict,
-    GenerateContentConfig,
-    GenerateContentConfigOrDict,
-    GenerateContentResponse,
-    Part,
-    PartUnion,
-)
-
 _logger = logging.getLogger(__name__)
 
 
@@ -87,7 +76,9 @@ def _to_chat_message(
     content: Content,
 ) -> ChatMessage:
     parts = content.parts or []
-    return ChatMessage(role="system", parts=[_to_part(part) for part in parts])
+    return ChatMessage(
+        role=content.role, parts=[_to_part(part) for part in parts]
+    )
 
 
 def _to_part(part: Part) -> MessagePart:

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

@@ -83,7 +83,7 @@ class UnknownPart(BaseModel):
 
 
 class ChatMessage(BaseModel):
-    role: str
+    role: str | None
     parts: List[MessagePart]
 
 

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

@@ -17,7 +17,10 @@
 import google.genai
 
 from opentelemetry._events import Event
-from opentelemetry.util.types import _ExtendedAttributes
+from opentelemetry.instrumentation.google_genai.uploader import (
+    upload_to_storage,
+)
+from opentelemetry.util.types import _ExtendedAttributes, AnyValue
 from opentelemetry.semconv._incubating.metrics import gen_ai_metrics
 from opentelemetry.semconv.schemas import Schemas
 
@@ -96,7 +99,7 @@ def log_completion_details(
         # input_tokens: int,
         # output_tokens: int,
     ) -> None:
-        _logger.debug("Recording user prompt.")
+        _logger.debug("Recording completion details event.")
         event_name = "gen_ai.completion.details"
 
         body = {
@@ -111,6 +114,44 @@ def log_completion_details(
         event = Event(event_name, body=body, attributes=attributes)
         self._event_logger.emit(event)
 
+    def log_completion_details_refs(
+        self,
+        *,
+        attributes: _ExtendedAttributes,
+        system_instructions: SystemMessage | None,
+        input_messages: InputMessages,
+        output_messages: OutputMessages,
+        response_id: str,
+        # request_model: str,
+        # response_model: str,
+        # input_tokens: int,
+        # output_tokens: int,
+    ) -> None:
+        _logger.debug("Recording completion details event as ref.")
+        event_name = "gen_ai.completion.details"
+
+        input_messages_ref = upload_to_storage(
+            f"{response_id}_input.json",
+            input_messages.model_dump(mode="json"),
+        )
+        output_messages_ref = upload_to_storage(
+            f"{response_id}_output.json",
+            output_messages.model_dump(mode="json"),
+        )
+        body = {
+            "gen_ai.input.messages_ref": input_messages_ref,
+            "gen_ai.output.messages_ref": output_messages_ref,
+        }
+        if system_instructions:
+            system_instructions_ref = upload_to_storage(
+                f"{response_id}_system_instructions.json",
+                system_instructions.model_dump(mode="json"),
+            )
+            body["gen_ai.system.instructions_ref"] = system_instructions_ref
+
+        event = Event(event_name, body=body, attributes=attributes)
+        self._event_logger.emit(event)
+
     def log_response_content(self, attributes, body):
         _logger.debug("Recording response.")
         event_name = "gen_ai.choice"

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

@@ -0,0 +1,102 @@
+# 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.
+
+"""Generic uploading module which could be part of a GenAI utils
+
+Has a global uploader object using `fsspec` under the hood to work with all sorts of local and
+cloud storage.
+"""
+
+from abc import ABC, abstractmethod
+import io
+import json
+import logging
+from pathlib import Path
+
+from fsspec import AbstractFileSystem
+from fsspec.asyn import AsyncFileSystem, AbstractBufferedFile
+
+from opentelemetry.trace import get_tracer
+from opentelemetry.util._once import Once
+from opentelemetry.util.types import AnyValue
+from google.cloud import storage
+from google.cloud.storage import transfer_manager
+
+
+_logger = logging.getLogger(__name__)
+_tracer = get_tracer(__name__)
+
+
+# TODO: try to make async or do background buffering + async
+
+
+class Uploader(ABC):
+    @abstractmethod
+    def upload(self, *, path: Path, value: AnyValue) -> str:
+        pass
+
+
+class GcsUploader(Uploader):
+    def __init__(
+        self, *, bucket_path: str, client: storage.Client | None = None
+    ) -> None:
+        self._client = client or storage.Client()
+        self._bucket_path = bucket_path
+
+    def upload(self, *, path: Path, value: AnyValue) -> str:
+        blob = self._client.bucket(self._bucket_path).blob(str(path))
+        with blob.open("w") as file:
+            json.dump(value, file)
+        _logger.debug("GcsUploader uploaded to %s", file)
+        return f"gs://{blob.bucket.name}/{blob.name}"
+
+
+class FsSpecUploader(Uploader):
+    def __init__(self, *, fs: AbstractFileSystem, base_path: Path) -> None:
+        self._fs = fs
+        self._base_path = base_path
+
+    @_tracer.start_as_current_span("fsspec_upload")
+    def upload(self, *, path: Path, value: AnyValue) -> str:
+        with self._fs.open(self._base_path / path, "wb") as file:
+            # Encode to utf8 bytes while writing to the file
+            text_wrapper = io.TextIOWrapper(
+                file, encoding="utf-8", write_through=True
+            )
+            json.dump(value, text_wrapper)
+        _logger.debug("Uploaded to %s", file)
+        return file.full_name
+
+
+_uploader: Uploader | None = None
+_uploader_once: Once = Once()
+
+
+def _get_uploader() -> Uploader | None:
+    return _uploader
+
+
+def set_uploader(uploader: Uploader) -> None:
+    def do_set():
+        global _uploader
+        _uploader = uploader
+
+    _uploader_once.do_once(do_set)
+
+
+def upload_to_storage(filename: str, value: AnyValue) -> str:
+    """Given an AnyValue, uploads it to storage (local filesystem, GCS, S3, etc.)"""
+    if not (uploader := _get_uploader()):
+        return "/dev/null"
+    return uploader.upload(path=Path(filename), value=value)