feat: preserve opaque_body provider state on prompt messages and LLM results

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(

src/graphon/model_runtime/README.md

@@ -22,3 +22,17 @@ 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).
+
+[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).

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,10 +259,11 @@ 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."""
-        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):

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,12 +138,22 @@ 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
+        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(
@@ -155,6 +168,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 +181,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 +190,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 +213,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,
         )

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:

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,