omni-agent-memory 0.1.0__py3-none-any.whl

omni/__init__.py

@@ -0,0 +1,10 @@
+"""OmniService — an independent, namespaced memory service.
+
+OmniMemory packaged as a standalone FastAPI service with a local-LLM (Ollama)
+memory pipeline. Clients (Claude Code via hooks + MCP, or any agent) push raw
+interactions to /ingest and fetch context-scoped memory from /retrieve.
+
+See external_docs/OmniMemory_design.md → "Claude Code Integration".
+"""
+
+__version__ = "0.1.0"

omni/cli.py

@@ -0,0 +1,289 @@
+"""omni — thin CLI client for OmniService (invoked by Claude Code hooks).
+
+Subcommands:
+    omni serve                         start the FastAPI service
+    omni mcp                           start the MCP stdio server
+    omni ingest   [--ns N] [--transcript P]   push new transcript turns (Stop hook)
+    omni retrieve [--ns N] [--query Q | --session-start]   fetch memory (SessionStart hook)
+    omni verify   [--ns N]             force a VERIFY flush
+    omni snapshot [--ns N]             print full assembled memory
+
+Hook usage reads the hook payload JSON from stdin (transcript_path, cwd, prompt).
+Namespace resolution order: --ns  >  stdin cwd  >  $CLAUDE_PROJECT_DIR  >  cwd.
+"""
+
+import argparse
+import json
+import os
+import sys
+from typing import List, Optional
+
+import httpx
+
+from omni import config, storage
+
+
+def _base_url() -> str:
+    return os.environ.get("OMNI_URL", f"http://{config.HOST}:{config.PORT}")
+
+
+def _read_stdin_json() -> dict:
+    if sys.stdin is None or sys.stdin.isatty():
+        return {}
+    raw = sys.stdin.read().strip()
+    if not raw:
+        return {}
+    try:
+        return json.loads(raw)
+    except Exception:
+        return {}
+
+
+def _resolve_ns(arg_ns: Optional[str], hook: dict) -> str:
+    return (arg_ns or hook.get("cwd")
+            or os.environ.get("CLAUDE_PROJECT_DIR") or os.getcwd())
+
+
+def _resolve_client(arg_client: Optional[str]) -> str:
+    return arg_client or config.CLIENT_ID
+
+
+# ---------------------------------------------------------------------------
+# Transcript parsing (Claude Code JSONL)
+# ---------------------------------------------------------------------------
+
+def _content_to_text(content) -> str:
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        chunks = []
+        for block in content:
+            if isinstance(block, dict):
+                if block.get("type") == "text" and block.get("text"):
+                    chunks.append(block["text"])
+                elif "content" in block and isinstance(block["content"], str):
+                    chunks.append(block["content"])
+        return "\n".join(chunks)
+    return ""
+
+
+def parse_transcript_turns(path: str) -> List[dict]:
+    """Extract ordered user/assistant text turns from a Claude Code transcript JSONL."""
+    turns: List[dict] = []
+    if not path or not os.path.exists(path):
+        return turns
+    with open(path) as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                entry = json.loads(line)
+            except Exception:
+                continue
+            msg = entry.get("message") if isinstance(entry, dict) else None
+            if not isinstance(msg, dict):
+                continue
+            role = msg.get("role")
+            if role not in ("user", "assistant"):
+                continue
+            text = _content_to_text(msg.get("content")).strip()
+            if text:
+                turns.append({"role": role, "content": text})
+    return turns
+
+
+# ---------------------------------------------------------------------------
+# Commands
+# ---------------------------------------------------------------------------
+
+def cmd_serve(args):
+    from omni.server import main as serve_main
+    serve_main()
+
+
+def cmd_mcp(args):
+    from omni.mcp_server import main as mcp_main
+    mcp_main()
+
+
+def cmd_ingest(args):
+    hook = _read_stdin_json()
+    ns = _resolve_ns(args.ns, hook)
+    client = _resolve_client(args.client)
+    transcript = args.transcript or hook.get("transcript_path")
+    if not transcript:
+        print(json.dumps({"ok": False, "error": "no transcript_path"}))
+        return
+
+    all_turns = parse_transcript_turns(transcript)
+
+    # Cursor: only ingest turns newer than what we already archived for this transcript.
+    d = storage.ns_dir(client, ns)
+    cursor = storage.load_cursor(d)
+    key = os.path.abspath(transcript)
+    already = int(cursor.get(key, 0))
+    new_turns = all_turns[already:]
+    if not new_turns:
+        print(json.dumps({"ok": True, "ingested": 0, "client_id": client, "namespace": ns}))
+        return
+
+    ts = hook.get("timestamp", "")
+    try:
+        resp = httpx.post(f"{_base_url()}/ingest", json={
+            "client_id": client, "namespace": ns, "turns": new_turns,
+            "timestamp": ts, "source": key,
+        }, timeout=30)
+        resp.raise_for_status()
+    except Exception as e:
+        print(json.dumps({"ok": False, "error": str(e)}))
+        return
+
+    cursor[key] = len(all_turns)
+    storage.save_cursor(d, cursor)
+    print(json.dumps({"ok": True, "ingested": len(new_turns),
+                      "client_id": client, "namespace": ns}))
+
+
+def cmd_retrieve(args):
+    hook = _read_stdin_json()
+    ns = _resolve_ns(args.ns, hook)
+    client = _resolve_client(args.client)
+    if args.session_start:
+        mode, query = "session-start", ""
+    else:
+        mode = "search"
+        query = args.query or hook.get("prompt", "")
+    try:
+        resp = httpx.post(f"{_base_url()}/retrieve", json={
+            "client_id": client, "namespace": ns, "query": query, "mode": mode,
+        }, timeout=30)
+        resp.raise_for_status()
+        context = resp.json().get("context", "")
+    except Exception as e:
+        # Never break the session on a memory miss.
+        sys.stderr.write(f"omni retrieve failed: {e}\n")
+        return
+
+    if args.session_start and not args.raw:
+        print(json.dumps({"hookSpecificOutput": {
+            "hookEventName": "SessionStart",
+            "additionalContext": f"# Project memory (OmniService)\n\n{context}",
+        }}))
+    else:
+        print(context)
+
+
+def cmd_verify(args):
+    hook = _read_stdin_json()
+    ns = _resolve_ns(args.ns, hook)
+    client = _resolve_client(args.client)
+    resp = httpx.post(f"{_base_url()}/verify", json={"client_id": client, "namespace": ns}, timeout=300)
+    print(json.dumps(resp.json()))
+
+
+def cmd_snapshot(args):
+    hook = _read_stdin_json()
+    ns = _resolve_ns(args.ns, hook)
+    client = _resolve_client(args.client)
+    resp = httpx.get(f"{_base_url()}/snapshot",
+                     params={"client_id": client, "namespace": ns}, timeout=30)
+    print(resp.json().get("text", ""))
+
+
+# Claude Code wiring. Capture runs at two points, both cursor-deduped so each transcript
+# turn is ingested exactly once: UserPromptSubmit (capture the user turn on submit — earliest
+# availability + crash resilience) and Stop (capture the completed exchange after each
+# assistant response). SessionStart seeds memory into context.
+_HOOK_BLOCK = {
+    "SessionStart": [{"hooks": [{"type": "command", "command": "omni retrieve --session-start"}]}],
+    "UserPromptSubmit": [{"hooks": [{"type": "command", "command": "omni ingest"}]}],
+    "Stop": [{"hooks": [{"type": "command", "command": "omni ingest"}]}],
+}
+_MCP_OMNI = {"command": "omni", "args": ["mcp"]}
+
+
+def _load_json(path: str) -> dict:
+    try:
+        with open(path) as f:
+            return json.load(f)
+    except Exception:
+        return {}
+
+
+def cmd_install_hooks(args):
+    """Idempotently wire Claude Code hooks + MCP into a project (merges, never clobbers)."""
+    proj = os.path.abspath(args.project or os.environ.get("CLAUDE_PROJECT_DIR") or os.getcwd())
+    sdir = os.path.join(proj, ".claude")
+    os.makedirs(sdir, exist_ok=True)
+    spath = os.path.join(sdir, "settings.json")
+    settings = _load_json(spath)
+    hooks = settings.setdefault("hooks", {})
+    added = []
+    for event, entries in _HOOK_BLOCK.items():
+        cur = hooks.setdefault(event, [])
+        cmd = entries[0]["hooks"][0]["command"]
+        present = any(cmd == h.get("command")
+                      for e in cur if isinstance(e, dict)
+                      for h in e.get("hooks", []) if isinstance(h, dict))
+        if not present:
+            cur.extend(entries)
+            added.append(event)
+    with open(spath, "w") as f:
+        json.dump(settings, f, indent=2)
+    mpath = os.path.join(proj, ".mcp.json")
+    mcfg = _load_json(mpath)
+    servers = mcfg.setdefault("mcpServers", {})
+    mcp_added = "omni" not in servers
+    servers.setdefault("omni", dict(_MCP_OMNI))
+    with open(mpath, "w") as f:
+        json.dump(mcfg, f, indent=2)
+    print(json.dumps({"ok": True, "project": proj, "settings": spath,
+                      "hooks_added": added, "mcp": mpath, "mcp_added": mcp_added}))
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(prog="omni", description="OmniService client")
+    sub = p.add_subparsers(dest="cmd", required=True)
+
+    sub.add_parser("serve", help="start the FastAPI service").set_defaults(func=cmd_serve)
+    sub.add_parser("mcp", help="start the MCP stdio server").set_defaults(func=cmd_mcp)
+
+    pi = sub.add_parser("ingest", help="push new transcript turns")
+    pi.add_argument("--ns")
+    pi.add_argument("--client", help="client id (default: $OMNI_CLIENT_ID or 'claude-code')")
+    pi.add_argument("--transcript")
+    pi.set_defaults(func=cmd_ingest)
+
+    pr = sub.add_parser("retrieve", help="fetch memory context")
+    pr.add_argument("--ns")
+    pr.add_argument("--client")
+    pr.add_argument("--query")
+    pr.add_argument("--session-start", action="store_true")
+    pr.add_argument("--raw", action="store_true", help="print context only (no hook JSON)")
+    pr.set_defaults(func=cmd_retrieve)
+
+    pv = sub.add_parser("verify", help="force a VERIFY flush")
+    pv.add_argument("--ns")
+    pv.add_argument("--client")
+    pv.set_defaults(func=cmd_verify)
+
+    ps = sub.add_parser("snapshot", help="print full assembled memory")
+    ps.add_argument("--ns")
+    ps.add_argument("--client")
+    ps.set_defaults(func=cmd_snapshot)
+
+    ph = sub.add_parser("install-hooks",
+                        help="wire Claude Code hooks + MCP into a project (idempotent)")
+    ph.add_argument("--project", help="project dir (default: $CLAUDE_PROJECT_DIR or cwd)")
+    ph.set_defaults(func=cmd_install_hooks)
+    return p
+
+
+def main():
+    args = build_parser().parse_args()
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

omni/config.py

@@ -0,0 +1,64 @@
+"""Configuration for OmniService.
+
+All values can be overridden via environment variables (OMNI_*). Defaults are
+chosen to sit alongside Ollama (:11434) without collisions.
+"""
+
+import os
+from pathlib import Path
+
+
+def _env(name: str, default: str) -> str:
+    return os.environ.get(name, default)
+
+
+# Storage root — service-owned, NOT under ~/.claude. One subdir per namespace.
+STORAGE_ROOT = Path(_env("OMNI_STORAGE_ROOT", str(Path.home() / ".omni"))).expanduser()
+
+# HTTP service
+HOST = _env("OMNI_HOST", "127.0.0.1")
+PORT = int(_env("OMNI_PORT", "11435"))
+
+# Client identity (required on every request; partitions storage per client).
+# CLI/MCP clients default to this; test runs use a distinct id (e.g. "test").
+CLIENT_ID = _env("OMNI_CLIENT_ID", "claude-code")
+
+# Local LLM (Ollama, OpenAI-compatible endpoint)
+OLLAMA_BASE_URL = _env("OMNI_OLLAMA_BASE_URL", "http://localhost:11434/v1")
+# Gemma 4 E4B. NOTE: the OpenAI /v1 path runs Ollama at its DEFAULT num_ctx (4096) and
+# silently truncates long prompts; a baked-in-context variant is required to use more.
+# Create once:  printf 'FROM gemma4:e4b\nPARAMETER num_ctx 32768\n' > Modelfile && \
+#               ollama create gemma4-ctx32k -f Modelfile
+EXTRACT_MODEL = _env("OMNI_EXTRACT_MODEL", "gemma4-ctx32k")   # EXTRACT / RELATE / COMPRESS
+VERIFY_MODEL = _env("OMNI_VERIFY_MODEL", "gemma4-ctx32k")     # VERIFY (stronger reasoning)
+
+# Ensemble answering (docs/ensemble_design.md): K context strategies -> answers -> rerank.
+ANSWER_MODEL = _env("OMNI_ANSWER_MODEL", "claude-code")
+RERANK_MODEL = _env("OMNI_RERANK_MODEL", "claude-code")
+ENSEMBLE_STRATEGIES = _env("OMNI_ENSEMBLE_STRATEGIES", "search,raw,chain,uncertain").split(",")
+
+# Fine-tuned student served via an OpenAI-compatible endpoint (training/serve_student.py).
+# When a memory-op model name == STUDENT_MODEL and STUDENT_URL is set, call_local routes
+# there instead of Ollama. Embeddings still use Ollama.
+STUDENT_MODEL = _env("OMNI_STUDENT_MODEL", "omnimem-ops")
+STUDENT_URL = _env("OMNI_STUDENT_URL", "")  # e.g. http://127.0.0.1:11437/v1
+
+# Background pipeline tuning
+# EXTRACT over a fact-dense session can emit many records; too small a cap truncates the
+# JSON mid-stream. Parsing is now salvage-tolerant, but a roomier cap avoids the loss.
+EXTRACT_MAX_TOKENS = int(_env("OMNI_EXTRACT_MAX_TOKENS", "4096"))
+VERIFY_MAX_TOKENS = int(_env("OMNI_VERIFY_MAX_TOKENS", "3072"))
+EXTRACT_WORKERS = int(_env("OMNI_EXTRACT_WORKERS", "4"))
+VERIFY_DEBOUNCE_SECONDS = float(_env("OMNI_VERIFY_DEBOUNCE_SECONDS", "20"))
+VERIFY_RAW_WINDOW = int(_env("OMNI_VERIFY_RAW_WINDOW", "8"))  # recent raw files VERIFY re-reads
+VERIFY_CHUNK = int(_env("OMNI_VERIFY_CHUNK", "6"))  # raw sessions per focused VERIFY call
+
+# Retrieval shaping
+RETRIEVE_TOP_PAGES = int(_env("OMNI_RETRIEVE_TOP_PAGES", "8"))
+RETRIEVE_TOP_HISTORY = int(_env("OMNI_RETRIEVE_TOP_HISTORY", "30"))
+
+# Vector index (semantic retrieval + Tr revision-chain reconstruction)
+VECTOR_ENABLED = _env("OMNI_VECTOR_ENABLED", "1") == "1"
+EMBED_MODEL = _env("OMNI_EMBED_MODEL", "nomic-embed-text")
+VECTOR_TOPK = int(_env("OMNI_VECTOR_TOPK", "12"))
+RECONSTRUCT_MAX_ENTITIES = int(_env("OMNI_RECONSTRUCT_MAX_ENTITIES", "12"))