""" SQLite Memory Store — sessions, papers, insights, artifacts, metrics, conversation. """ import sqlite3 import json import logging from pathlib import Path from typing import Optional from datetime import datetime logger = logging.getLogger(__name__) DB_PATH = "data/research_memory.db" class MemoryStore: def __init__(self, db_path: str = DB_PATH): self.db_path = Path(db_path) self.db_path.parent.mkdir(parents=True, exist_ok=True) self.conn = sqlite3.connect(str(self.db_path), check_same_thread=False) self.conn.row_factory = sqlite3.Row self._create_tables() logger.info(f"[MemoryStore] Connected to {self.db_path}") def _create_tables(self): self.conn.executescript(""" CREATE TABLE IF NOT EXISTS sessions ( session_id TEXT PRIMARY KEY, query TEXT NOT NULL, subtopics TEXT, filters TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS papers ( id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT NOT NULL, paper_id TEXT NOT NULL, title TEXT, authors TEXT, abstract TEXT, published TEXT, arxiv_url TEXT, categories TEXT, citation_count INTEGER DEFAULT 0, venue TEXT, final_score REAL DEFAULT 0.0, score_breakdown TEXT, UNIQUE(session_id, paper_id) ); CREATE TABLE IF NOT EXISTS insights ( id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT NOT NULL, paper_id TEXT NOT NULL, insight_type TEXT, content TEXT, created_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS artifacts ( id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT NOT NULL, artifact_type TEXT, content TEXT, file_path TEXT, created_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS metrics ( id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT NOT NULL, query_time_sec REAL, papers_retrieved INTEGER, papers_selected INTEGER, score_mean REAL, score_std REAL, user_rating INTEGER, created_at TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS conversation ( id INTEGER PRIMARY KEY AUTOINCREMENT, session_id TEXT NOT NULL, role TEXT, message TEXT, created_at TEXT NOT NULL ); """) self.conn.commit() # --- Sessions --- def create_session(self, session_id: str, query: str, subtopics: list = None): now = datetime.utcnow().isoformat() self.conn.execute(""" INSERT INTO sessions (session_id,query,subtopics,filters,created_at,updated_at) VALUES (?,?,?,?,?,?) ON CONFLICT(session_id) DO UPDATE SET query=excluded.query, subtopics=excluded.subtopics, updated_at=excluded.updated_at """, (session_id, query, json.dumps(subtopics or []), json.dumps({}), now, now)) self.conn.commit() def get_session(self, session_id: str) -> Optional[dict]: row = self.conn.execute( "SELECT * FROM sessions WHERE session_id=?", (session_id,) ).fetchone() return dict(row) if row else None def update_session_filters(self, session_id: str, filters: dict): self.conn.execute( "UPDATE sessions SET filters=?,updated_at=? WHERE session_id=?", (json.dumps(filters), datetime.utcnow().isoformat(), session_id) ) self.conn.commit() # --- Papers --- def save_papers(self, session_id: str, papers: list[dict]): for p in papers: self.conn.execute(""" INSERT OR REPLACE INTO papers (session_id,paper_id,title,authors,abstract,published,arxiv_url, categories,citation_count,venue,final_score,score_breakdown) VALUES (?,?,?,?,?,?,?,?,?,?,?,?) """, (session_id, p.get("paper_id",""), p.get("title",""), json.dumps(p.get("authors",[])), p.get("abstract",""), p.get("published",""), p.get("arxiv_url",""), json.dumps(p.get("categories",[])), p.get("citation_count",0), p.get("venue",""), p.get("final_score",0.0), json.dumps(p.get("score_breakdown",{})))) self.conn.commit() def get_papers(self, session_id: str) -> list[dict]: rows = self.conn.execute( "SELECT * FROM papers WHERE session_id=? ORDER BY final_score DESC", (session_id,) ).fetchall() results = [] for row in rows: d = dict(row) d["authors"] = json.loads(d["authors"] or "[]") d["categories"] = json.loads(d["categories"] or "[]") d["score_breakdown"] = json.loads(d["score_breakdown"] or "{}") results.append(d) return results # --- Insights --- def save_insight(self, session_id: str, paper_id: str, insight_type: str, content: str): self.conn.execute( "INSERT INTO insights (session_id,paper_id,insight_type,content,created_at) VALUES (?,?,?,?,?)", (session_id, paper_id, insight_type, content, datetime.utcnow().isoformat()) ) self.conn.commit() def get_insights(self, session_id: str) -> list[dict]: return [dict(r) for r in self.conn.execute( "SELECT * FROM insights WHERE session_id=?", (session_id,) ).fetchall()] # --- Artifacts --- def save_artifact(self, session_id: str, artifact_type: str, content: str, file_path: str = ""): self.conn.execute( "INSERT INTO artifacts (session_id,artifact_type,content,file_path,created_at) VALUES (?,?,?,?,?)", (session_id, artifact_type, content, file_path, datetime.utcnow().isoformat()) ) self.conn.commit() def get_artifact(self, session_id: str, artifact_type: str) -> Optional[str]: row = self.conn.execute( "SELECT content FROM artifacts WHERE session_id=? AND artifact_type=? ORDER BY id DESC LIMIT 1", (session_id, artifact_type) ).fetchone() return row["content"] if row else None # --- Metrics --- def save_metrics(self, session_id: str, metrics: dict): self.conn.execute(""" INSERT INTO metrics (session_id,query_time_sec,papers_retrieved,papers_selected, score_mean,score_std,user_rating,created_at) VALUES (?,?,?,?,?,?,?,?) """, (session_id, metrics.get("query_time_sec",0), metrics.get("papers_retrieved",0), metrics.get("papers_selected",0), metrics.get("score_mean",0), metrics.get("score_std",0), metrics.get("user_rating"), datetime.utcnow().isoformat())) self.conn.commit() def get_metrics(self, session_id: str) -> list[dict]: return [dict(r) for r in self.conn.execute( "SELECT * FROM metrics WHERE session_id=?", (session_id,) ).fetchall()] # --- Conversation --- def add_message(self, session_id: str, role: str, message: str): self.conn.execute( "INSERT INTO conversation (session_id,role,message,created_at) VALUES (?,?,?,?)", (session_id, role, message, datetime.utcnow().isoformat()) ) self.conn.commit() def get_conversation(self, session_id: str, last_n: int = 20) -> list[dict]: rows = self.conn.execute(""" SELECT role,message,created_at FROM conversation WHERE session_id=? ORDER BY id DESC LIMIT ? """, (session_id, last_n)).fetchall() return list(reversed([dict(r) for r in rows])) # --- Utility --- def list_sessions(self) -> list[dict]: return [dict(r) for r in self.conn.execute( "SELECT session_id,query,created_at FROM sessions ORDER BY created_at DESC" ).fetchall()] def close(self): self.conn.close()