Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/graphon/model_runtime/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,9 @@ 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).

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).

[Prompt filtering](../nodes/llm/llm_utils.py) preserves state on supported blocks
and assistant messages with state but no content; see the
[prompt tests](../../../tests/nodes/llm/test_llm_utils.py).
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,9 @@ def _merge_tool_call_delta(


@dataclass(slots=True)
class _LLMChunkAccumulator:
class LLMResultAccumulator:
"""Retain a complete result while a caller consumes and displays each chunk."""

content: str = ""
content_list: list[PromptMessageContentUnionTypes] = field(default_factory=list)
usage: LLMUsage = field(default_factory=LLMUsage.empty_usage)
Expand Down Expand Up @@ -242,7 +244,7 @@ def _build_llm_result_from_chunks(
A normalized `LLMResult` assembled from the consumed chunks.

"""
accumulator = _LLMChunkAccumulator()
accumulator = LLMResultAccumulator()
try:
accumulator.consume_all(chunks)
except Exception:
Expand Down
49 changes: 49 additions & 0 deletions tests/model_runtime/test_llm_result_accumulator.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
from graphon.model_runtime.entities.llm_entities import (
LLMResultChunk,
LLMResultChunkDelta,
)
from graphon.model_runtime.entities.message_entities import (
AssistantPromptMessage,
TextPromptMessageContent,
)
from graphon.model_runtime.model_providers.base import large_language_model


def test_accumulated_result_preserves_message_through_json() -> None:
accumulator = large_language_model.LLMResultAccumulator()
tool_call = AssistantPromptMessage.ToolCall(
id="call-1",
type="function",
function=AssistantPromptMessage.ToolCall.ToolCallFunction(
name="lookup",
arguments='{"id":9007199254740993}',
),
)
block = TextPromptMessageContent(data="", opaque_body={"signature": "sig"})
messages = [
AssistantPromptMessage(content="before"),
AssistantPromptMessage(
content=[block], tool_calls=[tool_call], opaque_body=False
),
AssistantPromptMessage(content="after"),
]
for index, message in enumerate(messages):
accumulator.consume(
LLMResultChunk(
model="test-model",
delta=LLMResultChunkDelta(index=index, message=message),
)
)

result = accumulator.to_result(model="test-model", prompt_messages=[])
restored_message = AssistantPromptMessage.model_validate_json(
result.message.model_dump_json()
)

assert restored_message.content == [
TextPromptMessageContent(data="before"),
TextPromptMessageContent(data="", opaque_body={"signature": "sig"}),
TextPromptMessageContent(data="after"),
]
assert restored_message.tool_calls == [tool_call]
assert restored_message.opaque_body is False
Loading