lm-deluge 0.0.87__py3-none-any.whl → 0.0.89__py3-none-any.whl

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

lm_deluge/tool/prefab/rlm/__init__.py

@@ -0,0 +1,296 @@
+"""
+RLM (Recursive Language Model) for lm-deluge.
+
+Enables models to process long contexts through a REPL environment
+with recursive LM calls, based on the RLM paper:
+https://alexzhang13.github.io/blog/2025/rlm/
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from lm_deluge.prompt import Conversation
+from lm_deluge.tool import Tool
+
+from .executor import RLMExecutionError, RLMExecutor
+from .parse import RLMSecurityError
+
+if TYPE_CHECKING:
+    from lm_deluge.api_requests.base import APIResponse
+    from lm_deluge.client import _LLMClient
+
+
+RLM_SYSTEM_PROMPT = """You have access to a long context stored in the variable `{context_var}`.
+You can write Python code to analyze this context using the `execute` tool.
+
+IMPORTANT RULES:
+1. You MUST use print() to see output. Bare expressions produce NO output.
+2. You MUST call final(answer) when you have the answer. This is required!
+
+Available in your code environment:
+- `{context_var}`: The full context as a string ({context_len:,} characters)
+- `lm(prompt)`: Make a recursive LLM call (runs in parallel when possible)
+- `final(answer)`: Signal completion with the given answer - YOU MUST CALL THIS!
+- `final_var(varname)`: Signal completion with a variable's value
+- Modules: `re`, `math`, `collections`, `json` (imports are allowed but optional)
+- From collections: `Counter`, `defaultdict`, `deque`, `namedtuple`, `OrderedDict`
+- Standard builtins: `len`, `str`, `int`, `list`, `dict`, `sum`, `sorted`, `map`, `filter`, etc.
+
+Example - count word occurrences:
+```python
+count = len(re.findall(r'\\bword\\b', {context_var}))
+print(f"Found {{count}} occurrences")
+final(count)
+```
+
+Example - use Counter:
+```python
+words = {context_var}.split()
+counts = Counter(words)
+print(counts.most_common(10))
+final(counts.most_common(10))
+```
+
+Example - analyze with lm() calls:
+```python
+chunks = [{context_var}[i:i+2000] for i in range(0, len({context_var}), 2000)][:3]
+summaries = [lm(f"Summarize: {{chunk}}") for chunk in chunks]
+combined = "\\n".join(str(s) for s in summaries)
+final(f"Summary:\\n{{combined}}")
+```
+
+Variables persist between execute() calls. Always call final() when you have the answer!
+"""
+
+
+class RLMManager:
+    """Manages RLM execution for a long context.
+
+    The RLMManager exposes a REPL-like interface as tools that allow an LLM
+    to analyze a long context by writing Python code.
+
+    Example:
+        >>> manager = RLMManager(
+        ...     context=long_document,
+        ...     client=LLMClient("gpt-4.1-mini"),  # For lm() calls
+        ... )
+        >>> main_client = LLMClient("gpt-4.1")
+        >>> conv = Conversation.system(manager.get_system_prompt())
+        >>> conv = conv.user("What are the main themes in this document?")
+        >>> conv, resp = await main_client.run_agent_loop(
+        ...     conv,
+        ...     tools=manager.get_tools(),
+        ... )
+        >>> if manager.is_complete:
+        ...     print(manager.final_answer)
+    """
+
+    def __init__(
+        self,
+        context: str,
+        client: _LLMClient,
+        context_var_name: str = "CONTEXT",
+        max_lm_calls_per_execution: int = 20,
+    ):
+        """Initialize the RLMManager.
+
+        Args:
+            context: The long context string to analyze
+            client: LLMClient for making recursive lm() calls
+            context_var_name: Variable name for the context (default: "CONTEXT")
+            max_lm_calls_per_execution: Maximum lm() calls allowed per execute() call
+        """
+        self.context = context
+        self.client = client
+        self.context_var_name = context_var_name
+        self.max_lm_calls_per_execution = max_lm_calls_per_execution
+
+        self.executor = RLMExecutor(
+            context=context,
+            client=client,
+            context_var_name=context_var_name,
+            max_lm_calls_per_execution=max_lm_calls_per_execution,
+        )
+
+        self._final_answer: str | None = None
+        self._tools: list[Tool] | None = None
+
+    async def _execute(self, code: str) -> str:
+        """Execute code against the context."""
+        try:
+            answer, is_final = await self.executor.execute(code)
+            if is_final:
+                self._final_answer = answer
+                # Truncate for display but keep full answer stored
+                display = answer[:1000] + "..." if len(answer) > 1000 else answer
+                return f"[FINAL ANSWER SET]\n{display}"
+            return answer
+        except RLMSecurityError as e:
+            return f"Security error: {e}"
+        except RLMExecutionError as e:
+            return f"Execution error: {e}"
+        except Exception as e:
+            return f"Unexpected error: {type(e).__name__}: {e}"
+
+    def get_system_prompt(self) -> str:
+        """Get the system prompt explaining the RLM environment."""
+        return RLM_SYSTEM_PROMPT.format(
+            context_var=self.context_var_name,
+            context_len=len(self.context),
+        )
+
+    def get_tools(self) -> list[Tool]:
+        """Get the tools for RLM execution."""
+        if self._tools is not None:
+            return self._tools
+
+        self._tools = [
+            Tool(
+                name="execute",
+                description=(
+                    f"Execute Python code to analyze the context. "
+                    f"The context ({len(self.context):,} chars) is available as `{self.context_var_name}`. "
+                    f"Use `lm(prompt)` for recursive LLM calls (parallel when possible), and "
+                    f"`final(answer)` or `final_var(varname)` to signal completion. "
+                    f"Variables persist between calls. "
+                    f"Modules available without import: re, math, collections, json."
+                ),
+                run=self._execute,
+                parameters={
+                    "code": {
+                        "type": "string",
+                        "description": "Python code to execute",
+                    }
+                },
+                required=["code"],
+            )
+        ]
+        return self._tools
+
+    @property
+    def is_complete(self) -> bool:
+        """Check if FINAL() was called."""
+        return self._final_answer is not None
+
+    @property
+    def final_answer(self) -> str | None:
+        """Get the final answer if set."""
+        return self._final_answer
+
+    def reset(self) -> None:
+        """Reset the RLM state."""
+        self.executor.reset()
+        self._final_answer = None
+
+
+@dataclass
+class RLMResult:
+    """Result from RLMPipeline."""
+
+    answer: str
+    conversation: Conversation
+    rounds_used: int
+    final_response: APIResponse
+
+
+class RLMPipeline:
+    """High-level pipeline for RLM processing.
+
+    A thin wrapper that takes a long context and question, sets up an RLMManager,
+    runs an agent loop until final() is called, and returns the result.
+
+    Example:
+        >>> pipeline = RLMPipeline(
+        ...     context=long_document,
+        ...     client=LLMClient("gpt-4.1"),  # Smart orchestrator
+        ...     lm_client=LLMClient("gpt-4.1-mini"),  # Cheaper model for lm() calls
+        ...     question="What are the main themes in this document?",
+        ... )
+        >>> result = await pipeline.run()
+        >>> print(result.answer)
+    """
+
+    def __init__(
+        self,
+        context: str,
+        client: _LLMClient,
+        question: str,
+        *,
+        lm_client: _LLMClient | None = None,
+        context_var_name: str = "CONTEXT",
+        max_rounds: int = 15,
+        max_lm_calls_per_execution: int = 20,
+    ):
+        """Initialize the RLMPipeline.
+
+        Args:
+            context: The long context string to analyze
+            client: LLMClient for the main agent (runs the execute loop)
+            question: The question to answer about the context
+            lm_client: LLMClient for lm() calls (defaults to same as client)
+            context_var_name: Variable name for the context (default: "CONTEXT")
+            max_rounds: Maximum agent loop rounds (default: 15)
+            max_lm_calls_per_execution: Maximum lm() calls per execute() call
+        """
+        self.context = context
+        self.client = client
+        self.lm_client = lm_client or client
+        self.question = question
+        self.context_var_name = context_var_name
+        self.max_rounds = max_rounds
+        self.max_lm_calls_per_execution = max_lm_calls_per_execution
+
+    async def run(self) -> RLMResult:
+        """Run the RLM pipeline until completion."""
+        manager = RLMManager(
+            context=self.context,
+            client=self.lm_client,
+            context_var_name=self.context_var_name,
+            max_lm_calls_per_execution=self.max_lm_calls_per_execution,
+        )
+
+        # Build conversation with system prompt and question
+        conv = Conversation.system(manager.get_system_prompt())
+        conv = conv.user(
+            f"Question to answer about the context:\n\n{self.question}\n\n"
+            "Use the execute tool to analyze the context and find the answer. "
+            "Start by peeking at the context structure, then use appropriate "
+            "techniques (regex, chunking, lm() calls) to find the answer. "
+            "Call final(answer) when you have the answer."
+        )
+
+        # Run agent loop
+        conv, resp = await self.client.run_agent_loop(
+            conv,
+            tools=manager.get_tools(),
+            max_rounds=self.max_rounds,
+        )
+
+        # Extract answer
+        if manager.is_complete:
+            answer = manager.final_answer or "No answer produced"
+        else:
+            # Model stopped without calling final() - use last response
+            answer = resp.completion or "No answer produced (final not called)"
+
+        # Count rounds used
+        rounds_used = sum(1 for m in conv.messages if m.role == "assistant")
+
+        return RLMResult(
+            answer=answer,
+            conversation=conv,
+            rounds_used=rounds_used,
+            final_response=resp,
+        )
+
+
+__all__ = [
+    "RLMManager",
+    "RLMPipeline",
+    "RLMResult",
+    "RLMExecutor",
+    "RLMExecutionError",
+    "RLMSecurityError",
+]