memory.py

Go to the documentation of this file.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# oc_chat/tools/memory.py
 
"""
Persistent memory via ChromaDB vector store.
 
Purpose
-------
Provides :func:`remember` and :func:`recall` for long-term semantic
memory across REPL sessions. Backed by a local ChromaDB persistent
client at ``./memory_store/`` with cosine-similarity search.
 
Falls back to the structured ``missing_dependency`` shape that the
dependency healer can intercept when ``chromadb`` is not installed.
 
Responsibilities
----------------
- :func:`_get_collection` — lazily build (or reload) the persistent
  collection on first access.
- :func:`remember` — embed and store a text chunk with optional
  metadata; returns the assigned id.
- :func:`recall` — semantic-similarity query returning the top-k
  matches with their metadata and distances.
 
Requires: ``pip install chromadb``
 
Diagnostics
-----------
Domain: OC_CHAT.TOOLS
Levels:
    L2 — lifecycle events
    L3 — semantic details
    L4 — deep tracing
 
Contracts
---------
- The collection is process-singleton — repeat calls reuse the same
  client/collection handles.
- Missing ``chromadb`` returns ``{"missing_dependency": "chromadb",
  ...}`` rather than raising, so the healer can offer to install.
- The on-disk store at ``memory_store/`` is opaque to callers; only
  the public ``remember`` / ``recall`` API is supported.
 
Dependencies
------------
- oc_diagnostics (structured debug logging via _dbg)
- oc_chat.tools.filesystem (lazy; _sandbox_root - places the memory store under the effective sandbox root)
- chromadb (lazy; the persistent vector store client and its cosine-similarity collection, optional)
 
See Also
--------
oc_chat/docs/conventions/CONTRACTS.md §3 Tool registry contract; §4 Tool return-
value conventions.
"""
from pathlib import Path
 
from oc_diagnostics import _dbg
 
def _store_dir() -> str:
    """Resolve the memory store dir under the effective sandbox root.
 
    Resolved at first use (not import time) so a session sandbox override set
    after import is honoured — keeping the long-term memory store aligned with
    the read/write sandbox (M8 F77; M8.0 WU-1) instead of the import-time
    process cwd.
    """
    from oc_chat.tools.filesystem import _sandbox_root
 
    return str(_sandbox_root() / "memory_store")
_COLLECTION_NAME = "oc_chat_memory"
_client = None
_collection = None
 
 
def _get_collection():
    """Return (or lazily create) the ChromaDB collection."""
    func = "_get_collection"
    global _client, _collection
    if _collection is not None:
        return _collection
    try:
        import chromadb
    except ImportError:
        raise ImportError("chromadb not installed. Run: pip install chromadb")
 
    store_dir = _store_dir()
    _dbg("OC_CHAT.TOOLS", func, f"first-time init at {store_dir}", 2)
    _client = chromadb.PersistentClient(path=store_dir)
    _collection = _client.get_or_create_collection(
        name=_COLLECTION_NAME,
        metadata={"hnsw:space": "cosine"},
    )
    return _collection
 
 
def remember(text: str, metadata: dict | None = None) -> dict:
    """Store a piece of information in long-term persistent memory."""
    try:
        col = _get_collection()
    except ImportError as e:
        return {"error": str(e)}
 
    import uuid
    doc_id = str(uuid.uuid4())
    col.add(documents=[text], ids=[doc_id], metadatas=[metadata or {}])
    return {"stored": True, "id": doc_id}
 
 
def recall(query: str, n_results: int = 5) -> dict:
    """Retrieve relevant memories using semantic similarity search."""
    try:
        col = _get_collection()
    except ImportError as e:
        return {"error": str(e)}
 
    count = col.count()
    if count == 0:
        return {"memories": [], "note": "Memory store is empty."}
 
    results = col.query(
        query_texts=[query],
        n_results=min(n_results, count),
    )
    memories = []
    for doc, meta in zip(
        results.get("documents", [[]])[0],
        results.get("metadatas", [[]])[0],
    ):
        memories.append({"text": doc, "metadata": meta})
 
    return {"memories": memories}
 
 
TOOLS = {
    "remember": remember,
    "recall":   recall,
}
 
SCHEMAS = [
    {
        "type": "function",
        "function": {
            "name": "remember",
            "description": "Store a piece of information in long-term persistent memory.",
            "parameters": {
                "type": "object",
                "properties": {
                    "text":     {"type": "string", "description": "Information to remember"},
                    "metadata": {"type": "object", "description": "Optional key-value metadata"},
                },
                "required": ["text"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "recall",
            "description": "Retrieve relevant memories using semantic similarity search.",
            "parameters": {
                "type": "object",
                "properties": {
                    "query":     {"type": "string",  "description": "Search query"},
                    "n_results": {"type": "integer", "description": "Number of memories to return (default 5)"},
                },
                "required": ["query"],
            },
        },
    },
]