From 12dcb962e5ba28df0a201be0da9ba98f2f5dce60 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9C=D0=B0=D1=80=D0=B0=D1=82=20=D0=A4=D0=B0=D0=B7=D1=8B?= =?UTF-8?q?=D0=BB=D0=BE=D0=B2?= Date: Fri, 15 May 2026 08:48:50 +0000 Subject: [PATCH] add file --- qdrant_client.py | 32 +++++++++----------------------- 1 file changed, 9 insertions(+), 23 deletions(-) diff --git a/qdrant_client.py b/qdrant_client.py index b48730c..aa9081f 100644 --- a/qdrant_client.py +++ b/qdrant_client.py @@ -1,36 +1,22 @@ -from qdrant_client import QdrantClient -from qdrant_client.http import models +from chromadb import Client as ChromaClient from langchain_ollama import OllamaEmbeddings -class QdrantStore: - def __init__(self, url="http://localhost:6333", collection_name="rag_collection"): - self.client = QdrantClient(url=url) +class ChromaStore: + def __init__(self, collection_name="rag_collection"): + self.client = ChromaClient() 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), - ) + if self.collection_name not in [c.name for c in self.client.get_collections()]: + self.client.create_collection(name=self.collection_name, metadata={}) 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))]) - ) + self.client.upsert(collection_name=self.collection_name, documents=documents, ids=[str(i) for i in range(len(documents))], metadatas=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] + results = self.client.query(collection_name=self.collection_name, query_text=query, n_results=limit, include_metadata=True) + return [r['metadata'] for r in results]