sajad-dev
197f70ee12
- 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.
118 lines
3.5 KiB
Python
118 lines
3.5 KiB
Python
"""
|
|
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
|
|
]
|