fix(model-runtime): complete streaming invocation callbacks

src/graphon/model_runtime/README.md

@@ -39,6 +39,10 @@ blocks retain their order and block state; see the
 [SlimLLM](../dsl/slim/llm.py) shares this aggregation for blocking ordinary and
 structured output calls, covered by the [Slim tests](../../../tests/dsl/test_slim_llm.py).
 
+After a `LargeLanguageModel.invoke` stream is fully consumed without errors, `on_after_invoke` receives a result containing all tool calls, assembled with the same merging rules as blocking results.
+Iteration failures, including before the first chunk, notify `on_invoke_error` once with the original exception before normal exception conversion and propagation, and skip `on_after_invoke`.
+See the [callback regression tests](../../../tests/model_runtime/test_model_dispatch.py).
+
 Hosts can display chunks while calling `consume(chunk)` on one [LLMResultAccumulator](model_providers/base/large_language_model.py) per invocation.
 After consuming the stream, call its `to_result(model=..., prompt_messages=...)` and use the returned `message` for replay or [JSON persistence](../../../tests/model_runtime/test_llm_result_accumulator.py).
 

src/graphon/model_runtime/callbacks/base_callback.py

@@ -31,7 +31,7 @@ def on_before_invoke(
         llm_instance: AIModel,
         model: str,
         credentials: dict,
-        prompt_messages: list[PromptMessage],
+        prompt_messages: Sequence[PromptMessage],
         model_parameters: dict,
         tools: list[PromptMessageTool] | None = None,
         stop: Sequence[str] | None = None,
@@ -123,7 +123,7 @@ def on_invoke_error(
         ex: Exception,
         model: str,
         credentials: dict,
-        prompt_messages: list[PromptMessage],
+        prompt_messages: Sequence[PromptMessage],
         model_parameters: dict,
         tools: list[PromptMessageTool] | None = None,
         stop: Sequence[str] | None = None,

src/graphon/model_runtime/callbacks/logging_callback.py

@@ -22,7 +22,7 @@ def on_before_invoke(
         llm_instance: AIModel,
         model: str,
         credentials: dict,
-        prompt_messages: list[PromptMessage],
+        prompt_messages: Sequence[PromptMessage],
         model_parameters: dict,
         tools: list[PromptMessageTool] | None = None,
         stop: Sequence[str] | None = None,
@@ -167,7 +167,7 @@ def on_invoke_error(
         ex: Exception,
         model: str,
         credentials: dict,
-        prompt_messages: list[PromptMessage],
+        prompt_messages: Sequence[PromptMessage],
         model_parameters: dict,
         tools: list[PromptMessageTool] | None = None,
         stop: Sequence[str] | None = None,

src/graphon/model_runtime/model_providers/base/large_language_model.py

@@ -181,12 +181,15 @@ def to_result(
 class _StreamingInvokeAccumulator:
     real_model: str
     message_content: list[PromptMessageContentUnionTypes] = field(default_factory=list)
+    tool_calls: list[AssistantPromptMessage.ToolCall] = field(default_factory=list)
     usage: LLMUsage | None = None
     system_fingerprint: str | None = None
     opaque_body: JsonValue = None
 
     def consume(self, chunk: LLMResultChunk) -> None:
         self._consume_content(chunk.delta.message.content)
+        if chunk.delta.message.tool_calls:
+            merge_tool_call_deltas(chunk.delta.message.tool_calls, self.tool_calls)
         self.real_model = chunk.model
         if chunk.delta.usage:
             self.usage = chunk.delta.usage
@@ -217,6 +220,7 @@ def to_result(
             prompt_messages=prompt_messages,
             message=AssistantPromptMessage(
                 content=self.message_content,
+                tool_calls=self.tool_calls,
                 opaque_body=self.opaque_body,
             ),
             usage=self.usage or LLMUsage.empty_usage(),
@@ -330,7 +334,7 @@ def invoke(
         self,
         model: str,
         credentials: dict,
-        prompt_messages: list[PromptMessage],
+        prompt_messages: Sequence[PromptMessage],
         model_parameters: dict | None = None,
         tools: list[PromptMessageTool] | None = None,
         stop: list[str] | None = None,
@@ -468,6 +472,18 @@ def _invoke_result_generator(
                 callbacks=callbacks,
             )
         except Exception as e:
+            self._trigger_invoke_error_callbacks(
+                model=model,
+                ex=e,
+                credentials=credentials,
+                prompt_messages=prompt_messages,
+                model_parameters=model_parameters,
+                tools=tools,
+                stop=stop,
+                stream=stream,
+                request_metadata=request_metadata,
+                callbacks=callbacks,
+            )
             raise self._transform_invoke_error(e) from e
 
         self._trigger_after_invoke_callbacks(
@@ -526,7 +542,7 @@ def get_num_tokens(
         self,
         model: str,
         credentials: dict,
-        prompt_messages: list[PromptMessage],
+        prompt_messages: Sequence[PromptMessage],
         tools: list[PromptMessageTool] | None = None,
     ) -> int:
         """Count prompt tokens for the given messages and optional tools."""
@@ -584,7 +600,7 @@ def _trigger_before_invoke_callbacks(
         self,
         model: str,
         credentials: dict,
-        prompt_messages: list[PromptMessage],
+        prompt_messages: Sequence[PromptMessage],
         model_parameters: dict,
         tools: list[PromptMessageTool] | None = None,
         stop: Sequence[str] | None = None,
@@ -710,7 +726,7 @@ def _trigger_invoke_error_callbacks(
         model: str,
         ex: Exception,
         credentials: dict,
-        prompt_messages: list[PromptMessage],
+        prompt_messages: Sequence[PromptMessage],
         model_parameters: dict,
         tools: list[PromptMessageTool] | None = None,
         stop: Sequence[str] | None = None,

tests/model_runtime/test_model_dispatch.py

@@ -245,13 +245,14 @@ def invoke_llm(
 class _RecordingCallback(Callback):
     def __init__(self) -> None:
         self.after_results: list[LLMResult] = []
+        self.invoke_errors: list[Exception] = []
 
     def on_before_invoke(
         self,
         llm_instance: AIModel,
         model: str,
         credentials: dict,
-        prompt_messages: list[PromptMessage],
+        prompt_messages: Sequence[PromptMessage],
         model_parameters: dict,
         tools: list[PromptMessageTool] | None = None,
         stop: Sequence[str] | None = None,
@@ -334,7 +335,7 @@ def on_invoke_error(
         ex: Exception,
         model: str,
         credentials: dict,
-        prompt_messages: list[PromptMessage],
+        prompt_messages: Sequence[PromptMessage],
         model_parameters: dict,
         tools: list[PromptMessageTool] | None = None,
         stop: Sequence[str] | None = None,
@@ -344,7 +345,6 @@ def on_invoke_error(
     ) -> None:
         _ = (
             llm_instance,
-            ex,
             model,
             credentials,
             prompt_messages,
@@ -355,6 +355,7 @@ def on_invoke_error(
             user,
             request_metadata,
         )
+        self.invoke_errors.append(ex)
 
 
 def _make_chunk(
@@ -363,12 +364,15 @@ def _make_chunk(
     content: str | list[PromptMessageContentUnionTypes] | None,
     usage: LLMUsage | None = None,
     system_fingerprint: str | None = None,
+    tool_calls: list[AssistantPromptMessage.ToolCall] | None = None,
 ) -> LLMResultChunk:
     return LLMResultChunk(
         model=model,
         delta=LLMResultChunkDelta(
             index=0,
-            message=AssistantPromptMessage(content=content),
+            message=AssistantPromptMessage(
+                content=content, tool_calls=tool_calls or []
+            ),
             usage=usage,
         ),
         system_fingerprint=system_fingerprint,
@@ -668,7 +672,7 @@ def test_large_language_model_accepts_llm_only_runtime_surface() -> None:
         model.invoke(
             model="fake-chat",
             credentials={},
-            prompt_messages=[UserPromptMessage(content="hello")],
+            prompt_messages=(UserPromptMessage(content="hello"),),
             stream=False,
         ),
     )
@@ -678,13 +682,16 @@ def test_large_language_model_accepts_llm_only_runtime_surface() -> None:
         model.get_num_tokens(
             model="fake-chat",
             credentials={},
-            prompt_messages=[UserPromptMessage(content="hello")],
+            prompt_messages=(UserPromptMessage(content="hello"),),
         )
         == 7
     )
 
 
-def test_large_language_model_stream_after_invoke_preserves_chunk_state() -> None:
+@pytest.mark.parametrize("split_tool_call", [False, True])
+def test_large_language_model_stream_after_invoke_preserves_chunk_state(
+    split_tool_call: bool,
+) -> None:
     provider = ProviderEntity(
         provider="test-provider",
         label=I18nObject(en_US="Test Provider"),
@@ -697,24 +704,49 @@ def test_large_language_model_stream_after_invoke_preserves_chunk_state() -> Non
     list_content: list[PromptMessageContentUnionTypes] = [
         TextPromptMessageContent(data=" there"),
     ]
+    tool_calls = [
+        AssistantPromptMessage.ToolCall(
+            id=f"call-{index}",
+            type="function",
+            function=AssistantPromptMessage.ToolCall.ToolCallFunction(
+                name="lookup", arguments=arguments
+            ),
+        )
+        for index, arguments in enumerate(['{ "city": "Paris" }', '{"city":"Rome"}'])
+    ]
+    first_chunk_tool_calls = [tool_calls[0].model_copy(deep=True)]
+    second_chunk_tool_calls = [tool_calls[1].model_copy(deep=True)]
+    if split_tool_call:
+        first_chunk_tool_calls[0].function.arguments = '{ "city": '
+        second_chunk_tool_calls.append(
+            AssistantPromptMessage.ToolCall(
+                id=tool_calls[0].id,
+                type="function",
+                function=AssistantPromptMessage.ToolCall.ToolCallFunction(
+                    name="", arguments='"Paris" }'
+                ),
+            )
+        )
     runtime = cast(
         "LLMModelRuntime",
         _StreamingLLMRuntimeStub([
             _make_chunk(
                 model="first-model",
                 content="hi",
                 system_fingerprint="fp-1",
+                tool_calls=first_chunk_tool_calls,
             ),
             _make_chunk(
                 model="second-model",
                 content=list_content,
                 usage=usage,
                 system_fingerprint="fp-2",
+                tool_calls=second_chunk_tool_calls,
             ),
         ]),
     )
     model = LargeLanguageModel(provider_schema=provider, model_runtime=runtime)
-    prompt_messages: list[PromptMessage] = [UserPromptMessage(content="hello")]
+    prompt_messages: Sequence[PromptMessage] = (UserPromptMessage(content="hello"),)
     callback = _RecordingCallback()
 
     result = model.invoke(
@@ -740,6 +772,7 @@ def test_large_language_model_stream_after_invoke_preserves_chunk_state() -> Non
     ]
     assert after_result.usage == usage
     assert after_result.system_fingerprint == "fp-2"
+    assert after_result.message.tool_calls == tool_calls
 
 
 @pytest.mark.parametrize("stream", [True, False])
@@ -831,7 +864,11 @@ def test_large_language_model_uses_last_opaque_body_per_invocation(
             assert type(blocks[1].opaque_body) is type(snapshot)
 
 
-def test_large_language_model_stream_errors_are_transformed() -> None:
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("emit_chunk", [False, True])
+def test_large_language_model_stream_failures_notify_callbacks(
+    stream: bool, emit_chunk: bool
+) -> None:
     provider = ProviderEntity(
         provider="test-provider",
         label=I18nObject(en_US="Test Provider"),
@@ -841,20 +878,26 @@ def test_large_language_model_stream_errors_are_transformed() -> None:
     runtime = cast(
         "LLMModelRuntime",
         _StreamingLLMRuntimeStub(
-            [_make_chunk(model="first-model", content="hi")],
+            [_make_chunk(model="first-model", content="hi")] if emit_chunk else [],
             fail_after_chunks=True,
         ),
     )
     model = LargeLanguageModel(provider_schema=provider, model_runtime=runtime)
-    result = model.invoke(
-        model="fake-chat",
-        credentials={},
-        prompt_messages=[UserPromptMessage(content="hello")],
-        stream=True,
-    )
+    callback = _RecordingCallback()
+
+    with pytest.raises(InvokeError, match="stream failed") as raised:
+        list(
+            model.invoke(
+                model="fake-chat",
+                credentials={},
+                prompt_messages=(UserPromptMessage(content="hello"),),
+                stream=stream,
+                callbacks=[callback],
+            )
+        )
 
-    with pytest.raises(InvokeError, match="stream failed"):
-        list(cast("Generator[LLMResultChunk, None, None]", result))
+    assert callback.invoke_errors == [raised.value.__cause__]
+    assert callback.after_results == []
 
 
 def test_text_embedding_model_accepts_embedding_only_runtime_surface() -> None: