mirror of
https://github.com/MODSetter/SurfSense.git
synced 2026-05-12 17:22:38 +02:00
Add chat tool streaming registry with shared, default, and connector tools.
This commit is contained in:
parent
ee16e1d5f9
commit
1392abf5b1
13 changed files with 275 additions and 0 deletions
|
|
@ -0,0 +1,23 @@
|
|||
"""Per-tool streaming: thinking-step and completion emission."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from app.tasks.chat.streaming.handlers.tools.emission_context import (
|
||||
ToolCompletionEmissionContext,
|
||||
)
|
||||
from app.tasks.chat.streaming.handlers.tools.registry import (
|
||||
iter_tool_completion_emission_frames,
|
||||
resolve_tool_completed_thinking_step,
|
||||
resolve_tool_start_thinking,
|
||||
)
|
||||
from app.tasks.chat.streaming.handlers.tools.shared.model import (
|
||||
ToolStartThinking,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"ToolCompletionEmissionContext",
|
||||
"ToolStartThinking",
|
||||
"iter_tool_completion_emission_frames",
|
||||
"resolve_tool_completed_thinking_step",
|
||||
"resolve_tool_start_thinking",
|
||||
]
|
||||
|
|
@ -0,0 +1,15 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator
|
||||
|
||||
from app.tasks.chat.streaming.handlers.tools.emission_context import (
|
||||
ToolCompletionEmissionContext,
|
||||
)
|
||||
|
||||
|
||||
def iter_completion_emission_frames(
|
||||
ctx: ToolCompletionEmissionContext,
|
||||
) -> Iterator[str]:
|
||||
out = ctx.tool_output
|
||||
payload = out if isinstance(out, dict) else {"result": out}
|
||||
yield ctx.emit_tool_output_card(payload)
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from app.tasks.chat.streaming.handlers.tools.default import (
|
||||
thinking as default_thinking,
|
||||
)
|
||||
from app.tasks.chat.streaming.handlers.tools.shared.model import (
|
||||
ToolStartThinking,
|
||||
)
|
||||
|
||||
|
||||
def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking:
|
||||
return default_thinking.resolve_start_thinking(tool_name, tool_input)
|
||||
|
||||
|
||||
def resolve_completed_thinking(
|
||||
tool_name: str, tool_output: Any, last_items: list[str],
|
||||
) -> tuple[str, list[str]]:
|
||||
return default_thinking.resolve_completed_thinking(
|
||||
tool_name, tool_output, last_items
|
||||
)
|
||||
|
|
@ -0,0 +1,31 @@
|
|||
from __future__ import annotations
|
||||
|
||||
SHARED_CONNECTOR_TOOLS: frozenset[str] = frozenset(
|
||||
{
|
||||
"create_calendar_event",
|
||||
"create_confluence_page",
|
||||
"create_dropbox_file",
|
||||
"create_gmail_draft",
|
||||
"create_google_drive_file",
|
||||
"create_jira_issue",
|
||||
"create_linear_issue",
|
||||
"create_notion_page",
|
||||
"create_onedrive_file",
|
||||
"delete_calendar_event",
|
||||
"delete_confluence_page",
|
||||
"delete_dropbox_file",
|
||||
"delete_google_drive_file",
|
||||
"delete_jira_issue",
|
||||
"delete_linear_issue",
|
||||
"delete_notion_page",
|
||||
"delete_onedrive_file",
|
||||
"send_gmail_email",
|
||||
"trash_gmail_email",
|
||||
"update_calendar_event",
|
||||
"update_confluence_page",
|
||||
"update_gmail_draft",
|
||||
"update_jira_issue",
|
||||
"update_linear_issue",
|
||||
"update_notion_page",
|
||||
}
|
||||
)
|
||||
|
|
@ -0,0 +1,3 @@
|
|||
"""Fallback tool package."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
|
@ -0,0 +1,24 @@
|
|||
"""Default tool-output card and a short completion terminal line."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterator
|
||||
|
||||
from app.tasks.chat.streaming.handlers.tools.emission_context import (
|
||||
ToolCompletionEmissionContext,
|
||||
)
|
||||
|
||||
|
||||
def iter_completion_emission_frames(
|
||||
ctx: ToolCompletionEmissionContext,
|
||||
) -> Iterator[str]:
|
||||
yield ctx.emit_tool_output_card(
|
||||
{
|
||||
"status": "completed",
|
||||
"result_length": len(str(ctx.tool_output)),
|
||||
},
|
||||
)
|
||||
yield ctx.streaming_service.format_terminal_info(
|
||||
f"Tool {ctx.tool_name} completed",
|
||||
"success",
|
||||
)
|
||||
|
|
@ -0,0 +1,23 @@
|
|||
"""Fallback thinking-step copy for unknown tools and connectors without custom UI."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from app.tasks.chat.streaming.handlers.tools.shared.model import (
|
||||
ToolStartThinking,
|
||||
)
|
||||
|
||||
|
||||
def resolve_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking:
|
||||
del tool_input
|
||||
title = tool_name.replace("_", " ").strip().capitalize() or tool_name
|
||||
return ToolStartThinking(title=title, items=[], include_items_on_frame=False)
|
||||
|
||||
|
||||
def resolve_completed_thinking(
|
||||
tool_name: str, tool_output: Any, last_items: list[str]
|
||||
) -> tuple[str, list[str]]:
|
||||
del tool_output
|
||||
title = tool_name.replace("_", " ").strip().capitalize() or tool_name
|
||||
return (title, last_items)
|
||||
|
|
@ -0,0 +1,34 @@
|
|||
"""Context for one tool-completion emission pass."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from app.tasks.chat.streaming.handlers.tool_output_frame import (
|
||||
emit_tool_output_available_frame,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ToolCompletionEmissionContext:
|
||||
"""Streaming service, tool output, and ids for completion frames."""
|
||||
|
||||
tool_name: str
|
||||
tool_call_id: str
|
||||
tool_output: Any
|
||||
streaming_service: Any
|
||||
content_builder: Any | None
|
||||
langchain_tool_call_id_holder: dict[str, str | None]
|
||||
stream_result: Any
|
||||
langgraph_config: dict[str, Any]
|
||||
staged_workspace_file_path: str | None
|
||||
|
||||
def emit_tool_output_card(self, payload: Any) -> str:
|
||||
return emit_tool_output_available_frame(
|
||||
streaming_service=self.streaming_service,
|
||||
content_builder=self.content_builder,
|
||||
langchain_id_holder=self.langchain_tool_call_id_holder,
|
||||
call_id=self.tool_call_id,
|
||||
output=payload,
|
||||
)
|
||||
|
|
@ -0,0 +1,88 @@
|
|||
"""Resolve thinking and emission modules by tool name."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
from collections.abc import Iterator
|
||||
from typing import Any
|
||||
|
||||
from app.tasks.chat.streaming.handlers.tools.connector.shared.tool_names import (
|
||||
SHARED_CONNECTOR_TOOLS,
|
||||
)
|
||||
from app.tasks.chat.streaming.handlers.tools.deliverables.tool_names import (
|
||||
DELIVERABLE_TOOLS,
|
||||
)
|
||||
from app.tasks.chat.streaming.handlers.tools.emission_context import (
|
||||
ToolCompletionEmissionContext,
|
||||
)
|
||||
from app.tasks.chat.streaming.handlers.tools.filesystem.tool_names import (
|
||||
FILESYSTEM_TOOLS,
|
||||
)
|
||||
from app.tasks.chat.streaming.handlers.tools.shared.model import (
|
||||
ToolStartThinking,
|
||||
)
|
||||
|
||||
_BASE = "app.tasks.chat.streaming.handlers.tools"
|
||||
_CONNECTOR_SHARED = "connector.shared"
|
||||
|
||||
_THINKING_ALIAS: dict[str, str] = {
|
||||
"execute_code": "filesystem.execute",
|
||||
}
|
||||
_EMISSION_ALIAS: dict[str, str] = {
|
||||
"edit_file": "filesystem.write_file",
|
||||
"execute_code": "filesystem.execute",
|
||||
}
|
||||
|
||||
|
||||
def _thinking_module(tool_name: str) -> str:
|
||||
if tool_name in SHARED_CONNECTOR_TOOLS:
|
||||
return _CONNECTOR_SHARED
|
||||
if tool_name in FILESYSTEM_TOOLS:
|
||||
return f"filesystem.{tool_name}"
|
||||
if tool_name in DELIVERABLE_TOOLS:
|
||||
return f"deliverables.{tool_name}"
|
||||
return _THINKING_ALIAS.get(tool_name, tool_name)
|
||||
|
||||
|
||||
def _emission_module(tool_name: str) -> str:
|
||||
if tool_name in _EMISSION_ALIAS:
|
||||
return _EMISSION_ALIAS[tool_name]
|
||||
if tool_name in SHARED_CONNECTOR_TOOLS:
|
||||
return _CONNECTOR_SHARED
|
||||
if tool_name in DELIVERABLE_TOOLS:
|
||||
return f"deliverables.{tool_name}"
|
||||
if tool_name in FILESYSTEM_TOOLS:
|
||||
return f"filesystem.{tool_name}"
|
||||
return tool_name
|
||||
|
||||
|
||||
def _import_thinking(tool_name: str):
|
||||
try:
|
||||
return importlib.import_module(f"{_BASE}.{_thinking_module(tool_name)}.thinking")
|
||||
except ModuleNotFoundError:
|
||||
return importlib.import_module(f"{_BASE}.default.thinking")
|
||||
|
||||
|
||||
def _import_emission(tool_name: str):
|
||||
try:
|
||||
return importlib.import_module(f"{_BASE}.{_emission_module(tool_name)}.emission")
|
||||
except ModuleNotFoundError:
|
||||
return importlib.import_module(f"{_BASE}.default.emission")
|
||||
|
||||
|
||||
def resolve_tool_start_thinking(tool_name: str, tool_input: Any) -> ToolStartThinking:
|
||||
return _import_thinking(tool_name).resolve_start_thinking(tool_name, tool_input)
|
||||
|
||||
|
||||
def resolve_tool_completed_thinking_step(
|
||||
tool_name: str, tool_output: Any, last_items: list[str]
|
||||
) -> tuple[str, list[str]]:
|
||||
return _import_thinking(tool_name).resolve_completed_thinking(
|
||||
tool_name, tool_output, last_items
|
||||
)
|
||||
|
||||
|
||||
def iter_tool_completion_emission_frames(
|
||||
ctx: ToolCompletionEmissionContext,
|
||||
) -> Iterator[str]:
|
||||
yield from _import_emission(ctx.tool_name).iter_completion_emission_frames(ctx)
|
||||
|
|
@ -0,0 +1,12 @@
|
|||
"""In-progress thinking-step title and bullet lines."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ToolStartThinking:
|
||||
title: str
|
||||
items: list[str]
|
||||
include_items_on_frame: bool = True
|
||||
Loading…
Add table
Add a link
Reference in a new issue