91 lines
2.7 KiB
Python
91 lines
2.7 KiB
Python
#!/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)
|