379 lines
14 KiB
Python
379 lines
14 KiB
Python
import json
|
|
import sqlite3
|
|
import threading
|
|
from datetime import datetime, timezone
|
|
|
|
import msgpack
|
|
import numpy as np
|
|
|
|
import config
|
|
|
|
DB_PATH = config.DATA_DIR / "jlens.db"
|
|
FRAMES_DIR = config.DATA_DIR / "frames"
|
|
|
|
SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS conversations (
|
|
id INTEGER PRIMARY KEY,
|
|
title TEXT NOT NULL DEFAULT '',
|
|
tags TEXT NOT NULL DEFAULT '[]',
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL
|
|
);
|
|
CREATE TABLE IF NOT EXISTS messages (
|
|
id INTEGER PRIMARY KEY,
|
|
conversation_id INTEGER NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
|
|
parent_id INTEGER REFERENCES messages(id),
|
|
role TEXT NOT NULL,
|
|
content TEXT NOT NULL,
|
|
meta TEXT,
|
|
frames_file TEXT,
|
|
created_at TEXT NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_messages_conv ON messages(conversation_id);
|
|
CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
|
|
content, content='messages', content_rowid='id'
|
|
);
|
|
CREATE TRIGGER IF NOT EXISTS messages_ai AFTER INSERT ON messages BEGIN
|
|
INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
|
|
END;
|
|
CREATE TRIGGER IF NOT EXISTS messages_ad AFTER DELETE ON messages BEGIN
|
|
INSERT INTO messages_fts(messages_fts, rowid, content)
|
|
VALUES ('delete', old.id, old.content);
|
|
END;
|
|
CREATE TRIGGER IF NOT EXISTS messages_au AFTER UPDATE OF content ON messages BEGIN
|
|
INSERT INTO messages_fts(messages_fts, rowid, content)
|
|
VALUES ('delete', old.id, old.content);
|
|
INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
|
|
END;
|
|
"""
|
|
|
|
|
|
def _now():
|
|
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
|
|
|
|
|
class Store:
|
|
def __init__(self):
|
|
config.DATA_DIR.mkdir(parents=True, exist_ok=True)
|
|
FRAMES_DIR.mkdir(parents=True, exist_ok=True)
|
|
self._local = threading.local()
|
|
conn = self._conn()
|
|
conn.executescript(SCHEMA)
|
|
conn.commit()
|
|
|
|
def _conn(self):
|
|
conn = getattr(self._local, "conn", None)
|
|
if conn is None:
|
|
conn = sqlite3.connect(DB_PATH)
|
|
conn.row_factory = sqlite3.Row
|
|
conn.execute("PRAGMA journal_mode=WAL")
|
|
conn.execute("PRAGMA foreign_keys=ON")
|
|
self._local.conn = conn
|
|
return conn
|
|
|
|
def create_conversation(self, title, tags=None):
|
|
conn = self._conn()
|
|
now = _now()
|
|
cur = conn.execute(
|
|
"INSERT INTO conversations (title, tags, created_at, updated_at) VALUES (?, ?, ?, ?)",
|
|
(title, json.dumps(tags or []), now, now),
|
|
)
|
|
conn.commit()
|
|
return cur.lastrowid
|
|
|
|
def update_conversation(self, conversation_id, title=None, tags=None):
|
|
conn = self._conn()
|
|
if title is not None:
|
|
conn.execute(
|
|
"UPDATE conversations SET title = ?, updated_at = ? WHERE id = ?",
|
|
(title, _now(), conversation_id),
|
|
)
|
|
if tags is not None:
|
|
conn.execute(
|
|
"UPDATE conversations SET tags = ?, updated_at = ? WHERE id = ?",
|
|
(json.dumps(tags), _now(), conversation_id),
|
|
)
|
|
conn.commit()
|
|
|
|
def delete_conversation(self, conversation_id):
|
|
conn = self._conn()
|
|
rows = conn.execute(
|
|
"SELECT frames_file FROM messages WHERE conversation_id = ? AND frames_file IS NOT NULL",
|
|
(conversation_id,),
|
|
).fetchall()
|
|
conn.execute("DELETE FROM conversations WHERE id = ?", (conversation_id,))
|
|
conn.commit()
|
|
for row in rows:
|
|
(FRAMES_DIR / row["frames_file"]).unlink(missing_ok=True)
|
|
|
|
def list_conversations(self, query=None, limit=200):
|
|
conn = self._conn()
|
|
if query:
|
|
hits = conn.execute(
|
|
"""
|
|
SELECT messages_fts.rowid AS mid,
|
|
snippet(messages_fts, 0, '[', ']', '…', 12) AS snip,
|
|
rank
|
|
FROM messages_fts
|
|
WHERE messages_fts MATCH ?
|
|
ORDER BY rank
|
|
LIMIT 500
|
|
""",
|
|
(query,),
|
|
).fetchall()
|
|
best = {}
|
|
for hit in hits:
|
|
row = conn.execute(
|
|
"SELECT conversation_id FROM messages WHERE id = ?", (hit["mid"],)
|
|
).fetchone()
|
|
if row and row["conversation_id"] not in best:
|
|
best[row["conversation_id"]] = hit["snip"]
|
|
rows = []
|
|
for cid, snip in list(best.items())[:limit]:
|
|
conv = conn.execute(
|
|
"""
|
|
SELECT c.id, c.title, c.tags, c.updated_at,
|
|
(SELECT count(*) FROM messages m WHERE m.conversation_id = c.id) AS n_messages
|
|
FROM conversations c WHERE c.id = ?
|
|
""",
|
|
(cid,),
|
|
).fetchone()
|
|
if conv:
|
|
rows.append(dict(conv, snippet=snip))
|
|
else:
|
|
rows = conn.execute(
|
|
"""
|
|
SELECT c.id, c.title, c.tags, c.updated_at,
|
|
count(m.id) AS n_messages, NULL AS snippet
|
|
FROM conversations c
|
|
LEFT JOIN messages m ON m.conversation_id = c.id
|
|
GROUP BY c.id
|
|
ORDER BY c.updated_at DESC
|
|
LIMIT ?
|
|
""",
|
|
(limit,),
|
|
).fetchall()
|
|
return [
|
|
dict(row, tags=json.loads(row["tags"]))
|
|
for row in (dict(r) for r in rows)
|
|
]
|
|
|
|
def get_conversation(self, conversation_id):
|
|
conn = self._conn()
|
|
conv = conn.execute(
|
|
"SELECT * FROM conversations WHERE id = ?", (conversation_id,)
|
|
).fetchone()
|
|
if conv is None:
|
|
raise ValueError(f"unknown conversation {conversation_id}")
|
|
rows = conn.execute(
|
|
"SELECT id, parent_id, role, content, meta, frames_file, created_at "
|
|
"FROM messages WHERE conversation_id = ? ORDER BY id",
|
|
(conversation_id,),
|
|
).fetchall()
|
|
messages = [
|
|
{
|
|
"id": row["id"],
|
|
"parent_id": row["parent_id"],
|
|
"role": row["role"],
|
|
"content": row["content"],
|
|
"meta": json.loads(row["meta"]) if row["meta"] else None,
|
|
"has_frames": row["frames_file"] is not None,
|
|
"created_at": row["created_at"],
|
|
}
|
|
for row in rows
|
|
]
|
|
return {
|
|
"id": conv["id"],
|
|
"title": conv["title"],
|
|
"tags": json.loads(conv["tags"]),
|
|
"created_at": conv["created_at"],
|
|
"updated_at": conv["updated_at"],
|
|
"messages": messages,
|
|
}
|
|
|
|
def add_message(self, conversation_id, parent_id, role, content, meta=None):
|
|
conn = self._conn()
|
|
cur = conn.execute(
|
|
"INSERT INTO messages (conversation_id, parent_id, role, content, meta, created_at) "
|
|
"VALUES (?, ?, ?, ?, ?, ?)",
|
|
(
|
|
conversation_id,
|
|
parent_id,
|
|
role,
|
|
content,
|
|
json.dumps(meta, ensure_ascii=False) if meta else None,
|
|
_now(),
|
|
),
|
|
)
|
|
conn.execute(
|
|
"UPDATE conversations SET updated_at = ? WHERE id = ?",
|
|
(_now(), conversation_id),
|
|
)
|
|
conn.commit()
|
|
return cur.lastrowid
|
|
|
|
def get_message(self, message_id):
|
|
conn = self._conn()
|
|
row = conn.execute(
|
|
"SELECT id, conversation_id, parent_id, role, content, meta, frames_file "
|
|
"FROM messages WHERE id = ?",
|
|
(message_id,),
|
|
).fetchone()
|
|
if row is None:
|
|
raise ValueError(f"unknown message {message_id}")
|
|
return dict(row)
|
|
|
|
def update_message(self, message_id, content, meta=None):
|
|
"""Rewrite a message's content (assistant edit, or a continuation
|
|
appending to it). The FTS index follows via the update trigger."""
|
|
conn = self._conn()
|
|
row = conn.execute(
|
|
"SELECT conversation_id FROM messages WHERE id = ?", (message_id,)
|
|
).fetchone()
|
|
if row is None:
|
|
raise ValueError(f"unknown message {message_id}")
|
|
if meta is not None:
|
|
conn.execute(
|
|
"UPDATE messages SET content = ?, meta = ? WHERE id = ?",
|
|
(content, json.dumps(meta, ensure_ascii=False), message_id),
|
|
)
|
|
else:
|
|
conn.execute(
|
|
"UPDATE messages SET content = ? WHERE id = ?", (content, message_id)
|
|
)
|
|
conn.execute(
|
|
"UPDATE conversations SET updated_at = ? WHERE id = ?",
|
|
(_now(), row["conversation_id"]),
|
|
)
|
|
conn.commit()
|
|
|
|
def path_to_root(self, message_id):
|
|
conn = self._conn()
|
|
path = []
|
|
current = message_id
|
|
while current is not None:
|
|
row = conn.execute(
|
|
"SELECT id, parent_id, role, content FROM messages WHERE id = ?",
|
|
(current,),
|
|
).fetchone()
|
|
if row is None:
|
|
break
|
|
path.append({"role": row["role"], "content": row["content"]})
|
|
current = row["parent_id"]
|
|
path.reverse()
|
|
return path
|
|
|
|
def save_frames(self, message_id, frames, layers, k):
|
|
vocab = {}
|
|
packed = []
|
|
for frame in frames:
|
|
vocab[frame["token_id"]] = frame["tok"]
|
|
entry = {
|
|
"pos": frame["pos"],
|
|
"phase": frame["phase"],
|
|
"token_id": frame["token_id"],
|
|
"layers": {},
|
|
}
|
|
for layer, d in frame["layers"].items():
|
|
for tid, s in zip(d["ids"], d["strs"]):
|
|
vocab[tid] = s
|
|
for tid, s in zip(d["m_ids"], d["m_strs"]):
|
|
vocab[tid] = s
|
|
entry["layers"][layer] = {
|
|
"ids": np.asarray(d["ids"], np.int32).tobytes(),
|
|
"p": np.asarray(d["p"], np.float16).tobytes(),
|
|
"m_ids": np.asarray(d["m_ids"], np.int32).tobytes(),
|
|
"m_p": np.asarray(d["m_p"], np.float16).tobytes(),
|
|
"m_rank": np.asarray(d["m_rank"], np.int32).tobytes(),
|
|
}
|
|
packed.append(entry)
|
|
blob = msgpack.packb(
|
|
{
|
|
"version": 1,
|
|
"k": k,
|
|
# generation id of the server-side residual store: lets pins
|
|
# keep working after a page reload (same server session); a
|
|
# restarted server simply reports the store as expired. Last
|
|
# frame: after a continuation merge it's the freshest gen.
|
|
"gen": frames[-1].get("gen") if frames else None,
|
|
"layers": [int(l) for l in layers],
|
|
"frames": packed,
|
|
"vocab": {str(t): s for t, s in vocab.items()},
|
|
}
|
|
)
|
|
filename = f"{message_id}.msgpack"
|
|
(FRAMES_DIR / filename).write_bytes(blob)
|
|
conn = self._conn()
|
|
conn.execute(
|
|
"UPDATE messages SET frames_file = ? WHERE id = ?", (filename, message_id)
|
|
)
|
|
conn.commit()
|
|
return filename
|
|
|
|
def load_frames(self, message_id):
|
|
conn = self._conn()
|
|
row = conn.execute(
|
|
"SELECT frames_file FROM messages WHERE id = ?", (message_id,)
|
|
).fetchone()
|
|
if row is None or row["frames_file"] is None:
|
|
raise ValueError(f"no frames for message {message_id}")
|
|
data = msgpack.unpackb((FRAMES_DIR / row["frames_file"]).read_bytes())
|
|
vocab = data["vocab"]
|
|
frames = []
|
|
for entry in data["frames"]:
|
|
frame = {
|
|
"type": "frame",
|
|
"phase": entry["phase"],
|
|
"pos": entry["pos"],
|
|
"token_id": entry["token_id"],
|
|
"tok": vocab.get(str(entry["token_id"]), ""),
|
|
"gen": data.get("gen"),
|
|
"layers": {},
|
|
}
|
|
for layer, d in entry["layers"].items():
|
|
ids = np.frombuffer(d["ids"], np.int32).tolist()
|
|
m_ids = np.frombuffer(d["m_ids"], np.int32).tolist()
|
|
frame["layers"][layer] = {
|
|
"ids": ids,
|
|
"p": [round(float(v), 5) for v in np.frombuffer(d["p"], np.float16)],
|
|
"strs": [vocab.get(str(t), "") for t in ids],
|
|
"m_ids": m_ids,
|
|
"m_p": [round(float(v), 5) for v in np.frombuffer(d["m_p"], np.float16)],
|
|
"m_rank": np.frombuffer(d["m_rank"], np.int32).tolist(),
|
|
"m_strs": [vocab.get(str(t), "") for t in m_ids],
|
|
}
|
|
frames.append(frame)
|
|
return {"k": data["k"], "layers": data["layers"], "frames": frames}
|
|
|
|
def export(self, conversation_id, fmt="json", include_frames=False):
|
|
conv = self.get_conversation(conversation_id)
|
|
if include_frames:
|
|
for message in conv["messages"]:
|
|
if message["has_frames"]:
|
|
try:
|
|
message["frames"] = self.load_frames(message["id"])
|
|
except ValueError:
|
|
pass
|
|
if fmt == "json":
|
|
return json.dumps(conv, ensure_ascii=False, indent=1), "application/json"
|
|
lines = [f"# {conv['title']}", ""]
|
|
if conv["tags"]:
|
|
lines.append(f"tags: {', '.join(conv['tags'])}")
|
|
lines.append("")
|
|
for message in conv["messages"]:
|
|
meta = message.get("meta") or {}
|
|
head = f"**{message['role']}** (#{message['id']}"
|
|
if message["parent_id"] is not None:
|
|
head += f" ← #{message['parent_id']}"
|
|
head += ")"
|
|
if meta.get("model_id"):
|
|
head += f" — {meta['model_id']} · {meta.get('quant') or meta.get('dtype')}"
|
|
lines.append(head)
|
|
lines.append("")
|
|
lines.append(message["content"])
|
|
lines.append("")
|
|
if include_frames and message.get("frames"):
|
|
lines.append(f"> {len(message['frames']['frames'])} lens frames (layers {message['frames']['layers']})")
|
|
lines.append("")
|
|
return "\n".join(lines), "text/markdown"
|