cortexlayer 0.1.0__py3-none-any.whl

cortexlayer/_engine/facts/prompts.py

@@ -0,0 +1,100 @@
+"""Fact-extraction prompts.
+
+The system prompt (``additive_extraction_prompt.txt``) and the user-prompt
+builder below are taken from Mem0 (https://github.com/mem0ai/mem0, Apache
+License 2.0; ``mem0/configs/prompts.py``, mem0ai 2.1.0), copied unchanged so
+Cortex's fact-memory engine starts from identical extraction behavior. See the
+repository ``NOTICE`` file. Changes: none to the prompt text; the builder is
+re-typed with type hints and its dates default to UTC today.
+"""
+
+from __future__ import annotations
+
+import json
+from datetime import datetime, timezone
+from importlib import resources
+from typing import Any, Dict, List, Optional, Union
+
+PAST_MESSAGE_TRUNCATION_LIMIT = 300
+
+_prompt_cache: Optional[str] = None
+
+
+def extraction_system_prompt() -> str:
+    """The additive (ADD-only) extraction system prompt, from package data."""
+    global _prompt_cache
+    if _prompt_cache is None:
+        _prompt_cache = (
+            resources.files(__package__)
+            .joinpath("additive_extraction_prompt.txt")
+            .read_text(encoding="utf-8")
+        )
+    return _prompt_cache
+
+
+def _truncate(text: str, limit: int = PAST_MESSAGE_TRUNCATION_LIMIT) -> str:
+    return text if len(text) <= limit else text[:limit] + "..."
+
+
+def _format_summary(summary: Union[None, str, Dict[str, Any]]) -> str:
+    if isinstance(summary, dict):
+        return summary.get("summary", "")
+    return summary or ""
+
+
+def _format_conversation_history(messages: Optional[List[Dict[str, Any]]]) -> str:
+    if not messages:
+        return ""
+    out = ""
+    for msg in messages:
+        role = msg.get("role", "")
+        content = msg.get("message") or msg.get("content", "")
+        if role and content:
+            out += f"{role}: {_truncate(content)}\n"
+    return out
+
+
+def _serialize_memories(memories: Optional[List[Dict[str, Any]]]) -> str:
+    return json.dumps(memories or [], ensure_ascii=False)
+
+
+def _format_new_messages(new_messages: Union[str, List[Dict[str, Any]], None]) -> str:
+    if isinstance(new_messages, str):
+        return new_messages
+    return json.dumps(new_messages or [], ensure_ascii=False)
+
+
+def _resolve_dates(current_date: Optional[str] = None, observation_date: Optional[str] = None):
+    if current_date is None:
+        current_date = datetime.now(timezone.utc).date().isoformat()
+    if observation_date is None:
+        observation_date = current_date
+    return current_date, observation_date
+
+
+def build_extraction_prompt(
+    *,
+    summary: Union[None, str, Dict[str, Any]] = None,
+    recently_extracted_memories: Optional[List[Dict[str, Any]]] = None,
+    existing_memories: Optional[List[Dict[str, Any]]] = None,
+    new_messages: Union[str, List[Dict[str, Any]], None] = None,
+    last_k_messages: Optional[List[Dict[str, Any]]] = None,
+    current_date: Optional[str] = None,
+    timestamp: Optional[str] = None,
+    custom_instructions: Optional[str] = None,
+) -> str:
+    """Build the user-side prompt (pairs with :func:`extraction_system_prompt`)."""
+    current_date, observation_date = _resolve_dates(current_date, timestamp)
+    sections = [
+        f"## Summary\n{_format_summary(summary)}",
+        f"## Last k Messages\n{_format_conversation_history(last_k_messages)}",
+        f"## Recently Extracted Memories\n{_serialize_memories(recently_extracted_memories)}",
+        f"## Existing Memories\n{_serialize_memories(existing_memories)}",
+        f"## New Messages\n{_format_new_messages(new_messages)}",
+        f"## Observation Date\n{observation_date}",
+        f"## Current Date\n{current_date}",
+    ]
+    if custom_instructions:
+        sections.append(f"## Custom Instructions\n{custom_instructions}")
+    sections.append("# Output:")
+    return "\n\n".join(sections)

cortexlayer/_engine/facts/scoring.py

@@ -0,0 +1,148 @@
+"""Search scoring for fact memory.
+
+Adapted from Mem0 (Apache License 2.0; mem0ai 2.1.0, ``mem0/utils/scoring.py``
+and ``Memory._compute_entity_boosts``). See the repository NOTICE file.
+
+Changes: Mem0 only computes BM25 when the vector store implements keyword
+search, and its Chroma connector does not, so in a Chroma deployment the BM25
+branch never runs. The default here is identical (no keyword term). Task 0079
+adds it as an opt-in: a dependency-free BM25 (no lemmatizer) whose score is
+normalised to [0, 1] by the IDF mass of the query terms the store knows, then
+fused as Mem0 does: ``combined = (semantic + keyword + entity) / max_possible``.
+"""
+
+from __future__ import annotations
+
+import math
+import re
+from typing import Any, Dict, List, Optional, Sequence
+
+ENTITY_BOOST_WEIGHT = 0.5
+ENTITY_MATCH_MIN_SIMILARITY = 0.5   # entity-store hits below this are ignored
+MAX_QUERY_ENTITIES = 8
+
+
+def distance_to_score(distance: float) -> float:
+    """Chroma L2 distance -> similarity in (0, 1], higher is closer (as Mem0)."""
+    return 1.0 / (1.0 + distance)
+
+
+def entity_boost(similarity: float, num_linked: int) -> float:
+    """Boost for the memories linked to one matched entity.
+
+    Scaled down for entities linked to very many memories (they carry less
+    information): ``sim * 0.5 / (1 + 0.001 * (n - 1)^2)``.
+    """
+    n = max(num_linked, 1)
+    weight = 1.0 / (1.0 + 0.001 * ((n - 1) ** 2))
+    return similarity * ENTITY_BOOST_WEIGHT * weight
+
+
+BM25_K1 = 1.2
+BM25_B = 0.75
+KEYWORD_MAX_EXTRA_CANDIDATES = 50   # keyword hits the semantic over-fetch missed
+
+_STOPWORDS = frozenset(
+    "a an and are as at be been but by can could did do does for from had has have he her his how i "
+    "if in is it its me my of on or our she so than that the their them then there these they this "
+    "those to was we were what when where which who whom why will with would you your".split()
+)
+_TOKEN = re.compile(r"[a-z0-9]+(?:'[a-z]+)?")
+
+
+def tokenize(text: str) -> List[str]:
+    """Lower-cased content words: possessive/plural endings folded, stopwords dropped.
+
+    Deliberately tiny (no stemmer/lemmatizer dependency); applied identically to
+    facts and queries, so consistency is what matters.
+    """
+    out: List[str] = []
+    for tok in _TOKEN.findall(text.lower()):
+        if tok.endswith("'s"):
+            tok = tok[:-2]
+        tok = tok.replace("'", "")
+        if tok in _STOPWORDS:
+            continue
+        if len(tok) > 4 and tok.endswith("ies"):
+            tok = tok[:-3] + "y"
+        elif len(tok) > 3 and tok.endswith("s") and not tok.endswith("ss"):
+            tok = tok[:-1]
+        if tok and tok not in _STOPWORDS:
+            out.append(tok)
+    return out
+
+
+def keyword_scores(query: str, documents: Dict[str, str]) -> Dict[str, float]:
+    """BM25 of ``query`` over ``documents`` (id -> text), normalised to (0, 1].
+
+    Raw BM25 is divided by the summed IDF of the query terms that occur in at
+    least one document, so a fact containing every (known) query term scores
+    ~1 and one containing only a very common term scores near 0 — comparable
+    with the semantic and entity terms whatever the store size. Query terms no
+    fact contains do not count against the others. Only facts with a nonzero
+    score are returned.
+    """
+    q_terms = list(dict.fromkeys(tokenize(query)))
+    if not q_terms or not documents:
+        return {}
+    tokens = {i: tokenize(t) for i, t in documents.items()}
+    n = len(tokens)
+    avg_len = (sum(len(t) for t in tokens.values()) / n) or 1.0
+    counts = {i: {} for i in tokens}
+    df = {t: 0 for t in q_terms}
+    wanted = set(q_terms)
+    for i, toks in tokens.items():
+        c = counts[i]
+        for t in toks:
+            if t in wanted:
+                c[t] = c.get(t, 0) + 1
+        for t in c:
+            df[t] += 1
+    idf = {t: math.log(1.0 + (n - d + 0.5) / (d + 0.5)) for t, d in df.items() if d}
+    total_idf = sum(idf.values())
+    if total_idf <= 0:
+        return {}
+    out: Dict[str, float] = {}
+    for i, c in counts.items():
+        if not c:
+            continue
+        norm = BM25_K1 * (1.0 - BM25_B + BM25_B * len(tokens[i]) / avg_len)
+        raw = sum(idf[t] * c[t] * (BM25_K1 + 1.0) / (c[t] + norm) for t in c)
+        score = min(raw / total_idf, 1.0)   # tf=1 at average length gives idf, so full coverage ~ 1
+        if score > 0:
+            out[i] = score
+    return out
+
+
+def score_and_rank(
+    semantic_results: List[Dict[str, Any]],
+    entity_boosts: Dict[str, float],
+    threshold: float,
+    top_k: int,
+    keyword: Optional[Dict[str, float]] = None,
+    keyword_weight: float = 1.0,
+) -> List[Dict[str, Any]]:
+    """``combined = (semantic + w*keyword + entity_boost) / max_possible``, top-k descending.
+
+    The threshold gates the *semantic* score before combining. ``max_possible``
+    is 1.0, +``keyword_weight`` when any keyword score is active, +0.5 when any
+    entity boost is. With no ``keyword`` this is exactly Mem0-with-Chroma.
+    """
+    keyword = keyword or {}
+    max_possible = 1.0 + (ENTITY_BOOST_WEIGHT if entity_boosts else 0.0)
+    if keyword:
+        max_possible += keyword_weight
+    scored: List[Dict[str, Any]] = []
+    for result in semantic_results:
+        mem_id = result.get("id")
+        if mem_id is None:
+            continue
+        semantic = result.get("score") or 0.0
+        if semantic < threshold:
+            continue
+        boost = entity_boosts.get(str(mem_id), 0.0)
+        kw = keyword_weight * keyword.get(str(mem_id), 0.0)
+        combined = min((semantic + kw + boost) / max_possible, 1.0)
+        scored.append({**result, "id": str(mem_id), "score": combined})
+    scored.sort(key=lambda r: r["score"], reverse=True)
+    return scored[:top_k]

cortexlayer/_engine/ingestion.py

@@ -0,0 +1,84 @@
+"""Ingestion: chunk raw text into small pages, extract entities, embed, insert.
+
+Pages are deliberately small — roughly one fact/statement each — because
+link-expansion works best on focused pages. Chunking: split on line breaks,
+then greedily pack sentences so no page exceeds ``MAX_PAGE_CHARS``. Entity
+extraction is delegated to an :class:`~cortexlayer._engine.nlp.NLP` object.
+
+Ingestion never runs the linking pass — linking stays a batch operation.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import List, Optional
+
+from chromadb.api.models.Collection import Collection
+
+from . import storage
+from .nlp import NLP
+
+MAX_PAGE_CHARS = 500
+
+_DATE_PREFIX_RE = re.compile(r"^\[.+?\]")
+
+
+def apply_timestamp(text: str, timestamp: str) -> str:
+    """Prefix ``[timestamp]`` onto undated lines only.
+
+    Lines already starting with a ``[date]``-style bracket prefix keep their
+    original date instead of getting a second one stacked on top.
+    """
+
+    def prefix_line(line: str) -> str:
+        if not line.strip():
+            return line
+        if _DATE_PREFIX_RE.match(line.lstrip()):
+            return line
+        return f"[{timestamp}] {line}"
+
+    return "\n".join(prefix_line(line) for line in text.splitlines())
+
+
+def chunk_text(text: str, nlp: NLP, max_chars: int = MAX_PAGE_CHARS) -> List[str]:
+    """Split raw text into page-sized chunks."""
+    pages: List[str] = []
+    for block in (b.strip() for b in text.splitlines()):
+        if not block:
+            continue
+        if len(block) <= max_chars:
+            pages.append(block)
+            continue
+        # Long turn: split into sentences and greedily repack.
+        current: List[str] = []
+        current_len = 0
+        for s in nlp.sentences(block):
+            if current and current_len + 1 + len(s) > max_chars:
+                pages.append(" ".join(current))
+                current, current_len = [], 0
+            current.append(s)
+            current_len += (1 if current_len else 0) + len(s)
+        if current:
+            pages.append(" ".join(current))
+    return pages
+
+
+def add_text(
+    collection: Collection,
+    text: str,
+    nlp: NLP,
+    timestamp: Optional[str] = None,
+) -> List[str]:
+    """Chunk → entities → insert. Returns the new page ids.
+
+    ``timestamp`` (e.g. ``"8 May, 2023"``) is prefixed to undated lines and
+    stored as each page's ``created_at``.
+    """
+    if timestamp:
+        text = apply_timestamp(text, timestamp)
+    return [
+        storage.insert_page(
+            collection, chunk, nlp.entities(chunk), created_at=timestamp or None
+        )
+        for chunk in chunk_text(text, nlp)
+    ]

cortexlayer/_engine/linking.py

@@ -0,0 +1,68 @@
+"""Linking pass: connect pages that share entities.
+
+Deterministic entity-overlap linking, no LLM judgment. Runs as a batch over the
+full collection — after bulk ingestion or on-demand via `memory_relink` — never
+per-insert. Re-running recomputes links from scratch, so it is idempotent.
+"""
+
+from __future__ import annotations
+
+from chromadb.api.models.Collection import Collection
+
+from . import storage
+
+PAGE_BATCH = 500
+
+
+def _normalize(entity: str) -> str:
+    return entity.strip().lower()
+
+
+def run_linking_pass(
+    collection: Collection,
+    common_entity_fraction: float = 0.3,
+    common_entity_min_pages: int = 5,
+) -> dict:
+    """Populate every page's `links` with IDs of pages sharing ≥1 entity.
+
+    Matching is exact on normalized (case-insensitive) entity strings.
+    Entities appearing in more than max(common_entity_min_pages,
+    common_entity_fraction * N) pages are ignored for linking — in dialogue,
+    speaker names show up nearly everywhere and would otherwise link each page
+    to nearly every other page (found in benchmark diagnosis).
+    Ignored entities stay stored on the page; they just don't create links.
+    Returns summary stats: {"pages": n, "links_written": m}.
+    """
+    pages: list[dict] = []
+    offset = 0
+    while True:
+        batch = storage.list_pages(collection, limit=PAGE_BATCH, offset=offset)
+        if not batch:
+            break
+        pages.extend(batch)
+        offset += len(batch)
+
+    # Inverted index: normalized entity -> page IDs containing it.
+    index: dict[str, list[str]] = {}
+    for page in pages:
+        for entity in page["entities"]:
+            index.setdefault(_normalize(entity), []).append(page["id"])
+
+    threshold = max(common_entity_min_pages, common_entity_fraction * len(pages))
+    linkable = {entity for entity, ids in index.items() if len(ids) <= threshold}
+
+    links_written = 0
+    for page in pages:
+        neighbors: list[str] = []
+        for entity in page["entities"]:
+            if _normalize(entity) not in linkable:
+                continue
+            for other_id in index.get(_normalize(entity), []):
+                if other_id != page["id"] and other_id not in neighbors:
+                    neighbors.append(other_id)
+        neighbors.sort()
+        if neighbors != sorted(page["links"]):
+            storage.update_links(collection, page["id"], neighbors)
+            links_written += 1
+
+    return {"pages": len(pages), "links_written": links_written}

cortexlayer/_engine/nlp.py

@@ -0,0 +1,162 @@
+"""Language processing behind ingestion: entity extraction + sentence splitting.
+
+Two implementations behind one tiny interface (``entities`` / ``sentences``):
+
+- :class:`SpacyNLP` — spaCy NER plus a proper-noun pass (the quality path; this
+  is exactly what the Cortex server uses).
+- :class:`RegexNLP` — dependency-free capitalised-phrase heuristics. Good
+  enough to try the library without downloading a model, but it finds fewer
+  entities, so the linking pass links fewer pages. ``"auto"`` falls back to it
+  (with a one-time warning) only when spaCy or its model is missing.
+"""
+
+from __future__ import annotations
+
+import re
+import warnings
+from typing import Any, List, Protocol
+
+from ..errors import LocalDependencyError, MissingModelError
+
+SPACY_MODEL = "en_core_web_sm"
+
+
+class NLP(Protocol):
+    def entities(self, text: str) -> List[str]: ...
+    def sentences(self, text: str) -> List[str]: ...
+
+
+def _ordered_unique_add(seen: set, out: List[str], name: str) -> None:
+    name = name.strip()
+    if name and name not in seen:
+        seen.add(name)
+        out.append(name)
+
+
+class SpacyNLP:
+    """spaCy-backed extraction. Loading is lazy and raises clear errors."""
+
+    def __init__(self, model: str = SPACY_MODEL) -> None:
+        self.model = model
+        self._nlp: Any = None
+
+    def _load(self) -> Any:
+        if self._nlp is None:
+            try:
+                import spacy
+            except ImportError as e:  # pragma: no cover — exercised via monkeypatch
+                raise LocalDependencyError(
+                    "spaCy is not installed. Run: pip install \"cortexlayer[local]\" "
+                    "(or use Memory(entity_extractor=\"regex\"))."
+                ) from e
+            try:
+                self._nlp = spacy.load(self.model)
+            except OSError as e:
+                raise MissingModelError(
+                    f"spaCy model '{self.model}' is not installed. Run: "
+                    f"python -m spacy download {self.model}  "
+                    "(or use Memory(entity_extractor=\"regex\") to run without one)."
+                ) from e
+        return self._nlp
+
+    def entities(self, text: str) -> List[str]:
+        """Ordered-unique entities/key terms.
+
+        spaCy NER first, plus a proper-noun fallback: the small model reliably
+        tags people/places but often misses titles and works of art (e.g.
+        "Inception"), which POS tagging still catches as PROPN. Both feed the
+        linking pass, where a missed shared term means a missed link.
+        """
+        doc = self._load()(text)
+        seen: set = set()
+        entities: List[str] = []
+        for ent in doc.ents:
+            _ordered_unique_add(seen, entities, ent.text)
+        lowered = [e.lower() for e in entities]
+        for token in doc:
+            if token.pos_ != "PROPN":
+                continue
+            word = token.text.strip()
+            if not word:
+                continue
+            # Skip tokens already covered by a collected entity ("Nolan" in
+            # "Christopher Nolan"), keep genuinely new key terms ("Inception").
+            if any(word.lower() in e or e in word.lower() for e in lowered):
+                continue
+            _ordered_unique_add(seen, entities, word)
+            lowered.append(word.lower())
+        return entities
+
+    def sentences(self, text: str) -> List[str]:
+        return [s.text.strip() for s in self._load()(text).sents if s.text.strip()]
+
+
+_CAP_RUN = re.compile(r"[A-Z][\w'’-]*(?:[ \t]+[A-Z][\w'’-]*)*")
+_SENTENCE_END = re.compile(r"(?<=[.!?])[\"')\]]*\s+(?=[\"'(\[]*[A-Z0-9])")
+_STOP = frozenset(
+    "I A An The This That These Those There Then Here He She It We They You Me "
+    "My Our His Her Their Its Your Mine Ours And But Or Nor So Yet If In On At "
+    "For To Of With By From As Is Are Was Were Be Been Am Do Does Did Have Has "
+    "Had Will Would Can Could Should May Might Must What When Where Who Whom "
+    "Which Why How Yes No Not Also Just Well Oh Hi Hello Thanks Thank Please "
+    "Okay Ok Maybe Sure Some Any All Each Every Both Either Neither One".split()
+)
+
+
+class RegexNLP:
+    """Dependency-free heuristics: capitalised runs minus function words."""
+
+    def entities(self, text: str) -> List[str]:
+        seen: set = set()
+        out: List[str] = []
+        for match in _CAP_RUN.finditer(text):
+            words = match.group(0).split()
+            while words and words[0] in _STOP:   # "The Godfather" -> "Godfather"
+                words.pop(0)
+            if words:
+                _ordered_unique_add(seen, out, " ".join(words))
+        return out
+
+    def sentences(self, text: str) -> List[str]:
+        return [s.strip() for s in _SENTENCE_END.split(text) if s.strip()]
+
+
+_warned_fallback = False
+
+
+def resolve_nlp(spec: Any = "auto", *, spacy_model: str = SPACY_MODEL) -> NLP:
+    """``"auto"`` | ``"spacy"`` | ``"regex"`` | any object with ``entities`` and
+    ``sentences`` methods.
+
+    ``"auto"`` prefers spaCy and falls back to the regex heuristics with a
+    one-time warning if spaCy or its model is unavailable — so a fresh
+    ``pip install "cortexlayer[local]"`` works even before the model download.
+    """
+    global _warned_fallback
+    if hasattr(spec, "entities") and hasattr(spec, "sentences"):
+        return spec
+    if spec == "regex":
+        return RegexNLP()
+    if spec == "spacy":
+        nlp = SpacyNLP(spacy_model)
+        nlp._load()          # fail now, with the helpful message
+        return nlp
+    if spec == "auto":
+        nlp = SpacyNLP(spacy_model)
+        try:
+            nlp._load()
+            return nlp
+        except LocalDependencyError as e:
+            if not _warned_fallback:
+                _warned_fallback = True
+                warnings.warn(
+                    f"cortexlayer: {e} Falling back to the simpler regex entity "
+                    "extractor — fewer entities means fewer links between pages.",
+                    RuntimeWarning,
+                    stacklevel=3,
+                )
+            return RegexNLP()
+    raise ValueError(
+        "entity_extractor must be 'auto', 'spacy', 'regex', or an object with "
+        f"entities() and sentences() methods, got {spec!r}"
+    )

cortexlayer/_engine/retrieval.py

@@ -0,0 +1,55 @@
+"""Retrieval pipeline: vector search + link-expansion.
+
+Searches the whole collection (no vault/cluster routing — deliberately excluded).
+Links only *add* candidates, never restrict. MVP expansion heuristic: pull the top 1
+linked page per seed, no relevance judgment.
+"""
+
+from __future__ import annotations
+
+from chromadb.api.models.Collection import Collection
+
+from . import storage
+
+DEFAULT_K = 4
+
+
+def retrieve(collection: Collection, query: str, k: int = DEFAULT_K) -> list[dict]:
+    """Return combined seed + link-expanded passages.
+
+    Each passage is ``{page_id, text, score, via, [linked_from]}``: seeds are
+    ``via="direct"`` (score = Chroma distance, lower = closer), expanded pages
+    are ``via="link"`` with the seed id in ``linked_from``. Seeds come first,
+    then expanded pages, deduplicated by page ID (a page that is both seed and
+    expansion keeps ``direct``).
+    """
+    return expand_links(collection, storage.query(collection, query, n_results=k))
+
+
+def expand_links(collection: Collection, seeds: list[dict]) -> list[dict]:
+    """Seeds (storage-shaped page dicts with ``score``) -> passages, adding the
+    top linked page of each seed. Split from :func:`retrieve` so other seed
+    sources (fact memory's entity-boosted search) reuse the same expansion."""
+    seen: set[str] = set()
+    passages: list[dict] = []
+
+    def _add(page_id: str, text: str, score: float, via: str,
+             linked_from: str | None = None) -> None:
+        if page_id not in seen:
+            seen.add(page_id)
+            passage: dict = {"page_id": page_id, "text": text,
+                             "score": score, "via": via}
+            if linked_from is not None:
+                passage["linked_from"] = linked_from
+            passages.append(passage)
+
+    for seed in seeds:
+        _add(seed["id"], seed["text"], seed.get("score", 0.0), "direct")
+    for seed in seeds:
+        # MVP heuristic: top 1 linked page per seed (links are stored sorted).
+        if seed["links"]:
+            target = storage.get_page(collection, seed["links"][0])
+            if target is not None:
+                _add(target["id"], target["text"], target.get("score", 0.0),
+                     "link", linked_from=seed["id"])
+    return passages