From a50c9648368fcf71ccf5bde07c1123d3493743c6 Mon Sep 17 00:00:00 2001 From: JQuery <3489161689@qq.com> Date: Mon, 31 Aug 2026 12:03:07 +0800 Subject: [PATCH] feat(memoria): improve RAG retrieval and evaluation --- README.md | 14 + ...60831_01_add_memory_revision_embeddings.py | 38 ++ .../memoria/server/cli/embedding_backfill.py | 174 ++++++ app/mneme/memoria/server/eval/README.md | 43 ++ app/mneme/memoria/server/eval/beir.py | 518 ++++++++++++++++++ .../memoria/server/memory/reconciliation.py | 30 + .../memoria/server/models/memory_revision.py | 10 + .../memoria/server/retrieval/documents.py | 11 +- app/mneme/memoria/server/retrieval/fusion.py | 16 +- app/mneme/memoria/server/retrieval/keyword.py | 7 +- .../memoria/server/retrieval/memories.py | 58 +- app/mneme/memoria/server/retrieval/profile.py | 3 + app/mneme/memoria/server/retrieval/vector.py | 10 +- app/mneme/memoria/server/runtime/retriever.py | 81 ++- .../memoria/server/services/embeddings.py | 12 + .../memoria/server/services/projections.py | 13 +- 16 files changed, 998 insertions(+), 40 deletions(-) create mode 100644 app/mneme/memoria/server/alembic/versions/20260831_01_add_memory_revision_embeddings.py create mode 100644 app/mneme/memoria/server/cli/embedding_backfill.py create mode 100644 app/mneme/memoria/server/eval/beir.py diff --git a/README.md b/README.md index 370feaf..8f67dfe 100644 --- a/README.md +++ b/README.md @@ -327,6 +327,20 @@ python -m app.mneme.memoria.server.cli.backfill \ --dry-run ``` +迁移后补齐历史 memory revision 的语义向量： + +```bash +python -m app.mneme.memoria.server.cli.embedding_backfill --memory --dry-run +python -m app.mneme.memoria.server.cli.embedding_backfill --memory --batch-size 100 +``` + +上下文化向量格式升级后重建当前文档 chunk 向量： + +```bash +python -m app.mneme.memoria.server.cli.embedding_backfill --documents --dry-run +python -m app.mneme.memoria.server.cli.embedding_backfill --documents --batch-size 100 +``` + 完整删除、重建和恢复流程见[部署指南中的 Memoria operations](deploy/DEPLOY.md#memoria-rebuild-and-deletion-operations)。 ## 项目结构 diff --git a/app/mneme/memoria/server/alembic/versions/20260831_01_add_memory_revision_embeddings.py b/app/mneme/memoria/server/alembic/versions/20260831_01_add_memory_revision_embeddings.py new file mode 100644 index 0000000..6c62443 --- /dev/null +++ b/app/mneme/memoria/server/alembic/versions/20260831_01_add_memory_revision_embeddings.py @@ -0,0 +1,38 @@ +"""add semantic embeddings to memory revisions + +Revision ID: 20260831_01 +Revises: 20260830_01 +Create Date: 2026-08-31 +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from pgvector.sqlalchemy import Vector + +from alembic import op +from app.mneme.memoria.server.config import settings + +revision: str = "20260831_01" +down_revision: str | Sequence[str] | None = "20260830_01" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.add_column( + "memory_revisions", + sa.Column("embedding", Vector(settings.EMBEDDING_DIMENSION), nullable=True), + ) + op.create_index( + "ix_memory_revisions_embedding_hnsw", + "memory_revisions", + ["embedding"], + postgresql_using="hnsw", + postgresql_ops={"embedding": "vector_cosine_ops"}, + ) + + +def downgrade() -> None: + op.drop_index("ix_memory_revisions_embedding_hnsw", table_name="memory_revisions") + op.drop_column("memory_revisions", "embedding") diff --git a/app/mneme/memoria/server/cli/embedding_backfill.py b/app/mneme/memoria/server/cli/embedding_backfill.py new file mode 100644 index 0000000..ee754d2 --- /dev/null +++ b/app/mneme/memoria/server/cli/embedding_backfill.py @@ -0,0 +1,174 @@ +"""Backfill Memoria memory and active-document embeddings.""" + +import argparse +import asyncio +import json + +from sqlalchemy import func, select, update + +from app.mneme.memoria.server.database import engine, open_read_session, open_write_session +from app.mneme.memoria.server.models.document_chunk import DocumentChunk +from app.mneme.memoria.server.models.document_projection import DocumentProjection +from app.mneme.memoria.server.models.memory_revision import MemoryRevision +from app.mneme.memoria.server.services.embeddings import ( + document_embedding_text, + embed_texts, + memory_embedding_text, +) + + +def _memory_filters(args: argparse.Namespace) -> list: + filters = [MemoryRevision.embedding.is_(None)] + if args.owner_id is not None: + filters.append(MemoryRevision.owner_id == args.owner_id) + if args.knowledge_base_id is not None: + filters.append(MemoryRevision.knowledge_base_id == args.knowledge_base_id) + return filters + + +async def memory_embedding_backfill(args: argparse.Namespace) -> dict: + filters = _memory_filters(args) + if args.resume_from: + filters.append(MemoryRevision.revision_id > args.resume_from) + if args.dry_run: + async with open_read_session() as db: + pending = await db.scalar(select(func.count()).select_from(MemoryRevision).where(*filters)) + return {"pending": int(pending or 0), "updated": 0, "dry_run": True} + + updated = 0 + resume_from = args.resume_from + while True: + statement = select( + MemoryRevision.revision_id, + MemoryRevision.subject, + MemoryRevision.predicate, + MemoryRevision.value, + ).where(*filters) + if resume_from: + statement = statement.where(MemoryRevision.revision_id > resume_from) + statement = statement.order_by(MemoryRevision.revision_id).limit(args.batch_size) + async with open_read_session() as db: + rows = (await db.execute(statement)).all() + if not rows: + break + + vectors = await embed_texts( + [memory_embedding_text(subject=row.subject, predicate=row.predicate, value=row.value) for row in rows] + ) + async with open_write_session() as db: + for row, vector in zip(rows, vectors, strict=True): + result = await db.execute( + update(MemoryRevision) + .where( + MemoryRevision.revision_id == row.revision_id, + MemoryRevision.embedding.is_(None), + ) + .values(embedding=vector) + ) + updated += result.rowcount or 0 + resume_from = rows[-1].revision_id + + return {"pending": 0, "updated": updated, "resume_from": resume_from, "dry_run": False} + + +def _document_filters(args: argparse.Namespace) -> list: + filters = [DocumentChunk.is_active.is_(True), DocumentProjection.status == "active"] + if args.owner_id is not None: + filters.append(DocumentChunk.owner_id == args.owner_id) + if args.knowledge_base_id is not None: + filters.append(DocumentChunk.knowledge_base_id == args.knowledge_base_id) + return filters + + +async def document_embedding_backfill(args: argparse.Namespace) -> dict: + cursor = int(args.resume_from or 0) + filters = _document_filters(args) + if args.dry_run: + async with open_read_session() as db: + eligible = await db.scalar( + select(func.count()) + .select_from(DocumentChunk) + .join(DocumentProjection) + .where(*filters, DocumentChunk.id > cursor) + ) + return {"eligible": int(eligible or 0), "updated": 0, "dry_run": True} + + updated = 0 + # ponytail: embedding version is not persisted; reruns rewrite active vectors. + # Add a version column only if representation upgrades become frequent. + while True: + statement = ( + select( + DocumentChunk.id, + DocumentChunk.content, + DocumentChunk.section_path, + DocumentProjection.file_name, + ) + .join(DocumentProjection) + .where(*filters, DocumentChunk.id > cursor) + .order_by(DocumentChunk.id) + .limit(args.batch_size) + ) + async with open_read_session() as db: + rows = (await db.execute(statement)).all() + if not rows: + break + + vectors = await embed_texts( + [ + document_embedding_text( + file_name=row.file_name, + section_path=row.section_path, + content=row.content, + ) + for row in rows + ] + ) + async with open_write_session() as db: + for row, vector in zip(rows, vectors, strict=True): + result = await db.execute( + update(DocumentChunk) + .where(DocumentChunk.id == row.id, DocumentChunk.is_active.is_(True)) + .values(embedding=vector) + ) + updated += result.rowcount or 0 + cursor = rows[-1].id + + return {"updated": updated, "resume_from": cursor, "dry_run": False} + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + mode = parser.add_mutually_exclusive_group(required=True) + mode.add_argument("--memory", action="store_true") + mode.add_argument("--documents", action="store_true") + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--owner-id", type=int) + parser.add_argument("--knowledge-base-id") + parser.add_argument("--resume-from") + parser.add_argument("--batch-size", type=int, default=100) + return parser + + +async def _main() -> None: + args = _parser().parse_args() + if args.batch_size <= 0: + raise SystemExit("--batch-size must be positive") + if args.owner_id is not None and args.owner_id <= 0: + raise SystemExit("--owner-id must be positive") + if args.documents and args.resume_from: + try: + resume_from = int(args.resume_from) + except ValueError as exc: + raise SystemExit("--resume-from must be a numeric chunk row ID") from exc + if resume_from < 0: + raise SystemExit("--resume-from must be non-negative") + try: + report = await document_embedding_backfill(args) if args.documents else await memory_embedding_backfill(args) + print(json.dumps(report, sort_keys=True)) + finally: + await engine.dispose() + + +if __name__ == "__main__": + asyncio.run(_main()) diff --git a/app/mneme/memoria/server/eval/README.md b/app/mneme/memoria/server/eval/README.md index ef344c8..ef0916c 100644 --- a/app/mneme/memoria/server/eval/README.md +++ b/app/mneme/memoria/server/eval/README.md @@ -105,3 +105,46 @@ that defaults to off. Operators can force the single-agent path immediately by setting `memory_agent.multi_agent.enabled` to `false` in `memoria.json`. Operators may stage the rollout with `rollout_percent` and `allowed_modes` in the same section. + +## Public retrieval benchmark + +`beir.py` runs a real offline retrieval benchmark on BEIR SciFact. It downloads +the official archive linked by the Hugging Face dataset card, verifies the +published MD5, deterministically samples test queries, and compares the +configured dense model with a TF-IDF lexical baseline, RRF hybrid retrieval, +and an optional configured CrossEncoder. Generated data and reports stay under +the git-ignored `storage/eval/` directory. + +The application image already contains the model dependencies, so no host +Python environment is required: + +```zsh +DOCKER=/Applications/Docker.app/Contents/Resources/bin/docker +"$DOCKER" run --rm \ + --user 0:0 \ + -e HOME=/tmp \ + -e HF_ENDPOINT=https://huggingface.co \ + -e RERANKER_ENABLED=true \ + -e BEIR_EVAL_GIT_COMMIT="$(git rev-parse HEAD)" \ + -e BEIR_EVAL_GIT_DIRTY="$(test -z "$(git status --porcelain)" && echo false || echo true)" \ + -v "$PWD:/app" \ + -w /app \ + reminder-app:local \ + python -m app.mneme.memoria.server.eval.beir \ + --sample-size 100 \ + --corpus-size 1000 \ + --top-k 10 \ + --candidate-k 50 \ + --rerank-k 20 \ + --reranker \ + --output storage/eval/beir/scifact-report.json +``` + +The report records the repository commit and dirty state, dataset checksum and +revision, sampled query IDs, resolved model source, configuration, timings, +per-pipeline Recall/MRR/NDCG, and failure samples. Embedding caches are bound to +the resolved model source and exact sampled document text. TF-IDF is only an offline lexical proxy; +use live evaluation against PostgreSQL when exact backend parity is required. +The default 1000-document corpus keeps every qrels-positive document for the +sampled queries and fills the remainder with fixed-seed distractors. This is a +repeatable development benchmark, not the official full-corpus BEIR score. diff --git a/app/mneme/memoria/server/eval/beir.py b/app/mneme/memoria/server/eval/beir.py new file mode 100644 index 0000000..4b5aed7 --- /dev/null +++ b/app/mneme/memoria/server/eval/beir.py @@ -0,0 +1,518 @@ +"""Run a reproducible BEIR SciFact retrieval benchmark. + +The benchmark compares the configured dense embedding model with a lightweight +lexical baseline, reciprocal-rank fusion, and an optional configured reranker. +It downloads the official BEIR archive linked by the Hugging Face dataset card. +""" + +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import os +import random +import subprocess +import sys +import time +import zipfile +from collections import defaultdict +from datetime import UTC, datetime +from pathlib import Path +from statistics import fmean +from typing import Any +from urllib.request import urlopen + +import numpy as np +from sklearn.feature_extraction.text import TfidfVectorizer + +from app.mneme.clients.embedding_client import get_embeddings, resolve_embedding_source +from app.mneme.clients.reranker_client import get_reranker +from app.mneme.conf.config import settings +from app.mneme.domains.eval.service import ( + calculate_mrr, + calculate_ndcg, + calculate_recall_at_k, +) +from app.mneme.memoria.server.retrieval.fusion import ( + FUSION_CANDIDATE_K, + LEXICAL_RRF_WEIGHT, + RRF_CONSTANT, +) + +DATASET_NAME = "BEIR SciFact" +DATASET_URL = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/scifact.zip" +DATASET_MD5 = "5f7d1de60b170fc8027bb7898e2efca1" +HF_DATASET_CARD = "https://huggingface.co/datasets/BeIR/scifact" +HF_DATASET_REVISION = "b3b5335604bf5ee3c4447671af975ea25143d4f5" + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", required=True, type=Path, help="JSON report path") + parser.add_argument( + "--data-dir", + type=Path, + default=Path("storage/eval/beir"), + help="download and extraction cache", + ) + parser.add_argument("--sample-size", type=int, default=100, help="test queries to evaluate") + parser.add_argument( + "--corpus-size", + type=int, + default=1000, + help="sampled corpus size; all selected-query positives are always retained", + ) + parser.add_argument("--seed", type=int, default=20260830, help="deterministic query sample seed") + parser.add_argument("--top-k", type=int, default=10, help="metric cutoff") + parser.add_argument("--candidate-k", type=int, default=50, help="dense/lexical candidates for RRF") + parser.add_argument("--rerank-k", type=int, default=20, help="hybrid candidates sent to reranker") + parser.add_argument("--embedding-batch-size", type=int, default=32) + parser.add_argument("--reranker-batch-size", type=int, default=16) + parser.add_argument("--reranker", action="store_true", help="run the configured CrossEncoder") + return parser + + +def _validate_args(args: argparse.Namespace) -> None: + for name in ( + "sample_size", + "corpus_size", + "top_k", + "candidate_k", + "rerank_k", + "embedding_batch_size", + "reranker_batch_size", + ): + if getattr(args, name) < 1: + raise ValueError(f"--{name.replace('_', '-')} must be positive") + if args.candidate_k < args.top_k: + raise ValueError("--candidate-k must be at least --top-k") + if args.rerank_k > args.candidate_k: + raise ValueError("--rerank-k cannot exceed --candidate-k") + + +def _digest(path: Path, algorithm: str = "sha256") -> str: + digest = hashlib.new(algorithm) + with path.open("rb") as source: + while chunk := source.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def _download_dataset(data_dir: Path) -> Path: + data_dir.mkdir(parents=True, exist_ok=True) + archive = data_dir / "scifact.zip" + dataset_dir = data_dir / "scifact" + if not archive.exists() or _digest(archive, "md5") != DATASET_MD5: + temporary = archive.with_suffix(".zip.part") + with urlopen(DATASET_URL, timeout=120) as response, temporary.open("wb") as target: + while chunk := response.read(1024 * 1024): + target.write(chunk) + if _digest(temporary, "md5") != DATASET_MD5: + temporary.unlink(missing_ok=True) + raise ValueError("downloaded SciFact archive failed the published MD5 check") + os.replace(temporary, archive) + + required = ( + dataset_dir / "corpus.jsonl", + dataset_dir / "queries.jsonl", + dataset_dir / "qrels" / "test.tsv", + ) + if not all(path.is_file() for path in required): + extraction_root = data_dir.resolve() + with zipfile.ZipFile(archive) as bundle: + for member in bundle.infolist(): + destination = (data_dir / member.filename).resolve() + if not destination.is_relative_to(extraction_root): + raise ValueError(f"unsafe archive member: {member.filename}") + bundle.extractall(data_dir) + if not all(path.is_file() for path in required): + raise FileNotFoundError("SciFact archive is missing corpus, queries, or test qrels") + return dataset_dir + + +def _load_jsonl(path: Path) -> dict[str, dict[str, Any]]: + rows: dict[str, dict[str, Any]] = {} + for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1): + if not line.strip(): + continue + row = json.loads(line) + row_id = str(row.get("_id", "")).strip() + if not row_id or row_id in rows: + raise ValueError(f"invalid or duplicate _id at {path}:{line_number}") + rows[row_id] = row + return rows + + +def _load_qrels(path: Path) -> dict[str, set[str]]: + result: dict[str, set[str]] = defaultdict(set) + with path.open(encoding="utf-8", newline="") as source: + for row in csv.DictReader(source, delimiter="\t"): + if int(row["score"]) > 0: + result[str(row["query-id"])].add(str(row["corpus-id"])) + return dict(result) + + +def _document_text(row: dict[str, Any]) -> str: + return "\n".join(part.strip() for part in (str(row.get("title", "")), str(row.get("text", ""))) if part.strip()) + + +def _model_identity(model_source: str) -> str: + path = Path(model_source) + if not path.exists(): + return model_source + resolved = path.resolve() + digest = hashlib.sha256(str(resolved).encode()) + for file_path in sorted(path for path in resolved.rglob("*") if path.is_file()): + stat = file_path.stat() + digest.update(str(file_path.relative_to(resolved)).encode()) + digest.update(f"\0{stat.st_size}\0{stat.st_mtime_ns}\0".encode()) + return digest.hexdigest() + + +def _embedding_cache_path(data_dir: Path, model_identity: str) -> Path: + key = hashlib.sha256(f"v2\0{DATASET_MD5}\0{model_identity}".encode()).hexdigest()[:12] + return data_dir / f"scifact-embeddings-{key}.npz" + + +def _document_vectors( + embeddings: Any, + doc_ids: list[str], + doc_texts: list[str], + cache_path: Path, +) -> tuple[np.ndarray, int, int]: + text_hashes = [hashlib.sha256(text.encode()).hexdigest() for text in doc_texts] + cached: dict[str, tuple[str, np.ndarray]] = {} + if cache_path.is_file(): + with np.load(cache_path, allow_pickle=False) as cache: + if {"doc_ids", "text_hashes", "vectors"}.difference(cache.files): + raise ValueError(f"invalid embedding cache: {cache_path}") + cached_ids = np.asarray(cache["doc_ids"], dtype=str).tolist() + cached_hashes = np.asarray(cache["text_hashes"], dtype=str).tolist() + cached_vectors = np.asarray(cache["vectors"], dtype=np.float32) + if ( + cached_vectors.ndim != 2 + or len(cached_ids) != len(cached_vectors) + or len(cached_hashes) != len(cached_vectors) + or len(set(cached_ids)) != len(cached_ids) + ): + raise ValueError(f"invalid embedding cache: {cache_path}") + cached = dict(zip(cached_ids, zip(cached_hashes, cached_vectors, strict=True), strict=True)) + + missing_indexes = [ + index + for index, (doc_id, text_hash) in enumerate(zip(doc_ids, text_hashes, strict=True)) + if doc_id not in cached or cached[doc_id][0] != text_hash + ] + if missing_indexes: + missing_vectors = np.asarray( + embeddings.embed_documents([doc_texts[index] for index in missing_indexes]), + dtype=np.float32, + ) + if missing_vectors.ndim != 2 or len(missing_vectors) != len(missing_indexes): + raise ValueError("embedding model returned an unexpected document vector shape") + if cached and missing_vectors.shape[1] != len(next(iter(cached.values()))[1]): + raise ValueError(f"embedding dimensions do not match cache: {cache_path}") + cached.update( + (doc_ids[index], (text_hashes[index], vector)) + for index, vector in zip(missing_indexes, missing_vectors, strict=True) + ) + temporary = cache_path.with_suffix(".npz.part") + with temporary.open("wb") as target: + np.savez_compressed( + target, + doc_ids=np.asarray(list(cached)), + text_hashes=np.asarray([item[0] for item in cached.values()]), + vectors=np.asarray([item[1] for item in cached.values()], dtype=np.float32), + ) + os.replace(temporary, cache_path) + + return ( + np.stack([cached[doc_id][1] for doc_id in doc_ids]), + len(doc_ids) - len(missing_indexes), + len(missing_indexes), + ) + + +def _top_ids(scores: np.ndarray, doc_ids: list[str], limit: int) -> list[str]: + limit = min(limit, len(doc_ids)) + candidates = np.argpartition(-scores, limit - 1)[:limit] + ranked = candidates[np.argsort(-scores[candidates], kind="stable")] + return [doc_ids[index] for index in ranked] + + +def _rrf( + left: list[str], + right: list[str], + limit: int, + *, + left_weight: float = 1.0, + right_weight: float = 1.0, +) -> list[str]: + scores: dict[str, float] = defaultdict(float) + for ranking, weight in ((left, left_weight), (right, right_weight)): + for rank, doc_id in enumerate(ranking, 1): + scores[doc_id] += weight / (RRF_CONSTANT + rank) + return sorted(scores, key=lambda doc_id: (-scores[doc_id], doc_id))[:limit] + + +def _metrics(rankings: dict[str, list[str]], qrels: dict[str, set[str]], top_k: int) -> dict[str, float]: + recalls: list[float] = [] + reciprocal_ranks: list[float] = [] + ndcgs: list[float] = [] + for query_id, ranking in rankings.items(): + expected = sorted(qrels[query_id]) + retrieved = ranking[:top_k] + recalls.append(calculate_recall_at_k(expected_ids=expected, retrieved_ids=retrieved)) + reciprocal_ranks.append(calculate_mrr(expected_ids=expected, retrieved_ids=retrieved)) + ndcgs.append(calculate_ndcg(expected_ids=expected, retrieved_ids=retrieved)) + return { + f"recall@{top_k}": round(fmean(recalls), 6), + f"mrr@{top_k}": round(fmean(reciprocal_ranks), 6), + f"ndcg@{top_k}": round(fmean(ndcgs), 6), + } + + +def _failures( + rankings: dict[str, list[str]], + qrels: dict[str, set[str]], + queries: dict[str, dict[str, Any]], + top_k: int, +) -> list[dict[str, Any]]: + failures = [] + for query_id, ranking in rankings.items(): + expected = qrels[query_id] + retrieved = ranking[:top_k] + missing = sorted(expected.difference(retrieved)) + if missing: + failures.append( + { + "query_id": query_id, + "query": str(queries[query_id].get("text", "")), + "expected": sorted(expected), + "missing": missing, + "retrieved": retrieved, + } + ) + return failures[:20] + + +def _git_state() -> dict[str, Any]: + environment_commit = os.getenv("BEIR_EVAL_GIT_COMMIT", "").strip() + environment_dirty = os.getenv("BEIR_EVAL_GIT_DIRTY", "").strip().lower() + if environment_commit: + return { + "commit": environment_commit, + "dirty": environment_dirty == "true" if environment_dirty else None, + } + try: + commit = subprocess.run( + ["git", "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + dirty = bool( + subprocess.run( + ["git", "status", "--porcelain"], + check=True, + capture_output=True, + text=True, + ).stdout.strip() + ) + return {"commit": commit, "dirty": dirty} + except (FileNotFoundError, subprocess.CalledProcessError): + return {"commit": None, "dirty": None} + + +def run(args: argparse.Namespace) -> dict[str, Any]: + _validate_args(args) + dataset_dir = _download_dataset(args.data_dir) + corpus = _load_jsonl(dataset_dir / "corpus.jsonl") + queries = _load_jsonl(dataset_dir / "queries.jsonl") + qrels = _load_qrels(dataset_dir / "qrels" / "test.tsv") + eligible_query_ids = sorted(set(queries).intersection(qrels)) + if args.sample_size > len(eligible_query_ids): + raise ValueError(f"--sample-size exceeds the {len(eligible_query_ids)} SciFact test queries") + query_ids = sorted(random.Random(args.seed).sample(eligible_query_ids, args.sample_size)) + positive_doc_ids = sorted({doc_id for query_id in query_ids for doc_id in qrels[query_id]}) + if args.corpus_size < len(positive_doc_ids): + raise ValueError( + f"--corpus-size must retain all {len(positive_doc_ids)} positive documents for the sampled queries" + ) + if args.corpus_size > len(corpus): + raise ValueError(f"--corpus-size exceeds the {len(corpus)} SciFact documents") + distractor_ids = sorted(set(corpus).difference(positive_doc_ids)) + sampled_distractors = random.Random(args.seed + 1).sample( + distractor_ids, + args.corpus_size - len(positive_doc_ids), + ) + doc_ids = sorted([*positive_doc_ids, *sampled_distractors]) + doc_texts = [_document_text(corpus[doc_id]) for doc_id in doc_ids] + query_texts = [str(queries[query_id].get("text", "")) for query_id in query_ids] + timings: dict[str, float] = {} + + embedding_model = str(settings.EMBEDDING_MODEL_PATH or settings.EMBEDDING_MODEL_NAME) + embeddings = get_embeddings() + embedding_source, _ = resolve_embedding_source() + embedding_cache = _embedding_cache_path(args.data_dir, _model_identity(embedding_source)) + embeddings.encode_kwargs["batch_size"] = args.embedding_batch_size + started = time.perf_counter() + doc_vectors, cache_hits, encoded_count = _document_vectors( + embeddings, + doc_ids, + doc_texts, + embedding_cache, + ) + timings["dense_corpus_encode_seconds"] = round(time.perf_counter() - started, 3) + timings["dense_corpus_cache_hits"] = cache_hits + timings["dense_corpus_encoded_count"] = encoded_count + started = time.perf_counter() + query_vectors = np.asarray( + embeddings.embed_documents(query_texts), + dtype=np.float32, + ) + timings["dense_query_encode_seconds"] = round(time.perf_counter() - started, 3) + started = time.perf_counter() + dense_scores = query_vectors @ doc_vectors.T + dense = { + query_id: _top_ids(dense_scores[index], doc_ids, args.candidate_k) for index, query_id in enumerate(query_ids) + } + timings["dense_search_seconds"] = round(time.perf_counter() - started, 3) + + # ponytail: TF-IDF is an offline lexical proxy; use a live PostgreSQL run when backend parity matters. + started = time.perf_counter() + vectorizer = TfidfVectorizer(lowercase=True, stop_words="english", ngram_range=(1, 2), min_df=2) + lexical_corpus = vectorizer.fit_transform(doc_texts) + lexical_queries = vectorizer.transform(query_texts) + timings["lexical_fit_seconds"] = round(time.perf_counter() - started, 3) + started = time.perf_counter() + lexical = { + query_id: _top_ids((lexical_queries[index] @ lexical_corpus.T).toarray().ravel(), doc_ids, args.candidate_k) + for index, query_id in enumerate(query_ids) + } + timings["lexical_search_seconds"] = round(time.perf_counter() - started, 3) + started = time.perf_counter() + hybrid = {query_id: _rrf(dense[query_id], lexical[query_id], args.candidate_k) for query_id in query_ids} + weighted_hybrid = { + query_id: _rrf( + dense[query_id], + lexical[query_id], + args.candidate_k, + right_weight=LEXICAL_RRF_WEIGHT, + ) + for query_id in query_ids + } + timings["hybrid_rrf_seconds"] = round(time.perf_counter() - started, 3) + + rankings = { + "dense": dense, + "lexical_tfidf": lexical, + "hybrid_rrf": hybrid, + "dense_weighted_rrf": weighted_hybrid, + } + reranker_model = None + if args.reranker: + reranker = get_reranker() + if reranker is None: + raise RuntimeError("--reranker requires RERANKER_ENABLED=true") + reranker_model = settings.RERANKER_MODEL_PATH or settings.RERANKER_MODEL_NAME + pairs: list[tuple[str, str]] = [] + candidates_by_pipeline: dict[str, dict[str, list[str]]] = {} + doc_index = {doc_id: index for index, doc_id in enumerate(doc_ids)} + for result_name, source in ( + ("dense_rerank", dense), + ("hybrid_rerank", weighted_hybrid), + ): + candidates_by_query: dict[str, list[str]] = {} + for query_id in query_ids: + candidates = source[query_id][: args.rerank_k] + candidates_by_query[query_id] = candidates + pairs.extend((str(queries[query_id]["text"]), doc_texts[doc_index[doc_id]]) for doc_id in candidates) + candidates_by_pipeline[result_name] = candidates_by_query + started = time.perf_counter() + scores = reranker.predict( + pairs, + batch_size=args.reranker_batch_size, + show_progress_bar=True, + ) + timings["rerank_seconds"] = round(time.perf_counter() - started, 3) + timings["rerank_pair_count"] = len(pairs) + offset = 0 + for result_name, candidates_by_query in candidates_by_pipeline.items(): + reranked: dict[str, list[str]] = {} + for query_id in query_ids: + candidates = candidates_by_query[query_id] + query_scores = scores[offset : offset + len(candidates)] + offset += len(candidates) + reranked[query_id] = [ + doc_id + for _, doc_id in sorted( + zip(query_scores, candidates, strict=True), + key=lambda item: -float(item[0]), + ) + ] + rankings[result_name] = reranked + + metrics = {name: _metrics(items, qrels, args.top_k) for name, items in rankings.items()} + return { + "schema_version": 1, + "created_at": datetime.now(UTC).isoformat(), + "repository": _git_state(), + "dataset": { + "name": DATASET_NAME, + "download_url": DATASET_URL, + "published_md5": DATASET_MD5, + "archive_sha256": _digest(args.data_dir / "scifact.zip"), + "hugging_face_card": HF_DATASET_CARD, + "hugging_face_revision": HF_DATASET_REVISION, + "full_corpus_count": len(corpus), + "evaluated_corpus_count": len(doc_ids), + "positive_document_count": len(positive_doc_ids), + "corpus_sampling": "all selected-query qrel positives plus fixed-seed distractors", + "test_query_count": len(eligible_query_ids), + "evaluated_query_count": len(query_ids), + "query_ids": query_ids, + "sample_seed": args.seed, + }, + "models": { + "embedding": embedding_model, + "embedding_source": embedding_source, + "reranker": reranker_model, + "lexical": "sklearn.TfidfVectorizer", + }, + "configuration": { + "top_k": args.top_k, + "corpus_size": args.corpus_size, + "candidate_k": args.candidate_k, + "rerank_k": args.rerank_k if args.reranker else None, + "embedding_batch_size": args.embedding_batch_size, + "embedding_cache": str(embedding_cache), + "reranker_batch_size": args.reranker_batch_size if args.reranker else None, + "rrf_constant": RRF_CONSTANT, + "production_candidate_k": FUSION_CANDIDATE_K, + "dense_rrf_weight": 1.0, + "lexical_rrf_weight": LEXICAL_RRF_WEIGHT, + }, + "metrics": metrics, + "timings": timings, + "failure_samples": {name: _failures(items, qrels, queries, args.top_k) for name, items in rankings.items()}, + } + + +def main(argv: list[str] | None = None) -> int: + args = _parser().parse_args(argv) + try: + report = run(args) + except Exception as exc: + print(f"BEIR evaluation failed: {exc}", file=sys.stderr) + return 1 + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps({"metrics": report["metrics"], "output": str(args.output)}, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/app/mneme/memoria/server/memory/reconciliation.py b/app/mneme/memoria/server/memory/reconciliation.py index 0f4c26d..1356bac 100644 --- a/app/mneme/memoria/server/memory/reconciliation.py +++ b/app/mneme/memoria/server/memory/reconciliation.py @@ -3,6 +3,7 @@ Slot locks, scoped evidence, fingerprints, and stale-candidate checks prevent lost updates and silent overwrites. """ +import logging from dataclasses import dataclass from datetime import datetime from typing import Any @@ -27,6 +28,7 @@ Sensitivity, ) from app.mneme.memoria.server.models.memory_revision import MemoryRevision +from app.mneme.memoria.server.observability.context import safe_log from app.mneme.memoria.server.repositories.memories import ( attach_candidate_evidence, attach_revision_evidence, @@ -41,6 +43,9 @@ utc_now, validate_evidence_scope, ) +from app.mneme.memoria.server.services.embeddings import embed_texts, memory_embedding_text + +logger = logging.getLogger(__name__) @dataclass(frozen=True) @@ -130,6 +135,23 @@ def _combined_sensitivity(current: str, additional: Sensitivity) -> CanonicalSen return max((current, additional), key=rank.__getitem__) # type: ignore[return-value] +async def _revision_embedding(*, subject: str, predicate: str, value: str) -> list[float] | None: + # ponytail: inline best-effort encoding avoids another worker; use an outbox + # projection if memory-write latency or lock contention becomes measurable. + try: + return (await embed_texts([memory_embedding_text(subject=subject, predicate=predicate, value=value)]))[0] + except Exception: + safe_log( + logger, + logging.WARNING, + "runtime_log", + phase="memory_revision_embedding", + status="degraded", + error_code="MEMORY_REVISION_EMBEDDING_FAILED", + ) + return None + + async def _create_memory( db: AsyncSession, *, @@ -146,6 +168,7 @@ async def _create_memory( actor: str, evidence_ids: list[str], ) -> CanonicalMemory: + embedding = await _revision_embedding(subject=subject, predicate=predicate, value=value) memory_id = new_id() revision_id = new_id() memory = CanonicalMemory( @@ -173,6 +196,7 @@ async def _create_memory( predicate=predicate, value=value, fingerprint=fingerprint, + embedding=embedding, reason=reason, actor=actor, ) @@ -450,6 +474,11 @@ async def revise_memory( owner_id=owner_id, knowledge_base_id=knowledge_base_id, ) + embedding = await _revision_embedding( + subject=normalized_subject, + predicate=normalized_predicate, + value=normalized_value, + ) now = utc_now() # Close the previous validity interval before making the new revision @@ -469,6 +498,7 @@ async def revise_memory( predicate=normalized_predicate, value=normalized_value, fingerprint=fingerprint, + embedding=embedding, valid_from=now, reason=reason, actor=actor, diff --git a/app/mneme/memoria/server/models/memory_revision.py b/app/mneme/memoria/server/models/memory_revision.py index defd3b1..d520ef0 100644 --- a/app/mneme/memoria/server/models/memory_revision.py +++ b/app/mneme/memoria/server/models/memory_revision.py @@ -4,10 +4,13 @@ """ from datetime import datetime +from typing import Any +from pgvector.sqlalchemy import Vector from sqlalchemy import CheckConstraint, DateTime, ForeignKey, Index, String, Text, UniqueConstraint, func, text from sqlalchemy.orm import Mapped, mapped_column +from app.mneme.memoria.server.config import settings from app.mneme.memoria.server.models.base import Base @@ -20,6 +23,12 @@ class MemoryRevision(Base): ), UniqueConstraint("revision_id", "memory_id"), Index("ix_memory_revisions_owner_scope", "owner_id", "knowledge_base_id"), + Index( + "ix_memory_revisions_embedding_hnsw", + "embedding", + postgresql_using="hnsw", + postgresql_ops={"embedding": "vector_cosine_ops"}, + ), Index( "uq_memory_revisions_open_memory", "memory_id", @@ -38,6 +47,7 @@ class MemoryRevision(Base): predicate: Mapped[str] = mapped_column(Text, nullable=False) value: Mapped[str] = mapped_column(Text, nullable=False) fingerprint: Mapped[str] = mapped_column(String(64), nullable=False) + embedding: Mapped[Any | None] = mapped_column(Vector(settings.EMBEDDING_DIMENSION), nullable=True) valid_from: Mapped[datetime] = mapped_column( DateTime(timezone=True), nullable=False, server_default=func.now() ) diff --git a/app/mneme/memoria/server/retrieval/documents.py b/app/mneme/memoria/server/retrieval/documents.py index 460e5fe..08078a3 100644 --- a/app/mneme/memoria/server/retrieval/documents.py +++ b/app/mneme/memoria/server/retrieval/documents.py @@ -5,7 +5,10 @@ from app.mneme.memoria.server.database import open_read_session from app.mneme.memoria.server.retrieval.contracts import RetrievalScope, RetrievedEvidence -from app.mneme.memoria.server.retrieval.fusion import reciprocal_rank_fusion +from app.mneme.memoria.server.retrieval.fusion import ( + LEXICAL_RRF_WEIGHT, + reciprocal_rank_fusion, +) from app.mneme.memoria.server.retrieval.keyword import search_keyword from app.mneme.memoria.server.retrieval.vector import search_vector @@ -33,4 +36,8 @@ async def search( async with open_read_session() as db: vector_hits = await search_vector(db, scope=scope, query=query, limit=top_k) keyword_hits = await search_keyword(db, scope=scope, query=query, limit=top_k) - return reciprocal_rank_fusion((vector_hits, keyword_hits), top_k=top_k) + return reciprocal_rank_fusion( + (vector_hits, keyword_hits), + top_k=top_k, + weights=(1.0, LEXICAL_RRF_WEIGHT), + ) diff --git a/app/mneme/memoria/server/retrieval/fusion.py b/app/mneme/memoria/server/retrieval/fusion.py index b13c697..91def43 100644 --- a/app/mneme/memoria/server/retrieval/fusion.py +++ b/app/mneme/memoria/server/retrieval/fusion.py @@ -8,32 +8,40 @@ from app.mneme.memoria.server.retrieval.contracts import DocumentSearchHit, RetrievedEvidence RRF_CONSTANT = 60 +FUSION_CANDIDATE_K = 50 +LEXICAL_RRF_WEIGHT = 0.3 def reciprocal_rank_fusion( rankings: Sequence[Sequence[DocumentSearchHit]], *, top_k: int, + weights: Sequence[float] | None = None, ) -> list[RetrievedEvidence]: """Fuse backend rankings without comparing their incompatible raw scores. - Each unique chunk contributes ``1 / (60 + rank)`` once per ranking. + Each unique chunk contributes ``weight / (60 + rank)`` once per ranking. + Omitted weights preserve equal fusion for existing callers. Deterministic chunk-ID tie breaking keeps replayed evaluations stable. """ if top_k <= 0: return [] + if weights is not None and len(weights) != len(rankings): + raise ValueError("weights must match the number of rankings") hits_by_chunk_id: dict[str, DocumentSearchHit] = {} scores_by_chunk_id: dict[str, float] = {} - for ranking in rankings: + for ranking_index, ranking in enumerate(rankings): + weight = weights[ranking_index] if weights is not None else 1.0 seen_chunk_ids: set[str] = set() for rank, hit in enumerate(ranking, start=1): if hit.chunk_id in seen_chunk_ids: continue seen_chunk_ids.add(hit.chunk_id) hits_by_chunk_id.setdefault(hit.chunk_id, hit) - scores_by_chunk_id[hit.chunk_id] = scores_by_chunk_id.get(hit.chunk_id, 0.0) + 1.0 / ( - RRF_CONSTANT + rank + scores_by_chunk_id[hit.chunk_id] = ( + scores_by_chunk_id.get(hit.chunk_id, 0.0) + + weight / (RRF_CONSTANT + rank) ) ranked_chunk_ids = sorted( diff --git a/app/mneme/memoria/server/retrieval/keyword.py b/app/mneme/memoria/server/retrieval/keyword.py index 73d4397..4af9521 100644 --- a/app/mneme/memoria/server/retrieval/keyword.py +++ b/app/mneme/memoria/server/retrieval/keyword.py @@ -9,6 +9,7 @@ from app.mneme.memoria.server.models.document_chunk import DocumentChunk from app.mneme.memoria.server.models.document_projection import DocumentProjection from app.mneme.memoria.server.retrieval.contracts import DocumentSearchHit, RetrievalScope +from app.mneme.memoria.server.retrieval.fusion import FUSION_CANDIDATE_K async def search_keyword( @@ -20,8 +21,8 @@ async def search_keyword( ) -> list[DocumentSearchHit]: """Return full-text matches for the query from active scoped document chunks. - PostgreSQL ``simple`` configuration favors exact lexical identifiers; - callers combine this ranking with dense-vector results for recall. + PostgreSQL ``simple`` configuration favors exact lexical identifiers. The + backend keeps a wider candidate pool for final dense-weighted fusion. """ if limit <= 0: return [] @@ -53,7 +54,7 @@ async def search_keyword( DocumentChunk.search_vector.op("@@")(text_query), ) .order_by(rank.desc(), DocumentChunk.chunk_id.asc()) - .limit(limit) + .limit(max(limit, FUSION_CANDIDATE_K)) ) rows = (await db.execute(statement)).mappings().all() return [ diff --git a/app/mneme/memoria/server/retrieval/memories.py b/app/mneme/memoria/server/retrieval/memories.py index 6548a8a..43c3dfd 100644 --- a/app/mneme/memoria/server/retrieval/memories.py +++ b/app/mneme/memoria/server/retrieval/memories.py @@ -6,7 +6,7 @@ from datetime import datetime from typing import Literal -from sqlalchemy import and_, func, or_, select +from sqlalchemy import and_, func, or_, select, text from app.mneme.memoria.server.database import open_read_session from app.mneme.memoria.server.models.canonical_memory import CanonicalMemory @@ -28,6 +28,7 @@ class MemoryRetriever: The retriever can select current or historical revisions and constrain memory types while preserving owner and optional knowledge-base scope. """ + async def search( self, *, @@ -39,6 +40,7 @@ async def search( memory_types: tuple[str, ...] | None = None, excluded_memory_types: tuple[str, ...] | None = None, evidence_type: Literal["memory", "profile"] = "memory", + query_embedding: list[float] | None = None, ) -> list[RetrievedEvidence]: """Search scoped memory revisions and return normalized evidence records. @@ -81,24 +83,51 @@ async def search( pattern = f"%{query.strip()}%" relevance = text_value.ilike(pattern) - statement = ( - select( + async with open_read_session() as db: + use_semantic = query_embedding is not None + if use_semantic: + missing_embedding = ( + select(MemoryRevision.revision_id) + .join(CanonicalMemory, MemoryRevision.memory_id == CanonicalMemory.memory_id) + .where(and_(*filters), MemoryRevision.embedding.is_(None)) + .limit(1) + ) + use_semantic = not bool(await db.scalar(select(missing_embedding.exists()))) + + semantic_distance = ( + MemoryRevision.embedding.cosine_distance(query_embedding).label("semantic_distance") + if use_semantic + else None + ) + columns = [ CanonicalMemory.memory_id, CanonicalMemory.memory_type, CanonicalMemory.confidence, + CanonicalMemory.retrieval_weight, MemoryRevision.revision_id, MemoryRevision.subject, MemoryRevision.predicate, MemoryRevision.value, MemoryRevision.valid_from, MemoryRevision.valid_to, + ] + if semantic_distance is not None: + await db.execute(text("SET LOCAL hnsw.ef_search = 100")) + await db.execute(text("SET LOCAL hnsw.iterative_scan = strict_order")) + columns.append(semantic_distance) + ordering = [semantic_distance.asc()] if semantic_distance is not None else [relevance.desc()] + statement = ( + select(*columns) + .join(MemoryRevision, MemoryRevision.memory_id == CanonicalMemory.memory_id) + .where(and_(*filters)) + .order_by( + *ordering, + CanonicalMemory.retrieval_weight.desc(), + CanonicalMemory.confidence.desc(), + MemoryRevision.valid_from.desc(), + ) + .limit(top_k) ) - .join(MemoryRevision, MemoryRevision.memory_id == CanonicalMemory.memory_id) - .where(and_(*filters)) - .order_by(relevance.desc(), CanonicalMemory.confidence.desc(), MemoryRevision.valid_from.desc()) - .limit(top_k) - ) - async with open_read_session() as db: rows = (await db.execute(statement)).mappings().all() return [ @@ -107,9 +136,18 @@ async def search( source_type=evidence_type, source_id=row["memory_id"], content=f"{row['subject']} {row['predicate']} {row['value']}", - score=float(row["confidence"]), + score=( + 1.0 - float(row["semantic_distance"]) + if row.get("semantic_distance") is not None + else float(row["confidence"]) + ), metadata={ "memory_type": row["memory_type"], + "confidence": float(row["confidence"]), + "retrieval_weight": float(row.get("retrieval_weight", 1.0)), + "semantic_score": ( + 1.0 - float(row["semantic_distance"]) if row.get("semantic_distance") is not None else None + ), "valid_from": _isoformat(row["valid_from"]), "valid_to": _isoformat(row["valid_to"]), }, diff --git a/app/mneme/memoria/server/retrieval/profile.py b/app/mneme/memoria/server/retrieval/profile.py index e148242..afeb870 100644 --- a/app/mneme/memoria/server/retrieval/profile.py +++ b/app/mneme/memoria/server/retrieval/profile.py @@ -15,6 +15,7 @@ class ProfileRetriever: The wrapper intentionally reuses MemoryRetriever so profile queries keep identical revision, validity, and authorization behavior. """ + def __init__(self, memories: MemoryRetriever | None = None) -> None: self._memories = memories or MemoryRetriever() @@ -25,6 +26,7 @@ async def search( knowledge_base_id: str | None, query: str, top_k: int, + query_embedding: list[float] | None = None, ) -> list[RetrievedEvidence]: """Retrieve current profile facts, preferences, goals, and constraints. @@ -39,4 +41,5 @@ async def search( temporal_scope="current", memory_types=PROFILE_MEMORY_TYPES, evidence_type="profile", + query_embedding=query_embedding, ) diff --git a/app/mneme/memoria/server/retrieval/vector.py b/app/mneme/memoria/server/retrieval/vector.py index a7e7cd6..0fd3513 100644 --- a/app/mneme/memoria/server/retrieval/vector.py +++ b/app/mneme/memoria/server/retrieval/vector.py @@ -3,12 +3,13 @@ Owner, knowledge-base, projection, and active-state filters are applied before ranking and limiting. """ -from sqlalchemy import select +from sqlalchemy import select, text from sqlalchemy.ext.asyncio import AsyncSession from app.mneme.memoria.server.models.document_chunk import DocumentChunk from app.mneme.memoria.server.models.document_projection import DocumentProjection from app.mneme.memoria.server.retrieval.contracts import DocumentSearchHit, RetrievalScope +from app.mneme.memoria.server.retrieval.fusion import FUSION_CANDIDATE_K from app.mneme.memoria.server.services.embeddings import embed_texts @@ -21,7 +22,8 @@ async def search_vector( ) -> list[DocumentSearchHit]: """Return the nearest active document chunks in the authorized scope. - Query embeddings are normalized by the embedding service. Owner, + The backend keeps a wider candidate pool for final fusion. Query embeddings + are normalized by the embedding service. Owner, knowledge-base, and projection filters are part of the SQL statement so unauthorized candidates never participate in nearest-neighbor ranking. """ @@ -29,6 +31,8 @@ async def search_vector( return [] query_embedding = (await embed_texts([query]))[0] + await db.execute(text("SET LOCAL hnsw.ef_search = 100")) + await db.execute(text("SET LOCAL hnsw.iterative_scan = strict_order")) distance = DocumentChunk.embedding.cosine_distance(query_embedding) statement = ( select( @@ -54,7 +58,7 @@ async def search_vector( DocumentProjection.status == "active", ) .order_by(distance.asc(), DocumentChunk.chunk_id.asc()) - .limit(limit) + .limit(max(limit, FUSION_CANDIDATE_K)) ) rows = (await db.execute(statement)).mappings().all() return [ diff --git a/app/mneme/memoria/server/runtime/retriever.py b/app/mneme/memoria/server/runtime/retriever.py index 0ead98d..8944f15 100644 --- a/app/mneme/memoria/server/runtime/retriever.py +++ b/app/mneme/memoria/server/runtime/retriever.py @@ -14,6 +14,7 @@ from app.mneme.memoria.server.retrieval.profile import PROFILE_MEMORY_TYPES, ProfileRetriever from app.mneme.memoria.server.retrieval.relations import RelationRetriever from app.mneme.memoria.server.runtime.contracts import RetrievalRequest +from app.mneme.memoria.server.services.embeddings import embed_texts RRF_CONSTANT = 60 logger = logging.getLogger(__name__) @@ -25,6 +26,7 @@ class ScopedEvidenceRetriever: Each source receives the same owner and knowledge-base scope. A source failure is logged and isolated; all-source failure remains terminal. """ + def __init__( self, *, @@ -37,6 +39,8 @@ def __init__( self._memories_factory = memories_factory self._profile_factory = profile_factory self._relations_factory = relations_factory + self._embed_memory_queries = memories_factory is MemoryRetriever + self._embed_profile_queries = profile_factory is ProfileRetriever async def retrieve(self, request: RetrievalRequest) -> list[RetrievedEvidence]: """Execute selected evidence-source searches concurrently and fuse results. @@ -50,11 +54,36 @@ async def retrieve(self, request: RetrievalRequest) -> list[RetrievedEvidence]: """ if not request.plan.uses_private_sources: return [] + if request.plan.document and request.knowledge_base_id is None: + raise ValueError("document retrieval requires a knowledge base scope") searches: list[tuple[str, Awaitable[list[RetrievedEvidence]]]] = [] + + async def embed_memory_query() -> list[float] | None: + try: + return (await embed_texts([request.question]))[0] + except asyncio.CancelledError: + raise + except Exception: + safe_log( + logger, + logging.WARNING, + "answer_phase", + phase="retrieve_memory_embedding", + status="failed", + error_code="AGENT_RETRIEVAL_EMBEDDING_FAILED", + ) + return None + + memory_embedding_task = ( + asyncio.create_task(embed_memory_query()) + if ( + (request.plan.memory and self._embed_memory_queries) + or (request.plan.profile and self._embed_profile_queries) + ) + else None + ) if request.plan.document: - if request.knowledge_base_id is None: - raise ValueError("document retrieval requires a knowledge base scope") searches.append( ( "document", @@ -69,29 +98,49 @@ async def retrieve(self, request: RetrievalRequest) -> list[RetrievedEvidence]: ) ) if request.plan.memory: + + async def search_memories() -> list[RetrievedEvidence]: + values = dict( + owner_id=request.owner_id, + knowledge_base_id=request.knowledge_base_id, + query=request.question, + top_k=request.top_k, + temporal_scope=request.temporal_scope, + excluded_memory_types=PROFILE_MEMORY_TYPES if request.plan.profile else None, + ) + if self._embed_memory_queries: + assert memory_embedding_task is not None + query_embedding = await memory_embedding_task + if query_embedding is not None: + values["query_embedding"] = query_embedding + return await self._memories_factory().search(**values) + searches.append( ( "memory", - self._memories_factory().search( - owner_id=request.owner_id, - knowledge_base_id=request.knowledge_base_id, - query=request.question, - top_k=request.top_k, - temporal_scope=request.temporal_scope, - excluded_memory_types=PROFILE_MEMORY_TYPES if request.plan.profile else None, - ), + search_memories(), ) ) if request.plan.profile: + + async def search_profile() -> list[RetrievedEvidence]: + values = dict( + owner_id=request.owner_id, + knowledge_base_id=request.knowledge_base_id, + query=request.question, + top_k=request.top_k, + ) + if self._embed_profile_queries: + assert memory_embedding_task is not None + query_embedding = await memory_embedding_task + if query_embedding is not None: + values["query_embedding"] = query_embedding + return await self._profile_factory().search(**values) + searches.append( ( "profile", - self._profile_factory().search( - owner_id=request.owner_id, - knowledge_base_id=request.knowledge_base_id, - query=request.question, - top_k=request.top_k, - ), + search_profile(), ) ) if request.plan.relations: diff --git a/app/mneme/memoria/server/services/embeddings.py b/app/mneme/memoria/server/services/embeddings.py index 8c53cef..9c81878 100644 --- a/app/mneme/memoria/server/services/embeddings.py +++ b/app/mneme/memoria/server/services/embeddings.py @@ -31,6 +31,18 @@ def embedding_model_ready() -> bool: return _embedding_model.cache_info().currsize > 0 +def memory_embedding_text(*, subject: str, predicate: str, value: str) -> str: + return "\n".join((subject.strip(), predicate.strip(), value.strip())) + + +def document_embedding_text(*, file_name: str, section_path: list[str], content: str) -> str: + return "\n".join( + part + for part in (file_name.strip(), " > ".join(section_path).strip(), content.strip()) + if part + ) + + def preload_embedding_model_sync() -> None: _embedding_model() diff --git a/app/mneme/memoria/server/services/projections.py b/app/mneme/memoria/server/services/projections.py index 492f7a3..3d29e78 100644 --- a/app/mneme/memoria/server/services/projections.py +++ b/app/mneme/memoria/server/services/projections.py @@ -25,7 +25,7 @@ load_projection_for_update, store_projection_batch, ) -from app.mneme.memoria.server.services.embeddings import embed_texts +from app.mneme.memoria.server.services.embeddings import document_embedding_text, embed_texts class IncompleteProjectionError(RuntimeError): @@ -312,7 +312,16 @@ async def finalize_projection(projection_id: str) -> bool: raise await db.commit() - embeddings = await embed_texts([chunk.content for chunk in prepared.chunks]) + embeddings = await embed_texts( + [ + document_embedding_text( + file_name=projection.file_name, + section_path=chunk.section_path, + content=chunk.content, + ) + for chunk in prepared.chunks + ] + ) try: async with db.begin(): return await _activate_projection(db, prepared, embeddings)