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

ate/memory/store.py

@@ -0,0 +1,548 @@
+"""Core MemoryStore implementation."""
+
+import json
+import os
+from typing import List, Dict, Any, Optional
+
+import memvid_sdk
+
+from .search import SearchResult
+from .export import MemoryInfo
+from .embeddings import EmbeddingConfig, EmbeddingManager
+from . import reranker
+
+
+class MemoryStore:
+    """Main interface for memory operations using memvid-sdk backend."""
+    
+    def __init__(self, mem):
+        """Initialize with a memvid memory instance."""
+        self._mem = mem
+        self._embedding_config = None
+
+    @staticmethod
+    def _extract_metadata(frame_data):
+        """Extract metadata from a frame_data dict.
+
+        Decodes JSON-encoded values, filters out ``extractous_metadata``,
+        and returns a plain dict.  Returns ``{}`` for non-dict input or
+        when ``extra_metadata`` is missing.
+        """
+        metadata = {}
+        if not isinstance(frame_data, dict):
+            return metadata
+        for key, value in frame_data.get('extra_metadata', {}).items():
+            if key == 'extractous_metadata':
+                continue
+            try:
+                if isinstance(value, str) and (value.startswith('{') or value.startswith('"')):
+                    metadata[key] = json.loads(value)
+                else:
+                    metadata[key] = value
+            except json.JSONDecodeError:
+                metadata[key] = value
+        return metadata
+    
+    @classmethod
+    def create(cls, path: str, enable_vec: bool = True, enable_lex: bool = True,
+               embedding_config: Optional[EmbeddingConfig] = None) -> 'MemoryStore':
+        """Create a new memory file.
+        
+        Args:
+            path: Path where the .mv2 file will be created
+            enable_vec: Whether to enable vector indexing
+            enable_lex: Whether to enable lexical indexing
+            embedding_config: Optional embedding configuration. If None, auto-detects.
+            
+        Returns:
+            MemoryStore instance
+        """
+        if embedding_config is None:
+            embedding_config = EmbeddingManager.detect()
+        
+        mem = memvid_sdk.use("basic", path, mode="create", 
+                            enable_vec=enable_vec, enable_lex=enable_lex)
+        
+        store = cls(mem)
+        store._embedding_config = embedding_config
+        return store
+    
+    @classmethod
+    def from_context(cls) -> 'MemoryStore':
+        """Open the active memory from context. Auto-initializes on first use."""
+        from .context import ContextManager
+        ctx = ContextManager.get_context()
+        return cls.open(ctx.path)
+
+    @classmethod
+    def from_context_or_path(cls, path: Optional[str] = None) -> 'MemoryStore':
+        """Open from explicit path or fall back to active context."""
+        if path:
+            return cls.open(path)
+        return cls.from_context()
+
+    @classmethod
+    def open(cls, path: str) -> 'MemoryStore':
+        """Open an existing memory file.
+        
+        Args:
+            path: Path to the existing .mv2 file
+            
+        Returns:
+            MemoryStore instance
+            
+        Raises:
+            FileNotFoundError: If the file doesn't exist
+        """
+        try:
+            mem = memvid_sdk.use("basic", path, mode="open")
+        except (FileNotFoundError, OSError):
+            raise FileNotFoundError(f"Memory file not found: {path}")
+        except Exception as e:
+            # Check if it's a file-not-found from memvid
+            if not os.path.exists(path):
+                raise FileNotFoundError(f"Memory file not found: {path}")
+            raise
+        
+        store = cls(mem)
+        # TODO: Read embedding config from metadata when implemented
+        store._embedding_config = EmbeddingManager.detect()
+        return store
+    
+    def close(self):
+        """Close the memory store."""
+        if self._mem:
+            self._mem.close()
+    
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - ensures cleanup."""
+        self.close()
+    
+    def add(self, text: str, tags: Optional[List[str]] = None, 
+            metadata: Optional[Dict[str, Any]] = None, title: Optional[str] = None,
+            enable_embedding: Optional[bool] = None) -> str:
+        """Add content to memory.
+        
+        Args:
+            text: The content text to store
+            tags: Optional list of tags
+            metadata: Optional metadata dictionary
+            title: Optional title for the content
+            enable_embedding: Whether to enable embedding for this entry
+            
+        Returns:
+            Frame ID as string
+            
+        Raises:
+            ValueError: If text is empty
+            TypeError: If text is None
+        """
+        if text is None:
+            raise TypeError("text cannot be None")
+        if text == "":
+            raise ValueError("text cannot be empty")
+            
+        # Prepare put arguments
+        put_args = {
+            'text': text,
+            'tags': tags,
+            'metadata': metadata,
+            'title': title
+        }
+        
+        # Add embedding args if specified
+        if enable_embedding is not None:
+            put_args['enable_embedding'] = enable_embedding
+            
+        frame_id = self._mem.put(**put_args)
+        return str(frame_id)
+    
+    def add_batch(self, items: List[Dict[str, Any]]) -> List[str]:
+        """Add multiple items to memory.
+        
+        Args:
+            items: List of dicts with 'text' key and optional 'tags', 'metadata', 'title'
+            
+        Returns:
+            List of frame IDs as strings
+            
+        Raises:
+            ValueError: If any item is missing 'text' key
+        """
+        if not items:
+            return []
+            
+        frame_ids = []
+        for item in items:
+            if 'text' not in item:
+                raise ValueError("Each item must have a 'text' key")
+            
+            text = item['text']
+            tags = item.get('tags')
+            metadata = item.get('metadata')
+            title = item.get('title')
+            
+            frame_id = self._mem.put(text=text, tags=tags, metadata=metadata, title=title)
+            frame_ids.append(str(frame_id))
+            
+        return frame_ids
+    
+    def search(self, query: str, top_k: int = 5, engine: Optional[str] = None) -> List[SearchResult]:
+        """Search memory for relevant content.
+        
+        Args:
+            query: Search query string
+            top_k: Maximum number of results to return
+            engine: Search engine ("auto" | "vec" | "lex" | "hybrid" | "rerank")
+            
+        Returns:
+            List of SearchResult objects ordered by relevance
+            
+        Raises:
+            ValueError: If query is empty
+        """
+        if not query.strip():
+            raise ValueError("query cannot be empty")
+            
+        # Handle rerank engine
+        if engine == "rerank":
+            return self._search_with_rerank(query, top_k)
+        
+        # Handle auto engine with potential rerank fallback
+        if engine == "auto":
+            # Check if vector index exists
+            try:
+                stats = self._mem.stats()
+                has_vec = stats.has_vec_index if hasattr(stats, 'has_vec_index') else stats.get('has_vec_index', False)
+                
+                if has_vec:
+                    # Use vector search if available
+                    engine = "vec"
+                else:
+                    # Check if LLM reranker is available
+                    reranker = self._get_reranker()
+                    if reranker is not None:
+                        return self._search_with_rerank(query, top_k)
+                    else:
+                        # Fallback to lexical
+                        engine = "lex"
+            except:
+                # Fallback to lexical on any error
+                engine = "lex"
+            
+        # Prepare find arguments for non-rerank engines
+        find_args = {'k': top_k}
+        # Note: real memvid find() does not accept 'engine' — it auto-selects
+        # based on available indices. We just call find() and read the engine
+        # from the response.
+            
+        result = self._mem.find(query, **find_args)
+        
+        # Handle both dict and object interfaces for result
+        if hasattr(result, 'hits'):
+            # Object interface (for mocking)
+            hits = result.hits
+            result_engine = getattr(result, 'engine', 'tantivy')
+        else:
+            # Dict interface (real implementation)
+            hits = result.get('hits', [])
+            result_engine = result.get('engine', 'tantivy')
+        
+        # Map engine names
+        if result_engine == 'tantivy':
+            mapped_engine = 'lex'
+        elif result_engine == 'vector':
+            mapped_engine = 'vec'
+        elif result_engine == 'hybrid':
+            mapped_engine = 'hybrid'
+        else:
+            mapped_engine = result_engine
+        
+        search_results = []
+        for hit in hits:
+            # Get metadata from the frame
+            # Handle both dict and object interfaces for hit
+            if hasattr(hit, 'frame_id'):
+                frame_id = hit.frame_id
+                snippet = hit.snippet
+                title = hit.title
+                score = hit.score
+                tags = getattr(hit, 'tags', [])
+                
+                # For mocked tests, uri might not be set properly
+                uri = getattr(hit, 'uri', f"mv2://frames/{frame_id}")
+            else:
+                frame_id = hit['frame_id']
+                snippet = hit['snippet']
+                title = hit.get('title')
+                score = hit['score']
+                tags = hit.get('tags', [])
+                uri = hit['uri']
+            
+            frame_data = self._mem.frame(uri)
+            metadata = {}
+            
+            # Handle both dict and object interfaces for frame_data
+            if hasattr(frame_data, 'metadata'):
+                # Object interface (for mocking)
+                if frame_data.metadata:
+                    metadata = frame_data.metadata
+            else:
+                # Dict interface (real implementation)
+                metadata = self._extract_metadata(frame_data)
+            
+            search_result = SearchResult(
+                frame_id=frame_id,
+                text=snippet,
+                title=title,
+                score=score,
+                tags=tags,
+                metadata=metadata,
+                engine=mapped_engine
+            )
+            search_results.append(search_result)
+        
+        # Sort by score descending
+        search_results.sort(key=lambda x: x.score, reverse=True)
+        return search_results
+
+    def _search_with_rerank(self, query: str, top_k: int) -> List[SearchResult]:
+        """Search using LLM reranking."""
+        # Phase 1: BM25 wide retrieval (get more candidates than needed)
+        candidate_count = min(top_k * 4, 50)
+        bm25_results = self._bm25_search(query, candidate_count)
+        
+        # If BM25 returns nothing (common with natural language queries),
+        # try fetching ALL frames as candidates for the LLM to evaluate
+        if not bm25_results:
+            bm25_results = self._get_all_frames(max_frames=50)
+        
+        # Phase 2: LLM re-ranking
+        reranker = self._get_reranker()
+        if reranker is None:
+            # Fallback to BM25 if no LLM available
+            return bm25_results[:top_k]
+        
+        return reranker.rerank(query, bm25_results, top_k=top_k)
+
+    def _bm25_search(self, query: str, top_k: int) -> List[SearchResult]:
+        """Perform BM25 search and return SearchResult objects."""
+        # Note: real memvid find() does not accept 'engine' kwarg.
+        # It auto-selects based on available indices (lex by default).
+        result = self._mem.find(query, k=top_k)
+        
+        # Handle both dict and object interfaces
+        if hasattr(result, 'hits'):
+            hits = result.hits
+        else:
+            hits = result.get('hits', [])
+        
+        search_results = []
+        for hit in hits:
+            # Handle both dict and object interfaces for hit
+            if hasattr(hit, 'frame_id'):
+                frame_id = hit.frame_id
+                snippet = hit.snippet
+                title = hit.title
+                score = hit.score
+                tags = getattr(hit, 'tags', [])
+                uri = getattr(hit, 'uri', f"mv2://frames/{frame_id}")
+            else:
+                frame_id = hit['frame_id']
+                snippet = hit['snippet']
+                title = hit.get('title')
+                score = hit['score']
+                tags = hit.get('tags', [])
+                uri = hit['uri']
+            
+            frame_data = self._mem.frame(uri)
+            metadata = {}
+            
+            # Handle both dict and object interfaces for frame_data
+            if hasattr(frame_data, 'metadata'):
+                if frame_data.metadata:
+                    metadata = frame_data.metadata
+            else:
+                metadata = self._extract_metadata(frame_data)
+            
+            search_result = SearchResult(
+                frame_id=frame_id,
+                text=snippet,
+                title=title,
+                score=score,
+                tags=tags,
+                metadata=metadata,
+                engine="lex"  # BM25 is lexical
+            )
+            search_results.append(search_result)
+        
+        return search_results
+
+    def _get_all_frames(self, max_frames: int = 50) -> List[SearchResult]:
+        """Get all frames as SearchResult objects for LLM re-ranking.
+        
+        Used when BM25 returns no candidates but we still want the LLM
+        to evaluate all stored memories against the query.
+        """
+        results = []
+        try:
+            timeline = self._mem.timeline()
+            if timeline is None:
+                return results
+            for entry in timeline[:max_frames]:
+                if isinstance(entry, dict):
+                    uri = entry.get('uri', f"mv2://frames/{entry.get('frame_id', len(results))}")
+                    frame_data = self._mem.frame(uri)
+                    text = entry.get('preview', '').split('\ntitle:')[0].split('\ntags:')[0].strip()
+                    if isinstance(frame_data, dict):
+                        title = frame_data.get('title')
+                        tags = frame_data.get('tags', [])
+                        metadata = self._extract_metadata(frame_data)
+                    else:
+                        title = getattr(frame_data, 'title', None)
+                        tags = getattr(frame_data, 'tags', [])
+                        metadata = getattr(frame_data, 'metadata', {}) or {}
+                    results.append(SearchResult(
+                        frame_id=entry.get('frame_id', len(results)),
+                        text=text,
+                        title=title,
+                        score=0.0,
+                        tags=tags,
+                        metadata=metadata,
+                        engine="lex"
+                    ))
+        except Exception:
+            pass
+        return results
+
+    def _get_reranker(self) -> Optional['reranker.LLMReranker']:
+        """Get a reranker instance if LLM provider is available."""
+        config = reranker.LLMReranker.detect()
+        if config is None:
+            return None
+        return reranker.LLMReranker(config)
+    
+    def export_jsonl(self, output_path: str) -> int:
+        """Export all memory content to JSONL file.
+        
+        Args:
+            output_path: Path to output JSONL file
+            
+        Returns:
+            Number of items exported
+        """
+        count = 0
+        with open(output_path, 'w', encoding='utf-8') as f:
+            # Try timeline-based iteration first (works with real memvid)
+            try:
+                timeline = self._mem.timeline()
+                if timeline is not None:
+                    for entry in timeline:
+                        if isinstance(entry, dict):
+                            uri = entry.get('uri', f"mv2://frames/{entry.get('frame_id', count)}")
+                            frame_data = self._mem.frame(uri)
+                            
+                            # Extract text — preview from timeline has the original text
+                            text = entry.get('preview', '').split('\ntitle:')[0].split('\ntags:')[0].strip()
+                            if isinstance(frame_data, dict):
+                                title = frame_data.get('title')
+                                tags = frame_data.get('tags', [])
+                                
+                                # Extract metadata from extra_metadata
+                                metadata = self._extract_metadata(frame_data)
+                            else:
+                                title = getattr(frame_data, 'title', None)
+                                tags = getattr(frame_data, 'tags', [])
+                                metadata = getattr(frame_data, 'metadata', {}) or {}
+                                text = getattr(frame_data, 'text', '') or getattr(frame_data, 'snippet', '')
+                            
+                            record = {
+                                'frame_id': entry.get('frame_id', count),
+                                'text': text,
+                                'title': title,
+                                'tags': tags,
+                                'metadata': metadata
+                            }
+                            f.write(json.dumps(record) + '\n')
+                            count += 1
+                    
+                    if count > 0:
+                        return count
+            except (AttributeError, TypeError):
+                pass  # Fall through to find-based approach
+            
+            # Fallback: use find with wildcard (for mocked tests)
+            result = self._mem.find("*", k=10000)
+            
+            if hasattr(result, 'hits'):
+                hits = result.hits
+            else:
+                hits = result.get('hits', [])
+            
+            for hit in hits:
+                if hasattr(hit, 'frame_id'):
+                    frame_id = hit.frame_id
+                    snippet = hit.snippet
+                    title = hit.title
+                    tags = getattr(hit, 'tags', [])
+                    uri = getattr(hit, 'uri', f"mv2://frames/{frame_id}")
+                else:
+                    frame_id = hit['frame_id']
+                    snippet = hit['snippet']
+                    title = hit.get('title')
+                    tags = hit.get('tags', [])
+                    uri = hit['uri']
+                
+                frame_data = self._mem.frame(uri)
+                metadata = {}
+                if hasattr(frame_data, 'metadata'):
+                    if frame_data.metadata:
+                        metadata = frame_data.metadata
+                elif isinstance(frame_data, dict):
+                    metadata = self._extract_metadata(frame_data)
+                
+                record = {
+                    'frame_id': frame_id,
+                    'text': snippet,
+                    'title': title,
+                    'tags': tags,
+                    'metadata': metadata
+                }
+                f.write(json.dumps(record) + '\n')
+                count += 1
+        
+        return count
+    
+    def info(self) -> MemoryInfo:
+        """Get information about the memory store.
+        
+        Returns:
+            MemoryInfo object with store statistics
+        """
+        stats = self._mem.stats()
+        
+        # Handle both dict and object interfaces for stats
+        if hasattr(stats, 'frame_count'):
+            # Object interface (for mocking)
+            return MemoryInfo(
+                path=self._mem.path,
+                frame_count=stats.frame_count,
+                size_bytes=stats.size_bytes,
+                has_lex_index=stats.has_lex_index,
+                has_vec_index=stats.has_vec_index,
+                has_time_index=stats.has_time_index
+            )
+        else:
+            # Dict interface (real implementation)
+            return MemoryInfo(
+                path=self._mem.path,
+                frame_count=stats['frame_count'],
+                size_bytes=stats['size_bytes'],
+                has_lex_index=stats['has_lex_index'],
+                has_vec_index=stats['has_vec_index'],
+                has_time_index=stats['has_time_index']
+            )

ate/recording/__init__.py

@@ -7,7 +7,7 @@ Records all interface method calls as transferable data that can be:
 - Used to train policies
 - Replayed on different hardware
 
-Example:
+Example - Basic recording:
     from ate.drivers import MechDogDriver
     from ate.recording import RecordingSession
 
@@ -24,21 +24,60 @@ Example:
     # Save the recording
     session.save("pickup_toy.demonstration")
 
-    # Later: upload to FoodforThought
-    session.upload()
+Example - Visual recording with object detection:
+    from ate.drivers import MechDogDriver
+    from ate.recording import VisualRecordingSession
+    from ate.detection import TrashDetector
+
+    dog = MechDogDriver(port="/dev/cu.usbserial-10", config=config)
+    dog.connect()
+
+    # Create detector for automatic object detection
+    detector = TrashDetector()
+
+    # Record with visual capture at 5 fps
+    with VisualRecordingSession(
+        dog,
+        name="trash_pickup",
+        capture_fps=5.0,
+        detector=detector,
+    ) as session:
+        dog.stand()
+        dog.walk(Vector3.forward(), speed=0.2)
+        time.sleep(5)
+
+    # Save recording and visual frames separately
+    session.save("trash_pickup.demonstration")
+    session.save_frames("trash_pickup_frames.json")
+
+    print(session.summary())
 """
 
 from .session import RecordingSession, RecordedCall
 from .wrapper import RecordingWrapper
 from .demonstration import Demonstration, load_demonstration
 from .upload import DemonstrationUploader, upload_demonstration
+from .visual import (
+    VisualRecordingSession,
+    RecordedFrame,
+    DetectionResult,
+    load_frames,
+)
 
 __all__ = [
+    # Core recording
     "RecordingSession",
     "RecordedCall",
     "RecordingWrapper",
+    # Demonstrations
     "Demonstration",
     "load_demonstration",
+    # Upload
     "DemonstrationUploader",
     "upload_demonstration",
+    # Visual recording with detection
+    "VisualRecordingSession",
+    "RecordedFrame",
+    "DetectionResult",
+    "load_frames",
 ]

ate/recording/session.py

@@ -54,14 +54,24 @@ class RecordedCall:
         # Handle common types
         if value is None or isinstance(value, (bool, int, float, str)):
             return value
+        if isinstance(value, bytes):
+            # Don't include large binary data in JSON - just note its size
+            if len(value) > 1000:
+                return {"__type__": "bytes", "__size__": len(value), "__truncated__": True}
+            import base64
+            return {"__type__": "bytes", "__b64__": base64.b64encode(value).decode('ascii')}
         if isinstance(value, (list, tuple)):
             return [self._serialize_value(v) for v in value]
         if isinstance(value, dict):
             return {k: self._serialize_value(v) for k, v in value.items()}
 
-        # Handle our interface types
+        # Handle our interface types (serialize fields individually to handle bytes)
         if hasattr(value, '__dataclass_fields__'):
-            return {"__type__": type(value).__name__, **asdict(value)}
+            result = {"__type__": type(value).__name__}
+            for field_name in value.__dataclass_fields__:
+                field_value = getattr(value, field_name)
+                result[field_name] = self._serialize_value(field_value)
+            return result
         if hasattr(value, 'to_dict'):
             return {"__type__": type(value).__name__, **value.to_dict()}