Files

671 lines
24 KiB
Python
Executable File
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
from pathlib import Path
from typing import Any, Dict, List, Optional
from pydantic import BaseModel
import sqlite3
import json
import math
import os
import time
def _cosine(a: List[float], b: List[float]) -> float:
"""Cosine similarity between two equal-length vectors. 0.0 on mismatch."""
if not a or not b or len(a) != len(b):
return 0.0
dot = sum(x * y for x, y in zip(a, b))
na = math.sqrt(sum(x * x for x in a))
nb = math.sqrt(sum(y * y for y in b))
if na == 0.0 or nb == 0.0:
return 0.0
return dot / (na * nb)
# -----------------------------
# Models
# -----------------------------
class MemoryItem(BaseModel):
id: str
section: str = "General"
text: str
tags: List[str] = []
position: int = 0
class MessageItem(BaseModel):
role: str # "user" or "assistant"
content: str
timestamp: float
model: Optional[str] = None
tokens: Optional[int] = None
class ConversationItem(BaseModel):
id: str
messages: List[MessageItem] = []
created_at: float
updated_at: float
title: Optional[str] = None
@property
def preview(self) -> str:
for msg in self.messages:
if msg.role == "user":
return msg.content[:80]
return "Empty conversation"
@property
def timestamp(self) -> float:
return self.created_at
# -----------------------------
# Persistent Store
# -----------------------------
class PersistentMemoryStore:
def __init__(self, db_path: Path):
self.db_path = db_path
os.makedirs(self.db_path.parent, exist_ok=True)
self._ensure_tables()
self._cache: Dict[str, MemoryItem] = self._load_all_memory()
# -----------------------------
# Internal helpers
# -----------------------------
def _connect(self):
conn = sqlite3.connect(self.db_path)
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA journal_mode=WAL;")
return conn
def _ensure_tables(self):
conn = self._connect()
cur = conn.cursor()
cur.execute("""
CREATE TABLE IF NOT EXISTS memory (
id TEXT PRIMARY KEY,
section TEXT NOT NULL DEFAULT 'General',
text TEXT NOT NULL,
tags TEXT
)
""")
# Migrate: add section column if it doesn't exist yet
try:
cur.execute("ALTER TABLE memory ADD COLUMN section TEXT NOT NULL DEFAULT 'General'")
except Exception:
pass
# Migrate: add position column for stable ordering
try:
cur.execute("ALTER TABLE memory ADD COLUMN position INTEGER NOT NULL DEFAULT 0")
except Exception:
pass
# Backfill positions for rows added before this column existed
cur.execute("SELECT COUNT(*) FROM memory WHERE position > 0")
if cur.fetchone()[0] == 0:
cur.execute("UPDATE memory SET position = rowid")
cur.execute("""
CREATE TABLE IF NOT EXISTS conversations (
id TEXT PRIMARY KEY,
created_at REAL NOT NULL,
updated_at REAL NOT NULL,
title TEXT
)
""")
# Migrate: add title column if it doesn't exist yet
try:
cur.execute("ALTER TABLE conversations ADD COLUMN title TEXT")
except Exception:
pass
cur.execute("""
CREATE TABLE IF NOT EXISTS messages (
id INTEGER PRIMARY KEY AUTOINCREMENT,
conversation_id TEXT NOT NULL,
role TEXT NOT NULL,
content TEXT NOT NULL,
timestamp REAL NOT NULL,
model TEXT,
tokens INTEGER,
FOREIGN KEY (conversation_id) REFERENCES conversations(id)
)
""")
for col, typedef in (("model", "TEXT"), ("tokens", "INTEGER")):
try:
cur.execute(f"ALTER TABLE messages ADD COLUMN {col} {typedef}")
except Exception:
pass
cur.execute("""
CREATE INDEX IF NOT EXISTS idx_messages_conversation_id
ON messages (conversation_id)
""")
# Semantic recall: one embedding vector per message, stored as JSON.
# Backfilled lazily by semantic_search_conversations so existing history
# gets indexed on first search.
cur.execute("""
CREATE TABLE IF NOT EXISTS message_vectors (
message_id INTEGER PRIMARY KEY,
embedding TEXT NOT NULL,
FOREIGN KEY (message_id) REFERENCES messages(id)
)
""")
cur.execute("""
CREATE TABLE IF NOT EXISTS settings (
key TEXT PRIMARY KEY,
value TEXT NOT NULL
)
""")
cur.execute(
"DELETE FROM settings WHERE key IN ('anthropic_api_key', 'escalation_model')"
)
conn.commit()
conn.close()
# -----------------------------
# Loaders
# -----------------------------
def _load_all_memory(self) -> Dict[str, MemoryItem]:
conn = self._connect()
cur = conn.cursor()
cur.execute("SELECT id, section, text, tags, position FROM memory ORDER BY position ASC, rowid ASC")
rows = cur.fetchall()
conn.close()
cache = {}
for row in rows:
try:
tags = json.loads(row["tags"]) if row["tags"] else []
except Exception:
tags = []
cache[row["id"]] = MemoryItem(
id=row["id"],
section=row["section"] or "General",
text=row["text"],
tags=tags,
position=row["position"] or 0,
)
return cache
# -----------------------------
# Memory API
# -----------------------------
def add(self, item: MemoryItem):
if not item.position:
conn = self._connect()
try:
row = conn.execute(
"SELECT position FROM memory WHERE id = ?", (item.id,)
).fetchone()
if row and row["position"]:
item.position = row["position"]
else:
row = conn.execute("SELECT MAX(position) AS max_position FROM memory").fetchone()
item.position = (row["max_position"] or 0) + 1
finally:
conn.close()
self._cache[item.id] = item
conn = self._connect()
try:
cur = conn.cursor()
cur.execute(
"INSERT OR REPLACE INTO memory (id, section, text, tags, position) VALUES (?, ?, ?, ?, ?)",
(item.id, item.section or "General", item.text, json.dumps(item.tags), item.position)
)
conn.commit()
finally:
conn.close()
def update(self, item: MemoryItem):
existing = self.get(item.id)
if existing:
item.position = existing.position
self.add(item)
def delete(self, item_id: str):
self._cache.pop(item_id, None)
conn = self._connect()
try:
cur = conn.cursor()
cur.execute("DELETE FROM memory WHERE id = ?", (item_id,))
conn.commit()
finally:
conn.close()
def get(self, item_id: str) -> Optional[MemoryItem]:
conn = self._connect()
try:
row = conn.execute(
"SELECT id, section, text, tags, position FROM memory WHERE id = ?",
(item_id,),
).fetchone()
finally:
conn.close()
if not row:
self._cache.pop(item_id, None)
return None
try:
tags = json.loads(row["tags"]) if row["tags"] else []
except Exception:
tags = []
item = MemoryItem(
id=row["id"], section=row["section"] or "General", text=row["text"],
tags=tags, position=row["position"] or 0,
)
self._cache[item_id] = item
return item
def all(self) -> List[MemoryItem]:
# Always read from DB — the memory service and backend run in separate processes
# with separate caches, so the cache can be stale for facts extracted by the
# memory service after this process started.
conn = self._connect()
cur = conn.cursor()
cur.execute("SELECT id, section, text, tags, position FROM memory ORDER BY position ASC, rowid ASC")
rows = cur.fetchall()
conn.close()
items = []
for row in rows:
try:
tags = json.loads(row["tags"]) if row["tags"] else []
except Exception:
tags = []
items.append(MemoryItem(
id=row["id"],
section=row["section"] or "General",
text=row["text"],
tags=tags,
position=row["position"] or 0,
))
return items
def reorder_section(self, section: str, ordered_ids: List[str]) -> bool:
"""Rewrite the order of items in a section using its existing position pool.
ordered_ids must contain exactly the ids currently in the section."""
section_norm = section or "General"
in_section = [i for i in self.all() if (i.section or "General") == section_norm]
if len(in_section) != len(ordered_ids):
return False
by_id = {i.id: i for i in in_section}
siblings = []
for id_ in ordered_ids:
if id_ not in by_id:
return False
siblings.append(by_id[id_])
positions = sorted([i.position for i in in_section])
conn = self._connect()
try:
cur = conn.cursor()
for s, new_pos in zip(siblings, positions):
s.position = new_pos
cur.execute("UPDATE memory SET position = ? WHERE id = ?", (new_pos, s.id))
conn.commit()
finally:
conn.close()
return True
# -----------------------------
# Conversation API
# -----------------------------
def create_conversation(self, conversation_id: str) -> ConversationItem:
now = time.time()
conn = self._connect()
try:
cur = conn.cursor()
cur.execute(
"INSERT OR IGNORE INTO conversations (id, created_at, updated_at) VALUES (?, ?, ?)",
(conversation_id, now, now)
)
conn.commit()
finally:
conn.close()
return ConversationItem(id=conversation_id, created_at=now, updated_at=now)
def set_conversation_title(self, conversation_id: str, title: str):
conn = self._connect()
try:
cur = conn.cursor()
cur.execute(
"UPDATE conversations SET title = ? WHERE id = ?",
(title, conversation_id),
)
conn.commit()
finally:
conn.close()
def add_message(self, conversation_id: str, role: str, content: str, model: Optional[str] = None, tokens: Optional[int] = None):
now = time.time()
conn = self._connect()
try:
cur = conn.cursor()
cur.execute(
"INSERT INTO messages (conversation_id, role, content, timestamp, model, tokens) VALUES (?, ?, ?, ?, ?, ?)",
(conversation_id, role, content, now, model, tokens)
)
message_id = cur.lastrowid
cur.execute(
"UPDATE conversations SET updated_at = ? WHERE id = ?",
(now, conversation_id)
)
conn.commit()
return message_id
finally:
conn.close()
def get_conversation(self, conversation_id: str) -> Optional[ConversationItem]:
conn = self._connect()
cur = conn.cursor()
cur.execute("SELECT * FROM conversations WHERE id = ?", (conversation_id,))
row = cur.fetchone()
if not row:
conn.close()
return None
cur.execute(
"SELECT role, content, timestamp, model, tokens FROM messages WHERE conversation_id = ? ORDER BY timestamp ASC",
(conversation_id,)
)
msg_rows = cur.fetchall()
conn.close()
messages = [MessageItem(role=r["role"], content=r["content"], timestamp=r["timestamp"], model=r["model"], tokens=r["tokens"]) for r in msg_rows]
return ConversationItem(
id=row["id"],
messages=messages,
created_at=row["created_at"],
updated_at=row["updated_at"],
title=row["title"] if "title" in row.keys() else None,
)
def all_conversations(self) -> List[ConversationItem]:
conn = self._connect()
cur = conn.cursor()
cur.execute("""
SELECT c.id, c.created_at, c.updated_at, c.title,
m.role, m.content, m.timestamp, m.model, m.tokens
FROM conversations c
LEFT JOIN messages m ON m.conversation_id = c.id
ORDER BY c.updated_at DESC, m.timestamp ASC
""")
rows = cur.fetchall()
conn.close()
convs: Dict[str, ConversationItem] = {}
order: list[str] = []
for row in rows:
cid = row["id"]
if cid not in convs:
convs[cid] = ConversationItem(
id=cid,
created_at=row["created_at"],
updated_at=row["updated_at"],
title=row["title"],
)
order.append(cid)
if row["role"] is not None:
convs[cid].messages.append(
MessageItem(role=row["role"], content=row["content"], timestamp=row["timestamp"], model=row["model"], tokens=row["tokens"])
)
return [convs[cid] for cid in order]
def delete_conversation(self, conversation_id: str):
conn = self._connect()
try:
cur = conn.cursor()
cur.execute("DELETE FROM messages WHERE conversation_id = ?", (conversation_id,))
cur.execute("DELETE FROM conversations WHERE id = ?", (conversation_id,))
conn.commit()
finally:
conn.close()
# -----------------------------
# Search API
# -----------------------------
def search_conversations(self, query: str, limit: int = 3) -> List[dict]:
"""Return up to `limit` conversations that contain the query string,
with full user+assistant exchange pairs around each match."""
if not query or not query.strip():
return []
q = query.strip().lower()
conn = self._connect()
cur = conn.cursor()
cur.execute("""
SELECT DISTINCT c.id, c.created_at, c.updated_at
FROM conversations c
JOIN messages m ON m.conversation_id = c.id
WHERE LOWER(m.content) LIKE ?
ORDER BY c.updated_at DESC
LIMIT ?
""", (f"%{q}%", limit))
rows = cur.fetchall()
results = []
for row in rows:
# Load all messages in order so we can find complete exchange pairs
cur.execute("""
SELECT role, content FROM messages
WHERE conversation_id = ?
ORDER BY timestamp ASC
""", (row["id"],))
all_msgs = [{"role": r["role"], "content": r["content"]} for r in cur.fetchall()]
# For each matching message, collect the full user+assistant pair around it
seen_pairs: set = set()
matches = []
for i, msg in enumerate(all_msgs):
if q not in msg["content"].lower():
continue
if msg["role"] == "user":
start, end = i, i + 1 if i + 1 < len(all_msgs) else i
else:
start, end = (i - 1 if i > 0 else i), i
if (start, end) in seen_pairs:
continue
seen_pairs.add((start, end))
for m in all_msgs[start:end + 1]:
matches.append({"role": m["role"], "content": m["content"][:500]})
if len(seen_pairs) >= 2:
break
if matches:
results.append({
"id": row["id"],
"updated_at": row["updated_at"],
"matches": matches,
})
conn.close()
return results
# nomic-embed-text is an asymmetric retrieval model: queries and stored
# documents must be embedded with these task prefixes or similarity collapses
# into noise. Cached vectors are document-embeddings (search_document:).
_EMBED_QUERY_PREFIX = "search_query: "
_EMBED_DOC_PREFIX = "search_document: "
async def semantic_search_conversations(
self, query: str, embed_fn, limit: int = 3, min_score: float = 0.6
) -> List[dict]:
"""Recall past conversations relevant to `query` using hybrid retrieval.
Combines semantic similarity (embeddings — finds reworded matches with no
shared keywords) with the existing lexical substring match (catches exact
terms the embedding underweights), unioned and deduped by conversation.
`embed_fn` is an async callable returning an embedding vector for a string
(typically OllamaManager.embed). Messages without a stored vector are
embedded and cached on first use (lazy backfill). If embeddings are
unavailable (no model / Ollama down) this degrades to pure lexical match,
so recall never silently breaks.
Returns the same shape as `search_conversations`: a list of
{id, updated_at, matches:[{role, content}]} with full user+assistant
pairs around each match.
"""
if not query or not query.strip():
return []
query_vec = await embed_fn(self._EMBED_QUERY_PREFIX + query.strip())
if not query_vec:
return self.search_conversations(query, limit=limit)
conn = self._connect()
cur = conn.cursor()
# Lazy backfill: embed any messages that don't have a vector yet.
cur.execute("""
SELECT m.id, m.content
FROM messages m
LEFT JOIN message_vectors v ON v.message_id = m.id
WHERE v.message_id IS NULL AND TRIM(m.content) != ''
""")
missing = cur.fetchall()
for row in missing:
vec = await embed_fn(self._EMBED_DOC_PREFIX + row["content"][:2000])
if vec:
cur.execute(
"INSERT OR REPLACE INTO message_vectors (message_id, embedding) VALUES (?, ?)",
(row["id"], json.dumps(vec)),
)
if missing:
conn.commit()
# Score every stored message against the query vector.
cur.execute("""
SELECT v.message_id, v.embedding, m.conversation_id
FROM message_vectors v
JOIN messages m ON m.id = v.message_id
""")
scored = []
for row in cur.fetchall():
try:
vec = json.loads(row["embedding"])
except Exception:
continue
score = _cosine(query_vec, vec)
if score >= min_score:
scored.append((score, row["message_id"], row["conversation_id"]))
scored.sort(reverse=True)
results: List[dict] = []
seen_convs: set = set()
for score, message_id, conv_id in scored:
if len(results) >= limit:
break
if conv_id in seen_convs:
continue
pair = self._exchange_pair(cur, conv_id, message_id)
if pair:
seen_convs.add(conv_id)
results.append(pair)
conn.close()
# Hybrid union: fill any remaining slots with lexical matches the
# embedding missed (e.g. exact proper nouns), skipping dupes.
if len(results) < limit:
for conv in self.search_conversations(query, limit=limit):
if conv["id"] not in seen_convs:
seen_convs.add(conv["id"])
results.append(conv)
if len(results) >= limit:
break
return results
def _exchange_pair(self, cur, conv_id: str, message_id: int) -> Optional[dict]:
"""Build a {id, updated_at, matches} record with the full user+assistant
pair surrounding `message_id`, in the shape search callers expect."""
cur.execute(
"SELECT id, role, content FROM messages WHERE conversation_id = ? ORDER BY timestamp ASC",
(conv_id,),
)
all_msgs = cur.fetchall()
idx = next((i for i, m in enumerate(all_msgs) if m["id"] == message_id), None)
if idx is None:
return None
if all_msgs[idx]["role"] == "user":
start, end = idx, min(idx + 1, len(all_msgs) - 1)
else:
start, end = max(idx - 1, 0), idx
matches = [
{"role": all_msgs[i]["role"], "content": all_msgs[i]["content"][:500]}
for i in range(start, end + 1)
]
cur.execute("SELECT updated_at FROM conversations WHERE id = ?", (conv_id,))
crow = cur.fetchone()
return {
"id": conv_id,
"updated_at": crow["updated_at"] if crow else 0,
"matches": matches,
}
# -----------------------------
# Settings API
# -----------------------------
_SETTINGS_DEFAULTS: Dict[str, Any] = {
"model": "",
"temperature": 0.7,
"system_prompt": "",
"timeout": 120,
# How long Ollama keeps the model resident in VRAM between messages.
# "30m"/"-1" (never unload)/"0" (unload now). Avoids cold-reload latency
# when you return to an idle chat. Empty → Ollama's 5-minute default.
"keep_alive": "30m",
# CPU/GPU offload: -1 = Auto (Ollama auto-fits layers to VRAM).
# 0100 = percent of model layers to force onto the GPU; the
# remainder runs on CPU. See OllamaManager.get_model_layers.
"gpu_offload": -1,
# Memory curator (the model that extracts facts after each exchange).
# Empty → same auto-selected model as chat. A dedicated model (e.g.
# "mistral:latest") gives better extraction but must share VRAM.
"memory_model": "mistral:latest",
# Curator CPU/GPU offload — same scale as gpu_offload above. Default 0
# (all CPU/RAM): OS-neutral and never evicts the chat model from a small
# GPU. Boxes with spare VRAM can set -1 (Auto) or a percent to use the GPU.
"memory_gpu_offload": 0,
# Similar-fact merge: when a newly extracted fact's embedding is at least
# this cosine-similar to an existing fact, UPDATE that fact in place
# instead of appending a duplicate ("edit with new info"). 0 disables
# (always append). Calibrated on nomic-embed-text: genuine updates
# (mileage/title/location changes) score 0.810.99, while distinct facts
# top out ~0.61 — so 0.80 catches updates and never merges unrelated
# facts. Lower to catch looser rephrases; raise toward 1.0 to be stricter.
"memory_merge_threshold": 0.80,
}
def get_settings(self) -> Dict[str, Any]:
conn = self._connect()
cur = conn.cursor()
cur.execute("SELECT key, value FROM settings")
rows = cur.fetchall()
conn.close()
result = dict(self._SETTINGS_DEFAULTS)
for row in rows:
try:
result[row["key"]] = json.loads(row["value"])
except Exception:
result[row["key"]] = row["value"]
return result
def update_settings(self, data: Dict[str, Any]):
conn = self._connect()
try:
cur = conn.cursor()
for key, value in data.items():
if key in self._SETTINGS_DEFAULTS:
cur.execute(
"INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)",
(key, json.dumps(value))
)
conn.commit()
finally:
conn.close()
# -----------------------------
# Store Instance
# -----------------------------
from ..nexus_config import MEMORY_DB
DB_PATH = MEMORY_DB
store = PersistentMemoryStore(DB_PATH)