41 lines
1.6 KiB
Python
41 lines
1.6 KiB
Python
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]
|