mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-05-12 09:12:40 +02:00
Expand orchestration gate coverage to resume and regenerate flows.
This commit is contained in:
parent
52593d88db
commit
0f40279d95
3 changed files with 128 additions and 8 deletions
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue