rfnry-memory 0.1.0__py3-none-any.whl

rfnry_memory/__init__.py

@@ -0,0 +1,153 @@
+from rfnry_memory.errors import (
+    ConfigurationError,
+    InputError,
+    MemoryException,
+    ModelError,
+    NotFoundError,
+    StorageError,
+)
+from rfnry_memory.fusion import RRF, Fusion, RelativeScore
+from rfnry_memory.index import HNSW, IVFFlat, VectorIndex
+from rfnry_memory.memory import Memory
+from rfnry_memory.models import (
+    Answer,
+    AnswerLike,
+    Audio,
+    Check,
+    CheckAnswer,
+    Choice,
+    ChoiceAnswer,
+    Chunk,
+    Chunking,
+    Content,
+    CorpusDescription,
+    CorpusFingerprint,
+    Counts,
+    Decided,
+    DecidedLike,
+    DecisionModel,
+    Document,
+    DocumentEmbeddingModel,
+    EmbeddingModel,
+    Entity,
+    EntityIn,
+    GenerationModel,
+    Hit,
+    Hits,
+    Image,
+    Ingested,
+    Judgement,
+    JudgementName,
+    Kind,
+    Language,
+    Match,
+    Message,
+    Neighborhood,
+    Part,
+    PillarReport,
+    Question,
+    QuestionLike,
+    Ref,
+    Relation,
+    RelationIn,
+    Remembered,
+    Reply,
+    ReplyLike,
+    Score,
+    ScoreAnswer,
+    Section,
+    Source,
+    SourceIn,
+    State,
+    Story,
+    Synced,
+    Text,
+    ToolCall,
+    ToolSpec,
+    Usage,
+    confidence,
+)
+from rfnry_memory.postgres import PostgresStorage
+from rfnry_memory.ranking import BM25, CoverDensity, TextRanking
+from rfnry_memory.tools import MemoryTools, Tool, Tools
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Answer",
+    "AnswerLike",
+    "Check",
+    "CheckAnswer",
+    "Choice",
+    "ChoiceAnswer",
+    "Chunk",
+    "Chunking",
+    "ConfigurationError",
+    "CorpusDescription",
+    "CorpusFingerprint",
+    "Counts",
+    "Decided",
+    "DecidedLike",
+    "DecisionModel",
+    "DocumentEmbeddingModel",
+    "EmbeddingModel",
+    "Entity",
+    "EntityIn",
+    "GenerationModel",
+    "HNSW",
+    "BM25",
+    "CoverDensity",
+    "TextRanking",
+    "RRF",
+    "RelativeScore",
+    "Fusion",
+    "Hit",
+    "Hits",
+    "IVFFlat",
+    "Ingested",
+    "InputError",
+    "Judgement",
+    "JudgementName",
+    "Kind",
+    "Match",
+    "Language",
+    "Memory",
+    "MemoryException",
+    "MemoryTools",
+    "Message",
+    "ModelError",
+    "Neighborhood",
+    "NotFoundError",
+    "PillarReport",
+    "PostgresStorage",
+    "Question",
+    "QuestionLike",
+    "Ref",
+    "Relation",
+    "RelationIn",
+    "Remembered",
+    "Reply",
+    "ReplyLike",
+    "Score",
+    "ScoreAnswer",
+    "Section",
+    "Source",
+    "SourceIn",
+    "State",
+    "StorageError",
+    "Story",
+    "Synced",
+    "Tool",
+    "ToolCall",
+    "ToolSpec",
+    "Tools",
+    "Usage",
+    "VectorIndex",
+    "confidence",
+    "Text",
+    "Image",
+    "Audio",
+    "Document",
+    "Part",
+    "Content",
+]

rfnry_memory/chunking/chunker.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+from rfnry_memory.chunking.splitter import RecursiveTextSplitter
+from rfnry_memory.chunking.structure import (
+    build_heading_spans,
+    find_atomic_regions,
+    section_path_at,
+)
+
+
+class Piece(BaseModel):
+    title: str | None
+    text: str
+
+
+def chunk_text(text: str, *, chunk_chars: int, overlap_chars: int) -> list[Piece]:
+    text = text.strip()
+    if not text:
+        return []
+    regions = find_atomic_regions(text)
+    spans = build_heading_spans(text, regions)
+    splitter = RecursiveTextSplitter(chunk_size=chunk_chars, chunk_overlap=overlap_chars)
+    pieces: list[Piece] = []
+
+    def free(start: int, end: int) -> None:
+        for piece, offset in splitter.split_text_with_offsets(text[start:end]):
+            pieces.append(Piece(title=section_path_at(start + offset, spans), text=piece))
+
+    cursor = 0
+    for region in regions:
+        free(cursor, region.start)
+        pieces.append(
+            Piece(title=section_path_at(region.start, spans), text=region.content.strip())
+        )
+        cursor = region.end
+    free(cursor, len(text))
+    return [piece for piece in pieces if piece.text]

rfnry_memory/chunking/splitter.py

@@ -0,0 +1,93 @@
+import re
+
+SEPARATORS = ["\n\n", "\n", ". ", " ", ""]
+
+
+def _split(text: str, separator: str) -> list[tuple[int, int]]:
+    if not separator:
+        return [(i, i + 1) for i in range(len(text))]
+    spans: list[tuple[int, int]] = []
+    prev = 0
+    for m in re.finditer(re.escape(separator), text):
+        spans.append((prev, m.start()))
+        prev = m.end()
+    spans.append((prev, len(text)))
+    return [(s, e) for s, e in spans if e > s]
+
+
+class RecursiveTextSplitter:
+    def __init__(self, chunk_size: int, chunk_overlap: int) -> None:
+        self._chunk_size = chunk_size
+        self._chunk_overlap = chunk_overlap
+
+    def split_text_with_offsets(self, text: str) -> list[tuple[str, int]]:
+        return self._split(text, SEPARATORS, 0)
+
+    def _emit(self, text: str, start: int, end: int, base: int) -> list[tuple[str, int]]:
+        chunk = text[start:end]
+        s = start + len(chunk) - len(chunk.lstrip())
+        e = end - (len(chunk) - len(chunk.rstrip()))
+        if e <= s:
+            return []
+        return [(text[s:e], base + s)]
+
+    def _merge(self, spans: list[tuple[int, int]], separator: str) -> list[tuple[int, int]]:
+        joiner = len(separator)
+        chunks: list[tuple[int, int]] = []
+        current: list[tuple[int, int]] = []
+        total = 0
+        for span in spans:
+            size = span[1] - span[0]
+            if total + size + (joiner if current else 0) > self._chunk_size and current:
+                chunks.append((current[0][0], current[-1][1]))
+                while total > self._chunk_overlap or (
+                    total + size + (joiner if current else 0) > self._chunk_size and total > 0
+                ):
+                    total -= current[0][1] - current[0][0] + (joiner if len(current) > 1 else 0)
+                    current = current[1:]
+            current.append(span)
+            total += size + (joiner if len(current) > 1 else 0)
+        if current:
+            chunks.append((current[0][0], current[-1][1]))
+        return chunks
+
+    def _split(self, text: str, separators: list[str], base: int) -> list[tuple[str, int]]:
+        if not text:
+            return []
+        if len(text) <= self._chunk_size:
+            return self._emit(text, 0, len(text), base)
+        separator = separators[-1]
+        remaining: list[str] = []
+        for i, sep in enumerate(separators):
+            if not sep:
+                separator = sep
+                break
+            if sep in text:
+                separator = sep
+                remaining = separators[i + 1 :]
+                break
+        found: list[tuple[str, int]] = []
+        good: list[tuple[int, int]] = []
+        for span in _split(text, separator):
+            piece = text[span[0] : span[1]]
+            if len(piece) < self._chunk_size:
+                good.append(span)
+                continue
+            for s, e in self._merge(good, separator):
+                found.extend(self._emit(text, s, e, base))
+            good = []
+            if remaining:
+                found.extend(self._split(piece, remaining, base + span[0]))
+            elif len(piece) > self._chunk_size:
+                step = max(1, self._chunk_size - self._chunk_overlap)
+                for start in range(0, len(piece), step):
+                    found.extend(
+                        self._emit(
+                            piece, start, min(start + self._chunk_size, len(piece)), base + span[0]
+                        )
+                    )
+            else:
+                found.extend(self._emit(text, span[0], span[1], base))
+        for s, e in self._merge(good, separator):
+            found.extend(self._emit(text, s, e, base))
+        return found

rfnry_memory/chunking/structure.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+_FENCE_OPEN_RE = re.compile(r"```[a-zA-Z0-9_+-]*")
+_TABLE_RE = re.compile(
+    r"(?:^\|[^\n]+\|\n)"
+    r"(?:^\|[\s\-:|]+\|\n)"
+    r"(?:^\|[^\n]+\|\n?)+",
+    re.MULTILINE,
+)
+_HEADING_RE = re.compile(r"^(#{1,6})\s+(.+?)$", re.MULTILINE)
+
+
+@dataclass
+class AtomicRegion:
+    start: int
+    end: int
+    content: str
+
+
+@dataclass
+class HeadingSpan:
+    start: int
+    end: int
+    path: tuple[str, ...]
+
+
+def _iter_code_fences(text: str) -> list[AtomicRegion]:
+    lines = text.splitlines(keepends=True)
+    starts: list[int] = []
+    pos = 0
+    for line in lines:
+        starts.append(pos)
+        pos += len(line)
+
+    regions: list[AtomicRegion] = []
+    i = 0
+    n = len(lines)
+    while i < n:
+        if _FENCE_OPEN_RE.fullmatch(lines[i].rstrip("\r\n")):
+            j = i + 1
+            while j < n and not lines[j].startswith("```"):
+                j += 1
+            if j < n:
+                start = starts[i]
+                end = starts[j] + 3
+                regions.append(AtomicRegion(start, end, text[start:end]))
+                i = j + 1
+                continue
+        i += 1
+    return regions
+
+
+def find_atomic_regions(text: str) -> list[AtomicRegion]:
+    regions: list[AtomicRegion] = []
+    regions.extend(_iter_code_fences(text))
+    for m in _TABLE_RE.finditer(text):
+        regions.append(AtomicRegion(m.start(), m.end(), m.group(0)))
+    regions.sort(key=lambda r: (r.start, -r.end))
+    merged: list[AtomicRegion] = []
+    for region in regions:
+        if merged and region.end <= merged[-1].end:
+            continue
+        if merged and region.start < merged[-1].end:
+            prev = merged[-1]
+            merged[-1] = AtomicRegion(prev.start, region.end, text[prev.start : region.end])
+            continue
+        merged.append(region)
+    return merged
+
+
+def build_heading_spans(
+    text: str,
+    exclude_regions: list[AtomicRegion] | None = None,
+) -> list[HeadingSpan]:
+    headings: list[tuple[int, int, str]] = []
+    exclude = exclude_regions or []
+    for m in _HEADING_RE.finditer(text):
+        offset = m.start()
+        if any(r.start <= offset < r.end for r in exclude):
+            continue
+        title = m.group(2).strip()
+        title = title.rstrip("#").rstrip() or title
+        headings.append((offset, len(m.group(1)), title))
+    if not headings:
+        return []
+    spans: list[HeadingSpan] = []
+    stack: list[tuple[int, str]] = []
+    for i, (offset, level, title) in enumerate(headings):
+        while stack and stack[-1][0] >= level:
+            stack.pop()
+        stack.append((level, title))
+        end = headings[i + 1][0] if i + 1 < len(headings) else len(text)
+        spans.append(
+            HeadingSpan(
+                start=offset,
+                end=end,
+                path=tuple(t for _, t in stack),
+            )
+        )
+    return spans
+
+
+def section_path_at(offset: int, spans: list[HeadingSpan]) -> str | None:
+    for s in spans:
+        if s.start <= offset < s.end:
+            return " > ".join(s.path)
+    return None

rfnry_memory/cli.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+import argparse
+import asyncio
+import sys
+from functools import partial
+from importlib.util import find_spec
+from typing import Any
+from urllib.parse import urlsplit
+
+from rfnry_memory.errors import MemoryException
+
+_WARNING = "take a snapshot of the database and stop app instances running the old version first"
+
+
+def _parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="rfnry-memory", description="prepare and upgrade the database rfnry-memory runs on"
+    )
+    commands = parser.add_subparsers(dest="command", required=True)
+
+    def command(name: str, summary: str) -> argparse.ArgumentParser:
+        sub = commands.add_parser(name, help=summary, description=summary)
+        sub.add_argument("--url", help="postgres connection string; without it, the PG* variables")
+        sub.add_argument("--schema", default="memory", help="the postgres schema (default: memory)")
+        return sub
+
+    create = command("create", "create the schema with this version's tables")
+    vectors = create.add_mutually_exclusive_group(required=True)
+    vectors.add_argument("--dimension", type=int, help="the embedding dimension")
+    vectors.add_argument("--no-vector", action="store_true", help="create no vector columns")
+    create.add_argument("--no-graph", action="store_true", help="do not install pg_trgm")
+    create.add_argument("--language", help="the default text search configuration of new corpora")
+    create.add_argument("--sql", action="store_true", help="print the statements and run nothing")
+    command("check", "say whether the schema fits this version")
+    upgrade = command("upgrade", "carry the schema to this version")
+    upgrade.add_argument("--sql", action="store_true", help="print the statements and run nothing")
+    upgrade.add_argument("--yes", action="store_true", help="run without asking")
+    change = command("set", "change a setting of the schema")
+    setting = change.add_mutually_exclusive_group(required=True)
+    setting.add_argument("--dimension", type=int, help="the embedding dimension")
+    setting.add_argument("--graph", action="store_true", help="install pg_trgm")
+    setting.add_argument("--language", help="the default text search configuration of new corpora")
+    change.add_argument("--sql", action="store_true", help="print the statements and run nothing")
+    change.add_argument("--yes", action="store_true", help="run without asking")
+    return parser
+
+
+def _scrubbed(text: str, url: str | None) -> str:
+    if not url:
+        return text
+    password = urlsplit(url).password
+    text = text.replace(url, "<url>")
+    return text.replace(f":{password}@", ":***@") if password else text
+
+
+def _quoted(name: str) -> str:
+    return '"' + name.replace('"', '""') + '"'
+
+
+async def _connect(args: argparse.Namespace) -> tuple[Any, str]:
+    import asyncpg
+
+    conn = await asyncpg.connect(args.url)
+    held = await conn.fetch(
+        "SELECT DISTINCT quote_ident(n.nspname) AS name FROM pg_extension e "
+        "JOIN pg_namespace n ON n.oid = e.extnamespace WHERE e.extname IN ('vector', 'pg_trgm') "
+        "AND n.nspname NOT IN ($1, 'public') ORDER BY 1",
+        args.schema,
+    )
+    path = [_quoted(args.schema), "public", *(r["name"] for r in held)]
+    setting = f"SET search_path TO {', '.join(path)}"
+    await conn.execute(setting)
+    return conn, setting
+
+
+def _confirmed(args: argparse.Namespace) -> bool:
+    if args.yes:
+        return True
+    print(_WARNING)
+    try:
+        answer = input("run these steps? [y/N] ")
+    except (EOFError, KeyboardInterrupt):
+        print()
+        return False
+    return answer.strip().lower() in ("y", "yes")
+
+
+async def _run(args: argparse.Namespace) -> int:
+    from rfnry_memory.postgres import migrations
+    from rfnry_memory.postgres.schema import _problems
+
+    conn, setting = await _connect(args)
+    try:
+        exists = await conn.fetchval(
+            "SELECT EXISTS (SELECT 1 FROM pg_namespace WHERE nspname = $1)", args.schema
+        )
+        if not exists and args.command != "create":
+            print(
+                f"cli: there is no schema {args.schema} — run rfnry-memory create --schema {args.schema}",
+                file=sys.stderr,
+            )
+            return 1
+        if args.command == "check":
+            problems, held = await _problems(conn, dimension=None, graph=False)
+            if problems:
+                for problem in problems:
+                    print(f"cli: {problem}", file=sys.stderr)
+                return 1
+            vectors = "no vectors" if held is None else f"vectors of {held} dimensions"
+            print(f"schema {args.schema} fits rfnry-memory: version {migrations.LATEST}, {vectors}")
+            return 0
+        before: list[str] = []
+        planner: migrations.Planner
+        if args.command == "create":
+            before = [f"CREATE SCHEMA IF NOT EXISTS {_quoted(args.schema)}"]
+            planner = partial(
+                migrations.plan_create,
+                dimension=None if args.no_vector else args.dimension,
+                graph=not args.no_graph,
+                language=args.language,
+            )
+        elif args.command == "upgrade":
+            planner = migrations.plan_upgrade
+        else:
+            planner = partial(
+                migrations.plan_set,
+                dimension=args.dimension,
+                graph=args.graph,
+                language=args.language,
+            )
+        shown = await migrations.apply(conn, planner, run=False, before=before)
+        if args.sql:
+            print(f"{setting};")
+            for statement in shown.statements:
+                print(f"{statement};")
+            return 0
+        if not shown.steps:
+            print(f"schema {args.schema}: nothing to do")
+            return 0
+        for step in shown.steps:
+            print(f"  {step}")
+        if args.command != "create" and not _confirmed(args):
+            print("nothing was run")
+            return 0
+        await migrations.apply(conn, planner, expected=shown, before=before)
+        print(f"schema {args.schema} is at version {migrations.LATEST}")
+        return 0
+    finally:
+        await conn.close()
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = _parser().parse_args(argv)
+    if args.schema in ("public", "information_schema") or args.schema.startswith("pg_"):
+        print(
+            f"cli: memory is not created in {args.schema!r} — name a schema of its own",
+            file=sys.stderr,
+        )
+        return 1
+    if args.command in ("create", "set") and args.dimension is not None and args.dimension < 1:
+        print("cli: a dimension is a positive int", file=sys.stderr)
+        return 1
+    if find_spec("asyncpg") is None:
+        print("cli: install rfnry-memory[postgres] to use this command", file=sys.stderr)
+        return 1
+    try:
+        return asyncio.run(_run(args))
+    except MemoryException as exc:
+        print(_scrubbed(str(exc), args.url), file=sys.stderr)
+        return 1
+    except Exception as exc:
+        print(_scrubbed(f"cli: {type(exc).__name__}: {exc}", args.url), file=sys.stderr)
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

rfnry_memory/errors.py

@@ -0,0 +1,22 @@
+class MemoryException(Exception):
+    pass
+
+
+class ConfigurationError(MemoryException):
+    pass
+
+
+class InputError(MemoryException):
+    pass
+
+
+class ModelError(MemoryException):
+    pass
+
+
+class StorageError(MemoryException):
+    pass
+
+
+class NotFoundError(MemoryException):
+    pass

rfnry_memory/filters.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import Any
+
+from rfnry_memory.errors import InputError
+
+RESERVED = frozenset({"corpus_id"})
+FACET_PREFIX = "metadata."
+
+
+def _valid_list(values: list[Any]) -> bool:
+    if all(isinstance(v, str) for v in values):
+        return True
+    return all(isinstance(v, int) and not isinstance(v, bool) for v in values)
+
+
+def validate_filters(filters: dict[str, Any] | None) -> None:
+    if not filters:
+        return
+    for key, condition in filters.items():
+        if not isinstance(key, str) or not key.strip():
+            raise InputError("filters: filter keys must be non-empty strings")
+        if key in RESERVED:
+            raise InputError(
+                f"filters: filter key {key!r} is reserved — pass it as a call argument"
+            )
+        if not key.startswith(FACET_PREFIX) or len(key) == len(FACET_PREFIX):
+            raise InputError(
+                f"filters: filter key {key!r} must address a facet as '{FACET_PREFIX}<key>'"
+            )
+        if isinstance(condition, list):
+            if not condition or not _valid_list(condition):
+                raise InputError(
+                    f"filters: filter {key!r} must be a non-empty list of str or of int"
+                )
+            continue
+        if isinstance(condition, str | int):
+            continue
+        raise InputError(
+            f"filters: filter {key!r} must be str, int, bool or a list of str or of int"
+        )
+
+
+def _lookup(metadata: dict[str, Any], path: str) -> tuple[Any, bool]:
+    current: Any = metadata
+    for part in path.split("."):
+        if not isinstance(current, dict) or part not in current:
+            return None, False
+        current = current[part]
+    return current, True
+
+
+def _same(a: Any, b: Any) -> bool:
+    if isinstance(a, bool) != isinstance(b, bool):
+        return False
+    return bool(a == b)
+
+
+def _condition_matches(value: Any, condition: Any) -> bool:
+    values = value if isinstance(value, list) else [value]
+    conditions = condition if isinstance(condition, list) else [condition]
+    return any(_same(v, c) for v in values for c in conditions)
+
+
+def within(value: Any, bound: Any) -> bool:
+    values = value if isinstance(value, list) else [value]
+    bounds = bound if isinstance(bound, list) else [bound]
+    return bool(values) and all(any(_same(v, b) for b in bounds) for v in values)
+
+
+def matches_filters(metadata: dict[str, Any], filters: dict[str, Any] | None) -> bool:
+    if not filters:
+        return True
+    for key, condition in filters.items():
+        value, found = _lookup(metadata, key[len(FACET_PREFIX) :])
+        if not found or not _condition_matches(value, condition):
+            return False
+    return True
+
+
+def stamp(metadata: dict[str, Any] | None, filters: dict[str, Any]) -> dict[str, Any]:
+    stamped = deepcopy(metadata or {})
+    for key, value in filters.items():
+        *parents, leaf = key[len(FACET_PREFIX) :].split(".")
+        node: Any = stamped
+        for part in parents:
+            node = node.setdefault(part, {})
+            if not isinstance(node, dict):
+                raise InputError(
+                    f"filters: metadata must agree with the bound filter {key!r} = {value!r}"
+                )
+        if leaf in node:
+            if not within(node[leaf], value):
+                raise InputError(
+                    f"filters: metadata must agree with the bound filter {key!r} = {value!r}"
+                )
+            continue
+        if isinstance(value, list) and len(value) > 1:
+            raise InputError(
+                f"filters: metadata must name one of the values bound to {key!r}: {value!r}"
+            )
+        node[leaf] = value[0] if isinstance(value, list) else value
+    return stamped