556 lines
22 KiB
Python
556 lines
22 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
from threading import RLock
|
||
from pathlib import Path
|
||
|
||
from app.core.config import AppConfig, load_app_config
|
||
from app.core.context_builder import ContextBuilder
|
||
from app.core.contracts import UserTask
|
||
from app.core.execution_engine import ExecutionEngine
|
||
from app.core.execution_scheduler import ExecutionScheduler
|
||
from app.core.async_router import AsyncRouter
|
||
from app.events.event_bus import EventBus
|
||
from app.events.event_store import SQLiteEventStore
|
||
from app.memory import MemoryInterface, MemoryStore, VectorIndex
|
||
from app.memory.write_policy import MemoryWritePolicy
|
||
from app.models import (
|
||
CoderAdapter,
|
||
CriticAdapter,
|
||
EmbeddingsAdapter,
|
||
OrchestratorAdapter,
|
||
create_adapter,
|
||
)
|
||
from app.models.async_adapters import AsyncOrchestratorAdapter, AsyncCriticAdapter, AsyncCoderAdapter
|
||
from app.permissions.approval_store import SQLiteApprovalStore
|
||
from app.core.permission_service import PermissionService
|
||
from app.runtime.runtime_loop import RuntimeLoop
|
||
from app.state.checkpoint_store import SQLiteCheckpointStore
|
||
from app.state.task_state_store import SQLiteTaskStateStore
|
||
from app.tools.file_read import FileReadTool
|
||
from app.tools.file_write import FileWriteTool
|
||
from app.tools.registry import ToolRegistry
|
||
from app.tools.sandbox import ToolSandbox
|
||
from app.tools.shell_exec import ShellExecTool
|
||
from app.tools.memory_tools import MemoryInsertTool, MemorySearchTool, MemoryListTool
|
||
|
||
|
||
class RuntimeController:
|
||
"""Composition root for the ducklm runtime."""
|
||
|
||
def __init__(self, base_dir: str | Path | None = None) -> None:
|
||
self.base_dir = Path(base_dir or Path(__file__).resolve().parents[2])
|
||
self.config: AppConfig = load_app_config(self.base_dir / "config")
|
||
|
||
self.event_bus = EventBus(
|
||
SQLiteEventStore(self.base_dir / "data" / "events" / "events.sqlite3")
|
||
)
|
||
self.task_state_store = SQLiteTaskStateStore(
|
||
self.base_dir / "data" / "state" / "task_state.sqlite3"
|
||
)
|
||
self.checkpoint_store = SQLiteCheckpointStore(
|
||
self.base_dir / "data" / "state" / "checkpoints.sqlite3"
|
||
)
|
||
self.approval_store = SQLiteApprovalStore(
|
||
self.base_dir / "data" / "permissions" / "approvals.sqlite3"
|
||
)
|
||
|
||
self._thinker: OrchestratorAdapter | None = None
|
||
self._json_compiler: OrchestratorAdapter | None = None
|
||
self._orchestrator: OrchestratorAdapter | None = None
|
||
self._coder: CoderAdapter | None = None
|
||
self._critic: CriticAdapter | None = None
|
||
self._sys_util: OrchestratorAdapter | None = None
|
||
self._model_cache: dict[tuple[object, ...], tuple[object, RLock]] = {}
|
||
self._memory_interface: MemoryInterface | None = None
|
||
self._memory_policy: MemoryWritePolicy | None = None
|
||
self.tool_registry = None
|
||
self.tool_sandbox = None
|
||
|
||
self._init_models()
|
||
self._init_memory()
|
||
|
||
runtime_config = self.config.runtime
|
||
|
||
self.tool_sandbox = ToolSandbox(
|
||
allowed_root=self.base_dir,
|
||
timeout_ms=runtime_config.step_timeout_ms,
|
||
)
|
||
|
||
self.tool_registry = self._create_tool_registry()
|
||
|
||
context_config = {
|
||
"max_context_tokens": runtime_config.max_context_tokens,
|
||
"context_budgets": runtime_config.context_budgets,
|
||
"reserve_for_generation_pct": runtime_config.reserve_for_generation_pct,
|
||
}
|
||
|
||
self.context_builder = ContextBuilder(
|
||
memory_interface=self._memory_interface,
|
||
tool_registry=self.tool_registry,
|
||
config=context_config,
|
||
)
|
||
|
||
self._prompts = self._load_prompts()
|
||
# ensure sys_util prompt is present in prompts dict for router
|
||
# ensure sys_util prompt is available to router (prompts.json may have "sys_util" key)
|
||
if "sys_util" not in self._prompts and "prompts" in self.config:
|
||
self._prompts["sys_util"] = self.config.get("sys_util")
|
||
|
||
self.context_builder = ContextBuilder(
|
||
memory_interface=self._memory_interface,
|
||
tool_registry=self.tool_registry,
|
||
config=context_config,
|
||
)
|
||
|
||
self.router = AsyncRouter(
|
||
thinker=None,
|
||
json_compiler=None,
|
||
intent_parser=None,
|
||
prompts=self._prompts,
|
||
event_bus=self.event_bus,
|
||
tool_registry=self.tool_registry,
|
||
retry_limit=runtime_config.orchestrator_retry_limit,
|
||
debug=runtime_config.debug if hasattr(runtime_config, 'debug') else False,
|
||
log_length=runtime_config.debug_orchestrator_log_length if hasattr(runtime_config, 'debug_orchestrator_log_length') else 500,
|
||
json_fix_retry_limit=runtime_config.json_fix_retry_limit if hasattr(runtime_config, 'json_fix_retry_limit') else 2,
|
||
json_fix_use_sys_util=runtime_config.json_fix_use_sys_util if hasattr(runtime_config, "json_fix_use_sys_util") else True,
|
||
intent_classifier=runtime_config.intent_classifier if hasattr(runtime_config, "intent_classifier") else "thinker",
|
||
)
|
||
|
||
self.permission_service = PermissionService(
|
||
config=self._load_permissions_config(),
|
||
)
|
||
|
||
self.execution_engine = ExecutionEngine(
|
||
event_bus=self.event_bus,
|
||
tool_registry=self.tool_registry,
|
||
permission_service=self.permission_service,
|
||
scheduler=ExecutionScheduler(
|
||
retry_limit=runtime_config.planner_retry_limit
|
||
),
|
||
critic=self._critic,
|
||
memory_policy=self._memory_policy,
|
||
memory_interface=self._memory_interface,
|
||
prompts=self._prompts,
|
||
)
|
||
|
||
self.runtime_loop = RuntimeLoop(
|
||
event_bus=self.event_bus,
|
||
task_state_store=self.task_state_store,
|
||
checkpoint_store=self.checkpoint_store,
|
||
context_builder=self.context_builder,
|
||
router=self.router,
|
||
execution_engine=self.execution_engine,
|
||
permission_service=self.permission_service,
|
||
memory_interface=self._memory_interface,
|
||
)
|
||
|
||
def _load_prompts(self) -> dict[str, str]:
|
||
prompts_dir = self.base_dir / "config" / "prompts"
|
||
prompts = {}
|
||
|
||
if prompts_dir.is_dir():
|
||
for md_file in prompts_dir.glob("*.md"):
|
||
role = md_file.stem
|
||
prompts[role] = md_file.read_text(encoding="utf-8")
|
||
|
||
if prompts:
|
||
return prompts
|
||
|
||
prompts_file = self.base_dir / "config" / "prompts.json"
|
||
if prompts_file.exists():
|
||
with open(prompts_file) as f:
|
||
return json.load(f)
|
||
return {}
|
||
|
||
def _load_permissions_config(self) -> dict:
|
||
permissions_file = self.base_dir / "config" / "permissions.json"
|
||
if not permissions_file.exists():
|
||
return {}
|
||
with permissions_file.open("r", encoding="utf-8") as handle:
|
||
return json.load(handle)
|
||
|
||
def _init_models(self) -> None:
|
||
try:
|
||
memory_config = self.config.runtime.memory_thresholds or {}
|
||
if memory_config:
|
||
self._memory_policy = MemoryWritePolicy(
|
||
store_threshold=memory_config.get("default_store_weight", 0.8),
|
||
)
|
||
print("Models policy ready")
|
||
except Exception as e:
|
||
print(f"Models init failed: {e}")
|
||
|
||
def load_models_at_startup(self) -> None:
|
||
"""Load all LLM models synchronously. Called from startup hook in executor."""
|
||
import os
|
||
os.chdir(str(self.base_dir / "models"))
|
||
|
||
try:
|
||
print("Loading thinker model...")
|
||
thinker_config = self.config.models.thinker or {}
|
||
if thinker_config.get("path"):
|
||
llm, lock = self._get_or_create_llm("thinker", thinker_config)
|
||
self._thinker = OrchestratorAdapter(llm, system_prompt=self._prompts.get("thinker"), lock=lock)
|
||
print(f"Thinker loaded: {self._thinker} (model: {thinker_config.get("path")})")
|
||
|
||
print("Loading json_compiler model...")
|
||
compiler_config = self.config.models.json_compiler or {}
|
||
if compiler_config.get("path"):
|
||
llm, lock = self._get_or_create_llm("json_compiler", compiler_config)
|
||
self._json_compiler = OrchestratorAdapter(llm, system_prompt=self._prompts.get("json_compiler"), lock=lock)
|
||
print(f"JSON Compiler loaded: {self._json_compiler} (model: {compiler_config.get("path")})")
|
||
|
||
print("Loading coder model...")
|
||
coder_config = self.config.models.coder or {}
|
||
if coder_config.get("path"):
|
||
llm, lock = self._get_or_create_llm("coder", coder_config)
|
||
self._coder = CoderAdapter(llm, system_prompt=self._prompts.get("coder"), lock=lock)
|
||
print(f"Coder loaded: {self._coder} (model: {coder_config.get("path")})")
|
||
|
||
print("Loading critic model...")
|
||
critic_config = self.config.models.critic or {}
|
||
if critic_config.get("path"):
|
||
llm, lock = self._get_or_create_llm("critic", critic_config)
|
||
self._critic = CriticAdapter(llm, system_prompt=self._prompts.get("critic"), lock=lock)
|
||
print(f"Critic loaded: {self._critic} (model: {critic_config.get("path")})")
|
||
|
||
print("Loading sys_util model...")
|
||
sys_util_config = self.config.models.sys_util or {}
|
||
if sys_util_config.get("path"):
|
||
llm, lock = self._get_or_create_llm("sys_util", sys_util_config)
|
||
self._sys_util = OrchestratorAdapter(llm, system_prompt=self._prompts.get("sys_util"), lock=lock)
|
||
print(f"Sys_util loaded: {self._sys_util} (model: {sys_util_config.get("path")})")
|
||
|
||
print("All models loaded successfully")
|
||
|
||
async_thinker = AsyncOrchestratorAdapter(self._thinker) if self._thinker else None
|
||
async_compiler = AsyncOrchestratorAdapter(self._json_compiler) if self._json_compiler else None
|
||
async_coder = AsyncCoderAdapter(self._coder) if self._coder else None
|
||
async_critic = AsyncCriticAdapter(self._critic) if self._critic else None
|
||
async_sys_util = AsyncOrchestratorAdapter(self._sys_util) if self._sys_util else None
|
||
|
||
self.router.set_thinker(async_thinker)
|
||
self.router.set_json_compiler(async_compiler)
|
||
self.router.set_sys_util(async_sys_util)
|
||
self.router.set_tool_registry(self.tool_registry)
|
||
if async_critic:
|
||
self.execution_engine.set_critic(async_critic)
|
||
if async_coder:
|
||
self.execution_engine.set_coder(async_coder)
|
||
|
||
except Exception as e:
|
||
print(f"Failed to load models at startup: {e}")
|
||
raise RuntimeError(f"Model loading failed: {e}") from e
|
||
|
||
def _model_cache_key(self, model_config: dict) -> tuple[object, ...]:
|
||
path = str((self.base_dir / "models" / model_config.get("path", "")).resolve())
|
||
return (
|
||
path,
|
||
model_config.get("backend", "cpu"),
|
||
model_config.get("n_gpu_layers", 0),
|
||
model_config.get("n_ctx", 4096),
|
||
)
|
||
|
||
def _get_or_create_llm(self, model_type: str, model_config: dict):
|
||
key = self._model_cache_key(model_config)
|
||
cached = self._model_cache.get(key)
|
||
if cached:
|
||
print(f"Reusing model instance: {model_config.get('path')} for {model_type}")
|
||
return cached
|
||
|
||
llm = create_adapter(model_type, model_config, self.base_dir / "models")
|
||
lock = RLock()
|
||
cached = (llm, lock)
|
||
self._model_cache[key] = cached
|
||
return cached
|
||
|
||
def _init_memory(self) -> None:
|
||
try:
|
||
emb_config = self.config.models.embeddings or {}
|
||
model_path = self.base_dir / emb_config.get("path", "models/all-MiniLM-L6-v2")
|
||
if not model_path.exists() and not Path(emb_config.get("path", "")).is_absolute():
|
||
model_path = self.base_dir / "models" / emb_config.get("path", "all-MiniLM-L6-v2")
|
||
if not model_path.exists():
|
||
print(f"Memory init skipped: embeddings model not found at {model_path}")
|
||
self._memory_interface = None
|
||
return
|
||
embeddings = EmbeddingsAdapter(
|
||
model_path=model_path,
|
||
embedding_dim=emb_config.get("embedding_dim", 384),
|
||
)
|
||
|
||
store = MemoryStore(
|
||
self.base_dir / "data" / "memory" / "memory.sqlite3"
|
||
)
|
||
vector_index = VectorIndex(
|
||
index_path=self.base_dir / "data" / "memory" / "index.bin",
|
||
embedding_dim=embeddings.embedding_dim,
|
||
)
|
||
|
||
self._memory_interface = MemoryInterface(store, vector_index, embeddings)
|
||
|
||
except Exception as e:
|
||
print(f"Memory init failed: {e}")
|
||
self._memory_interface = None
|
||
|
||
def _create_tool_registry(self) -> ToolRegistry:
|
||
from app.tools.registry import ToolRegistry
|
||
from app.tools.plugins.shell_exec import Tool as ShellExecTool
|
||
from app.tools.plugins.file_read import Tool as FileReadTool
|
||
from app.tools.plugins.file_write import Tool as FileWriteTool
|
||
from app.tools.plugins.memory_tools import Tool as MemoryTool
|
||
from app.tools.discover import ToolDiscovery
|
||
|
||
registry = ToolRegistry()
|
||
|
||
tool_init_map = {
|
||
"shell_exec": lambda m: ShellExecTool(self.tool_sandbox),
|
||
"file_read": lambda m: FileReadTool(self.tool_sandbox),
|
||
"file_write": lambda m: FileWriteTool(self.tool_sandbox),
|
||
"memory": lambda m: MemoryTool(self._memory_interface),
|
||
}
|
||
|
||
discovery = ToolDiscovery()
|
||
discovered = discovery.discover()
|
||
|
||
for name, data in discovered.items():
|
||
init_fn = tool_init_map.get(name)
|
||
if init_fn:
|
||
tool = init_fn(data.get("manifest", {}))
|
||
registry.register(tool)
|
||
registry._schemas[name] = {
|
||
"description": data.get("manifest", {}).get("description", ""),
|
||
"args_schema": data.get("manifest", {}).get("args_schema", {}),
|
||
"requires_permission": data.get("manifest", {}).get("requires_permission", False),
|
||
}
|
||
print(f"Registered tool: {name}")
|
||
else:
|
||
print(f"No init mapping for tool: {name} - skipping")
|
||
|
||
return registry
|
||
|
||
@property
|
||
def orchestrator(self) -> OrchestratorAdapter | None:
|
||
return self._orchestrator
|
||
|
||
@property
|
||
def coder(self) -> CoderAdapter | None:
|
||
return self._coder
|
||
|
||
@property
|
||
def critic(self) -> CriticAdapter | None:
|
||
return self._critic
|
||
|
||
@property
|
||
def memory_interface(self) -> MemoryInterface | None:
|
||
return self._memory_interface
|
||
|
||
def _ensure_orchestrator(self) -> OrchestratorAdapter | None:
|
||
if self._orchestrator is not None:
|
||
return self._orchestrator
|
||
try:
|
||
orch_config = self.config.models.orchestrator or {}
|
||
if orch_config.get("path"):
|
||
llm, lock = self._get_or_create_llm("orchestrator", orch_config)
|
||
self._orchestrator = OrchestratorAdapter(llm, lock=lock)
|
||
except Exception as e:
|
||
print(f"Orchestrator load failed: {e}")
|
||
return self._orchestrator
|
||
|
||
def _ensure_critic(self) -> CriticAdapter | None:
|
||
if self._critic is not None:
|
||
return self._critic
|
||
try:
|
||
critic_config = self.config.models.critic or {}
|
||
if critic_config.get("path"):
|
||
llm, lock = self._get_or_create_llm("critic", critic_config)
|
||
self._critic = CriticAdapter(llm, lock=lock)
|
||
except Exception as e:
|
||
print(f"Critic load failed: {e}")
|
||
return self._critic
|
||
|
||
def handle_task(self, task: UserTask) -> dict[str, object]:
|
||
return self.runtime_loop.run_task(task)
|
||
|
||
def resolve_permission(self, task_id: str, decision: str) -> dict[str, object]:
|
||
return self.runtime_loop.resolve_permission(
|
||
task_id=task_id, decision=decision
|
||
)
|
||
|
||
def resolve_secret(self, task_id: str, secret: str) -> dict[str, object]:
|
||
return self.runtime_loop.resolve_secret(
|
||
task_id=task_id, secret=secret
|
||
)
|
||
|
||
def resolve_password(self, task_id: str, password: str) -> dict[str, object]:
|
||
return self.runtime_loop.resolve_password(
|
||
task_id=task_id, password=password
|
||
)
|
||
|
||
def handle_critic_feedback(
|
||
self,
|
||
feedback: str,
|
||
task_id: str | None = None,
|
||
session_id: str | None = None,
|
||
feedback_type: str | None = None,
|
||
severity: str | None = None,
|
||
correction: str | None = None,
|
||
remember: bool = True,
|
||
retry: bool = False,
|
||
assistant_answer: str | None = None,
|
||
correctness_override: float | None = None,
|
||
usefulness_override: float | None = None,
|
||
safety_override: float | None = None,
|
||
) -> dict[str, object]:
|
||
target_task_id = task_id
|
||
target_session_id = session_id
|
||
|
||
if not target_session_id and not target_task_id:
|
||
return {
|
||
"status": "error",
|
||
"message": "Either task_id or session_id must be provided",
|
||
}
|
||
|
||
state = self.task_state_store.get_task(target_task_id) if target_task_id else None
|
||
if not target_session_id and state:
|
||
target_session_id = state.get("session_id")
|
||
|
||
if not target_task_id and target_session_id:
|
||
recent_tasks = self.task_state_store.get_session_tasks(target_session_id, limit=1)
|
||
if recent_tasks:
|
||
target_task_id = recent_tasks[0]["task_id"]
|
||
|
||
min_weight = 0.3
|
||
max_weight = 0.95
|
||
user_weight = 0.9
|
||
|
||
final_weight = max(min_weight, min(max_weight, user_weight))
|
||
|
||
task_input = state.get("task_input") if state else None
|
||
last_directive = state.get("last_directive") if state else None
|
||
feedback_type = feedback_type or "other"
|
||
severity = severity or "major"
|
||
|
||
lesson = self._build_feedback_lesson(
|
||
feedback_type=feedback_type,
|
||
severity=severity,
|
||
feedback=feedback,
|
||
correction=correction,
|
||
task_input=task_input,
|
||
)
|
||
|
||
metadata = {
|
||
"feedback_text": feedback,
|
||
"feedback_type": feedback_type,
|
||
"severity": severity,
|
||
"correction": correction,
|
||
"assistant_answer": assistant_answer,
|
||
"task_input": task_input,
|
||
"last_directive": last_directive,
|
||
"overrides": {
|
||
"correctness": correctness_override,
|
||
"usefulness": usefulness_override,
|
||
"safety": safety_override,
|
||
},
|
||
"source": "user",
|
||
}
|
||
|
||
feedback_text = lesson
|
||
if correctness_override is not None:
|
||
feedback_text += f" | Correctness corrected to: {correctness_override}"
|
||
if usefulness_override is not None:
|
||
feedback_text += f" | Usefulness corrected to: {usefulness_override}"
|
||
if safety_override is not None:
|
||
feedback_text += f" | Safety corrected to: {safety_override}"
|
||
|
||
retry_result = None
|
||
stored = False
|
||
store_error = None
|
||
try:
|
||
if remember and self._memory_interface:
|
||
self._memory_interface.insert(
|
||
text=feedback_text,
|
||
kind="critique",
|
||
source="user",
|
||
task_id=target_task_id,
|
||
session_id=target_session_id,
|
||
weight=final_weight,
|
||
metadata=metadata,
|
||
)
|
||
stored = True
|
||
elif remember and not self._memory_interface:
|
||
store_error = "Memory not available"
|
||
except Exception as e:
|
||
store_error = str(e)
|
||
|
||
if retry and task_input:
|
||
retry_input = self._build_retry_input(
|
||
task_input=task_input,
|
||
feedback=feedback,
|
||
feedback_type=feedback_type,
|
||
correction=correction,
|
||
)
|
||
retry_task = UserTask(
|
||
session_id=target_session_id or "feedback-retry",
|
||
input=retry_input,
|
||
context={
|
||
"feedback_retry": True,
|
||
"original_task_id": target_task_id,
|
||
"feedback_type": feedback_type,
|
||
"severity": severity,
|
||
"correction": correction,
|
||
},
|
||
)
|
||
retry_result = self.handle_task(retry_task)
|
||
|
||
status = "ok" if stored or not remember else "error"
|
||
return {
|
||
"status": status,
|
||
"message": "Feedback saved" if stored else (store_error or "Feedback accepted"),
|
||
"stored": stored,
|
||
"task_id": target_task_id,
|
||
"session_id": target_session_id,
|
||
"lesson": lesson,
|
||
"retry_result": retry_result,
|
||
}
|
||
|
||
def _build_feedback_lesson(
|
||
self,
|
||
feedback_type: str,
|
||
severity: str,
|
||
feedback: str,
|
||
correction: str | None,
|
||
task_input: str | None,
|
||
) -> str:
|
||
parts = [
|
||
"User critique lesson.",
|
||
f"Error type: {feedback_type}.",
|
||
f"Severity: {severity}.",
|
||
]
|
||
if task_input:
|
||
parts.append(f"Original task: {task_input}")
|
||
if feedback:
|
||
parts.append(f"What was wrong: {feedback}")
|
||
if correction:
|
||
parts.append(f"Preferred correction: {correction}")
|
||
return " | ".join(parts)
|
||
|
||
def _build_retry_input(
|
||
self,
|
||
task_input: str,
|
||
feedback: str,
|
||
feedback_type: str,
|
||
correction: str | None,
|
||
) -> str:
|
||
retry_input = (
|
||
f"Повтори задачу с учетом обратной связи.\n"
|
||
f"Исходная задача: {task_input}\n"
|
||
f"Тип ошибки: {feedback_type}\n"
|
||
f"Что было неверно: {feedback}\n"
|
||
)
|
||
if correction:
|
||
retry_input += f"Как должно быть: {correction}\n"
|
||
return retry_input
|