Add vector store module with Qdrant + Ollama embeddings

vector_store.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+"""Vector store module for RAG agent using Qdrant and Ollama embeddings."""
+
+import uuid
+from typing import List, Optionl
+
+from langchain_community.embeddings import OllamaEmbeddings
+from langchain_community.vectorstores import Qdrant
+from langchain_core.documents import Document
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from qdrant_client import QdrantClient
+from qdrant_client.models import Distance, VectorParams
+
+COLLECTION_NAME = "rag_knowledge_base"
+OLLAMA_BASE_URL = "http://localhost:11434"
+EMBED_MODEL = "nomic-embed-text"
+EMBED_DIMENSION = 768
+
+def get_embeddings() -> OllamaEmbeddings:
+    """Create and return Ollama embeddings instance."""
+    return OllamaEmbeddings(
+        model=EMBED_MODEL
+        base_url=OLLAMA_BASE_URL,
+    )
+
+
+def get_qdrant_client() -> QdrantClient:
+    """Create and return Qdrant client (in-memory for local dev)."""
+    return QdrantClient(":memory:")
+
+
+def create_collection(client: QdrantClient) -> None:
+    """Create Qdrant collection if doesn't exist."""
+    existing = [c.name for c in client.get_collections().collections]
+    if COLLECTION_NAME not in existing:
+        client.create_collection(
+            collection_name=COLLECTION_NAME,
+            vectors_config=VectorParams(
+                size=EMBED_DIMENSION,
+                distance=Distance.COSINE,
+            ),
+        )
+
+
+def get_vector_store(client: Optional[QdrantClient] = None) -> Qdrant:
+    """Initialize and return Qdrant vector store."""
+    if client is None:
+        client = get_qdrant_client()
+    create_collection(client)
+    embeddings = get_embeddings()
+    return Qdrant(
+        client=client,
+        collection_name=COLLECTION_NAME,
+        embeddings=embeddings,
+    )
+
+
+def chunk_documents(
+    documents: List[Document],
+    chunk_size: int = 512,
+    chunk_overlap: int = 50,
+) -> List[Document]:
+    """Split documents into chunks using RecursiveCharacterTextSplitter."""
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size=chunk_size,
+        chunk_overlap=chunk_overlap,
+        length_function=len,
+        is_separator_regex=False,
+    )
+    return splitter.split_documents(documents)
+
+
+def add_documents_to_store(
+    vector_store: Qdrant,
+    documents: List[Document],
+) -> List[str]:
+    """Add documents to vector store, return list of IDs."""
+    chunks = chunk_documents(documents)
+    ids = [str(uuid.uuid4()) for _ in chunks]
+    vector_store.add_documents(documents=chunks, ids=ids)
+    return ids
+
+
+def search_store(
+    vector_store: Qdrant,
+    query: str,
+    max_results: int = 5,
+) -> List[Document]:
+    """Search vector store for relevant documents."""
+    return vector_store.similarity_search(query, k=max_results)