omega-code 0.4.0__py3-none-any.whl

omega/memory/store.py

@@ -0,0 +1,321 @@
+import os
+import re
+import secrets
+import sqlite3
+import time
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Any
+
+Node = dict[str, Any]
+
+GLOBAL_DIR = Path.home() / ".omega" / "memory"
+
+SCOPES = {"global", "project"}
+TYPES = {"fact", "preference", "decision", "entity", "file_note", "open_question"}
+VOLATILITIES = {"stable", "medium", "volatile"}
+SENSITIVITIES = {"normal", "personal", "sensitive"}
+RELATIONS = {"relates_to", "supersedes", "depends_on", "part_of", "mentions", "contradicts"}
+
+SCHEMA = """
+CREATE TABLE IF NOT EXISTS nodes (
+  id TEXT PRIMARY KEY,
+  type TEXT NOT NULL,
+  title TEXT NOT NULL,
+  body TEXT NOT NULL,
+  confidence REAL DEFAULT 0.8,
+  volatility TEXT DEFAULT 'stable',
+  sensitivity TEXT DEFAULT 'normal',
+  importance REAL DEFAULT 0.5,
+  superseded_by TEXT REFERENCES nodes(id),
+  source_session_id TEXT,
+  created REAL, updated REAL, last_accessed REAL, access_count INTEGER DEFAULT 0
+);
+CREATE TABLE IF NOT EXISTS edges (
+  src TEXT REFERENCES nodes(id), dst TEXT REFERENCES nodes(id),
+  relation TEXT,
+  created REAL,
+  PRIMARY KEY (src, dst, relation)
+);
+CREATE VIRTUAL TABLE IF NOT EXISTS nodes_fts USING fts5(
+  title, body, content='nodes', content_rowid='rowid');
+CREATE TRIGGER IF NOT EXISTS nodes_ai AFTER INSERT ON nodes BEGIN
+  INSERT INTO nodes_fts(rowid, title, body) VALUES (new.rowid, new.title, new.body);
+END;
+CREATE TRIGGER IF NOT EXISTS nodes_ad AFTER DELETE ON nodes BEGIN
+  INSERT INTO nodes_fts(nodes_fts, rowid, title, body) VALUES ('delete', old.rowid, old.title, old.body);
+END;
+CREATE TRIGGER IF NOT EXISTS nodes_au AFTER UPDATE ON nodes BEGIN
+  INSERT INTO nodes_fts(nodes_fts, rowid, title, body) VALUES ('delete', old.rowid, old.title, old.body);
+  INSERT INTO nodes_fts(rowid, title, body) VALUES (new.rowid, new.title, new.body);
+END;
+CREATE TABLE IF NOT EXISTS meta (
+  last_consolidated REAL, nodes_since_consolidation INTEGER
+);
+"""
+
+
+def project_dir(cwd: str | None = None) -> Path:
+    return Path(cwd or os.getcwd()) / ".omega"
+
+
+def _db_path(scope: str, cwd: str | None = None) -> Path:
+    if scope not in SCOPES:
+        raise ValueError(f"invalid scope {scope!r}; expected one of {sorted(SCOPES)}")
+    if scope == "global":
+        return GLOBAL_DIR / "memory.db"
+    return project_dir(cwd) / "memory.db"
+
+
+def db_exists(scope: str, cwd: str | None = None) -> bool:
+    return _db_path(scope, cwd).exists()
+
+
+def _bootstrap_project(cwd: str | None = None) -> None:
+    root = Path(cwd or os.getcwd())
+    project_dir(cwd).mkdir(parents=True, exist_ok=True)
+    # Never touch .gitignore outside an actual repo -- a bare cwd or $HOME
+    # isn't ours to modify.
+    if not (root / ".git").is_dir():
+        return
+    gitignore = root / ".gitignore"
+    line = ".omega/"
+    existing = gitignore.read_text() if gitignore.exists() else ""
+    if line in existing.splitlines():
+        return
+    with gitignore.open("a") as f:
+        if existing and not existing.endswith("\n"):
+            f.write("\n")
+        f.write(line + "\n")
+
+
+def connect(scope: str, cwd: str | None = None) -> sqlite3.Connection:
+    path = _db_path(scope, cwd)
+    if scope == "project" and not path.exists():
+        _bootstrap_project(cwd)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    conn = sqlite3.connect(path)
+    conn.row_factory = sqlite3.Row
+    conn.execute("PRAGMA journal_mode=WAL")
+    conn.executescript(SCHEMA)
+    if conn.execute("SELECT COUNT(*) FROM meta").fetchone()[0] == 0:
+        conn.execute("INSERT INTO meta (last_consolidated, nodes_since_consolidation) "
+                     "VALUES (0, 0)")
+        conn.commit()
+    return conn
+
+
+def _validate(**kwargs: str | None) -> None:
+    checks = {"type": TYPES, "volatility": VOLATILITIES, "sensitivity": SENSITIVITIES,
+              "relation": RELATIONS}
+    for field, value in kwargs.items():
+        if value is not None and value not in checks[field]:
+            raise ValueError(f"invalid {field} {value!r}; expected one of {sorted(checks[field])}")
+
+
+def write_node(scope: str, type: str, title: str, body: str, confidence: float = 0.8,
+               volatility: str = "stable", sensitivity: str = "normal", importance: float = 0.5,
+               source_session_id: str | None = None, cwd: str | None = None) -> str:
+    _validate(type=type, volatility=volatility, sensitivity=sensitivity)
+    conn = connect(scope, cwd)
+    try:
+        node_id = secrets.token_hex(4)
+        now = time.time()
+        conn.execute(
+            "INSERT INTO nodes (id, type, title, body, confidence, volatility, sensitivity, "
+            "importance, source_session_id, created, updated, last_accessed, access_count) "
+            "VALUES (?,?,?,?,?,?,?,?,?,?,?,?,0)",
+            (node_id, type, title, body, confidence, volatility, sensitivity, importance,
+             source_session_id, now, now, now))
+        conn.commit()
+        return node_id
+    finally:
+        conn.close()
+
+
+def add_edge(scope: str, src: str, dst: str, relation: str, cwd: str | None = None) -> None:
+    _validate(relation=relation)
+    conn = connect(scope, cwd)
+    try:
+        conn.execute("INSERT OR IGNORE INTO edges (src, dst, relation, created) VALUES (?,?,?,?)",
+                     (src, dst, relation, time.time()))
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def get(scope: str, id_or_title: str, cwd: str | None = None) -> Node | None:
+    conn = connect(scope, cwd)
+    try:
+        row = conn.execute("SELECT * FROM nodes WHERE id = ?", (id_or_title,)).fetchone()
+        if row is None:
+            row = conn.execute(
+                "SELECT * FROM nodes WHERE lower(title) = lower(?) ORDER BY updated DESC LIMIT 1",
+                (id_or_title,)).fetchone()
+        return dict(row) if row else None
+    finally:
+        conn.close()
+
+
+def _fts_query_variants(query: str) -> Iterator[str]:
+    yield query
+    terms = re.findall(r"\w+", query)
+    if terms:
+        yield " ".join(f'"{t}"' for t in terms)
+
+
+def search(scope: str, query: str, type: str | None = None, limit: int = 8,
+          include_superseded: bool = False, cwd: str | None = None) -> list[Node]:
+    _validate(type=type)
+    conn = connect(scope, cwd)
+    try:
+        sql = ("SELECT n.*, bm25(nodes_fts) AS rank FROM nodes_fts "
+               "JOIN nodes n ON n.rowid = nodes_fts.rowid WHERE nodes_fts MATCH ?")
+        extra: list[str] = []
+        if type is not None:
+            sql += " AND n.type = ?"
+            extra.append(type)
+        if not include_superseded:
+            sql += " AND n.superseded_by IS NULL"
+        sql += " ORDER BY rank LIMIT ?"
+
+        rows = None
+        for variant in _fts_query_variants(query):
+            try:
+                # bm25: smaller (more negative) is a better match, so plain
+                # ascending order (the default) already ranks best-first.
+                rows = conn.execute(sql, (variant, *extra, limit)).fetchall()
+                break
+            except sqlite3.OperationalError:
+                continue
+        return [dict(r) for r in rows] if rows is not None else []
+    finally:
+        conn.close()
+
+
+def neighbors(scope: str, id: str, depth: int = 1, cwd: str | None = None) -> list[Node]:
+    depth = min(max(depth, 0), 2)
+    conn = connect(scope, cwd)
+    try:
+        visited = {id}
+        frontier = [id]
+        results: list[Node] = []
+        seen: set[str] = set()
+        for _ in range(depth):
+            next_frontier = []
+            for nid in frontier:
+                rows = conn.execute(
+                    "SELECT dst AS other, relation, 'out' AS direction FROM edges WHERE src = ? "
+                    "UNION ALL "
+                    "SELECT src AS other, relation, 'in' AS direction FROM edges WHERE dst = ?",
+                    (nid, nid)).fetchall()
+                for r in rows:
+                    other = r["other"]
+                    if other in visited:
+                        continue
+                    node = conn.execute("SELECT * FROM nodes WHERE id = ?", (other,)).fetchone()
+                    if node is None:
+                        continue
+                    visited.add(other)
+                    next_frontier.append(other)
+                    if other not in seen:
+                        seen.add(other)
+                        d = dict(node)
+                        d["relation"], d["direction"] = r["relation"], r["direction"]
+                        results.append(d)
+            frontier = next_frontier
+        return results
+    finally:
+        conn.close()
+
+
+def touch(scope: str, id: str, cwd: str | None = None) -> None:
+    conn = connect(scope, cwd)
+    try:
+        conn.execute("UPDATE nodes SET last_accessed = ?, access_count = access_count + 1 "
+                     "WHERE id = ?", (time.time(), id))
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def mark_superseded(scope: str, old_id: str, new_id: str, cwd: str | None = None) -> None:
+    conn = connect(scope, cwd)
+    try:
+        conn.execute("UPDATE nodes SET superseded_by = ?, updated = ? WHERE id = ?",
+                     (new_id, time.time(), old_id))
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def retag(scope: str, id: str, cwd: str | None = None, **fields: str | float | None) -> None:
+    fields = {k: v for k, v in fields.items() if k in ("volatility", "importance")
+              and v is not None}
+    if not fields:
+        return
+    if "volatility" in fields:
+        _validate(volatility=str(fields["volatility"]))
+    conn = connect(scope, cwd)
+    try:
+        set_clause = ", ".join(f"{k} = ?" for k in fields)
+        conn.execute(f"UPDATE nodes SET {set_clause}, updated = ? WHERE id = ?",
+                     (*fields.values(), time.time(), id))
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def recent(scope: str, since: float | None = None, limit: int = 50,
+          cwd: str | None = None) -> list[Node]:
+    conn = connect(scope, cwd)
+    try:
+        sql = "SELECT * FROM nodes WHERE superseded_by IS NULL"
+        params: list[float] = []
+        if since is not None:
+            sql += " AND created >= ?"
+            params.append(since)
+        sql += " ORDER BY created DESC LIMIT ?"
+        params.append(limit)
+        return [dict(r) for r in conn.execute(sql, params).fetchall()]
+    finally:
+        conn.close()
+
+
+def all_nodes(scope: str, include_superseded: bool = False, cwd: str | None = None) -> list[Node]:
+    conn = connect(scope, cwd)
+    try:
+        sql = "SELECT * FROM nodes"
+        if not include_superseded:
+            sql += " WHERE superseded_by IS NULL"
+        return [dict(r) for r in conn.execute(sql).fetchall()]
+    finally:
+        conn.close()
+
+
+def bump_since_consolidation(scope: str, cwd: str | None = None) -> None:
+    conn = connect(scope, cwd)
+    try:
+        conn.execute("UPDATE meta SET nodes_since_consolidation = nodes_since_consolidation + 1")
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def reset_consolidation(scope: str, cwd: str | None = None) -> None:
+    conn = connect(scope, cwd)
+    try:
+        conn.execute("UPDATE meta SET nodes_since_consolidation = 0, last_consolidated = ?",
+                     (time.time(),))
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def since_consolidation(scope: str, cwd: str | None = None) -> int:
+    conn = connect(scope, cwd)
+    try:
+        row = conn.execute("SELECT nodes_since_consolidation FROM meta").fetchone()
+        return row[0] if row else 0
+    finally:
+        conn.close()

omega/memory/tools.py

@@ -0,0 +1,175 @@
+import re
+
+from .. import tools
+from . import store
+from .store import Node
+
+S = {"type": "string"}
+I = {"type": "integer"}
+F = {"type": "number"}
+
+# Safety net, same shape as permissions.FORBIDDEN_PATTERNS: mislabeling a
+# node's sensitivity must not be the only thing standing between a leaked
+# secret and the preamble every future turn gets injected with.
+SENSITIVE_PATTERNS = [
+    (re.compile(r"\bsk-[A-Za-z0-9]{10,}"), "looks like an API key (sk-...)"),
+    (re.compile(r"\bghp_[A-Za-z0-9]{10,}"), "looks like a GitHub token (ghp_...)"),
+    (re.compile(r"\bxox[bp]-[A-Za-z0-9-]{10,}"), "looks like a Slack token (xox[bp]-...)"),
+    (re.compile(r"\bAKIA[A-Z0-9]{12,}"), "looks like an AWS access key (AKIA...)"),
+    (re.compile(r"-----BEGIN [A-Z ]*PRIVATE KEY"), "contains a private key block"),
+    (re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}"), "contains an email address"),
+    (re.compile(r"\b(?:\+?\d{1,3}[-.\s]?)?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}\b"),
+     "contains a phone number"),
+    (re.compile(r"\b(password|passwd|secret)\s*[:=]", re.I), "contains a password/secret assignment"),
+    # No `/` in the class and letters+digits both required: file paths and
+    # long snake_case identifiers are the normal content of a coding agent's
+    # memory and must not be mistaken for secrets.
+    (re.compile(r"(?=[A-Za-z0-9+=_-]*\d)(?=[A-Za-z0-9+=_-]*[A-Za-z])[A-Za-z0-9+=_-]{32,}"),
+     "contains a long token-like blob"),
+]
+
+
+def _scan_sensitive(*texts: str) -> list[str]:
+    blob = "\n".join(texts)
+    reasons = []
+    for pattern, reason in SENSITIVE_PATTERNS:
+        if pattern.search(blob) and reason not in reasons:
+            reasons.append(reason)
+    return reasons
+
+
+def _resolve(id_or_title: str) -> tuple[str, Node] | tuple[None, None]:
+    """Look up a node across both scopes, project first. Never triggers the
+    project db's .gitignore bootstrap by probing a db that doesn't exist yet."""
+    for scope in ("project", "global"):
+        if scope == "project" and not store.db_exists("project"):
+            continue
+        node = store.get(scope, id_or_title)
+        if node:
+            return scope, node
+    return None, None
+
+
+@tools.tool(
+    "remember",
+    "Save a durable fact about the user or their projects to long-term memory "
+    "so it survives into FUTURE sessions. The current conversation is already "
+    "remembered without this. type: fact|preference|decision|entity|file_note|"
+    "open_question. scope: project (default) or global (true across projects). "
+    "volatility: stable|medium|volatile -- volatile facts (current bug, mood, "
+    "this week's focus) are never auto-injected, only recallable. sensitivity: "
+    "normal|personal|sensitive -- personal facts are rarely injected, sensitive "
+    "never. confidence and importance are 0-1.",
+    {"title": S, "body": S, "type": S, "scope": S, "confidence": F,
+     "volatility": S, "sensitivity": S, "importance": F,
+     "relates_to": {"type": "array", "items": S}},
+    ["title", "body"], mutates=True)
+def _remember(title: str, body: str, type: str = "fact", scope: str = "project",
+             confidence: float = 0.8, volatility: str = "stable", sensitivity: str = "normal",
+             importance: float = 0.5, relates_to: list[str] | None = None) -> str:
+    reasons = _scan_sensitive(title, body)
+    if reasons:
+        sensitivity = "sensitive"
+    try:
+        node_id = store.write_node(scope, type, title, body, confidence=confidence,
+                                   volatility=volatility, sensitivity=sensitivity,
+                                   importance=importance, source_session_id=tools.SESSION_ID)
+    except ValueError as e:
+        return f"error: {e}"
+
+    unresolved: list[str] = []
+    for ref in relates_to or []:
+        target = store.get(scope, ref)
+        if target is None:
+            unresolved.append(ref)
+        else:
+            store.add_edge(scope, node_id, target["id"], "relates_to")
+    store.bump_since_consolidation(scope)
+
+    msg = f"remembered [{node_id}] {title!r} ({scope})"
+    if reasons:
+        msg += f"; sensitivity forced to 'sensitive': {'; '.join(reasons)}"
+    if unresolved:
+        msg += f"; could not resolve relates_to: {', '.join(unresolved)}"
+    return msg
+
+
+@tools.tool(
+    "recall",
+    "Search LONG-TERM memory notes saved across sessions with the `remember` "
+    "tool. This is NOT the conversation: everything said in this session, "
+    "including a resumed one, is already in your context above. Do not call "
+    "this to answer questions about what was just discussed.",
+    {"query": S, "scope": S, "type": S, "depth": I}, ["query"])
+def _recall(query: str, scope: str = "both", type: str | None = None, depth: int = 1) -> str:
+    if scope not in ("both", "project", "global"):
+        return f"error: invalid scope {scope!r}; expected 'both', 'project' or 'global'"
+    scopes = ["project", "global"] if scope == "both" else [scope]
+
+    blocks: list[str] = []
+    for sc in scopes:
+        if sc == "project" and not store.db_exists("project"):
+            continue
+        try:
+            hits = store.search(sc, query, type=type, limit=8)
+        except ValueError as e:
+            return f"error: {e}"
+        for node in hits:
+            store.touch(sc, node["id"])
+            lines = [f"[{node['id']}] {node['type']} · {sc} · conf {node['confidence']} · "
+                     f"{node['volatility']} · importance {node['importance']}",
+                     node["title"], node["body"]]
+            for nb in store.neighbors(sc, node["id"], depth=depth):
+                lines.append(f"  → {nb['relation']} ({nb['direction']}): "
+                             f"[{nb['id']}] {nb['title']}")
+            blocks.append("\n".join(lines))
+
+    if not blocks:
+        return "(no matching memories)"
+    return tools.truncate("\n\n".join(blocks))
+
+
+@tools.tool(
+    "supersede",
+    "Replace an outdated memory node with a corrected one. The old node stops "
+    "appearing in recall/preamble by default but stays queryable by id.",
+    {"old": S, "new_body": S, "confidence": F}, ["old", "new_body"], mutates=True)
+def _supersede(old: str, new_body: str, confidence: float | None = None) -> str:
+    scope, node = _resolve(old)
+    if node is None or scope is None:
+        return f"error: no memory node found matching {old!r}"
+    new_id = store.write_node(
+        scope, node["type"], node["title"], new_body,
+        confidence=confidence if confidence is not None else node["confidence"],
+        volatility=node["volatility"], sensitivity=node["sensitivity"],
+        importance=node["importance"], source_session_id=tools.SESSION_ID)
+    store.add_edge(scope, new_id, node["id"], "supersedes")
+    store.mark_superseded(scope, node["id"], new_id)
+    store.bump_since_consolidation(scope)
+    return f"superseded [{node['id']}] with [{new_id}] ({scope})"
+
+
+@tools.tool(
+    "link",
+    "Add an explicit relation between two existing memory nodes (contradicts, "
+    "depends_on, part_of, mentions, relates_to).",
+    {"a": S, "b": S, "relation": S}, ["a", "b", "relation"], mutates=True)
+def _link(a: str, b: str, relation: str) -> str:
+    scope_a, node_a = _resolve(a)
+    if node_a is None or scope_a is None:
+        return f"error: no memory node found matching {a!r}"
+    # Edges live inside a single sqlite db, so both ends must be in the same
+    # scope -- look b up in a's scope directly rather than resolving it
+    # independently, which could silently pick the wrong-scope node.
+    node_b = store.get(scope_a, b)
+    if node_b is None:
+        scope_b, other = _resolve(b)
+        if other is None:
+            return f"error: no memory node found matching {b!r}"
+        return (f"error: {a!r} is in scope {scope_a!r} but {b!r} is in scope "
+                f"{scope_b!r}; link requires both nodes in the same scope")
+    try:
+        store.add_edge(scope_a, node_a["id"], node_b["id"], relation)
+    except ValueError as e:
+        return f"error: {e}"
+    return f"linked [{node_a['id']}] --{relation}--> [{node_b['id']}] ({scope_a})"

omega/migrate.py

@@ -0,0 +1,40 @@
+"""One-time compatibility copy from the pre-rename `.rig` layout to `.omega`.
+
+Runs from `__main__.cli()` before anything else touches config, permissions,
+or sessions, so every other module can assume `~/.omega` (and `<cwd>/.omega`)
+already hold whatever `.rig` had -- without ever deleting the old directory,
+in case something still points at it."""
+import os
+import shutil
+from pathlib import Path
+
+_PRESERVE_MODE = ("config.json", "permissions.json")
+
+
+def _copy(old: Path, new: Path) -> bool:
+    if new.exists() or not old.is_dir():
+        return False
+    shutil.copytree(old, new)
+    for name in _PRESERVE_MODE:
+        dst = new / name
+        src = old / name
+        if dst.exists() and src.exists():
+            dst.chmod(src.stat().st_mode & 0o777)
+    return True
+
+
+def migrate_home(home: Path | None = None) -> bool:
+    home = home or Path.home()
+    return _copy(home / ".rig", home / ".omega")
+
+
+def migrate_project(cwd: str | None = None) -> bool:
+    root = Path(cwd or os.getcwd())
+    return _copy(root / ".rig", root / ".omega")
+
+
+def run(cwd: str | None = None, home: Path | None = None) -> None:
+    if migrate_home(home):
+        from rich.console import Console
+        Console().print("[dim]migrated ~/.rig → ~/.omega (the old directory was left in place)[/dim]")
+    migrate_project(cwd)