Add Qdrant and ChromaDB support to the project
- Added Qdrant service to both docker-compose files for production and development. - Updated environment variables in .env.example and settings.py to include Qdrant configuration. - Included necessary dependencies for Qdrant and ChromaDB in requirements.txt. - Updated .gitignore to exclude ChromaDB data files.
This commit is contained in:
@@ -0,0 +1,117 @@
|
||||
"""
|
||||
Qdrant Vector Store — ذخیره و جستجوی وکتورها
|
||||
"""
|
||||
from qdrant_client import QdrantClient
|
||||
from qdrant_client.http import models as qmodels
|
||||
|
||||
from .client import get_qdrant_client
|
||||
from .config import load_rag_config, RAGConfig
|
||||
|
||||
|
||||
class QdrantVectorStore:
|
||||
"""
|
||||
ذخیره و جستجوی documents در Qdrant.
|
||||
"""
|
||||
|
||||
def __init__(self, config: RAGConfig | None = None):
|
||||
self.config = config or load_rag_config()
|
||||
self.qdrant = self.config.qdrant
|
||||
self._client: QdrantClient | None = None
|
||||
|
||||
@property
|
||||
def client(self) -> QdrantClient:
|
||||
if self._client is None:
|
||||
self._client = get_qdrant_client(self.qdrant)
|
||||
return self._client
|
||||
|
||||
def ensure_collection(self, recreate: bool = False) -> None:
|
||||
"""
|
||||
اطمینان از وجود collection با نام و اندازه مناسب.
|
||||
"""
|
||||
name = self.qdrant.collection_name
|
||||
size = self.qdrant.vector_size
|
||||
|
||||
try:
|
||||
self.client.get_collection(name)
|
||||
if recreate:
|
||||
self.client.delete_collection(name)
|
||||
self.client.create_collection(
|
||||
collection_name=name,
|
||||
vectors_config=qmodels.VectorParams(
|
||||
size=size,
|
||||
distance=qmodels.Distance.COSINE,
|
||||
),
|
||||
)
|
||||
except Exception:
|
||||
self.client.create_collection(
|
||||
collection_name=name,
|
||||
vectors_config=qmodels.VectorParams(
|
||||
size=size,
|
||||
distance=qmodels.Distance.COSINE,
|
||||
),
|
||||
)
|
||||
|
||||
def add_documents(
|
||||
self,
|
||||
ids: list[str],
|
||||
embeddings: list[list[float]],
|
||||
documents: list[str],
|
||||
metadatas: list[dict] | None = None,
|
||||
) -> int:
|
||||
"""
|
||||
افزودن documents به collection.
|
||||
metadata فقط str, int, float, bool پشتیبانی میشود.
|
||||
"""
|
||||
self.ensure_collection()
|
||||
metas = metadatas or [{}] * len(ids)
|
||||
|
||||
def _serialize(m: dict) -> dict:
|
||||
out = {}
|
||||
for k, v in m.items():
|
||||
if v is None:
|
||||
continue
|
||||
if isinstance(v, (str, int, float, bool)):
|
||||
out[k] = v
|
||||
else:
|
||||
out[k] = str(v)
|
||||
return out
|
||||
|
||||
payloads = [
|
||||
{"text": doc, "doc_id": sid, **_serialize(m)}
|
||||
for doc, m, sid in zip(documents, metas, ids)
|
||||
]
|
||||
|
||||
self.client.upsert(
|
||||
collection_name=self.qdrant.collection_name,
|
||||
points=[
|
||||
qmodels.PointStruct(id=pid, vector=emb, payload=pl)
|
||||
for pid, emb, pl in zip(ids, embeddings, payloads)
|
||||
],
|
||||
)
|
||||
return len(ids)
|
||||
|
||||
def search(
|
||||
self,
|
||||
query_vector: list[float],
|
||||
limit: int = 5,
|
||||
score_threshold: float | None = None,
|
||||
) -> list[dict]:
|
||||
"""
|
||||
جستجوی شباهت بر اساس query vector.
|
||||
"""
|
||||
results = self.client.search(
|
||||
collection_name=self.qdrant.collection_name,
|
||||
query_vector=query_vector,
|
||||
limit=limit,
|
||||
score_threshold=score_threshold,
|
||||
)
|
||||
|
||||
return [
|
||||
{
|
||||
"id": str(r.id),
|
||||
"score": r.score,
|
||||
"text": r.payload.get("text", ""),
|
||||
"metadata": {k: v for k, v in r.payload.items() if k != "text"},
|
||||
}
|
||||
for r in results
|
||||
]
|
||||
Reference in New Issue
Block a user