From f9df3ddfd25e882c09d82960b771e41ef61e5f3c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9C=D0=B0=D1=80=D0=B0=D1=82=20=D0=A4=D0=B0=D0=B7=D1=8B?= =?UTF-8?q?=D0=BB=D0=BE=D0=B2?= Date: Tue, 12 May 2026 11:49:45 +0000 Subject: [PATCH] Add vector store module with Qdrant + Ollama embeddings --- vector_store.py | 90 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 vector_store.py diff --git a/vector_store.py b/vector_store.py new file mode 100644 index 0000000..546d666 --- /dev/null +++ b/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)