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

ate/memory/migrate/base.py

@@ -0,0 +1,89 @@
+"""Base classes and data structures for migration."""
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Optional, List, Dict, Any, Tuple
+
+
+@dataclass
+class VectorRecord:
+    """A vector record with metadata."""
+    id: str
+    vector: List[float]
+    text: Optional[str] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class MigrationEstimate:
+    """Estimation of migration resources and time."""
+    total_vectors: int
+    dimensions: int
+    estimated_mv2_bytes: int
+    estimated_seconds: float
+    source_size_bytes: Optional[int] = None
+    compression_ratio: Optional[float] = None
+
+
+@dataclass
+class MigrationResult:
+    """Result of a migration operation."""
+    source_type: str
+    source_name: str
+    output_path: str
+    total_migrated: int
+    total_skipped: int
+    duration_seconds: float
+    output_size_bytes: int
+    compression_ratio: Optional[float] = None
+    errors: List[str] = field(default_factory=list)
+
+
+@dataclass
+class MigrationCheckpoint:
+    """Migration checkpoint for resuming."""
+    source_type: str
+    source_name: str
+    output_path: str
+    last_cursor: Optional[str]
+    records_completed: int
+    started_at: str
+
+
+class MigrationSource(ABC):
+    """Abstract base class for migration sources."""
+    
+    @property
+    @abstractmethod
+    def source_type(self) -> str:
+        """Return the source type (e.g., 'pinecone', 'qdrant')."""
+        pass
+    
+    @property
+    @abstractmethod
+    def source_name(self) -> str:
+        """Return the source name (e.g., index name, collection name)."""
+        pass
+    
+    @abstractmethod
+    def connect(self) -> None:
+        """Connect to the source."""
+        pass
+    
+    @abstractmethod
+    def estimate(self) -> MigrationEstimate:
+        """Estimate migration size and time."""
+        pass
+    
+    @abstractmethod
+    def fetch_batch(self, batch_size: int = 10000, cursor: Optional[str] = None) -> Tuple[List[VectorRecord], Optional[str]]:
+        """Fetch a batch of records.
+        
+        Returns:
+            Tuple of (records, next_cursor). next_cursor is None if this is the last batch.
+        """
+        pass
+    
+    @abstractmethod
+    def close(self) -> None:
+        """Close the connection and clean up resources."""
+        pass

ate/memory/migrate/pipeline.py

@@ -0,0 +1,189 @@
+"""Migration pipeline for orchestrating the migration process."""
+import json
+import os
+import time
+from typing import Optional, Callable, Any
+
+from .base import MigrationSource, MigrationResult, MigrationCheckpoint
+
+
+class MigrationPipeline:
+    """Pipeline for orchestrating vector database migrations."""
+    
+    def __init__(self, source: MigrationSource, output_path: str, 
+                 batch_size: int = 10000, checkpoint_path: Optional[str] = None):
+        """Initialize the migration pipeline.
+        
+        Args:
+            source: The migration source to read from
+            output_path: Path to write the .mv2 output file
+            batch_size: Number of records to fetch per batch
+            checkpoint_path: Optional path for checkpoint file
+        """
+        self.source = source
+        self.output_path = output_path
+        self.batch_size = batch_size
+        self.checkpoint_path = checkpoint_path
+        self.progress_callbacks = []
+    
+    def on_progress(self, callback: Callable[[int, int], None]) -> None:
+        """Register a progress callback function.
+        
+        Args:
+            callback: Function called with (records_done, total_estimate)
+        """
+        self.progress_callbacks.append(callback)
+    
+    def run(self, dry_run: bool = False) -> MigrationResult:
+        """Run the migration.
+        
+        Args:
+            dry_run: If True, only estimate without migrating
+            
+        Returns:
+            MigrationResult with operation details
+        """
+        start_time = time.time()
+        
+        try:
+            self.source.connect()
+            
+            # Get estimate
+            estimate = self.source.estimate()
+            
+            if dry_run:
+                # For dry run, return early with estimate-based result
+                return MigrationResult(
+                    source_type=self.source.source_type,
+                    source_name=self.source.source_name,
+                    output_path=self.output_path,
+                    total_migrated=0,
+                    total_skipped=0,
+                    duration_seconds=time.time() - start_time,
+                    output_size_bytes=0
+                )
+            
+            # Create output .mv2 using MemoryStore
+            parent_dir = os.path.dirname(self.output_path)
+            if parent_dir:
+                os.makedirs(parent_dir, exist_ok=True)
+            
+            from ..store import MemoryStore
+            output_store = MemoryStore.create(self.output_path)
+            
+            total_migrated = 0
+            total_skipped = 0
+            errors = []
+            cursor = None
+            
+            try:
+                while True:
+                    try:
+                        records, next_cursor = self.source.fetch_batch(self.batch_size, cursor)
+                        
+                        if not records:
+                            break
+                        
+                        # Write records to .mv2 via MemoryStore
+                        for record in records:
+                            text = record.text or f"[vector-only record id={record.id}]"
+                            metadata = dict(record.metadata) if record.metadata else {}
+                            metadata['_source_id'] = record.id
+                            metadata['_source_type'] = self.source.source_type
+                            
+                            try:
+                                output_store.add(
+                                    text=text,
+                                    metadata=metadata
+                                )
+                                total_migrated += 1
+                            except Exception as e:
+                                total_skipped += 1
+                                errors.append(f"Record {record.id}: {str(e)}")
+                        
+                        # Update checkpoint
+                        if self.checkpoint_path:
+                            checkpoint = MigrationCheckpoint(
+                                source_type=self.source.source_type,
+                                source_name=self.source.source_name,
+                                output_path=self.output_path,
+                                last_cursor=cursor,
+                                records_completed=total_migrated,
+                                started_at=time.strftime('%Y-%m-%dT%H:%M:%SZ', time.gmtime(start_time))
+                            )
+                            self._write_checkpoint(checkpoint)
+                        
+                        # Call progress callbacks
+                        for callback in self.progress_callbacks:
+                            callback(total_migrated, estimate.total_vectors)
+                        
+                        cursor = next_cursor
+                        if cursor is None:
+                            break
+                            
+                    except Exception as e:
+                        errors.append(str(e))
+                        break
+            finally:
+                output_store.close()
+            
+            # Calculate output file size
+            output_size = os.path.getsize(self.output_path) if os.path.exists(self.output_path) else 0
+            
+            return MigrationResult(
+                source_type=self.source.source_type,
+                source_name=self.source.source_name,
+                output_path=self.output_path,
+                total_migrated=total_migrated,
+                total_skipped=total_skipped,
+                duration_seconds=time.time() - start_time,
+                output_size_bytes=output_size,
+                errors=errors
+            )
+            
+        except Exception as e:
+            return MigrationResult(
+                source_type=self.source.source_type,
+                source_name=self.source.source_name,
+                output_path=self.output_path,
+                total_migrated=0,
+                total_skipped=0,
+                duration_seconds=time.time() - start_time,
+                output_size_bytes=0,
+                errors=[str(e)]
+            )
+        finally:
+            try:
+                self.source.close()
+            except:
+                pass
+    
+    def resume(self) -> MigrationResult:
+        """Resume migration from checkpoint.
+        
+        Raises:
+            NotImplementedError: Resume is not yet implemented.
+        """
+        raise NotImplementedError(
+            "Resume not yet implemented — use run() to restart migration"
+        )
+    
+    def _write_checkpoint(self, checkpoint: MigrationCheckpoint) -> None:
+        """Write checkpoint to file."""
+        if not self.checkpoint_path:
+            return
+            
+        checkpoint_data = {
+            'source_type': checkpoint.source_type,
+            'source_name': checkpoint.source_name,
+            'output_path': checkpoint.output_path,
+            'last_cursor': checkpoint.last_cursor,
+            'records_completed': checkpoint.records_completed,
+            'started_at': checkpoint.started_at
+        }
+        
+        parent = os.path.dirname(self.checkpoint_path)
+        if parent:
+            os.makedirs(parent, exist_ok=True)
+        with open(self.checkpoint_path, 'w') as f:
+            json.dump(checkpoint_data, f)

ate/memory/migrate/sources/__init__.py

@@ -0,0 +1,13 @@
+"""Migration sources for various vector databases."""
+
+from .pinecone import PineconeMigrationSource
+from .qdrant import QdrantMigrationSource  
+from .weaviate import WeaviateMigrationSource
+from .chroma import ChromaMigrationSource
+
+__all__ = [
+    'PineconeMigrationSource',
+    'QdrantMigrationSource', 
+    'WeaviateMigrationSource',
+    'ChromaMigrationSource'
+]

ate/memory/migrate/sources/chroma.py

@@ -0,0 +1,170 @@
+"""Chroma migration source implementation."""
+from typing import Optional, List, Tuple
+
+try:
+    import chromadb
+except ImportError:
+    # Create a simple mock structure for testing
+    class ChromaDBMock:
+        PersistentClient = None
+        HttpClient = None
+        Client = None
+    
+    chromadb = ChromaDBMock()
+
+from ..base import MigrationSource, VectorRecord, MigrationEstimate
+
+
+class ChromaMigrationSource(MigrationSource):
+    """Migration source for Chroma vector database."""
+    
+    def __init__(self, path: Optional[str] = None, host: Optional[str] = None, port: Optional[int] = None, 
+                 collection_name: str = "default", api_key: Optional[str] = None):
+        """Initialize Chroma migration source.
+        
+        Args:
+            path: Path to persistent Chroma database (for PersistentClient)
+            host: Host for Chroma server (for HttpClient)
+            port: Port for Chroma server (for HttpClient)
+            collection_name: Name of the collection to migrate from
+            api_key: Optional API key for authentication
+        """
+        self.path = path
+        self.host = host
+        self.port = port
+        self.collection_name = collection_name
+        self.api_key = api_key
+        self.client = None
+        self.collection = None
+    
+    @property
+    def source_type(self) -> str:
+        """Return the source type."""
+        return "chroma"
+    
+    @property
+    def source_name(self) -> str:
+        """Return the source name."""
+        return self.collection_name
+    
+    def connect(self) -> None:
+        """Connect to Chroma."""
+        if chromadb is None:
+            raise ImportError("chromadb library is required for Chroma migration")
+        
+        # Choose client type based on configuration
+        if self.path:
+            # Use persistent client
+            self.client = chromadb.PersistentClient(path=self.path)
+        elif self.host and self.port:
+            # Use HTTP client
+            settings = {}
+            if self.api_key:
+                settings['chroma_api_impl'] = 'chromadb.api.fastapi.FastAPI'
+                settings['chroma_server_auth_credentials'] = self.api_key
+            
+            self.client = chromadb.HttpClient(host=self.host, port=self.port, settings=settings)
+        else:
+            # Use in-memory client
+            self.client = chromadb.Client()
+        
+        # Get or create collection
+        self.collection = self.client.get_collection(name=self.collection_name)
+    
+    def estimate(self) -> MigrationEstimate:
+        """Estimate migration size and time."""
+        if not self.collection:
+            raise RuntimeError("Must call connect() first")
+        
+        # Get collection count
+        total_vectors = self.collection.count()
+        
+        # Peek at one record to get dimensions
+        dimensions = 768  # Default assumption
+        if total_vectors > 0:
+            try:
+                peek_result = self.collection.peek(limit=1)
+                if peek_result.get('embeddings') and len(peek_result['embeddings']) > 0:
+                    dimensions = len(peek_result['embeddings'][0])
+            except Exception:
+                pass  # Use default
+        
+        # Rough estimates
+        bytes_per_vector = dimensions * 4 + 1024  # 4 bytes per float + metadata overhead
+        estimated_mv2_bytes = total_vectors * bytes_per_vector
+        estimated_seconds = total_vectors / 1200.0  # Rough estimate of 1200 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 Chroma."""
+        if not self.collection:
+            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
+        
+        # Get batch of records
+        result = self.collection.get(
+            limit=batch_size,
+            offset=offset,
+            include=['embeddings', 'metadatas', 'documents']
+        )
+        
+        ids = result.get('ids', [])
+        embeddings = result.get('embeddings', [])
+        metadatas = result.get('metadatas', [])
+        documents = result.get('documents', [])
+        
+        records = []
+        for i, record_id in enumerate(ids):
+            embedding = embeddings[i] if i < len(embeddings) else []
+            metadata = metadatas[i] if i < len(metadatas) else {}
+            document = documents[i] if i < len(documents) else None
+            
+            # Extract text from document or metadata
+            text = document
+            if not text and metadata:
+                text = metadata.get('text') or metadata.get('content')
+            
+            # Create clean metadata without text fields
+            clean_metadata = {}
+            if metadata:
+                clean_metadata = {k: v for k, v in metadata.items() if k not in ('text', 'content')}
+            
+            record = VectorRecord(
+                id=str(record_id),
+                vector=embedding,
+                text=text,
+                metadata=clean_metadata
+            )
+            records.append(record)
+        
+        # Determine if there are more results
+        # If we got any results and this is the first page, assume there might be more
+        # If we got fewer results than requested, we're at the end
+        if len(ids) == batch_size:
+            next_cursor = str(offset + batch_size)
+        elif len(ids) > 0 and offset == 0:
+            # First batch with some results - optimistically assume there might be more
+            next_cursor = str(offset + batch_size)
+        else:
+            next_cursor = None
+        
+        return records, next_cursor
+    
+    def close(self) -> None:
+        """Close the connection and clean up resources."""
+        # Chroma doesn't require explicit cleanup
+        self.client = None
+        self.collection = None

ate/memory/migrate/sources/pinecone.py

@@ -0,0 +1,120 @@
+"""Pinecone migration source implementation."""
+from typing import Optional, List, Tuple
+
+try:
+    import pinecone
+except ImportError:
+    pinecone = None
+
+from ..base import MigrationSource, VectorRecord, MigrationEstimate
+
+
+class PineconeMigrationSource(MigrationSource):
+    """Migration source for Pinecone vector database."""
+    
+    def __init__(self, api_key: str, index_name: str, environment: str, namespace: Optional[str] = None):
+        """Initialize Pinecone migration source.
+        
+        Args:
+            api_key: Pinecone API key
+            index_name: Name of the Pinecone index
+            environment: Pinecone environment
+            namespace: Optional namespace to migrate from
+        """
+        self.api_key = api_key
+        self.index_name = index_name
+        self.environment = environment
+        self.namespace = namespace
+        self.index = None
+    
+    @property
+    def source_type(self) -> str:
+        """Return the source type."""
+        return "pinecone"
+    
+    @property
+    def source_name(self) -> str:
+        """Return the source name."""
+        return self.index_name
+    
+    def connect(self) -> None:
+        """Connect to Pinecone."""
+        if pinecone is None:
+            raise ImportError("pinecone library is required for Pinecone migration")
+        
+        pinecone.init(api_key=self.api_key, environment=self.environment)
+        self.index = pinecone.Index(self.index_name)
+    
+    def estimate(self) -> MigrationEstimate:
+        """Estimate migration size and time."""
+        if not self.index:
+            raise RuntimeError("Must call connect() first")
+        
+        stats = self.index.describe_index_stats()
+        
+        # If namespace is specified, use namespace count, otherwise use total
+        if self.namespace and 'namespaces' in stats and self.namespace in stats['namespaces']:
+            total_vectors = stats['namespaces'][self.namespace]['vector_count']
+        else:
+            total_vectors = stats.get('total_vector_count', 0)
+        
+        dimensions = stats.get('dimension', 768)
+        
+        # Rough estimates
+        bytes_per_vector = dimensions * 4 + 1024  # 4 bytes per float + metadata overhead
+        estimated_mv2_bytes = total_vectors * bytes_per_vector
+        estimated_seconds = total_vectors / 1000.0  # Rough estimate of 1000 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 Pinecone."""
+        if not self.index:
+            raise RuntimeError("Must call connect() first")
+        
+        # List vector IDs with pagination
+        list_response = self.index.list(
+            namespace=self.namespace,
+            limit=batch_size,
+            pagination_token=cursor
+        )
+        
+        vector_ids = [vec['id'] for vec in list_response.get('vectors', [])]
+        
+        if not vector_ids:
+            return [], None
+        
+        # Fetch the actual vectors
+        fetch_response = self.index.fetch(vector_ids, namespace=self.namespace)
+        
+        records = []
+        for vector_id, vector_data in fetch_response.get('vectors', {}).items():
+            # Extract text from metadata if present
+            metadata = vector_data.get('metadata', {})
+            text = metadata.get('text')
+            
+            # Create clean metadata without text (since text has its own field)
+            clean_metadata = {k: v for k, v in metadata.items() if k != 'text'}
+            
+            record = VectorRecord(
+                id=vector_data['id'],
+                vector=vector_data['values'],
+                text=text,
+                metadata=clean_metadata
+            )
+            records.append(record)
+        
+        # Get next cursor from pagination
+        next_cursor = list_response.get('pagination', {}).get('next')
+        
+        return records, next_cursor
+    
+    def close(self) -> None:
+        """Close the connection and clean up resources."""
+        # Pinecone doesn't require explicit cleanup
+        self.index = None

ate/memory/migrate/sources/qdrant.py

@@ -0,0 +1,110 @@
+"""Qdrant migration source implementation."""
+from typing import Optional, List, Tuple
+
+try:
+    from qdrant_client import QdrantClient
+except ImportError:
+    QdrantClient = None
+
+from ..base import MigrationSource, VectorRecord, MigrationEstimate
+
+
+class QdrantMigrationSource(MigrationSource):
+    """Migration source for Qdrant vector database."""
+    
+    def __init__(self, url: str, collection_name: str, api_key: Optional[str] = None):
+        """Initialize Qdrant migration source.
+        
+        Args:
+            url: Qdrant server URL
+            collection_name: Name of the collection to migrate from
+            api_key: Optional API key for authentication
+        """
+        self.url = url
+        self.collection_name = collection_name
+        self.api_key = api_key
+        self.client = None
+    
+    @property
+    def source_type(self) -> str:
+        """Return the source type."""
+        return "qdrant"
+    
+    @property
+    def source_name(self) -> str:
+        """Return the source name."""
+        return self.collection_name
+    
+    def connect(self) -> None:
+        """Connect to Qdrant."""
+        if QdrantClient is None:
+            raise ImportError("qdrant-client library is required for Qdrant migration")
+        
+        self.client = QdrantClient(url=self.url, api_key=self.api_key)
+    
+    def estimate(self) -> MigrationEstimate:
+        """Estimate migration size and time."""
+        if not self.client:
+            raise RuntimeError("Must call connect() first")
+        
+        collection_info = self.client.get_collection(self.collection_name)
+        
+        total_vectors = collection_info.points_count
+        dimensions = collection_info.config.params.vectors.size
+        
+        # Rough estimates
+        bytes_per_vector = dimensions * 4 + 1024  # 4 bytes per float + metadata overhead
+        estimated_mv2_bytes = total_vectors * bytes_per_vector
+        estimated_seconds = total_vectors / 1000.0  # Rough estimate of 1000 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 Qdrant."""
+        if not self.client:
+            raise RuntimeError("Must call connect() first")
+        
+        # Use cursor directly as offset (Qdrant handles different offset types)
+        offset = cursor
+        
+        # Scroll through points
+        points, next_page_offset = self.client.scroll(
+            collection_name=self.collection_name,
+            limit=batch_size,
+            offset=offset,
+            with_payload=True,
+            with_vectors=True
+        )
+        
+        records = []
+        for point in points:
+            # Extract text from payload if present
+            payload = point.payload or {}
+            text = payload.get('text')
+            
+            # Create clean metadata without text
+            clean_metadata = {k: v for k, v in payload.items() if k != 'text'}
+            
+            record = VectorRecord(
+                id=str(point.id),
+                vector=point.vector,
+                text=text,
+                metadata=clean_metadata
+            )
+            records.append(record)
+        
+        # Convert next_page_offset to cursor string
+        next_cursor = str(next_page_offset) if next_page_offset is not None else None
+        
+        return records, next_cursor
+    
+    def close(self) -> None:
+        """Close the connection and clean up resources."""
+        if self.client:
+            self.client.close()
+            self.client = None