diff --git a/.sampo/changesets/openai-stream-tool-calls-without-index.md b/.sampo/changesets/openai-stream-tool-calls-without-index.md new file mode 100644 index 000000000..1bc8537d9 --- /dev/null +++ b/.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. diff --git a/posthog/ai/openai/_streaming.py b/posthog/ai/openai/_streaming.py index 2ee95ec91..ea1035ea3 100644 --- a/posthog/ai/openai/_streaming.py +++ b/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 diff --git a/posthog/ai/openai/openai_converter.py b/posthog/ai/openai/openai_converter.py index 0f391fb9f..70e9bd38c 100644 --- a/posthog/ai/openai/openai_converter.py +++ b/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: diff --git a/posthog/test/ai/openai/test_openai_converter.py b/posthog/test/ai/openai/test_openai_converter.py index e23ecac87..ba5079f2c 100644 --- a/posthog/test/ai/openai/test_openai_converter.py +++ b/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", "{}"), + ] diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index 5a093b650..e3d3263d7 100644 --- a/references/public_api_snapshot.txt +++ b/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]]]