Working retrieval with the cli

services/rag/langchain/vector_storage.py

@@ -4,7 +4,7 @@ import os
 from typing import Optional
 
 from dotenv import load_dotenv
-from langchain_community.vectorstores import Qdrant
+from langchain_qdrant import QdrantVectorStore
 from langchain_core.documents import Document
 from langchain_ollama import OllamaEmbeddings
 from qdrant_client import QdrantClient
@@ -23,7 +23,7 @@ OLLAMA_EMBEDDING_MODEL = os.getenv("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text")
 
 def initialize_vector_store(
     collection_name: str = "documents_langchain", recreate_collection: bool = False
-) -> Qdrant:
+) -> QdrantVectorStore:
     """
     Initialize and return a Qdrant vector store with Ollama embeddings.
 
@@ -34,19 +34,18 @@ def initialize_vector_store(
     Returns:
         Initialized Qdrant vector store
     """
-    # Initialize Qdrant client
-    client = QdrantClient(
-        host=QDRANT_HOST,
-        port=QDRANT_REST_PORT,
-    )
-
     # Initialize Ollama embeddings
     embeddings = OllamaEmbeddings(
         model=OLLAMA_EMBEDDING_MODEL,
         base_url="http://localhost:11434",  # Default Ollama URL
     )
 
-    # Check if collection exists, if not create it
+    # Check if collection exists and create if needed
+    client = QdrantClient(
+        host=QDRANT_HOST,
+        port=QDRANT_REST_PORT,
+    )
+
     collection_exists = False
     try:
         client.get_collection(collection_name)
@@ -63,7 +62,6 @@ def initialize_vector_store(
     if not collection_exists:
         # Create collection using the Qdrant client directly
         from qdrant_client.http.models import Distance, VectorParams
-        import numpy as np
 
         # First, we need to determine the embedding size by creating a sample embedding
         sample_embedding = embeddings.embed_query("sample text for dimension detection")
@@ -75,25 +73,18 @@ def initialize_vector_store(
             vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE),
         )
 
-        # Now create the Qdrant instance connected to the newly created collection
-        vector_store = Qdrant(
-            client=client,
-            collection_name=collection_name,
-            embeddings=embeddings,
-        )
-    else:
-        # Collection exists, just connect to it
-        vector_store = Qdrant(
-            client=client,
-            collection_name=collection_name,
-            embeddings=embeddings,
-        )
+    # Create the Qdrant instance connected to the collection
+    vector_store = QdrantVectorStore(
+        client=client,
+        collection_name=collection_name,
+        embedding=embeddings,
+    )
 
     return vector_store
 
 
 def add_documents_to_vector_store(
-    vector_store: Qdrant, documents: list[Document], batch_size: int = 10
+    vector_store: QdrantVectorStore, documents: list[Document], batch_size: int = 10
 ) -> None:
     """
     Add documents to the vector store.
@@ -109,7 +100,7 @@ def add_documents_to_vector_store(
         vector_store.add_documents(batch)
 
 
-def search_vector_store(vector_store: Qdrant, query: str, top_k: int = 5) -> list:
+def search_vector_store(vector_store: QdrantVectorStore, query: str, top_k: int = 5) -> list:
     """
     Search the vector store for similar documents.
 
@@ -138,7 +129,7 @@ OPENROUTER_EMBEDDING_MODEL = os.getenv("OPENROUTER_EMBEDDING_MODEL", "openai/tex
 
 def initialize_vector_store_with_openrouter(
     collection_name: str = "documents_langchain"
-) -> Qdrant:
+) -> QdrantVectorStore:
     # Initialize Qdrant client
     client = QdrantClient(
         host=QDRANT_HOST,
@@ -153,7 +144,7 @@ def initialize_vector_store_with_openrouter(
     )
 
     # Create or get the vector store
-    vector_store = Qdrant(
+    vector_store = QdrantVectorStore(
         client=client,
         collection_name=collection_name,
         embeddings=embeddings,