1609 lines
64 KiB
Python
1609 lines
64 KiB
Python
"""Chat API routes — multi-turn conversation with Agent via REST and WebSocket."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hmac
|
|
import json
|
|
import logging
|
|
from typing import Any, TYPE_CHECKING
|
|
|
|
import os
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from fastapi import (
|
|
APIRouter,
|
|
File,
|
|
HTTPException,
|
|
Request,
|
|
UploadFile,
|
|
WebSocket,
|
|
WebSocketDisconnect,
|
|
)
|
|
from fastapi.responses import FileResponse
|
|
from pydantic import BaseModel
|
|
|
|
from agentkit.chat.skill_routing import ExecutionMode
|
|
from agentkit.core.phase import default_policy, policy_from_config
|
|
from agentkit.core.protocol import CancellationToken
|
|
from agentkit.core.react import ReActEngine
|
|
from agentkit.server._fallback_chain import execute_with_fallback_chain
|
|
from agentkit.session.manager import SessionManager
|
|
from agentkit.session.models import MessageRole, SessionStatus
|
|
from agentkit.tools.advance_phase import AdvancePhaseTool
|
|
|
|
if TYPE_CHECKING:
|
|
from agentkit.llm.gateway import LLMGateway
|
|
from agentkit.server.config import ServerConfig
|
|
from agentkit.tools.base import Tool
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
router = APIRouter(prefix="/chat", tags=["chat"])
|
|
|
|
|
|
# ── Request/Response schemas ──────────────────────────────────────────
|
|
|
|
|
|
class CreateSessionRequest(BaseModel):
|
|
agent_name: str
|
|
metadata: dict[str, Any] | None = None
|
|
|
|
|
|
class SendMessageRequest(BaseModel):
|
|
content: str
|
|
role: str = "user"
|
|
# Optional execution mode override. "plan_exec" → 501 (KTD4: WebSocket only).
|
|
execution_mode: str | None = None
|
|
|
|
|
|
class SessionResponse(BaseModel):
|
|
session_id: str
|
|
agent_name: str
|
|
status: str
|
|
metadata: dict[str, Any]
|
|
created_at: str
|
|
updated_at: str
|
|
|
|
|
|
class MessageResponse(BaseModel):
|
|
message_id: str
|
|
session_id: str
|
|
role: str
|
|
content: str
|
|
tool_call_id: str | None = None
|
|
agent_name: str | None = None
|
|
created_at: str
|
|
|
|
|
|
# ── Chat WebSocket connection manager ─────────────────────────────────
|
|
|
|
|
|
class ChatConnectionManager:
|
|
"""Track active WebSocket connections per session_id."""
|
|
|
|
def __init__(self) -> None:
|
|
# session_id -> list of (websocket, pending_replies)
|
|
self._connections: dict[str, list[tuple[WebSocket, dict[str, asyncio.Future]]]] = {}
|
|
|
|
def add(self, session_id: str, ws: WebSocket, pending: dict[str, asyncio.Future]) -> None:
|
|
self._connections.setdefault(session_id, []).append((ws, pending))
|
|
|
|
def remove(self, session_id: str, ws: WebSocket) -> None:
|
|
conns = self._connections.get(session_id)
|
|
if conns is None:
|
|
return
|
|
self._connections[session_id] = [(w, p) for w, p in conns if w is not ws]
|
|
if not self._connections[session_id]:
|
|
del self._connections[session_id]
|
|
|
|
def get_connections(self, session_id: str) -> list[tuple[WebSocket, dict[str, asyncio.Future]]]:
|
|
return self._connections.get(session_id, [])
|
|
|
|
async def send_json(self, session_id: str, message: dict) -> None:
|
|
"""Send a JSON message to all connections for a session."""
|
|
conns = self._connections.get(session_id, [])
|
|
stale: list[WebSocket] = []
|
|
for ws, _ in conns:
|
|
try:
|
|
await ws.send_json(message)
|
|
except (ConnectionError, RuntimeError, asyncio.TimeoutError):
|
|
stale.append(ws)
|
|
for ws in stale:
|
|
self.remove(session_id, ws)
|
|
|
|
|
|
chat_manager = ChatConnectionManager()
|
|
|
|
|
|
# U4: Active team sessions — maps session_id to the ExpertTeam currently executing.
|
|
# When a message arrives during team execution, it is routed as an intervention
|
|
# instead of starting a new chat task. Populated by _execute_team_collab.
|
|
_active_teams: dict[str, "object"] = {}
|
|
|
|
|
|
def _register_active_team(session_id: str, team: "object") -> None:
|
|
"""Register an active team for a session (intervention routing)."""
|
|
_active_teams[session_id] = team
|
|
|
|
|
|
def _unregister_active_team(session_id: str) -> None:
|
|
"""Unregister the active team for a session."""
|
|
_active_teams.pop(session_id, None)
|
|
|
|
|
|
def _get_active_team(session_id: str) -> "object | None":
|
|
"""Get the active team for a session, if any."""
|
|
return _active_teams.get(session_id)
|
|
|
|
|
|
# ── Helper ────────────────────────────────────────────────────────────
|
|
|
|
|
|
_VALID_TEAM_EVENT_TYPES = frozenset(
|
|
{
|
|
"team_formed",
|
|
"expert_step",
|
|
"expert_result",
|
|
"expert_result_chunk",
|
|
"expert_result_chunk_reset",
|
|
"plan_update",
|
|
"team_synthesis",
|
|
"team_synthesis_chunk",
|
|
"team_dissolved",
|
|
"plan_step",
|
|
"phase_started",
|
|
"phase_completed",
|
|
"phase_failed",
|
|
"replanning",
|
|
# PM Collaboration 模式事件 (U1-U4)
|
|
"collaboration_contract_defined",
|
|
"collaboration_notice",
|
|
"review_result",
|
|
"risk_flagged",
|
|
# Board Meeting 模式事件
|
|
"board_started",
|
|
"expert_speech",
|
|
"expert_speech_chunk",
|
|
"round_summary",
|
|
"user_intervention",
|
|
"board_concluded",
|
|
}
|
|
)
|
|
|
|
|
|
async def emit_team_event(websocket: WebSocket, event_type: str, data: dict) -> None:
|
|
"""Emit a team-related WebSocket event.
|
|
|
|
Supported event types:
|
|
team_formed — Team assembled with expert list and plan.
|
|
data: {team_id, status, experts, plan_phases, lead_expert}
|
|
expert_step — An expert is executing a step.
|
|
data: {expert_id, expert_name, expert_color, content, step}
|
|
expert_result — An expert produced a result.
|
|
data: {expert_id, expert_name, expert_color, content}
|
|
plan_update — Team plan phases updated.
|
|
data: {plan_phases}
|
|
team_synthesis — Final synthesis from the lead expert.
|
|
data: {content}
|
|
team_dissolved — Team dissolved after completion.
|
|
data: {team_id}
|
|
"""
|
|
if event_type not in _VALID_TEAM_EVENT_TYPES:
|
|
logger.warning(f"emit_team_event: invalid event_type '{event_type}'")
|
|
return
|
|
await websocket.send_json(
|
|
{
|
|
"type": event_type,
|
|
"data": data,
|
|
}
|
|
)
|
|
|
|
|
|
def _get_session_manager(request: Request) -> SessionManager:
|
|
return request.app.state.session_manager
|
|
|
|
|
|
async def _execute_board_meeting(
|
|
websocket: WebSocket,
|
|
session_id: str,
|
|
content: str,
|
|
sm: SessionManager,
|
|
) -> bool:
|
|
"""Intercept @board prefix and execute a board meeting discussion.
|
|
|
|
Returns True if the input was handled as a board meeting (caller should return),
|
|
False if the input should continue through the normal chat pipeline.
|
|
|
|
Flow:
|
|
1. Resolve @board routing via BoardRouter
|
|
2. Create BoardTeam with expert configs
|
|
3. Register handoff_transport handler to relay events to WebSocket
|
|
4. Execute BoardOrchestrator
|
|
5. Send final conclusion as final_answer
|
|
6. Persist user topic + final summary to session history
|
|
"""
|
|
from agentkit.experts.board_router import BoardRouter
|
|
from agentkit.experts.board import BoardTeam
|
|
from agentkit.experts.board_orchestrator import BoardOrchestrator
|
|
|
|
app_state = websocket.app.state
|
|
|
|
# Resolve ExpertTemplateRegistry from app.state (loaded at startup)
|
|
template_registry = getattr(app_state, "expert_template_registry", None)
|
|
if template_registry is None:
|
|
from agentkit.experts.registry import ExpertTemplateRegistry
|
|
|
|
template_registry = ExpertTemplateRegistry()
|
|
|
|
board_router = BoardRouter(template_registry=template_registry)
|
|
routing_result = board_router.resolve(content)
|
|
|
|
if not routing_result.matched:
|
|
return False # Not a @board input, continue normal pipeline
|
|
|
|
if not routing_result.topic:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "error",
|
|
"data": {"message": "私董会需要一个讨论主题，例如：@board 如何看待 AI 未来"},
|
|
}
|
|
)
|
|
return True
|
|
|
|
# Resolve expert configs from specified experts or default template
|
|
expert_configs = board_router.resolve_expert_configs(routing_result.specified_experts)
|
|
if not expert_configs:
|
|
await websocket.send_json(
|
|
{"type": "error", "data": {"message": "无法解析私董会成员，请检查专家名称或模板配置"}}
|
|
)
|
|
return True
|
|
|
|
# Read board config from server_config if available; user-specified rounds take precedence
|
|
max_rounds = routing_result.max_rounds
|
|
if max_rounds is None:
|
|
max_rounds = 5
|
|
server_config = getattr(app_state, "server_config", None)
|
|
if server_config is not None:
|
|
board_cfg = getattr(server_config, "board", None) or {}
|
|
if isinstance(board_cfg, dict):
|
|
max_rounds = int(board_cfg.get("max_rounds", 5))
|
|
|
|
# Create BoardTeam
|
|
team = BoardTeam(
|
|
pool=app_state.agent_pool,
|
|
template_registry=template_registry,
|
|
max_rounds=max_rounds,
|
|
)
|
|
|
|
# Register handoff_transport handler to relay board events to WebSocket
|
|
async def _relay_board_event(message: dict) -> None:
|
|
msg_type = message.get("type")
|
|
if not msg_type:
|
|
return
|
|
# Strip internal fields, keep only event data
|
|
event_data = {k: v for k, v in message.items() if k != "type"}
|
|
await emit_team_event(websocket, msg_type, event_data)
|
|
# Persist board events so a page reload can reconstruct the
|
|
# discussion (otherwise the user only sees the final conclusion
|
|
# and loses every expert speech). We persist on the terminal
|
|
# "board_started" / "expert_speech" / "round_summary" /
|
|
# "board_concluded" events — chunks are intermediate and would
|
|
# explode the history size. The "board_started" payload carries
|
|
# the expert list (name/avatar/color/is_moderator/persona) so a
|
|
# reload can rebuild the boardState and the sidebar can show the
|
|
# "私董会" badge — without it, every restored conversation
|
|
# appears as a plain chat with 0 experts.
|
|
#
|
|
# We also stash the rendering hint (message_type + expert identity)
|
|
# in Message.metadata so a reload can still show the speech as a
|
|
# proper board_speech card instead of a plain assistant bubble.
|
|
experts_data = event_data.get("experts")
|
|
board_started_text = (
|
|
f"私董会开始：{event_data.get('topic', '')}"
|
|
if event_data.get("topic")
|
|
else "私董会开始"
|
|
)
|
|
persistable: dict[str, tuple[str, str, dict[str, object] | None]] = {
|
|
"board_started": (
|
|
"assistant",
|
|
board_started_text,
|
|
{
|
|
"message_type": "board_started",
|
|
"board_started": {
|
|
"team_id": event_data.get("team_id"),
|
|
"topic": event_data.get("topic"),
|
|
"experts": experts_data,
|
|
"max_rounds": event_data.get("max_rounds"),
|
|
},
|
|
},
|
|
),
|
|
"expert_speech": (
|
|
"assistant",
|
|
event_data.get("content", ""),
|
|
{
|
|
"message_type": "board_speech",
|
|
"expert_name": event_data.get("expert_name"),
|
|
"expert_avatar": event_data.get("expert_avatar"),
|
|
"expert_color": event_data.get("expert_color"),
|
|
"board_round": event_data.get("round"),
|
|
"board_role": event_data.get("role"),
|
|
},
|
|
),
|
|
"round_summary": (
|
|
"assistant",
|
|
event_data.get("content", ""),
|
|
{
|
|
"message_type": "board_summary",
|
|
"expert_name": event_data.get("moderator_name"),
|
|
"expert_avatar": event_data.get("moderator_avatar"),
|
|
"expert_color": event_data.get("moderator_color"),
|
|
"board_round": event_data.get("round"),
|
|
"board_role": "summary",
|
|
},
|
|
),
|
|
"board_concluded": (
|
|
"assistant",
|
|
event_data.get("summary") or "私董会已结束",
|
|
{
|
|
"message_type": "board_conclusion",
|
|
"board_round": event_data.get("total_rounds"),
|
|
"board_conclusion": {
|
|
"summary": event_data.get("summary", ""),
|
|
"decision_advice": event_data.get("decision_advice", ""),
|
|
"total_rounds": event_data.get("total_rounds", 0),
|
|
"consensus_points": event_data.get("consensus_points", []),
|
|
"dissent_points": event_data.get("dissent_points", []),
|
|
},
|
|
},
|
|
),
|
|
}
|
|
if msg_type in persistable:
|
|
role, content, meta = persistable[msg_type]
|
|
if content:
|
|
try:
|
|
await sm.append_message(
|
|
session_id=session_id,
|
|
role=MessageRole(role),
|
|
content=content,
|
|
metadata=meta,
|
|
)
|
|
except Exception as persist_err: # noqa: BLE001
|
|
# Persistence is best-effort; never let a save error
|
|
# break the live stream.
|
|
logger.warning(f"Failed to persist {msg_type} to session store: {persist_err}")
|
|
|
|
team.handoff_transport.register_handler(team.team_channel, _relay_board_event)
|
|
|
|
# Append user topic to session history — store the resolved topic, not
|
|
# the raw "@board:experts rounds=5 ..." syntax. Frontend renders the
|
|
# original prefix as a structured bubble; persisting it raw means the
|
|
# bubble shows the full prefix even on history reload.
|
|
await sm.append_message(
|
|
session_id=session_id,
|
|
role=MessageRole.USER,
|
|
content=f"@board {routing_result.topic}",
|
|
)
|
|
|
|
try:
|
|
await team.create_board(topic=routing_result.topic, expert_configs=expert_configs)
|
|
orchestrator = BoardOrchestrator(team=team)
|
|
result = await orchestrator.execute(routing_result.topic)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Board meeting failed for session {session_id}: {e}", exc_info=True)
|
|
await websocket.send_json(
|
|
{"type": "error", "data": {"message": f"私董会执行失败: {str(e)[:200]}"}}
|
|
)
|
|
try:
|
|
await team.dissolve()
|
|
except (RuntimeError, asyncio.TimeoutError, ConnectionError):
|
|
pass
|
|
return True
|
|
finally:
|
|
# dissolve() already clears handlers via handoff_transport.close()
|
|
pass
|
|
|
|
# Build final answer text from conclusion
|
|
summary = result.get("summary", "")
|
|
decision_advice = result.get("decision_advice", "")
|
|
consensus_points = result.get("consensus_points", []) or []
|
|
dissent_points = result.get("dissent_points", []) or []
|
|
total_rounds = result.get("total_rounds", 0)
|
|
|
|
final_parts: list[str] = []
|
|
if summary:
|
|
final_parts.append(f"## 讨论总结\n\n{summary}")
|
|
if decision_advice:
|
|
final_parts.append(f"## 决策建议\n\n{decision_advice}")
|
|
if consensus_points:
|
|
final_parts.append("## 共识点\n\n" + "\n".join(f"- {p}" for p in consensus_points))
|
|
if dissent_points:
|
|
final_parts.append("## 分歧点\n\n" + "\n".join(f"- {p}" for p in dissent_points))
|
|
final_parts.append(f"\n\n_共进行 {total_rounds} 轮讨论_")
|
|
|
|
final_content = "\n\n".join(final_parts)
|
|
|
|
await websocket.send_json(
|
|
{
|
|
"type": "final_answer",
|
|
"content": final_content,
|
|
"is_final": True,
|
|
}
|
|
)
|
|
|
|
# Persist final summary as assistant message
|
|
await sm.append_message(
|
|
session_id=session_id,
|
|
role=MessageRole.ASSISTANT,
|
|
content=final_content,
|
|
agent_name="board_meeting",
|
|
)
|
|
|
|
# Dissolve the team to release expert agents
|
|
try:
|
|
await team.dissolve()
|
|
except (RuntimeError, asyncio.TimeoutError, ConnectionError) as e:
|
|
logger.warning(f"Board team dissolve failed: {e}")
|
|
|
|
return True
|
|
|
|
|
|
async def _execute_team_collab(
|
|
websocket: WebSocket,
|
|
session_id: str,
|
|
content: str,
|
|
sm: SessionManager,
|
|
) -> bool:
|
|
"""Intercept @team prefix and execute a pipeline team collaboration.
|
|
|
|
Returns True if the input was handled as a team collaboration (caller should return),
|
|
False if the input should continue through the normal chat pipeline.
|
|
|
|
Flow:
|
|
1. Resolve @team routing via ExpertTeamRouter
|
|
2. Create ExpertTeam with lead + member configs
|
|
3. Register handoff_transport handler to relay events to WebSocket
|
|
4. Execute TeamOrchestrator (pipeline mode)
|
|
5. Send final synthesis as final_answer
|
|
6. Persist user task + final result to session history
|
|
"""
|
|
from agentkit.experts.router import ExpertTeamRouter
|
|
from agentkit.experts.team import ExpertTeam
|
|
from agentkit.experts.orchestrator import TeamOrchestrator
|
|
|
|
app_state = websocket.app.state
|
|
|
|
# Resolve ExpertTemplateRegistry from app.state (loaded at startup)
|
|
template_registry = getattr(app_state, "expert_template_registry", None)
|
|
if template_registry is None:
|
|
from agentkit.experts.registry import ExpertTemplateRegistry
|
|
|
|
template_registry = ExpertTemplateRegistry()
|
|
|
|
team_router = ExpertTeamRouter(template_registry=template_registry)
|
|
routing_result = team_router.resolve(content)
|
|
|
|
if not routing_result.matched:
|
|
return False # Not a @team input, continue normal pipeline
|
|
|
|
if not routing_result.task_content:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "error",
|
|
"data": {"message": "团队任务需要一个描述，例如：@team 开发用户登录功能"},
|
|
}
|
|
)
|
|
return True
|
|
|
|
# Resolve expert configs from specified experts or default dev_team template
|
|
expert_configs = team_router.resolve_expert_configs(routing_result.specified_experts)
|
|
if not expert_configs:
|
|
await websocket.send_json(
|
|
{"type": "error", "data": {"message": "无法解析团队成员，请检查专家名称或模板配置"}}
|
|
)
|
|
return True
|
|
|
|
# Split configs: first is lead, rest are members (V2 verification)
|
|
lead_config = expert_configs[0]
|
|
member_configs = expert_configs[1:] if len(expert_configs) > 1 else []
|
|
|
|
# Create ExpertTeam
|
|
team = ExpertTeam(
|
|
pool=app_state.agent_pool,
|
|
template_registry=template_registry,
|
|
redis_client=getattr(app_state, "working_redis_client", None),
|
|
)
|
|
|
|
# Register handoff_transport handler to relay team events to WebSocket
|
|
async def _relay_team_event(message: dict) -> None:
|
|
msg_type = message.get("type")
|
|
if not msg_type:
|
|
return
|
|
# Strip internal fields, keep only event data
|
|
event_data = {k: v for k, v in message.items() if k != "type"}
|
|
await emit_team_event(websocket, msg_type, event_data)
|
|
|
|
team.handoff_transport.register_handler(team.team_channel, _relay_team_event)
|
|
|
|
try:
|
|
# Append user task to session history (inside try so we can compensate on error)
|
|
await sm.append_message(
|
|
session_id=session_id,
|
|
role=MessageRole.USER,
|
|
content=content,
|
|
)
|
|
|
|
await team.create_team(lead_config=lead_config, member_configs=member_configs)
|
|
# U7: Create checkpoint manager for crash recovery
|
|
from agentkit.orchestrator.checkpoint import PipelineCheckpoint
|
|
|
|
checkpoint = PipelineCheckpoint(
|
|
redis_client=getattr(app_state, "working_redis_client", None)
|
|
)
|
|
orchestrator = TeamOrchestrator(team=team, checkpoint=checkpoint)
|
|
# U4: Register active team so WS messages during execution route as interventions
|
|
_register_active_team(session_id, team)
|
|
result = await orchestrator.execute(routing_result.task_content)
|
|
except asyncio.CancelledError:
|
|
logger.info(f"Team collaboration cancelled for session {session_id}")
|
|
await websocket.send_json({"type": "error", "data": {"message": "团队协作已取消"}})
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Team collaboration failed for session {session_id}: {e}", exc_info=True)
|
|
await websocket.send_json(
|
|
{"type": "error", "data": {"message": f"团队协作执行失败: {str(e)[:200]}"}}
|
|
)
|
|
return True
|
|
finally:
|
|
# U4: Always unregister the active team first so subsequent messages
|
|
# don't route to a dissolving team.
|
|
_unregister_active_team(session_id)
|
|
# Always dissolve the team and remove handler to avoid leaks
|
|
try:
|
|
await team.dissolve()
|
|
except (RuntimeError, asyncio.TimeoutError, ConnectionError) as e:
|
|
logger.warning(f"Team dissolve failed: {e}")
|
|
# dissolve() already clears handlers via handoff_transport.close()
|
|
|
|
# Build final answer text from synthesis result
|
|
final_result = result.get("result") or {}
|
|
final_content = (
|
|
final_result.get("content", "") if isinstance(final_result, dict) else str(final_result)
|
|
)
|
|
|
|
if not final_content:
|
|
# Fallback: use phase results if synthesis is empty
|
|
phase_results = result.get("phase_results") or {}
|
|
if phase_results:
|
|
parts = []
|
|
# Build a phase_id -> phase_name lookup from the plan
|
|
phase_names = {}
|
|
plan_obj = result.get("plan")
|
|
if plan_obj:
|
|
for ph in plan_obj.phases:
|
|
phase_names[ph.id] = ph.name
|
|
for phase_id, pr in phase_results.items():
|
|
if isinstance(pr, dict) and "content" in pr:
|
|
parts.append(
|
|
f"### {phase_names.get(phase_id, pr.get('phase_name', phase_id))}\n\n{pr['content']}"
|
|
)
|
|
final_content = "\n\n".join(parts) if parts else "团队执行完成，但未生成最终结果。"
|
|
else:
|
|
final_content = "团队执行完成，但未生成最终结果。"
|
|
|
|
await websocket.send_json(
|
|
{
|
|
"type": "final_answer",
|
|
"content": final_content,
|
|
"is_final": True,
|
|
}
|
|
)
|
|
|
|
# Persist final synthesis as assistant message
|
|
await sm.append_message(
|
|
session_id=session_id,
|
|
role=MessageRole.ASSISTANT,
|
|
content=final_content,
|
|
agent_name="team_collab",
|
|
)
|
|
|
|
return True
|
|
|
|
|
|
def _session_to_response(session) -> SessionResponse:
|
|
return SessionResponse(
|
|
session_id=session.session_id,
|
|
agent_name=session.agent_name,
|
|
status=session.status.value,
|
|
metadata=session.metadata,
|
|
created_at=session.created_at.isoformat(),
|
|
updated_at=session.updated_at.isoformat(),
|
|
)
|
|
|
|
|
|
def _message_to_response(msg) -> MessageResponse:
|
|
return MessageResponse(
|
|
message_id=msg.message_id,
|
|
session_id=msg.session_id,
|
|
role=msg.role.value,
|
|
content=msg.content,
|
|
tool_call_id=msg.tool_call_id,
|
|
agent_name=msg.agent_name,
|
|
created_at=msg.created_at.isoformat(),
|
|
)
|
|
|
|
|
|
def _build_phase_engine(
|
|
*,
|
|
server_config: ServerConfig | None,
|
|
llm_gateway: LLMGateway,
|
|
execution_mode: ExecutionMode,
|
|
base_tools: list[Tool],
|
|
session_id: str = "",
|
|
) -> tuple[ReActEngine | None, list[Tool] | None, str | None]:
|
|
"""Build a PLAN_EXEC engine with PhasePolicy + AdvancePhaseTool.
|
|
|
|
Encapsulates the WS path's phase_policy construction so the REST path
|
|
can reuse it without duplicating config-lookup + policy_from_config +
|
|
AdvancePhaseTool registration. KTD5: PLAN_EXEC bypasses the fallback
|
|
chain — callers must NOT route the returned engine through
|
|
``execute_with_fallback_chain``.
|
|
|
|
Args:
|
|
server_config: ``app.state.server_config`` (or None for tests).
|
|
llm_gateway: ``app.state.llm_gateway``.
|
|
execution_mode: routing.execution_mode (WS) or PLAN_EXEC (REST).
|
|
base_tools: routing.tools (WS) or agent tool list (REST).
|
|
session_id: included in log lines for traceability only.
|
|
|
|
Returns ``(engine, tools_with_advance_phase, error_message)``:
|
|
- execution_mode != PLAN_EXEC → ``(None, None, None)`` (fall back to REACT).
|
|
- plan_exec.enabled=False → ``(None, None, None)`` (fall back to REACT).
|
|
- phase policy construction failed → ``(None, None, error_message)``.
|
|
- PLAN_EXEC engaged → ``(engine, tools_with_advance_phase, None)``.
|
|
"""
|
|
if execution_mode != ExecutionMode.PLAN_EXEC:
|
|
return (None, None, None)
|
|
|
|
plan_exec_cfg = getattr(server_config, "plan_exec", None) or {}
|
|
if plan_exec_cfg.get("enabled", True) is False:
|
|
logger.info(
|
|
"PLAN_EXEC disabled by config (plan_exec.enabled=False), "
|
|
"falling back to REACT for session %s",
|
|
session_id,
|
|
)
|
|
return (None, None, None)
|
|
|
|
try:
|
|
phase_policy = policy_from_config(plan_exec_cfg)
|
|
if phase_policy is None:
|
|
# Empty config (no `plan_exec:` section) → use KTD5 defaults.
|
|
phase_policy = default_policy()
|
|
except (ValueError, TypeError, KeyError) as e:
|
|
logger.error(
|
|
"PLAN_EXEC phase policy construction failed for session %s: %s",
|
|
session_id,
|
|
e,
|
|
)
|
|
return (None, None, f"phase policy error: {str(e)[:200]}")
|
|
|
|
engine = ReActEngine(
|
|
llm_gateway=llm_gateway,
|
|
phase_policy=phase_policy,
|
|
)
|
|
advance_phase_tool = AdvancePhaseTool(engine=engine)
|
|
tools_with_advance_phase = list(base_tools) + [advance_phase_tool]
|
|
return (engine, tools_with_advance_phase, None)
|
|
|
|
|
|
# ── REST endpoints ────────────────────────────────────────────────────
|
|
|
|
|
|
@router.get("/sessions", response_model=list[SessionResponse])
|
|
async def list_sessions(req: Request):
|
|
"""List all chat sessions."""
|
|
sm = _get_session_manager(req)
|
|
sessions = await sm.list_sessions()
|
|
return [_session_to_response(s) for s in sessions]
|
|
|
|
|
|
@router.post("/sessions", response_model=SessionResponse)
|
|
async def create_session(request: CreateSessionRequest, req: Request):
|
|
"""Create a new chat session bound to an Agent."""
|
|
sm = _get_session_manager(req)
|
|
session = await sm.create_session(
|
|
agent_name=request.agent_name,
|
|
metadata=request.metadata,
|
|
)
|
|
return _session_to_response(session)
|
|
|
|
|
|
@router.get("/sessions/{session_id}", response_model=SessionResponse)
|
|
async def get_session(session_id: str, req: Request):
|
|
"""Get session information."""
|
|
sm = _get_session_manager(req)
|
|
session = await sm.get_session(session_id)
|
|
if session is None:
|
|
raise HTTPException(status_code=404, detail=f"Session '{session_id}' not found")
|
|
return _session_to_response(session)
|
|
|
|
|
|
@router.get("/sessions/{session_id}/messages", response_model=list[MessageResponse])
|
|
async def get_messages(session_id: str, req: Request, limit: int | None = None, offset: int = 0):
|
|
"""Get conversation history for a session."""
|
|
sm = _get_session_manager(req)
|
|
session = await sm.get_session(session_id)
|
|
if session is None:
|
|
raise HTTPException(status_code=404, detail=f"Session '{session_id}' not found")
|
|
messages = await sm.get_messages(session_id, limit=limit, offset=offset)
|
|
return [_message_to_response(m) for m in messages]
|
|
|
|
|
|
@router.post("/sessions/{session_id}/messages", response_model=MessageResponse)
|
|
async def send_message(session_id: str, request: SendMessageRequest, req: Request):
|
|
"""Send a message to the Agent (synchronous mode — waits for full reply)."""
|
|
sm = _get_session_manager(req)
|
|
session = await sm.get_session(session_id)
|
|
if session is None:
|
|
raise HTTPException(status_code=404, detail=f"Session '{session_id}' not found")
|
|
if session.status == SessionStatus.CLOSED:
|
|
raise HTTPException(status_code=400, detail=f"Session '{session_id}' is closed")
|
|
|
|
# U3: PLAN_EXEC via REST — non-streaming, bypasses the fallback chain
|
|
# (KTD5: PLAN_EXEC and execute_with_fallback_chain are mutually exclusive).
|
|
# When plan_exec is disabled by config, falls through to the REACT path below.
|
|
if request.execution_mode == "plan_exec":
|
|
# Resolve the Agent early — PLAN_EXEC needs its tool list + system prompt.
|
|
pool = req.app.state.agent_pool
|
|
agent = pool.get_agent(session.agent_name)
|
|
if agent is None:
|
|
raise HTTPException(status_code=404, detail=f"Agent '{session.agent_name}' not found")
|
|
|
|
plan_exec_engine, plan_exec_tools, plan_exec_error = _build_phase_engine(
|
|
server_config=getattr(req.app.state, "server_config", None),
|
|
llm_gateway=req.app.state.llm_gateway,
|
|
execution_mode=ExecutionMode.PLAN_EXEC,
|
|
base_tools=agent._tool_registry.list_tools() if agent._tool_registry else [],
|
|
session_id=session_id,
|
|
)
|
|
if plan_exec_error is not None:
|
|
raise HTTPException(status_code=500, detail=plan_exec_error)
|
|
if plan_exec_engine is not None:
|
|
# PLAN_EXEC engaged — append user msg, execute non-streaming, return.
|
|
await sm.append_message(
|
|
session_id=session_id,
|
|
role=MessageRole.USER,
|
|
content=request.content,
|
|
)
|
|
chat_messages = await sm.get_chat_messages(session_id)
|
|
system_prompt = getattr(agent, "_system_prompt", None) or (
|
|
agent.get_system_prompt() if hasattr(agent, "get_system_prompt") else None
|
|
)
|
|
try:
|
|
plan_exec_result = await plan_exec_engine.execute(
|
|
messages=chat_messages,
|
|
tools=plan_exec_tools,
|
|
model=agent.get_model()
|
|
if hasattr(agent, "get_model")
|
|
else getattr(agent, "_llm_model", "default"),
|
|
agent_name=agent.name,
|
|
system_prompt=system_prompt,
|
|
)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"PLAN_EXEC execution error for session {session_id}: {e}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
assistant_msg = await sm.append_message(
|
|
session_id=session_id,
|
|
role=MessageRole.ASSISTANT,
|
|
content=plan_exec_result.output,
|
|
agent_name=agent.name,
|
|
)
|
|
return _message_to_response(assistant_msg)
|
|
# else: plan_exec.enabled=False → fall through to REACT path below.
|
|
|
|
# Append user message
|
|
await sm.append_message(
|
|
session_id=session_id,
|
|
role=MessageRole.USER,
|
|
content=request.content,
|
|
)
|
|
|
|
# Get full conversation history for the Agent
|
|
chat_messages = await sm.get_chat_messages(session_id)
|
|
|
|
# Resolve the Agent
|
|
pool = req.app.state.agent_pool
|
|
agent = pool.get_agent(session.agent_name)
|
|
if agent is None:
|
|
raise HTTPException(status_code=404, detail=f"Agent '{session.agent_name}' not found")
|
|
|
|
# Execute the Agent
|
|
try:
|
|
# Reuse Agent's ReActEngine if available (U2: Chat pipeline optimization)
|
|
react_engine = getattr(agent, "_react_engine", None)
|
|
if react_engine is None:
|
|
react_engine = ReActEngine(llm_gateway=req.app.state.llm_gateway)
|
|
else:
|
|
react_engine.reset()
|
|
tools = agent._tool_registry.list_tools() if agent._tool_registry else []
|
|
system_prompt = getattr(agent, "_system_prompt", None) or (
|
|
agent.get_system_prompt() if hasattr(agent, "get_system_prompt") else None
|
|
)
|
|
# G7/U3: Three-tier fallback chain (main → Recovery → Emergency).
|
|
# Wired only here (KTD5); CLI / ReWOO / Reflexion internal ReAct bypass.
|
|
server_config = getattr(req.app.state, "server_config", None)
|
|
fallback_chain_cfg = (
|
|
getattr(server_config, "fallback_chain", None) if server_config else None
|
|
)
|
|
chat_result = await execute_with_fallback_chain(
|
|
react_engine=react_engine,
|
|
llm_gateway=req.app.state.llm_gateway,
|
|
messages=chat_messages,
|
|
tools=tools,
|
|
model=agent.get_model()
|
|
if hasattr(agent, "get_model")
|
|
else getattr(agent, "_llm_model", "default"),
|
|
agent_name=agent.name,
|
|
system_prompt=system_prompt,
|
|
fallback_chain_config=fallback_chain_cfg,
|
|
)
|
|
|
|
# Append assistant reply
|
|
assistant_msg = await sm.append_message(
|
|
session_id=session_id,
|
|
role=MessageRole.ASSISTANT,
|
|
content=chat_result.output,
|
|
agent_name=agent.name,
|
|
)
|
|
response = _message_to_response(assistant_msg)
|
|
# Attach structured error payload when Emergency tier fired.
|
|
if chat_result.error_struct is not None:
|
|
response_dict = (
|
|
response.model_dump() if hasattr(response, "model_dump") else dict(response)
|
|
)
|
|
response_dict["error_struct"] = chat_result.error_struct
|
|
response_dict["fallback_status"] = chat_result.status
|
|
return response_dict
|
|
return response
|
|
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Chat execution error for session {session_id}: {e}")
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.delete("/sessions/{session_id}")
|
|
async def close_session(session_id: str, req: Request):
|
|
"""Close a chat session."""
|
|
sm = _get_session_manager(req)
|
|
session = await sm.close_session(session_id)
|
|
if session is None:
|
|
raise HTTPException(status_code=404, detail=f"Session '{session_id}' not found")
|
|
return {"status": "closed", "session_id": session_id}
|
|
|
|
|
|
# ── WebSocket endpoint ────────────────────────────────────────────────
|
|
|
|
|
|
async def _resolve_ws_dept_context(
|
|
websocket: WebSocket,
|
|
) -> tuple[str | None, list[str], Path | None]:
|
|
"""Resolve user_id, department_ids, and db_path for quota enforcement.
|
|
|
|
Reads ``websocket.state.current_user`` (set by :class:`AuthMiddleware`
|
|
when JWT auth via ``?token=`` query param succeeds). If no user context
|
|
is available (API key auth), returns ``(None, [], None)`` — quota is
|
|
not enforced for API key clients.
|
|
|
|
Returns ``(user_id, department_ids, db_path)``:
|
|
- ``user_id``: the authenticated user's id, or ``None``.
|
|
- ``department_ids``: the user's active department ids (empty for
|
|
admins and API key clients).
|
|
- ``db_path``: the auth DB path (needed for quota checks), or ``None``
|
|
if no user context is available.
|
|
"""
|
|
current_user: dict[str, Any] | None = getattr(websocket.state, "current_user", None)
|
|
if current_user is None:
|
|
return None, [], None
|
|
|
|
user_id = current_user.get("user_id")
|
|
role = current_user.get("role")
|
|
|
|
# Admins and API-key clients (user_id=None) — no quota enforcement.
|
|
if role == "admin" or not user_id:
|
|
return user_id, [], None
|
|
|
|
# Regular user: look up their active department ids.
|
|
from agentkit.server.admin.context import _fetch_user_department_ids
|
|
|
|
db_path = getattr(websocket.app.state, "auth_db_path", None)
|
|
if db_path is None:
|
|
db_path_resolved = None
|
|
else:
|
|
db_path_resolved = Path(db_path)
|
|
try:
|
|
department_ids = await _fetch_user_department_ids(db_path_resolved, user_id)
|
|
except (ConnectionError, OSError, asyncio.TimeoutError, ValueError, KeyError, RuntimeError):
|
|
logger.exception(
|
|
"Failed to fetch department ids for WebSocket user %s — fail-closed",
|
|
user_id,
|
|
)
|
|
# Fail-closed: return empty list with db_path=None so gateway
|
|
# skips quota enforcement but we log the failure. The user
|
|
# gets a degraded experience rather than a security bypass.
|
|
return user_id, [], None
|
|
return user_id, department_ids, db_path_resolved
|
|
|
|
return user_id, [], None
|
|
|
|
|
|
@router.websocket("/ws/{session_id}")
|
|
async def chat_websocket(websocket: WebSocket, session_id: str) -> None:
|
|
"""WebSocket endpoint for real-time chat with streaming.
|
|
|
|
Client → Server messages:
|
|
{"type": "message", "content": "..."} — Send a user message
|
|
{"type": "cancel"} — Cancel current execution
|
|
{"type": "ping"} — Heartbeat
|
|
|
|
Server → Client messages:
|
|
{"type": "connected", "session_id": "..."} — Connection confirmed
|
|
{"type": "token", "content": "..."} — LLM token streaming
|
|
{"type": "step", "data": {...}} — ReAct step event
|
|
{"type": "ask_human", "question": "...", "request_id": "..."} — Agent asks user
|
|
{"type": "final_answer", "content": "..."} — Agent's final reply
|
|
{"type": "error", "data": {"message": "..."}} — Error occurred
|
|
{"type": "pong"} — Heartbeat response
|
|
|
|
Expert Team events (Server → Client):
|
|
{"type": "team_formed", "data": {...}} — Team assembled
|
|
{"type": "expert_step", "data": {...}} — Expert executing step
|
|
{"type": "expert_result", "data": {...}} — Expert produced result
|
|
{"type": "plan_update", "data": {...}} — Plan phases updated
|
|
{"type": "team_synthesis", "data": {...}} — Final synthesis
|
|
{"type": "team_dissolved", "data": {...}} — Team dissolved
|
|
"""
|
|
# Authentication
|
|
configured_api_key: str | None = None
|
|
if hasattr(websocket.app.state, "server_config") and websocket.app.state.server_config:
|
|
configured_api_key = websocket.app.state.server_config.api_key
|
|
if configured_api_key is None and hasattr(websocket.app.state, "api_key"):
|
|
configured_api_key = websocket.app.state.api_key
|
|
|
|
if configured_api_key:
|
|
provided = websocket.query_params.get("api_key")
|
|
if not provided or not hmac.compare_digest(provided, configured_api_key):
|
|
await websocket.accept()
|
|
await websocket.send_json({"type": "error", "data": {"message": "Invalid api_key"}})
|
|
await websocket.close(code=4001, reason="Invalid api_key")
|
|
return
|
|
|
|
await websocket.accept()
|
|
|
|
# Validate session
|
|
sm: SessionManager = websocket.app.state.session_manager
|
|
session = await sm.get_session(session_id)
|
|
if session is None:
|
|
await websocket.send_json(
|
|
{"type": "error", "data": {"message": f"Session '{session_id}' not found"}}
|
|
)
|
|
await websocket.close(code=1000, reason="Session not found")
|
|
return
|
|
|
|
if session.status == SessionStatus.CLOSED:
|
|
await websocket.send_json({"type": "error", "data": {"message": "Session is closed"}})
|
|
await websocket.close(code=1000, reason="Session closed")
|
|
return
|
|
|
|
# Track pending replies for AskHumanTool and confirmations
|
|
pending_replies: dict[str, asyncio.Future] = {}
|
|
pending_confirmations: dict[str, asyncio.Future] = {}
|
|
chat_manager.add(session_id, websocket, pending_replies)
|
|
|
|
cancellation_token = CancellationToken()
|
|
|
|
# Per-session concurrency guard to prevent unlimited task creation (DoS mitigation)
|
|
_MAX_CONCURRENT_TASKS = 4
|
|
active_tasks: set[asyncio.Task] = set()
|
|
|
|
try:
|
|
await websocket.send_json({"type": "connected", "session_id": session_id})
|
|
|
|
# Listen for client messages
|
|
while True:
|
|
try:
|
|
raw = await asyncio.wait_for(websocket.receive_text(), timeout=300.0)
|
|
except asyncio.TimeoutError:
|
|
await websocket.send_json({"type": "pong"})
|
|
continue
|
|
|
|
try:
|
|
msg = json.loads(raw)
|
|
except json.JSONDecodeError:
|
|
continue
|
|
|
|
msg_type = msg.get("type")
|
|
|
|
if msg_type == "message":
|
|
content = msg.get("content", "")
|
|
model = msg.get("model") # Optional model override from frontend
|
|
|
|
# U4: If a team is currently executing for this session, route
|
|
# the message as an intervention instead of a new chat task.
|
|
active_team = _get_active_team(session_id)
|
|
if active_team is not None:
|
|
try:
|
|
await active_team.add_user_intervention(content)
|
|
await websocket.send_json(
|
|
{
|
|
"type": "team_intervention_ack",
|
|
"data": {"content": content},
|
|
}
|
|
)
|
|
except (asyncio.QueueFull, RuntimeError, ConnectionError) as e:
|
|
logger.warning(f"Failed to enqueue intervention: {e}")
|
|
await websocket.send_json(
|
|
{
|
|
"type": "error",
|
|
"data": {"message": f"干预消息入队失败: {e}"},
|
|
}
|
|
)
|
|
continue
|
|
|
|
# Create a fresh CancellationToken for each message
|
|
message_token = CancellationToken()
|
|
|
|
# Guard against unlimited concurrent tasks
|
|
# Clean up completed tasks first
|
|
active_tasks.difference_update(t for t in active_tasks if t.done())
|
|
if len(active_tasks) >= _MAX_CONCURRENT_TASKS:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "error",
|
|
"data": {
|
|
"message": "Too many concurrent requests. Please wait for the current task to complete."
|
|
},
|
|
}
|
|
)
|
|
continue
|
|
|
|
# Run in background task so the WebSocket receive loop stays free
|
|
# to process confirmation_reply / reply messages while the agent
|
|
# is waiting for user confirmation (otherwise deadlock).
|
|
task = asyncio.create_task(
|
|
_handle_chat_message(
|
|
websocket,
|
|
session_id,
|
|
content,
|
|
sm,
|
|
message_token,
|
|
pending_replies,
|
|
pending_confirmations,
|
|
model_override=model,
|
|
)
|
|
)
|
|
active_tasks.add(task)
|
|
task.add_done_callback(active_tasks.discard)
|
|
|
|
elif msg_type == "reply":
|
|
# Reply to AskHumanTool
|
|
request_id = msg.get("request_id")
|
|
reply_content = msg.get("content", "")
|
|
if request_id and request_id in pending_replies:
|
|
pending_replies[request_id].set_result(reply_content)
|
|
|
|
elif msg_type == "confirmation_reply":
|
|
# Reply to confirmation request
|
|
confirmation_id = msg.get("confirmation_id")
|
|
approved = msg.get("approved", False)
|
|
logger.info(
|
|
f"Received confirmation_reply: id={confirmation_id!r}, approved={approved}"
|
|
)
|
|
if confirmation_id and confirmation_id in pending_confirmations:
|
|
pending_confirmations[confirmation_id].set_result(approved)
|
|
logger.info(f"Confirmation {confirmation_id} set_result({approved})")
|
|
else:
|
|
logger.warning(
|
|
f"Confirmation {confirmation_id!r} not found in pending_confirmations"
|
|
)
|
|
|
|
elif msg_type == "cancel":
|
|
cancellation_token.cancel()
|
|
await websocket.send_json({"type": "result", "data": {"status": "cancelled"}})
|
|
|
|
elif msg_type == "ping":
|
|
await websocket.send_json({"type": "pong"})
|
|
|
|
except WebSocketDisconnect:
|
|
logger.debug(f"Chat WebSocket disconnected for session {session_id}")
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Chat WebSocket error for session {session_id}: {e}")
|
|
try:
|
|
await websocket.send_json({"type": "error", "data": {"message": str(e)}})
|
|
except (ConnectionError, RuntimeError, asyncio.TimeoutError):
|
|
pass
|
|
finally:
|
|
# Clean up pending futures
|
|
for fut in pending_replies.values():
|
|
if not fut.done():
|
|
fut.cancel()
|
|
for fut in pending_confirmations.values():
|
|
if not fut.done():
|
|
fut.cancel()
|
|
chat_manager.remove(session_id, websocket)
|
|
|
|
|
|
async def _handle_chat_message(
|
|
websocket: WebSocket,
|
|
session_id: str,
|
|
content: str,
|
|
sm: SessionManager,
|
|
cancellation_token: CancellationToken,
|
|
pending_replies: dict[str, asyncio.Future],
|
|
pending_confirmations: dict[str, asyncio.Future] | None = None,
|
|
model_override: str | None = None,
|
|
) -> None:
|
|
"""Handle a user message: append to session, execute Agent, stream events.
|
|
|
|
Uses RequestPreprocessor for minimal preprocessing: @skill prefix + greeting regex + REACT.
|
|
|
|
Board Meeting mode: @board prefix is intercepted before RequestPreprocessor
|
|
and routed to BoardOrchestrator for multi-round group discussion.
|
|
|
|
Team Collaboration mode: @team prefix is intercepted before RequestPreprocessor
|
|
and routed to TeamOrchestrator for pipeline-based expert collaboration.
|
|
"""
|
|
from agentkit.chat.request_preprocessor import RequestPreprocessor
|
|
|
|
# Board Meeting mode: intercept @board prefix before any other preprocessing
|
|
if await _execute_board_meeting(websocket, session_id, content, sm):
|
|
return
|
|
|
|
# Team Collaboration mode: intercept @team prefix before any other preprocessing
|
|
if await _execute_team_collab(websocket, session_id, content, sm):
|
|
return
|
|
|
|
# Resolve Agent first (needed for default tools/prompt)
|
|
pool = websocket.app.state.agent_pool
|
|
session = await sm.get_session(session_id)
|
|
if session is None:
|
|
await websocket.send_json({"type": "error", "data": {"message": "Session lost"}})
|
|
return
|
|
|
|
agent = pool.get_agent(session.agent_name)
|
|
if agent is None:
|
|
await websocket.send_json(
|
|
{"type": "error", "data": {"message": f"Agent '{session.agent_name}' not found"}}
|
|
)
|
|
return
|
|
|
|
# Default execution parameters from agent
|
|
default_tools = agent._tool_registry.list_tools() if agent._tool_registry else []
|
|
default_system_prompt = getattr(agent, "_system_prompt", None) or (
|
|
agent.get_system_prompt() if hasattr(agent, "get_system_prompt") else None
|
|
)
|
|
default_model = (
|
|
agent.get_model()
|
|
if hasattr(agent, "get_model")
|
|
else getattr(agent, "_llm_model", "default")
|
|
)
|
|
|
|
# Resolve skill routing using RequestPreprocessor
|
|
skill_registry = getattr(websocket.app.state, "skill_registry", None)
|
|
request_preprocessor: RequestPreprocessor = websocket.app.state.request_preprocessor
|
|
|
|
routing = await request_preprocessor.preprocess(
|
|
content=content,
|
|
skill_registry=skill_registry,
|
|
default_tools=default_tools,
|
|
default_system_prompt=default_system_prompt,
|
|
default_model=default_model,
|
|
default_agent_name=agent.name,
|
|
)
|
|
|
|
# Debug: log tools that will be passed to ReActEngine
|
|
tool_names = [t.name for t in routing.tools]
|
|
logger.info(
|
|
f"Chat {session_id}: resolved {len(routing.tools)} tools: {tool_names}, model={routing.model}, skill={routing.skill_name}"
|
|
)
|
|
|
|
# Apply model override from frontend selector
|
|
if model_override:
|
|
routing.model = model_override
|
|
|
|
# Notify frontend about skill match
|
|
if routing.matched:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "skill_match",
|
|
"data": {
|
|
"skill": routing.skill_name,
|
|
"method": routing.match_method,
|
|
"confidence": routing.match_confidence,
|
|
},
|
|
}
|
|
)
|
|
|
|
# Append user message (use clean_content if @skill: prefix was stripped)
|
|
await sm.append_message(
|
|
session_id=session_id, role=MessageRole.USER, content=routing.clean_content
|
|
)
|
|
|
|
# Get full conversation history
|
|
chat_messages = await sm.get_chat_messages(session_id)
|
|
|
|
# Handle DIRECT_CHAT: direct LLM call, no ReAct loop
|
|
if routing.execution_mode == ExecutionMode.DIRECT_CHAT:
|
|
direct_messages = []
|
|
if routing.system_prompt:
|
|
direct_messages.append({"role": "system", "content": routing.system_prompt})
|
|
direct_messages.extend(chat_messages)
|
|
# Resolve department context for quota enforcement (U4/KTD-3).
|
|
ws_user_id, ws_dept_ids, ws_db_path = await _resolve_ws_dept_context(websocket)
|
|
try:
|
|
response = await websocket.app.state.llm_gateway.chat(
|
|
messages=direct_messages,
|
|
model=routing.model or "default",
|
|
agent_name=agent.name,
|
|
task_type="chat",
|
|
user_id=ws_user_id,
|
|
department_ids=ws_dept_ids if ws_dept_ids else None,
|
|
db_path=ws_db_path,
|
|
)
|
|
final_content = response.content or ""
|
|
if not final_content or not final_content.strip():
|
|
from agentkit.core.fallback import EMPTY_LLM_RESPONSE
|
|
|
|
final_content = EMPTY_LLM_RESPONSE
|
|
await websocket.send_json(
|
|
{
|
|
"type": "final_answer",
|
|
"content": final_content,
|
|
"is_final": True,
|
|
}
|
|
)
|
|
await sm.append_message(
|
|
session_id=session_id,
|
|
role=MessageRole.ASSISTANT,
|
|
content=final_content,
|
|
agent_name=agent.name,
|
|
)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as e:
|
|
# Check if this is a QuotaExceededError (U4: WebSocket quota).
|
|
from agentkit.llm.gateway import QuotaExceededError
|
|
|
|
if isinstance(e, QuotaExceededError):
|
|
logger.warning(
|
|
"WebSocket DIRECT_CHAT quota exceeded for session %s: %s",
|
|
session_id,
|
|
e,
|
|
)
|
|
await websocket.send_json(
|
|
{
|
|
"type": "error",
|
|
"data": {
|
|
"message": "quota exceeded",
|
|
"quota_info": {
|
|
"department_id": e.department_id,
|
|
"quota_type": e.quota_type,
|
|
"period": e.period,
|
|
"limit": e.limit,
|
|
"current": e.current,
|
|
},
|
|
},
|
|
}
|
|
)
|
|
else:
|
|
logger.error(f"Chat DIRECT_CHAT error for session {session_id}: {e}")
|
|
await websocket.send_json({"type": "error", "data": {"message": str(e)[:200]}})
|
|
return
|
|
|
|
# U4/G6: PLAN_EXEC — build PhasePolicy from server config.
|
|
# KTD5 (Wave 2): fallback chain NOT applied to PLAN_EXEC — phase policy and
|
|
# fallback chain are mutually exclusive. PLAN_EXEC uses its own engine.
|
|
# U3: logic extracted into _build_phase_engine so REST can reuse it.
|
|
plan_exec_engine, plan_exec_tools, plan_exec_error = _build_phase_engine(
|
|
server_config=getattr(websocket.app.state, "server_config", None),
|
|
llm_gateway=websocket.app.state.llm_gateway,
|
|
execution_mode=routing.execution_mode,
|
|
base_tools=routing.tools,
|
|
session_id=session_id,
|
|
)
|
|
if plan_exec_error is not None:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "error",
|
|
# Truncate to 200 chars to match nearby error paths and
|
|
# avoid leaking config internals (see chat.py:1090, 1320).
|
|
"data": {"message": plan_exec_error},
|
|
}
|
|
)
|
|
return
|
|
|
|
# Handle advanced execution modes: REWOO/REFLEXION/TEAM_COLLAB
|
|
# still fall back to REACT with a warning. PLAN_EXEC is handled above.
|
|
if routing.execution_mode not in (
|
|
ExecutionMode.REACT,
|
|
ExecutionMode.SKILL_REACT,
|
|
ExecutionMode.PLAN_EXEC,
|
|
):
|
|
logger.warning(
|
|
f"Execution mode {routing.execution_mode.value} not implemented "
|
|
f"in chat WebSocket path, falling back to REACT"
|
|
)
|
|
|
|
# Execute Agent with streaming
|
|
# Reuse Agent's ReActEngine if available (U2: Chat pipeline optimization).
|
|
# PLAN_EXEC creates a fresh engine with phase_policy set (cannot reuse the
|
|
# agent's _react_engine — it has no policy).
|
|
if plan_exec_engine is not None:
|
|
react_engine = plan_exec_engine
|
|
routing.tools = plan_exec_tools
|
|
else:
|
|
react_engine = getattr(agent, "_react_engine", None)
|
|
if react_engine is None:
|
|
react_engine = ReActEngine(llm_gateway=websocket.app.state.llm_gateway)
|
|
else:
|
|
react_engine.reset()
|
|
|
|
# Create confirmation handler that sends request to frontend and waits for reply
|
|
# Use the same dict object — do NOT use `or {}` because an empty dict is falsy
|
|
# and would create a new dict, breaking the shared state with the WS loop.
|
|
_pending_confirmations = pending_confirmations if pending_confirmations is not None else {}
|
|
|
|
async def _confirmation_handler(confirmation_id: str, command: str, reason: str) -> bool:
|
|
"""Send confirmation request to frontend via WebSocket and wait for user reply."""
|
|
# Send confirmation request to frontend
|
|
await websocket.send_json(
|
|
{
|
|
"type": "confirmation_request",
|
|
"data": {
|
|
"confirmation_id": confirmation_id,
|
|
"command": command,
|
|
"reason": reason,
|
|
},
|
|
}
|
|
)
|
|
|
|
# Create a Future and wait for the user's reply
|
|
loop = asyncio.get_running_loop()
|
|
future: asyncio.Future[bool] = loop.create_future()
|
|
_pending_confirmations[confirmation_id] = future
|
|
logger.info(f"Confirmation request {confirmation_id} sent, waiting for reply")
|
|
|
|
try:
|
|
# Wait up to 5 minutes for user confirmation
|
|
result = await asyncio.wait_for(future, timeout=300.0)
|
|
logger.info(f"Confirmation request {confirmation_id} resolved: {result}")
|
|
# Immediately notify frontend of the result so the card updates
|
|
# without waiting for the tool to re-execute
|
|
await websocket.send_json(
|
|
{
|
|
"type": "confirmation_result",
|
|
"data": {"confirmation_id": confirmation_id, "approved": result},
|
|
}
|
|
)
|
|
return result
|
|
except asyncio.TimeoutError:
|
|
logger.warning(f"Confirmation request {confirmation_id} timed out")
|
|
return False
|
|
except asyncio.CancelledError:
|
|
logger.warning(f"Confirmation request {confirmation_id} cancelled")
|
|
return False
|
|
finally:
|
|
_pending_confirmations.pop(confirmation_id, None)
|
|
|
|
logger.info(
|
|
f"Chat session {session_id}: executing with {len(routing.tools)} tools, model={routing.model}, skill={routing.skill_name}"
|
|
)
|
|
|
|
try:
|
|
final_content = ""
|
|
token_buffer: list[str] = []
|
|
# Track phase transitions for phase_changed events (PLAN_EXEC only).
|
|
# For non-PLAN_EXEC modes, current_phase is always None → no events.
|
|
prev_phase = react_engine.current_phase
|
|
async for event in react_engine.execute_stream(
|
|
messages=chat_messages,
|
|
tools=routing.tools,
|
|
model=routing.model,
|
|
agent_name=routing.agent_name,
|
|
system_prompt=routing.system_prompt,
|
|
cancellation_token=cancellation_token,
|
|
confirmation_handler=_confirmation_handler,
|
|
):
|
|
if event.event_type == "final_answer":
|
|
# Flush any buffered tokens as a single write
|
|
if token_buffer:
|
|
await websocket.send_json({"type": "token", "content": "".join(token_buffer)})
|
|
token_buffer.clear()
|
|
# Then send final answer
|
|
final_content = event.data.get("output", "")
|
|
if not final_content or not final_content.strip():
|
|
from agentkit.core.fallback import EMPTY_LLM_RESPONSE
|
|
|
|
final_content = EMPTY_LLM_RESPONSE
|
|
await websocket.send_json(
|
|
{
|
|
"type": "final_answer",
|
|
"content": final_content,
|
|
"is_final": True,
|
|
}
|
|
)
|
|
elif event.event_type == "token":
|
|
# Buffer tokens instead of sending immediately
|
|
token_buffer.append(event.data.get("content", ""))
|
|
elif event.event_type == "thinking":
|
|
# If we have buffered tokens, convert them to a thinking event
|
|
if token_buffer:
|
|
buffered_text = "".join(token_buffer)
|
|
token_buffer.clear()
|
|
await websocket.send_json({"type": "thinking", "content": buffered_text})
|
|
# Also send the thinking event content
|
|
thinking_msg = event.data.get("message", "")
|
|
if thinking_msg:
|
|
await websocket.send_json({"type": "thinking", "content": thinking_msg})
|
|
elif event.event_type == "tool_call":
|
|
# Convert buffered tokens to thinking (they were "thinking" text before tool call)
|
|
if token_buffer:
|
|
buffered_text = "".join(token_buffer)
|
|
token_buffer.clear()
|
|
await websocket.send_json({"type": "thinking", "content": buffered_text})
|
|
await websocket.send_json(
|
|
{
|
|
"type": "step",
|
|
"data": {
|
|
"event_type": event.event_type,
|
|
"step": event.step,
|
|
"data": event.data,
|
|
},
|
|
}
|
|
)
|
|
elif event.event_type == "confirmation_request":
|
|
pass
|
|
elif event.event_type == "confirmation_result":
|
|
await websocket.send_json(
|
|
{
|
|
"type": "confirmation_result",
|
|
"data": event.data,
|
|
}
|
|
)
|
|
elif event.event_type == "phase_violation":
|
|
# Wave 4 U2: forward phase violations to the client so the
|
|
# frontend can surface them in the PhaseIndicator UI (alongside
|
|
# the LLM reinjection that already happens via the tool_result
|
|
# error dict).
|
|
await websocket.send_json(
|
|
{
|
|
"type": "phase_violation",
|
|
"data": event.data,
|
|
}
|
|
)
|
|
else:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "step",
|
|
"data": {
|
|
"event_type": event.event_type,
|
|
"step": event.step,
|
|
"data": event.data,
|
|
},
|
|
}
|
|
)
|
|
|
|
# U4/G6: emit phase_changed event when the phase state machine
|
|
# transitions (PLAN_EXEC only). For non-PLAN_EXEC modes,
|
|
# current_phase is always None → this branch never fires.
|
|
curr_phase = react_engine.current_phase
|
|
if curr_phase != prev_phase:
|
|
await websocket.send_json(
|
|
{
|
|
"type": "phase_changed",
|
|
"data": {
|
|
"phase": curr_phase.value if curr_phase else None,
|
|
"previous": prev_phase.value if prev_phase else None,
|
|
},
|
|
}
|
|
)
|
|
prev_phase = curr_phase
|
|
|
|
# Append assistant reply to session
|
|
if final_content:
|
|
await sm.append_message(
|
|
session_id=session_id,
|
|
role=MessageRole.ASSISTANT,
|
|
content=final_content,
|
|
agent_name=agent.name,
|
|
)
|
|
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Chat execution error for session {session_id}: {e}")
|
|
# Show meaningful error to user, but avoid leaking full stack traces
|
|
error_msg = str(e)
|
|
# Truncate very long error messages
|
|
if len(error_msg) > 200:
|
|
error_msg = error_msg[:200] + "..."
|
|
await websocket.send_json({"type": "error", "data": {"message": error_msg}})
|
|
|
|
|
|
# ── File upload ───────────────────────────────────────────────────────
|
|
|
|
|
|
MAX_UPLOAD_SIZE = 10 * 1024 * 1024 # 10 MB
|
|
UPLOAD_DIR = Path(os.environ.get("AGENTKIT_UPLOAD_DIR", "data/uploads"))
|
|
|
|
|
|
def _ensure_upload_dir() -> Path:
|
|
UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
|
|
return UPLOAD_DIR
|
|
|
|
|
|
def _sanitize_filename(name: str) -> str:
|
|
"""Remove path separators and keep only safe characters."""
|
|
name = name.replace("\\", "_").replace("/", "_")
|
|
return "".join(c for c in name if c.isalnum() or c in "._-").strip(".")
|
|
|
|
|
|
@router.post("/upload")
|
|
async def upload_chat_file(file: UploadFile = File(...)) -> dict[str, Any]:
|
|
"""Upload a file to be referenced in chat messages.
|
|
|
|
Returns metadata including a public download URL.
|
|
"""
|
|
if file.size is not None and file.size > MAX_UPLOAD_SIZE:
|
|
raise HTTPException(status_code=413, detail="File exceeds 10 MB limit")
|
|
|
|
original_name = file.filename or "unnamed"
|
|
safe_name = _sanitize_filename(original_name) or "unnamed"
|
|
ext = Path(safe_name).suffix
|
|
stored_name = f"{uuid.uuid4().hex}{ext}"
|
|
upload_dir = _ensure_upload_dir()
|
|
file_path = upload_dir / stored_name
|
|
|
|
try:
|
|
contents = await file.read()
|
|
if len(contents) > MAX_UPLOAD_SIZE:
|
|
raise HTTPException(status_code=413, detail="File exceeds 10 MB limit")
|
|
file_path.write_bytes(contents)
|
|
except HTTPException:
|
|
raise
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc:
|
|
logger.error(f"Failed to save uploaded file: {exc}")
|
|
raise HTTPException(status_code=500, detail="Failed to save file") from exc
|
|
finally:
|
|
await file.close()
|
|
|
|
return {
|
|
"filename": original_name,
|
|
"stored_name": stored_name,
|
|
"content_type": file.content_type or "application/octet-stream",
|
|
"size": file_path.stat().st_size,
|
|
"download_url": f"/api/v1/chat/uploads/{stored_name}",
|
|
}
|
|
|
|
|
|
@router.get("/uploads/{filename}")
|
|
async def download_chat_file(filename: str) -> FileResponse:
|
|
"""Download an uploaded chat file by its stored filename."""
|
|
upload_dir = _ensure_upload_dir()
|
|
safe_filename = _sanitize_filename(filename)
|
|
file_path = upload_dir / safe_filename
|
|
if not file_path.exists() or not file_path.is_file():
|
|
raise HTTPException(status_code=404, detail="File not found")
|
|
return FileResponse(file_path, filename=safe_filename)
|