apunkt/SurfSense
1
0
Fork 0
mirror of https://github.com/MODSetter/SurfSense.git synced 2026-05-12 09:12:40 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

build streaming contexts for chat resume and regenerate paths

Browse source
This commit is contained in:
CREDO23 2026-05-07 17:57:27 +02:00
parent a04b2e88bd
commit 52895e37e9
6 changed files with 633 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

100
surfsense_backend/tests/unit/tasks/chat/streaming/test_orchestrator_stream_chat.py
View file

@ -8,6 +8,7 @@ from typing import Any
import pytest
from app.tasks.chat.streaming.orchestration import StreamingContext
from app.tasks.chat.streaming.orchestration import orchestrator
from app.tasks.chat.streaming.orchestration.orchestrator import (
stream_chat,
stream_regenerate,
@ -138,3 +139,102 @@ async def test_stream_regenerate_uses_streaming_context_path() -> None:
"text_delta:text-1:g",
"text_end:text-1",
]
async def test_stream_chat_builds_streaming_context_when_not_provided() -> None:
service = _StreamingService()
agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("b")}}])
async def _fake_builder(**kwargs: Any) -> StreamingContext:
del kwargs
return StreamingContext(
agent=agent,
config={"configurable": {"thread_id": "thread-b"}},
input_data={"messages": []},
streaming_service=service,
)
old = orchestrator.build_chat_streaming_context
orchestrator.build_chat_streaming_context = _fake_builder
try:
frames = await _collect(
stream_chat(
user_query="q",
search_space_id=1,
chat_id=3,
)
)
finally:
orchestrator.build_chat_streaming_context = old
assert frames == [
"text_start:text-1",
"text_delta:text-1:b",
"text_end:text-1",
]
async def test_stream_resume_builds_streaming_context_when_not_provided() -> None:
service = _StreamingService()
agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("u")}}])
async def _fake_builder(**kwargs: Any) -> StreamingContext:
del kwargs
return StreamingContext(
agent=agent,
config={"configurable": {"thread_id": "thread-u"}},
input_data={"messages": []},
streaming_service=service,
)
old = orchestrator.build_resume_streaming_context
orchestrator.build_resume_streaming_context = _fake_builder
try:
frames = await _collect(
stream_resume(
chat_id=9,
search_space_id=1,
decisions=[],
)
)
finally:
orchestrator.build_resume_streaming_context = old
assert frames == [
"text_start:text-1",
"text_delta:text-1:u",
"text_end:text-1",
]
async def test_stream_regenerate_builds_streaming_context_when_not_provided() -> None:
service = _StreamingService()
agent = _Agent([{"event": "on_chat_model_stream", "data": {"chunk": _Chunk("x")}}])
async def _fake_builder(**kwargs: Any) -> StreamingContext:
del kwargs
return StreamingContext(
agent=agent,
config={"configurable": {"thread_id": "thread-x"}},
input_data={"messages": []},
streaming_service=service,
)
old = orchestrator.build_regenerate_streaming_context
orchestrator.build_regenerate_streaming_context = _fake_builder
try:
frames = await _collect(
stream_regenerate(
user_query="q",
search_space_id=1,
chat_id=2,
)
)
finally:
orchestrator.build_regenerate_streaming_context = old
assert frames == [
"text_start:text-1",
"text_delta:text-1:x",
"text_end:text-1",
]