docsif 0.1.0__py3-none-any.whl

sif/models/embedding.py

@@ -0,0 +1,71 @@
+"""Pydantic models for Embedding operations."""
+
+from enum import Enum
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class ModelType(str, Enum):
+    """Type of embedding model."""
+
+    GGUF = "gguf"
+    SENTENCE_TRANSFORMERS = "sentence_transformers"
+    OPENAI = "openai"
+    HUGGINGFACE = "huggingface"
+    MODELSCOPE = "modelscope"
+
+
+class EmbeddingConfig(BaseModel):
+    """Configuration for embedding models."""
+
+    model_type: ModelType = Field(ModelType.MODELSCOPE, description="Model type")
+    model_path: str | None = Field(None, description="Path to model file")
+    model_name: str = Field("Qwen/Qwen3-Embedding-0.6B", description="Model name or identifier")
+
+    # Model parameters
+    embedding_dim: int = Field(1024, ge=1, description="Embedding dimension")
+    max_tokens: int = Field(512, ge=1, description="Maximum tokens per input")
+    batch_size: int = Field(32, ge=1, description="Batch size for inference")
+
+    # GGUF specific
+    n_gpu_layers: int = Field(0, ge=0, description="Number of GPU layers")
+    n_ctx: int = Field(2048, ge=512, description="Context size")
+
+    # API keys (for remote models)
+    api_key: str | None = Field(None, exclude=True)
+    api_base: str | None = None
+
+    # Caching
+    cache_embeddings: bool = Field(True, description="Cache embeddings")
+    cache_dir: str | None = None
+
+
+class EmbeddingModelInfo(BaseModel):
+    """Information about an embedding model."""
+
+    model_config = ConfigDict(from_attributes=True)
+
+    id: str = Field(..., description="Model ID")
+    name: str = Field(..., description="Model name")
+    model_type: ModelType = Field(...)
+    embedding_dim: int = Field(..., ge=1)
+    max_tokens: int = Field(..., ge=1)
+    loaded: bool = Field(False, description="Whether model is loaded")
+    device: str | None = None
+
+
+class EmbeddingRequest(BaseModel):
+    """Request for generating embeddings."""
+
+    texts: list[str] = Field(..., min_length=1, description="Texts to embed")
+    normalize: bool = Field(True, description="Normalize embeddings")
+
+
+class EmbeddingResponse(BaseModel):
+    """Response with generated embeddings."""
+
+    embeddings: list[list[float]] = Field(..., description="Generated embeddings")
+    model_id: str = Field(..., description="Model used")
+    dimensions: int = Field(..., ge=1)
+    total_tokens: int = Field(..., ge=0)
+    processing_time_ms: float = Field(..., ge=0)

sif/models/search.py

@@ -0,0 +1,92 @@
+"""Pydantic models for Search operations."""
+
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class SearchType(str, Enum):
+    """Type of search to perform."""
+
+    BM25 = "bm25"
+    VECTOR = "vector"
+    HYBRID = "hybrid"
+
+
+class SearchOptions(BaseModel):
+    """Options for search queries."""
+
+    limit: int = Field(10, ge=1, le=100, description="Maximum results")
+    offset: int = Field(0, ge=0, description="Result offset")
+    threshold: float = Field(0.0, ge=0, le=1, description="Minimum score threshold")
+    include_chunks: bool = Field(True, description="Include matched chunks")
+    include_metadata: bool = Field(True, description="Include document metadata")
+    highlight_matches: bool = Field(True, description="Highlight matching terms")
+
+    # BM25 options
+    bm25_k1: float = Field(1.5, ge=0.1, le=3.0)
+    bm25_b: float = Field(0.75, ge=0.1, le=1.0)
+
+    # Vector options
+    vector_weight: float = Field(0.7, ge=0, le=1)
+
+    # Hybrid options
+    bm25_weight: float = Field(0.3, ge=0, le=1)
+    rrf_k: int = Field(60, ge=1, description="RRF fusion parameter")
+
+    # Query expansion
+    expand_query: bool = Field(False, description="Enable query expansion")
+    expansion_factor: int = Field(3, ge=1, le=10)
+
+    # Reranking
+    rerank: bool = Field(False, description="Enable reranking")
+    rerank_top_k: int = Field(20, ge=1, le=100)
+
+
+class SearchQuery(BaseModel):
+    """Model for search query."""
+
+    query: str = Field(..., min_length=1, description="Search query")
+    collection_ids: list[str] | None = Field(None, description="Filter by collections")
+    search_type: SearchType = Field(SearchType.HYBRID, description="Search type")
+    options: SearchOptions = Field(default_factory=SearchOptions)
+    filters: dict[str, Any] = Field(default_factory=dict, description="Additional filters")
+
+
+class SearchResult(BaseModel):
+    """Model for individual search result."""
+
+    model_config = ConfigDict(from_attributes=True)
+
+    document_id: str = Field(..., description="Document ID")
+    document_path: str = Field(..., description="Document path")
+    document_title: str | None = None
+    score: float = Field(..., ge=0, le=1)
+
+    # Component scores (for hybrid search)
+    bm25_score: float | None = None
+    vector_score: float | None = None
+
+    # Result data
+    content_preview: str | None = None
+    matched_chunks: list[dict[str, Any]] = Field(default_factory=list)
+    highlights: list[str] = Field(default_factory=list)
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    context_description: str | None = None
+
+
+class SearchResponse(BaseModel):
+    """Model for search response."""
+
+    query: str = Field(..., description="Original query")
+    search_type: SearchType = Field(...)
+    total: int = Field(..., ge=0, description="Total matching documents")
+    results: list[SearchResult] = Field(default_factory=list)
+
+    # Performance metrics
+    search_time_ms: float = Field(..., ge=0)
+
+    # Query expansion info
+    expanded_query: str | None = None
+    expansion_terms: list[str] = Field(default_factory=list)

sif/search/__init__.py

@@ -0,0 +1,22 @@
+"""Search functionality for SIF.
+
+This package implements various search strategies:
+- BM25 full-text search via SQLite FTS5
+- Vector semantic search via sqlite-vec
+- Hybrid search combining multiple strategies
+- RRF (Reciprocal Rank Fusion) for result merging
+"""
+
+from sif.search.bm25 import BM25Searcher
+from sif.search.hybrid import HybridSearcher, SearchPipeline
+from sif.search.rrf import RRFFusion
+from sif.search.vector import VectorSearcher
+
+
+__all__ = [
+    "BM25Searcher",
+    "HybridSearcher",
+    "RRFFusion",
+    "SearchPipeline",
+    "VectorSearcher",
+]

sif/search/benchmark.py

@@ -0,0 +1,110 @@
+"""Benchmark evaluation for search quality."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Callable
+
+from sif.core.models import SearchResult
+from sif.utils.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+
+def precision_at_k(relevance: Sequence[int], k: int) -> float:
+    """Precision@K: relevant items in top-K / K."""
+    if k == 0:
+        return 0.0
+    return sum(relevance[:k]) / k
+
+
+def recall_at_k(relevance: Sequence[int], total_relevant: int, k: int) -> float:
+    """Recall@K: relevant items retrieved / total relevant."""
+    if total_relevant == 0:
+        return 0.0
+    return sum(relevance[:k]) / total_relevant
+
+
+def reciprocal_rank(relevance: Sequence[int]) -> float:
+    """Reciprocal rank: 1 / rank of first relevant item."""
+    for i, rel in enumerate(relevance, start=1):
+        if rel:
+            return 1.0 / i
+    return 0.0
+
+
+def mean_reciprocal_rank(all_relevance: list[Sequence[int]]) -> float:
+    """Mean Reciprocal Rank across queries."""
+    rr_scores = [reciprocal_rank(r) for r in all_relevance]
+    return sum(rr_scores) / len(rr_scores) if rr_scores else 0.0
+
+
+class SearchEvaluator:
+    """Evaluate search quality against a benchmark fixture."""
+
+    def __init__(self, fixture: dict) -> None:
+        """Initialize with fixture data.
+
+        Args:
+            fixture: Dictionary with keys:
+                - queries: list of {query: str, relevant_docids: list[str], collections?: list[str]}
+        """
+        self.fixture = fixture
+        self._validate_fixture()
+
+    def _validate_fixture(self) -> None:
+        """Validate fixture format."""
+        if "queries" not in self.fixture:
+            raise ValueError("Fixture must contain 'queries' key")
+        for i, item in enumerate(self.fixture["queries"]):
+            if "query" not in item:
+                raise ValueError(f"Query {i} missing 'query' field")
+            if "relevant_docids" not in item:
+                raise ValueError(f"Query {i} missing 'relevant_docids' field")
+
+    def evaluate(
+        self,
+        search_fn: Callable[[str], list[SearchResult]],
+        k_values: list[int] | None = None,
+    ) -> dict[str, float]:
+        """Evaluate search quality.
+
+        Args:
+            search_fn: Function that takes a query string and returns SearchResult list
+            k_values: List of k values for precision/recall@k. Default: [1, 5, 10]
+
+        Returns:
+            Dictionary of averaged metrics
+        """
+        k_values = k_values or [1, 5, 10]
+        metrics: dict[str, list[float]] = {}
+        all_relevance: list[Sequence[int]] = []
+
+        for query_item in self.fixture["queries"]:
+            query = query_item["query"]
+            relevant_ids = set(query_item["relevant_docids"])
+
+            # Run search
+            results = search_fn(query)
+            result_ids = [r.document_id for r in results]
+
+            # Build relevance vector
+            relevance = [1 if rid in relevant_ids else 0 for rid in result_ids]
+            all_relevance.append(relevance)
+
+            # Compute metrics for each k
+            for k in k_values:
+                metrics.setdefault(f"precision@{k}", []).append(precision_at_k(relevance, k))
+                metrics.setdefault(f"recall@{k}", []).append(
+                    recall_at_k(relevance, len(relevant_ids), k)
+                )
+
+        # Average across queries
+        averaged: dict[str, float] = {}
+        for key, values in metrics.items():
+            averaged[key] = sum(values) / len(values) if values else 0.0
+        averaged["mrr"] = mean_reciprocal_rank(all_relevance)
+        averaged["num_queries"] = float(len(self.fixture["queries"]))
+
+        return averaged

sif/search/bm25.py

@@ -0,0 +1,216 @@
+"""BM25 full-text search implementation."""
+
+from __future__ import annotations
+
+import os
+import sqlite3
+
+from sif.core.models import SearchOptions, SearchResult
+
+
+class BM25Searcher:
+    """BM25 full-text search using SQLite FTS5."""
+
+    def __init__(self, db: sqlite3.Connection) -> None:
+        self.db = db
+
+    def search(self, query: str, options: SearchOptions | None = None) -> list[SearchResult]:
+        """Search documents using BM25."""
+        if options is None:
+            options = SearchOptions()
+
+        # Build the FTS query
+        fts_query = self._build_fts_query(query)
+
+        # Build collection filter
+        collection_filter = ""
+        params = [fts_query]
+
+        if options.collection_ids:
+            placeholders = ", ".join(["?"] * len(options.collection_ids))
+            collection_filter = f"AND d.collection_id IN ({placeholders})"
+            params.extend(options.collection_ids)
+
+        # Execute search
+        sql = f"""
+            SELECT
+                d.id as document_id,
+                d.title,
+                d.path,
+                c.name as collection_name,
+                rank as score
+            FROM documents_fts fts
+            JOIN documents d ON fts.rowid = d.rowid
+            JOIN collections c ON d.collection_id = c.id
+            WHERE documents_fts MATCH ? {collection_filter}
+            ORDER BY rank
+            LIMIT ? OFFSET ?
+        """
+        params.extend([options.limit, options.offset])
+
+        cursor = self.db.execute(sql, params)
+        results = []
+
+        for rank, row in enumerate(cursor.fetchall(), 1):
+            # Convert rank to score (lower rank = higher score)
+            score = 1.0 / (1.0 + abs(row["score"]))
+
+            if score < options.min_score:
+                continue
+
+            result = SearchResult(
+                document_id=row["document_id"],
+                title=row["title"] or "",
+                path=row["path"],
+                collection_name=row["collection_name"],
+                score=score,
+                rank=rank,
+            )
+
+            if options.include_content:
+                result.content = self._get_document_content(row["document_id"])
+
+            if options.include_highlights:
+                result.highlights = self._get_highlights(
+                    row["document_id"], query, options.max_highlights
+                )
+
+            results.append(result)
+
+        return self._attach_contexts(results)
+
+    def _attach_contexts(self, results: list[SearchResult]) -> list[SearchResult]:
+        """Attach path context descriptions to search results via batch query."""
+        if not results:
+            return results
+        paths = list({r.path for r in results})
+        placeholders = ", ".join(["?"] * len(paths))
+        sql = f"""
+            SELECT target_id, content FROM contexts
+            WHERE context_type = 'path' AND target_id IN ({placeholders})
+        """
+        cursor = self.db.execute(sql, paths)
+        # Normalize keys for cross-platform matching (macOS /private/tmp, etc.)
+        context_map = {
+            os.path.realpath(row["target_id"]): row["content"] for row in cursor.fetchall()
+        }
+        for result in results:
+            result.context_description = context_map.get(os.path.realpath(result.path))
+        return results
+
+    def search_chunks(
+        self, query: str, options: SearchOptions | None = None
+    ) -> list[tuple[str, str, float]]:
+        """Search document chunks using BM25.
+
+        Returns list of (document_id, chunk_content, score) tuples.
+        """
+        if options is None:
+            options = SearchOptions()
+
+        fts_query = self._build_fts_query(query)
+
+        sql = """
+            SELECT
+                dc.document_id,
+                dc.content,
+                rank as score
+            FROM chunks_fts fts
+            JOIN document_chunks dc ON fts.rowid = dc.rowid
+            WHERE chunks_fts MATCH ?
+            ORDER BY rank
+            LIMIT ? OFFSET ?
+        """
+
+        cursor = self.db.execute(sql, [fts_query, options.limit, options.offset])
+        results = []
+
+        for row in cursor.fetchall():
+            score = 1.0 / (1.0 + abs(row["score"]))
+            if score >= options.min_score:
+                results.append((row["document_id"], row["content"], score))
+
+        return results
+
+    def _build_fts_query(self, query: str) -> str:
+        """Build FTS5 query from user query.
+
+        Handles:
+        - Multiple terms (AND search)
+        - Phrases (quoted strings)
+        - Prefix search (term*)
+        """
+        # Simple implementation: just join terms with AND
+        terms = query.split()
+        if not terms:
+            return "*"
+
+        # Use NEAR for better relevance
+        if len(terms) == 1:
+            return f"{terms[0]}*"
+
+        # For multiple terms, use AND
+        return " AND ".join(f"{term}*" for term in terms)
+
+    def _get_document_content(self, document_id: str) -> str | None:
+        """Get document content."""
+        cursor = self.db.execute("SELECT content FROM documents WHERE id = ?", (document_id,))
+        row = cursor.fetchone()
+        return row[0] if row else None
+
+    def _get_highlights(self, document_id: str, query: str, max_highlights: int = 3) -> list[str]:
+        """Get highlighted snippets for a document."""
+        # Get chunks for the document
+        cursor = self.db.execute(
+            """
+            SELECT content FROM document_chunks
+            WHERE document_id = ?
+            ORDER BY sequence
+            """,
+            (document_id,),
+        )
+
+        chunks = [row[0] for row in cursor.fetchall()]
+        if not chunks:
+            # Get full document content
+            content = self._get_document_content(document_id)
+            if content:
+                chunks = [content[:1000]]  # First 1000 chars
+
+        # Find chunks containing query terms
+        query_terms = [t.lower() for t in query.split()]
+        highlights = []
+
+        for chunk in chunks:
+            chunk_lower = chunk.lower()
+            if any(term in chunk_lower for term in query_terms):
+                # Extract snippet around the match
+                snippet = self._extract_snippet(chunk, query_terms)
+                if snippet:
+                    highlights.append(snippet)
+                    if len(highlights) >= max_highlights:
+                        break
+
+        return highlights
+
+    def _extract_snippet(self, text: str, query_terms: list[str], context: int = 50) -> str:
+        """Extract a snippet around query matches."""
+        text_lower = text.lower()
+
+        for term in query_terms:
+            pos = text_lower.find(term)
+            if pos >= 0:
+                start = max(0, pos - context)
+                end = min(len(text), pos + len(term) + context)
+                snippet = text[start:end]
+
+                # Add ellipsis if truncated
+                if start > 0:
+                    snippet = "..." + snippet
+                if end < len(text):
+                    snippet = snippet + "..."
+
+                return snippet.strip()
+
+        # Return first part if no match found
+        return text[:200].strip() + "..." if len(text) > 200 else text.strip()

sif/search/expansion.py

@@ -0,0 +1,157 @@
+"""Query expansion for improving search recall."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from sif.utils.logging import get_logger
+
+
+logger = get_logger(__name__)
+
+if TYPE_CHECKING:
+    from sif.embedding.manager import EmbeddingManager
+
+
+class QueryExpansion:
+    """Query expansion to improve search recall.
+
+    Query expansion adds related terms to the original query
+    to help find more relevant documents.
+
+    Methods:
+    - Synonym expansion: Add synonyms for query terms
+    - Embedding expansion: Find similar terms using embeddings
+    - LLM expansion: Use language model to generate related terms
+    """
+
+    def __init__(
+        self,
+        embedding_manager: EmbeddingManager | None = None,
+        expansion_factor: int = 3,
+    ) -> None:
+        """Initialize query expansion.
+
+        Args:
+            embedding_manager: Optional embedding manager for semantic expansion
+            expansion_factor: Number of expansion terms to add per query term
+        """
+        self._embedding_manager = embedding_manager
+        self._expansion_factor = expansion_factor
+
+    def expand(self, query: str, intent: str | None = None) -> list[str]:
+        """Expand a query with related terms.
+
+        Args:
+            query: Original query
+            intent: Optional intent hint to prepend to each variant
+
+        Returns:
+            List of query variants: [original_query, expanded_variant_1, ...]
+        """
+        expansion_terms = self._get_expansion_terms(query)
+
+        variants = [query]
+        variants.extend(f"{query} {term}" for term in expansion_terms[: self._expansion_factor])
+
+        if intent:
+            variants = [f"{intent}: {v}" for v in variants]
+
+        if not variants:
+            return [query]
+
+        logger.debug(f"Expanded query: '{query}' -> {variants}")
+        return variants
+
+    def _get_expansion_terms(self, query: str) -> list[str]:
+        """Get expansion terms using embedding-based pseudo-relevance feedback.
+
+        Args:
+            query: Original query
+
+        Returns:
+            List of expansion terms sorted by similarity
+        """
+        if self._embedding_manager is None:
+            return []
+
+        try:
+            # Embed the query
+            query_embedding = self._embedding_manager.embed_single(query)
+
+            # Generate candidate terms: individual words from query + common synonyms
+            words = query.split()
+            candidates = set(words)
+
+            # Simple synonym map for common terms (expandable)
+            synonym_map = {
+                "python": ["py", "cpython"],
+                "search": ["find", "lookup", "query"],
+                "index": ["catalog", "register"],
+                "document": ["doc", "file", "page"],
+                "config": ["configuration", "settings"],
+            }
+
+            for word in words:
+                word_lower = word.lower()
+                if word_lower in synonym_map:
+                    candidates.update(synonym_map[word_lower])
+
+            # Remove query words from candidates (we only want expansion terms)
+            query_words_lower = {w.lower() for w in words}
+            candidates = candidates - query_words_lower
+
+            if not candidates:
+                return []
+
+            candidate_list = list(candidates)
+            candidate_embeddings = self._embedding_manager.embed(candidate_list)
+
+            # Compute cosine similarities
+            query_vec = np.array(query_embedding)
+            similarities = []
+            for emb in candidate_embeddings.embeddings:
+                cand_vec = np.array(emb)
+                sim = np.dot(query_vec, cand_vec) / (
+                    np.linalg.norm(query_vec) * np.linalg.norm(cand_vec) + 1e-8
+                )
+                similarities.append(sim)
+
+            # Return top-k most similar terms that are NOT already in the query
+            scored = [
+                (term, sim)
+                for term, sim in zip(candidate_list, similarities)
+                if term.lower() not in query_words_lower
+            ]
+            scored.sort(key=lambda x: x[1], reverse=True)
+
+            return [term for term, _ in scored[: self._expansion_factor]]
+        except Exception as e:
+            logger.warning(f"Query expansion failed: {e}")
+            return []
+
+    def expand_batch(self, queries: list[str]) -> list[str]:
+        """Expand multiple queries and return a flat list of unique variants.
+
+        Args:
+            queries: List of queries
+
+        Returns:
+            Flat list of unique expanded query variants
+        """
+        all_variants: list[str] = []
+        for query in queries:
+            variants = self.expand(query)
+            all_variants.extend(variants)
+
+        # Deduplicate while preserving order
+        seen: set[str] = set()
+        unique_variants: list[str] = []
+        for v in all_variants:
+            if v not in seen:
+                seen.add(v)
+                unique_variants.append(v)
+
+        return unique_variants