cortexlayer 0.1.0__py3-none-any.whl

cortexlayer/memory.py

@@ -0,0 +1,372 @@
+"""``Memory`` — the embedded, in-process Cortex engine (no server).
+
+>>> from cortexlayer import Memory                                   # doctest: +SKIP
+>>> m = Memory()                                                     # doctest: +SKIP
+>>> m.add("I moved to Lisbon in March", user_id="alice")             # doctest: +SKIP
+>>> m.search("where does alice live?", user_id="alice")              # doctest: +SKIP
+
+Same method names and result types as :class:`~cortexlayer.CortexClient`, so
+switching between local and hosted is a change of constructor. Memories live in
+an embedded Chroma store under ``data_dir``; each ``user_id`` gets its own
+collection, so users can never see each other's pages.
+
+Needs ``pip install "cortexlayer[local]"``. Importing this module does not —
+the engine is only loaded when a ``Memory`` is constructed.
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+from typing import Any, Callable, Dict, List, Optional
+
+from ._validate import need_int, need_str
+from .errors import CortexConfigError, InvalidRequestError, LocalDependencyError, NotFoundError
+from .types import Answer, AddResult, Page, PageList, SearchResult
+
+DEFAULT_USER_ID = "default"
+DATA_DIR_ENV = "CORTEXLAYER_DATA_DIR"
+MAX_SEARCH_LIMIT = 20
+MAX_LIST_LIMIT = 500
+BACKENDS = ("raw", "facts")
+_CONFIG_KEYS = {
+    "backend", "data_dir", "entity_extractor", "spacy_model",
+    "default_user_id", "auto_relink", "llm", "embedder", "custom_instructions",
+    "observation_date_from_timestamp", "keyword_scoring",
+}
+
+
+def _default_data_dir() -> str:
+    return os.environ.get(DATA_DIR_ENV) or os.path.join(
+        os.path.expanduser("~"), ".cortexlayer"
+    )
+
+
+def _load_engine(backend: str = "raw") -> Any:
+    """Import the engine, or explain how to install it."""
+    try:
+        from ._engine import compression, ingestion, linking, nlp, retrieval, shaping, storage
+        if backend == "facts":
+            from ._engine.facts import backends as fact_backends
+            from ._engine.facts import engine as fact_engine
+    except ImportError as e:
+        raise LocalDependencyError(
+            "The embedded engine needs extra packages. Run: "
+            'pip install "cortexlayer[local]"  '
+            f"(missing: {getattr(e, 'name', None) or e})."
+        ) from e
+
+    class _Engine:
+        pass
+
+    eng = _Engine()
+    for mod in (compression, ingestion, linking, nlp, retrieval, shaping, storage):
+        setattr(eng, mod.__name__.rsplit(".", 1)[-1], mod)
+    if backend == "facts":
+        eng.fact_backends, eng.fact_engine = fact_backends, fact_engine
+    return eng
+
+
+class Memory:
+    """Embedded memory store.
+
+    Args:
+        data_dir: where the embedded store lives (default ``~/.cortexlayer``,
+            or ``$CORTEXLAYER_DATA_DIR``). It is created on first use.
+        entity_extractor: ``"auto"`` (default: spaCy if available, else a
+            simpler regex extractor with a one-time warning), ``"spacy"``,
+            ``"regex"``, or your own object with ``entities(text)`` and
+            ``sentences(text)`` methods. Entities drive linking, so a weaker
+            extractor means fewer links.
+        spacy_model: spaCy model name (default ``en_core_web_sm``).
+        default_user_id: used when a call omits ``user_id``.
+        auto_relink: re-run the linking pass after every ``add``. Off by
+            default — linking is a batch pass over all pages, so prefer
+            calling :meth:`relink` after adding several memories.
+        backend: ``"raw"`` (default) stores your text as small pages — adding
+            never needs an LLM. ``"facts"`` asks an LLM to distil each ``add``
+            into self-contained facts (Mem0-style), stores those, and boosts
+            search by shared entities; it needs an LLM (``llm=``, Ollama by
+            default) and an embedder (``embedder=``, Chroma's built-in by
+            default).
+        llm: (``"facts"``) ``None`` = Ollama at ``$OLLAMA_HOST``; a dict such
+            as ``{"model": "qwen3.5:9b", "host": "..."}``; a callable
+            ``fn(system, user) -> str``; or any object with ``generate()``.
+        embedder: (``"facts"``) ``None``/``"chroma"`` (default), ``"ollama"``,
+            a dict like ``{"provider": "ollama", "model": "qwen3-embedding:8b"}``,
+            or an object with ``embed_batch(texts, action)`` and ``name``. A
+            store remembers its embedder and refuses a different one.
+        custom_instructions: (``"facts"``) extra extraction rules appended to
+            the prompt.
+        observation_date_from_timestamp: (``"facts"``) pass ``add(timestamp=)``
+            to the extractor as the conversation's Observation Date, so
+            "last week" resolves against that date. Off by default, which
+            matches Mem0 (it resolves relative dates against today).
+        keyword_scoring: (``"facts"``) fuse a BM25 keyword score into search so
+            a fact that literally contains a query word is not outranked by
+            generic facts with slightly closer vectors. ``True`` (weight 1, as
+            Mem0's design; the default) or a float weight; ``False`` is plain
+            semantic + entity scoring, identical to Mem0 on a Chroma store.
+    """
+
+    def __init__(
+        self,
+        data_dir: Optional[str] = None,
+        *,
+        entity_extractor: Any = "auto",
+        spacy_model: str = "en_core_web_sm",
+        default_user_id: str = DEFAULT_USER_ID,
+        auto_relink: bool = False,
+        backend: str = "raw",
+        llm: Any = None,
+        embedder: Any = None,
+        custom_instructions: Optional[str] = None,
+        observation_date_from_timestamp: bool = False,
+        keyword_scoring: Any = True,
+    ) -> None:
+        if backend not in BACKENDS:
+            raise CortexConfigError(f"backend must be one of {BACKENDS}, got {backend!r}")
+        self._backend = backend
+        self._e = _load_engine(backend)
+        self._data_dir = os.path.abspath(data_dir or _default_data_dir())
+        self._default_user = self._check_uid(default_user_id)
+        self._auto_relink = bool(auto_relink)
+        self._nlp = self._e.nlp.resolve_nlp(entity_extractor, spacy_model=spacy_model)
+        self._client = self._e.storage.get_client(self._data_dir)
+        self._collections: Dict[str, Any] = {}
+        self._lock = threading.Lock()
+        self._facts = None
+        if backend == "facts":
+            self._facts = self._e.fact_engine.FactEngine(
+                self._client, self._data_dir,
+                self._e.fact_backends.resolve_llm(llm),
+                self._e.fact_backends.resolve_embedder(embedder),
+                self._nlp, custom_instructions=custom_instructions,
+                observation_date_from_timestamp=observation_date_from_timestamp,
+                keyword_weight=float(keyword_scoring),
+            )
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "Memory":
+        """Build from a plain dict (keys: ``data_dir``, ``entity_extractor``,
+        ``spacy_model``, ``default_user_id``, ``auto_relink``, ``backend``,
+        ``llm``, ``embedder``, ``custom_instructions``).
+        """
+        if not isinstance(config, dict):
+            raise CortexConfigError("config must be a dict")
+        unknown = set(config) - _CONFIG_KEYS
+        if unknown:
+            raise CortexConfigError(f"unknown config keys: {', '.join(sorted(unknown))}")
+        backend = config.get("backend", "raw")
+        if backend not in BACKENDS:
+            raise CortexConfigError(
+                f"backend {backend!r} is not available; choose one of {BACKENDS}"
+            )
+        return cls(**config)
+
+    def __repr__(self) -> str:
+        return f"Memory(data_dir={self._data_dir!r}, backend={self._backend!r})"
+
+    # --- internals ---
+
+    def _check_uid(self, uid: str) -> str:
+        try:
+            return self._e.storage.validate_user_id(uid)
+        except ValueError as e:
+            raise InvalidRequestError(str(e)) from e
+
+    def _uid(self, user_id: Optional[str]) -> str:
+        """``None`` means the instance's default user."""
+        return self._check_uid(self._default_user if user_id is None else user_id)
+
+    def _col(self, user_id: Optional[str]):
+        uid = self._uid(user_id)
+        with self._lock:
+            col = self._collections.get(uid)
+            if col is None:
+                col = self._collections[uid] = (
+                    self._facts.collection(uid) if self._facts is not None
+                    else self._e.storage.get_collection(self._client, uid)
+                )
+            return col
+
+    def collection(self, user_id: Optional[str] = None) -> Any:
+        """The user's underlying Chroma collection (raw pages, or extracted
+        facts for ``backend="facts"``; both share one page schema).
+
+        For embedders that serve or browse the store directly, e.g. the Cortex
+        server. Treat it as read-only: writing through it bypasses embedding,
+        entity extraction and the facts pipeline."""
+        return self._col(user_id)
+
+    def _passages(self, user_id, col, query: str, limit: int, expand_links: bool) -> List[dict]:
+        """Seeds (+ optional link expansion) as ``{page_id, text, score, via, ...}``."""
+        if self._facts is not None:
+            seeds = self._facts.seeds(self._uid(user_id), query, limit)
+        else:
+            seeds = self._e.storage.query(col, query, n_results=limit)
+        if expand_links:
+            return self._e.retrieval.expand_links(col, seeds)
+        return [
+            {"page_id": p["id"], "text": p["text"], "score": p.get("score", 0.0), "via": "direct"}
+            for p in seeds
+        ]
+
+    def _all_pages(self, col) -> List[dict]:
+        return self._e.shaping.scan_all(col, self._e.storage.list_pages)
+
+    # --- memory ---
+
+    def add(
+        self,
+        text: str,
+        *,
+        user_id: Optional[str] = None,
+        timestamp: Optional[str] = None,
+    ) -> AddResult:
+        """Save ``text`` as memory. Long text is chunked into several small
+        pages (one fact each works best). ``timestamp`` (e.g. ``"8 May, 2023"``)
+        is stored as the date. Does not relink unless ``auto_relink`` is on."""
+        need_str(text, "text")
+        if timestamp is not None:
+            need_str(timestamp, "timestamp")
+        col = self._col(user_id)
+        if self._facts is not None:
+            ids = [f["id"] for f in self._facts.add(self._uid(user_id), text, timestamp)]
+        else:
+            ids = self._e.ingestion.add_text(col, text, self._nlp, timestamp=timestamp)
+        if self._auto_relink and ids:
+            self._e.linking.run_linking_pass(col)
+        return AddResult(page_ids=ids)
+
+    def search(
+        self,
+        query: str,
+        *,
+        user_id: Optional[str] = None,
+        limit: int = 4,
+        expand_links: bool = True,
+    ) -> List[SearchResult]:
+        """Semantic search. With ``expand_links`` (default) related pages are
+        pulled in via links — those results have ``via == "link"``. An empty
+        store returns ``[]``."""
+        need_str(query, "query")
+        need_int(limit, "limit", 1, MAX_SEARCH_LIMIT)
+        if not isinstance(expand_links, bool):
+            raise InvalidRequestError("expand_links must be True or False")
+        col = self._col(user_id)
+        if col.count() == 0:
+            return []
+        passages = self._passages(user_id, col, query, limit, expand_links)
+        return [self._e.shaping.search_result(p) for p in passages]
+
+    def answer(
+        self,
+        query: str,
+        *,
+        user_id: Optional[str] = None,
+        limit: int = 4,
+        model: Optional[str] = None,
+        chat: Optional[Callable[[str, str], str]] = None,
+    ) -> Answer:
+        """Retrieve, then have an LLM distil a short direct answer plus the ids
+        of the pages that support it.
+
+        Needs an LLM: by default a local Ollama at ``$OLLAMA_HOST``
+        (``http://localhost:11434``). Pass ``chat=fn(prompt, model) -> str`` to
+        use any other model. An unreachable LLM raises (this call has no
+        degraded mode — use :meth:`search` for raw passages).
+        """
+        need_str(query, "query")
+        need_int(limit, "limit", 1, MAX_SEARCH_LIMIT)
+        col = self._col(user_id)
+        if col.count() == 0:
+            return Answer(answer="", source_page_ids=[])
+        passages = self._passages(user_id, col, query, limit, True)
+        kwargs: Dict[str, Any] = {"_chat": chat}
+        if model:
+            kwargs["model"] = model
+        out = self._e.compression.compress(query, passages, **kwargs)
+        return Answer(answer=out["answer"], source_page_ids=list(out["source_page_ids"]))
+
+    def get(self, id: str, *, user_id: Optional[str] = None) -> Page:
+        """One page by id. Unknown ids — including another user's — raise
+        :class:`NotFoundError`."""
+        need_str(id, "id")
+        col = self._col(user_id)
+        stored = self._e.storage.get_page(col, id)
+        if stored is None:
+            raise NotFoundError(f"page not found: {id}", status=404)
+        out, incoming = self._e.shaping.link_context(self._all_pages(col))
+        return self._e.shaping.page(stored, out, incoming)
+
+    def get_all(
+        self,
+        *,
+        user_id: Optional[str] = None,
+        query: Optional[str] = None,
+        limit: int = 50,
+        offset: int = 0,
+    ) -> PageList:
+        """Browse pages; ``query`` is a case-insensitive substring filter (use
+        :meth:`search` for semantic search)."""
+        need_int(limit, "limit", 1, MAX_LIST_LIMIT)
+        need_int(offset, "offset", 0)
+        if query is not None:
+            need_str(query, "query")
+        col = self._col(user_id)
+        everything = self._all_pages(col)
+        needle = query.strip().lower() if query else None
+        matches = [
+            p for p in everything if needle is None or needle in p.get("text", "").lower()
+        ]
+        out, incoming = self._e.shaping.link_context(everything)
+        window = matches[offset:offset + limit]
+        return PageList(
+            pages=[self._e.shaping.page(p, out, incoming) for p in window],
+            total=len(matches), limit=limit, offset=offset,
+        )
+
+    def update(self, id: str, text: str, *, user_id: Optional[str] = None) -> None:
+        """Replace a page's text (re-embeds, re-extracts entities, keeps links).
+        Links can go stale — call :meth:`relink` after batch edits."""
+        need_str(id, "id")
+        need_str(text, "text")
+        col = self._col(user_id)
+        if self._e.storage.get_page(col, id) is None:
+            raise NotFoundError(f"page not found: {id}", status=404)
+        if self._facts is not None:
+            self._facts.update(self._uid(user_id), id, text)
+        else:
+            self._e.storage.update_page_text(col, id, text, self._nlp.entities(text))
+
+    def delete(self, id: str, *, user_id: Optional[str] = None) -> None:
+        """Delete a page. Links from other pages that pointed at it are
+        skipped by retrieval and cleaned up by the next :meth:`relink`."""
+        need_str(id, "id")
+        col = self._col(user_id)
+        if self._e.storage.get_page(col, id) is None:
+            raise NotFoundError(f"page not found: {id}", status=404)
+        self._e.storage.delete_page(col, id)
+
+    def delete_all(self, *, user_id: Optional[str] = None) -> int:
+        """Delete every page of one user. Returns how many were removed."""
+        uid = self._uid(user_id)
+        col = self._col(uid)
+        removed = col.count()
+        with self._lock:
+            self._collections.pop(uid, None)
+            if self._facts is not None:
+                self._facts.delete_all(uid)
+            else:
+                self._client.delete_collection(col.name)
+        return removed
+
+    def relink(self, *, user_id: Optional[str] = None) -> Dict[str, Any]:
+        """Re-run the batch linking pass (idempotent). Returns
+        ``{"pages": n, "links_written": m}``."""
+        return self._e.linking.run_linking_pass(self._col(user_id))
+
+    def count(self, *, user_id: Optional[str] = None) -> int:
+        """Number of pages stored for one user."""
+        return self._col(user_id).count()

cortexlayer/types.py

@@ -0,0 +1,199 @@
+"""Result types. Plain frozen dataclasses (no pydantic) built tolerantly from
+server JSON: unknown fields are ignored so a newer server never breaks an
+older client."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+
+def _s(d: Dict[str, Any], key: str, default: str = "") -> str:
+    v = d.get(key, default)
+    return v if isinstance(v, str) else default
+
+
+def _list(d: Dict[str, Any], key: str) -> List[str]:
+    v = d.get(key) or []
+    return [x for x in v if isinstance(x, str)] if isinstance(v, list) else []
+
+
+@dataclass(frozen=True)
+class SearchResult:
+    """One hit from :meth:`CortexClient.search`.
+
+    ``text`` is the full memory text when the source provides it (local
+    ``Memory`` always does); ``snippet`` is its first ~200 characters.
+
+    ``via`` is ``"direct"`` (vector hit) or ``"link"`` (pulled in by
+    link-expansion; ``linked_from`` is the seed page that led to it). ``score``
+    semantics depend on the backend (a distance, lower = closer, on raw; the
+    fused semantic + keyword + entity score, higher = closer, on facts) —
+    compare within one store or server, not across.
+    """
+
+    id: str
+    title: str
+    snippet: str
+    score: float
+    via: str = "direct"
+    linked_from: Optional[str] = None
+    source: str = "private"
+    text: str = ""   # the full memory text (local ``Memory`` fills it; a server may omit it)
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "SearchResult":
+        lf = d.get("linked_from")
+        return cls(
+            id=_s(d, "id"),
+            title=_s(d, "title"),
+            snippet=_s(d, "snippet"),
+            score=float(d.get("score") or 0.0),
+            via=_s(d, "via", "direct"),
+            linked_from=lf if isinstance(lf, str) else None,
+            source=_s(d, "source", "private"),
+            text=_s(d, "text"),
+        )
+
+
+@dataclass(frozen=True)
+class Page:
+    """A stored memory page. ``links`` point out, ``linked_from`` point in."""
+
+    id: str
+    title: str
+    content: str
+    snippet: str = ""
+    links: List[str] = field(default_factory=list)
+    linked_from: List[str] = field(default_factory=list)
+    degree: int = 0
+    created_at: str = ""
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "Page":
+        return cls(
+            id=_s(d, "id"),
+            title=_s(d, "title"),
+            content=_s(d, "content"),
+            snippet=_s(d, "snippet"),
+            links=_list(d, "links"),
+            linked_from=_list(d, "linked_from"),
+            degree=int(d.get("degree") or 0),
+            created_at=_s(d, "created_at"),
+        )
+
+
+@dataclass(frozen=True)
+class PageList:
+    """One window of :meth:`CortexClient.get_all` (``total`` counts all matches)."""
+
+    pages: List[Page]
+    total: int
+    limit: int
+    offset: int
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "PageList":
+        return cls(
+            pages=[Page.from_dict(p) for p in d.get("pages") or []],
+            total=int(d.get("total") or 0),
+            limit=int(d.get("limit") or 0),
+            offset=int(d.get("offset") or 0),
+        )
+
+
+@dataclass(frozen=True)
+class AddResult:
+    """Ids of the pages created by :meth:`CortexClient.add` (long text is
+    chunked into several small pages)."""
+
+    page_ids: List[str]
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "AddResult":
+        return cls(page_ids=_list(d, "page_ids"))
+
+
+@dataclass(frozen=True)
+class Account:
+    user_id: str
+    account_name: str
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "Account":
+        return cls(user_id=_s(d, "user_id"), account_name=_s(d, "account_name"))
+
+
+@dataclass(frozen=True)
+class UsageBucket:
+    key: str
+    calls: int
+    errors: int
+    avg_latency_ms: Optional[float] = None
+    label: Optional[str] = None
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "UsageBucket":
+        avg = d.get("avg_latency_ms")
+        lab = d.get("label")
+        return cls(
+            key=_s(d, "key"),
+            calls=int(d.get("calls") or 0),
+            errors=int(d.get("errors") or 0),
+            avg_latency_ms=float(avg) if isinstance(avg, (int, float)) else None,
+            label=lab if isinstance(lab, str) else None,
+        )
+
+
+@dataclass(frozen=True)
+class Usage:
+    """Your own API usage, aggregated (``group_by``: day | key | operation)."""
+
+    group_by: str
+    since: str
+    until: str
+    total_calls: int
+    total_errors: int
+    buckets: List[UsageBucket]
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "Usage":
+        return cls(
+            group_by=_s(d, "group_by"),
+            since=_s(d, "since"),
+            until=_s(d, "until"),
+            total_calls=int(d.get("total_calls") or 0),
+            total_errors=int(d.get("total_errors") or 0),
+            buckets=[UsageBucket.from_dict(b) for b in d.get("buckets") or []],
+        )
+
+
+@dataclass(frozen=True)
+class Graph:
+    """Raw graph payload (``nodes`` / ``edges`` as the server sends them)."""
+
+    nodes: List[Dict[str, Any]]
+    edges: List[Dict[str, Any]]
+    truncated: bool = False
+    stale: bool = False
+    start: Optional[str] = None
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "Graph":
+        start = d.get("start")
+        return cls(
+            nodes=list(d.get("nodes") or []),
+            edges=list(d.get("edges") or []),
+            truncated=bool(d.get("truncated", False)),
+            stale=bool(d.get("stale", False)),
+            start=start if isinstance(start, str) else None,
+        )
+
+
+@dataclass(frozen=True)
+class Answer:
+    """A short direct answer distilled from retrieved pages
+    (:meth:`Memory.answer`; needs an LLM — Ollama by default)."""
+
+    answer: str
+    source_page_ids: List[str]