Files
agent-s-rag-pamyatyu/src/vector_store.py
T
2026-07-01 14:05:59 +03:00

42 lines
1.4 KiB
Python

"""
Vector store implementation using Qdrant.
"""
from qdrant_client import QdrantClient
from qdrant_client.http import models as qdrant_models
from langchain.vectorstores import Qdrant
from config import QDRANT_HOST, QDRANT_PORT, QDRANT_COLLECTION_NAME
from embeddings import get_ollama_embeddings
def get_qdrant_client() -> QdrantClient:
"""
Creates a Qdrant client connected to the local Qdrant instance.
"""
return QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
def ensure_collection(client: QdrantClient, collection_name: str, vector_size: int = 768) -> None:
"""
Ensures that the specified collection exists in Qdrant.
If it does not exist, it will be created with the given vector size.
"""
if not client.has_collection(collection_name):
client.recreate_collection(
collection_name=collection_name,
vectors_config=qdrant_models.VectorParams(
size=vector_size,
distance="Cosine"
)
)
def get_vector_store() -> Qdrant:
"""
Returns a Qdrant vector store instance ready for use with LangChain.
"""
client = get_qdrant_client()
ensure_collection(client, QDRANT_COLLECTION_NAME)
embeddings = get_ollama_embeddings()
return Qdrant(
client=client,
collection_name=QDRANT_COLLECTION_NAME,
embeddings=embeddings
)