Cleanup init commit
This commit is contained in:
Extraltodeus
+378
@@ -0,0 +1,378 @@
|
||||
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"
|
||||
Reference in New Issue
Block a user