Expand orchestration gate coverage to resume and regenerate flows.

This commit is contained in:
CREDO23 2026-05-07 16:18:29 +02:00
parent 52593d88db
commit 0f40279d95
3 changed files with 128 additions and 8 deletions

View file

@ -84,15 +84,20 @@ async def test_stream_agent_events_emits_text_lifecycle_and_updates_result() ->
]
assert result.accumulated_text == "Hello world"
assert result.agent_called_update_memory is False
assert agent.calls[0][1]["version"] == "v2"
async def test_stream_agent_events_passes_runtime_context_to_agent() -> None:
service = _StreamingService()
agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("x")}}])
class _ContextAwareAgent:
async def astream_events(self, input_data: Any, **kwargs: Any):
del input_data
text = "ctx-ok" if kwargs.get("context") else "ctx-missing"
yield {"event": "on_chat_model_stream", "data": {"chunk": _Chunk(text)}}
agent = _ContextAwareAgent()
result = StreamOutput()
_ = await _collect(
frames = await _collect(
stream_agent_events(
agent=agent,
config={"configurable": {"thread_id": "t-2"}},
@ -103,5 +108,8 @@ async def test_stream_agent_events_passes_runtime_context_to_agent() -> None:
)
)
assert agent.calls
assert agent.calls[0][1]["context"] == {"mentioned_document_ids": [1, 2]}
assert frames == [
"text_start:text-1",
"text_delta:text-1:ctx-ok",
"text_end:text-1",
]

View file

@ -8,7 +8,11 @@ from typing import Any
import pytest
from app.tasks.chat.streaming.orchestration import StreamExecutionInput
from app.tasks.chat.streaming.orchestration.orchestrator import stream_chat
from app.tasks.chat.streaming.orchestration.orchestrator import (
stream_chat,
stream_regenerate,
stream_resume,
)
pytestmark = pytest.mark.unit
@ -84,5 +88,53 @@ async def test_stream_chat_uses_orchestration_input_path() -> None:
"text_delta:text-1:!",
"text_end:text-1",
]
assert agent.calls
assert agent.calls[0][1]["version"] == "v2"
async def test_stream_resume_uses_orchestration_input_path() -> None:
service = _StreamingService()
agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("r")}}])
frames = await _collect(
stream_resume(
chat_id=9,
search_space_id=1,
decisions=[],
orchestration_input=StreamExecutionInput(
agent=agent,
config={"configurable": {"thread_id": "thread-r"}},
input_data={"messages": []},
streaming_service=service,
),
)
)
assert frames == [
"text_start:text-1",
"text_delta:text-1:r",
"text_end:text-1",
]
async def test_stream_regenerate_uses_orchestration_input_path() -> None:
service = _StreamingService()
agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("g")}}])
frames = await _collect(
stream_regenerate(
user_query="q",
search_space_id=1,
chat_id=2,
orchestration_input=StreamExecutionInput(
agent=agent,
config={"configurable": {"thread_id": "thread-g"}},
input_data={"messages": []},
streaming_service=service,
),
)
)
assert frames == [
"text_start:text-1",
"text_delta:text-1:g",
"text_end:text-1",
]