| """Base AI Agent — common functionality for all agents. |
| |
| Each agent: |
| - Has a name, role, and specialization |
| - Runs in a background thread (persistent) |
| - Picks up goals from the goal memory |
| - Uses the LLM (via harness) to generate responses |
| - Reports progress and results |
| - Can be paused/resumed |
| - Tracks stats (goals completed, steps executed, errors) |
| """ |
|
|
| from __future__ import annotations |
|
|
| import logging |
| import threading |
| import time |
| from dataclasses import dataclass, field |
| from typing import Any, Callable |
|
|
| from ..memory.goal_memory import GoalMemory, Goal |
| from ..memory.persistent import PersistentMemory |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| @dataclass |
| class AgentStats: |
| """Stats for a single agent.""" |
| goals_assigned: int = 0 |
| goals_completed: int = 0 |
| goals_failed: int = 0 |
| steps_executed: int = 0 |
| errors: int = 0 |
| uptime_s: float = 0.0 |
| last_active: float = 0.0 |
|
|
|
|
| class BaseAgent: |
| """Base class for all AI agents. |
| |
| Subclasses implement `process_goal()` which is called when a goal |
| is assigned to this agent. |
| """ |
|
|
| def __init__(self, name: str, role: str, description: str, |
| goal_memory: GoalMemory, |
| persistent_memory: PersistentMemory | None = None, |
| generate_fn: Callable[[str], str] | None = None, |
| poll_interval_s: float = 5.0) -> None: |
| self.name = name |
| self.role = role |
| self.description = description |
| self.goal_memory = goal_memory |
| self.persistent_memory = persistent_memory |
| self._generate_fn = generate_fn |
| self._poll_interval = poll_interval_s |
|
|
| self._running = False |
| self._paused = False |
| self._thread: threading.Thread | None = None |
| self._current_goal: Goal | None = None |
| self._stats = AgentStats() |
| self._start_time = 0.0 |
|
|
| def set_generate_fn(self, fn: Callable[[str], str]) -> None: |
| """Set the function used to generate LLM responses.""" |
| self._generate_fn = fn |
|
|
| def start(self) -> None: |
| """Start the agent in a background thread.""" |
| if self._running: |
| return |
| self._running = True |
| self._paused = False |
| self._start_time = time.time() |
| self._thread = threading.Thread(target=self._run_loop, daemon=True, name=f"agent-{self.name}") |
| self._thread.start() |
| logger.info("Agent '%s' started (%s)", self.name, self.role) |
|
|
| def stop(self) -> None: |
| """Stop the agent.""" |
| self._running = False |
| if self._thread: |
| self._thread.join(timeout=10) |
| logger.info("Agent '%s' stopped", self.name) |
|
|
| def pause(self) -> None: |
| """Pause the agent (doesn't pick up new goals).""" |
| self._paused = True |
| logger.info("Agent '%s' paused", self.name) |
|
|
| def resume(self) -> None: |
| """Resume the agent.""" |
| self._paused = False |
| logger.info("Agent '%s' resumed", self.name) |
|
|
| def _run_loop(self) -> None: |
| """Main agent loop — continuously picks up and processes goals.""" |
| while self._running: |
| try: |
| if self._paused: |
| time.sleep(self._poll_interval) |
| continue |
|
|
| |
| goals = self.goal_memory.get_goals_for_agent(self.name) |
| active = [g for g in goals if g.status in ("in_progress", "planning")] |
|
|
| if not active: |
| |
| pending = self.goal_memory.list_goals(status="pending") |
| for g in pending: |
| if self._can_handle(g): |
| self.goal_memory.assign_agent(g.id, self.name) |
| self._stats.goals_assigned += 1 |
| active = [g] |
| break |
|
|
| if not active: |
| time.sleep(self._poll_interval) |
| continue |
|
|
| |
| goal = active[0] |
| self._current_goal = goal |
| self._stats.last_active = time.time() |
|
|
| result = self.process_goal(goal) |
|
|
| if result.get("success"): |
| self._stats.steps_executed += 1 |
| |
| self.goal_memory.execute_step( |
| goal.id, result.get("output", ""), success=True |
| ) |
| else: |
| self._stats.errors += 1 |
| self.goal_memory.execute_step( |
| goal.id, result.get("error", "Unknown error"), success=False |
| ) |
|
|
| |
| updated = self.goal_memory.get_goal(goal.id) |
| if updated and updated.status == "completed": |
| self._stats.goals_completed += 1 |
| if self.persistent_memory: |
| self.persistent_memory.add_episodic( |
| "event", f"Agent {self.name} completed goal: {goal.title}", |
| importance=0.8, tags=["goal", "completed", self.name] |
| ) |
| logger.info("Agent '%s' completed goal: %s", self.name, goal.title) |
|
|
| self._current_goal = None |
|
|
| except Exception as e: |
| logger.error("Agent '%s' error: %s", self.name, e) |
| self._stats.errors += 1 |
| time.sleep(self._poll_interval) |
|
|
| self._stats.uptime_s = time.time() - self._start_time |
|
|
| def _can_handle(self, goal: Goal) -> bool: |
| """Check if this agent can handle a goal. Override in subclasses.""" |
| return True |
|
|
| def process_goal(self, goal: Goal) -> dict[str, Any]: |
| """Process a goal step. Must be implemented by subclasses. |
| |
| Returns: {"success": bool, "output": str, "error": str} |
| """ |
| raise NotImplementedError |
|
|
| def _generate(self, prompt: str) -> str: |
| """Generate a response using the LLM.""" |
| if self._generate_fn: |
| return self._generate_fn(prompt) |
| return "(no generation function available)" |
|
|
| def get_status(self) -> dict[str, Any]: |
| """Get agent status.""" |
| return { |
| "name": self.name, |
| "role": self.role, |
| "description": self.description, |
| "running": self._running, |
| "paused": self._paused, |
| "current_goal": self._current_goal.title if self._current_goal else None, |
| "stats": { |
| "goals_assigned": self._stats.goals_assigned, |
| "goals_completed": self._stats.goals_completed, |
| "goals_failed": self._stats.goals_failed, |
| "steps_executed": self._stats.steps_executed, |
| "errors": self._stats.errors, |
| "uptime_s": round(self._stats.uptime_s, 1), |
| "last_active": self._stats.last_active, |
| }, |
| } |
|
|
