534 lines
21 KiB
Python
534 lines
21 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
from typing import Any
|
||
|
||
from app.core.contracts import ExecutionDirective
|
||
from app.core.intent_parser import IntentParser
|
||
from app.events.event_bus import EventBus
|
||
from app.events.event_types import (
|
||
ORCHESTRATOR_CALLED,
|
||
ORCHESTRATOR_FALLBACK_USED,
|
||
ORCHESTRATOR_RETRY,
|
||
ORCHESTRATOR_RESULT,
|
||
ORCHESTRATOR_UNAVAILABLE,
|
||
THINKER_CALLED,
|
||
THINKER_RESULT,
|
||
JSON_COMPILER_CALLED,
|
||
JSON_COMPILER_RESULT,
|
||
)
|
||
from app.models.async_adapters import AsyncOrchestratorAdapter
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class AsyncRouter:
|
||
"""Async router using Thinker + JSON Compiler pipeline."""
|
||
|
||
def __init__(
|
||
self,
|
||
thinker: AsyncOrchestratorAdapter | None = None,
|
||
json_compiler: AsyncOrchestratorAdapter | None = None,
|
||
intent_parser: IntentParser | None = None,
|
||
prompts: dict[str, str] | None = None,
|
||
event_bus: EventBus | None = None,
|
||
tool_registry=None,
|
||
retry_limit: int = 2,
|
||
debug: bool = False,
|
||
log_length: int = 500,
|
||
json_fix_retry_limit: int = 2,
|
||
json_fix_use_sys_util: bool = True,
|
||
intent_classifier: str = "thinker",
|
||
) -> None:
|
||
self._thinker = thinker
|
||
self._json_compiler = json_compiler
|
||
self._intent_classifier = intent_classifier
|
||
self._sys_util = None
|
||
self._intent_parser = intent_parser or IntentParser()
|
||
self._prompts = prompts or {}
|
||
self._event_bus = event_bus
|
||
self._tool_registry = tool_registry
|
||
self._retry_limit = retry_limit
|
||
self._debug = debug
|
||
self._log_length = log_length
|
||
self._json_fix_retry_limit = json_fix_retry_limit
|
||
self._json_fix_use_sys_util = json_fix_use_sys_util
|
||
self._orchestrator = None # Set separately if needed for classification
|
||
|
||
def set_event_bus(self, event_bus: EventBus) -> None:
|
||
self._event_bus = event_bus
|
||
|
||
def set_thinker(self, thinker: AsyncOrchestratorAdapter) -> None:
|
||
self._thinker = thinker
|
||
|
||
def set_json_compiler(self, json_compiler: AsyncOrchestratorAdapter) -> None:
|
||
self._json_compiler = json_compiler
|
||
|
||
def set_sys_util(self, sys_util: AsyncOrchestratorAdapter) -> None:
|
||
self._sys_util = sys_util
|
||
|
||
def set_orchestrator(self, orchestrator: AsyncOrchestratorAdapter) -> None:
|
||
self._orchestrator = orchestrator
|
||
|
||
def set_tool_registry(self, tool_registry) -> None:
|
||
self._tool_registry = tool_registry
|
||
|
||
async def decide(
|
||
self,
|
||
state: dict[str, Any],
|
||
context: dict[str, Any],
|
||
task_id: str | None = None,
|
||
session_id: str | None = None,
|
||
) -> ExecutionDirective:
|
||
task_context = context.get("task_context", {})
|
||
requested_tool = task_context.get("requested_tool")
|
||
task_summary = str(context.get("task_summary", ""))
|
||
|
||
if requested_tool:
|
||
self._emit_event(
|
||
ORCHESTRATOR_RESULT,
|
||
{"reason": "explicit_tool_request", "tool": requested_tool},
|
||
task_id,
|
||
session_id,
|
||
)
|
||
return ExecutionDirective(
|
||
type="tool",
|
||
payload={
|
||
"tool": requested_tool,
|
||
"args": task_context.get("tool_args", {}),
|
||
},
|
||
requires_permission=requested_tool in {"shell_exec", "file_write"},
|
||
confidence=0.9,
|
||
reason="Task context explicitly requested a tool execution.",
|
||
)
|
||
|
||
parsed_intent = self._intent_parser.parse(task_summary)
|
||
if parsed_intent:
|
||
self._emit_event(
|
||
ORCHESTRATOR_RESULT,
|
||
{"reason": "deterministic_intent_parser", "directive": parsed_intent.model_dump(mode="json")},
|
||
task_id,
|
||
session_id,
|
||
)
|
||
return parsed_intent
|
||
|
||
if self._thinker is None:
|
||
fallback = self._fallback_directive(task_summary)
|
||
self._emit_event(
|
||
ORCHESTRATOR_FALLBACK_USED,
|
||
{"reason": "thinker_unavailable", "directive": fallback.model_dump(mode="json")},
|
||
task_id,
|
||
session_id,
|
||
)
|
||
return fallback
|
||
|
||
if self._json_compiler is None:
|
||
fallback = self._fallback_directive(task_summary)
|
||
self._emit_event(
|
||
ORCHESTRATOR_FALLBACK_USED,
|
||
{"reason": "json_compiler_unavailable", "directive": fallback.model_dump(mode="json")},
|
||
task_id,
|
||
session_id,
|
||
)
|
||
return fallback
|
||
|
||
mode_hint = await self._classify_intent(task_summary)
|
||
thinker_prompt = self._build_thinker_prompt(task_summary, context, mode_hint)
|
||
|
||
for thinker_attempt in range(self._retry_limit + 1):
|
||
if thinker_attempt > 0:
|
||
self._emit_event(
|
||
ORCHESTRATOR_RETRY,
|
||
{"attempt": thinker_attempt, "prompt": thinker_prompt},
|
||
task_id,
|
||
session_id,
|
||
)
|
||
thinker_prompt = self._add_thinker_feedback(thinker_prompt, last_thinker_error, thinker_attempt)
|
||
|
||
self._emit_event(
|
||
THINKER_CALLED,
|
||
{"attempt": thinker_attempt, "mode": mode_hint},
|
||
task_id,
|
||
session_id,
|
||
)
|
||
|
||
try:
|
||
thinker_result = await self._thinker.generate(thinker_prompt)
|
||
except Exception as e:
|
||
logger.warning(f"Thinker generate failed: {e}")
|
||
last_thinker_error = str(e)
|
||
continue
|
||
|
||
logger.info(f"Thinker result (attempt {thinker_attempt + 1}): {thinker_result}")
|
||
self._emit_event(
|
||
THINKER_RESULT,
|
||
{"result": thinker_result, "attempt": thinker_attempt},
|
||
task_id,
|
||
session_id,
|
||
)
|
||
|
||
if mode_hint == "conversation" and self._looks_like_tool_plan(thinker_result):
|
||
mode_hint = "execution"
|
||
self._emit_event(
|
||
ORCHESTRATOR_FALLBACK_USED,
|
||
{"reason": "thinker_proposed_tool_plan_despite_conversation_hint"},
|
||
task_id,
|
||
session_id,
|
||
)
|
||
|
||
if self._is_simple_response(thinker_result):
|
||
json_compiler_prompt = self._build_json_compiler_prompt(thinker_result)
|
||
else:
|
||
json_compiler_prompt = self._build_json_compiler_prompt(thinker_result)
|
||
|
||
for compiler_attempt in range(self._json_fix_retry_limit + 1):
|
||
self._emit_event(
|
||
JSON_COMPILER_CALLED,
|
||
{"attempt": compiler_attempt, "plan": thinker_result},
|
||
task_id,
|
||
session_id,
|
||
)
|
||
|
||
try:
|
||
compiler_result = await self._json_compiler.generate(json_compiler_prompt)
|
||
except Exception as e:
|
||
logger.warning(f"JSON Compiler generate failed: {e}")
|
||
compiler_result = None
|
||
|
||
if compiler_result:
|
||
logger.info(f"JSON Compiler result (attempt {compiler_attempt + 1}): {compiler_result}")
|
||
self._emit_event(
|
||
JSON_COMPILER_RESULT,
|
||
{"result": compiler_result, "attempt": compiler_attempt},
|
||
task_id,
|
||
session_id,
|
||
)
|
||
|
||
directive = self._validate_directive(compiler_result, mode_hint) if compiler_result else None
|
||
if directive is not None:
|
||
directive = self._guard_rail_check(directive)
|
||
self._emit_event(
|
||
ORCHESTRATOR_RESULT,
|
||
{"directive": directive.model_dump(mode="json"), "thinker_attempt": thinker_attempt, "compiler_attempt": compiler_attempt},
|
||
task_id,
|
||
session_id,
|
||
)
|
||
return directive
|
||
|
||
if compiler_result:
|
||
logger.warning(f"JSON Compiler validation failed, attempting fix (attempt {compiler_attempt + 1})")
|
||
fix_result = await self._fix_invalid_json(compiler_result, compiler_attempt, task_id, session_id)
|
||
if fix_result:
|
||
fixed_directive = self._validate_directive(fix_result, mode_hint)
|
||
if fixed_directive is not None:
|
||
fixed_directive = self._guard_rail_check(fixed_directive)
|
||
self._emit_event(
|
||
ORCHESTRATOR_RESULT,
|
||
{"directive": fixed_directive.model_dump(mode="json"), "fixed": True},
|
||
task_id,
|
||
session_id,
|
||
)
|
||
return fixed_directive
|
||
|
||
last_thinker_error = f"JSON Compiler failed after {self._json_fix_retry_limit + 1} attempts"
|
||
|
||
self._emit_event(
|
||
ORCHESTRATOR_UNAVAILABLE,
|
||
{"reason": "retry_exhausted", "last_error": last_thinker_error},
|
||
task_id,
|
||
session_id,
|
||
)
|
||
raise RuntimeError(f"Thinker/Compiler pipeline failed after {self._retry_limit + 1} attempts")
|
||
|
||
def _fallback_directive(self, task_summary: str) -> ExecutionDirective:
|
||
parsed = self._intent_parser.parse(task_summary)
|
||
if parsed:
|
||
return parsed
|
||
|
||
return ExecutionDirective(
|
||
type="respond",
|
||
payload={"text": f"Runtime accepted task: {task_summary}"},
|
||
requires_permission=False,
|
||
confidence=0.4,
|
||
reason="Fallback response because local orchestration models are not loaded.",
|
||
)
|
||
|
||
def _is_simple_response(self, thinker_result: str) -> bool:
|
||
result_lower = thinker_result.lower().strip()
|
||
return result_lower.startswith("ответ:") or result_lower.startswith("response:") or "не нужно" in result_lower
|
||
|
||
def _extract_conversation_response(self, thinker_result: str) -> str:
|
||
"""Extract text response from thinker result for conversation mode."""
|
||
result_lower = thinker_result.lower()
|
||
|
||
# Skip the ПЛАН lines, just get the ОТВЕТ part
|
||
lines = thinker_result.split('\n')
|
||
response_lines = []
|
||
capture = False
|
||
|
||
for line in lines:
|
||
if line.strip().lower().startswith('ответ:') or line.strip().lower().startswith('response:'):
|
||
capture = True
|
||
response_lines.append(line)
|
||
elif capture and line.strip():
|
||
# Check if this is a new ПЛАН or step
|
||
if line.strip().lower().startswith('план') or line.strip().lower().startswith('step'):
|
||
break
|
||
response_lines.append(line)
|
||
|
||
if response_lines:
|
||
return '\n'.join(response_lines).replace('ответ:', '').replace('response:', '').strip()
|
||
|
||
# Fallback: return first few sentences
|
||
sentences = thinker_result.split('.')[:3]
|
||
return '. '.join(sentences).strip()
|
||
|
||
def _looks_like_tool_plan(self, thinker_result: str) -> bool:
|
||
result = thinker_result.lower()
|
||
tool_names = set()
|
||
if self._tool_registry:
|
||
tool_names = set(self._tool_registry.list_names())
|
||
tool_markers = {"shell_exec", "file_read", "file_write", "memory", *tool_names}
|
||
plan_markers = ("план:", "шаг", "step", "tool", "инструмент")
|
||
return any(marker in result for marker in tool_markers) and any(marker in result for marker in plan_markers)
|
||
|
||
def _build_thinker_prompt(
|
||
self, task_summary: str, context: dict[str, Any], mode_hint: str
|
||
) -> str:
|
||
base_prompt = self._prompts.get("thinker", "")
|
||
memory_context = context.get("memory_context", [])
|
||
|
||
tools_json = "[]"
|
||
if self._tool_registry:
|
||
schemas = self._tool_registry.list_schemas()
|
||
tools_json = json.dumps(schemas, ensure_ascii=False, indent=2)
|
||
|
||
prompt_lines = [
|
||
base_prompt,
|
||
"",
|
||
f"Task: {task_summary}",
|
||
f"Mode hint: {mode_hint}",
|
||
]
|
||
|
||
if memory_context:
|
||
memory_text = "\n".join([f"- {m.get('text', '')}" for m in memory_context[:5]])
|
||
prompt_lines.append(f"\nRelevant memory:\n{memory_text}")
|
||
|
||
session_history = context.get("session_history", [])
|
||
if session_history:
|
||
history_text = "\n".join([f"- {h.get('text', '')}" for h in session_history[:3]])
|
||
prompt_lines.append(f"\nPrevious requests in this session:\n{history_text}")
|
||
|
||
prompt_lines.extend([
|
||
"",
|
||
f"AVAILABLE TOOLS (JSON):",
|
||
tools_json,
|
||
"",
|
||
])
|
||
|
||
return "\n".join(prompt_lines)
|
||
|
||
def _build_json_compiler_prompt(self, thinker_result: str) -> str:
|
||
base_prompt = self._prompts.get("json_compiler", "")
|
||
|
||
prompt_lines = [
|
||
base_prompt,
|
||
"",
|
||
"Thinker's plan:",
|
||
thinker_result,
|
||
"",
|
||
]
|
||
|
||
return "\n".join(prompt_lines)
|
||
|
||
def _determine_mode_from_context(self, context: dict[str, Any]) -> str:
|
||
"""Legacy method - kept for compatibility"""
|
||
task_summary = str(context.get("task_summary", "")).lower()
|
||
keywords = ["запусти", "выполни", "создай", "напиши", "удали", "run", "execute", "create"]
|
||
for kw in keywords:
|
||
if kw in task_summary:
|
||
return "execution"
|
||
return "conversation"
|
||
|
||
async def _classify_intent(self, task_summary: str) -> str:
|
||
"""LLM-based intent classification"""
|
||
if self._intent_classifier == "orchestrator" and self._orchestrator:
|
||
classifier_model = self._orchestrator
|
||
else:
|
||
classifier_model = self._thinker
|
||
|
||
if not classifier_model:
|
||
logger.warning("No classifier model available, using default")
|
||
return "conversation"
|
||
|
||
classification_prompt = f"""Классифицируй запрос пользователя: "{task_summary}"
|
||
|
||
Классы:
|
||
- execution: чтобы ответить, агенту нужно обратиться к локальной среде, файлам, shell, tools, памяти, сети или выполнить проверку/операцию. Это включает вопросы о текущем состоянии ПК, установленных пакетах, файлах, процессах, времени работы, обновлениях, логах.
|
||
- conversation: можно ответить сразу из диалога и общих знаний, без проверки локальной среды и без tools.
|
||
- clarification_needed: нельзя понять, что именно пользователь хочет.
|
||
|
||
Верни ровно один токен без рассуждений: execution или conversation или clarification_needed"""
|
||
|
||
try:
|
||
result = await classifier_model.generate(classification_prompt)
|
||
classification = self._extract_classification(result)
|
||
if classification:
|
||
logger.info(f"Intent classified: {classification} for task: {task_summary}")
|
||
return classification
|
||
|
||
logger.warning(f"Invalid classification result: {result}, defaulting to conversation")
|
||
return "conversation"
|
||
except Exception as e:
|
||
logger.warning(f"Intent classification failed: {e}, defaulting to conversation")
|
||
return "conversation"
|
||
|
||
def _extract_classification(self, raw_result: str) -> str | None:
|
||
result = raw_result.strip().lower()
|
||
allowed = {"execution", "conversation", "clarification_needed"}
|
||
if result in allowed:
|
||
return result
|
||
|
||
result = re.sub(r"<think>.*?</think>", " ", result, flags=re.DOTALL)
|
||
if (
|
||
"shell_exec" in result
|
||
or "execute command" in result
|
||
or "command execution" in result
|
||
or "use the tool" in result
|
||
or "use a tool" in result
|
||
):
|
||
return "execution"
|
||
tokens = re.findall(r"\b(execution|conversation|clarification_needed)\b", result)
|
||
if tokens:
|
||
return tokens[-1]
|
||
|
||
first_word = result.split()[0] if result.split() else ""
|
||
if first_word in allowed:
|
||
return first_word
|
||
|
||
return None
|
||
|
||
def _validate_directive(self, output: str, mode_hint: str) -> ExecutionDirective | None:
|
||
if not output:
|
||
return None
|
||
|
||
try:
|
||
json_start = output.find("{")
|
||
json_end = output.rfind("}") + 1
|
||
if json_start < 0 or json_end <= 0:
|
||
return None
|
||
|
||
json_str = output[json_start:json_end]
|
||
data = json.loads(json_str)
|
||
|
||
if "type" not in data:
|
||
return None
|
||
|
||
msg_type = data.get("type", "")
|
||
payload = data.get("payload", {})
|
||
|
||
if msg_type == "step" and "tool" in payload:
|
||
tool = payload.get("tool", "")
|
||
args = payload.get("args", {})
|
||
payload = {"tool": tool, "args": args}
|
||
|
||
if msg_type == "plan":
|
||
payload = {"steps": payload.get("steps", [])}
|
||
|
||
return ExecutionDirective(
|
||
type=msg_type,
|
||
payload=payload,
|
||
confidence=data.get("confidence", 0.9),
|
||
reason=data.get("reason", ""),
|
||
)
|
||
except (json.JSONDecodeError, ValueError, TypeError) as e:
|
||
logger.warning(f"Directive JSON validation failed: {e}")
|
||
return None
|
||
|
||
def _guard_rail_check(self, directive: ExecutionDirective) -> ExecutionDirective:
|
||
tool_name = directive.payload.get("tool", "")
|
||
if tool_name in {"shell_exec", "file_write", "file_delete"}:
|
||
return ExecutionDirective(
|
||
type=directive.type,
|
||
payload=directive.payload,
|
||
requires_permission=True,
|
||
confidence=directive.confidence,
|
||
reason=directive.reason,
|
||
)
|
||
return directive
|
||
|
||
def _add_thinker_feedback(self, prompt: str, error: str, attempt: int) -> str:
|
||
feedback = f"\n[ATTEMPT {attempt + 1} FAILED: {error}]\n"
|
||
feedback += "Provide a valid semantic plan.\n"
|
||
return prompt + feedback
|
||
|
||
def _emit_event(
|
||
self,
|
||
event_type: str,
|
||
payload: dict[str, Any],
|
||
task_id: str | None,
|
||
session_id: str | None,
|
||
) -> None:
|
||
if self._event_bus and task_id:
|
||
from app.core.contracts import RuntimeEvent
|
||
event = RuntimeEvent(
|
||
task_id=task_id,
|
||
session_id=session_id or "unknown",
|
||
sequence=self._event_bus.next_sequence(task_id),
|
||
type=event_type,
|
||
payload=payload,
|
||
)
|
||
self._event_bus.publish(event)
|
||
|
||
SYS_UTIL_PROMPT = None
|
||
|
||
async def _fix_invalid_json(self, invalid_result: str, attempt: int, task_id: str | None, session_id: str | None) -> str | None:
|
||
"""Try to fix invalid JSON using sys_util model."""
|
||
if not self._sys_util:
|
||
return None
|
||
|
||
first_brace = invalid_result.find('{')
|
||
last_brace = invalid_result.rfind('}')
|
||
if first_brace < 0 or last_brace <= first_brace:
|
||
return None
|
||
|
||
truncated_json = invalid_result[first_brace:last_brace + 1]
|
||
|
||
error_msg = ""
|
||
try:
|
||
json.loads(truncated_json)
|
||
except json.JSONDecodeError as e:
|
||
error_msg = str(e)
|
||
|
||
sys_util_prompt = (
|
||
self._prompts.get("sys_util")
|
||
if self._prompts
|
||
else self.SYS_UTIL_PROMPT or (
|
||
"You are a STRICT JSON repair engine. "
|
||
"Your job is ONLY to fix invalid JSON syntax. "
|
||
"You MUST output valid JSON or nothing else."
|
||
)
|
||
)
|
||
fix_prompt = f"""{sys_util_prompt}
|
||
|
||
{error_msg}
|
||
|
||
Fixed JSON:"""
|
||
|
||
try:
|
||
logger.info(f"JSON fix using sys_util model (attempt {attempt + 1})")
|
||
fixed_result = await self._sys_util.generate(fix_prompt)
|
||
|
||
fixed_first = fixed_result.find('{')
|
||
fixed_last = fixed_result.rfind('}')
|
||
if fixed_first >= 0 and fixed_last > fixed_first:
|
||
return fixed_result[fixed_first:fixed_last + 1]
|
||
|
||
return None
|
||
|
||
except Exception as e:
|
||
logger.warning(f"JSON fix failed: {e}")
|
||
return None
|