feat(retrieval): add BGE-M3 sparse document search

.env-example

@@ -42,6 +42,9 @@ EMBEDDING_MODEL_PATH=
 EMBEDDING_CACHE_DIR=./storage/model_cache/sentence_transformers
 EMBEDDING_LOCAL_FILES_ONLY=false
 EMBEDDING_PRELOAD_ON_STARTUP=false
+EMBEDDING_SPARSE_ENABLED=false
+EMBEDDING_SPARSE_HEAD_PATH=
+EMBEDDING_SPARSE_DIMENSION=250002
 RERANKER_ENABLED=false
 RERANKER_MODEL_NAME=BAAI/bge-reranker-v2-m3
 RERANKER_MODEL_PATH=

README.md

@@ -191,6 +191,8 @@ Mneme 使用两层配置：
 | `NEO4J_URI` / `NEO4J_PASSWORD` | 图谱服务连接 |
 | `MEMORY_AGENT_EMBEDDING_MODEL_NAME` | Memory Agent embedding 模型，默认 `BAAI/bge-m3` |
 | `MEMORY_AGENT_EMBEDDING_LOCAL_FILES_ONLY` | 是否只从本地缓存加载模型 |
+| `MEMORY_AGENT_EMBEDDING_SPARSE_ENABLED` | 是否启用 BGE-M3 sparse 文档检索；启用后需先回填活动 chunk |
+| `MEMORY_AGENT_EMBEDDING_SPARSE_HEAD_PATH` | 可选的本地 `sparse_linear.pt` 路径；为空时从 BGE-M3 缓存解析 |
 | `APP_HOST_PORT` | 宿主机监听地址，生产环境建议绑定 `127.0.0.1:8000` |
 
 ### Agent 配置入口
@@ -341,6 +343,10 @@ 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
 ```
 
+启用 BGE-M3 sparse 检索时，先执行数据库迁移，再设置
+`MEMORY_AGENT_EMBEDDING_SPARSE_ENABLED=true` 并运行上述 document backfill。任一 owner/KB
+仍有活动 chunk 缺少 sparse 向量时，检索会继续使用原有 dense + keyword 路径。
+
 完整删除、重建和恢复流程见[部署指南中的 Memoria operations](deploy/DEPLOY.md#memoria-rebuild-and-deletion-operations)。
 
 ## 项目结构

app/mneme/memoria/server/alembic/versions/20260901_01_add_document_sparse_embeddings.py

@@ -0,0 +1,41 @@
+"""add BGE-M3 sparse embeddings to document chunks
+
+Revision ID: 20260901_01
+Revises: 20260831_01
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from pgvector.sqlalchemy import SPARSEVEC
+
+from alembic import op
+from app.mneme.memoria.server.config import settings
+
+revision: str = "20260901_01"
+down_revision: str | Sequence[str] | None = "20260831_01"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+    op.add_column(
+        "document_chunks",
+        sa.Column(
+            "sparse_embedding",
+            SPARSEVEC(settings.EMBEDDING_SPARSE_DIMENSION),
+            nullable=True,
+        ),
+    )
+    op.create_index(
+        "ix_document_chunks_sparse_embedding_hnsw",
+        "document_chunks",
+        ["sparse_embedding"],
+        postgresql_using="hnsw",
+        postgresql_ops={"sparse_embedding": "sparsevec_ip_ops"},
+    )
+
+
+def downgrade() -> None:
+    op.drop_index("ix_document_chunks_sparse_embedding_hnsw", table_name="document_chunks")
+    op.drop_column("document_chunks", "sparse_embedding")

app/mneme/memoria/server/cli/embedding_backfill.py

@@ -3,6 +3,7 @@
 import argparse
 import asyncio
 import json
+from typing import Any
 
 from sqlalchemy import func, select, update
 
@@ -13,6 +14,7 @@
 from app.mneme.memoria.server.services.embeddings import (
     document_embedding_text,
     embed_texts,
+    embed_texts_with_sparse,
     memory_embedding_text,
 )
 
@@ -114,7 +116,7 @@ async def document_embedding_backfill(args: argparse.Namespace) -> dict:
         if not rows:
             break
 
-        vectors = await embed_texts(
+        vectors, sparse_vectors = await embed_texts_with_sparse(
             [
                 document_embedding_text(
                     file_name=row.file_name,
@@ -125,11 +127,14 @@ async def document_embedding_backfill(args: argparse.Namespace) -> dict:
             ]
         )
         async with open_write_session() as db:
-            for row, vector in zip(rows, vectors, strict=True):
+            for row, vector, sparse_vector in zip(rows, vectors, sparse_vectors, strict=True):
+                values: dict[str, Any] = {"embedding": vector}
+                if sparse_vector is not None:
+                    values["sparse_embedding"] = sparse_vector
                 result = await db.execute(
                     update(DocumentChunk)
                     .where(DocumentChunk.id == row.id, DocumentChunk.is_active.is_(True))
-                    .values(embedding=vector)
+                    .values(**values)
                 )
                 updated += result.rowcount or 0
         cursor = rows[-1].id

app/mneme/memoria/server/config.py

@@ -30,6 +30,9 @@ class MemoryAgentSettings(BaseSettings):
     EMBEDDING_CACHE_DIR: str = "storage/model_cache/sentence_transformers"
     EMBEDDING_LOCAL_FILES_ONLY: bool = False
     EMBEDDING_PRELOAD_ON_STARTUP: bool = False
+    EMBEDDING_SPARSE_ENABLED: bool = False
+    EMBEDDING_SPARSE_HEAD_PATH: str = ""
+    EMBEDDING_SPARSE_DIMENSION: int = 250002
     EXTRACTION_LLM_BASE_URL: str = ""
     EXTRACTION_LLM_API_KEY: SecretStr = SecretStr("")
     EXTRACTION_LLM_MODEL: str = ""

app/mneme/memoria/server/models/document_chunk.py

@@ -5,7 +5,7 @@
 
 from typing import Any
 
-from pgvector.sqlalchemy import Vector
+from pgvector.sqlalchemy import SPARSEVEC, Vector
 from sqlalchemy import Boolean, Computed, ForeignKey, Index, Integer, String, Text, UniqueConstraint, text
 from sqlalchemy.dialects.postgresql import JSONB, TSVECTOR
 from sqlalchemy.orm import Mapped, mapped_column
@@ -38,6 +38,12 @@ class DocumentChunk(Base):
             "search_vector",
             postgresql_using="gin",
         ),
+        Index(
+            "ix_document_chunks_sparse_embedding_hnsw",
+            "sparse_embedding",
+            postgresql_using="hnsw",
+            postgresql_ops={"sparse_embedding": "sparsevec_ip_ops"},
+        ),
         Index(
             "ix_document_chunks_active_scope",
             "owner_id",
@@ -68,4 +74,8 @@ class DocumentChunk(Base):
     page_no: Mapped[int | None] = mapped_column(Integer)
     section_path: Mapped[list[str]] = mapped_column(JSONB, nullable=False, default=list)
     embedding: Mapped[Any] = mapped_column(Vector(settings.EMBEDDING_DIMENSION), nullable=False)
+    sparse_embedding: Mapped[Any | None] = mapped_column(
+        SPARSEVEC(settings.EMBEDDING_SPARSE_DIMENSION),
+        nullable=True,
+    )
     is_active: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default="false")

app/mneme/memoria/server/retrieval/documents.py

@@ -6,11 +6,19 @@
 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 (
+    BGE_M3_DENSE_SCORE_WEIGHT,
+    BGE_M3_KEYWORD_SCORE_WEIGHT,
+    BGE_M3_SPARSE_SCORE_WEIGHT,
     DENSE_SCORE_WEIGHT,
     normalized_score_fusion,
 )
 from app.mneme.memoria.server.retrieval.keyword import search_keyword
+from app.mneme.memoria.server.retrieval.sparse import search_sparse, sparse_scope_ready
 from app.mneme.memoria.server.retrieval.vector import search_vector
+from app.mneme.memoria.server.services.embeddings import (
+    embed_texts_with_sparse,
+    sparse_embeddings_enabled,
+)
 
 
 class DocumentRetriever:
@@ -33,9 +41,45 @@ async def search(
         if top_k <= 0:
             return []
 
+        sparse_embedding = None
+        if sparse_embeddings_enabled():
+            dense_embeddings, sparse_embeddings = await embed_texts_with_sparse([query])
+            dense_embedding = dense_embeddings[0]
+            sparse_embedding = sparse_embeddings[0]
+
         async with open_read_session() as db:
-            vector_hits = await search_vector(db, scope=scope, query=query, limit=top_k)
+            vector_hits = (
+                await search_vector(
+                    db,
+                    scope=scope,
+                    query=query,
+                    limit=top_k,
+                    query_embedding=dense_embedding,
+                )
+                if sparse_embedding is not None
+                else await search_vector(db, scope=scope, query=query, limit=top_k)
+            )
             keyword_hits = await search_keyword(db, scope=scope, query=query, limit=top_k)
+            sparse_hits = (
+                await search_sparse(
+                    db,
+                    scope=scope,
+                    query_embedding=sparse_embedding,
+                    limit=top_k,
+                )
+                if sparse_embedding is not None and await sparse_scope_ready(db, scope=scope)
+                else []
+            )
+        if sparse_hits:
+            return normalized_score_fusion(
+                (vector_hits, sparse_hits, keyword_hits),
+                top_k=top_k,
+                weights=(
+                    BGE_M3_DENSE_SCORE_WEIGHT,
+                    BGE_M3_SPARSE_SCORE_WEIGHT,
+                    BGE_M3_KEYWORD_SCORE_WEIGHT,
+                ),
+            )
         return normalized_score_fusion(
             (vector_hits, keyword_hits),
             top_k=top_k,

app/mneme/memoria/server/retrieval/fusion.py

@@ -11,6 +11,9 @@
 FUSION_CANDIDATE_K = 200
 LEXICAL_RRF_WEIGHT = 0.55
 DENSE_SCORE_WEIGHT = 0.55
+BGE_M3_DENSE_SCORE_WEIGHT = 0.41
+BGE_M3_SPARSE_SCORE_WEIGHT = 0.54
+BGE_M3_KEYWORD_SCORE_WEIGHT = 0.05
 
 
 def reciprocal_rank_fusion(

app/mneme/memoria/server/retrieval/sparse.py

@@ -0,0 +1,86 @@
+"""Retrieve document chunks by BGE-M3 sparse inner-product similarity."""
+
+from pgvector import SparseVector
+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
+
+
+async def sparse_scope_ready(db: AsyncSession, *, scope: RetrievalScope) -> bool:
+    missing = (
+        select(DocumentChunk.id)
+        .join(DocumentProjection)
+        .where(
+            DocumentChunk.owner_id == scope.owner_id,
+            DocumentChunk.knowledge_base_id == scope.knowledge_base_id,
+            DocumentChunk.is_active.is_(True),
+            DocumentChunk.sparse_embedding.is_(None),
+            DocumentProjection.owner_id == scope.owner_id,
+            DocumentProjection.knowledge_base_id == scope.knowledge_base_id,
+            DocumentProjection.status == "active",
+        )
+        .limit(1)
+    )
+    return (await db.scalar(missing)) is None
+
+
+async def search_sparse(
+    db: AsyncSession,
+    *,
+    scope: RetrievalScope,
+    query_embedding: SparseVector,
+    limit: int,
+) -> list[DocumentSearchHit]:
+    if limit <= 0 or not query_embedding.indices():
+        return []
+
+    await db.execute(text("SET LOCAL hnsw.ef_search = 100"))
+    await db.execute(text("SET LOCAL hnsw.iterative_scan = strict_order"))
+    distance = DocumentChunk.sparse_embedding.max_inner_product(query_embedding)
+    similarity = (-distance).label("score")
+    statement = (
+        select(
+            DocumentChunk.chunk_id,
+            DocumentChunk.document_id,
+            DocumentChunk.document_version,
+            DocumentChunk.chunk_index,
+            DocumentChunk.content,
+            DocumentChunk.page_no,
+            DocumentChunk.section_path,
+            DocumentProjection.file_name,
+            similarity,
+        )
+        .join(DocumentProjection)
+        .where(
+            DocumentChunk.owner_id == scope.owner_id,
+            DocumentChunk.knowledge_base_id == scope.knowledge_base_id,
+            DocumentChunk.is_active.is_(True),
+            DocumentChunk.sparse_embedding.is_not(None),
+            DocumentProjection.owner_id == scope.owner_id,
+            DocumentProjection.knowledge_base_id == scope.knowledge_base_id,
+            DocumentProjection.status == "active",
+        )
+        .order_by(distance.asc(), DocumentChunk.chunk_id.asc())
+        .limit(max(limit, FUSION_CANDIDATE_K))
+    )
+    rows = (await db.execute(statement)).mappings().all()
+    return [
+        DocumentSearchHit(
+            chunk_id=row["chunk_id"],
+            document_id=row["document_id"],
+            content=row["content"],
+            score=float(row["score"]),
+            metadata={
+                "document_version": row["document_version"],
+                "file_name": row["file_name"],
+                "chunk_index": row["chunk_index"],
+                "page_no": row["page_no"],
+                "section_path": row["section_path"],
+            },
+        )
+        for row in rows
+    ]

app/mneme/memoria/server/retrieval/vector.py

@@ -19,6 +19,7 @@ async def search_vector(
     scope: RetrievalScope,
     query: str,
     limit: int,
+    query_embedding: list[float] | None = None,
 ) -> list[DocumentSearchHit]:
     """Return the nearest active document chunks in the authorized scope.
 
@@ -30,7 +31,8 @@ async def search_vector(
     if limit <= 0:
         return []
 
-    query_embedding = (await embed_texts([query]))[0]
+    if query_embedding is None:
+        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)