add on chunk events to genai spans (#200)

* add on chunk events to genai spans

* small update in doc comment

* wrap add event in try except

Author: dinmukhamedm

src/lmnr/opentelemetry_lib/opentelemetry/instrumentation/google_genai/__init__.py

@@ -384,6 +384,10 @@ def _build_from_streaming_response(
     aggregated_usage_metadata = defaultdict(int)
     model_version = None
     for chunk in response:
+        try:
+            span.add_event("llm.content.completion.chunk")
+        except Exception:
+            pass
         # Important: do all processing in a separate sync function, that is
         # wrapped in @dont_throw. If we did it here, the @dont_throw on top of
         # this function would not be able to catch the errors, as they are
@@ -434,6 +438,10 @@ async def _abuild_from_streaming_response(
     aggregated_usage_metadata = defaultdict(int)
     model_version = None
     async for chunk in response:
+        try:
+            span.add_event("llm.content.completion.chunk")
+        except Exception:
+            pass
         # Important: do all processing in a separate sync function, that is
         # wrapped in @dont_throw. If we did it here, the @dont_throw on top of
         # this function would not be able to catch the errors, as they are

src/lmnr/sdk/laminar.py

@@ -757,12 +757,12 @@ def set_span_attributes(
         instrumentation.
         Example:
         ```python
-        with L.start_as_current_span(
+        with Laminar.start_as_current_span(
             name="my_span_name", input=input["messages"], span_type="LLM"
         ):
             response = await my_custom_call_to_openai(input)
-            L.set_span_output(response["choices"][0]["message"]["content"])
-            L.set_span_attributes({
+            Laminar.set_span_output(response["choices"][0]["message"]["content"])
+            Laminar.set_span_attributes({
                 Attributes.PROVIDER: 'openai',
                 Attributes.REQUEST_MODEL: input["model"],
                 Attributes.RESPONSE_MODEL: response["model"],

tests/test_google_genai.py

@@ -879,8 +879,10 @@ def test_google_genai_streaming(span_exporter: InMemorySpanExporter):
         ],
     )
     final_response = ""
+    chunk_count = 0
     for chunk in stream:
         final_response += chunk.text or ""
+        chunk_count += 1
 
     spans = span_exporter.get_finished_spans()
     assert len(spans) == 1
@@ -917,6 +919,8 @@ def test_google_genai_streaming(span_exporter: InMemorySpanExporter):
     assert span.attributes["gen_ai.usage.input_tokens"] == 7
     assert span.attributes["gen_ai.usage.output_tokens"] == 166
     assert span.attributes["llm.usage.total_tokens"] == 175  # 173 + 2 (thinking tokens)
+    assert len(span.events) == chunk_count
+    assert all(event.name == "llm.content.completion.chunk" for event in span.events)
 
 
 @pytest.mark.vcr