add qdrant_client
This commit is contained in:
@@ -0,0 +1,36 @@
|
|||||||
|
from qdrant_client import QdrantClient
|
||||||
|
from qdrant_client.http import models
|
||||||
|
from langchain_ollama import OllamaEmbeddings
|
||||||
|
|
||||||
|
class QdrantStore:
|
||||||
|
def __init__(self, url="http://localhost:6333", collection_name="rag_collection"):
|
||||||
|
self.client = QdrantClient(url=url)
|
||||||
|
self.collection_name = collection_name
|
||||||
|
self._ensure_collection()
|
||||||
|
|
||||||
|
def _ensure_collection(self):
|
||||||
|
if self.collection_name not in self.client.get_collections().collections:
|
||||||
|
self.client.recreate_collection(
|
||||||
|
collection_name=self.collection_name,
|
||||||
|
vectors_config=models.VectorParams(size=384, distance=models.Distance.COSINE),
|
||||||
|
)
|
||||||
|
|
||||||
|
def add_documents(self, documents, titles):
|
||||||
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||||
|
vectors = embeddings.embed_documents(documents)
|
||||||
|
payload = [{"title": t} for t in titles]
|
||||||
|
self.client.upsert(
|
||||||
|
collection_name=self.collection_name,
|
||||||
|
points=models.Batch(points=[models.PointStruct(id=i, vector=v, payload=p) for i, (v, p) in enumerate(zip(vectors, payload))])
|
||||||
|
)
|
||||||
|
|
||||||
|
def search(self, query, limit=5):
|
||||||
|
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||||
|
query_vector = embeddings.embed_query(query)
|
||||||
|
results = self.client.search(
|
||||||
|
collection_name=self.collection_name,
|
||||||
|
query_vector=query_vector,
|
||||||
|
limit=limit,
|
||||||
|
with_payload=True,
|
||||||
|
)
|
||||||
|
return [r.payload for r in results]
|
||||||
Reference in New Issue
Block a user