add a simple unit test

Author: keith-decker

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

@@ -18,14 +18,15 @@
 from uuid import UUID
 
 from opentelemetry._events import get_event_logger
-from opentelemetry._logs import get_logger
 from opentelemetry.metrics import get_meter
 from opentelemetry.semconv.schemas import Schemas
 from opentelemetry.trace import get_tracer
 
 from .data import ChatGeneration, Error, Message
 from .exporters import SpanMetricEventExporter, SpanMetricExporter
 from .types import LLMInvocation
+
+# TODO: Get the tool version for emitting spans, use GenAI Utils for now
 from .version import __version__
 
 
@@ -60,20 +61,11 @@ def __init__(self, exporter_type_full: bool = True, **kwargs):
             schema_url=Schemas.V1_28_0.value,
         )
 
-        logger_provider = kwargs.get("logger_provider")
-        self._logger = get_logger(
-            __name__,
-            __version__,
-            logger_provider=logger_provider,
-            schema_url=Schemas.V1_28_0.value,
-        )
-
         self._exporter = (
             SpanMetricEventExporter(
                 tracer=self._tracer,
                 meter=self._meter,
                 event_logger=self._event_logger,
-                logger=self._event_logger,
             )
             if exporter_type_full
             else SpanMetricExporter(tracer=self._tracer, meter=self._meter)
@@ -97,7 +89,7 @@ def start_llm(
         )
         with self._lock:
             self._llm_registry[invocation.run_id] = invocation
-        self._exporter.init_llm(invocation)
+        self._exporter.init(invocation)
 
     def stop_llm(
         self,
@@ -110,7 +102,7 @@ def stop_llm(
         invocation.end_time = time.time()
         invocation.chat_generations = chat_generations
         invocation.attributes.update(attributes)
-        self._exporter.export_llm(invocation)
+        self._exporter.export(invocation)
         return invocation
 
     def fail_llm(
@@ -119,8 +111,8 @@ def fail_llm(
         with self._lock:
             invocation = self._llm_registry.pop(run_id)
         invocation.end_time = time.time()
-        invocation.attributes.update(attributes)
-        self._exporter.error_llm(error, invocation)
+        invocation.attributes.update(**attributes)
+        self._exporter.error(error, invocation)
         return invocation
 
 

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

@@ -6,7 +6,6 @@ class Message:
     content: str
     type: str
     name: str
-    tool_call_id: str
 
 
 @dataclass

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

@@ -179,10 +179,9 @@ def _end_span(self, run_id: UUID):
         state = self.spans[run_id]
         for child_id in state.children:
             child_state = self.spans.get(child_id)
-            if child_state and child_state.span._end_time is None:
+            if child_state:
                 child_state.span.end()
-        if state.span._end_time is None:
-            state.span.end()
+        state.span.end()
 
     def init(self, invocation: LLMInvocation):
         if (
@@ -481,9 +480,8 @@ def export(self, invocation: LLMInvocation):
 
             for index, message in enumerate(invocation.messages):
                 content = message.content
-                message_type = message.type
                 span.set_attribute(f"gen_ai.prompt.{index}.content", content)
-                span.set_attribute(f"gen_ai.prompt.{index}.role", message_type)
+                span.set_attribute(f"gen_ai.prompt.{index}.role", message.type)
 
             for index, chat_generation in enumerate(
                 invocation.chat_generations

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

@@ -0,0 +1,78 @@
+from uuid import uuid4
+
+import pytest
+
+from opentelemetry import trace
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
+    InMemorySpanExporter,
+)
+from opentelemetry.util.genai.api import (
+    llm_start,
+    llm_stop,
+)
+from opentelemetry.util.genai.types import (
+    ChatGeneration,
+    Message,
+)
+
+
+@pytest.fixture
+def telemetry_setup():
+    """Set up telemetry providers for testing"""
+    # Set up in-memory span exporter to capture spans
+    memory_exporter = InMemorySpanExporter()
+    tracer_provider = TracerProvider()
+    tracer_provider.add_span_processor(SimpleSpanProcessor(memory_exporter))
+
+    # Set the tracer provider
+    trace.set_tracer_provider(tracer_provider)
+
+    yield memory_exporter
+
+    # Cleanup
+    memory_exporter.clear()
+    # Reset to default tracer provider
+    trace.set_tracer_provider(trace.NoOpTracerProvider())
+
+
+def test_llm_start_and_stop_creates_span(telemetry_setup):
+    memory_exporter = telemetry_setup
+
+    run_id = uuid4()
+    message = Message(content="hello world", type="Human", name="message name")
+    chat_generation = ChatGeneration(content="hello back", type="AI")
+
+    # Start and stop LLM invocation
+    llm_start(
+        [message], run_id=run_id, custom_attr="value", system="test-system"
+    )
+    invocation = llm_stop(
+        run_id, chat_generations=[chat_generation], extra="info"
+    )
+
+    # Get the spans that were created
+    spans = memory_exporter.get_finished_spans()
+
+    # Verify span was created
+    assert len(spans) == 1
+    span = spans[0]
+
+    # Verify span properties
+    assert span.name == "test-system.chat"
+    assert span.kind == trace.SpanKind.CLIENT
+
+    # Verify span attributes
+    assert span.attributes.get("gen_ai.operation.name") == "chat"
+    assert span.attributes.get("gen_ai.system") == "test-system"
+    # Add more attribute checks as needed
+
+    # Verify span timing
+    assert span.start_time > 0
+    assert span.end_time > span.start_time
+
+    # Verify invocation data
+    assert invocation.run_id == run_id
+    assert invocation.attributes.get("custom_attr") == "value"
+    assert invocation.attributes.get("extra") == "info"