rakam-systems-vectorstore 0.1.1rc7__py3-none-any.whl

rakam_systems_vectorstore/components/vectorstore/pg_vector_store.py

@@ -0,0 +1,827 @@
+import os
+import time
+from functools import lru_cache
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+import dotenv
+import numpy as np
+from django.contrib.postgres.search import SearchQuery
+from django.contrib.postgres.search import SearchRank
+from django.contrib.postgres.search import SearchVector
+from django.db import connection
+from django.db import transaction
+from openai import OpenAI
+from sentence_transformers import SentenceTransformer
+
+from rakam_systems_core.ai_utils import logging
+from rakam_systems_core.ai_core.interfaces.vectorstore import VectorStore
+from rakam_systems_vectorstore.components.vectorstore.pg_models import Collection
+from rakam_systems_vectorstore.components.vectorstore.pg_models import NodeEntry
+from rakam_systems_vectorstore.core import Node
+from rakam_systems_vectorstore.core import NodeMetadata
+from rakam_systems_vectorstore.core import VSFile
+
+logger = logging.getLogger(__name__)
+
+# Load environment variables
+dotenv.load_dotenv()
+api_key = os.getenv("OPENAI_API_KEY")
+
+
+class PgVectorStore(VectorStore):
+    """
+    A class for managing collection-based vector stores using pgvector and Django ORM.
+    Enhanced for better semantic search performance with hybrid search, re-ranking, and caching.
+
+    Note: Vector columns are created without dimension constraints, allowing flexibility
+    to use different embedding models without needing to alter the database schema.
+    """
+
+    def __init__(
+        self,
+        name: str = "pg_vector_store",
+        config=None,
+        embedding_model: str = "Snowflake/snowflake-arctic-embed-m",
+        use_embedding_api: bool = False,
+        api_model: str = "text-embedding-3-small",
+    ) -> None:
+        """
+        Initializes the PgVectorStore with the specified embedding model.
+
+        :param name: Name of the vector store component.
+        :param config: Configuration object.
+        :param embedding_model: Pre-trained SentenceTransformer model name.
+        :param use_embedding_api: Whether to use OpenAI's embedding API instead of local model.
+        :param api_model: OpenAI API model to use for embeddings if use_embedding_api is True.
+        """
+        super().__init__(name=name, config=config)
+        self._ensure_pgvector_extension()
+        self.use_embedding_api = use_embedding_api
+
+        if self.use_embedding_api:
+            self.client = OpenAI(api_key=api_key)
+            self.api_model = api_model
+            sample_embedding = self._get_api_embedding("Sample text")
+            self.embedding_dim = len(sample_embedding)
+        else:
+            self.embedding_model = SentenceTransformer(
+                embedding_model, trust_remote_code=True
+            )
+            self.embedding_dim = self.embedding_model.get_sentence_embedding_dimension()
+
+        logger.info(
+            f"Initialized PgVectorStore with embedding dimension: {self.embedding_dim}"
+        )
+
+    def _ensure_pgvector_extension(self) -> None:
+        """
+        Ensures that the pgvector extension is installed in the PostgreSQL database.
+        """
+        with connection.cursor() as cursor:
+            try:
+                cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;")
+                logger.info("Ensured pgvector extension is installed")
+            except Exception as e:
+                logger.error(f"Failed to create pgvector extension: {e}")
+                raise
+
+    def _get_api_embedding(self, text: str) -> List[float]:
+        """
+        Gets embedding from OpenAI API.
+
+        :param text: Text to embed
+        :return: Embedding vector
+        """
+        try:
+            response = self.client.embeddings.create(
+                input=[text], model=self.api_model)
+            return response.data[0].embedding
+        except Exception as e:
+            logger.error(f"Failed to get API embedding: {e}")
+            raise
+
+    @lru_cache(maxsize=1000)
+    def predict_embeddings(self, query: str) -> np.ndarray:
+        """
+        Predicts embeddings for a given query using the embedding model.
+        Caches results to reduce redundant computations.
+
+        :param query: Query string to encode.
+        :return: Normalized embedding vector for the query.
+        """
+        logger.debug(f"Predicting embeddings for query: {query}")
+        start_time = time.time()
+
+        if self.use_embedding_api:
+            query_embedding = self._get_api_embedding(query)
+            query_embedding = np.array(query_embedding, dtype="float32")
+        else:
+            query_embedding = self.embedding_model.encode(query)
+            query_embedding = np.array(query_embedding, dtype="float32")
+
+        # Normalize embedding for cosine similarity
+        norm = np.linalg.norm(query_embedding)
+        if norm > 0:
+            query_embedding = query_embedding / norm
+        else:
+            logger.warning(f"Zero norm encountered for query: {query}")
+
+        logger.debug(
+            f"Embedding generation took {time.time() - start_time:.2f} seconds"
+        )
+        return query_embedding
+
+    def get_embeddings(
+        self, sentences: List[str], parallel: bool = True, batch_size: int = 8
+    ) -> np.ndarray:
+        """
+        Generates embeddings for a list of sentences with normalization.
+
+        :param sentences: List of sentences to encode.
+        :param parallel: Whether to use parallel processing (default is True).
+        :param batch_size: Batch size for processing (default is 8).
+        :return: Normalized embedding vectors for the sentences.
+        """
+        logger.info(f"Generating embeddings for {len(sentences)} sentences")
+        start = time.time()
+
+        if self.use_embedding_api:
+            all_embeddings = []
+            for i in range(0, len(sentences), batch_size):
+                batch = sentences[i: i + batch_size]
+                response = self.client.embeddings.create(
+                    input=batch, model=self.api_model
+                )
+                batch_embeddings = [data.embedding for data in response.data]
+                all_embeddings.extend(batch_embeddings)
+            embeddings = np.array(all_embeddings, dtype="float32")
+        else:
+            if parallel:
+                os.environ["TOKENIZERS_PARALLELISM"] = "false"
+                pool = self.embedding_model.start_multi_process_pool(
+                    target_devices=["cpu"] * 5
+                )
+                embeddings = self.embedding_model.encode_multi_process(
+                    sentences, pool, batch_size=batch_size
+                )
+                self.embedding_model.stop_multi_process_pool(pool)
+            else:
+                os.environ["TOKENIZERS_PARALLELISM"] = "true"
+                embeddings = self.embedding_model.encode(
+                    sentences,
+                    batch_size=batch_size,
+                    show_progress_bar=True,
+                    convert_to_tensor=True,
+                )
+                embeddings = embeddings.cpu().detach().numpy()
+
+        # Normalize embeddings
+        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
+        norms[norms == 0] = 1  # Avoid division by zero
+        embeddings = embeddings / norms
+
+        logger.info(
+            f"Time taken to encode {len(sentences)} items: {time.time() - start:.2f} seconds"
+        )
+        return embeddings
+
+    def get_or_create_collection(self, collection_name: str) -> Collection:
+        """
+        Gets or creates a collection with the specified name.
+
+        :param collection_name: Name of the collection.
+        :return: Collection object.
+        """
+        collection, created = Collection.objects.get_or_create(
+            name=collection_name, defaults={
+                "embedding_dim": self.embedding_dim}
+        )
+        logger.info(
+            f"{'Created new' if created else 'Using existing'} collection: {collection_name}"
+        )
+        return collection
+
+    def _rerank_results(
+        self,
+        query: str,
+        results: List[Tuple[Dict, str, float]],
+        suggested_nodes: List[Node],
+        top_k: int,
+    ) -> Tuple[Dict, List[Node]]:
+        """
+        Re-ranks search results using a combination of vector similarity and keyword relevance.
+
+        :param query: The search query.
+        :param results: Initial search results (metadata, content, distance).
+        :param suggested_nodes: List of Node objects.
+        :param top_k: Number of results to return after re-ranking.
+        :return: Tuple of re-ranked results dictionary and updated suggested_nodes.
+        """
+        logger.debug(f"Re-ranking {len(results)} results for query: {query}")
+
+        # Perform full-text search to get keyword relevance scores
+        search_query = SearchQuery(query, config="english")
+        queryset = NodeEntry.objects.filter(
+            collection__name="document_collection",
+            node_id__in=[int(res[0]["node_id"]) for res in results],
+        ).annotate(
+            rank=SearchRank(SearchVector(
+                "content", config="english"), search_query)
+        )
+
+        # Combine vector distance and keyword rank
+        reranked_results = []
+        node_id_to_rank = {node.node_id: node.rank for node in queryset}
+        for metadata, content, distance in results:
+            node_id = metadata["node_id"]
+            keyword_score = node_id_to_rank.get(node_id, 0.0)
+            # Combine scores (adjust weights as needed)
+            combined_score = 0.7 * (1 - distance) + 0.3 * keyword_score
+            reranked_results.append((metadata, content, combined_score))
+
+        # Sort by combined score and take top_k
+        reranked_results = sorted(reranked_results, key=lambda x: x[2], reverse=True)[
+            :top_k
+        ]
+        valid_suggestions = {
+            str(res[0]["node_id"]): res for res in reranked_results}
+
+        # Update suggested_nodes to match re-ranked order
+        node_id_order = [res[0]["node_id"] for res in reranked_results]
+        updated_nodes = sorted(
+            suggested_nodes,
+            key=lambda node: node_id_order.index(node.metadata.node_id)
+            if node.metadata.node_id in node_id_order
+            else len(node_id_order),
+        )[:top_k]
+
+        logger.debug(f"Re-ranked to {len(valid_suggestions)} results")
+        return valid_suggestions, updated_nodes
+
+    def search(
+        self,
+        collection_name: str,
+        query: str,
+        distance_type: str = "cosine",
+        number: int = 5,
+        meta_data_filters: Optional[Dict[str, Any]] = None,
+        hybrid_search: bool = True,
+    ) -> Tuple[Dict, List[Node]]:
+        """
+        Retrieve relevant documents from the vector store using hybrid search and re-ranking.
+
+        :param collection_name: Name of the collection to search.
+        :param query: Search query.
+        :param distance_type: Distance metric ("cosine", "l2", "dot").
+        :param number: Number of results to return.
+        :param meta_data_filters: Dictionary of metadata filters (e.g., {"is_validated": True}).
+        :param hybrid_search: Whether to use hybrid search combining vector and keyword search.
+        :return: Tuple of search results (dictionary) and suggested nodes.
+        """
+        logger.info(
+            f"Searching in collection: {collection_name} for query: '{query}'")
+
+        try:
+            collection = Collection.objects.get(name=collection_name)
+        except Collection.DoesNotExist:
+            logger.error(f"No collection found with name: {collection_name}")
+            raise ValueError(
+                f"No collection found with name: {collection_name}")
+
+        # Generate query embedding
+        query_embedding = self.predict_embeddings(query)
+
+        # Build base queryset
+        queryset = NodeEntry.objects.filter(collection=collection)
+
+        # Apply metadata filters
+        if meta_data_filters:
+            for key, value in meta_data_filters.items():
+                queryset = queryset.filter(
+                    **{f"custom_metadata__{key}": value})
+
+        # Construct SQL query for vector search
+        if distance_type == "cosine":
+            distance_operator = "<=>"
+        elif distance_type == "l2":
+            distance_operator = "<->"
+        elif distance_type == "dot":
+            distance_operator = "<#>"
+        else:
+            logger.error(f"Unsupported distance type: {distance_type}")
+            raise ValueError(f"Unsupported distance type: {distance_type}")
+
+        # Request more results for hybrid search and re-ranking
+        search_buffer_factor = 2 if hybrid_search else 1
+        limit = number * search_buffer_factor
+        embedding_str = "[" + ",".join(str(x) for x in query_embedding) + "]"
+
+        sql_query = f"""
+            SELECT
+                node_id,
+                content,
+                source_file_uuid,
+                position,
+                custom_metadata,
+                embedding {distance_operator} %s::vector AS distance
+            FROM
+                {NodeEntry._meta.db_table}
+            WHERE
+                collection_id = %s
+            ORDER BY
+                distance
+            LIMIT
+                %s
+        """
+
+        # Execute vector search
+        with connection.cursor() as cursor:
+            cursor.execute(sql_query, [embedding_str, collection.id, limit])
+            results = cursor.fetchall()
+            columns = [col[0] for col in cursor.description]
+
+        # Process vector search results
+        valid_suggestions = {}
+        suggested_nodes = []
+        seen_texts = set()
+
+        for row in results:
+            result_dict = dict(zip(columns, row))
+            node_id = result_dict["node_id"]
+            content = result_dict["content"]
+            distance = result_dict["distance"]
+
+            if content not in seen_texts:
+                seen_texts.add(content)
+                custom_metadata = result_dict["custom_metadata"] or {}
+                if isinstance(custom_metadata, str):
+                    try:
+                        import json
+
+                        custom_metadata = json.loads(custom_metadata)
+                    except (json.JSONDecodeError, TypeError):
+                        custom_metadata = {}
+
+                metadata = NodeMetadata(
+                    source_file_uuid=result_dict["source_file_uuid"],
+                    position=result_dict["position"],
+                    custom=custom_metadata,
+                )
+                metadata.node_id = node_id
+                node = Node(content=content, metadata=metadata)
+                node.embedding = result_dict.get("embedding")
+                suggested_nodes.append(node)
+
+                valid_suggestions[str(node_id)] = (
+                    {
+                        "node_id": node_id,
+                        "source_file_uuid": result_dict["source_file_uuid"],
+                        "position": result_dict["position"],
+                        "custom": custom_metadata,
+                    },
+                    content,
+                    float(distance),
+                )
+
+        # Perform hybrid search and re-ranking if enabled
+        if hybrid_search:
+            valid_suggestions, suggested_nodes = self._rerank_results(
+                query, list(valid_suggestions.values()
+                            ), suggested_nodes, number
+            )
+
+        logger.info(f"Search returned {len(valid_suggestions)} results")
+        return valid_suggestions, suggested_nodes
+
+    @transaction.atomic
+    def create_collection_from_files(
+        self, collection_name: str, files: List[VSFile]
+    ) -> None:
+        """
+        Creates a collection from a list of VSFile objects.
+
+        :param collection_name: Name of the collection to create.
+        :param files: List of VSFile objects containing nodes.
+        """
+        logger.info(f"Creating collection: {collection_name} from files")
+        nodes = [node for file in files for node in file.nodes]
+        self.create_collection_from_nodes(collection_name, nodes)
+
+    @transaction.atomic
+    def create_collection_from_nodes(
+        self, collection_name: str, nodes: List[Node]
+    ) -> None:
+        """
+        Creates a collection from a list of nodes.
+
+        :param collection_name: Name of the collection to create.
+        :param nodes: List of Node objects.
+        """
+        if not nodes:
+            logger.warning(
+                f"Cannot create collection '{collection_name}' because nodes list is empty"
+            )
+            return
+
+        # Filter out nodes with None or empty content (these would cause embedding errors)
+        original_count = len(nodes)
+        nodes = [node for node in nodes if node.content is not None and str(
+            node.content).strip()]
+
+        if len(nodes) < original_count:
+            logger.warning(
+                f"Filtered out {original_count - len(nodes)} nodes with empty/None content")
+
+        if not nodes:
+            logger.warning(
+                f"No valid nodes for collection '{collection_name}' after filtering")
+            return
+
+        total_nodes = len(nodes)
+        logger.info(
+            f"Creating collection: {collection_name} with {total_nodes} nodes")
+
+        start_time = time.time()
+        collection = self.get_or_create_collection(collection_name)
+        NodeEntry.objects.filter(collection=collection).delete()
+
+        # Generate embeddings
+        embed_start = time.time()
+        text_chunks = [str(node.content) for node in nodes]
+        embeddings = self.get_embeddings(text_chunks, parallel=False)
+        embed_time = time.time() - embed_start
+        logger.info(
+            f"Embeddings generated in {embed_time:.2f}s ({total_nodes/embed_time:.0f} nodes/s)")
+
+        # Prepare node entries
+        prep_start = time.time()
+        node_entries = [
+            NodeEntry(
+                collection=collection,
+                content=node.content,
+                embedding=embeddings[i].tolist(),
+                source_file_uuid=node.metadata.source_file_uuid,
+                position=node.metadata.position,
+                custom_metadata=node.metadata.custom or {},
+            )
+            for i, node in enumerate(nodes)
+        ]
+        prep_time = time.time() - prep_start
+        logger.info(f"Node entries prepared in {prep_time:.2f}s")
+
+        # Bulk insert
+        insert_start = time.time()
+        created_entries = NodeEntry.objects.bulk_create(node_entries)
+        insert_time = time.time() - insert_start
+        logger.info(
+            f"Bulk insert completed in {insert_time:.2f}s ({total_nodes/insert_time:.0f} nodes/s)")
+
+        for i, node in enumerate(nodes):
+            node.metadata.node_id = created_entries[i].node_id
+
+        total_time = time.time() - start_time
+        logger.info(
+            f"Created collection '{collection_name}' with {len(created_entries)} nodes in {total_time:.2f}s"
+        )
+
+    @transaction.atomic
+    def add_nodes(self, collection_name: str, nodes: List[Node]) -> None:
+        """
+        Adds nodes to an existing collection.
+
+        :param collection_name: Name of the collection to update.
+        :param nodes: List of Node objects to be added.
+        """
+        if not nodes:
+            logger.warning("No nodes to add")
+            return
+
+        # Filter out nodes with None or empty content (these would cause embedding errors)
+        original_count = len(nodes)
+        nodes = [node for node in nodes if node.content is not None and str(
+            node.content).strip()]
+
+        if len(nodes) < original_count:
+            logger.warning(
+                f"Filtered out {original_count - len(nodes)} nodes with empty/None content")
+
+        if not nodes:
+            logger.warning("No valid nodes to add after filtering")
+            return
+
+        total_nodes = len(nodes)
+        logger.info(
+            f"Adding {total_nodes} nodes to collection: {collection_name}")
+
+        start_time = time.time()
+        try:
+            collection = Collection.objects.get(name=collection_name)
+        except Collection.DoesNotExist:
+            raise ValueError(
+                f"No collection found with name: {collection_name}")
+
+        # Generate embeddings
+        embed_start = time.time()
+        text_chunks = [str(node.content) for node in nodes]
+        embeddings = self.get_embeddings(text_chunks, parallel=False)
+        embed_time = time.time() - embed_start
+        logger.info(
+            f"Embeddings generated in {embed_time:.2f}s ({total_nodes/embed_time:.0f} nodes/s)")
+
+        # Prepare node entries
+        prep_start = time.time()
+        node_entries = [
+            NodeEntry(
+                collection=collection,
+                content=node.content,
+                embedding=embeddings[i].tolist(),
+                source_file_uuid=node.metadata.source_file_uuid,
+                position=node.metadata.position,
+                custom_metadata=node.metadata.custom or {},
+            )
+            for i, node in enumerate(nodes)
+        ]
+        prep_time = time.time() - prep_start
+        logger.info(f"Node entries prepared in {prep_time:.2f}s")
+
+        # Bulk insert
+        insert_start = time.time()
+        created_entries = NodeEntry.objects.bulk_create(node_entries)
+        insert_time = time.time() - insert_start
+        logger.info(
+            f"Bulk insert completed in {insert_time:.2f}s ({total_nodes/insert_time:.0f} nodes/s)")
+
+        for i, node in enumerate(nodes):
+            node.metadata.node_id = created_entries[i].node_id
+
+        total_time = time.time() - start_time
+        logger.info(
+            f"Added {len(created_entries)} nodes to collection '{collection_name}' in {total_time:.2f}s"
+        )
+
+    @transaction.atomic
+    def delete_nodes(self, collection_name: str, node_ids: List[int]) -> None:
+        """
+        Deletes nodes from an existing collection.
+
+        :param collection_name: Name of the collection to update.
+        :param node_ids: List of node IDs to be deleted.
+        """
+        if not node_ids:
+            logger.warning("No node IDs to delete")
+            return
+
+        logger.info(
+            f"Deleting {len(node_ids)} nodes from collection: {collection_name}"
+        )
+        try:
+            collection = Collection.objects.get(name=collection_name)
+        except Collection.DoesNotExist:
+            raise ValueError(
+                f"No collection found with name: {collection_name}")
+
+        existing_ids = set(
+            NodeEntry.objects.filter(
+                collection=collection, node_id__in=node_ids
+            ).values_list("node_id", flat=True)
+        )
+        missing_ids = set(node_ids) - existing_ids
+        if missing_ids:
+            logger.warning(
+                f"Node ID(s) {missing_ids} not found in collection {collection_name}"
+            )
+
+        deleted_count, _ = NodeEntry.objects.filter(
+            collection=collection, node_id__in=existing_ids
+        ).delete()
+        logger.info(
+            f"Deleted {deleted_count} nodes from collection '{collection_name}'"
+        )
+
+    @transaction.atomic
+    def add_files(self, collection_name: str, files: List[VSFile]) -> None:
+        """
+        Adds file nodes to the specified collection.
+
+        :param collection_name: Name of the collection to update.
+        :param files: List of VSFile objects whose nodes are to be added.
+        """
+        logger.info(f"Adding files to collection: {collection_name}")
+        all_nodes = [node for file in files for node in file.nodes]
+        self.add_nodes(collection_name, all_nodes)
+
+    @transaction.atomic
+    def delete_files(self, collection_name: str, files: List[VSFile]) -> None:
+        """
+        Deletes file nodes from the specified collection.
+
+        :param collection_name: Name of the collection to update.
+        :param files: List of VSFile objects whose nodes are to be deleted.
+        """
+        logger.info(f"Deleting files from collection: {collection_name}")
+        node_ids_to_delete = [
+            node.metadata.node_id
+            for file in files
+            for node in file.nodes
+            if node.metadata.node_id
+        ]
+        if node_ids_to_delete:
+            self.delete_nodes(collection_name, node_ids_to_delete)
+        else:
+            logger.warning("No node IDs found in provided files")
+
+    def list_collections(self) -> List[str]:
+        """
+        Lists all available collections.
+
+        :return: List of collection names.
+        """
+        return list(Collection.objects.values_list("name", flat=True))
+
+    def get_collection_info(self, collection_name: str) -> Dict[str, Any]:
+        """
+        Gets information about a collection.
+
+        :param collection_name: Name of the collection.
+        :return: Dictionary containing collection information.
+        """
+        try:
+            collection = Collection.objects.get(name=collection_name)
+        except Collection.DoesNotExist:
+            raise ValueError(
+                f"No collection found with name: {collection_name}")
+
+        node_count = NodeEntry.objects.filter(collection=collection).count()
+        return {
+            "name": collection.name,
+            "embedding_dim": collection.embedding_dim,
+            "node_count": node_count,
+            "created_at": collection.created_at,
+            "updated_at": collection.updated_at,
+        }
+
+    @transaction.atomic
+    def delete_collection(self, collection_name: str) -> None:
+        """
+        Deletes a collection and all its nodes.
+
+        :param collection_name: Name of the collection to delete.
+        """
+        try:
+            collection = Collection.objects.get(name=collection_name)
+        except Collection.DoesNotExist:
+            raise ValueError(
+                f"No collection found with name: {collection_name}")
+
+        node_count = NodeEntry.objects.filter(hourly=collection).count()
+        collection.delete()
+        logger.info(
+            f"Deleted collection '{collection_name}' with {node_count} nodes")
+
+    # VectorStore interface methods
+    def add(self, vectors: List[List[float]], metadatas: List[Dict[str, Any]]) -> Any:
+        """
+        Adds vectors with metadata to the default collection.
+        This method implements the VectorStore interface.
+
+        :param vectors: List of embedding vectors to add.
+        :param metadatas: List of metadata dictionaries for each vector.
+        :return: List of node IDs that were created.
+        """
+        if not vectors or not metadatas:
+            logger.warning("Empty vectors or metadatas provided to add()")
+            return []
+
+        if len(vectors) != len(metadatas):
+            raise ValueError(
+                "Number of vectors must match number of metadatas")
+
+        # Get or create default collection
+        collection_name = metadatas[0].get(
+            "collection_name", "default_collection")
+        collection = self.get_or_create_collection(collection_name)
+
+        # Create nodes from vectors and metadatas
+        node_entries = []
+        for i, (vector, metadata) in enumerate(zip(vectors, metadatas)):
+            content = metadata.get("content", "")
+            source_file_uuid = metadata.get("source_file_uuid", "")
+            position = metadata.get("position", i)
+            custom_metadata = {
+                k: v
+                for k, v in metadata.items()
+                if k not in ["content", "source_file_uuid", "position", "collection_name"]
+            }
+
+            node_entries.append(
+                NodeEntry(
+                    collection=collection,
+                    content=content,
+                    embedding=vector,
+                    source_file_uuid=source_file_uuid,
+                    position=position,
+                    custom_metadata=custom_metadata,
+                )
+            )
+
+        created_entries = NodeEntry.objects.bulk_create(node_entries)
+        node_ids = [entry.node_id for entry in created_entries]
+        logger.info(
+            f"Added {len(node_ids)} vectors to collection '{collection_name}'")
+        return node_ids
+
+    def query(
+        self, vector: List[float], top_k: int = 5, **kwargs
+    ) -> List[Dict[str, Any]]:
+        """
+        Queries the vector store for similar vectors.
+        This method implements the VectorStore interface.
+
+        :param vector: Query vector.
+        :param top_k: Number of results to return.
+        :param kwargs: Additional parameters (collection_name, distance_type, meta_data_filters).
+        :return: List of dictionaries containing search results.
+        """
+        collection_name = kwargs.get("collection_name", "default_collection")
+        distance_type = kwargs.get("distance_type", "cosine")
+        meta_data_filters = kwargs.get("meta_data_filters")
+
+        try:
+            collection = Collection.objects.get(name=collection_name)
+        except Collection.DoesNotExist:
+            logger.warning(f"Collection '{collection_name}' not found")
+            return []
+
+        # Convert vector to numpy array
+        query_embedding = np.array(vector, dtype="float32")
+
+        # Normalize if using cosine distance
+        if distance_type == "cosine":
+            norm = np.linalg.norm(query_embedding)
+            if norm > 0:
+                query_embedding = query_embedding / norm
+
+        # Build queryset
+        queryset = NodeEntry.objects.filter(collection=collection)
+
+        # Apply metadata filters
+        if meta_data_filters:
+            for key, value in meta_data_filters.items():
+                queryset = queryset.filter(
+                    **{f"custom_metadata__{key}": value})
+
+        # Determine distance operator
+        if distance_type == "cosine":
+            distance_operator = "<=>"
+        elif distance_type == "l2":
+            distance_operator = "<->"
+        elif distance_type == "dot":
+            distance_operator = "<#>"
+        else:
+            raise ValueError(f"Unsupported distance type: {distance_type}")
+
+        embedding_str = "[" + ",".join(str(x) for x in query_embedding) + "]"
+
+        sql_query = f"""
+            SELECT
+                node_id,
+                content,
+                source_file_uuid,
+                position,
+                custom_metadata,
+                embedding {distance_operator} %s::vector AS distance
+            FROM
+                {NodeEntry._meta.db_table}
+            WHERE
+                collection_id = %s
+            ORDER BY
+                distance
+            LIMIT
+                %s
+        """
+
+        # Execute query
+        with connection.cursor() as cursor:
+            cursor.execute(sql_query, [embedding_str, collection.id, top_k])
+            results = cursor.fetchall()
+            columns = [col[0] for col in cursor.description]
+
+        # Format results
+        formatted_results = []
+        for row in results:
+            result_dict = dict(zip(columns, row))
+            formatted_results.append({
+                "node_id": result_dict["node_id"],
+                "content": result_dict["content"],
+                "source_file_uuid": result_dict["source_file_uuid"],
+                "position": result_dict["position"],
+                "metadata": result_dict["custom_metadata"] or {},
+                "distance": float(result_dict["distance"]),
+            })
+
+        logger.info(f"Query returned {len(formatted_results)} results")
+        return formatted_results