ducklm/app/runtime/runtime_controller.py

644 lines
26 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

from __future__ import annotations
import json
from concurrent.futures import Future, ThreadPoolExecutor
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.command_analyzer import CommandAnalyzer
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.recall import MemoryRecallService
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._background_executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="ducklm-task")
self._background_tasks: dict[str, Future[dict[str, object]]] = {}
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,
command_timeout_ms=runtime_config.shell_command_timeout_ms,
idle_timeout_ms=runtime_config.shell_idle_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.command_analyzer = CommandAnalyzer(self.permission_service)
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,
recovery_limit=runtime_config.tool_retry_limit,
critic_retry_limit=runtime_config.critic_retry_limit,
command_analyzer=self.command_analyzer,
)
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)
# Create MemoryRecallService using the configured model (default: sys_util)
# Reuses already-loaded async adapter — no duplicate model loading
recall_model_name = self.config.runtime.recall_model
recall_async_model = {
"sys_util": async_sys_util,
"thinker": async_thinker,
"json_compiler": async_compiler,
"critic": async_critic,
"coder": async_coder,
}.get(recall_model_name, async_sys_util)
self._recall_service = MemoryRecallService(
memory_interface=self._memory_interface,
recall_model=recall_async_model,
)
self.runtime_loop.set_recall_service(self._recall_service)
print(f"MemoryRecallService initialized with model: {recall_model_name}")
# Set memory policy in runtime loop
self.runtime_loop.set_memory_policy(self._memory_policy)
print(f"MemoryWritePolicy set: {self._memory_policy is not None}")
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 submit_task(self, task: UserTask) -> dict[str, object]:
self._background_tasks[task.task_id] = self._background_executor.submit(
self.handle_task,
task,
)
return {"task_id": task.task_id, "status": "accepted"}
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 submit_permission_resolution(self, task_id: str, decision: str) -> dict[str, object]:
if not self.task_state_store.get_task(task_id):
return {"task_id": task_id, "status": "failed", "result": {"error": "Unknown task_id"}}
self._background_tasks[task_id] = self._background_executor.submit(
self.resolve_permission,
task_id,
decision,
)
return {"task_id": task_id, "status": "accepted"}
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 submit_secret_resolution(self, task_id: str, secret: str) -> dict[str, object]:
if not self.task_state_store.get_task(task_id):
return {"task_id": task_id, "status": "failed", "result": {"error": "Unknown task_id"}}
self._background_tasks[task_id] = self._background_executor.submit(
self.resolve_secret,
task_id,
secret,
)
return {"task_id": task_id, "status": "accepted"}
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 resolve_review(self, task_id: str, decision: str, correction: str | None = None) -> dict[str, object]:
return self.runtime_loop.resolve_review(
task_id=task_id,
decision=decision,
correction=correction,
)
def submit_review_resolution(self, task_id: str, decision: str, correction: str | None = None) -> dict[str, object]:
if not self.task_state_store.get_task(task_id):
return {"task_id": task_id, "status": "failed", "result": {"error": "Unknown task_id"}}
self._background_tasks[task_id] = self._background_executor.submit(
self.resolve_review,
task_id,
decision,
correction,
)
return {"task_id": task_id, "status": "accepted"}
def submit_password_resolution(self, task_id: str, password: str) -> dict[str, object]:
if not self.task_state_store.get_task(task_id):
return {"task_id": task_id, "status": "failed", "result": {"error": "Unknown task_id"}}
self._background_tasks[task_id] = self._background_executor.submit(
self.resolve_password,
task_id,
password,
)
return {"task_id": task_id, "status": "accepted"}
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