465 lines
19 KiB
Python
465 lines
19 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,
|
|
correctness_override: float | None = None,
|
|
usefulness_override: float | None = None,
|
|
safety_override: float | None = None,
|
|
) -> dict[str, object]:
|
|
if not self._memory_interface:
|
|
return {"status": "error", "message": "Memory not available"}
|
|
|
|
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",
|
|
}
|
|
|
|
if not target_session_id and target_task_id:
|
|
state = self.task_state_store.get_task(target_task_id)
|
|
if 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))
|
|
|
|
metadata = {
|
|
"feedback_text": feedback,
|
|
"overrides": {
|
|
"correctness": correctness_override,
|
|
"usefulness": usefulness_override,
|
|
"safety": safety_override,
|
|
},
|
|
"source": "user",
|
|
}
|
|
|
|
feedback_text = f"User feedback: {feedback}"
|
|
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}"
|
|
|
|
try:
|
|
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,
|
|
)
|
|
return {
|
|
"status": "ok",
|
|
"message": "Feedback saved",
|
|
"task_id": target_task_id,
|
|
"session_id": target_session_id,
|
|
}
|
|
except Exception as e:
|
|
return {"status": "error", "message": str(e)}
|