Stamp gen ai refs on spans and logs

Author: aabmass

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

@@ -18,14 +18,16 @@
 import logging
 import posixpath
 import threading
+from collections.abc import Mapping
 from concurrent.futures import Future, ThreadPoolExecutor
 from dataclasses import asdict, dataclass
 from functools import partial
 from os import environ
-from typing import Any, Callable, Literal, TextIO, cast
+from typing import Any, Callable, Final, Literal, TextIO, cast
 from uuid import uuid4
 
 from opentelemetry._logs import LogRecord
+from opentelemetry.semconv._incubating.attributes import gen_ai_attributes
 from opentelemetry.trace import Span
 from opentelemetry.util.genai import types
 from opentelemetry.util.genai.environment_variables import (
@@ -39,6 +41,18 @@
 except ImportError:
     fsspec = None
 
+
+GEN_AI_INPUT_MESSAGES_REF: Final = (
+    gen_ai_attributes.GEN_AI_INPUT_MESSAGES + "_ref"
+)
+GEN_AI_OUTPUT_MESSAGES_REF: Final = (
+    gen_ai_attributes.GEN_AI_OUTPUT_MESSAGES + "_ref"
+)
+GEN_AI_SYSTEM_INSTRUCTIONS_REF: Final = (
+    gen_ai_attributes.GEN_AI_SYSTEM_INSTRUCTIONS + "_ref"
+)
+
+
 _logger = logging.getLogger(__name__)
 
 
@@ -201,12 +215,43 @@ def to_dict(
                 },
             )
 
-            # TODO: stamp the refs on telemetry
+            # stamp the refs on telemetry
+            references = {
+                GEN_AI_INPUT_MESSAGES_REF: ref_names.inputs_ref,
+                GEN_AI_OUTPUT_MESSAGES_REF: ref_names.outputs_ref,
+                GEN_AI_SYSTEM_INSTRUCTIONS_REF: ref_names.system_instruction_ref,
+            }
+            if span:
+                self._set_in_span(span, references)
+            if log_record:
+                self._set_in_log(log_record, references)
 
         def shutdown(self) -> None:
             # TODO: support timeout
             self._executor.shutdown()
 
+        @staticmethod
+        def _set_in_span(span: Span, references: dict[str, str]) -> None:
+            # stamp the refs on telemetry
+            span.set_attributes(references)
+
+        @staticmethod
+        def _set_in_log(
+            log_record: LogRecord, references: dict[str, str]
+        ) -> None:
+            # set in both attributes and body until they are consolidated
+            # https://github.com/open-telemetry/semantic-conventions/issues/1870
+            log_record.attributes = {
+                **(log_record.attributes or {}),
+                **references,
+            }
+
+            if log_record.body is None or isinstance(log_record.body, Mapping):
+                log_record.body = {
+                    **(log_record.body or {}),
+                    **references,
+                }
+
     def fsspec_upload_hook() -> UploadHook:
         base_path = environ.get(OTEL_INSTRUMENTATION_GENAI_UPLOAD_BASE_PATH)
         if not base_path:

util/opentelemetry-util-genai/tests/test_fsspec_upload.py

@@ -25,8 +25,8 @@
 from unittest.mock import MagicMock, patch
 
 import fsspec
-from fsspec.implementations.memory import MemoryFileSystem
 
+from opentelemetry._logs import LogRecord
 from opentelemetry.test.test_base import TestBase
 from opentelemetry.util.genai import types
 from opentelemetry.util.genai._fsspec_upload import (
@@ -195,9 +195,6 @@ def test_upload(self):
 
 
 class TestFsspecUploadHookIntegration(TestBase):
-    def setUp(self):
-        MemoryFileSystem.store.clear()
-
     def assert_fsspec_equal(self, path: str, value: str) -> None:
         with fsspec.open(path, "r") as file:
             self.assertEqual(file.read(), value)
@@ -207,12 +204,87 @@ def test_upload_completions(self):
             uploader=FsspecUploader(),
             base_path=BASE_PATH,
         )
+        tracer = self.tracer_provider.get_tracer(__name__)
+        log_record = LogRecord()
+
+        with tracer.start_as_current_span("chat mymodel") as span:
+            hook.upload(
+                inputs=FAKE_INPUTS,
+                outputs=FAKE_OUTPUTS,
+                system_instruction=FAKE_SYSTEM_INSTRUCTION,
+                span=span,
+                log_record=log_record,
+            )
+        hook.shutdown()
+
+        finished_spans = self.get_finished_spans()
+        self.assertEqual(len(finished_spans), 1)
+        span = finished_spans[0]
+
+        # span attributes, log attributes, and log body have refs
+        for attributes in [
+            span.attributes,
+            log_record.attributes,
+            log_record.body,
+        ]:
+            for ref_key in [
+                "gen_ai.input.messages_ref",
+                "gen_ai.output.messages_ref",
+                "gen_ai.system_instructions_ref",
+            ]:
+                self.assertIn(ref_key, attributes)
+
+        self.assert_fsspec_equal(
+            span.attributes["gen_ai.input.messages_ref"],
+            '[{"role":"user","parts":[{"content":"What is the capital of France?","type":"text"}]}]',
+        )
+        self.assert_fsspec_equal(
+            span.attributes["gen_ai.output.messages_ref"],
+            '[{"role":"assistant","parts":[{"content":"Paris","type":"text"}],"finish_reason":"stop"}]',
+        )
+        self.assert_fsspec_equal(
+            span.attributes["gen_ai.system_instructions_ref"],
+            '[{"content":"You are a helpful assistant.","type":"text"}]',
+        )
+
+    @staticmethod
+    def upload_with_log(log_record: LogRecord):
+        hook = FsspecUploadHook(
+            uploader=FsspecUploader(),
+            base_path=BASE_PATH,
+        )
+
         hook.upload(
             inputs=FAKE_INPUTS,
             outputs=FAKE_OUTPUTS,
             system_instruction=FAKE_SYSTEM_INSTRUCTION,
+            log_record=log_record,
         )
-
-        fs = fsspec.open(BASE_PATH).fs
-        self.assertEqual(len(fs.ls(BASE_PATH)), 3)
-        # TODO: test stamped telemetry
+        hook.shutdown()
+
+    def test_stamps_empty_log(self):
+        log_record = LogRecord()
+        self.upload_with_log(log_record)
+
+        # stamp on both body and attributes
+        self.assertIn("gen_ai.input.messages_ref", log_record.attributes)
+        self.assertIn("gen_ai.output.messages_ref", log_record.attributes)
+        self.assertIn("gen_ai.system_instructions_ref", log_record.attributes)
+        self.assertIn("gen_ai.input.messages_ref", log_record.body)
+        self.assertIn("gen_ai.output.messages_ref", log_record.body)
+        self.assertIn("gen_ai.system_instructions_ref", log_record.body)
+
+    def test_stamps_log_with_map_body(self):
+        log_record = LogRecord(body={"hello": "world"})
+        self.upload_with_log(log_record)
+
+        # stamp on both body and attributes, preserving existing
+        self.assertEqual(log_record.body["hello"], "world")
+        self.assertIn("gen_ai.input.messages_ref", log_record.body)
+        self.assertIn("gen_ai.output.messages_ref", log_record.body)
+        self.assertIn("gen_ai.system_instructions_ref", log_record.body)
+
+    def test_doesnt_stamp_log_string_body(self):
+        log_record = LogRecord(body="hello world")
+        self.upload_with_log(log_record)
+        self.assertEqual(log_record.body, "hello world")