foodforthought-cli 0.2.8__py3-none-any.whl → 0.3.0__py3-none-any.whl

ate/memory/reranker.py

@@ -0,0 +1,353 @@
+"""LLM re-ranking search engine for ate memory."""
+
+import json
+import os
+import requests
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional, List
+
+from .search import SearchResult
+
+
+@dataclass
+class RerankConfig:
+    """Configuration for LLM re-ranking."""
+    provider: str          # "anthropic" | "openai" | "google" | "ollama"
+    model: Optional[str] = None   # Model override (e.g. "claude-haiku-3.5")
+    api_key: Optional[str] = None # Explicit key (overrides env)
+
+    def __post_init__(self):
+        """Validate provider and set default models after initialization."""
+        valid_providers = ["anthropic", "openai", "google", "ollama"]
+        if self.provider not in valid_providers:
+            raise ValueError(f"Invalid provider: {self.provider}. Must be one of {valid_providers}")
+
+        # Set default models based on provider
+        if self.model is None:
+            if self.provider == "anthropic":
+                self.model = "claude-haiku-3.5-latest"
+            elif self.provider == "openai":
+                self.model = "gpt-4o-mini"
+            elif self.provider == "google":
+                self.model = "gemini-2.0-flash"
+            elif self.provider == "ollama":
+                self.model = "llama3.2"
+
+
+class LLMProvider(ABC):
+    """Abstract LLM provider for re-ranking."""
+
+    @abstractmethod
+    def complete(self, prompt: str, max_tokens: int = 200) -> str:
+        """Send a completion request and return the text response."""
+        pass
+
+
+class AnthropicProvider(LLMProvider):
+    """Anthropic Claude provider."""
+    
+    def __init__(self, model: str, api_key: str):
+        self.model = model
+        self.api_key = api_key
+
+    def complete(self, prompt: str, max_tokens: int = 200) -> str:
+        """Send a completion request to Anthropic API."""
+        headers = {
+            'x-api-key': self.api_key,
+            'anthropic-version': '2023-06-01',
+            'content-type': 'application/json'
+        }
+        
+        data = {
+            'model': self.model,
+            'max_tokens': max_tokens,
+            'messages': [{'role': 'user', 'content': prompt}]
+        }
+        
+        response = requests.post(
+            'https://api.anthropic.com/v1/messages',
+            headers=headers,
+            json=data,
+            timeout=5
+        )
+        response.raise_for_status()
+        
+        result = response.json()
+        return result['content'][0]['text']
+
+
+class OpenAIProvider(LLMProvider):
+    """OpenAI provider."""
+    
+    def __init__(self, model: str, api_key: str):
+        self.model = model
+        self.api_key = api_key
+
+    def complete(self, prompt: str, max_tokens: int = 200) -> str:
+        """Send a completion request to OpenAI API."""
+        headers = {
+            'Authorization': f'Bearer {self.api_key}',
+            'content-type': 'application/json'
+        }
+        
+        data = {
+            'model': self.model,
+            'max_tokens': max_tokens,
+            'messages': [{'role': 'user', 'content': prompt}]
+        }
+        
+        response = requests.post(
+            'https://api.openai.com/v1/chat/completions',
+            headers=headers,
+            json=data,
+            timeout=5
+        )
+        response.raise_for_status()
+        
+        result = response.json()
+        return result['choices'][0]['message']['content']
+
+
+class GoogleProvider(LLMProvider):
+    """Google Gemini provider."""
+    
+    def __init__(self, model: str, api_key: str):
+        self.model = model
+        self.api_key = api_key
+
+    def complete(self, prompt: str, max_tokens: int = 200) -> str:
+        """Send a completion request to Google Gemini API."""
+        headers = {
+            'content-type': 'application/json'
+        }
+        
+        data = {
+            'contents': [{'parts': [{'text': prompt}]}],
+            'generationConfig': {'maxOutputTokens': max_tokens}
+        }
+        
+        url = f'https://generativelanguage.googleapis.com/v1beta/models/{self.model}:generateContent?key={self.api_key}'
+        response = requests.post(url, headers=headers, json=data, timeout=5)
+        response.raise_for_status()
+        
+        result = response.json()
+        return result['candidates'][0]['content']['parts'][0]['text']
+
+
+class OllamaLLMProvider(LLMProvider):
+    """Ollama local provider."""
+    
+    def __init__(self, model: str, host: str = "http://localhost:11434"):
+        self.model = model
+        self.host = host
+
+    def complete(self, prompt: str, max_tokens: int = 200) -> str:
+        """Send a completion request to Ollama API."""
+        headers = {
+            'content-type': 'application/json'
+        }
+        
+        data = {
+            'model': self.model,
+            'prompt': prompt,
+            'stream': False
+        }
+        
+        response = requests.post(
+            f'{self.host}/api/generate',
+            headers=headers,
+            json=data,
+            timeout=30  # Ollama can be slower
+        )
+        response.raise_for_status()
+        
+        result = response.json()
+        return result['response']
+
+
+class LLMReranker:
+    """Re-ranks BM25 search results using an LLM for semantic understanding."""
+
+    def __init__(self, config: RerankConfig):
+        """Initialize with LLM provider config."""
+        self.config = config
+
+    @staticmethod
+    def detect() -> Optional[RerankConfig]:
+        """Auto-detect LLM provider from env vars.
+        Detection order: ANTHROPIC_API_KEY → OPENAI_API_KEY → GOOGLE_API_KEY → Ollama
+        """
+        # Check Anthropic first (highest priority)
+        anthropic_key = os.environ.get('ANTHROPIC_API_KEY')
+        if anthropic_key:
+            return RerankConfig(
+                provider="anthropic",
+                api_key=anthropic_key
+            )
+
+        # Check OpenAI second
+        openai_key = os.environ.get('OPENAI_API_KEY')
+        if openai_key:
+            return RerankConfig(
+                provider="openai",
+                api_key=openai_key
+            )
+
+        # Check Google third
+        google_key = os.environ.get('GOOGLE_API_KEY')
+        if google_key:
+            return RerankConfig(
+                provider="google",
+                api_key=google_key
+            )
+
+        # Check Ollama fourth (local service)
+        ollama_host = os.environ.get('OLLAMA_HOST', 'http://localhost:11434')
+        if LLMReranker._is_ollama_available(ollama_host):
+            return RerankConfig(
+                provider="ollama"
+            )
+
+        # No providers available
+        return None
+
+    @staticmethod
+    def _is_ollama_available(host: str) -> bool:
+        """Check if Ollama is reachable at the given host."""
+        try:
+            response = requests.get(f'{host}/api/tags', timeout=2)
+            return response.status_code == 200
+        except:
+            return False
+
+    def _get_provider(self) -> LLMProvider:
+        """Get the appropriate LLM provider instance."""
+        if self.config.provider == "anthropic":
+            return AnthropicProvider(self.config.model, self.config.api_key)
+        elif self.config.provider == "openai":
+            return OpenAIProvider(self.config.model, self.config.api_key)
+        elif self.config.provider == "google":
+            return GoogleProvider(self.config.model, self.config.api_key)
+        elif self.config.provider == "ollama":
+            ollama_host = os.environ.get('OLLAMA_HOST', 'http://localhost:11434')
+            return OllamaLLMProvider(self.config.model, ollama_host)
+        else:
+            raise ValueError(f"Unknown provider: {self.config.provider}")
+
+    def _build_rerank_prompt(self, query: str, candidates: List[SearchResult], top_k: int) -> str:
+        """Build the rerank prompt for the LLM."""
+        lines = [
+            "You are a memory retrieval system. Given a query and a list of memory snippets,",
+            "return the indices of the most relevant snippets, ranked by relevance.",
+            "",
+            f"Query: \"{query}\"",
+            "",
+            "Snippets:"
+        ]
+        
+        # Add numbered candidates (truncated to keep prompt reasonable)
+        for i, candidate in enumerate(candidates):
+            title_part = f" — {candidate.title}" if candidate.title else ""
+            text_truncated = candidate.text[:200] + "..." if len(candidate.text) > 200 else candidate.text
+            lines.append(f"[{i}] {text_truncated}{title_part}")
+        
+        lines.extend([
+            "",
+            f"Return ONLY a JSON array of indices in order of relevance, e.g. [3, 0, 7, 1]",
+            f"Return at most {top_k} indices. Only include relevant results.",
+            "If none are relevant, return []."
+        ])
+        
+        return "\n".join(lines)
+
+    def rerank(self, query: str, candidates: List[SearchResult], top_k: int = 5) -> List[SearchResult]:
+        """Re-rank candidates using the LLM.
+
+        Sends a structured prompt to the LLM with the query and candidate texts.
+        Returns re-ordered results with updated scores (1.0 = best match).
+
+        Args:
+            query: Original search query
+            candidates: BM25 search results to re-rank
+            top_k: Number of results to return after re-ranking
+
+        Returns:
+            List of SearchResult objects re-ordered by semantic relevance
+        """
+        # Handle empty candidates
+        if not candidates:
+            return []
+
+        try:
+            # Get LLM provider
+            provider = self._get_provider()
+            
+            # Build prompt
+            prompt = self._build_rerank_prompt(query, candidates, top_k)
+            
+            # Get LLM response
+            response = provider.complete(prompt, max_tokens=200)
+            
+            # Parse response
+            try:
+                indices = json.loads(response.strip())
+            except json.JSONDecodeError:
+                # Retry once with stricter prompt
+                strict_prompt = prompt + "\n\nIMPORTANT: Return ONLY valid JSON array format like [1, 3, 0]. No other text."
+                try:
+                    response = provider.complete(strict_prompt, max_tokens=200)
+                    indices = json.loads(response.strip())
+                except (json.JSONDecodeError, Exception):
+                    # Final fallback to original order
+                    return self._fallback_to_original(candidates, top_k)
+            
+            # Validate indices
+            if not isinstance(indices, list):
+                return self._fallback_to_original(candidates, top_k)
+            
+            # Handle empty result
+            if not indices:
+                return []
+            
+            # Reorder candidates based on LLM ranking
+            reranked = []
+            for rank, idx in enumerate(indices[:top_k]):
+                if isinstance(idx, int) and 0 <= idx < len(candidates):
+                    candidate = candidates[idx]
+                    # Normalize scores: first=1.0, linearly decreasing
+                    score = 1.0 - (rank / len(indices)) if len(indices) > 1 else 1.0
+                    
+                    # Create new result with updated score and engine
+                    reranked_result = SearchResult(
+                        frame_id=candidate.frame_id,
+                        text=candidate.text,
+                        title=candidate.title,
+                        score=score,
+                        tags=candidate.tags,
+                        metadata=candidate.metadata,
+                        engine="rerank"
+                    )
+                    reranked.append(reranked_result)
+            
+            return reranked
+            
+        except Exception:
+            # Fallback to original order on any error
+            return self._fallback_to_original(candidates, top_k)
+
+    def _fallback_to_original(self, candidates: List[SearchResult], top_k: int) -> List[SearchResult]:
+        """Fallback to original BM25 order with rerank engine label."""
+        fallback_results = []
+        for candidate in candidates[:top_k]:
+            fallback_result = SearchResult(
+                frame_id=candidate.frame_id,
+                text=candidate.text,
+                title=candidate.title,
+                score=candidate.score,  # Keep original scores
+                tags=candidate.tags,
+                metadata=candidate.metadata,
+                engine="rerank"  # Still label as rerank even though it fell back
+            )
+            fallback_results.append(fallback_result)
+        return fallback_results

ate/memory/search.py

@@ -0,0 +1,26 @@
+"""Search operations and result structures."""
+
+from dataclasses import dataclass
+from typing import List, Dict, Any, Optional
+
+
+@dataclass
+class SearchResult:
+    """Result from a memory search operation.
+    
+    Attributes:
+        frame_id: Unique identifier for the memory frame
+        text: The original content text
+        title: Title of the content (if provided)
+        score: Relevance score (higher = more relevant)
+        tags: List of tags associated with this content
+        metadata: Dictionary of additional metadata
+        engine: Search engine used ("lex" | "vec" | "hybrid")
+    """
+    frame_id: int
+    text: str
+    title: Optional[str]
+    score: float
+    tags: List[str]
+    metadata: Dict[str, Any]
+    engine: str = "lex"