File size: 7,113 Bytes
0e3d4b8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 | """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
# Find a goal assigned to this agent
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:
# Try to pick up an unassigned goal that matches our role
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
# Process the first active goal
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
# Advance the goal
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
)
# Check if goal is completed
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,
},
}
|