File size: 5,300 Bytes
71b4454 | 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 | """
Agent Memory Store
SQLite-backed persistent memory for agents: episodic events (goals, actions,
observations, final answers) and key/value semantic facts, scoped per agent_id
and optionally per task_id. Thread-safe; safe to share across agents.
"""
from __future__ import annotations
import json
import sqlite3
import threading
import time
import uuid
from contextlib import contextmanager
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional
DEFAULT_DB_PATH = Path(__file__).resolve().parents[2] / "data" / "agent_memory.db"
_SCHEMA = """
CREATE TABLE IF NOT EXISTS episodic_memory (
id TEXT PRIMARY KEY,
agent_id TEXT NOT NULL,
task_id TEXT,
role TEXT NOT NULL,
content TEXT NOT NULL,
metadata TEXT,
created_at REAL NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_episodic_agent ON episodic_memory(agent_id, created_at);
CREATE INDEX IF NOT EXISTS idx_episodic_task ON episodic_memory(task_id);
CREATE TABLE IF NOT EXISTS semantic_memory (
agent_id TEXT NOT NULL,
key TEXT NOT NULL,
value TEXT NOT NULL,
updated_at REAL NOT NULL,
PRIMARY KEY (agent_id, key)
);
"""
@dataclass(slots=True)
class MemoryEvent:
id: str
agent_id: str
role: str
content: str
task_id: Optional[str] = None
metadata: dict[str, Any] = field(default_factory=dict)
created_at: float = field(default_factory=time.time)
class AgentMemoryStore:
"""Shared SQLite-backed memory store used by every agent in the runtime."""
def __init__(self, db_path: Path | str = DEFAULT_DB_PATH):
self.db_path = Path(db_path)
self.db_path.parent.mkdir(parents=True, exist_ok=True)
self._lock = threading.Lock()
with self._connect() as conn:
conn.executescript(_SCHEMA)
@contextmanager
def _connect(self):
conn = sqlite3.connect(self.db_path, timeout=30)
try:
yield conn
conn.commit()
finally:
conn.close()
def record_event(
self,
agent_id: str,
role: str,
content: str,
task_id: Optional[str] = None,
metadata: Optional[dict[str, Any]] = None,
) -> MemoryEvent:
event = MemoryEvent(
id=str(uuid.uuid4()),
agent_id=agent_id,
role=role,
content=content,
task_id=task_id,
metadata=metadata or {},
)
with self._lock, self._connect() as conn:
conn.execute(
"INSERT INTO episodic_memory (id, agent_id, task_id, role, content, metadata, created_at) "
"VALUES (?, ?, ?, ?, ?, ?, ?)",
(
event.id,
event.agent_id,
event.task_id,
event.role,
event.content,
json.dumps(event.metadata),
event.created_at,
),
)
return event
def recent_events(
self, agent_id: str, limit: int = 20, task_id: Optional[str] = None
) -> list[MemoryEvent]:
query = (
"SELECT id, agent_id, task_id, role, content, metadata, created_at "
"FROM episodic_memory WHERE agent_id = ?"
)
params: list[Any] = [agent_id]
if task_id:
query += " AND task_id = ?"
params.append(task_id)
query += " ORDER BY created_at DESC LIMIT ?"
params.append(limit)
with self._lock, self._connect() as conn:
rows = conn.execute(query, params).fetchall()
events = [
MemoryEvent(
id=r[0], agent_id=r[1], task_id=r[2], role=r[3], content=r[4],
metadata=json.loads(r[5]) if r[5] else {}, created_at=r[6],
)
for r in rows
]
events.reverse()
return events
def set_fact(self, agent_id: str, key: str, value: Any) -> None:
with self._lock, self._connect() as conn:
conn.execute(
"INSERT INTO semantic_memory (agent_id, key, value, updated_at) VALUES (?, ?, ?, ?) "
"ON CONFLICT(agent_id, key) DO UPDATE SET value = excluded.value, updated_at = excluded.updated_at",
(agent_id, key, json.dumps(value), time.time()),
)
def get_fact(self, agent_id: str, key: str, default: Any = None) -> Any:
with self._lock, self._connect() as conn:
row = conn.execute(
"SELECT value FROM semantic_memory WHERE agent_id = ? AND key = ?",
(agent_id, key),
).fetchone()
return json.loads(row[0]) if row else default
def all_facts(self, agent_id: str) -> dict[str, Any]:
with self._lock, self._connect() as conn:
rows = conn.execute(
"SELECT key, value FROM semantic_memory WHERE agent_id = ?", (agent_id,)
).fetchall()
return {k: json.loads(v) for k, v in rows}
def clear_agent(self, agent_id: str) -> None:
with self._lock, self._connect() as conn:
conn.execute("DELETE FROM episodic_memory WHERE agent_id = ?", (agent_id,))
conn.execute("DELETE FROM semantic_memory WHERE agent_id = ?", (agent_id,))
|