vector-inspector 0.2.6__py3-none-any.whl → 0.3.1__py3-none-any.whl

vector_inspector/core/connection_manager.py

@@ -0,0 +1,277 @@
+"""Connection manager for handling multiple vector database connections."""
+
+import uuid
+from typing import Dict, Optional, List, Any
+from enum import Enum
+from PySide6.QtCore import QObject, Signal
+
+from .connections.base_connection import VectorDBConnection
+
+
+class ConnectionState(Enum):
+    """Possible connection states."""
+    DISCONNECTED = "disconnected"
+    CONNECTING = "connecting"
+    CONNECTED = "connected"
+    ERROR = "error"
+
+
+class ConnectionInstance:
+    """Represents a single active connection with its state and context."""
+    
+    def __init__(
+        self,
+        connection_id: str,
+        name: str,
+        provider: str,
+        connection: VectorDBConnection,
+        config: Dict[str, Any]
+    ):
+        """
+        Initialize a connection instance.
+        
+        Args:
+            connection_id: Unique connection identifier
+            name: User-friendly connection name
+            provider: Provider type (chromadb, qdrant, etc.)
+            connection: The actual connection object
+            config: Connection configuration dict
+        """
+        self.id = connection_id
+        self.name = name
+        self.provider = provider
+        self.connection = connection
+        self.config = config
+        self.state = ConnectionState.DISCONNECTED
+        self.active_collection: Optional[str] = None
+        self.collections: List[str] = []
+        self.error_message: Optional[str] = None
+        
+    def get_display_name(self) -> str:
+        """Get a display-friendly connection name."""
+        return f"{self.name} ({self.provider})"
+    
+    def get_breadcrumb(self) -> str:
+        """Get breadcrumb showing connection > collection."""
+        if self.active_collection:
+            return f"{self.name} > {self.active_collection}"
+        return self.name
+
+
+class ConnectionManager(QObject):
+    """Manages multiple vector database connections and saved profiles.
+    
+    Signals:
+        connection_opened: Emitted when a new connection is opened (connection_id)
+        connection_closed: Emitted when a connection is closed (connection_id)
+        connection_state_changed: Emitted when connection state changes (connection_id, state)
+        active_connection_changed: Emitted when active connection changes (connection_id or None)
+        active_collection_changed: Emitted when active collection changes (connection_id, collection_name or None)
+        collections_updated: Emitted when collections list is updated (connection_id, collections)
+    """
+    
+    # Signals
+    connection_opened = Signal(str)  # connection_id
+    connection_closed = Signal(str)  # connection_id
+    connection_state_changed = Signal(str, ConnectionState)  # connection_id, state
+    active_connection_changed = Signal(object)  # connection_id or None
+    active_collection_changed = Signal(str, object)  # connection_id, collection_name or None
+    collections_updated = Signal(str, list)  # connection_id, collections
+    
+    MAX_CONNECTIONS = 10  # Limit to prevent resource exhaustion
+    
+    def __init__(self):
+        """Initialize the connection manager."""
+        super().__init__()
+        self._connections: Dict[str, ConnectionInstance] = {}
+        self._active_connection_id: Optional[str] = None
+        
+    def create_connection(
+        self,
+        name: str,
+        provider: str,
+        connection: VectorDBConnection,
+        config: Dict[str, Any]
+    ) -> str:
+        """
+        Create a new connection instance (not yet connected).
+        
+        Args:
+            name: User-friendly connection name
+            provider: Provider type
+            connection: The connection object
+            config: Connection configuration
+            
+        Returns:
+            The connection ID
+            
+        Raises:
+            RuntimeError: If maximum connections limit reached
+        """
+        if len(self._connections) >= self.MAX_CONNECTIONS:
+            raise RuntimeError(f"Maximum number of connections ({self.MAX_CONNECTIONS}) reached")
+        
+        connection_id = str(uuid.uuid4())
+        instance = ConnectionInstance(connection_id, name, provider, connection, config)
+        self._connections[connection_id] = instance
+        
+        # Set as active if it's the first connection
+        if len(self._connections) == 1:
+            self._active_connection_id = connection_id
+            self.active_connection_changed.emit(connection_id)
+        
+        # Don't emit connection_opened yet - wait until actually connected
+        return connection_id
+    
+    def mark_connection_opened(self, connection_id: str):
+        """
+        Mark a connection as opened (after successful connection).
+        
+        Args:
+            connection_id: ID of connection that opened
+        """
+        if connection_id in self._connections:
+            self.connection_opened.emit(connection_id)
+    
+    def get_connection(self, connection_id: str) -> Optional[ConnectionInstance]:
+        """Get a connection instance by ID."""
+        return self._connections.get(connection_id)
+    
+    def get_active_connection(self) -> Optional[ConnectionInstance]:
+        """Get the currently active connection instance."""
+        if self._active_connection_id:
+            return self._connections.get(self._active_connection_id)
+        return None
+    
+    def get_active_connection_id(self) -> Optional[str]:
+        """Get the currently active connection ID."""
+        return self._active_connection_id
+    
+    def set_active_connection(self, connection_id: str) -> bool:
+        """
+        Set the active connection.
+        
+        Args:
+            connection_id: ID of connection to make active
+            
+        Returns:
+            True if successful, False if connection not found
+        """
+        if connection_id not in self._connections:
+            return False
+        
+        self._active_connection_id = connection_id
+        self.active_connection_changed.emit(connection_id)
+        return True
+    
+    def close_connection(self, connection_id: str) -> bool:
+        """
+        Close and remove a connection.
+        
+        Args:
+            connection_id: ID of connection to close
+            
+        Returns:
+            True if successful, False if connection not found
+        """
+        instance = self._connections.get(connection_id)
+        if not instance:
+            return False
+        
+        # Disconnect the connection
+        try:
+            instance.connection.disconnect()
+        except Exception as e:
+            print(f"Error disconnecting: {e}")
+        
+        # Remove from connections dict
+        del self._connections[connection_id]
+        
+        # If this was the active connection, set a new one or None
+        if self._active_connection_id == connection_id:
+            if self._connections:
+                # Set first available connection as active
+                self._active_connection_id = next(iter(self._connections.keys()))
+                self.active_connection_changed.emit(self._active_connection_id)
+            else:
+                self._active_connection_id = None
+                self.active_connection_changed.emit(None)
+        
+        self.connection_closed.emit(connection_id)
+        return True
+    
+    def update_connection_state(self, connection_id: str, state: ConnectionState, error: Optional[str] = None):
+        """
+        Update the state of a connection.
+        
+        Args:
+            connection_id: ID of connection
+            state: New connection state
+            error: Optional error message if state is ERROR
+        """
+        instance = self._connections.get(connection_id)
+        if instance:
+            instance.state = state
+            if error:
+                instance.error_message = error
+            else:
+                instance.error_message = None
+            self.connection_state_changed.emit(connection_id, state)
+    
+    def update_collections(self, connection_id: str, collections: List[str]):
+        """
+        Update the collections list for a connection.
+        
+        Args:
+            connection_id: ID of connection
+            collections: List of collection names
+        """
+        instance = self._connections.get(connection_id)
+        if instance:
+            instance.collections = collections
+            self.collections_updated.emit(connection_id, collections)
+    
+    def set_active_collection(self, connection_id: str, collection_name: Optional[str]):
+        """
+        Set the active collection for a connection.
+        
+        Args:
+            connection_id: ID of connection
+            collection_name: Name of collection to make active, or None
+        """
+        instance = self._connections.get(connection_id)
+        if instance:
+            instance.active_collection = collection_name
+            self.active_collection_changed.emit(connection_id, collection_name)
+    
+    def get_all_connections(self) -> List[ConnectionInstance]:
+        """Get list of all connection instances."""
+        return list(self._connections.values())
+    
+    def get_connection_count(self) -> int:
+        """Get the number of active connections."""
+        return len(self._connections)
+    
+    def close_all_connections(self):
+        """Close all connections. Typically called on application exit."""
+        connection_ids = list(self._connections.keys())
+        for conn_id in connection_ids:
+            self.close_connection(conn_id)
+    
+    def rename_connection(self, connection_id: str, new_name: str) -> bool:
+        """
+        Rename a connection.
+        
+        Args:
+            connection_id: ID of connection
+            new_name: New name for the connection
+            
+        Returns:
+            True if successful, False if connection not found
+        """
+        instance = self._connections.get(connection_id)
+        if instance:
+            instance.name = new_name
+            return True
+        return False
+

vector_inspector/core/connections/__init__.py

@@ -3,5 +3,6 @@
 from .base_connection import VectorDBConnection
 from .chroma_connection import ChromaDBConnection
 from .qdrant_connection import QdrantConnection
+from .pinecone_connection import PineconeConnection
 
-__all__ = ["VectorDBConnection", "ChromaDBConnection", "QdrantConnection"]
+__all__ = ["VectorDBConnection", "ChromaDBConnection", "QdrantConnection", "PineconeConnection"]

vector_inspector/core/connections/base_connection.py

@@ -229,5 +229,46 @@ class VectorDBConnection(ABC):
             {"name": "in", "server_side": True},
             {"name": "not in", "server_side": True},
             {"name": "contains", "server_side": False},
-            {"name": "not contains", "server_side": False},
         ]
+    
+    def get_embedding_model(self, collection_name: str, connection_id: Optional[str] = None) -> Optional[str]:
+        """
+        Get the embedding model used for a collection.
+        
+        Retrieves the model name from:
+        1. Collection-level metadata (if supported)
+        2. Vector metadata (_embedding_model field)
+        3. User settings (for collections we can't modify)
+        
+        Args:
+            collection_name: Name of collection
+            connection_id: Optional connection ID for settings lookup
+            
+        Returns:
+            Model name string (e.g., "sentence-transformers/all-MiniLM-L6-v2") or None
+        """
+        try:
+            # First try to get from collection-level metadata
+            info = self.get_collection_info(collection_name)
+            if info and info.get("embedding_model"):
+                return info["embedding_model"]
+            
+            # Fall back to checking a sample vector's metadata
+            data = self.get_all_items(collection_name, limit=1, offset=0)
+            if data and data.get("metadatas") and len(data["metadatas"]) > 0:
+                metadata = data["metadatas"][0]
+                if "_embedding_model" in metadata:
+                    return metadata["_embedding_model"]
+            
+            # Finally, check user settings (for collections we can't modify)
+            if connection_id:
+                from ...services.settings_service import SettingsService
+                settings = SettingsService()
+                model_info = settings.get_embedding_model(connection_id, collection_name)
+                if model_info:
+                    return model_info["model"]
+            
+            return None
+        except Exception as e:
+            print(f"Failed to get embedding model: {e}")
+            return None

vector_inspector/core/connections/chroma_connection.py

@@ -6,10 +6,44 @@ from pathlib import Path
 import chromadb
 from chromadb.api import ClientAPI
 from chromadb.api.models.Collection import Collection
+from chromadb import Documents, EmbeddingFunction, Embeddings
 
 from .base_connection import VectorDBConnection
 
 
+class DimensionAwareEmbeddingFunction(EmbeddingFunction):
+    """Embedding function that selects model based on collection's expected dimension."""
+    
+    def __init__(self, expected_dimension: int):
+        """Initialize with expected dimension (model loaded lazily on first use)."""
+        self.expected_dimension = expected_dimension
+        self.model = None
+        self.model_name = None
+        self.model_type = None
+        self._initialized = False
+    
+    def _ensure_model_loaded(self):
+        """Lazy load the embedding model on first use."""
+        if self._initialized:
+            return
+        
+        from ..embedding_utils import get_embedding_model_for_dimension
+        print(f"[ChromaDB] Loading embedding model for {self.expected_dimension}d vectors...")
+        self.model, self.model_name, self.model_type = get_embedding_model_for_dimension(self.expected_dimension)
+        print(f"[ChromaDB] Using {self.model_type} model '{self.model_name}' for {self.expected_dimension}d embeddings")
+        self._initialized = True
+    
+    def __call__(self, input: Documents) -> Embeddings:
+        """Embed documents using the dimension-appropriate model."""
+        self._ensure_model_loaded()
+        from ..embedding_utils import encode_text
+        embeddings = []
+        for text in input:
+            embedding = encode_text(text, self.model, self.model_type)
+            embeddings.append(embedding)
+        return embeddings
+
+
 class ChromaDBConnection(VectorDBConnection):
     """Manages connection to ChromaDB and provides query interface."""
     
@@ -90,12 +124,47 @@ class ChromaDBConnection(VectorDBConnection):
             print(f"Failed to list collections: {e}")
             return []
     
-    def get_collection(self, name: str) -> Optional[Collection]:
-        """
-        Get or create a collection.
+    def _get_collection_basic(self, name: str) -> Optional[Collection]:
+        """Get collection without custom embedding function (for info lookup)."""
+        if not self._client:
+            return None
+        try:
+            return self._client.get_collection(name=name)
+        except Exception as e:
+            return None
+    
+    def _get_embedding_function_for_collection(self, name: str) -> Optional[EmbeddingFunction]:
+        """Get the appropriate embedding function for a collection based on its dimension."""
+        # Get basic collection to check dimension
+        basic_col = self._get_collection_basic(name)
+        if not basic_col:
+            return None
+        
+        try:
+            # Get a sample to determine vector dimension
+            sample = basic_col.get(limit=1, include=["embeddings"])
+            embeddings = sample.get("embeddings") if sample else None
+            # Avoid numpy array truthiness issues - check is not None explicitly
+            if embeddings is not None and len(embeddings) > 0:
+                first_embedding = embeddings[0]
+                # Check if embedding exists and has content
+                if first_embedding is not None and len(first_embedding) > 0:
+                    vector_dim = len(first_embedding)
+                    print(f"[ChromaDB] Collection '{name}' has {vector_dim}d vectors")
+                    return DimensionAwareEmbeddingFunction(vector_dim)
+        except Exception as e:
+            print(f"[ChromaDB] Failed to determine embedding function: {e}")
+            import traceback
+            traceback.print_exc()
+        
+        return None
+    
+    def get_collection(self, name: str, embedding_function: Optional[EmbeddingFunction] = None) -> Optional[Collection]:
+        """Get a collection (without overriding existing embedding function).
         
         Args:
             name: Collection name
+            embedding_function: Optional custom embedding function (ignored if collection exists)
             
         Returns:
             Collection object or None if failed
@@ -103,7 +172,9 @@ class ChromaDBConnection(VectorDBConnection):
         if not self._client:
             return None
         try:
-            self._current_collection = self._client.get_or_create_collection(name=name)
+            # Just get the collection without trying to override embedding function
+            # This avoids conflicts with existing collections
+            self._current_collection = self._client.get_collection(name=name)
             return self._current_collection
         except Exception as e:
             print(f"Failed to get collection: {e}")
@@ -119,7 +190,7 @@ class ChromaDBConnection(VectorDBConnection):
         Returns:
             Dictionary with collection info
         """
-        collection = self.get_collection(name)
+        collection = self._get_collection_basic(name)
         if not collection:
             return None
         
@@ -141,27 +212,37 @@ class ChromaDBConnection(VectorDBConnection):
             # ChromaDB uses cosine distance by default (or can be configured)
             # Try to get metadata from collection if available
             distance_metric = "Cosine (default)"
+            embedding_model = None
             try:
                 # ChromaDB collections may have metadata about distance function
                 col_metadata = collection.metadata
-                if col_metadata and "hnsw:space" in col_metadata:
-                    space = col_metadata["hnsw:space"]
-                    if space == "l2":
-                        distance_metric = "Euclidean (L2)"
-                    elif space == "ip":
-                        distance_metric = "Inner Product"
-                    elif space == "cosine":
-                        distance_metric = "Cosine"
+                if col_metadata:
+                    if "hnsw:space" in col_metadata:
+                        space = col_metadata["hnsw:space"]
+                        if space == "l2":
+                            distance_metric = "Euclidean (L2)"
+                        elif space == "ip":
+                            distance_metric = "Inner Product"
+                        elif space == "cosine":
+                            distance_metric = "Cosine"
+                    # Get embedding model if stored
+                    if "embedding_model" in col_metadata:
+                        embedding_model = col_metadata["embedding_model"]
             except:
                 pass  # Use default if unable to determine
             
-            return {
+            result = {
                 "name": name,
                 "count": count,
                 "metadata_fields": metadata_fields,
                 "vector_dimension": vector_dimension,
                 "distance_metric": distance_metric,
             }
+            
+            if embedding_model:
+                result["embedding_model"] = embedding_model
+            
+            return result
         except Exception as e:
             print(f"Failed to get collection info: {e}")
             return None
@@ -189,10 +270,22 @@ class ChromaDBConnection(VectorDBConnection):
         Returns:
             Query results or None if failed
         """
+        print(f"[ChromaDB] query_collection called for '{collection_name}'")
         collection = self.get_collection(collection_name)
         if not collection:
+            print(f"[ChromaDB] Failed to get collection '{collection_name}'")
             return None
         
+        # If query_texts provided, we need to manually embed them with dimension-aware model
+        if query_texts and not query_embeddings:
+            embedding_function = self._get_embedding_function_for_collection(collection_name)
+            if embedding_function:
+                print(f"[ChromaDB] Manually embedding query texts with dimension-aware model")
+                query_embeddings = embedding_function(query_texts)
+                query_texts = None  # Use embeddings instead of texts
+            else:
+                print(f"[ChromaDB] Warning: Could not determine embedding function, using collection's default")
+        
         try:
             results = collection.query(
                 query_texts=query_texts,
@@ -205,6 +298,8 @@ class ChromaDBConnection(VectorDBConnection):
             return cast(Dict[str, Any], results)
         except Exception as e:
             print(f"Query failed: {e}")
+            import traceback
+            traceback.print_exc()
             return None
     
     def get_all_items(
@@ -368,8 +463,34 @@ class ChromaDBConnection(VectorDBConnection):
 
     # Implement base connection uniform APIs
     def create_collection(self, name: str, vector_size: int, distance: str = "Cosine") -> bool:
-        """Create a collection. Chroma doesn't require vector size at creation."""
-        return self.get_collection(name) is not None
+        """Create a collection. If it doesn't exist, attempt to create it using Chroma client APIs."""
+        if not self._client:
+            return False
+
+        try:
+            # Prefer get_or_create_collection if available
+            if hasattr(self._client, "get_or_create_collection"):
+                col = self._client.get_or_create_collection(name=name)
+                self._current_collection = col
+                return True
+
+            # Fallback to create_collection/create and then fetch
+            if hasattr(self._client, "create_collection"):
+                try:
+                    self._client.create_collection(name=name)
+                except Exception:
+                    # Some clients may raise if already exists; ignore
+                    pass
+                col = self._client.get_collection(name=name)
+                self._current_collection = col
+                return col is not None
+
+            # As a last resort, check if collection exists
+            col = self.get_collection(name)
+            return col is not None
+        except Exception as e:
+            print(f"Failed to create collection: {e}")
+            return False
 
     def get_items(self, name: str, ids: List[str]) -> Dict[str, Any]:
         """Retrieve items by IDs."""