mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-05-12 09:12:40 +02:00
wire orchestrator streaming context path and align event relay outputs
This commit is contained in:
CREDO23
parent
0f40279d95
commit
a04b2e88bd
8 changed files with 94 additions and 109 deletions
|
|
@ -1,11 +1,11 @@
|
|||
"""Composable orchestration pieces for chat streaming."""
|
||||
|
||||
from app.tasks.chat.streaming.orchestration.event_stream import stream_agent_events
|
||||
from app.tasks.chat.streaming.orchestration.input import StreamExecutionInput
|
||||
from app.tasks.chat.streaming.orchestration.output import StreamOutput
|
||||
from app.tasks.chat.streaming.orchestration.event_stream import stream_output
|
||||
from app.tasks.chat.streaming.orchestration.input import StreamingContext
|
||||
from app.tasks.chat.streaming.orchestration.output import StreamingResult
|
||||
|
||||
__all__ = [
|
||||
"StreamExecutionInput",
|
||||
"StreamOutput",
|
||||
"stream_agent_events",
|
||||
"StreamingContext",
|
||||
"StreamingResult",
|
||||
"stream_output",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -6,18 +6,18 @@ from collections.abc import AsyncIterator
|
|||
from typing import Any
|
||||
|
||||
from app.agents.new_chat.feature_flags import get_flags
|
||||
from app.tasks.chat.streaming.orchestration.output import StreamOutput
|
||||
from app.tasks.chat.streaming.orchestration.output import StreamingResult
|
||||
from app.tasks.chat.streaming.relay.event_relay import EventRelay
|
||||
from app.tasks.chat.streaming.relay.state import AgentEventRelayState
|
||||
|
||||
|
||||
async def stream_agent_events(
|
||||
async def stream_output(
|
||||
*,
|
||||
agent: Any,
|
||||
config: dict[str, Any],
|
||||
input_data: Any,
|
||||
streaming_service: Any,
|
||||
result: StreamOutput,
|
||||
result: StreamingResult,
|
||||
step_prefix: str = "thinking",
|
||||
initial_step_id: str | None = None,
|
||||
initial_step_title: str = "",
|
||||
|
|
|
|||
|
|
@ -7,8 +7,8 @@ from typing import Any
|
|||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StreamExecutionInput:
|
||||
"""Container for dependencies required by ``stream_agent_events``."""
|
||||
class StreamingContext:
|
||||
"""Container for dependencies required by ``stream_output``."""
|
||||
|
||||
agent: Any
|
||||
config: dict[str, Any]
|
||||
|
|
|
|||
|
|
@ -1,9 +1,4 @@
|
|||
"""Top-level chat streaming entrypoints.
|
||||
|
||||
For now these orchestrator functions are thin compatibility wrappers around the
|
||||
current ``stream_new_chat`` / ``stream_resume_chat`` implementations. Routing
|
||||
calls through this module lets us cut over to the fully modular event relay in
|
||||
one place later without touching API routes again.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -14,9 +9,47 @@ from typing import Any, Literal
|
|||
from app.agents.new_chat.filesystem_selection import FilesystemSelection
|
||||
from app.db import ChatVisibility
|
||||
from app.tasks.chat.stream_new_chat import stream_new_chat, stream_resume_chat
|
||||
from app.tasks.chat.streaming.orchestration.event_stream import stream_agent_events
|
||||
from app.tasks.chat.streaming.orchestration.input import StreamExecutionInput
|
||||
from app.tasks.chat.streaming.orchestration.output import StreamOutput
|
||||
from app.tasks.chat.streaming.orchestration.event_stream import stream_output
|
||||
from app.tasks.chat.streaming.orchestration.input import StreamingContext
|
||||
from app.tasks.chat.streaming.orchestration.output import StreamingResult
|
||||
|
||||
|
||||
def _build_streaming_result(
|
||||
*,
|
||||
chat_id: int,
|
||||
request_id: str | None,
|
||||
filesystem_selection: FilesystemSelection | None,
|
||||
suffix: str,
|
||||
) -> StreamingResult:
|
||||
return StreamingResult(
|
||||
request_id=request_id,
|
||||
turn_id=f"{chat_id}:{suffix}",
|
||||
filesystem_mode=(filesystem_selection.mode.value if filesystem_selection else "cloud"),
|
||||
client_platform=(
|
||||
filesystem_selection.client_platform.value if filesystem_selection else "web"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _stream_output_with_streaming_context(
|
||||
*,
|
||||
streaming_context: StreamingContext,
|
||||
result: StreamingResult,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
async for frame in stream_output(
|
||||
agent=streaming_context.agent,
|
||||
config=streaming_context.config,
|
||||
input_data=streaming_context.input_data,
|
||||
streaming_service=streaming_context.streaming_service,
|
||||
result=result,
|
||||
step_prefix=streaming_context.step_prefix,
|
||||
initial_step_id=streaming_context.initial_step_id,
|
||||
initial_step_title=streaming_context.initial_step_title,
|
||||
initial_step_items=streaming_context.initial_step_items,
|
||||
content_builder=streaming_context.content_builder,
|
||||
runtime_context=streaming_context.runtime_context,
|
||||
):
|
||||
yield frame
|
||||
|
||||
|
||||
async def stream_chat(
|
||||
|
|
@ -37,34 +70,19 @@ async def stream_chat(
|
|||
filesystem_selection: FilesystemSelection | None = None,
|
||||
request_id: str | None = None,
|
||||
user_image_data_urls: list[str] | None = None,
|
||||
orchestration_input: StreamExecutionInput | None = None,
|
||||
streaming_context: StreamingContext | None = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Stream a new chat turn through the current production pipeline."""
|
||||
if orchestration_input is not None:
|
||||
result = StreamOutput(
|
||||
if streaming_context is not None:
|
||||
result = _build_streaming_result(
|
||||
chat_id=chat_id,
|
||||
request_id=request_id,
|
||||
turn_id=f"{chat_id}:orchestrator",
|
||||
filesystem_mode=(
|
||||
filesystem_selection.mode.value if filesystem_selection else "cloud"
|
||||
),
|
||||
client_platform=(
|
||||
filesystem_selection.client_platform.value
|
||||
if filesystem_selection
|
||||
else "web"
|
||||
),
|
||||
filesystem_selection=filesystem_selection,
|
||||
suffix="orchestrator",
|
||||
)
|
||||
async for frame in stream_agent_events(
|
||||
agent=orchestration_input.agent,
|
||||
config=orchestration_input.config,
|
||||
input_data=orchestration_input.input_data,
|
||||
streaming_service=orchestration_input.streaming_service,
|
||||
async for frame in _stream_output_with_streaming_context(
|
||||
streaming_context=streaming_context,
|
||||
result=result,
|
||||
step_prefix=orchestration_input.step_prefix,
|
||||
initial_step_id=orchestration_input.initial_step_id,
|
||||
initial_step_title=orchestration_input.initial_step_title,
|
||||
initial_step_items=orchestration_input.initial_step_items,
|
||||
content_builder=orchestration_input.content_builder,
|
||||
runtime_context=orchestration_input.runtime_context,
|
||||
):
|
||||
yield frame
|
||||
return
|
||||
|
|
@ -101,34 +119,19 @@ async def stream_resume(
|
|||
filesystem_selection: FilesystemSelection | None = None,
|
||||
request_id: str | None = None,
|
||||
disabled_tools: list[str] | None = None,
|
||||
orchestration_input: StreamExecutionInput | None = None,
|
||||
streaming_context: StreamingContext | None = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Resume an interrupted chat turn through the current production pipeline."""
|
||||
if orchestration_input is not None:
|
||||
result = StreamOutput(
|
||||
if streaming_context is not None:
|
||||
result = _build_streaming_result(
|
||||
chat_id=chat_id,
|
||||
request_id=request_id,
|
||||
turn_id=f"{chat_id}:orchestrator-resume",
|
||||
filesystem_mode=(
|
||||
filesystem_selection.mode.value if filesystem_selection else "cloud"
|
||||
),
|
||||
client_platform=(
|
||||
filesystem_selection.client_platform.value
|
||||
if filesystem_selection
|
||||
else "web"
|
||||
),
|
||||
filesystem_selection=filesystem_selection,
|
||||
suffix="orchestrator-resume",
|
||||
)
|
||||
async for frame in stream_agent_events(
|
||||
agent=orchestration_input.agent,
|
||||
config=orchestration_input.config,
|
||||
input_data=orchestration_input.input_data,
|
||||
streaming_service=orchestration_input.streaming_service,
|
||||
async for frame in _stream_output_with_streaming_context(
|
||||
streaming_context=streaming_context,
|
||||
result=result,
|
||||
step_prefix=orchestration_input.step_prefix,
|
||||
initial_step_id=orchestration_input.initial_step_id,
|
||||
initial_step_title=orchestration_input.initial_step_title,
|
||||
initial_step_items=orchestration_input.initial_step_items,
|
||||
content_builder=orchestration_input.content_builder,
|
||||
runtime_context=orchestration_input.runtime_context,
|
||||
):
|
||||
yield frame
|
||||
return
|
||||
|
|
@ -166,34 +169,19 @@ async def stream_regenerate(
|
|||
request_id: str | None = None,
|
||||
user_image_data_urls: list[str] | None = None,
|
||||
flow: Literal["new", "regenerate"] = "regenerate",
|
||||
orchestration_input: StreamExecutionInput | None = None,
|
||||
streaming_context: StreamingContext | None = None,
|
||||
) -> AsyncGenerator[str, None]:
|
||||
"""Regenerate an assistant turn through the current production pipeline."""
|
||||
if orchestration_input is not None:
|
||||
result = StreamOutput(
|
||||
if streaming_context is not None:
|
||||
result = _build_streaming_result(
|
||||
chat_id=chat_id,
|
||||
request_id=request_id,
|
||||
turn_id=f"{chat_id}:orchestrator-regenerate",
|
||||
filesystem_mode=(
|
||||
filesystem_selection.mode.value if filesystem_selection else "cloud"
|
||||
),
|
||||
client_platform=(
|
||||
filesystem_selection.client_platform.value
|
||||
if filesystem_selection
|
||||
else "web"
|
||||
),
|
||||
filesystem_selection=filesystem_selection,
|
||||
suffix="orchestrator-regenerate",
|
||||
)
|
||||
async for frame in stream_agent_events(
|
||||
agent=orchestration_input.agent,
|
||||
config=orchestration_input.config,
|
||||
input_data=orchestration_input.input_data,
|
||||
streaming_service=orchestration_input.streaming_service,
|
||||
async for frame in _stream_output_with_streaming_context(
|
||||
streaming_context=streaming_context,
|
||||
result=result,
|
||||
step_prefix=orchestration_input.step_prefix,
|
||||
initial_step_id=orchestration_input.initial_step_id,
|
||||
initial_step_title=orchestration_input.initial_step_title,
|
||||
initial_step_items=orchestration_input.initial_step_items,
|
||||
content_builder=orchestration_input.content_builder,
|
||||
runtime_context=orchestration_input.runtime_context,
|
||||
):
|
||||
yield frame
|
||||
return
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ from typing import Any
|
|||
|
||||
|
||||
@dataclass
|
||||
class StreamOutput:
|
||||
class StreamingResult:
|
||||
accumulated_text: str = ""
|
||||
is_interrupted: bool = False
|
||||
interrupt_value: dict[str, Any] | None = None
|
||||
|
|
@ -27,6 +27,3 @@ class StreamOutput:
|
|||
assistant_message_id: int | None = None
|
||||
content_builder: Any | None = field(default=None, repr=False)
|
||||
|
||||
|
||||
# Backwards-compatible alias while imports migrate.
|
||||
StreamResult = StreamOutput
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ from app.tasks.chat.streaming.handlers.custom_event_dispatch import (
|
|||
)
|
||||
from app.tasks.chat.streaming.handlers.tool_end import iter_tool_end_frames
|
||||
from app.tasks.chat.streaming.handlers.tool_start import iter_tool_start_frames
|
||||
from app.tasks.chat.streaming.orchestration.output import StreamOutput
|
||||
from app.tasks.chat.streaming.orchestration.output import StreamingResult
|
||||
from app.tasks.chat.streaming.relay.state import AgentEventRelayState
|
||||
from app.tasks.chat.streaming.relay.thinking_step_completion import (
|
||||
complete_active_thinking_step,
|
||||
|
|
@ -52,7 +52,7 @@ class EventRelay:
|
|||
events: AsyncIterator[dict[str, Any]],
|
||||
*,
|
||||
state: AgentEventRelayState,
|
||||
result: StreamOutput,
|
||||
result: StreamingResult,
|
||||
step_prefix: str = "thinking",
|
||||
content_builder: Any | None = None,
|
||||
config: dict[str, Any] | None = None,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue