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

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

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