jonandClaude Opus 4.8
492020b547
- Conversations bind to a project on creation; RAG scopes to the conversation's project, not the global setting. - Action tools (web_search/fetch_url/remember) are withheld unless allow_action_tools is enabled (off by default). Settings toggle. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1066 lines
41 KiB
Python
1066 lines
41 KiB
Python
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional
|
||
from pydantic import BaseModel
|
||
|
||
from ..nexus_config import DEFAULT_MEMORY_MODEL
|
||
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.vec_enabled = self._probe_vec()
|
||
self._ensure_tables()
|
||
self._cache: Dict[str, MemoryItem] = self._load_all_memory()
|
||
|
||
# -----------------------------
|
||
# Internal helpers
|
||
# -----------------------------
|
||
@staticmethod
|
||
def _probe_vec() -> bool:
|
||
"""True if this host can load the sqlite-vec extension. Some Python
|
||
builds ship SQLite with loadable extensions disabled — those fall back
|
||
to the brute-force cosine scan, so recall never depends on this."""
|
||
try:
|
||
import sqlite_vec
|
||
c = sqlite3.connect(":memory:")
|
||
c.enable_load_extension(True)
|
||
sqlite_vec.load(c)
|
||
c.execute("SELECT vec_version()")
|
||
c.close()
|
||
return True
|
||
except Exception:
|
||
return False
|
||
|
||
def _connect(self):
|
||
conn = sqlite3.connect(self.db_path)
|
||
conn.row_factory = sqlite3.Row
|
||
conn.execute("PRAGMA journal_mode=WAL;")
|
||
if self.vec_enabled:
|
||
try:
|
||
import sqlite_vec
|
||
conn.enable_load_extension(True)
|
||
sqlite_vec.load(conn)
|
||
conn.enable_load_extension(False)
|
||
except Exception:
|
||
pass
|
||
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
|
||
# Migrate: bind a conversation to a project ("" = unscoped).
|
||
try:
|
||
cur.execute("ALTER TABLE conversations ADD COLUMN project_id TEXT NOT NULL DEFAULT ''")
|
||
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)
|
||
)
|
||
""")
|
||
|
||
# RAG: one row per chunk. doc_id groups the chunks of a single uploaded
|
||
# document; embedding is a JSON vector (search_document-prefixed), set at
|
||
# ingest so retrieval needs no backfill.
|
||
cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS documents (
|
||
id TEXT PRIMARY KEY,
|
||
doc_id TEXT NOT NULL,
|
||
title TEXT NOT NULL,
|
||
chunk_idx INTEGER NOT NULL,
|
||
text TEXT NOT NULL,
|
||
embedding TEXT,
|
||
created_at REAL NOT NULL
|
||
)
|
||
""")
|
||
cur.execute("CREATE INDEX IF NOT EXISTS idx_documents_doc_id ON documents (doc_id)")
|
||
|
||
# Projects / workspaces: group documents so RAG can scope to one set.
|
||
cur.execute("""
|
||
CREATE TABLE IF NOT EXISTS projects (
|
||
id TEXT PRIMARY KEY,
|
||
name TEXT NOT NULL,
|
||
created_at REAL NOT NULL
|
||
)
|
||
""")
|
||
# documents.project_id — "" (or missing) means unscoped / All.
|
||
try:
|
||
cur.execute("ALTER TABLE documents ADD COLUMN project_id TEXT NOT NULL DEFAULT ''")
|
||
except Exception:
|
||
pass
|
||
|
||
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, project_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, project_id) VALUES (?, ?, ?, ?)",
|
||
(conversation_id, now, now, project_id or "")
|
||
)
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
return ConversationItem(id=conversation_id, created_at=now, updated_at=now)
|
||
|
||
def conversation_project(self, conversation_id: str) -> Optional[str]:
|
||
"""The conversation's bound project ('' = unscoped), or None if it doesn't
|
||
exist yet — lets a new chat inherit the current workspace."""
|
||
conn = self._connect()
|
||
row = conn.execute(
|
||
"SELECT project_id FROM conversations WHERE id = ?", (conversation_id,)
|
||
).fetchone()
|
||
conn.close()
|
||
return None if row is None else (row["project_id"] or "")
|
||
|
||
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)),
|
||
)
|
||
self._vec_upsert_msg(conn, row["id"], vec) # mirror into the ANN index
|
||
if missing:
|
||
conn.commit()
|
||
|
||
# Rank messages by similarity. Fast path: the sqlite-vec ANN index over an
|
||
# over-fetch (conversation dedup below thins it); else brute-force cosine.
|
||
scored = None
|
||
if self.vec_enabled:
|
||
scored = self._vec_search_messages(conn, query_vec, max(limit * 5, 20), min_score)
|
||
if scored is None:
|
||
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
|
||
|
||
# --- message vector index (sqlite-vec) — same pattern as documents --------
|
||
def _ensure_vec_msgs(self, conn, dim: int) -> None:
|
||
conn.execute(
|
||
f"CREATE VIRTUAL TABLE IF NOT EXISTS vec_messages "
|
||
f"USING vec0(embedding float[{dim}] distance_metric=cosine)"
|
||
)
|
||
|
||
def _vec_upsert_msg(self, conn, message_id: int, vec: list) -> None:
|
||
if not (self.vec_enabled and vec):
|
||
return
|
||
try:
|
||
import sqlite_vec
|
||
self._ensure_vec_msgs(conn, len(vec))
|
||
conn.execute("DELETE FROM vec_messages WHERE rowid = ?", (message_id,))
|
||
conn.execute(
|
||
"INSERT INTO vec_messages(rowid, embedding) VALUES (?, ?)",
|
||
(message_id, sqlite_vec.serialize_float32(vec)),
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
def _backfill_vec_msgs(self, conn, dim: int) -> None:
|
||
"""Index any message_vectors rows missing from vec_messages."""
|
||
try:
|
||
import sqlite_vec
|
||
self._ensure_vec_msgs(conn, dim)
|
||
rows = conn.execute(
|
||
"SELECT mv.message_id AS mid, mv.embedding AS emb FROM message_vectors mv "
|
||
"LEFT JOIN vec_messages v ON v.rowid = mv.message_id WHERE v.rowid IS NULL"
|
||
).fetchall()
|
||
for r in rows:
|
||
try:
|
||
vec = json.loads(r["emb"])
|
||
if len(vec) == dim:
|
||
conn.execute(
|
||
"INSERT INTO vec_messages(rowid, embedding) VALUES (?, ?)",
|
||
(r["mid"], sqlite_vec.serialize_float32(vec)),
|
||
)
|
||
except Exception:
|
||
pass
|
||
conn.commit()
|
||
except Exception:
|
||
pass
|
||
|
||
def _vec_search_messages(self, conn, query_vec: list, fetch: int, min_score: float):
|
||
"""Top message hits via the vec index as sorted [(score, msg_id, conv_id)],
|
||
or None to fall back to the brute-force scan. Stale rows for deleted
|
||
messages are dropped by the inner join, so they never surface."""
|
||
try:
|
||
import sqlite_vec
|
||
self._backfill_vec_msgs(conn, len(query_vec))
|
||
rows = conn.execute(
|
||
"SELECT v.rowid AS mid, v.distance AS distance, m.conversation_id AS cid "
|
||
"FROM vec_messages v JOIN messages m ON m.id = v.rowid "
|
||
"WHERE v.embedding MATCH ? ORDER BY v.distance LIMIT ?",
|
||
(sqlite_vec.serialize_float32(query_vec), fetch),
|
||
).fetchall()
|
||
return [
|
||
(1.0 - r["distance"], r["mid"], r["cid"])
|
||
for r in rows if (1.0 - r["distance"]) >= min_score
|
||
] # distance-asc == score-desc, already sorted
|
||
except Exception:
|
||
return None
|
||
|
||
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,
|
||
}
|
||
|
||
# -----------------------------
|
||
# Documents (RAG)
|
||
# -----------------------------
|
||
@staticmethod
|
||
def _chunk_text(text: str, size: int = 800, overlap: int = 120) -> List[str]:
|
||
"""Pack paragraphs into ~`size`-char chunks with a char `overlap` carried
|
||
across boundaries, so a passage spanning two chunks still matches. Any
|
||
single paragraph larger than `size` (common in PDFs with few blank lines)
|
||
is hard-split into overlapping windows first.
|
||
ponytail: char-based, not token-based — fine for local recall; move to a
|
||
token splitter only if chunk sizes start hurting the context budget."""
|
||
units: List[str] = []
|
||
for para in (p.strip() for p in text.split("\n\n")):
|
||
if not para:
|
||
continue
|
||
if len(para) <= size:
|
||
units.append(para)
|
||
else:
|
||
step = max(1, size - overlap)
|
||
units.extend(para[i:i + size] for i in range(0, len(para), step))
|
||
|
||
chunks: List[str] = []
|
||
buf = ""
|
||
for u in units:
|
||
if buf and len(buf) + len(u) + 2 > size:
|
||
chunks.append(buf)
|
||
tail = buf[-overlap:] if overlap else "" # overlap seed for the next chunk
|
||
buf = f"{tail}\n\n{u}" if tail else u
|
||
else:
|
||
buf = f"{buf}\n\n{u}" if buf else u
|
||
if buf:
|
||
chunks.append(buf)
|
||
return chunks
|
||
|
||
async def add_document(self, title: str, content: str, embed_fn, project_id: str = "") -> dict:
|
||
"""Chunk, embed, and store a document. Returns {doc_id, chunks}."""
|
||
import uuid as _uuid
|
||
doc_id = str(_uuid.uuid4())
|
||
pieces = self._chunk_text(content)
|
||
now = time.time()
|
||
conn = self._connect()
|
||
cur = conn.cursor()
|
||
for i, piece in enumerate(pieces):
|
||
vec = await embed_fn(self._EMBED_DOC_PREFIX + piece[:2000])
|
||
cur.execute(
|
||
"INSERT INTO documents (id, doc_id, title, chunk_idx, text, embedding, created_at, project_id)"
|
||
" VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
|
||
(str(_uuid.uuid4()), doc_id, title, i, piece,
|
||
json.dumps(vec) if vec else None, now, project_id or ""),
|
||
)
|
||
self._vec_upsert(conn, cur.lastrowid, vec) # mirror into the ANN index
|
||
conn.commit()
|
||
conn.close()
|
||
return {"doc_id": doc_id, "title": title, "chunks": len(pieces)}
|
||
|
||
# --- projects / workspaces --------------------------------------------------
|
||
def create_project(self, name: str) -> dict:
|
||
import uuid as _uuid
|
||
pid = str(_uuid.uuid4())
|
||
conn = self._connect()
|
||
conn.execute("INSERT INTO projects (id, name, created_at) VALUES (?, ?, ?)",
|
||
(pid, name.strip(), time.time()))
|
||
conn.commit()
|
||
conn.close()
|
||
return {"id": pid, "name": name.strip()}
|
||
|
||
def list_projects(self) -> List[dict]:
|
||
conn = self._connect()
|
||
rows = conn.execute("""
|
||
SELECT p.id, p.name, p.created_at,
|
||
(SELECT COUNT(DISTINCT doc_id) FROM documents d WHERE d.project_id = p.id) AS docs
|
||
FROM projects p ORDER BY p.created_at ASC
|
||
""").fetchall()
|
||
conn.close()
|
||
return [dict(r) for r in rows]
|
||
|
||
def delete_project(self, project_id: str) -> bool:
|
||
"""Delete a project; its documents survive but become unscoped ("")."""
|
||
conn = self._connect()
|
||
conn.execute("UPDATE documents SET project_id = '' WHERE project_id = ?", (project_id,))
|
||
cur = conn.execute("DELETE FROM projects WHERE id = ?", (project_id,))
|
||
deleted = cur.rowcount
|
||
conn.commit()
|
||
conn.close()
|
||
return deleted > 0
|
||
|
||
# --- vector index (sqlite-vec) — accelerator over the JSON embedding column,
|
||
# --- with brute-force cosine below as the guaranteed fallback ---------------
|
||
def _ensure_vec_docs(self, conn, dim: int) -> None:
|
||
conn.execute(
|
||
f"CREATE VIRTUAL TABLE IF NOT EXISTS vec_documents "
|
||
f"USING vec0(embedding float[{dim}] distance_metric=cosine)"
|
||
)
|
||
|
||
def _vec_upsert(self, conn, rowid: int, vec: list) -> None:
|
||
"""Best-effort mirror of one chunk's vector into the vec index."""
|
||
if not (self.vec_enabled and vec):
|
||
return
|
||
try:
|
||
import sqlite_vec
|
||
self._ensure_vec_docs(conn, len(vec))
|
||
conn.execute("DELETE FROM vec_documents WHERE rowid = ?", (rowid,))
|
||
conn.execute(
|
||
"INSERT INTO vec_documents(rowid, embedding) VALUES (?, ?)",
|
||
(rowid, sqlite_vec.serialize_float32(vec)),
|
||
)
|
||
except Exception:
|
||
pass # the index is an accelerator, never a requirement
|
||
|
||
def _backfill_vec(self, conn, dim: int) -> None:
|
||
"""Index any document chunks missing from vec_documents (older rows, or
|
||
rows written while the extension was unavailable)."""
|
||
try:
|
||
import sqlite_vec
|
||
self._ensure_vec_docs(conn, dim)
|
||
rows = conn.execute(
|
||
"SELECT d.rowid AS rid, d.embedding AS emb FROM documents d "
|
||
"LEFT JOIN vec_documents v ON v.rowid = d.rowid "
|
||
"WHERE v.rowid IS NULL AND d.embedding IS NOT NULL"
|
||
).fetchall()
|
||
for r in rows:
|
||
try:
|
||
vec = json.loads(r["emb"])
|
||
if len(vec) == dim:
|
||
conn.execute(
|
||
"INSERT INTO vec_documents(rowid, embedding) VALUES (?, ?)",
|
||
(r["rid"], sqlite_vec.serialize_float32(vec)),
|
||
)
|
||
except Exception:
|
||
pass
|
||
conn.commit()
|
||
except Exception:
|
||
pass
|
||
|
||
def _vec_search(self, query_vec: list, limit: int, min_score: float):
|
||
"""KNN over the vec index. Returns hits, or None to signal fall back to
|
||
the brute-force scan (e.g. extension error or dimension mismatch)."""
|
||
try:
|
||
import sqlite_vec
|
||
conn = self._connect()
|
||
self._backfill_vec(conn, len(query_vec))
|
||
rows = conn.execute(
|
||
"SELECT d.title AS title, d.text AS text, v.distance AS distance "
|
||
"FROM vec_documents v JOIN documents d ON d.rowid = v.rowid "
|
||
"WHERE v.embedding MATCH ? ORDER BY v.distance LIMIT ?",
|
||
(sqlite_vec.serialize_float32(query_vec), limit),
|
||
).fetchall()
|
||
conn.close()
|
||
# sqlite-vec cosine distance = 1 - cosine similarity
|
||
return [
|
||
{"title": r["title"], "text": r["text"], "score": 1.0 - r["distance"]}
|
||
for r in rows if (1.0 - r["distance"]) >= min_score
|
||
]
|
||
except Exception:
|
||
return None
|
||
|
||
def list_documents(self, project_id: Optional[str] = None) -> List[dict]:
|
||
"""All documents, or just one project's when project_id is given."""
|
||
conn = self._connect()
|
||
cur = conn.cursor()
|
||
if project_id is not None:
|
||
cur.execute("""
|
||
SELECT doc_id, title, COUNT(*) AS chunks, MIN(created_at) AS created_at
|
||
FROM documents WHERE project_id = ?
|
||
GROUP BY doc_id, title ORDER BY created_at DESC
|
||
""", (project_id,))
|
||
else:
|
||
cur.execute("""
|
||
SELECT doc_id, title, COUNT(*) AS chunks, MIN(created_at) AS created_at
|
||
FROM documents GROUP BY doc_id, title ORDER BY created_at DESC
|
||
""")
|
||
rows = cur.fetchall()
|
||
conn.close()
|
||
return [dict(r) for r in rows]
|
||
|
||
def get_document(self, doc_id: str) -> List[dict]:
|
||
"""Ordered chunks of one document: [{chunk_idx, text}]."""
|
||
conn = self._connect()
|
||
cur = conn.cursor()
|
||
cur.execute(
|
||
"SELECT chunk_idx, text FROM documents WHERE doc_id = ? ORDER BY chunk_idx",
|
||
(doc_id,),
|
||
)
|
||
rows = cur.fetchall()
|
||
conn.close()
|
||
return [dict(r) for r in rows]
|
||
|
||
def delete_document(self, doc_id: str) -> bool:
|
||
conn = self._connect()
|
||
cur = conn.cursor()
|
||
if self.vec_enabled:
|
||
try:
|
||
for r in cur.execute("SELECT rowid FROM documents WHERE doc_id = ?", (doc_id,)).fetchall():
|
||
conn.execute("DELETE FROM vec_documents WHERE rowid = ?", (r["rowid"],))
|
||
except Exception:
|
||
pass
|
||
cur.execute("DELETE FROM documents WHERE doc_id = ?", (doc_id,))
|
||
deleted = cur.rowcount
|
||
conn.commit()
|
||
conn.close()
|
||
return deleted > 0
|
||
|
||
async def search_documents(
|
||
self, query: str, embed_fn, limit: int = 3, min_score: float = 0.6,
|
||
project_id: Optional[str] = None,
|
||
) -> List[dict]:
|
||
"""Top-`limit` document chunks most similar to `query`. Returns
|
||
[{title, text, score}]. Empty on no query / embeddings down.
|
||
When project_id is given, only that project's docs are searched.
|
||
ponytail: scoped search uses the brute-force path (easy SQL filter, few
|
||
docs per project); the vec index accelerates the unscoped "All" case."""
|
||
if not query or not query.strip():
|
||
return []
|
||
query_vec = await embed_fn(self._EMBED_QUERY_PREFIX + query.strip())
|
||
if not query_vec:
|
||
return []
|
||
# Fast path (unscoped only): the sqlite-vec ANN index.
|
||
if not project_id and self.vec_enabled:
|
||
hits = self._vec_search(query_vec, limit, min_score)
|
||
if hits is not None:
|
||
return hits
|
||
conn = self._connect()
|
||
cur = conn.cursor()
|
||
if project_id:
|
||
cur.execute("SELECT title, text, embedding FROM documents WHERE embedding IS NOT NULL AND project_id = ?", (project_id,))
|
||
else:
|
||
cur.execute("SELECT title, text, embedding FROM documents WHERE embedding IS NOT NULL")
|
||
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({"title": row["title"], "text": row["text"], "score": score})
|
||
conn.close()
|
||
scored.sort(key=lambda d: d["score"], reverse=True)
|
||
return scored[:limit]
|
||
|
||
# -----------------------------
|
||
# Settings API
|
||
# -----------------------------
|
||
_SETTINGS_DEFAULTS: Dict[str, Any] = {
|
||
"model": "",
|
||
# Qwen3-style reasoning. Off by default: the hidden <think> block is pure
|
||
# latency for chat/memory. Turn on for hard multi-step problems.
|
||
"think": False,
|
||
"temperature": 0.7,
|
||
# Context window (tokens Ollama keeps in view). 0 → Ollama's model default.
|
||
"num_ctx": 0,
|
||
# RAG retrieval: how many document chunks to inject, and the minimum
|
||
# cosine similarity (0-1) a chunk must clear to count as relevant.
|
||
"rag_top_k": 3,
|
||
"rag_min_score": 0.6,
|
||
# Active project/workspace; "" = all documents (unscoped).
|
||
"active_project": "",
|
||
# Consent gate for tools that act (web_search/fetch_url/remember). Off by
|
||
# default: a playbook can list them, but they only run when this is on.
|
||
"allow_action_tools": False,
|
||
"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).
|
||
# 0–100 = 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": DEFAULT_MEMORY_MODEL,
|
||
# 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.81–0.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) |
GitNexus · part of the Nexus stack