diff --git a/vector_store.py b/vector_store.py new file mode 100644 index 0000000..e671ba6 --- /dev/null +++ b/vector_store.py @@ -0,0 +1,40 @@ +from qdrant_client import QdrantClient +from qdrant_client.http import models as _models +from langchain_ollama import OllamaEmbeddings +import uuid +from typing import List, Tuple + +class VectorStore: + def __init__(self, collection_name: str = "rag_collection", host: str = "localhost", port: int = 6333): + self.client = QdrantClient(host=host, port=port) + self.collection_name = collection_name + existing = [c.name for c in self.client.get_collections().collections] + if collection_name not in existing: + self.client.recreate_collection( + collection_name=collection_name, + vectors_config={"distance": "Cosine"}, + ) + self.embeddings = OllamaEmbeddings(model="nomic-embed-text") + + def add_documents(self, documents: List[str]) -> None: + vectors = self.embeddings.embed_documents(documents) + points = [] + for text, vector in zip(documents, vectors): + points.append( + _models.PointStruct( + id=uuid.uuid4().hex, + vector=vector, + payload={"text": text}, + ) + ) + self.client.upsert(collection_name=self.collection_name, points=points) + + def search(self, query: str, max_results: int = 5) -> List[Tuple[str, float]]: + query_vector = self.embeddings.embed_query(query) + results = self.client.search( + collection_name=self.collection_name, + query_vector=query_vector, + limit=max_results, + with_payload=True, + ) + return [(p.payload["text"], p.score) for p in results]