Feat/qdrant vdb (#1649)

* feat: Qdrant vector search support

Signed-off-by: Anush008 <anushshetty90@gmail.com>

* fix: modify env

* fix: fix the old version problem

* fix: For older versions

* perf: minor perf

---------

Signed-off-by: Anush008 <anushshetty90@gmail.com>
Co-authored-by: Anush008 <anushshetty90@gmail.com>
Co-authored-by: Junyan Qin <rockchinq@gmail.com>
This commit is contained in:
Guanchao Wang
2025-09-12 12:41:16 +08:00
committed by GitHub
parent 345c8b113f
commit 6f98feaaf1
6 changed files with 142 additions and 17 deletions
+9 -9
View File
@@ -24,23 +24,23 @@ class Retriever(base_service.BaseService):
extra_args={}, # TODO: add extra args extra_args={}, # TODO: add extra args
) )
chroma_results = await self.ap.vector_db_mgr.vector_db.search(kb_id, query_embedding[0], k) vector_results = await self.ap.vector_db_mgr.vector_db.search(kb_id, query_embedding[0], k)
# 'ids' is always returned by ChromaDB, even if not explicitly in 'include' # 'ids' shape mirrors the Chroma-style response contract for compatibility
matched_chroma_ids = chroma_results.get('ids', [[]])[0] matched_vector_ids = vector_results.get('ids', [[]])[0]
distances = chroma_results.get('distances', [[]])[0] distances = vector_results.get('distances', [[]])[0]
chroma_metadatas = chroma_results.get('metadatas', [[]])[0] vector_metadatas = vector_results.get('metadatas', [[]])[0]
if not matched_chroma_ids: if not matched_vector_ids:
self.ap.logger.info('No relevant chunks found in Chroma.') self.ap.logger.info('No relevant chunks found in vector database.')
return [] return []
result: list[retriever_entities.RetrieveResultEntry] = [] result: list[retriever_entities.RetrieveResultEntry] = []
for i, id in enumerate(matched_chroma_ids): for i, id in enumerate(matched_vector_ids):
entry = retriever_entities.RetrieveResultEntry( entry = retriever_entities.RetrieveResultEntry(
id=id, id=id,
metadata=chroma_metadatas[i], metadata=vector_metadatas[i],
distance=distances[i], distance=distances[i],
) )
result.append(entry) result.append(entry)
+14 -2
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
from ..core import app from ..core import app
from .vdb import VectorDatabase from .vdb import VectorDatabase
from .vdbs.chroma import ChromaVectorDatabase from .vdbs.chroma import ChromaVectorDatabase
from .vdbs.qdrant import QdrantVectorDatabase
class VectorDBManager: class VectorDBManager:
@@ -13,6 +14,17 @@ class VectorDBManager:
self.ap = ap self.ap = ap
async def initialize(self): async def initialize(self):
# 初始化 Chroma 向量数据库(可扩展为多种实现) kb_config = self.ap.instance_config.data.get('vdb')
if self.vector_db is None: if kb_config:
if kb_config.get('use') == 'chroma':
self.vector_db = ChromaVectorDatabase(self.ap)
self.ap.logger.info('Initialized Chroma vector database backend.')
elif kb_config.get('use') == 'qdrant':
self.vector_db = QdrantVectorDatabase(self.ap)
self.ap.logger.info('Initialized Qdrant vector database backend.')
else:
self.vector_db = ChromaVectorDatabase(self.ap)
self.ap.logger.warning('No valid vector database backend configured, defaulting to Chroma.')
else:
self.vector_db = ChromaVectorDatabase(self.ap) self.vector_db = ChromaVectorDatabase(self.ap)
self.ap.logger.warning('No vector database backend configured, defaulting to Chroma.')
+5 -4
View File
@@ -14,24 +14,25 @@ class VectorDatabase(abc.ABC):
metadatas: list[dict[str, Any]], metadatas: list[dict[str, Any]],
documents: list[str], documents: list[str],
) -> None: ) -> None:
"""向指定 collection 添加向量数据。""" """Add vector data to the specified collection."""
pass pass
@abc.abstractmethod @abc.abstractmethod
async def search(self, collection: str, query_embedding: np.ndarray, k: int = 5) -> Dict[str, Any]: async def search(self, collection: str, query_embedding: np.ndarray, k: int = 5) -> Dict[str, Any]:
"""在指定 collection 中检索最相似的向量。""" """Search for the most similar vectors in the specified collection."""
pass pass
@abc.abstractmethod @abc.abstractmethod
async def delete_by_file_id(self, collection: str, file_id: str) -> None: async def delete_by_file_id(self, collection: str, file_id: str) -> None:
"""根据 file_id 删除指定 collection 中的向量。""" """Delete vectors from the specified collection by file_id."""
pass pass
@abc.abstractmethod @abc.abstractmethod
async def get_or_create_collection(self, collection: str): async def get_or_create_collection(self, collection: str):
"""获取或创建 collection""" """Get or create collection."""
pass pass
@abc.abstractmethod @abc.abstractmethod
async def delete_collection(self, collection: str): async def delete_collection(self, collection: str):
"""Delete collection."""
pass pass
+104
View File
@@ -0,0 +1,104 @@
from __future__ import annotations
from typing import Any, Dict, List
from qdrant_client import AsyncQdrantClient, models
from pkg.core import app
from pkg.vector.vdb import VectorDatabase
class QdrantVectorDatabase(VectorDatabase):
def __init__(self, ap: app.Application):
self.ap = ap
url = self.ap.instance_config.data['vdb']['qdrant']['url']
host = self.ap.instance_config.data['vdb']['qdrant']['host']
port = self.ap.instance_config.data['vdb']['qdrant']['port']
api_key = self.ap.instance_config.data['vdb']['qdrant']['api_key']
if url:
self.client = AsyncQdrantClient(url=url, api_key=api_key)
else:
self.client = AsyncQdrantClient(host=host, port=int(port), api_key=api_key)
self._collections: set[str] = set()
async def _ensure_collection(self, collection: str, vector_size: int) -> None:
if collection in self._collections:
return
exists = await self.client.collection_exists(collection)
if exists:
self._collections.add(collection)
return
await self.client.create_collection(
collection_name=collection,
vectors_config=models.VectorParams(size=vector_size, distance=models.Distance.COSINE),
)
self._collections.add(collection)
self.ap.logger.info(f"Qdrant collection '{collection}' created with dim={vector_size}.")
async def get_or_create_collection(self, collection: str):
# Qdrant requires vector size to create a collection; no-op here.
pass
async def add_embeddings(
self,
collection: str,
ids: List[str],
embeddings_list: List[List[float]],
metadatas: List[Dict[str, Any]],
) -> None:
if not embeddings_list:
return
await self._ensure_collection(collection, len(embeddings_list[0]))
points = [
models.PointStruct(id=ids[i], vector=embeddings_list[i], payload=metadatas[i]) for i in range(len(ids))
]
await self.client.upsert(collection_name=collection, points=points)
self.ap.logger.info(f"Added {len(ids)} embeddings to Qdrant collection '{collection}'.")
async def search(self, collection: str, query_embedding: list[float], k: int = 5) -> dict[str, Any]:
exists = await self.client.collection_exists(collection)
if not exists:
return {'ids': [[]], 'metadatas': [[]], 'distances': [[]]}
hits = (
await self.client.query_points(
collection_name=collection,
query=query_embedding,
limit=k,
with_payload=True,
)
).points
ids = [str(hit.id) for hit in hits]
metadatas = [hit.payload or {} for hit in hits]
# Qdrant's score is similarity; convert to a pseudo-distance for consistency
distances = [1 - float(hit.score) if hit.score is not None else 1.0 for hit in hits]
results = {'ids': [ids], 'metadatas': [metadatas], 'distances': [distances]}
self.ap.logger.info(f"Qdrant search in '{collection}' returned {len(results.get('ids', [[]])[0])} results.")
return results
async def delete_by_file_id(self, collection: str, file_id: str) -> None:
exists = await self.client.collection_exists(collection)
if not exists:
return
await self.client.delete(
collection_name=collection,
points_selector=models.Filter(
must=[models.FieldCondition(key='file_id', match=models.MatchValue(value=file_id))]
),
)
self.ap.logger.info(f"Deleted embeddings from Qdrant collection '{collection}' with file_id: {file_id}")
async def delete_collection(self, collection: str):
try:
await self.client.delete_collection(collection)
self._collections.discard(collection)
self.ap.logger.info(f"Qdrant collection '{collection}' deleted.")
except Exception:
self.ap.logger.warning(f"Qdrant collection '{collection}' not found.")
+3 -2
View File
@@ -1,9 +1,9 @@
[project] [project]
name = "langbot" name = "langbot"
version = "4.2.2" version = "4.2.2"
description = "高稳定、支持扩展、多模态 - 大模型原生即时通信机器人平台" description = "Easy-to-use global IM bot platform designed for LLM era"
readme = "README.md" readme = "README.md"
requires-python = ">=3.10.1" requires-python = ">=3.10.1,<4.0"
dependencies = [ dependencies = [
"aiocqhttp>=1.4.4", "aiocqhttp>=1.4.4",
"aiofiles>=24.1.0", "aiofiles>=24.1.0",
@@ -60,6 +60,7 @@ dependencies = [
"html2text>=2024.2.26", "html2text>=2024.2.26",
"langchain>=0.2.0", "langchain>=0.2.0",
"chromadb>=0.4.24", "chromadb>=0.4.24",
"qdrant-client (>=1.15.1,<2.0.0)",
] ]
keywords = [ keywords = [
"bot", "bot",
+7
View File
@@ -20,3 +20,10 @@ system:
jwt: jwt:
expire: 604800 expire: 604800
secret: '' secret: ''
vdb:
use: chroma
qdrant:
url: ''
host: localhost
port: 6333
api_key: ''