From 29f7470f63bf72ae9dcee58c4066d4b9627d9b62 Mon Sep 17 00:00:00 2001 From: hpiclaranet Date: Sat, 22 Aug 2026 10:51:57 +0100 Subject: [PATCH 1/3] feat: preserve opaque_body provider state on prompt messages and LLM results Plugin LLM results carry provider round-trip state in opaque_body (e.g. signed/encrypted thinking blocks, Gemini thought_signature), matching the dify_plugin SDK entities. Graphon dropped that state in two places: - AssistantPromptMessage / PromptMessageContent had no opaque_body field, so daemon -> core Pydantic validation silently discarded it. - _LLMChunkAccumulator and _StreamingInvokeAccumulator rebuilt the assistant message from content and tool calls only. Add the field (same name, type and default as the plugin SDK) and carry the last non-None snapshot through both accumulators. Refs: #260 --- .../entities/message_entities.py | 4 +- .../base/large_language_model.py | 14 +++- ..._large_language_model_non_stream_result.py | 65 ++++++++++++++- tests/model_runtime/test_message_entities.py | 79 ++++++++++++++++++- .../test_streaming_invoke_accumulator.py | 67 ++++++++++++++++ 5 files changed, 225 insertions(+), 4 deletions(-) create mode 100644 tests/model_runtime/test_streaming_invoke_accumulator.py diff --git a/src/graphon/model_runtime/entities/message_entities.py b/src/graphon/model_runtime/entities/message_entities.py index 15e31c50..a022a931 100644 --- a/src/graphon/model_runtime/entities/message_entities.py +++ b/src/graphon/model_runtime/entities/message_entities.py @@ -5,7 +5,7 @@ from enum import StrEnum, auto from typing import Annotated, Any, Literal, Self -from pydantic import BaseModel, Field, field_serializer, field_validator +from pydantic import BaseModel, Field, JsonValue, field_serializer, field_validator class PromptMessageRole(StrEnum): @@ -59,6 +59,7 @@ class PromptMessageContent(ABC, BaseModel): """Model class for prompt message content.""" type: PromptMessageContentType + opaque_body: JsonValue | None = None class TextPromptMessageContent(PromptMessageContent): @@ -258,6 +259,7 @@ def transform_id_to_str(cls, value: Any) -> str: role: PromptMessageRole = PromptMessageRole.ASSISTANT tool_calls: list[ToolCall] = [] + opaque_body: JsonValue | None = None def is_empty(self) -> bool: """Check whether the assistant message has no content or tool calls.""" 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 deae9fba..76036248 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 @@ -4,6 +4,8 @@ from collections.abc import Callable, Generator, Iterator, Mapping, Sequence from dataclasses import dataclass, field +from pydantic import JsonValue + from graphon.model_runtime.callbacks.base_callback import Callback from graphon.model_runtime.callbacks.logging_callback import LoggingCallback from graphon.model_runtime.entities.llm_entities import ( @@ -122,6 +124,7 @@ class _LLMChunkAccumulator: usage: LLMUsage = field(default_factory=LLMUsage.empty_usage) system_fingerprint: str | None = None tool_calls: list[AssistantPromptMessage.ToolCall] = field(default_factory=list) + opaque_body: JsonValue | None = None def consume_all(self, chunks: Iterator[LLMResultChunk]) -> None: for chunk in chunks: @@ -135,6 +138,8 @@ def consume(self, chunk: LLMResultChunk) -> None: self.usage = chunk.delta.usage if chunk.system_fingerprint: self.system_fingerprint = chunk.system_fingerprint + if chunk.delta.message.opaque_body is not None: + self.opaque_body = chunk.delta.message.opaque_body def _consume_content(self, chunk: LLMResultChunk) -> None: content = chunk.delta.message.content @@ -155,6 +160,7 @@ def to_result( message=AssistantPromptMessage( content=self.content or self.content_list, tool_calls=self.tool_calls, + opaque_body=self.opaque_body, ), usage=self.usage, system_fingerprint=self.system_fingerprint, @@ -167,6 +173,7 @@ class _StreamingInvokeAccumulator: message_content: list[PromptMessageContentUnionTypes] = field(default_factory=list) usage: LLMUsage | None = None system_fingerprint: str | None = None + opaque_body: JsonValue | None = None def consume(self, chunk: LLMResultChunk) -> None: self._consume_content(chunk.delta.message.content) @@ -175,6 +182,8 @@ def consume(self, chunk: LLMResultChunk) -> None: self.usage = chunk.delta.usage if chunk.system_fingerprint: self.system_fingerprint = chunk.system_fingerprint + if chunk.delta.message.opaque_body is not None: + self.opaque_body = chunk.delta.message.opaque_body def _consume_content( self, @@ -196,7 +205,10 @@ def to_result( return LLMResult( model=self.real_model, prompt_messages=prompt_messages, - message=AssistantPromptMessage(content=self.message_content), + message=AssistantPromptMessage( + content=self.message_content, + opaque_body=self.opaque_body, + ), usage=self.usage or LLMUsage.empty_usage(), system_fingerprint=self.system_fingerprint, ) diff --git a/tests/model_runtime/test_large_language_model_non_stream_result.py b/tests/model_runtime/test_large_language_model_non_stream_result.py index ddd3664a..341a21a6 100644 --- a/tests/model_runtime/test_large_language_model_non_stream_result.py +++ b/tests/model_runtime/test_large_language_model_non_stream_result.py @@ -1,4 +1,5 @@ from collections.abc import Iterator +from typing import Any from graphon.model_runtime.entities.llm_entities import ( LLMResult, @@ -24,8 +25,13 @@ def _make_chunk( tool_calls: list[AssistantPromptMessage.ToolCall] | None = None, usage: LLMUsage | None = None, system_fingerprint: str | None = None, + opaque_body: Any | None = None, ) -> LLMResultChunk: - message = AssistantPromptMessage(content=content, tool_calls=tool_calls or []) + message = AssistantPromptMessage( + content=content, + tool_calls=tool_calls or [], + opaque_body=opaque_body, + ) delta = LLMResultChunkDelta(index=0, message=message, usage=usage) return LLMResultChunk( model=model, @@ -151,6 +157,63 @@ def test__normalize_non_stream_runtime_result__empty_iterator_defaults() -> None assert result.system_fingerprint is None +def test_non_stream_result_preserves_opaque_body() -> None: + prompt_messages = [UserPromptMessage(content="hi")] + opaque_body = { + "assistant_blocks": [{"type": "thinking", "signature": "sig-1"}], + } + chunk = _make_chunk( + content="hello", + usage=LLMUsage.empty_usage(), + opaque_body=opaque_body, + ) + + result = normalize_non_stream_runtime_result( + model="test-model", + prompt_messages=prompt_messages, + result=iter([chunk]), + ) + + assert result.message.opaque_body == opaque_body + + +def test_non_stream_result_opaque_body_last_non_none_wins() -> None: + prompt_messages = [UserPromptMessage(content="hi")] + chunks = iter([ + _make_chunk( + content="a", + usage=LLMUsage.empty_usage(), + opaque_body={"snapshot": 1}, + ), + _make_chunk(content="b", usage=LLMUsage.empty_usage()), + _make_chunk( + content="c", + usage=LLMUsage.empty_usage(), + opaque_body={"snapshot": 2}, + ), + ]) + + result = normalize_non_stream_runtime_result( + model="test-model", + prompt_messages=prompt_messages, + result=chunks, + ) + + assert result.message.opaque_body == {"snapshot": 2} + + +def test_non_stream_result_opaque_body_defaults_to_none() -> None: + prompt_messages = [UserPromptMessage(content="hi")] + + result = normalize_non_stream_runtime_result( + model="test-model", + prompt_messages=prompt_messages, + result=iter([_make_chunk(content="hello", usage=LLMUsage.empty_usage())]), + ) + + assert result.message.opaque_body is None + + def test__normalize_non_stream_runtime_result__accumulates_all_chunks() -> None: prompt_messages = [UserPromptMessage(content="hi")] diff --git a/tests/model_runtime/test_message_entities.py b/tests/model_runtime/test_message_entities.py index 293477ed..de49299c 100644 --- a/tests/model_runtime/test_message_entities.py +++ b/tests/model_runtime/test_message_entities.py @@ -1,4 +1,10 @@ +from graphon.model_runtime.entities.llm_entities import ( + LLMResultChunk, + LLMResultChunkDelta, + LLMUsage, +) from graphon.model_runtime.entities.message_entities import ( + AssistantPromptMessage, ImagePromptMessageContent, TextPromptMessageContent, UserPromptMessage, @@ -29,5 +35,76 @@ def test_prompt_message_normalizes_dict_content_items_for_serialization() -> Non assert isinstance(message.content, list) assert isinstance(message.content[0], TextPromptMessageContent) assert message.model_dump(mode="json")["content"] == [ - {"type": "text", "data": "hello"}, + {"type": "text", "data": "hello", "opaque_body": None}, ] + + +def test_assistant_prompt_message_opaque_body_defaults_to_none() -> None: + message = AssistantPromptMessage(content="ok") + + assert message.opaque_body is None + + +def test_assistant_prompt_message_opaque_body_survives_json_round_trip() -> None: + opaque_body = { + "assistant_blocks": [{"type": "thinking", "signature": "sig-1"}], + } + message = AssistantPromptMessage(content="ok", opaque_body=opaque_body) + + restored = AssistantPromptMessage.model_validate_json(message.model_dump_json()) + + assert restored.opaque_body == opaque_body + + +def test_assistant_prompt_message_accepts_payload_without_opaque_body() -> None: + message = AssistantPromptMessage.model_validate( + {"role": "assistant", "content": "ok", "tool_calls": []}, + ) + + assert message.opaque_body is None + + +def test_prompt_message_content_opaque_body_survives_json_round_trip() -> None: + content = TextPromptMessageContent( + data="hello", + opaque_body={"thought_signature": "sig-1"}, + ) + message = AssistantPromptMessage(content=[content]) + + restored = AssistantPromptMessage.model_validate_json(message.model_dump_json()) + + assert isinstance(restored.content, list) + assert isinstance(restored.content[0], TextPromptMessageContent) + assert restored.content[0].opaque_body == {"thought_signature": "sig-1"} + + +def test_prompt_message_content_accepts_dict_with_opaque_body() -> None: + message = UserPromptMessage.model_validate( + {"content": [{"type": "text", "data": "hello", "opaque_body": {"key": "v"}}]}, + ) + + assert isinstance(message.content, list) + assert isinstance(message.content[0], TextPromptMessageContent) + assert message.content[0].opaque_body == {"key": "v"} + + +def test_llm_result_chunk_json_round_trip_preserves_opaque_body() -> None: + """Simulate the plugin daemon -> core deserialization boundary.""" + opaque_body = { + "assistant_blocks": [ + {"type": "thinking", "thinking": "...", "signature": "sig-1"}, + {"type": "redacted_thinking", "data": "enc-1"}, + ], + } + chunk = LLMResultChunk( + model="test-model", + delta=LLMResultChunkDelta( + index=0, + message=AssistantPromptMessage(content="ok", opaque_body=opaque_body), + usage=LLMUsage.empty_usage(), + ), + ) + + restored = LLMResultChunk.model_validate_json(chunk.model_dump_json()) + + assert restored.delta.message.opaque_body == opaque_body diff --git a/tests/model_runtime/test_streaming_invoke_accumulator.py b/tests/model_runtime/test_streaming_invoke_accumulator.py new file mode 100644 index 00000000..4366ae0a --- /dev/null +++ b/tests/model_runtime/test_streaming_invoke_accumulator.py @@ -0,0 +1,67 @@ +from typing import Any + +from graphon.model_runtime.entities.llm_entities import ( + LLMResultChunk, + LLMResultChunkDelta, + LLMUsage, +) +from graphon.model_runtime.entities.message_entities import ( + AssistantPromptMessage, + UserPromptMessage, +) +from graphon.model_runtime.model_providers.base.large_language_model import ( + _StreamingInvokeAccumulator, +) + + +def _make_chunk( + *, + content: str | None, + opaque_body: Any | None = None, + usage: LLMUsage | None = None, +) -> LLMResultChunk: + message = AssistantPromptMessage( + content=content, + tool_calls=[], + opaque_body=opaque_body, + ) + delta = LLMResultChunkDelta(index=0, message=message, usage=usage) + return LLMResultChunk(model="test-model", delta=delta) + + +def test_streaming_invoke_accumulator_preserves_opaque_body() -> None: + accumulator = _StreamingInvokeAccumulator(real_model="test-model") + opaque_body = { + "assistant_blocks": [{"type": "thinking", "signature": "sig-1"}], + } + + accumulator.consume(_make_chunk(content="hello")) + accumulator.consume(_make_chunk(content=" world", opaque_body=opaque_body)) + + result = accumulator.to_result(prompt_messages=[UserPromptMessage(content="hi")]) + + assert result.message.opaque_body == opaque_body + + +def test_streaming_invoke_accumulator_opaque_body_defaults_to_none() -> None: + accumulator = _StreamingInvokeAccumulator(real_model="test-model") + + accumulator.consume(_make_chunk(content="hello")) + + result = accumulator.to_result(prompt_messages=[UserPromptMessage(content="hi")]) + + assert result.message.opaque_body is None + + +def test_streaming_invoke_accumulator_none_chunk_does_not_clobber_snapshot() -> None: + accumulator = _StreamingInvokeAccumulator(real_model="test-model") + opaque_body = { + "assistant_blocks": [{"type": "redacted_thinking", "data": "enc-1"}], + } + + accumulator.consume(_make_chunk(content="a", opaque_body=opaque_body)) + accumulator.consume(_make_chunk(content="b")) + + result = accumulator.to_result(prompt_messages=[UserPromptMessage(content="hi")]) + + assert result.message.opaque_body == opaque_body From a6358c1ba18727a3dd88ee6a051b26f45b1b8c16 Mon Sep 17 00:00:00 2001 From: WH-2099 Date: Wed, 23 Sep 2026 09:58:06 +0000 Subject: [PATCH 2/3] fix(model-runtime): retain opaque state in mixed content --- src/graphon/model_runtime/README.md | 11 +++ .../entities/message_entities.py | 4 +- .../base/large_language_model.py | 10 ++- tests/model_runtime/test_message_entities.py | 14 +++ tests/model_runtime/test_model_dispatch.py | 90 +++++++++++++++++++ .../test_streaming_invoke_accumulator.py | 2 +- 6 files changed, 127 insertions(+), 4 deletions(-) diff --git a/src/graphon/model_runtime/README.md b/src/graphon/model_runtime/README.md index bf94d662..727b521e 100644 --- a/src/graphon/model_runtime/README.md +++ b/src/graphon/model_runtime/README.md @@ -22,3 +22,14 @@ for invocation. [Entities](entities/) describe model parameters and credential forms for host UIs. The [model dispatch tests](../../../tests/model_runtime/test_model_dispatch.py) verify that provider discovery and individual capabilities remain independent. + +## Preserve provider message state + +Use `opaque_body` on [assistant messages and content items](entities/message_entities.py) for provider state that must survive JSON serialization, such as signatures or encrypted reasoning. +The field accepts `JsonValue` and defaults to `None`; the adapter owns its meaning. +An assistant message with any non-`None` state is nonempty even without visible content or tool calls. +See the [message entity tests](../../../tests/model_runtime/test_message_entities.py) for serialization and empty-message behavior. + +When consuming chunks, [LargeLanguageModel](model_providers/base/large_language_model.py) retains the last non-`None` assistant snapshot within each invocation for both the non-stream result and the result passed to `on_after_invoke`. +Adapters should emit complete snapshots because each replaces the previous value, including empty collections, empty strings, zero, and `False`. +Aggregation preserves the order of mixed strings and content blocks, including each block's state; see the [model dispatch tests](../../../tests/model_runtime/test_model_dispatch.py). diff --git a/src/graphon/model_runtime/entities/message_entities.py b/src/graphon/model_runtime/entities/message_entities.py index a022a931..00b10173 100644 --- a/src/graphon/model_runtime/entities/message_entities.py +++ b/src/graphon/model_runtime/entities/message_entities.py @@ -262,8 +262,8 @@ def transform_id_to_str(cls, value: Any) -> str: opaque_body: JsonValue | None = None def is_empty(self) -> bool: - """Check whether the assistant message has no content or tool calls.""" - return super().is_empty() and not self.tool_calls + """Check whether the assistant message has no content, calls, or state.""" + return super().is_empty() and not self.tool_calls and self.opaque_body is None class SystemPromptMessage(PromptMessage): 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 76036248..5debbd05 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 @@ -143,9 +143,17 @@ def consume(self, chunk: LLMResultChunk) -> None: def _consume_content(self, chunk: LLMResultChunk) -> None: content = chunk.delta.message.content + if not content: + return if isinstance(content, str): - self.content += content + if self.content_list: + self.content_list.append(TextPromptMessageContent(data=content)) + else: + self.content += content elif isinstance(content, list): + if self.content: + self.content_list.append(TextPromptMessageContent(data=self.content)) + self.content = "" self.content_list.extend(content) def to_result( diff --git a/tests/model_runtime/test_message_entities.py b/tests/model_runtime/test_message_entities.py index de49299c..e1776585 100644 --- a/tests/model_runtime/test_message_entities.py +++ b/tests/model_runtime/test_message_entities.py @@ -1,3 +1,6 @@ +import pytest +from pydantic import JsonValue + from graphon.model_runtime.entities.llm_entities import ( LLMResultChunk, LLMResultChunkDelta, @@ -108,3 +111,14 @@ def test_llm_result_chunk_json_round_trip_preserves_opaque_body() -> None: restored = LLMResultChunk.model_validate_json(chunk.model_dump_json()) assert restored.delta.message.opaque_body == opaque_body + + +@pytest.mark.parametrize("opaque_body", [{}, [], "", 0, False, {"signature": "s"}]) +def test_assistant_with_opaque_body_is_not_empty(opaque_body: JsonValue) -> None: + message = AssistantPromptMessage(opaque_body=opaque_body) + restored = AssistantPromptMessage.model_validate_json(message.model_dump_json()) + + assert restored.opaque_body == opaque_body + assert type(restored.opaque_body) is type(opaque_body) + assert not restored.is_empty() + assert AssistantPromptMessage().is_empty() diff --git a/tests/model_runtime/test_model_dispatch.py b/tests/model_runtime/test_model_dispatch.py index 4a83620b..b17db45b 100644 --- a/tests/model_runtime/test_model_dispatch.py +++ b/tests/model_runtime/test_model_dispatch.py @@ -5,6 +5,7 @@ from typing import IO, Any, cast import pytest +from pydantic import JsonValue from graphon.model_runtime.callbacks.base_callback import Callback from graphon.model_runtime.entities.common_entities import I18nObject @@ -741,6 +742,95 @@ def test_large_language_model_stream_after_invoke_preserves_chunk_state() -> Non assert after_result.system_fingerprint == "fp-2" +@pytest.mark.parametrize("stream", [True, False]) +@pytest.mark.parametrize("mixed_content", [True, False]) +@pytest.mark.parametrize("opaque_body", [{"signature": "s"}, {}, [], "", 0, False]) +def test_large_language_model_uses_last_opaque_body_per_invocation( + stream: bool, + mixed_content: bool, + opaque_body: JsonValue, +) -> None: + provider = ProviderEntity( + provider="test-provider", + label=I18nObject(en_US="Test Provider"), + supported_model_types=[ModelType.LLM], + configurate_methods=[], + ) + runtime = _StreamingLLMRuntimeStub([]) + model = LargeLanguageModel( + provider_schema=provider, + model_runtime=cast("LLMModelRuntime", runtime), + ) + callback = _RecordingCallback() + for snapshot in (opaque_body, None): + content_block = TextPromptMessageContent(data="", opaque_body=snapshot) + chunk_contents: list[str | list[PromptMessageContentUnionTypes]] = ( + ["answer", [content_block], "tail", ""] + if mixed_content + else ["answer", "", "", ""] + ) + runtime._chunks = tuple( + LLMResultChunk.model_validate({ + "model": "test-model", + "delta": { + "index": index, + "message": { + "content": chunk_contents[index], + "opaque_body": chunk_snapshot, + }, + }, + }) + for index, chunk_snapshot in enumerate( + [{"partial": True}, snapshot, None, None] + if snapshot is not None + else [None], + ) + ) + result = model.invoke( + model="test-model", + credentials={}, + prompt_messages=[UserPromptMessage(content="hello")], + stream=stream, + callbacks=[callback], + ) + if isinstance(result, LLMResult): + assert result.message.opaque_body == snapshot + assert type(result.message.opaque_body) is type(snapshot) + if mixed_content and snapshot is not None: + assert result.message.content == [ + TextPromptMessageContent(data="answer"), + content_block, + TextPromptMessageContent(data="tail"), + ] + assert isinstance(result.message.content, list) + assert result.message.content[1].opaque_body == snapshot + assert type(result.message.content[1].opaque_body) is type(snapshot) + else: + chunks = list(result) + assert chunks[-1].delta.message.opaque_body is None + if mixed_content and snapshot is not None: + assert [ + chunk.delta.message.content for chunk in chunks + ] == chunk_contents + blocks = chunks[1].delta.message.content + assert isinstance(blocks, list) + assert blocks[0].opaque_body == snapshot + assert type(blocks[0].opaque_body) is type(snapshot) + + assert callback.after_results[-1].message.opaque_body == snapshot + assert type(callback.after_results[-1].message.opaque_body) is type(snapshot) + if mixed_content and snapshot is not None: + blocks = callback.after_results[-1].message.content + assert blocks == [ + TextPromptMessageContent(data="answer"), + content_block, + TextPromptMessageContent(data="tail"), + ] + assert isinstance(blocks, list) + assert blocks[1].opaque_body == snapshot + assert type(blocks[1].opaque_body) is type(snapshot) + + def test_large_language_model_stream_errors_are_transformed() -> None: provider = ProviderEntity( provider="test-provider", diff --git a/tests/model_runtime/test_streaming_invoke_accumulator.py b/tests/model_runtime/test_streaming_invoke_accumulator.py index 4366ae0a..b1598685 100644 --- a/tests/model_runtime/test_streaming_invoke_accumulator.py +++ b/tests/model_runtime/test_streaming_invoke_accumulator.py @@ -53,7 +53,7 @@ def test_streaming_invoke_accumulator_opaque_body_defaults_to_none() -> None: assert result.message.opaque_body is None -def test_streaming_invoke_accumulator_none_chunk_does_not_clobber_snapshot() -> None: +def test_streaming_invoke_accumulator_preserves_opaque_body_after_none() -> None: accumulator = _StreamingInvokeAccumulator(real_model="test-model") opaque_body = { "assistant_blocks": [{"type": "redacted_thinking", "data": "enc-1"}], From ab3dceaa2466127444d0862738e63adf230e4848 Mon Sep 17 00:00:00 2001 From: WH-2099 Date: Thu, 24 Sep 2026 10:15:50 +0000 Subject: [PATCH 3/3] fix(llm): preserve opaque state across Slim and prompt filtering --- src/graphon/dsl/slim/llm.py | 88 +++++++++------------------ src/graphon/model_runtime/README.md | 3 + src/graphon/nodes/llm/llm_utils.py | 1 + tests/dsl/test_slim_llm.py | 93 ++++++++++++++++++++++++++++- tests/nodes/llm/test_llm_utils.py | 54 +++++++++++++++++ 5 files changed, 178 insertions(+), 61 deletions(-) diff --git a/src/graphon/dsl/slim/llm.py b/src/graphon/dsl/slim/llm.py index 33ad9205..f58e522d 100644 --- a/src/graphon/dsl/slim/llm.py +++ b/src/graphon/dsl/slim/llm.py @@ -3,7 +3,7 @@ import json import logging from collections.abc import Generator, Iterable, Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import Any, Literal, overload, override from pydantic import StrictStr, TypeAdapter, ValidationError @@ -31,7 +31,7 @@ ) from graphon.model_runtime.entities.model_entities import AIModelEntity from graphon.model_runtime.model_providers.base.large_language_model import ( - merge_tool_call_deltas, + normalize_non_stream_runtime_result, ) from graphon.model_runtime.model_providers.base.tokenizers.gpt2_tokenizer import ( GPT2Tokenizer, @@ -85,16 +85,6 @@ def finalize(self, *, expect_structured_output: bool) -> Mapping[str, Any] | Non return None -@dataclass(slots=True) -class _CollectedLLMResult: - content_text: str = "" - content_parts: list[Any] = field(default_factory=list) - usage: LLMUsage = field(default_factory=LLMUsage.empty_usage) - tool_calls: list[AssistantPromptMessage.ToolCall] = field(default_factory=list) - structured_output: Mapping[str, Any] | None = None - system_fingerprint: str | None = None - - class SlimLLM(LLMProtocol): """Standard Slim-backed LLM runtime. @@ -502,67 +492,45 @@ def _collect_llm_result( chunks: Iterable[LLMResultChunk], expect_structured_output: bool, ) -> LLMResult: - collected = _CollectedLLMResult() + system_fingerprint: str | None = None structured_output_accumulator = ( _StructuredOutputAccumulator() if expect_structured_output else None ) - for chunk in chunks: - _accumulate_llm_chunk( - collected=collected, - chunk=chunk, - structured_output_accumulator=structured_output_accumulator, - ) + def collect_structured_output_and_fingerprint() -> Generator[ + LLMResultChunk, None, None + ]: + nonlocal system_fingerprint + for chunk in chunks: + _consume_structured_output_chunk( + chunk=chunk, + accumulator=structured_output_accumulator, + ) + if chunk.system_fingerprint is not None: + system_fingerprint = chunk.system_fingerprint + yield chunk - collected.structured_output = _finalize_structured_output( + result = normalize_non_stream_runtime_result( + model=model, + prompt_messages=prompt_messages, + result=collect_structured_output_and_fingerprint(), + ) + result.system_fingerprint = system_fingerprint + structured_output = _finalize_structured_output( accumulator=structured_output_accumulator, expect_structured_output=expect_structured_output, ) - assistant_message = AssistantPromptMessage( - content=collected.content_text or collected.content_parts, - tool_calls=collected.tool_calls, - ) - if collected.structured_output is not None: + if structured_output is not None: return LLMResultWithStructuredOutput( model=model, prompt_messages=list(prompt_messages), - message=assistant_message, - usage=collected.usage, - system_fingerprint=collected.system_fingerprint, - structured_output=collected.structured_output, + message=result.message, + usage=result.usage, + system_fingerprint=system_fingerprint, + structured_output=structured_output, ) - return LLMResult( - model=model, - prompt_messages=list(prompt_messages), - message=assistant_message, - usage=collected.usage, - system_fingerprint=collected.system_fingerprint, - ) - - -def _accumulate_llm_chunk( - *, - collected: _CollectedLLMResult, - chunk: LLMResultChunk, - structured_output_accumulator: _StructuredOutputAccumulator | None = None, -) -> None: - delta_message = chunk.delta.message - if isinstance(delta_message.content, str): - collected.content_text += delta_message.content - elif isinstance(delta_message.content, list): - collected.content_parts.extend(delta_message.content) - - if delta_message.tool_calls: - merge_tool_call_deltas(delta_message.tool_calls, collected.tool_calls) - if chunk.delta.usage is not None: - collected.usage = chunk.delta.usage - _consume_structured_output_chunk( - chunk=chunk, - accumulator=structured_output_accumulator, - ) - if chunk.system_fingerprint is not None: - collected.system_fingerprint = chunk.system_fingerprint + return result def _consume_structured_output_chunk( diff --git a/src/graphon/model_runtime/README.md b/src/graphon/model_runtime/README.md index f4d7155e..276889dc 100644 --- a/src/graphon/model_runtime/README.md +++ b/src/graphon/model_runtime/README.md @@ -33,3 +33,6 @@ See the [message entity tests](../../../tests/model_runtime/test_message_entitie When consuming chunks, [LargeLanguageModel](model_providers/base/large_language_model.py) retains the last non-`None` assistant snapshot within each invocation for both the non-stream result and the result passed to `on_after_invoke`. Adapters should emit complete snapshots because each replaces the previous value, including empty collections, empty strings, zero, and `False`. Aggregation preserves the order of mixed strings and content blocks, including each block's state; see the [model dispatch tests](../../../tests/model_runtime/test_model_dispatch.py). + +[SlimLLM](../dsl/slim/llm.py) uses the same aggregation for blocking ordinary and structured output calls; see the [Slim tests](../../../tests/dsl/test_slim_llm.py). +[Prompt filtering](../nodes/llm/llm_utils.py) preserves state on supported content blocks and keeps assistant messages with non-`None` state even when their content list is empty; see the [prompt tests](../../../tests/nodes/llm/test_llm_utils.py). diff --git a/src/graphon/nodes/llm/llm_utils.py b/src/graphon/nodes/llm/llm_utils.py index 7ee3b4aa..a853c90a 100644 --- a/src/graphon/nodes/llm/llm_utils.py +++ b/src/graphon/nodes/llm/llm_utils.py @@ -215,6 +215,7 @@ def _filter_prompt_messages( if ( len(prompt_message_content) == 1 and prompt_message_content[0].type == PromptMessageContentType.TEXT + and prompt_message_content[0].opaque_body is None ): prompt_message.content = prompt_message_content[0].data else: diff --git a/tests/dsl/test_slim_llm.py b/tests/dsl/test_slim_llm.py index 4e3ee35d..af455d27 100644 --- a/tests/dsl/test_slim_llm.py +++ b/tests/dsl/test_slim_llm.py @@ -9,11 +9,16 @@ from typing import Any import pytest +from pydantic import JsonValue import graphon.dsl.slim.llm as slim_llm_module from graphon.dsl.slim import SlimClientConfig, SlimClientError, SlimLLM from graphon.model_runtime.entities.llm_entities import LLMResult -from graphon.model_runtime.entities.message_entities import SystemPromptMessage +from graphon.model_runtime.entities.message_entities import ( + AssistantPromptMessage, + SystemPromptMessage, + TextPromptMessageContent, +) from graphon.model_runtime.model_providers.base.tokenizers import gpt2_tokenizer @@ -209,6 +214,92 @@ def test_slim_llm_passes_merged_parameters_and_json_schema( assert json.loads(model_parameters["json_schema"]) == schema +@pytest.mark.parametrize("use_structured_output", [False, True]) +@pytest.mark.parametrize("opaque_body", [{"signature": "s"}, {}, [], "", 0, False]) +def test_slim_llm_collects_mixed_content_and_latest_opaque_body( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + use_structured_output: bool, + opaque_body: JsonValue, +) -> None: + client = _patch_recording_slim_client(monkeypatch) + block = TextPromptMessageContent(data="", opaque_body={"signature": "block"}) + chunks: list[dict[str, Any]] = [ + { + "delta": { + "message": {"content": content, "opaque_body": chunk_opaque_body} + }, + "structured_output": {"ok": True} if index == 1 else None, + } + for index, (content, chunk_opaque_body) in enumerate([ + ("before", {"partial": True}), + ([block.model_dump()], None), + ("after", None), + ("", opaque_body), + ("", None), + ]) + ] + for index, arguments in enumerate(['{"city":', '"Paris"}']): + chunks[index]["delta"]["message"]["tool_calls"] = [ + { + "id": "call-1" if index == 0 else "", + "type": "function", + "function": { + "name": "weather" if index == 0 else "", + "arguments": arguments, + }, + }, + ] + chunks[0]["system_fingerprint"] = "first" + chunks[-2]["system_fingerprint"] = "" + chunks[-1]["delta"]["usage"] = {"prompt_tokens": 3, "completion_tokens": 2} + monkeypatch.setattr(client, "invoke_chunks", lambda **_: iter(chunks)) + llm = _build_llm(tmp_path) + + if use_structured_output: + result = llm.invoke_llm_with_structured_output( + prompt_messages=[], + json_schema={"type": "object"}, + model_parameters={}, + stop=None, + stream=False, + ) + assert result.structured_output == {"ok": True} + else: + result = llm.invoke_llm( + prompt_messages=[], + model_parameters={}, + tools=None, + stop=None, + stream=False, + ) + + assert result.message == AssistantPromptMessage( + content=[ + TextPromptMessageContent(data="before"), + block, + TextPromptMessageContent(data="after"), + ], + tool_calls=[ + AssistantPromptMessage.ToolCall( + id="call-1", + type="function", + function=AssistantPromptMessage.ToolCall.ToolCallFunction( + name="weather", arguments='{"city":"Paris"}' + ), + ), + ], + opaque_body=opaque_body, + ) + assert type(result.message.opaque_body) is type(opaque_body) + assert ( + result.usage.prompt_tokens, + result.usage.completion_tokens, + result.usage.total_tokens, + ) == (3, 2, 5) + assert result.system_fingerprint == "" + + def test_slim_llm_preserves_slim_client_errors( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, diff --git a/tests/nodes/llm/test_llm_utils.py b/tests/nodes/llm/test_llm_utils.py index b899bb9f..39035203 100644 --- a/tests/nodes/llm/test_llm_utils.py +++ b/tests/nodes/llm/test_llm_utils.py @@ -1,14 +1,18 @@ from unittest.mock import MagicMock import pytest +from pydantic import JsonValue from graphon.model_runtime.entities.message_entities import ( + AssistantPromptMessage, ImagePromptMessageContent, PromptMessageRole, + TextPromptMessageContent, ) from graphon.nodes.llm import llm_utils from graphon.nodes.llm.entities import LLMNodeChatModelMessage from graphon.nodes.llm.exc import NoPromptFoundError +from graphon.prompt_entities import MemoryConfig from graphon.template_rendering import TemplateRenderError from ...helpers import build_variable_pool @@ -44,6 +48,56 @@ def test_fetch_prompt_messages_renders_basic_variables_and_context() -> None: assert stop == ["done"] +@pytest.mark.parametrize("text", ["answer", ""]) +@pytest.mark.parametrize("opaque_body", [{"signature": "s"}, {}, [], "", 0, False]) +def test_fetch_prompt_messages_preserves_history_content_opaque_body( + text: str, + opaque_body: JsonValue, +) -> None: + block = TextPromptMessageContent(data=text, opaque_body=opaque_body) + history_message = AssistantPromptMessage(content=[block]) + memory = MagicMock() + memory.get_history_prompt_messages.return_value = [history_message] + model_instance = _model_instance() + model_instance.get_model_schema().model_properties = {} + + prompt_messages, _ = llm_utils.fetch_prompt_messages( + prompt_template=[], + sys_files=[], + memory=memory, + memory_config=MemoryConfig(window=MemoryConfig.WindowConfig(enabled=False)), + model_instance=model_instance, + vision_detail=ImagePromptMessageContent.DETAIL.HIGH, + variable_pool=build_variable_pool(), + jinja2_variables=[], + ) + + assert len(prompt_messages) == 1 + assert prompt_messages[0].content == [block] + + +def test_fetch_prompt_messages_preserves_opaque_body_without_content() -> None: + memory = MagicMock() + memory.get_history_prompt_messages.return_value = [ + AssistantPromptMessage(content=[], opaque_body={}), + ] + model_instance = _model_instance() + model_instance.get_model_schema().model_properties = {} + + prompt_messages, _ = llm_utils.fetch_prompt_messages( + prompt_template=[], + sys_files=[], + memory=memory, + memory_config=MemoryConfig(window=MemoryConfig.WindowConfig(enabled=False)), + model_instance=model_instance, + vision_detail=ImagePromptMessageContent.DETAIL.HIGH, + variable_pool=build_variable_pool(), + jinja2_variables=[], + ) + + assert prompt_messages == [AssistantPromptMessage(content=[], opaque_body={})] + + def test_fetch_prompt_messages_rejects_empty_jinja_prompt() -> None: with pytest.raises(NoPromptFoundError): llm_utils.fetch_prompt_messages(