polymath-agent 0.4.0__py3-none-any.whl

polymath/context_store.py

@@ -0,0 +1,336 @@
+"""
+SQLite-backed unified context store.
+All models share one thread per session.
+"""
+from __future__ import annotations
+
+import sqlite3
+import uuid
+from contextlib import contextmanager
+from datetime import datetime, timezone, timedelta
+from pathlib import Path
+
+from polymath.config import DB_FILE
+from polymath.adapters.base import Message
+
+# Single persistent connection — WAL mode for async safety, timeout so
+# writers wait instead of immediately raising OperationalError: database is locked
+_db_conn: sqlite3.Connection | None = None
+
+
+def _conn() -> sqlite3.Connection:
+    global _db_conn
+    DB_FILE.parent.mkdir(parents=True, exist_ok=True)
+    if _db_conn is None:
+        _db_conn = sqlite3.connect(
+            str(DB_FILE),
+            timeout=30,
+            check_same_thread=False,
+        )
+        _db_conn.row_factory = sqlite3.Row
+        _db_conn.execute("PRAGMA journal_mode=WAL")
+        _db_conn.execute("PRAGMA synchronous=NORMAL")
+        _db_conn.commit()
+    return _db_conn
+
+
+def init_db() -> None:
+    with _conn() as conn:
+        conn.executescript("""
+        CREATE TABLE IF NOT EXISTS sessions (
+            id          TEXT PRIMARY KEY,
+            title       TEXT,
+            created_at  TEXT,
+            updated_at  TEXT,
+            project     TEXT DEFAULT ''
+        );
+        -- migrate: add project column if upgrading from older db
+        CREATE TABLE IF NOT EXISTS _migrations (key TEXT PRIMARY KEY);
+
+
+        CREATE TABLE IF NOT EXISTS messages (
+            id          INTEGER PRIMARY KEY AUTOINCREMENT,
+            session_id  TEXT NOT NULL,
+            role        TEXT NOT NULL,
+            content     TEXT NOT NULL,
+            model       TEXT DEFAULT '',
+            step        TEXT DEFAULT '',
+            token_count INTEGER DEFAULT 0,
+            created_at  TEXT NOT NULL,
+            FOREIGN KEY (session_id) REFERENCES sessions(id)
+        );
+
+        CREATE TABLE IF NOT EXISTS summaries (
+            session_id          TEXT PRIMARY KEY,
+            summary             TEXT NOT NULL,
+            summarized_through  INTEGER,
+            updated_at          TEXT,
+            FOREIGN KEY (session_id) REFERENCES sessions(id)
+        );
+        """)
+        # Safe migration for existing databases
+        try:
+            conn.execute("ALTER TABLE sessions ADD COLUMN project TEXT DEFAULT ''")
+        except Exception:
+            pass  # column already exists
+
+
+# ── Sessions ──────────────────────────────────────────────────────────────────
+
+def new_session(title: str = "") -> str:
+    sid = str(uuid.uuid4())[:8]
+    now = datetime.now(timezone.utc).isoformat()
+    with _conn() as conn:
+        conn.execute(
+            "INSERT INTO sessions (id, title, created_at, updated_at) VALUES (?,?,?,?)",
+            (sid, title or f"Session {sid}", now, now),
+        )
+    return sid
+
+
+def list_sessions() -> list[dict]:
+    with _conn() as conn:
+        rows = conn.execute(
+            "SELECT id, title, created_at, updated_at FROM sessions ORDER BY updated_at DESC LIMIT 20"
+        ).fetchall()
+    return [dict(r) for r in rows]
+
+
+def get_session(sid: str) -> dict | None:
+    with _conn() as conn:
+        row = conn.execute("SELECT * FROM sessions WHERE id=?", (sid,)).fetchone()
+    return dict(row) if row else None
+
+
+def touch_session(sid: str, title: str = "") -> None:
+    now = datetime.now(timezone.utc).isoformat()
+    with _conn() as conn:
+        if title:
+            conn.execute("UPDATE sessions SET updated_at=?, title=? WHERE id=?", (now, title, sid))
+        else:
+            conn.execute("UPDATE sessions SET updated_at=? WHERE id=?", (now, sid))
+
+
+def set_session_project(sid: str, project: str) -> None:
+    """Tag an existing session to a project."""
+    now = datetime.now(timezone.utc).isoformat()
+    with _conn() as conn:
+        conn.execute(
+            "UPDATE sessions SET project=?, updated_at=? WHERE id=?",
+            (project, now, sid),
+        )
+
+
+# ── Messages ──────────────────────────────────────────────────────────────────
+
+def add_message(
+    session_id: str,
+    role: str,
+    content: str,
+    model: str = "",
+    step: str = "",
+    token_count: int = 0,
+) -> int:
+    now = datetime.now(timezone.utc).isoformat()
+    with _conn() as conn:
+        cur = conn.execute(
+            "INSERT INTO messages (session_id,role,content,model,step,token_count,created_at) VALUES (?,?,?,?,?,?,?)",
+            (session_id, role, content, model, step, token_count, now),
+        )
+        touch_session(session_id)
+        return cur.lastrowid
+
+
+def get_messages(session_id: str, limit: int = 0) -> list[Message]:
+    with _conn() as conn:
+        if limit:
+            rows = conn.execute(
+                "SELECT * FROM messages WHERE session_id=? ORDER BY id DESC LIMIT ?",
+                (session_id, limit),
+            ).fetchall()
+            rows = list(reversed(rows))
+        else:
+            rows = conn.execute(
+                "SELECT * FROM messages WHERE session_id=? ORDER BY id ASC",
+                (session_id,),
+            ).fetchall()
+    return [Message(role=r["role"], content=r["content"], model=r["model"], step=r["step"]) for r in rows]
+
+
+def get_all_messages_raw(session_id: str) -> list[dict]:
+    with _conn() as conn:
+        rows = conn.execute(
+            "SELECT * FROM messages WHERE session_id=? ORDER BY id ASC", (session_id,)
+        ).fetchall()
+    return [dict(r) for r in rows]
+
+
+def last_message_id(session_id: str) -> int:
+    with _conn() as conn:
+        row = conn.execute(
+            "SELECT MAX(id) as max_id FROM messages WHERE session_id=?", (session_id,)
+        ).fetchone()
+    return row["max_id"] or 0
+
+
+# ── Summaries ─────────────────────────────────────────────────────────────────
+
+def get_summary(session_id: str) -> dict | None:
+    with _conn() as conn:
+        row = conn.execute("SELECT * FROM summaries WHERE session_id=?", (session_id,)).fetchone()
+    return dict(row) if row else None
+
+
+def save_summary(session_id: str, summary: str, through_id: int) -> None:
+    now = datetime.now(timezone.utc).isoformat()
+    with _conn() as conn:
+        conn.execute(
+            """INSERT INTO summaries (session_id, summary, summarized_through, updated_at)
+               VALUES (?,?,?,?)
+               ON CONFLICT(session_id) DO UPDATE SET
+                   summary=excluded.summary,
+                   summarized_through=excluded.summarized_through,
+                   updated_at=excluded.updated_at""",
+            (session_id, summary, through_id, now),
+        )
+
+
+def auto_name_session(session_id: str, first_message: str) -> None:
+    """Generate a human-readable session name from the first user message."""
+    words = first_message.strip().split()
+    name = " ".join(words[:6])
+    if len(words) > 6:
+        name += "..."
+    touch_session(session_id, title=name)
+
+
+def get_recent_sessions(project: str = "", limit: int = 10) -> list[dict]:
+    """Return sessions ordered by last active, optionally filtered by project."""
+    with _conn() as conn:
+        if project:
+            rows = conn.execute(
+                "SELECT * FROM sessions WHERE project=? ORDER BY updated_at DESC LIMIT ?",
+                (project, limit),
+            ).fetchall()
+        else:
+            rows = conn.execute(
+                "SELECT * FROM sessions ORDER BY updated_at DESC LIMIT ?",
+                (limit,),
+            ).fetchall()
+    return [dict(r) for r in rows]
+
+
+def get_last_session_for_project(project: str) -> dict | None:
+    """Return the most recently active session tagged to a project."""
+    with _conn() as conn:
+        row = conn.execute(
+            "SELECT * FROM sessions WHERE project=? ORDER BY updated_at DESC LIMIT 1",
+            (project,),
+        ).fetchone()
+    return dict(row) if row else None
+
+
+def was_recent(session: dict, hours: int = 24) -> bool:
+    """Return True if session was active within the last N hours."""
+    try:
+        updated = datetime.fromisoformat(session["updated_at"])
+        if updated.tzinfo is None:
+            updated = updated.replace(tzinfo=timezone.utc)
+        return datetime.now(timezone.utc) - updated < timedelta(hours=hours)
+    except Exception:
+        return False
+
+
+def time_ago(session: dict) -> str:
+    """Human-readable time since last activity."""
+    try:
+        updated = datetime.fromisoformat(session["updated_at"])
+        if updated.tzinfo is None:
+            updated = updated.replace(tzinfo=timezone.utc)
+        diff = datetime.now(timezone.utc) - updated
+        minutes = int(diff.total_seconds() / 60)
+        if minutes < 60:
+            return f"{minutes}m ago"
+        hours = minutes // 60
+        if hours < 24:
+            return f"{hours}h ago"
+        days = hours // 24
+        return f"{days}d ago"
+    except Exception:
+        return "unknown"
+
+
+def new_session_for_project(project: str) -> str:
+    """Create a new session tagged to a project."""
+    sid = str(uuid.uuid4())[:8]
+    now = datetime.now(timezone.utc).isoformat()
+    with _conn() as conn:
+        conn.execute(
+            "INSERT INTO sessions (id, title, created_at, updated_at, project) VALUES (?,?,?,?,?)",
+            (sid, f"Session {sid}", now, now, project),
+        )
+    return sid
+
+
+def session_cost(session_id: str) -> tuple[int, float]:
+    """Returns (total_tokens, estimated_usd). Rough estimate only."""
+    msgs = get_all_messages_raw(session_id)
+    total_tokens = sum(m.get("token_count", 0) for m in msgs)
+    usd = total_tokens * 0.003 / 1000
+    return total_tokens, usd
+
+
+def delete_session(session_id: str) -> bool:
+    """Delete a session and all its messages. Returns True if found and deleted."""
+    with _conn() as conn:
+        row = conn.execute("SELECT id FROM sessions WHERE id=?", (session_id,)).fetchone()
+        if not row:
+            return False
+        conn.execute("DELETE FROM messages WHERE session_id=?", (session_id,))
+        conn.execute("DELETE FROM summaries WHERE session_id=?", (session_id,))
+        conn.execute("DELETE FROM sessions WHERE id=?", (session_id,))
+    return True
+
+
+def export_session_md(session_id: str) -> str:
+    """Export session to a Markdown file in SESSIONS_DIR and return the path."""
+    from polymath.config import SESSIONS_DIR
+    SESSIONS_DIR.mkdir(parents=True, exist_ok=True)
+
+    session = get_session(session_id)
+    messages = get_messages(session_id)
+    if not session:
+        return ""
+
+    filename = f"{session_id}_{session['title'].replace(' ', '_')}.md"
+    file_path = SESSIONS_DIR / filename
+
+    lines = [
+        f"# {session['title']}",
+        f"Session ID: {session_id}",
+        f"Created: {session['created_at']}",
+        f"Updated: {session['updated_at']}",
+        "\n---\n",
+    ]
+
+    for m in messages:
+        role = m.role.upper()
+        if m.model:
+            role = f"{role} ({m.model})"
+        
+        lines.append(f"### {role}")
+        if m.step:
+            lines.append(f"*Step: {m.step}*")
+        
+        lines.append(f"\n{m.content}\n")
+        
+        if m.tool_calls:
+            lines.append("#### Tool Calls:")
+            for tc in m.tool_calls:
+                lines.append(f"- `{tc.name}({tc.arguments})`")
+            lines.append("\n")
+
+    with open(file_path, "w") as f:
+        f.write("\n".join(lines))
+    
+    return str(file_path)