fix(ai): keep streamed OpenAI tool calls that have no index

.sampo/changesets/openai-stream-tool-calls-without-index.md

@@ -0,0 +1,5 @@
+---
+pypi/posthog: patch
+---
+
+Streamed OpenAI-compatible chat completions whose tool call deltas carry no `index` are no longer dropped from `$ai_output`. A delta with a new id starts a call, and one without an id continues the latest call.

posthog/ai/openai/_streaming.py

@@ -1,7 +1,7 @@
 """Sync-neutral state accumulation for OpenAI streaming endpoints."""
 
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from ..types import StreamingEventData, TokenUsage
 from ..utils import merge_usage_stats, _responses_stop_reason
@@ -52,7 +52,7 @@ class _ChatCompletionsStreamState:
 
     usage_stats: TokenUsage = field(default_factory=lambda: TokenUsage())
     output: List[Any] = field(default_factory=list)
-    _tool_calls: Dict[int, Dict[str, Any]] = field(default_factory=dict)
+    _tool_calls: Dict[Union[int, str], Dict[str, Any]] = field(default_factory=dict)
     model: Optional[str] = None
     stop_reason: Optional[str] = None
     service_tier: Optional[str] = None

posthog/ai/openai/openai_converter.py

@@ -6,7 +6,7 @@
 Chat Completions API and Responses API formats.
 """
 
-from typing import Any, Dict, List, Optional, cast
+from typing import Any, Dict, List, Optional, Union, cast
 
 from posthog.ai.media import to_plain
 from posthog.ai.types import (
@@ -732,7 +732,7 @@ def extract_openai_tool_calls_from_chunk(chunk: Any) -> Optional[List[Dict[str,
 
 
 def accumulate_openai_tool_calls(
-    accumulated_tool_calls: Dict[int, Dict[str, Any]],
+    accumulated_tool_calls: Dict[Union[int, str], Dict[str, Any]],
     chunk_tool_calls: List[Dict[str, Any]],
 ) -> None:
     """
@@ -741,6 +741,8 @@ def accumulate_openai_tool_calls(
     OpenAI sends tool calls incrementally:
     - First chunk has id, type, function.name and partial function.arguments
     - Subsequent chunks have more function.arguments
+    Deltas without an index are matched by id, start a call when they carry an id
+    or a function name, and otherwise continue the latest call.
 
     Args:
         accumulated_tool_calls: Dictionary mapping index to accumulated tool call data
@@ -749,7 +751,25 @@ def accumulate_openai_tool_calls(
     for tool_call_delta in chunk_tool_calls:
         index = tool_call_delta.get("index")
         if index is None:
-            continue
+            # Some OpenAI-compatible providers omit the index. A delta with a new id
+            # or a function name starts a call, and one with neither continues the
+            # latest call. Those calls get keys of their own, so they can never
+            # collide with an index the provider sends later.
+            call_id = tool_call_delta.get("id")
+            index = next(
+                (
+                    key
+                    for key, tc in accumulated_tool_calls.items()
+                    if call_id and tc["id"] == call_id
+                ),
+                None,
+            )
+            if index is None:
+                has_name = bool((tool_call_delta.get("function") or {}).get("name"))
+                if accumulated_tool_calls and not call_id and not has_name:
+                    index = next(reversed(accumulated_tool_calls))
+                else:
+                    index = f"no-index-{len(accumulated_tool_calls)}"
 
         # Initialize tool call if first time seeing this index
         if index not in accumulated_tool_calls:

posthog/test/ai/openai/test_openai_converter.py

@@ -9,8 +9,12 @@
 except ImportError:
     OPENAI_AVAILABLE = False
 
-from posthog.ai.openai._streaming import _ResponsesStreamState
+from posthog.ai.openai._streaming import (
+    _ChatCompletionsStreamState,
+    _ResponsesStreamState,
+)
 from posthog.ai.openai.openai_converter import (
+    accumulate_openai_tool_calls,
     extract_openai_stop_reason,
     format_openai_input,
 )
@@ -257,3 +261,131 @@ def test_responses_stream_records_every_terminal_stop_reason(
     )
 
     assert state.stop_reason == expected
+
+
+def _tool_call_chunk(*tool_calls):
+    return types.SimpleNamespace(
+        choices=[
+            types.SimpleNamespace(
+                delta=types.SimpleNamespace(content=None, tool_calls=list(tool_calls)),
+                finish_reason=None,
+            )
+        ],
+        usage=None,
+    )
+
+
+def _tool_call_delta(index=None, id=None, name=None, arguments=None):
+    return types.SimpleNamespace(
+        index=index,
+        id=id,
+        type="function" if id else None,
+        function=types.SimpleNamespace(name=name, arguments=arguments),
+    )
+
+
+def test_stream_tool_calls_without_index_are_kept():
+    # Some OpenAI-compatible providers send each call whole, with no index.
+    state = _ChatCompletionsStreamState()
+    state.process_chunk(
+        _tool_call_chunk(_tool_call_delta(id="call_1", name="a", arguments="{}"))
+    )
+    state.process_chunk(
+        _tool_call_chunk(_tool_call_delta(id="call_2", name="b", arguments='{"x": 1}'))
+    )
+
+    assert [(tc["id"], tc["function"]["name"]) for tc in state.tool_calls] == [
+        ("call_1", "a"),
+        ("call_2", "b"),
+    ]
+    assert state.tool_calls[1]["function"]["arguments"] == '{"x": 1}'
+
+
+def test_stream_tool_calls_without_index_join_argument_fragments():
+    tool_calls = {}
+    accumulate_openai_tool_calls(
+        tool_calls, [{"id": "call_1", "function": {"name": "a", "arguments": '{"q":'}}]
+    )
+    accumulate_openai_tool_calls(tool_calls, [{"function": {"arguments": ' "hi"}'}}])
+    accumulate_openai_tool_calls(
+        tool_calls, [{"id": "call_1", "function": {"arguments": ""}}]
+    )
+
+    assert list(tool_calls.values()) == [
+        {
+            "id": "call_1",
+            "type": "function",
+            "function": {"name": "a", "arguments": '{"q": "hi"}'},
+        }
+    ]
+
+
+def test_stream_tool_calls_with_index_are_unchanged():
+    tool_calls = {}
+    accumulate_openai_tool_calls(
+        tool_calls,
+        [
+            {"index": 0, "id": "call_1", "function": {"name": "a", "arguments": "{"}},
+            {"index": 1, "id": "call_2", "function": {"name": "b", "arguments": "{}"}},
+        ],
+    )
+    accumulate_openai_tool_calls(
+        tool_calls, [{"index": 0, "function": {"arguments": "}"}}]
+    )
+
+    assert tool_calls[0]["function"]["arguments"] == "{}"
+    assert tool_calls[1]["function"]["name"] == "b"
+
+
+def test_stream_tool_calls_without_index_or_id_start_a_call_per_name():
+    tool_calls = {}
+    accumulate_openai_tool_calls(
+        tool_calls, [{"function": {"name": "a", "arguments": '{"x": 1}'}}]
+    )
+    accumulate_openai_tool_calls(
+        tool_calls, [{"function": {"name": "b", "arguments": '{"y": 2}'}}]
+    )
+
+    assert [
+        (tc["function"]["name"], tc["function"]["arguments"])
+        for tc in tool_calls.values()
+    ] == [("a", '{"x": 1}'), ("b", '{"y": 2}')]
+
+
+def test_indexless_call_does_not_corrupt_later_indexed_call():
+    # A call without an index must not take a key that a later real index reuses.
+    tool_calls = {}
+    accumulate_openai_tool_calls(
+        tool_calls, [{"id": "call_b", "function": {"name": "b", "arguments": "{}"}}]
+    )
+    accumulate_openai_tool_calls(
+        tool_calls,
+        [{"index": 0, "id": "call_a", "function": {"name": "a", "arguments": "{}"}}],
+    )
+
+    assert [(tc["id"], tc["function"]["arguments"]) for tc in tool_calls.values()] == [
+        ("call_b", "{}"),
+        ("call_a", "{}"),
+    ]
+
+
+def test_indexless_calls_keep_their_order_around_indexed_ones():
+    tool_calls = {}
+    accumulate_openai_tool_calls(
+        tool_calls,
+        [
+            {"index": 0, "id": "call_a", "function": {"name": "a", "arguments": "{"}},
+            {"id": "call_b", "function": {"name": "b", "arguments": "{}"}},
+            {"index": 0, "function": {"arguments": "}"}},
+            {"id": "call_c", "function": {"name": "c", "arguments": "{}"}},
+        ],
+    )
+
+    assert [
+        (tc["id"], tc["function"]["name"], tc["function"]["arguments"])
+        for tc in tool_calls.values()
+    ] == [
+        ("call_a", "a", "{}"),
+        ("call_b", "b", "{}"),
+        ("call_c", "c", "{}"),
+    ]

references/public_api_snapshot.txt

@@ -1300,7 +1300,7 @@ function posthog.ai.media.bytes_to_base64(data: bytes) -> str
 function posthog.ai.media.ensure_serializable(obj: Any) -> Any
 function posthog.ai.media.normalize_part_keys(d: dict) -> dict
 function posthog.ai.media.to_plain(obj: Any) -> Any
-function posthog.ai.openai.openai_converter.accumulate_openai_tool_calls(accumulated_tool_calls: Dict[int, Dict[str, Any]], chunk_tool_calls: List[Dict[str, Any]]) -> None
+function posthog.ai.openai.openai_converter.accumulate_openai_tool_calls(accumulated_tool_calls: Dict[Union[int, str], Dict[str, Any]], chunk_tool_calls: List[Dict[str, Any]]) -> None
 function posthog.ai.openai.openai_converter.extract_openai_content_from_chunk(chunk: Any, provider_type: str = 'chat') -> Optional[Any]
 function posthog.ai.openai.openai_converter.extract_openai_stop_reason(response: Any) -> Optional[str]
 function posthog.ai.openai.openai_converter.extract_openai_tool_calls_from_chunk(chunk: Any) -> Optional[List[Dict[str, Any]]]