fix genai streaming parts by merging texts

pyproject.toml

@@ -6,7 +6,7 @@
 
 [project]
 name = "lmnr"
-version = "0.7.16"
+version = "0.7.17"
 description = "Python SDK for Laminar"
 authors = [
   { name = "lmnr.ai", email = "[email protected]" }
@@ -39,6 +39,7 @@ dependencies = [
   "httpx (>=0.24.0)",
   "orjson (>=3.0.0)",
   "packaging (>=22.0)",
+  "opentelemetry-instrumentation-threading>=0.57b0",
 ]
 # poetry would auto-generate these based on requires-python, but other build
 # systems don't, so we need to specify them manually.
@@ -145,4 +146,4 @@ build-backend = "uv_build"
 members = ["examples/fastapi-app"]
 
 [tool.ruff]
-target-version = "py310"
+target-version = "py310"

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

@@ -21,6 +21,7 @@
 from .utils import (
     dont_throw,
     get_content,
+    merge_text_parts,
     process_content_union,
     process_stream_chunk,
     role_from_content_union,
@@ -403,7 +404,7 @@ def _build_from_streaming_response(
             candidates=[
                 {
                     "content": {
-                        "parts": final_parts,
+                        "parts": merge_text_parts(final_parts),
                         "role": role,
                     },
                 }
@@ -453,7 +454,7 @@ async def _abuild_from_streaming_response(
             candidates=[
                 {
                     "content": {
-                        "parts": final_parts,
+                        "parts": merge_text_parts(final_parts),
                         "role": role,
                     },
                 }

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

@@ -40,6 +40,56 @@ class ProcessChunkResult(TypedDict):
     model_version: str | None
 
 
+def merge_text_parts(
+    parts: list[types.PartDict | types.File | types.Part | str],
+) -> list[types.Part]:
+    if not parts:
+        return []
+
+    merged_parts: list[types.Part] = []
+    accumulated_text = ""
+
+    for part in parts:
+        # Handle string input - treat as text
+        if isinstance(part, str):
+            accumulated_text += part
+        # Handle File objects - they are not text, so don't merge
+        elif isinstance(part, types.File):
+            # Flush any accumulated text first
+            if accumulated_text:
+                merged_parts.append(types.Part(text=accumulated_text))
+                accumulated_text = ""
+            # Add the File as-is (wrapped in a Part if needed)
+            # Note: File objects should be passed through as-is in the original part
+            merged_parts.append(part)
+        # Handle Part and PartDict (dicts)
+        else:
+            part_dict = to_dict(part)
+
+            # Check if this is a text part
+            if part_dict.get("text") is not None:
+                accumulated_text += part_dict.get("text")
+            else:
+                # Non-text part (inline_data, function_call, etc.)
+                # Flush any accumulated text first
+                if accumulated_text:
+                    merged_parts.append(types.Part(text=accumulated_text))
+                    accumulated_text = ""
+
+                # Add the non-text part as-is
+                if isinstance(part, types.Part):
+                    merged_parts.append(part)
+                elif isinstance(part, dict):
+                    # Convert dict to Part object
+                    merged_parts.append(types.Part(**part_dict))
+
+    # Don't forget to add any remaining accumulated text
+    if accumulated_text:
+        merged_parts.append(types.Part(text=accumulated_text))
+
+    return merged_parts
+
+
 def set_span_attribute(span: Span, name: str, value: Any):
     if value is not None and value != "":
         span.set_attribute(name, value)

src/lmnr/version.py

@@ -3,7 +3,7 @@
 from packaging import version
 
 
-__version__ = "0.7.16"
+__version__ = "0.7.17"
 PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}"
 
 

tests/cassettes/test_google_genai/test_google_genai_streaming.yaml

@@ -0,0 +1,74 @@
+interactions:
+- request:
+    body: '{"contents": [{"parts": [{"text": "Write a short poem about cats"}], "role":
+      "user"}]}'
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate, zstd
+      connection:
+      - keep-alive
+      content-length:
+      - '86'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+      user-agent:
+      - google-genai-sdk/1.34.0 gl-python/3.13.5
+      x-goog-api-client:
+      - google-genai-sdk/1.34.0 gl-python/3.13.5
+    method: POST
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:streamGenerateContent?alt=sse
+  response:
+    body:
+      string: "data: {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \"A silent\"}],\"role\":
+        \"model\"},\"index\": 0}],\"usageMetadata\": {\"promptTokenCount\": 7,\"candidatesTokenCount\":
+        2,\"totalTokenCount\": 9,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\":
+        7}]},\"modelVersion\": \"gemini-2.5-flash-lite\",\"responseId\": \"QoveaJGbIeunkdUP8abQ8Aw\"}\r\n\r\ndata:
+        {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" stalk, a velvet
+        paw,\\nA hunter's grace, without a flaw.\\nEmerald\"}],\"role\": \"model\"},\"index\":
+        0}],\"usageMetadata\": {\"promptTokenCount\": 7,\"candidatesTokenCount\":
+        21,\"totalTokenCount\": 28,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\":
+        7}]},\"modelVersion\": \"gemini-2.5-flash-lite\",\"responseId\": \"QoveaJGbIeunkdUP8abQ8Aw\"}\r\n\r\ndata:
+        {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \" eyes, that softly
+        gleam,\\nLost in a dream, a furry dream.\\n\\nA gentle purr, a rumbling sound,\\nWhen
+        happy hearts are to be found.\\nThey stretch and yawn, a lazy art,\\n\"}],\"role\":
+        \"model\"},\"index\": 0}],\"usageMetadata\": {\"promptTokenCount\": 7,\"candidatesTokenCount\":
+        68,\"totalTokenCount\": 75,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\":
+        7}]},\"modelVersion\": \"gemini-2.5-flash-lite\",\"responseId\": \"QoveaJGbIeunkdUP8abQ8Aw\"}\r\n\r\ndata:
+        {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \"And steal away a
+        human heart.\"}],\"role\": \"model\"},\"finishReason\": \"STOP\",\"index\":
+        0}],\"usageMetadata\": {\"promptTokenCount\": 7,\"candidatesTokenCount\":
+        75,\"totalTokenCount\": 82,\"promptTokensDetails\": [{\"modality\": \"TEXT\",\"tokenCount\":
+        7}]},\"modelVersion\": \"gemini-2.5-flash-lite\",\"responseId\": \"QoveaJGbIeunkdUP8abQ8Aw\"}\r\n\r\n"
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      Content-Disposition:
+      - attachment
+      Content-Type:
+      - text/event-stream
+      Date:
+      - Thu, 02 Oct 2025 14:25:06 GMT
+      Server:
+      - scaffolding on HTTPServer2
+      Server-Timing:
+      - gfet4t7; dur=311
+      Transfer-Encoding:
+      - chunked
+      Vary:
+      - Origin
+      - X-Origin
+      - Referer
+      X-Content-Type-Options:
+      - nosniff
+      X-Frame-Options:
+      - SAMEORIGIN
+      X-XSS-Protection:
+      - '0'
+    status:
+      code: 200
+      message: OK
+version: 1

tests/test_google_genai.py

@@ -831,6 +831,62 @@ def test_google_genai_error(span_exporter: InMemorySpanExporter):
     assert "google.genai.errors.ClientError" in event.attributes["exception.stacktrace"]
 
 
+@pytest.mark.vcr
+def test_google_genai_streaming(span_exporter: InMemorySpanExporter):
+    client = Client(api_key="123")
+
+    stream = client.models.generate_content_stream(
+        model="gemini-2.5-flash-lite",
+        contents=[
+            {
+                "role": "user",
+                "parts": [
+                    {"text": "Write a short poem about cats"},
+                ],
+            }
+        ],
+    )
+    final_response = ""
+    for chunk in stream:
+        final_response += chunk.text or ""
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span = spans[0]
+    assert span.name == "gemini.generate_content_stream"
+    assert (
+        final_response
+        == """A silent stalk, a velvet paw,
+A hunter's grace, without a flaw.
+Emerald eyes, that softly gleam,
+Lost in a dream, a furry dream.
+
+A gentle purr, a rumbling sound,
+When happy hearts are to be found.
+They stretch and yawn, a lazy art,
+And steal away a human heart."""
+    )
+
+    assert json.loads(span.attributes["gen_ai.prompt.0.content"]) == [
+        {
+            "text": "Write a short poem about cats",
+            "type": "text",
+        }
+    ]
+    assert span.attributes["gen_ai.prompt.0.role"] == "user"
+    assert json.loads(span.attributes["gen_ai.completion.0.content"]) == [
+        {
+            "text": final_response,
+            "type": "text",
+        }
+    ]
+
+    assert span.attributes["gen_ai.completion.0.role"] == "model"
+    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)
+
+
 @pytest.mark.vcr
 def test_google_genai_no_tokens(span_exporter: InMemorySpanExporter):
     client = Client(api_key="123")