From f767440f424307021d5c5dc6a06c9757ca81d450 Mon Sep 17 00:00:00 2001 From: WH-2099 Date: Wed, 28 Oct 2026 13:30:27 +0000 Subject: [PATCH 1/2] fix(model-runtime): complete streaming invocation callbacks --- src/graphon/model_runtime/README.md | 4 ++ .../base/large_language_model.py | 18 ++++- tests/model_runtime/test_model_dispatch.py | 69 +++++++++++++++---- 3 files changed, 77 insertions(+), 14 deletions(-) diff --git a/src/graphon/model_runtime/README.md b/src/graphon/model_runtime/README.md index 8a72f974..fc62471e 100644 --- a/src/graphon/model_runtime/README.md +++ b/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). diff --git a/src/graphon/model_runtime/model_providers/base/large_language_model.py b/src/graphon/model_runtime/model_providers/base/large_language_model.py index ac2eb602..109b30dd 100644 --- a/src/graphon/model_runtime/model_providers/base/large_language_model.py +++ b/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(), @@ -441,7 +445,7 @@ def _invoke_result_generator( model: str, result: Generator[LLMResultChunk, None, None], credentials: dict, - prompt_messages: Sequence[PromptMessage], + prompt_messages: list[PromptMessage], model_parameters: dict, tools: list[PromptMessageTool] | None = None, stop: Sequence[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( diff --git a/tests/model_runtime/test_model_dispatch.py b/tests/model_runtime/test_model_dispatch.py index b17db45b..656a2def 100644 --- a/tests/model_runtime/test_model_dispatch.py +++ b/tests/model_runtime/test_model_dispatch.py @@ -245,6 +245,7 @@ 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, @@ -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, @@ -684,7 +688,10 @@ def test_large_language_model_accepts_llm_only_runtime_surface() -> None: ) -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,6 +704,29 @@ 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([ @@ -704,12 +734,14 @@ def test_large_language_model_stream_after_invoke_preserves_chunk_state() -> Non 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, ), ]), ) @@ -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: From bb1ed6bcf269a9f85c0510909ff311a002bb2a45 Mon Sep 17 00:00:00 2001 From: WH-2099 Date: Thu, 29 Oct 2026 06:30:39 +0000 Subject: [PATCH 2/2] fix(model-runtime): accept prompt message sequences --- src/graphon/model_runtime/callbacks/base_callback.py | 4 ++-- .../model_runtime/callbacks/logging_callback.py | 4 ++-- .../model_providers/base/large_language_model.py | 10 +++++----- tests/model_runtime/test_model_dispatch.py | 12 ++++++------ 4 files changed, 15 insertions(+), 15 deletions(-) diff --git a/src/graphon/model_runtime/callbacks/base_callback.py b/src/graphon/model_runtime/callbacks/base_callback.py index 8179b38e..967e5d36 100644 --- a/src/graphon/model_runtime/callbacks/base_callback.py +++ b/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, diff --git a/src/graphon/model_runtime/callbacks/logging_callback.py b/src/graphon/model_runtime/callbacks/logging_callback.py index 28fd0bcb..39d0b282 100644 --- a/src/graphon/model_runtime/callbacks/logging_callback.py +++ b/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, diff --git a/src/graphon/model_runtime/model_providers/base/large_language_model.py b/src/graphon/model_runtime/model_providers/base/large_language_model.py index 109b30dd..835a996a 100644 --- a/src/graphon/model_runtime/model_providers/base/large_language_model.py +++ b/src/graphon/model_runtime/model_providers/base/large_language_model.py @@ -334,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, @@ -445,7 +445,7 @@ def _invoke_result_generator( model: str, result: Generator[LLMResultChunk, None, None], credentials: dict, - prompt_messages: list[PromptMessage], + prompt_messages: Sequence[PromptMessage], model_parameters: dict, tools: list[PromptMessageTool] | None = None, stop: Sequence[str] | None = None, @@ -542,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.""" @@ -600,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, @@ -726,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, diff --git a/tests/model_runtime/test_model_dispatch.py b/tests/model_runtime/test_model_dispatch.py index 656a2def..f188e422 100644 --- a/tests/model_runtime/test_model_dispatch.py +++ b/tests/model_runtime/test_model_dispatch.py @@ -252,7 +252,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, @@ -335,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, @@ -672,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, ), ) @@ -682,7 +682,7 @@ 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 ) @@ -746,7 +746,7 @@ def test_large_language_model_stream_after_invoke_preserves_chunk_state( ]), ) 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( @@ -890,7 +890,7 @@ def test_large_language_model_stream_failures_notify_callbacks( model.invoke( model="fake-chat", credentials={}, - prompt_messages=[UserPromptMessage(content="hello")], + prompt_messages=(UserPromptMessage(content="hello"),), stream=stream, callbacks=[callback], )