Add vector_store.py

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]