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
|