310 lines
15 KiB
Python
310 lines
15 KiB
Python
|
"""Qdrant REST adapter preserving the application's original chunk IDs and metadata.
|
|||
|
|
|
|||
|
|
The small collection interface matches the store's existing snapshot/rollback code.
|
|||
|
|
REST calls have deadlines; backend errors never fall back to a stale local index.
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
import hashlib
|
|||
|
|
import math
|
|||
|
|
import threading
|
|||
|
|
import uuid
|
|||
|
|
import httpx
|
|||
|
|
|
|||
|
|
|
|||
|
|
class QdrantError(RuntimeError):
|
|||
|
|
def __init__(self, message, *, status_code=None):
|
|||
|
|
super().__init__(message)
|
|||
|
|
self.status_code = status_code
|
|||
|
|
|
|||
|
|
|
|||
|
|
class QdrantClient:
|
|||
|
|
def __init__(self, url, api_key, prefix='yawan_', *, http_client=None):
|
|||
|
|
if not prefix or not all(c.isalnum() or c in '_-' for c in prefix):
|
|||
|
|
raise ValueError('Qdrant 知识库命名空间无效')
|
|||
|
|
self.prefix = prefix
|
|||
|
|
self._creation_lock = threading.Lock()
|
|||
|
|
self.http = http_client or httpx.Client(base_url=url.rstrip('/'),
|
|||
|
|
headers={'api-key': api_key}, timeout=httpx.Timeout(25, connect=5),
|
|||
|
|
limits=httpx.Limits(max_connections=8, max_keepalive_connections=4), trust_env=False)
|
|||
|
|
|
|||
|
|
def request(self, method, path, body=None, *, missing_ok=False):
|
|||
|
|
try:
|
|||
|
|
response = self.http.request(method, path, json=body)
|
|||
|
|
except httpx.TimeoutException as exc:
|
|||
|
|
raise QdrantError('Qdrant 知识索引请求超时，请检查向量服务') from exc
|
|||
|
|
except httpx.HTTPError as exc:
|
|||
|
|
raise QdrantError('Qdrant 知识索引连接失败，请检查向量服务') from exc
|
|||
|
|
if missing_ok and response.status_code == 404:
|
|||
|
|
return None
|
|||
|
|
if not response.is_success:
|
|||
|
|
# No body: it can contain stored user text or reflected credentials.
|
|||
|
|
raise QdrantError(f'Qdrant 知识索引操作失败（HTTP {response.status_code}）', status_code=response.status_code)
|
|||
|
|
try:
|
|||
|
|
payload = response.json()
|
|||
|
|
return payload.get('result', payload)
|
|||
|
|
except (ValueError, AttributeError) as exc:
|
|||
|
|
raise QdrantError('Qdrant 知识索引返回了无效数据') from exc
|
|||
|
|
|
|||
|
|
def get_collection(self, name):
|
|||
|
|
collection = QdrantCollection(self, name)
|
|||
|
|
result = self.request('GET', collection.path, missing_ok=True)
|
|||
|
|
if result is None:
|
|||
|
|
raise QdrantError(f'知识索引集合不存在：{name}')
|
|||
|
|
collection.configure(result)
|
|||
|
|
return collection
|
|||
|
|
|
|||
|
|
def get_or_create_collection(self, name):
|
|||
|
|
collection = QdrantCollection(self, name)
|
|||
|
|
with self._creation_lock:
|
|||
|
|
description = self.request('GET', collection.path, missing_ok=True)
|
|||
|
|
if description is None:
|
|||
|
|
try:
|
|||
|
|
self.request('PUT', collection.path, {
|
|||
|
|
'vectors': {'size': 512, 'distance': 'Cosine'},
|
|||
|
|
'on_disk_payload': True,
|
|||
|
|
'optimizers_config': {'default_segment_number': 2},
|
|||
|
|
'hnsw_config': {'m': 16, 'ef_construct': 128}})
|
|||
|
|
except QdrantError as exc:
|
|||
|
|
# Another process can win after our GET; only accept a
|
|||
|
|
# confirmed, compatible collection, never hide server errors.
|
|||
|
|
if exc.status_code not in {400, 409}:
|
|||
|
|
raise
|
|||
|
|
description = self.request('GET', collection.path, missing_ok=True)
|
|||
|
|
if description is None:
|
|||
|
|
raise exc
|
|||
|
|
else:
|
|||
|
|
description = {'config': {'params': {'vectors': {'size': 512, 'distance': 'Cosine'}}}}
|
|||
|
|
collection.configure(description)
|
|||
|
|
self.request('PUT', collection.path + '/index?wait=true',
|
|||
|
|
{'field_name': 'meta.file', 'field_schema': 'keyword'})
|
|||
|
|
else:
|
|||
|
|
collection.configure(description)
|
|||
|
|
return collection
|
|||
|
|
|
|||
|
|
def delete_collection(self, name):
|
|||
|
|
return self.request('DELETE', QdrantCollection(self, name).path, missing_ok=True)
|
|||
|
|
|
|||
|
|
def close(self):
|
|||
|
|
self.http.close()
|
|||
|
|
|
|||
|
|
|
|||
|
|
class QdrantCollection:
|
|||
|
|
def __init__(self, client, name):
|
|||
|
|
if not name or not all(c.isalnum() or c in '_-' for c in name):
|
|||
|
|
raise ValueError('知识索引集合名称无效')
|
|||
|
|
self.client, self.name = client, name
|
|||
|
|
self.path = '/collections/' + client.prefix + name
|
|||
|
|
self.distance = 'Cosine'
|
|||
|
|
|
|||
|
|
def configure(self, description):
|
|||
|
|
try:
|
|||
|
|
vectors = description['config']['params']['vectors']
|
|||
|
|
except (KeyError, TypeError):
|
|||
|
|
vectors = None
|
|||
|
|
if not isinstance(vectors, dict) or vectors.get('size') != 512 or vectors.get('distance') not in {'Cosine', 'Euclid'}:
|
|||
|
|
raise QdrantError('知识索引集合必须使用 512 维 Cosine 或 Euclid 向量')
|
|||
|
|
self.distance = vectors['distance']
|
|||
|
|
|
|||
|
|
def point_id(self, original_id):
|
|||
|
|
return str(uuid.uuid5(uuid.NAMESPACE_URL, self.client.prefix + self.name + ':' + str(original_id)))
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def filter(where):
|
|||
|
|
if not where:
|
|||
|
|
return None
|
|||
|
|
if not isinstance(where, dict) or set(where) != {'file'}:
|
|||
|
|
raise ValueError('不支持的知识索引过滤条件')
|
|||
|
|
value = where['file']
|
|||
|
|
if isinstance(value, dict):
|
|||
|
|
if set(value) != {'$in'} or not isinstance(value['$in'], (list, tuple, set)) or any(not isinstance(v, str) for v in value['$in']):
|
|||
|
|
raise ValueError('知识索引文件过滤条件无效')
|
|||
|
|
match = {'any': list(value['$in'])}
|
|||
|
|
elif isinstance(value, str):
|
|||
|
|
match = {'value': value}
|
|||
|
|
else:
|
|||
|
|
raise ValueError('知识索引文件过滤条件无效')
|
|||
|
|
return {'must': [{'key': 'meta.file', 'match': match}]}
|
|||
|
|
|
|||
|
|
def count(self, where=None):
|
|||
|
|
body = {'exact': True}
|
|||
|
|
filtering = self.filter(where)
|
|||
|
|
if filtering is not None:
|
|||
|
|
body['filter'] = filtering
|
|||
|
|
result = self.client.request('POST', self.path + '/points/count', body)
|
|||
|
|
return int(result['count'])
|
|||
|
|
|
|||
|
|
def _pages(self, where, include, size):
|
|||
|
|
if isinstance(size, bool) or not isinstance(size, int) or not 1 <= size <= 4000:
|
|||
|
|
raise ValueError('索引分页大小无效')
|
|||
|
|
offset = None
|
|||
|
|
seen_offsets = set()
|
|||
|
|
while True:
|
|||
|
|
body = {'limit': size, 'with_vector': 'embeddings' in include,
|
|||
|
|
'with_payload': {'include': ['chunk_id', 'meta'] + (['text'] if 'documents' in include else [])}}
|
|||
|
|
filtering = self.filter(where)
|
|||
|
|
if filtering is not None:
|
|||
|
|
body['filter'] = filtering
|
|||
|
|
if offset is not None:
|
|||
|
|
body['offset'] = offset
|
|||
|
|
result = self.client.request('POST', self.path + '/points/scroll', body)
|
|||
|
|
yield result.get('points', [])
|
|||
|
|
new_offset = result.get('next_page_offset')
|
|||
|
|
if new_offset is None:
|
|||
|
|
break
|
|||
|
|
if new_offset in seen_offsets:
|
|||
|
|
raise QdrantError('Qdrant 分页游标未前进')
|
|||
|
|
seen_offsets.add(new_offset)
|
|||
|
|
offset = new_offset
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _format(points, include):
|
|||
|
|
result = {'ids': [], 'documents': [], 'metadatas': [], 'embeddings': []}
|
|||
|
|
for point in points:
|
|||
|
|
payload = point.get('payload') or {}
|
|||
|
|
if not isinstance(payload.get('chunk_id'), str) or not payload['chunk_id']:
|
|||
|
|
raise QdrantError('知识索引原始片段编号缺失')
|
|||
|
|
result['ids'].append(payload['chunk_id'])
|
|||
|
|
if 'documents' in include:
|
|||
|
|
if not isinstance(payload.get('text'), str):
|
|||
|
|
raise QdrantError('知识索引原文缺失')
|
|||
|
|
result['documents'].append(payload['text'])
|
|||
|
|
if 'metadatas' in include:
|
|||
|
|
if not isinstance(payload.get('meta'), dict):
|
|||
|
|
raise QdrantError('知识索引来源元数据缺失')
|
|||
|
|
result['metadatas'].append(payload['meta'])
|
|||
|
|
if 'embeddings' in include:
|
|||
|
|
vector = point.get('vector')
|
|||
|
|
if not isinstance(vector, list) or len(vector) != 512 or any(not isinstance(v, (int, float)) or not math.isfinite(v) for v in vector):
|
|||
|
|
raise QdrantError('知识索引向量维度无效')
|
|||
|
|
result['embeddings'].append(vector)
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
def iter_records(self, where=None, include=None, page_size=256):
|
|||
|
|
include = ['documents', 'metadatas'] if include is None else include
|
|||
|
|
seen = set()
|
|||
|
|
for page in self._pages(where, include, page_size):
|
|||
|
|
self._check_scope(page, where)
|
|||
|
|
result = self._format(page, include)
|
|||
|
|
for identifier in result['ids']:
|
|||
|
|
if identifier in seen:
|
|||
|
|
raise QdrantError('知识索引分页返回重复片段，请重新读取')
|
|||
|
|
seen.add(identifier)
|
|||
|
|
yield result
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _check_scope(points, where):
|
|||
|
|
if not where:
|
|||
|
|
return
|
|||
|
|
allowed = where['file']
|
|||
|
|
for point in points:
|
|||
|
|
name = (point.get('payload', {}).get('meta') or {}).get('file')
|
|||
|
|
if not (name in allowed['$in'] if isinstance(allowed, dict) else name == allowed):
|
|||
|
|
raise QdrantError('知识索引返回了文件范围外的片段')
|
|||
|
|
|
|||
|
|
def get(self, *, ids=None, where=None, include=None, limit=None, offset=0):
|
|||
|
|
include = ['documents', 'metadatas'] if include is None else include
|
|||
|
|
if isinstance(offset, bool) or not isinstance(offset, int) or offset < 0 or (limit is not None and (isinstance(limit, bool) or not isinstance(limit, int) or limit < 0)):
|
|||
|
|
raise ValueError('索引分页范围无效')
|
|||
|
|
if limit == 0 or ids == []:
|
|||
|
|
return self._format([], include)
|
|||
|
|
if ids is not None:
|
|||
|
|
self.filter(where)
|
|||
|
|
found = {}
|
|||
|
|
for start in range(0, len(ids), 128):
|
|||
|
|
batch = ids[start:start + 128]
|
|||
|
|
points = self.client.request('POST', self.path + '/points',
|
|||
|
|
{'ids': [self.point_id(identifier) for identifier in batch],
|
|||
|
|
'with_payload': True, 'with_vector': 'embeddings' in include})
|
|||
|
|
for point in points:
|
|||
|
|
original_id = (point.get('payload') or {}).get('chunk_id')
|
|||
|
|
if original_id not in batch:
|
|||
|
|
raise QdrantError('知识索引返回了请求范围外的片段')
|
|||
|
|
if where:
|
|||
|
|
name = (point.get('payload', {}).get('meta') or {}).get('file')
|
|||
|
|
allowed = where['file']
|
|||
|
|
if not (name in allowed['$in'] if isinstance(allowed, dict) else name == allowed):
|
|||
|
|
continue
|
|||
|
|
found[original_id] = point
|
|||
|
|
selected = [found[identifier] for identifier in ids if identifier in found]
|
|||
|
|
return self._format(selected[offset:None if limit is None else offset + limit], include)
|
|||
|
|
retained = []
|
|||
|
|
visited = 0
|
|||
|
|
for page in self._pages(where, include, min(256, max(1, limit or 256))):
|
|||
|
|
self._check_scope(page, where)
|
|||
|
|
for point in page:
|
|||
|
|
if visited >= offset:
|
|||
|
|
retained.append(point)
|
|||
|
|
visited += 1
|
|||
|
|
if limit is not None and len(retained) >= limit:
|
|||
|
|
return self._format(retained, include)
|
|||
|
|
return self._format(retained, include)
|
|||
|
|
|
|||
|
|
def add(self, *, ids, documents, embeddings, metadatas):
|
|||
|
|
if not (len(ids) == len(documents) == len(embeddings) == len(metadatas)):
|
|||
|
|
raise ValueError('索引编号、正文、向量和元数据数量不一致')
|
|||
|
|
# Validate the entire input before the first batch is written.
|
|||
|
|
if any(not isinstance(identifier, str) or not identifier for identifier in ids) or len(set(ids)) != len(ids):
|
|||
|
|
raise ValueError('索引片段编号为空或重复')
|
|||
|
|
for text, vector, meta in zip(documents, embeddings, metadatas):
|
|||
|
|
if not isinstance(text, str) or not text.strip() or not isinstance(meta, dict):
|
|||
|
|
raise ValueError('索引正文或元数据无效')
|
|||
|
|
if not isinstance(vector, (list, tuple)) or len(vector) != 512 or any(not isinstance(v, (int, float)) or not math.isfinite(v) for v in vector):
|
|||
|
|
raise ValueError('索引向量必须是 512 维有限数值')
|
|||
|
|
for start in range(0, len(ids), 128):
|
|||
|
|
points = [{'id': self.point_id(ids[i]), 'vector': embeddings[i],
|
|||
|
|
'payload': {'chunk_id': ids[i], 'text': documents[i], 'meta': metadatas[i]}}
|
|||
|
|
for i in range(start, min(len(ids), start + 128))]
|
|||
|
|
self.client.request('PUT', self.path + '/points?wait=true', {'points': points})
|
|||
|
|
|
|||
|
|
def delete(self, *, where=None, ids=None):
|
|||
|
|
if ids is not None and where is not None:
|
|||
|
|
raise ValueError('删除知识索引不可同时指定两种范围')
|
|||
|
|
if ids is not None:
|
|||
|
|
if not ids:
|
|||
|
|
return None
|
|||
|
|
body = {'points': [self.point_id(identifier) for identifier in ids]}
|
|||
|
|
elif where:
|
|||
|
|
body = {'filter': self.filter(where)}
|
|||
|
|
else:
|
|||
|
|
raise ValueError('删除知识索引必须指定文件或片段编号')
|
|||
|
|
return self.client.request('POST', self.path + '/points/delete?wait=true', body)
|
|||
|
|
|
|||
|
|
def query(self, *, query_embeddings, n_results=6, where=None):
|
|||
|
|
if isinstance(n_results, bool) or not isinstance(n_results, int) or n_results < 1:
|
|||
|
|
raise ValueError('知识索引检索数量无效')
|
|||
|
|
result = {'ids': [], 'documents': [], 'metadatas': [], 'distances': []}
|
|||
|
|
for vector in query_embeddings:
|
|||
|
|
body = {'query': vector, 'limit': n_results, 'with_payload': True,
|
|||
|
|
'with_vector': False, 'params': {'hnsw_ef': max(64, min(256, n_results * 2))}}
|
|||
|
|
filtering = self.filter(where)
|
|||
|
|
if filtering is not None:
|
|||
|
|
body['filter'] = filtering
|
|||
|
|
points = self.client.request('POST', self.path + '/points/query', body).get('points', [])
|
|||
|
|
self._check_scope(points, where)
|
|||
|
|
formatted = self._format(points, ['documents', 'metadatas'])
|
|||
|
|
for key in ('ids', 'documents', 'metadatas'):
|
|||
|
|
result[key].append(formatted[key])
|
|||
|
|
# Chroma's existing default is squared L2. Our embeddings are normalized.
|
|||
|
|
result['distances'].append([max(0., 2. * (1. - float(point['score']))) if self.distance == 'Cosine'
|
|||
|
|
else float(point['score']) ** 2 for point in points])
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
|
|||
|
|
_client = None
|
|||
|
|
_client_signature = None
|
|||
|
|
_lock = threading.Lock()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_client():
|
|||
|
|
from .backend import qdrant_settings
|
|||
|
|
global _client, _client_signature
|
|||
|
|
url, key, prefix = qdrant_settings()
|
|||
|
|
signature = (url, hashlib.sha256(key.encode()).hexdigest(), prefix)
|
|||
|
|
with _lock:
|
|||
|
|
if _client is None or _client_signature != signature:
|
|||
|
|
if _client is not None:
|
|||
|
|
_client.close()
|
|||
|
|
_client = QdrantClient(url, key, prefix)
|
|||
|
|
_client_signature = signature
|
|||
|
|
return _client
|