lm-deluge 0.0.67__py3-none-any.whl → 0.0.88__py3-none-any.whl

lm_deluge/tool/prefab/full_text_search/__init__.py

@@ -0,0 +1,285 @@
+"""Full text search prefab tool using Tantivy."""
+
+import json
+import tempfile
+from pathlib import Path
+from typing import Annotated, Any
+
+from lm_deluge.tool import Tool
+
+from .tantivy_index import SearchResult, TantivySearch
+
+
+class FullTextSearchManager:
+    """
+    Full-text search tools using Tantivy.
+
+    Provides two tools:
+    - search: Search the corpus and get document IDs + previews
+    - fetch: Get the full contents of specific documents by ID
+
+    Args:
+        corpus: List of document dicts to index. Each dict must have an "id" field.
+        search_fields: List of field names to search. If None, searches all fields.
+        preview_fields: Fields to include in search result previews.
+        index_path: Path to store the Tantivy index. If None, uses a temp directory.
+        search_tool_name: Name for the search tool (default: "search")
+        fetch_tool_name: Name for the fetch tool (default: "fetch")
+        max_results: Maximum number of search results to return (default: 10)
+        include_fields: Fields to include in the index (searchable). If None, includes all.
+        exclude_fields: Fields to exclude from the index (not searchable).
+
+    Example:
+        ```python
+        corpus = [
+            {"id": "1", "title": "Hello World", "content": "This is a test document."},
+            {"id": "2", "title": "Another Doc", "content": "More content here."},
+        ]
+        manager = FullTextSearchManager(
+            corpus=corpus,
+            search_fields=["title", "content"],
+            preview_fields=["title"],
+        )
+        tools = manager.get_tools()
+        ```
+    """
+
+    def __init__(
+        self,
+        corpus: list[dict[str, Any]],
+        *,
+        search_fields: list[str] | None = None,
+        preview_fields: list[str] | None = None,
+        index_path: str | Path | None = None,
+        search_tool_name: str = "search",
+        fetch_tool_name: str = "fetch",
+        max_results: int = 10,
+        include_fields: list[str] | None = None,
+        exclude_fields: list[str] | None = None,
+        deduplicate_by: str | None = None,
+    ):
+        # Initialize _temp_dir early to avoid __del__ issues
+        self._temp_dir: str | None = None
+
+        self.corpus = corpus
+        self.search_fields = search_fields
+        self.preview_fields = preview_fields
+        self.search_tool_name = search_tool_name
+        self.fetch_tool_name = fetch_tool_name
+        self.max_results = max_results
+        self._tools: list[Tool] | None = None
+
+        # Validate corpus
+        if not corpus:
+            raise ValueError("Corpus cannot be empty")
+
+        # Ensure all documents have an id field
+        for i, doc in enumerate(corpus):
+            if "id" not in doc:
+                raise ValueError(f"Document at index {i} is missing 'id' field")
+
+        # Set up index path
+        if index_path is None:
+            self._temp_dir = tempfile.mkdtemp(prefix="tantivy_")
+            self._index_path = Path(self._temp_dir)
+        else:
+            self._temp_dir = None
+            self._index_path = Path(index_path)
+
+        # Determine search fields from corpus if not provided
+        if search_fields is None:
+            # Use all string fields except 'id'
+            sample = corpus[0]
+            self.search_fields = [
+                k for k, v in sample.items() if k != "id" and isinstance(v, str)
+            ]
+        else:
+            self.search_fields = search_fields
+
+        # Initialize Tantivy index
+        self._index = TantivySearch(
+            index_path=str(self._index_path),
+            include_fields=include_fields,
+            exclude_fields=exclude_fields,
+        )
+
+        # Build the index
+        self._index.build_index(
+            records=corpus,
+            deduplicate_by=deduplicate_by,
+        )
+
+        # Cache documents for efficient fetch
+        self._index.cache_documents_for_fetch(corpus, id_column="id")
+
+    def _format_preview(self, result: SearchResult) -> dict[str, Any]:
+        """Format a search result for preview."""
+        preview: dict[str, Any] = {
+            "id": result.id,
+            "score": round(result.score, 4),
+        }
+
+        # Add preview fields
+        if self.preview_fields:
+            for field in self.preview_fields:
+                if field in result.content:
+                    value = result.content[field]
+                    # Truncate long values
+                    if isinstance(value, str) and len(value) > 200:
+                        value = value[:200] + "..."
+                    preview[field] = value
+        else:
+            # Include all fields with truncation
+            for field, value in result.content.items():
+                if field == "id":
+                    continue
+                if isinstance(value, str) and len(value) > 200:
+                    value = value[:200] + "..."
+                preview[field] = value
+
+        return preview
+
+    def _search(
+        self,
+        query: Annotated[str, "Search query to find relevant documents"],
+        limit: Annotated[int, "Maximum number of results to return"] = 10,
+    ) -> str:
+        """
+        Search the corpus for documents matching the query.
+
+        Returns a list of document previews with IDs and scores.
+        Use the fetch tool to get full document contents.
+        """
+        try:
+            # Use the search fields
+            assert self.search_fields is not None
+            results = self._index.search(
+                queries=[query],
+                fields=self.search_fields,
+                limit=min(limit, self.max_results),
+                escape=True,
+            )
+
+            previews = [self._format_preview(r) for r in results]
+
+            return json.dumps(
+                {
+                    "status": "success",
+                    "query": query,
+                    "num_results": len(previews),
+                    "results": previews,
+                },
+                indent=2,
+            )
+
+        except Exception as e:
+            return json.dumps({"status": "error", "error": str(e)})
+
+    def _fetch(
+        self,
+        document_ids: Annotated[
+            list[str], "List of document IDs to fetch full contents for"
+        ],
+    ) -> str:
+        """
+        Fetch the full contents of documents by their IDs.
+
+        Use search first to find relevant document IDs.
+        """
+        try:
+            documents, found_ids, missing_ids = self._index.fetch(document_ids)
+
+            result: dict[str, Any] = {
+                "status": "success",
+                "found": len(found_ids),
+                "missing": len(missing_ids),
+                "documents": documents,
+            }
+
+            if missing_ids:
+                result["missing_ids"] = missing_ids
+
+            return json.dumps(result, indent=2)
+
+        except Exception as e:
+            return json.dumps({"status": "error", "error": str(e)})
+
+    def get_tools(self) -> list[Tool]:
+        """Return the search and fetch tools."""
+        if self._tools is not None:
+            return self._tools
+
+        search_tool = Tool.from_function(self._search, name=self.search_tool_name)
+        fetch_tool = Tool.from_function(self._fetch, name=self.fetch_tool_name)
+
+        # Update descriptions for clarity
+        search_tool = search_tool.model_copy(
+            update={
+                "description": (
+                    "Search the document corpus for relevant results. "
+                    "Returns document IDs, relevance scores, and previews. "
+                    "Use the fetch tool to get full document contents."
+                )
+            }
+        )
+
+        fetch_tool = fetch_tool.model_copy(
+            update={
+                "description": (
+                    "Fetch the full contents of documents by their IDs. "
+                    "Use after searching to get complete document text."
+                )
+            }
+        )
+
+        self._tools = [search_tool, fetch_tool]
+        return self._tools
+
+    def search(
+        self, query: str, limit: int = 10, fields: list[str] | None = None
+    ) -> list[SearchResult]:
+        """
+        Direct search method for programmatic use.
+
+        Args:
+            query: Search query string
+            limit: Maximum number of results
+            fields: Fields to search (defaults to self.search_fields)
+
+        Returns:
+            List of SearchResult objects
+        """
+        search_fields = fields or self.search_fields
+        assert search_fields is not None
+        return self._index.search(
+            queries=[query],
+            fields=search_fields,
+            limit=min(limit, self.max_results),
+            escape=True,
+        )
+
+    def fetch(self, document_ids: list[str]) -> list[dict[str, Any]]:
+        """
+        Direct fetch method for programmatic use.
+
+        Args:
+            document_ids: List of document IDs to fetch
+
+        Returns:
+            List of document dicts
+        """
+        documents, _, _ = self._index.fetch(document_ids)
+        return documents
+
+    def __del__(self):
+        """Clean up temp directory if used."""
+        if self._temp_dir is not None:
+            import shutil
+
+            try:
+                shutil.rmtree(self._temp_dir, ignore_errors=True)
+            except Exception:
+                pass
+
+
+__all__ = ["FullTextSearchManager", "SearchResult"]

lm_deluge/tool/prefab/full_text_search/tantivy_index.py

@@ -0,0 +1,396 @@
+import os
+import re
+import shutil
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Iterable, Set, TypedDict
+
+from lenlp import normalizer
+from tantivy import Document, Index, SchemaBuilder
+from tqdm.auto import tqdm
+
+# Pattern to match field-like constructs (X:Y)
+FIELDLIKE = re.compile(r'(?<!\\)\b([^\s:"]+)\s*:\s*([^\s")]+)')
+
+
+class SearchResultContent(TypedDict):
+    title: str
+    url: str
+    description: str
+    keywords: str  # comma-delimited :)
+    content: str
+
+
+@dataclass
+class SearchResult:
+    id: str
+    score: float
+    content: SearchResultContent | dict[str, Any]  # for our purposes should be SRC
+
+
+class SearchIndex(ABC):
+    @abstractmethod
+    def build_index(
+        self,
+        records: list[dict],
+        deduplicate_by: str | None = None,
+    ):
+        raise NotImplementedError
+
+    @abstractmethod
+    def search(
+        self, queries: list[str], fields: list[str], limit: int = 8, escape: bool = True
+    ) -> list[SearchResult]:
+        raise NotImplementedError
+
+
+def deduplicate_exact(records: list[dict], deduplicate_by: str) -> list[dict]:
+    seen = set()
+    deduplicated = []
+    for record in tqdm(records):
+        key = normalizer.normalize(record[deduplicate_by])
+        if key not in seen:
+            deduplicated.append(record)
+            seen.add(key)
+    print(
+        f"Deduplication removed {len(records) - len(deduplicated)} of {len(records)} records."
+    )
+    return deduplicated
+
+
+def _fix_unbalanced_quotes(s: str) -> str:
+    """Keep balanced quotes; if an odd number of quotes, drop the strays.
+
+    Handles both double quotes (") and single quotes (') since Tantivy
+    treats both as phrase delimiters.
+    """
+    # Handle double quotes
+    if s.count('"') % 2 != 0:
+        s = s.replace('"', " ")
+
+    # Handle single quotes
+    if s.count("'") % 2 != 0:
+        s = s.replace("'", " ")
+
+    return s
+
+
+def _quote_nonfield_colons(
+    q: str, known_fields: Set[str], json_prefixes: Set[str]
+) -> str:
+    """Quote X:Y patterns when X is not a known field name.
+
+    This prevents strings like "3:12" from being interpreted as field queries.
+    Preserves legitimate field queries like "title:roof".
+    """
+
+    def repl(m: re.Match) -> str:
+        raw_prefix = m.group(1)
+        # Handle escaped dots (a\.b)
+        norm = raw_prefix.replace(r"\.", ".")
+        top = norm.split(".", 1)[0]
+
+        if norm in known_fields or top in json_prefixes:
+            return m.group(0)  # Genuine field query; keep as-is
+        return f'"{m.group(0)}"'  # Protect accidental field-like token
+
+    return FIELDLIKE.sub(repl, q)
+
+
+def _strip_syntax_to_words(q: str) -> str:
+    """Last-resort sanitizer: remove operators but keep words/numbers.
+
+    This is used as a fallback when all else fails to ensure we never crash.
+    """
+    q = re.sub(r'([+\-!(){}\[\]^"~*?:\\\/]|&&|\|\|)', " ", q)
+    q = re.sub(r"\s+", " ", q).strip()
+    return q or "*"  # Never return empty (match all if nothing left)
+
+
+def sanitize_query_for_tantivy(
+    query: str, known_fields: Iterable[str], json_field_prefixes: Iterable[str] = ()
+) -> str:
+    """Context-aware query sanitization for Tantivy.
+
+    Uses heuristics to detect whether special characters are intended as
+    Tantivy syntax or just regular text:
+    - Parentheses: only syntax if AND/OR/&&/|| present
+    - Square brackets: only syntax if TO or IN present
+    - Curly braces: only syntax if TO present
+
+    Always quotes X:Y patterns when X is not a known field, and fixes
+    unbalanced quotes.
+    """
+    # Step 1: Detect if special syntax is actually being used
+    has_boolean_ops = bool(re.search(r"\b(AND|OR)\b|&&|\|\|", query))
+    has_range_keyword = bool(re.search(r"\bTO\b", query))
+    has_set_keyword = bool(re.search(r"\bIN\b", query))
+
+    # Step 2: Remove syntax chars that are likely just regular text
+    if not has_boolean_ops:
+        # No boolean operators, so parentheses aren't for grouping
+        query = query.replace("(", " ").replace(")", " ")
+
+    if not has_range_keyword:
+        # No ranges, so curly braces aren't for exclusive bounds
+        query = query.replace("{", " ").replace("}", " ")
+
+    if not (has_range_keyword or has_set_keyword):
+        # No ranges or sets, so square brackets aren't query syntax
+        query = query.replace("[", " ").replace("]", " ")
+
+    # Step 3: Fix unbalanced quotes (both " and ')
+    query = _fix_unbalanced_quotes(query)
+
+    # Step 4: Quote non-field colons (e.g., "3:12")
+    query = _quote_nonfield_colons(query, set(known_fields), set(json_field_prefixes))
+
+    return query
+
+
+class TantivySearch(SearchIndex):
+    def __init__(
+        self,
+        index_path: str,
+        include_fields: list[str] | None = None,
+        exclude_fields: list[str] | None = None,
+    ):
+        self.index_path = Path(index_path)
+        self.schema = None
+        self.index = None
+        self.include_fields = include_fields
+        self.exclude_fields = exclude_fields
+        self.searchable_fields: set[str] | None = (
+            None  # Will be set during schema creation
+        )
+        self.id_to_record: dict[str, dict] = {}  # Cache for efficient document fetching
+
+    def create_schema(self, sample_record: dict):
+        schema_builder = SchemaBuilder()
+
+        # Determine which fields should be searchable
+        all_fields = set(sample_record.keys()) - {"id"}
+
+        if self.include_fields is not None and self.exclude_fields is not None:
+            # Both provided: searchable = include - exclude
+            self.searchable_fields = set(self.include_fields) - set(self.exclude_fields)
+        elif self.include_fields is not None:
+            # Only include provided: use only those
+            self.searchable_fields = set(self.include_fields)
+        elif self.exclude_fields is not None:
+            # Only exclude provided: use all except those
+            self.searchable_fields = all_fields - set(self.exclude_fields)
+        else:
+            # Neither provided: all fields searchable
+            # Set explicitly so we have field names for query sanitization
+            self.searchable_fields = all_fields
+
+        # Add all fields as text fields (they're all indexed, but we'll control search access)
+        for key in sample_record.keys():
+            if key == "id":
+                continue
+            schema_builder.add_text_field(key, stored=True, tokenizer_name="en_stem")
+        schema_builder.add_text_field(
+            "id", stored=True, tokenizer_name="raw", index_option="basic"
+        )
+        self.schema = schema_builder.build()
+
+    @classmethod
+    def from_index(cls, index_path: str):
+        index = cls(index_path)
+        index.load_index()
+        return index
+
+    def load_index(self):
+        if not self.index_path.exists():
+            raise FileNotFoundError(f"Index not found at {self.index_path}")
+        self.index = Index.open(str(self.index_path))
+
+    def _get_schema_field_names(self) -> set[str]:
+        """Get field names for query sanitization."""
+        # searchable_fields is always set after create_schema() is called
+        return self.searchable_fields if self.searchable_fields is not None else set()
+
+    def build_index(
+        self,
+        records: list[dict],
+        deduplicate_by: str | None = None,
+    ):
+        if not records:
+            raise ValueError("No records to index")
+
+        if not self.schema:
+            self.create_schema(records[0])
+        assert self.schema is not None, "Schema not created"
+
+        # Create index
+        if self.index_path.exists():
+            shutil.rmtree(self.index_path)
+        os.makedirs(self.index_path, exist_ok=True)
+        self.index = Index(self.schema, str(self.index_path))
+
+        # Deduplicate if requested
+        if deduplicate_by is not None:
+            records = deduplicate_exact(records, deduplicate_by)
+
+        # Index documents
+        writer = self.index.writer()
+        for record in tqdm(records):
+            writer.add_document(Document(**{k: [str(v)] for k, v in record.items()}))
+        writer.commit()
+        writer.wait_merging_threads()
+
+    def search(
+        self, queries: list[str], fields: list[str], limit: int = 8, escape: bool = True
+    ) -> list[SearchResult]:
+        assert self.index is not None, "Index not built"
+
+        # Validate that requested fields are searchable
+        if self.searchable_fields is not None:
+            non_searchable = set(fields) - self.searchable_fields
+            if non_searchable:
+                raise ValueError(
+                    f"Cannot search non-searchable fields: {sorted(non_searchable)}. "
+                    f"Searchable fields are: {sorted(self.searchable_fields)}"
+                )
+
+        self.index.reload()
+        searcher = self.index.searcher()
+        results = []
+
+        # Get field names for query sanitization
+        known_fields = self._get_schema_field_names()
+
+        for query in queries:
+            if escape:
+                # Three-tier fallback strategy
+                # Tier 1: Try sanitized query with context-aware cleaning
+                sanitized = sanitize_query_for_tantivy(query, known_fields)
+                try:
+                    query_obj = self.index.parse_query(sanitized, fields)
+                except ValueError:
+                    # Tier 2: Bag-of-words fallback - strip all operators
+                    bag_of_words = _strip_syntax_to_words(sanitized)
+                    try:
+                        query_obj = self.index.parse_query(bag_of_words, fields)
+                    except ValueError:
+                        # This should never happen, but if it does, match nothing
+                        query_obj = self.index.parse_query("", fields)
+            else:
+                # User disabled escaping - use query as-is
+                try:
+                    query_obj = self.index.parse_query(query, fields)
+                except Exception as e:
+                    print(f"Error parsing query: {e}")
+                    query_obj = None
+
+            if query_obj:
+                hits = searcher.search(query_obj, limit=limit).hits
+            else:
+                hits = []
+
+            for score, doc_address in hits:
+                doc = searcher.doc(doc_address)
+                content = {k: v[0] for k, v in doc.to_dict().items()}
+                results.append(
+                    SearchResult(
+                        id=content["id"],
+                        score=score,
+                        content=content,
+                    )
+                )
+
+        # Rank fusion using reciprocal rank fusion
+        doc_scores = {}
+        for rank, result in enumerate(results, 1):
+            doc_key = str(result.id)
+            if doc_key not in doc_scores:
+                doc_scores[doc_key] = 0
+            doc_scores[doc_key] += 1 / (60 + rank)
+
+        # Sort and deduplicate
+        unique_results = {}
+        for result in results:
+            doc_key = str(result.id)
+            if doc_key not in unique_results:
+                unique_results[doc_key] = result
+
+        sorted_results = sorted(
+            unique_results.values(),
+            key=lambda x: doc_scores[str(x.id)],
+            reverse=True,
+        )
+
+        return sorted_results[:limit]
+
+    def cache_documents_for_fetch(
+        self, documents: list[dict], id_column: str = "id"
+    ) -> None:
+        """Cache documents for efficient retrieval by ID.
+
+        This builds an in-memory lookup table for O(1) document fetching.
+        Should be called after build_index() with the original document data.
+
+        Args:
+            documents: List of document records (typically original unfiltered data)
+            id_column: Name of the ID field in the documents (default: "id")
+        """
+        self.id_to_record = {}
+        for record in documents:
+            doc_id = str(record.get(id_column, ""))
+            self.id_to_record[doc_id] = record
+
+    def fetch(self, document_ids: list[str]) -> tuple[list[dict], list[str], list[str]]:
+        """Fetch documents by ID using cached lookup.
+
+        Returns full original documents from the cache built by cache_documents_for_fetch().
+        If cache is empty, falls back to querying the Tantivy index (slower, returns only indexed fields).
+
+        Args:
+            document_ids: List of document IDs to retrieve
+
+        Returns:
+            Tuple of (documents, found_ids, missing_ids) where:
+            - documents: List of document dicts with all fields
+            - found_ids: List of IDs that were successfully found
+            - missing_ids: List of IDs that were not found
+        """
+        results = []
+        found_ids = []
+        missing_ids = []
+
+        # Use cached lookup if available
+        if self.id_to_record:
+            for doc_id in document_ids:
+                if doc_id in self.id_to_record:
+                    results.append(self.id_to_record[doc_id])
+                    found_ids.append(doc_id)
+                else:
+                    missing_ids.append(doc_id)
+            return results, found_ids, missing_ids
+
+        # Fallback: query Tantivy index (slower, returns only indexed fields)
+        assert self.index is not None, "Index not built and no cached documents"
+
+        self.index.reload()
+        searcher = self.index.searcher()
+
+        for doc_id in document_ids:
+            query_str = f'id:"{doc_id}"'
+            try:
+                query = self.index.parse_query(query_str, ["id"])
+                hits = searcher.search(query, limit=1).hits
+
+                if hits:
+                    _, doc_address = hits[0]
+                    doc = searcher.doc(doc_address)
+                    content = {k: v[0] for k, v in doc.to_dict().items()}
+                    results.append(content)
+                    found_ids.append(doc_id)
+                else:
+                    missing_ids.append(doc_id)
+            except Exception:
+                missing_ids.append(doc_id)
+
+        return results, found_ids, missing_ids