Projets/ObsiGate
1
1
Fork 0
Code Issues 16 Pull Requests Actions Packages Projects Releases 16 Wiki Activity

feat(search): recherche semantique - embeddings vectoriels + hybride RRF (#70)

Browse Source
This commit is contained in:
bruno
2026-09-12 00:12:20 -04:00
parent 45be3125de
commit 31d4be8015
19 changed files with 1388 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/main.py
+10 -2
View File
@@ -98,6 +98,7 @@ from backend.search import (
suggest_tags,
suggest_titles,
)
from backend.semantic_search import init_semantic_index
from backend.services.backups import diff_backup as service_diff_backup
from backend.services.backups import get_backup_dir as service_get_backup_dir
from backend.services.backups import list_backup_files as service_list_backup_files
@@ -281,7 +282,8 @@ class AdvancedSearchResultItem(BaseModel):
path: str = Field(description="Relative file path")
title: str = Field(description="File title")
tags: list[str] = Field(description="File tags")
score: float = Field(description="TF-IDF relevance score")
score: float = Field(description="TF-IDF relevance score (or fused RRF score in semantic mode)")
semantic_score: float = Field(default=0.0, description="Cosine similarity from the semantic index (0 when unavailable)")
snippet: str = Field(description="Content excerpt with <mark> highlights")
modified: str = Field(description="ISO 8601 modification timestamp")
extension: str = Field(default="", description="File extension")
@@ -301,6 +303,7 @@ class AdvancedSearchResponse(BaseModel):
limit: int = Field(description="Page size")
facets: SearchFacets = Field(description="Faceted counts by tag and vault")
query_time_ms: float = Field(default=0, description="Server-side query time in milliseconds")
semantic_available: bool = Field(default=False, description="True when the semantic (embedding) index is ready")
class TitleSuggestion(BaseModel):
@@ -694,6 +697,8 @@ async def lifespan(app: FastAPI):
# would freeze HTTP responses if run in the async event loop.
loop = asyncio.get_running_loop()
await loop.run_in_executor(_search_executor, init_inverted_index)
# Build the semantic (embedding) index in the same background thread pool.
await loop.run_in_executor(_search_executor, init_semantic_index)
# Scan for plugins in all vaults
logger.info("Scanning for plugins...")
@@ -2589,6 +2594,7 @@ async def api_advanced_search(
created: str | None = Query(None, description="Created date filter (>date, <date, date..date)"),
modified: str | None = Query(None, description="Modified date filter (>date, <date, date..date, <Nd)"),
size: str | None = Query(None, description="Size filter (>size, <size, size..size, e.g. >1MB, <10KB)"),
semantic: bool = Query(False, description="Fuse TF-IDF with semantic embeddings (RRF)"),
current_user=Depends(require_auth),
):
"""Advanced full-text search with TF-IDF scoring, facets, and pagination.
@@ -2605,6 +2611,8 @@ async def api_advanced_search(
- Remaining text is scored using TF-IDF with accent normalization.
- Toggles: case_sensitive, whole_word, regex
- Path filters: include_paths, exclude_paths (glob patterns)
- ``semantic=true`` — fuse the TF-IDF ranking with the semantic (embedding)
ranking via Reciprocal Rank Fusion and expose ``semantic_score`` per result.
Results include ``<mark>``-highlighted snippets and faceted tag/vault counts.
"""
@@ -2615,7 +2623,7 @@ async def api_advanced_search(
limit=limit, offset=offset, sort=sort,
case_sensitive=case_sensitive, whole_word=whole_word, regex=regex,
include_paths=include_paths, exclude_paths=exclude_paths,
created=created, modified=modified, size=size),
created=created, modified=modified, size=size, semantic=semantic),
)
backend/requirements-semantic.txt
+18
View File
@@ -0,0 +1,18 @@
# ObsiGate — Optional dependencies for semantic search (#70)
#
# These are NOT required: the semantic search engine degrades gracefully to a
# dependency-free hashing embedder and a pure-Python cosine store when they are
# absent. Install this file to enable the full local model + fast vector index:
#
# pip install -r backend/requirements-semantic.txt
#
# NOTE: sentence-transformers pulls in PyTorch (large download). If you only
# want the vector acceleration, install numpy + faiss-cpu and configure an
# external embedding endpoint instead (OBSIGATE_EMBEDDING_*).
# Local embedding model (all-MiniLM-L6-v2, ~80 MB, CPU)
sentence-transformers>=2.2.0
# Vector storage / similarity search
numpy>=1.24.0
faiss-cpu>=1.7.4
backend/search.py
+171 -54
View File
@@ -4,12 +4,14 @@ import re
import time
import unicodedata
from collections import defaultdict
from collections.abc import Callable
from typing import Any
from snowballstemmer import stemmer as _snowball_stemmer
from sortedcontainers import SortedList
from backend import indexer as _indexer
from backend import semantic_search as _semantic
from backend.indexer import index
logger = logging.getLogger("obsigate.search")
@@ -655,6 +657,10 @@ def _on_index_change_hook(action: str, vault_name: str, path: str, file_info: di
inv.remove_document(vault_name, path)
except Exception as e:
logger.warning(f"Inverted index incremental update failed ({action} {vault_name}/{path}): {e}")
try:
_semantic.on_index_change(action, vault_name, path, file_info)
except Exception as e:
logger.warning(f"Semantic index incremental update failed ({action} {vault_name}/{path}): {e}")
# Register the hook with indexer (indexer is already imported at top of file)
@@ -1102,6 +1108,7 @@ def advanced_search(
created: str | None = None,
modified: str | None = None,
size: str | None = None,
semantic: bool = False,
) -> dict[str, Any]:
"""Advanced full-text search with TF-IDF scoring, facets, and pagination.
@@ -1121,10 +1128,13 @@ def advanced_search(
limit: Max results per page.
offset: Pagination offset.
sort_by: ``"relevance"`` or ``"modified"``.
semantic: When True, fuse the TF-IDF ranking with the semantic
(embedding) ranking via Reciprocal Rank Fusion and expose a
``semantic_score`` per result.
Returns:
Dict with ``results``, ``total``, ``offset``, ``limit``, ``facets``,
``query_time_ms``.
``query_time_ms`` and ``semantic_available``.
"""
t0 = time.monotonic()
query = query.strip() if query else ""
@@ -1182,58 +1192,63 @@ def advanced_search(
# ------------------------------------------------------------------
# Step 2: Apply filters on candidate set
# ------------------------------------------------------------------
if effective_vault != "all":
candidates &= inv.vault_docs.get(effective_vault, set())
if all_tags and has_terms:
for t in all_tags:
candidates &= inv.tag_docs.get(t.lower(), set())
if parsed["title"]:
norm_title_filter = normalize_text(parsed["title"])
candidates = {
dk for dk in candidates
if norm_title_filter in normalize_text(inv.doc_info[dk].get("title", ""))
}
if parsed["path"]:
norm_path_filter = normalize_text(parsed["path"])
candidates = {
dk for dk in candidates
if norm_path_filter in normalize_text(inv.doc_info[dk].get("path", ""))
}
if parsed["ext"]:
ext_filter = parsed["ext"]
candidates = {
dk for dk in candidates
if (
inv.doc_info[dk].get("path", "").rsplit("/", 1)[-1].lower() == ext_filter
or inv.doc_info[dk].get("path", "").rsplit("/", 1)[-1].lower().endswith(f".{ext_filter}")
)
}
# Date and size filters (from query operators or API params)
date_range_created = _parse_date_range(created or parsed.get("created"))
if date_range_created:
candidates = {
dk for dk in candidates
if _matches_date_range(inv.doc_info[dk].get("created"), date_range_created)
}
date_range_modified = _parse_date_range(modified or parsed.get("modified"))
if date_range_modified:
candidates = {
dk for dk in candidates
if _matches_date_range(inv.doc_info[dk].get("modified"), date_range_modified)
}
size_range = _parse_size_range(size or parsed.get("size"))
if size_range:
candidates = {
dk for dk in candidates
if _matches_size_range(inv.doc_info[dk].get("size", 0), size_range)
}
def _apply_metadata_filters(docs: set) -> set:
"""Restrict a document-key set to the query's metadata filters."""
if effective_vault != "all":
docs &= inv.vault_docs.get(effective_vault, set())
if all_tags:
for t in all_tags:
docs &= inv.tag_docs.get(t.lower(), set())
if parsed["title"]:
norm_title_filter = normalize_text(parsed["title"])
docs = {
dk for dk in docs
if norm_title_filter in normalize_text(inv.doc_info[dk].get("title", ""))
}
if parsed["path"]:
norm_path_filter = normalize_text(parsed["path"])
docs = {
dk for dk in docs
if norm_path_filter in normalize_text(inv.doc_info[dk].get("path", ""))
}
if parsed["ext"]:
ext_filter = parsed["ext"]
docs = {
dk for dk in docs
if (
inv.doc_info[dk].get("path", "").rsplit("/", 1)[-1].lower() == ext_filter
or inv.doc_info[dk].get("path", "").rsplit("/", 1)[-1].lower().endswith(f".{ext_filter}")
)
}
if date_range_created:
docs = {
dk for dk in docs
if _matches_date_range(inv.doc_info[dk].get("created"), date_range_created)
}
if date_range_modified:
docs = {
dk for dk in docs
if _matches_date_range(inv.doc_info[dk].get("modified"), date_range_modified)
}
if size_range:
docs = {
dk for dk in docs
if _matches_size_range(inv.doc_info[dk].get("size", 0), size_range)
}
return docs
candidates = _apply_metadata_filters(candidates)
# ------------------------------------------------------------------
# Step 3: Score only the candidates (not all N documents)
@@ -1312,16 +1327,22 @@ def advanced_search(
"title": file_info["title"],
"tags": file_info.get("tags", []),
"score": round(score, 4),
"semantic_score": 0.0,
"snippet": snippet,
"modified": file_info.get("modified", ""),
"extension": file_info.get("extension", file_info.get("path", "").rsplit(".", 1)[-1] if "." in file_info.get("path", "") else ""),
}
scored_results.append((score, result))
# Facets
facet_vaults[vault_name] = facet_vaults.get(vault_name, 0) + 1
for tag in file_info.get("tags", []):
facet_tags[tag] = facet_tags.get(tag, 0) + 1
# ------------------------------------------------------------------
# Step 4: Optional semantic fusion (RRF with the TF-IDF ranking)
# ------------------------------------------------------------------
semantic_available = _semantic.get_semantic_index().is_ready()
if semantic and has_terms and not regex:
scored_results = _fuse_semantic_results(
scored_results, query, effective_vault, inv, _apply_metadata_filters,
include_paths, exclude_paths, limit,
)
# Sort
if sort_by == "modified":
@@ -1329,6 +1350,12 @@ def advanced_search(
else:
scored_results.sort(key=lambda x: -x[0])
# Facets are recomputed from the final result set (covers semantic-only docs)
for _, result in scored_results:
facet_vaults[result["vault"]] = facet_vaults.get(result["vault"], 0) + 1
for tag in result.get("tags", []):
facet_tags[tag] = facet_tags.get(tag, 0) + 1
total = len(scored_results)
page = scored_results[offset: offset + limit]
elapsed_ms = round((time.monotonic() - t0) * 1000, 1)
@@ -1343,9 +1370,99 @@ def advanced_search(
"vaults": dict(sorted(facet_vaults.items(), key=lambda x: -x[1])),
},
"query_time_ms": elapsed_ms,
"semantic_available": semantic_available,
}
def _fuse_semantic_results(
scored_results: list[tuple[float, dict[str, Any]]],
query: str,
vault_filter: str,
inv: InvertedIndex,
apply_metadata_filters: Callable[[set], set],
include_paths: str | None,
exclude_paths: str | None,
limit: int,
) -> list[tuple[float, dict[str, Any]]]:
"""Fuse the TF-IDF ranking with the semantic ranking using RRF.
Documents found only by the semantic engine are materialized from the
inverted index metadata (with a plain, non-highlighted snippet). The
returned tuples carry the fused score, and every result dict gets a
``semantic_score`` (cosine similarity, 0.0 when absent).
Args:
scored_results: Existing ``(tfidf_score, result_dict)`` tuples.
query: Raw free-text query.
vault_filter: Effective vault filter.
inv: Inverted index (document metadata source).
apply_metadata_filters: Callable restricting a doc-key set to the
query's tag/title/path/ext/date/size filters.
include_paths: Include glob patterns (or None).
exclude_paths: Exclude glob patterns (or None).
limit: Requested page size (drives how many semantic hits to fetch).
Returns:
New ``(fused_score, result_dict)`` list (unsorted).
"""
sem_hits = _semantic.semantic_search_docs(query, vault_filter=vault_filter, top_k=max(limit * 5, 200))
if not sem_hits:
return scored_results
# Semantic candidates must satisfy the same metadata + path filters.
semantic_universe = apply_metadata_filters(set(inv.doc_info.keys()))
semantic_universe = {
dk for dk in semantic_universe
if _passes_path_filters(inv.doc_info[dk].get("path", ""), include_paths, exclude_paths)
}
sem_scores: dict[str, float] = {}
sem_ranked: list[str] = []
for doc_key, similarity in sem_hits:
if doc_key not in semantic_universe:
continue
sem_scores[doc_key] = similarity
sem_ranked.append(doc_key)
if not sem_ranked:
return scored_results
existing: dict[str, dict[str, Any]] = {}
lexical_ranked: list[str] = []
for _, lex_result in sorted(scored_results, key=lambda item: -item[0]):
key = f"{lex_result['vault']}::{lex_result['path']}"
existing[key] = lex_result
lexical_ranked.append(key)
fused = _semantic.rrf_fuse([lexical_ranked, sem_ranked])
merged: list[tuple[float, dict[str, Any]]] = []
for doc_key, fused_score in fused.items():
result = existing.get(doc_key)
if result is None:
file_info = inv.doc_info.get(doc_key)
if file_info is None:
continue
content = file_info.get("content", "")
result = {
"vault": inv.doc_vault[doc_key],
"path": file_info["path"],
"title": file_info["title"],
"tags": file_info.get("tags", []),
"score": 0.0,
"semantic_score": 0.0,
"snippet": _escape_html(content[:200].strip()) if content else "",
"modified": file_info.get("modified", ""),
"extension": file_info.get(
"extension",
file_info.get("path", "").rsplit(".", 1)[-1] if "." in file_info.get("path", "") else "",
),
}
result["semantic_score"] = round(sem_scores.get(doc_key, 0.0), 4)
merged.append((fused_score, result))
return merged
# ---------------------------------------------------------------------------
# Suggestion helpers
# ---------------------------------------------------------------------------
backend/semantic_search.py
+612
View File
@@ -0,0 +1,612 @@
"""ObsiGate — Semantic search: embeddings, vector store and hybrid retrieval.
This module adds a *semantic* layer on top of the existing TF-IDF search. Each
document is split into overlapping chunks, each chunk is converted into a dense
vector, and queries are matched by cosine similarity. Results are combined with
the lexical ranking through Reciprocal Rank Fusion (RRF).
Design goals
------------
* **Zero mandatory dependency.** ``sentence-transformers`` (local model),
``numpy`` and ``faiss`` are *optional*. They are imported lazily and, when
missing, the module falls back to a deterministic pure-Python hashing embedder
and a pure-Python cosine store. The feature therefore degrades gracefully and
the default CI (which only installs ``backend/requirements.txt``) keeps working.
* **Plug-in providers.** Embeddings can come from the local
``all-MiniLM-L6-v2`` model, from an OpenAI-compatible ``/embeddings`` endpoint
(configured via env vars), or from the deterministic fallback.
* **Incremental.** The index is updated document-by-document from the indexer
change hook (file watcher + API mutations), never rebuilt on each search.
Optional extras are listed in ``backend/requirements-semantic.txt``.
"""
from __future__ import annotations
import hashlib
import logging
import math
import os
import re
import threading
from abc import ABC, abstractmethod
from collections import Counter
from itertools import pairwise
from typing import Any
logger = logging.getLogger("obsigate.semantic")
# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
EMBEDDING_DIM = 384 # all-MiniLM-L6-v2 output dimension
CHUNK_TOKENS = 512 # target chunk size (whitespace tokens)
CHUNK_OVERLAP_TOKENS = 64 # overlap between consecutive chunks
DEFAULT_TOP_K = 200 # max documents returned by a semantic query
RRF_K = 60 # Reciprocal Rank Fusion smoothing constant
MAX_QUERY_CHARS = 2000 # guard against pathological queries
_WORD_RE = re.compile(r"[\w]+", re.UNICODE)
# ---------------------------------------------------------------------------
# Tokenization / chunking
# ---------------------------------------------------------------------------
def _simple_tokens(text: str) -> list[str]:
"""Split *text* into lowercase word tokens (keeps accents)."""
return _WORD_RE.findall(text.lower())
def chunk_text(
text: str,
chunk_tokens: int = CHUNK_TOKENS,
overlap: int = CHUNK_OVERLAP_TOKENS,
) -> list[str]:
"""Split *text* into overlapping windows of roughly *chunk_tokens* words.
Args:
text: Raw document text.
chunk_tokens: Target number of whitespace tokens per chunk.
overlap: Number of tokens shared by two consecutive chunks.
Returns:
A list of chunk strings. Empty input yields an empty list.
"""
if not text or not text.strip():
return []
if chunk_tokens <= 0:
chunk_tokens = CHUNK_TOKENS
overlap = max(0, min(overlap, chunk_tokens - 1))
words = text.split()
if len(words) <= chunk_tokens:
return [" ".join(words)]
step = max(1, chunk_tokens - overlap)
chunks: list[str] = []
for start in range(0, len(words), step):
window = words[start:start + chunk_tokens]
if not window:
break
chunks.append(" ".join(window))
if start + chunk_tokens >= len(words):
break
return chunks
# ---------------------------------------------------------------------------
# Embedding providers
# ---------------------------------------------------------------------------
class EmbeddingProvider(ABC):
"""Base class for embedding backends."""
name: str = "base"
def __init__(self, dimension: int = EMBEDDING_DIM) -> None:
self.dimension = dimension
@abstractmethod
def encode(self, texts: list[str]) -> list[list[float]]:
"""Return one L2-normalized vector per input text."""
def encode_one(self, text: str) -> list[float]:
"""Convenience wrapper returning the vector for a single text."""
vectors = self.encode([text])
return vectors[0] if vectors else [0.0] * self.dimension
class HashEmbeddingProvider(EmbeddingProvider):
"""Deterministic, dependency-free hashing embedder.
This is a *lexical* fallback: it hashes word unigrams, word bigrams and
character trigrams into fixed-size signed buckets (the "hashing trick"),
then L2-normalizes the result. It captures shared vocabulary and
morphological variants (``backup``/``backups``), so it already improves
recall over exact TF-IDF matching, but it does not understand synonyms the
way a real transformer model does.
"""
name = "hash"
def _add_feature(self, vec: list[float], key: str, weight: float) -> None:
digest = hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest()
h = int.from_bytes(digest, "big")
idx = h % self.dimension
sign = 1.0 if (h >> 63) & 1 else -1.0
vec[idx] += sign * weight
def _encode_one(self, text: str) -> list[float]:
vec = [0.0] * self.dimension
tokens = _simple_tokens(text)
if not tokens:
return vec
tf = Counter(tokens)
for token, count in tf.items():
weight = 1.0 + math.log(count)
self._add_feature(vec, "w:" + token, weight)
for gram in _char_ngrams(token, 3):
self._add_feature(vec, "g:" + gram, weight * 0.5)
for first, second in pairwise(tokens):
self._add_feature(vec, "b:" + first + "_" + second, 0.5)
norm = math.sqrt(sum(v * v for v in vec))
if norm > 0.0:
vec = [v / norm for v in vec]
return vec
def encode(self, texts: list[str]) -> list[list[float]]:
return [self._encode_one(t or "") for t in texts]
def _char_ngrams(token: str, n: int) -> list[str]:
"""Return padded character n-grams for *token* (bounded to avoid blow-up)."""
if len(token) < n:
return [token]
if len(token) > 24:
token = token[:24]
return [token[i:i + n] for i in range(len(token) - n + 1)]
class SentenceTransformerProvider(EmbeddingProvider):
"""Local ``all-MiniLM-L6-v2`` embeddings via ``sentence-transformers``."""
name = "sentence-transformers"
def __init__(self, model_name: str = "all-MiniLM-L6-v2") -> None:
super().__init__(EMBEDDING_DIM)
self.model_name = model_name
self._model: Any | None = None
@staticmethod
def is_available() -> bool:
try:
import sentence_transformers # noqa: F401
except Exception:
return False
return True
def _get_model(self) -> Any:
if self._model is None:
from sentence_transformers import SentenceTransformer
self._model = SentenceTransformer(self.model_name)
return self._model
def encode(self, texts: list[str]) -> list[list[float]]:
if not texts:
return []
model = self._get_model()
vectors = model.encode(texts, normalize_embeddings=True)
return [[float(x) for x in vec] for vec in vectors]
class RemoteEmbeddingProvider(EmbeddingProvider):
"""OpenAI-compatible ``/embeddings`` endpoint (API key based)."""
name = "remote"
def __init__(
self,
base_url: str,
api_key: str,
model: str = "text-embedding-3-small",
dimension: int = EMBEDDING_DIM,
) -> None:
super().__init__(dimension)
self.base_url = base_url.rstrip("/")
self.api_key = api_key
self.model = model
def encode(self, texts: list[str]) -> list[list[float]]:
if not texts:
return []
import httpx
response = httpx.post(
f"{self.base_url}/embeddings",
headers={"Authorization": f"Bearer {self.api_key}"},
json={"model": self.model, "input": texts},
timeout=30.0,
)
response.raise_for_status()
payload = response.json()
data = sorted(payload.get("data", []), key=lambda item: item.get("index", 0))
return [self._normalize([float(x) for x in item["embedding"]]) for item in data]
@staticmethod
def _normalize(vec: list[float]) -> list[float]:
norm = math.sqrt(sum(v * v for v in vec))
if norm > 0.0:
return [v / norm for v in vec]
return vec
_provider: EmbeddingProvider | None = None
_provider_lock = threading.Lock()
def _build_provider() -> EmbeddingProvider:
"""Select the best available provider (respecting ``OBSIGATE_EMBEDDING_PROVIDER``)."""
requested = os.getenv("OBSIGATE_EMBEDDING_PROVIDER", "auto").strip().lower()
if requested in ("auto", "local", "sentence-transformers") and SentenceTransformerProvider.is_available():
model = os.getenv("OBSIGATE_EMBEDDING_MODEL", "all-MiniLM-L6-v2")
return SentenceTransformerProvider(model)
remote_key = os.getenv("OBSIGATE_EMBEDDING_API_KEY", "")
remote_url = os.getenv("OBSIGATE_EMBEDDING_BASE_URL", "")
if requested in ("auto", "remote") and remote_key and remote_url:
model = os.getenv("OBSIGATE_EMBEDDING_MODEL", "text-embedding-3-small")
dim = int(os.getenv("OBSIGATE_EMBEDDING_DIM", str(EMBEDDING_DIM)))
return RemoteEmbeddingProvider(remote_url, remote_key, model, dim)
if requested == "remote":
logger.warning(
"OBSIGATE_EMBEDDING_PROVIDER=remote but OBSIGATE_EMBEDDING_API_KEY/BASE_URL missing; using hash fallback"
)
return HashEmbeddingProvider()
def get_embedding_provider() -> EmbeddingProvider:
"""Return the cached embedding provider (built on first use)."""
global _provider
with _provider_lock:
if _provider is None:
_provider = _build_provider()
logger.info("Semantic embedding provider: %s (dim=%d)", _provider.name, _provider.dimension)
return _provider
def reset_embedding_provider() -> None:
"""Forget the cached provider (used by tests and config reloads)."""
global _provider
with _provider_lock:
_provider = None
# ---------------------------------------------------------------------------
# Vector store
# ---------------------------------------------------------------------------
class VectorStore:
"""In-memory vector store with optional numpy / faiss acceleration.
Vectors are always kept as Python lists (source of truth). A numpy matrix
and/or a faiss ``IndexFlatIP`` are built lazily and invalidated on mutation.
All vectors are expected to be L2-normalized, so the inner product equals
the cosine similarity.
"""
def __init__(self, dimension: int = EMBEDDING_DIM) -> None:
self.dimension = dimension
self._keys: list[str] = []
self._chunks: list[str] = []
self._vectors: list[list[float]] = []
self._dirty = True
self._numpy: Any | None = None
self._numpy_checked = False
self._matrix: Any | None = None
self._faiss: Any | None = None
self._faiss_checked = False
self._faiss_index: Any | None = None
def __len__(self) -> int:
return len(self._vectors)
def clear(self) -> None:
self._keys = []
self._chunks = []
self._vectors = []
self._dirty = True
def add(self, key: str, chunk: str, vector: list[float]) -> None:
self._keys.append(key)
self._chunks.append(chunk)
self._vectors.append(vector)
self._dirty = True
def remove_document(self, key: str) -> None:
"""Remove every chunk belonging to *key*."""
kept = [(k, c, v) for k, c, v in zip(self._keys, self._chunks, self._vectors) if k != key]
if len(kept) == len(self._vectors):
return
self._keys = [k for k, _, _ in kept]
self._chunks = [c for _, c, _ in kept]
self._vectors = [v for _, _, v in kept]
self._dirty = True
# -- optional accelerators -------------------------------------------------
def _get_numpy(self) -> Any | None:
if not self._numpy_checked:
self._numpy_checked = True
try:
import numpy as np
self._numpy = np
except Exception:
self._numpy = None
return self._numpy
def _get_faiss(self) -> Any | None:
if not self._faiss_checked:
self._faiss_checked = True
try:
import faiss
self._faiss = faiss
except Exception:
self._faiss = None
return self._faiss
def _rebuild_accelerators(self) -> None:
self._dirty = False
np = self._get_numpy()
if np is None or not self._vectors:
self._matrix = None
self._faiss_index = None
return
self._matrix = np.asarray(self._vectors, dtype="float32")
faiss = self._get_faiss()
if faiss is not None:
index = faiss.IndexFlatIP(self.dimension)
index.add(self._matrix)
self._faiss_index = index
else:
self._faiss_index = None
def search(self, query_vector: list[float], top_k: int = DEFAULT_TOP_K) -> list[tuple[str, float]]:
"""Return ``(doc_key, cosine_similarity)`` pairs sorted by similarity."""
if not self._vectors:
return []
top_k = max(1, min(top_k, len(self._vectors)))
if self._dirty:
self._rebuild_accelerators()
np = self._get_numpy()
if np is not None and self._faiss_index is not None and self._matrix is not None:
query = np.asarray([query_vector], dtype="float32")
scores, indices = self._faiss_index.search(query, top_k)
return [
(self._keys[int(idx)], float(score))
for score, idx in zip(scores[0], indices[0])
if idx >= 0
]
if np is not None and self._matrix is not None:
query = np.asarray(query_vector, dtype="float32")
scores = self._matrix @ query
order = np.argsort(scores)[::-1][:top_k]
return [(self._keys[int(i)], float(scores[int(i)])) for i in order]
scored = [(self._keys[i], _dot(self._vectors[i], query_vector)) for i in range(len(self._vectors))]
scored.sort(key=lambda item: item[1], reverse=True)
return scored[:top_k]
def chunk_of(self, index: int) -> str:
"""Return the stored chunk text at *index* (used by diagnostics/tests)."""
return self._chunks[index]
def _dot(a: list[float], b: list[float]) -> float:
"""Dot product for two equal-length vectors."""
return sum(x * y for x, y in zip(a, b))
# ---------------------------------------------------------------------------
# Reciprocal Rank Fusion
# ---------------------------------------------------------------------------
def rrf_fuse(rankings: list[list[str]], k: int = RRF_K) -> dict[str, float]:
"""Fuse several ranked key lists into a single score map.
``score(key) = Σ_rankings 1 / (k + rank(key))`` where ``rank`` is
1-based. Documents ranked highly by several methods rise to the top.
Args:
rankings: Ordered lists of document keys (best first).
k: RRF smoothing constant.
Returns:
Mapping ``doc_key -> fused score`` (insertion order is unspecified).
"""
scores: dict[str, float] = {}
for ranking in rankings:
seen: set[str] = set()
for rank, key in enumerate(ranking):
if key in seen:
continue
seen.add(key)
scores[key] = scores.get(key, 0.0) + 1.0 / (k + rank + 1)
return scores
# ---------------------------------------------------------------------------
# Semantic index
# ---------------------------------------------------------------------------
class SemanticIndex:
"""Holds document chunk embeddings and answers similarity queries."""
def __init__(self, provider: EmbeddingProvider | None = None) -> None:
self.provider = provider
self.store = VectorStore(provider.dimension if provider else EMBEDDING_DIM)
self.doc_keys: set[str] = set()
self._ready = False
self._lock = threading.Lock()
def is_ready(self) -> bool:
"""Return True once a full rebuild has completed."""
return self._ready
def is_stale(self) -> bool:
"""Alias used by callers that check index freshness."""
return not self._ready
def _ensure_provider(self) -> EmbeddingProvider:
if self.provider is None:
self.provider = get_embedding_provider()
self.store = VectorStore(self.provider.dimension)
return self.provider
@staticmethod
def _document_text(file_info: dict[str, Any]) -> str:
title = file_info.get("title", "") or ""
content = file_info.get("content", "") or ""
return (title + "\n\n" + content).strip()
def _embed_document(self, doc_key: str, file_info: dict[str, Any]) -> None:
text = self._document_text(file_info)
if not text:
return
provider = self._ensure_provider()
chunks = chunk_text(text)
if not chunks:
return
vectors = provider.encode(chunks)
for chunk, vector in zip(chunks, vectors):
self.store.add(doc_key, chunk, vector)
self.doc_keys.add(doc_key)
def rebuild(self) -> None:
"""Rebuild the whole index from the global in-memory index."""
from backend.indexer import index
provider = self._ensure_provider()
with self._lock:
self.store = VectorStore(provider.dimension)
self.doc_keys = set()
for vault_name, vault_data in index.items():
for file_info in vault_data.get("files", []):
doc_key = f"{vault_name}::{file_info.get('path', '')}"
try:
self._embed_document(doc_key, file_info)
except Exception as exc:
logger.warning("Semantic embedding failed for %s: %s", doc_key, exc)
self._ready = True
logger.info(
"Semantic index built: %d documents, %d chunks (provider=%s)",
len(self.doc_keys),
len(self.store),
provider.name,
)
def add_document(self, vault_name: str, path: str, file_info: dict[str, Any]) -> None:
"""Add or refresh a single document (no-op until the index is ready)."""
if not self._ready or not file_info:
return
doc_key = f"{vault_name}::{path}"
with self._lock:
self.store.remove_document(doc_key)
self.doc_keys.discard(doc_key)
try:
self._embed_document(doc_key, file_info)
except Exception as exc:
logger.warning("Semantic embedding failed for %s: %s", doc_key, exc)
def remove_document(self, vault_name: str, path: str) -> None:
"""Remove a single document (no-op until the index is ready)."""
if not self._ready:
return
doc_key = f"{vault_name}::{path}"
with self._lock:
self.store.remove_document(doc_key)
self.doc_keys.discard(doc_key)
def search(
self,
query: str,
vault_filter: str = "all",
top_k: int = DEFAULT_TOP_K,
) -> list[tuple[str, float]]:
"""Return ``(doc_key, best_chunk_similarity)`` pairs, best first."""
if not self._ready or not query or not query.strip():
return []
provider = self._ensure_provider()
query_vector = provider.encode_one(query[:MAX_QUERY_CHARS])
hits = self.store.search(query_vector, top_k=max(top_k * 4, top_k))
best: dict[str, float] = {}
for doc_key, score in hits:
if vault_filter != "all" and not doc_key.startswith(vault_filter + "::"):
continue
if doc_key not in best or score > best[doc_key]:
best[doc_key] = score
ordered = sorted(best.items(), key=lambda item: item[1], reverse=True)
return ordered[:top_k]
_semantic_index: SemanticIndex | None = None
_index_lock = threading.Lock()
def get_semantic_index() -> SemanticIndex:
"""Return the process-wide semantic index (created on first access)."""
global _semantic_index
with _index_lock:
if _semantic_index is None:
_semantic_index = SemanticIndex()
return _semantic_index
def reset_semantic_index() -> None:
"""Drop the singleton index (tests)."""
global _semantic_index
with _index_lock:
_semantic_index = None
def init_semantic_index() -> None:
"""Force a full semantic index build. Called after ``build_index`` on startup."""
from backend.indexer import index
if any(vdata.get("files") for vdata in index.values()):
get_semantic_index().rebuild()
def on_index_change(action: str, vault_name: str, path: str, file_info: dict[str, Any]) -> None:
"""Incremental hook registered with the indexer change notifier."""
index_obj = get_semantic_index()
if action == "add" and file_info:
index_obj.add_document(vault_name, path, file_info)
elif action == "remove":
index_obj.remove_document(vault_name, path)
def semantic_search_docs(
query: str,
vault_filter: str = "all",
top_k: int = DEFAULT_TOP_K,
) -> list[tuple[str, float]]:
"""Convenience wrapper around :meth:`SemanticIndex.search`."""
return get_semantic_index().search(query, vault_filter=vault_filter, top_k=top_k)
def semantic_status() -> dict[str, Any]:
"""Return provider/index diagnostics for the API and the UI."""
index_obj = get_semantic_index()
provider = index_obj.provider or get_embedding_provider()
return {
"available": index_obj.is_ready(),
"provider": provider.name,
"dimension": provider.dimension,
"documents": len(index_obj.doc_keys),
"chunks": len(index_obj.store),
}
backend/services/search.py
+5 -2
View File
@@ -56,11 +56,13 @@ def advanced_search_vaults(
created: str | None = None,
modified: str | None = None,
size: str | None = None,
semantic: bool = False,
) -> dict[str, Any]:
"""Advanced full-text search (TF-IDF, facets, operators).
No permission filtering is applied: callers that need it (the tool layer)
filter the ``results`` list themselves.
When ``semantic`` is True, the TF-IDF ranking is fused with the embedding
(semantic) ranking via RRF. No permission filtering is applied: callers
that need it (the tool layer) filter the ``results`` list themselves.
"""
from backend.search import advanced_search
@@ -79,6 +81,7 @@ def advanced_search_vaults(
created=created,
modified=modified,
size=size,
semantic=semantic,
)