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

ate/memory/migrate/sources/weaviate.py

@@ -0,0 +1,160 @@
+"""Weaviate migration source implementation."""
+from typing import Optional, List, Tuple, Dict, Any
+
+try:
+    import weaviate
+    WEAVIATE_AVAILABLE = True
+except ImportError:
+    # Create a simple mock structure for testing
+    class WeaviateMock:
+        class Client:
+            pass
+        class auth:
+            class AuthApiKey:
+                def __init__(self, **kwargs):
+                    pass
+    
+    weaviate = WeaviateMock()
+    WEAVIATE_AVAILABLE = False
+
+from ..base import MigrationSource, VectorRecord, MigrationEstimate
+
+
+class WeaviateMigrationSource(MigrationSource):
+    """Migration source for Weaviate vector database."""
+    
+    def __init__(self, host: str, class_name: str, api_key: Optional[str] = None, auth_config: Optional[Dict[str, Any]] = None):
+        """Initialize Weaviate migration source.
+        
+        Args:
+            host: Weaviate server host URL
+            class_name: Name of the class to migrate from
+            api_key: Optional API key for authentication
+            auth_config: Optional authentication configuration
+        """
+        self.host = host
+        self.class_name = class_name
+        self.api_key = api_key
+        self.auth_config = auth_config or {}
+        self.client = None
+    
+    @property
+    def source_type(self) -> str:
+        """Return the source type."""
+        return "weaviate"
+    
+    @property
+    def source_name(self) -> str:
+        """Return the source name."""
+        return self.class_name
+    
+    def connect(self) -> None:
+        """Connect to Weaviate."""
+        if weaviate.Client is None:
+            raise ImportError("weaviate-client library is required for Weaviate migration")
+        
+        # Build client configuration
+        if self.api_key:
+            auth = weaviate.auth.AuthApiKey(api_key=self.api_key)
+            self.client = weaviate.Client(url=self.host, auth_client_secret=auth)
+        else:
+            self.client = weaviate.Client(url=self.host)
+    
+    def estimate(self) -> MigrationEstimate:
+        """Estimate migration size and time."""
+        if not self.client:
+            raise RuntimeError("Must call connect() first")
+        
+        # Get aggregate count using fluent API
+        result = (self.client.query
+                     .aggregate(self.class_name)
+                     .with_meta_count()
+                     .do())
+        
+        aggregate_data = result.get('data', {}).get('Aggregate', {}).get(self.class_name, [])
+        
+        if aggregate_data:
+            total_vectors = aggregate_data[0].get('meta', {}).get('count', 0)
+        else:
+            total_vectors = 0
+        
+        # Get class schema to determine vector dimensions
+        schema = self.client.schema.get()
+        dimensions = 768  # Default assumption
+        
+        for class_def in schema.get('classes', []):
+            if class_def['class'] == self.class_name:
+                # Look for vectorizer configuration
+                if 'vectorizer' in class_def:
+                    # This is a simplified assumption; real implementation would inspect vectorizer config
+                    dimensions = 1536 if 'openai' in class_def['vectorizer'].lower() else 768
+                break
+        
+        # Rough estimates
+        bytes_per_vector = dimensions * 4 + 2048  # 4 bytes per float + metadata overhead
+        estimated_mv2_bytes = total_vectors * bytes_per_vector
+        estimated_seconds = total_vectors / 800.0  # Rough estimate of 800 vectors/second
+        
+        return MigrationEstimate(
+            total_vectors=total_vectors,
+            dimensions=dimensions,
+            estimated_mv2_bytes=estimated_mv2_bytes,
+            estimated_seconds=estimated_seconds
+        )
+    
+    def fetch_batch(self, batch_size: int = 10000, cursor: Optional[str] = None) -> Tuple[List[VectorRecord], Optional[str]]:
+        """Fetch a batch of records from Weaviate."""
+        if not self.client:
+            raise RuntimeError("Must call connect() first")
+        
+        # Parse cursor as offset if provided
+        offset = 0
+        if cursor:
+            try:
+                offset = int(cursor)
+            except (ValueError, TypeError):
+                offset = 0
+        
+        # Query with additional vector data
+        result = (self.client.query
+                 .get(self.class_name)
+                 .with_additional(['id', 'vector'])
+                 .with_limit(batch_size)
+                 .with_offset(offset)
+                 .do())
+        
+        if 'errors' in result:
+            raise Exception(f"Weaviate query error: {result['errors']}")
+        
+        objects = result.get('data', {}).get('Get', {}).get(self.class_name, [])
+        
+        records = []
+        for obj in objects:
+            additional = obj.get('_additional', {})
+            obj_id = additional.get('id', 'unknown')
+            vector = additional.get('vector', [])
+            
+            # Extract text and metadata
+            properties = {k: v for k, v in obj.items() if not k.startswith('_')}
+            text = properties.get('text') or properties.get('content')
+            
+            # Create clean metadata without text
+            clean_metadata = {k: v for k, v in properties.items() if k not in ('text', 'content')}
+            
+            record = VectorRecord(
+                id=obj_id,
+                vector=vector,
+                text=text,
+                metadata=clean_metadata
+            )
+            records.append(record)
+        
+        # Determine if there are more results
+        next_cursor = str(offset + batch_size) if len(objects) == batch_size else None
+        
+        return records, next_cursor
+    
+    def close(self) -> None:
+        """Close the connection and clean up resources."""
+        # Weaviate client doesn't require explicit cleanup
+        self.client = None

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"