File size: 6,191 Bytes
296a506 | 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 | """
Frox AI — Memory Tool
A local-first, self-contained memory store: JSON-persisted facts per
user, retrieved by cosine similarity over embeddings from Morph's own
model (ctx.engine.embed() — see inference/engine/morph_engine.py).
No external vector DB required, so this works standalone in this repo.
The production backend architecture document specs a Postgres +
Qdrant version of this same idea for multi-instance deployments
(Section 5: Memory System) — this module is the reference
implementation / local-dev equivalent, not a competing design.
"""
from __future__ import annotations
import json
import math
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional
from tools.registry import tool, ToolContext
@dataclass
class MemoryItem:
id: str
user_id: str
text: str
category: str
embedding: List[float]
created_at: float
def to_dict(self, include_embedding: bool = False) -> dict:
d = {"id": self.id, "user_id": self.user_id, "text": self.text,
"category": self.category, "created_at": self.created_at}
if include_embedding:
d["embedding"] = self.embedding
return d
def _cosine(a: List[float], b: List[float]) -> float:
dot = sum(x * y for x, y in zip(a, b))
norm_a = math.sqrt(sum(x * x for x in a))
norm_b = math.sqrt(sum(y * y for y in b))
if norm_a == 0 or norm_b == 0:
return 0.0
return dot / (norm_a * norm_b)
class MemoryStore:
"""
Per-process memory store, JSON-persisted to disk. Embeddings come
from whatever MorphInferenceEngine is passed in at save time — the
store itself has no model dependency, so it's easy to unit test
with fake embeddings.
"""
def __init__(self, path: Optional[str] = None):
self.path = Path(path) if path else None
self._items: Dict[str, MemoryItem] = {}
if self.path and self.path.exists():
self._load()
def _load(self):
try:
data = json.loads(self.path.read_text())
for item in data:
self._items[item["id"]] = MemoryItem(**item)
except (json.JSONDecodeError, TypeError, KeyError):
pass # start fresh rather than crash on a corrupted file
def _persist(self):
if self.path is None:
return
self.path.parent.mkdir(parents=True, exist_ok=True)
self.path.write_text(json.dumps(
[i.to_dict(include_embedding=True) for i in self._items.values()], indent=2
))
def add(self, user_id: str, text: str, embedding: List[float], category: str = "fact") -> MemoryItem:
item = MemoryItem(
id=str(uuid.uuid4()), user_id=user_id, text=text,
category=category, embedding=embedding, created_at=time.time(),
)
self._items[item.id] = item
self._persist()
return item
def search(self, user_id: str, query_embedding: List[float], k: int = 5) -> List[MemoryItem]:
candidates = [i for i in self._items.values() if i.user_id == user_id]
scored = sorted(candidates, key=lambda i: _cosine(i.embedding, query_embedding), reverse=True)
return scored[:k]
def list_all(self, user_id: str) -> List[MemoryItem]:
return [i for i in self._items.values() if i.user_id == user_id]
def delete(self, memory_id: str) -> bool:
if memory_id in self._items:
del self._items[memory_id]
self._persist()
return True
return False
def clear(self, user_id: str):
to_delete = [i.id for i in self._items.values() if i.user_id == user_id]
for mid in to_delete:
del self._items[mid]
self._persist()
@tool(
name="memory_save",
description="Save a fact about the user for later conversations",
timeout=10.0,
)
def memory_save(ctx: ToolContext, text: str, category: str = "fact") -> dict:
"""
Args:
text: The fact to remember, e.g. "User prefers Python over JavaScript".
category: One of fact | preference | project | person (freeform, for organization).
Plain `def`, not `async def`: engine.embed() is synchronous and
GPU-bound — thread-offloaded by the registry.
"""
if ctx.memory_store is None:
raise RuntimeError("No memory_store configured in ToolContext")
if ctx.engine is None:
raise RuntimeError("No engine configured in ToolContext (needed to embed the memory)")
embedding = ctx.engine.embed(text)
item = ctx.memory_store.add(ctx.user_id, text, embedding, category)
return {"saved": True, "memory_id": item.id, "text": text}
@tool(
name="memory_search",
description="Search the user's saved memories for relevant facts",
timeout=10.0,
)
def memory_search(ctx: ToolContext, query: str, k: int = 5) -> dict:
"""
Args:
query: What to search for, e.g. "what programming language does the user like".
k: Max number of memories to return.
Plain `def`, not `async def`: engine.embed() is synchronous and
GPU-bound — thread-offloaded by the registry.
"""
if ctx.memory_store is None:
raise RuntimeError("No memory_store configured in ToolContext")
if ctx.engine is None:
raise RuntimeError("No engine configured in ToolContext (needed to embed the query)")
query_embedding = ctx.engine.embed(query)
results = ctx.memory_store.search(ctx.user_id, query_embedding, k=k)
return {"query": query, "memories": [r.to_dict() for r in results]}
@tool(
name="memory_forget",
description="Delete a previously saved memory by its id",
timeout=5.0,
)
def memory_forget(ctx: ToolContext, memory_id: str) -> dict:
"""
Args:
memory_id: The id returned by memory_save.
Plain `def` for consistency with memory_save/memory_search in this
module — no async I/O here either.
"""
if ctx.memory_store is None:
raise RuntimeError("No memory_store configured in ToolContext")
deleted = ctx.memory_store.delete(memory_id)
return {"deleted": deleted, "memory_id": memory_id}
|