204 lines
11 KiB
Python
204 lines
11 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Upload one temporary acceptance document, verify retrieval, then delete it.
|
||
|
|
|
||
|
|
Service mutations use only 127.0.0.1:18080 and the KB created by this run.
|
||
|
|
Milvus access, when enabled, is restricted to read-only entities/query calls.
|
||
|
|
"""
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import time
|
||
|
|
import urllib.error
|
||
|
|
import urllib.request
|
||
|
|
import uuid
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
from milvus_verify_acceptance import (ACCEPTANCE, BASE, MODES, NoRedirect,
|
||
|
|
VerificationError, inspect_results,
|
||
|
|
read_object, search_payload)
|
||
|
|
|
||
|
|
|
||
|
|
class SmokeClient:
|
||
|
|
def __init__(self):
|
||
|
|
self.opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), NoRedirect())
|
||
|
|
self.token = ''
|
||
|
|
self.kb_id = None
|
||
|
|
self.knowledge_id = None
|
||
|
|
|
||
|
|
def request(self, route, method='GET', payload=None, raw=None, content_type=None):
|
||
|
|
kb = '/api/v1/knowledge-bases/' + str(self.kb_id)
|
||
|
|
knowledge = '/api/v1/knowledge/' + str(self.knowledge_id)
|
||
|
|
allowed = {('POST', '/api/v1/auth/login'), ('POST', '/api/v1/knowledge-bases')}
|
||
|
|
if self.kb_id:
|
||
|
|
allowed.update({('POST', kb + '/knowledge/file'), ('POST', kb + '/hybrid-search'),
|
||
|
|
('DELETE', kb), ('GET', kb)})
|
||
|
|
if self.knowledge_id:
|
||
|
|
allowed.add(('GET', knowledge))
|
||
|
|
if (method, route) not in allowed:
|
||
|
|
raise VerificationError('Request is outside the temporary resource allowlist')
|
||
|
|
headers = {'Accept': 'application/json'}
|
||
|
|
if self.token:
|
||
|
|
headers['Authorization'] = 'Bearer ' + self.token
|
||
|
|
if payload is not None:
|
||
|
|
raw = json.dumps(payload).encode()
|
||
|
|
content_type = 'application/json'
|
||
|
|
if content_type:
|
||
|
|
headers['Content-Type'] = content_type
|
||
|
|
started = time.monotonic()
|
||
|
|
request = urllib.request.Request(BASE + route, data=raw, headers=headers, method=method)
|
||
|
|
try:
|
||
|
|
try:
|
||
|
|
response = self.opener.open(request, timeout=45)
|
||
|
|
except urllib.error.HTTPError as error:
|
||
|
|
response = error
|
||
|
|
with response:
|
||
|
|
status = response.code
|
||
|
|
body = response.read(8 * 1024 * 1024)
|
||
|
|
value = json.loads(body) if status in (200, 201, 202) else None
|
||
|
|
except (OSError, ValueError):
|
||
|
|
raise VerificationError('Acceptance request failed or timed out') from None
|
||
|
|
return status, value, round((time.monotonic() - started) * 1000, 2)
|
||
|
|
|
||
|
|
def count_vectors(self, collection):
|
||
|
|
if collection not in ('yawan_acceptance_hybrid_512', 'weknora_embeddings_512') or not self.kb_id:
|
||
|
|
raise VerificationError('Milvus query is outside the allowed collections')
|
||
|
|
request = urllib.request.Request('http://127.0.0.1:19530/v2/vectordb/entities/query',
|
||
|
|
data=json.dumps({'collectionName': collection,
|
||
|
|
'filter': 'knowledge_base_id == ' + json.dumps(self.kb_id),
|
||
|
|
'outputFields': ['count(*)'], 'consistencyLevel': 'Strong'}).encode(),
|
||
|
|
headers={'Content-Type': 'application/json'})
|
||
|
|
try:
|
||
|
|
with self.opener.open(request, timeout=30) as response:
|
||
|
|
value = json.load(response)
|
||
|
|
if value.get('code') != 0:
|
||
|
|
raise VerificationError('Milvus count query failed')
|
||
|
|
return int(value['data'][0]['count(*)'])
|
||
|
|
except (OSError, ValueError, KeyError, IndexError):
|
||
|
|
raise VerificationError('Milvus count query failed or returned an invalid shape') from None
|
||
|
|
|
||
|
|
|
||
|
|
def run(args, report):
|
||
|
|
credentials = read_object(ACCEPTANCE / 'migration-token.json', 'Migration credentials')
|
||
|
|
state = read_object(ACCEPTANCE / 'legacy-migration-state.json', 'Migration state')
|
||
|
|
client = SmokeClient()
|
||
|
|
status, login, _ = client.request('/api/v1/auth/login', 'POST',
|
||
|
|
{'email': credentials['email'], 'password': credentials['password']})
|
||
|
|
if status != 200 or login.get('active_tenant', {}).get('id') != state['tenant_id']:
|
||
|
|
raise VerificationError('Login failed or used an unexpected tenant')
|
||
|
|
client.token = login['token']
|
||
|
|
suffix = uuid.uuid4().hex[:12]
|
||
|
|
kb_name = '临时验收-新上传检索-' + suffix
|
||
|
|
try:
|
||
|
|
status, value, latency = client.request('/api/v1/knowledge-bases', 'POST', {
|
||
|
|
'name': kb_name, 'description': 'Temporary acceptance-only upload smoke test',
|
||
|
|
'type': 'document', 'embedding_model_id': state['model_id'],
|
||
|
|
'summary_model_id': '',
|
||
|
|
'chunking_config': {'chunk_size': 512, 'chunk_overlap': 0},
|
||
|
|
'indexing_strategy': {'vector_enabled': True, 'keyword_enabled': True,
|
||
|
|
'wiki_enabled': False, 'graph_enabled': False},
|
||
|
|
'storage_provider_config': {'provider': 'local'},
|
||
|
|
'auto_tag_config': {'enabled': False},
|
||
|
|
})
|
||
|
|
if status != 201 or not isinstance(value.get('data'), dict):
|
||
|
|
raise VerificationError('Temporary knowledge base creation failed')
|
||
|
|
client.kb_id = value['data']['id']
|
||
|
|
report.update(kb_id=client.kb_id, create_ms=latency)
|
||
|
|
migrated_ids = {row['id'] for row in state['knowledge_bases'].values()}
|
||
|
|
if client.kb_id in migrated_ids:
|
||
|
|
client.kb_id = None
|
||
|
|
raise VerificationError('Created KB unexpectedly matched protected migration IDs')
|
||
|
|
if args.verify_milvus:
|
||
|
|
report['vectors_before'] = {name: client.count_vectors(name) for name in
|
||
|
|
('yawan_acceptance_hybrid_512', 'weknora_embeddings_512')}
|
||
|
|
if any(report['vectors_before'].values()):
|
||
|
|
raise VerificationError('Fresh temporary KB unexpectedly already has vectors')
|
||
|
|
boundary = 'acceptance' + uuid.uuid4().hex
|
||
|
|
document = ('# 资料归档演练\n\n编号：' + suffix + '\n\n'
|
||
|
|
'资料归档采用三步流程：先登记文件名称，再核对目录分类，最后确认文件可以检索。'
|
||
|
|
'检索测试使用资料归档和目录核对两个关键词。完成演练后删除临时资料，不影响其他知识库。\n')
|
||
|
|
raw = ('--' + boundary + '\r\nContent-Disposition: form-data; name="file"; '
|
||
|
|
'filename="acceptance-upload-' + suffix + '.md"\r\n'
|
||
|
|
'Content-Type: text/markdown; charset=utf-8\r\n\r\n' + document +
|
||
|
|
'\r\n--' + boundary + '--\r\n').encode('utf-8')
|
||
|
|
status, value, latency = client.request('/api/v1/knowledge-bases/' + client.kb_id + '/knowledge/file',
|
||
|
|
'POST', raw=raw, content_type='multipart/form-data; boundary=' + boundary)
|
||
|
|
if status not in (200, 201) or not isinstance(value.get('data'), dict):
|
||
|
|
raise VerificationError('Temporary file upload failed')
|
||
|
|
client.knowledge_id = value['data']['id']
|
||
|
|
report.update(knowledge_id=client.knowledge_id, upload_ms=latency, file_bytes=len(document.encode('utf-8')))
|
||
|
|
started = time.monotonic()
|
||
|
|
while time.monotonic() - started < args.index_timeout:
|
||
|
|
status, value, _ = client.request('/api/v1/knowledge/' + client.knowledge_id)
|
||
|
|
parse_status = value.get('data', {}).get('parse_status') if isinstance(value, dict) else None
|
||
|
|
if parse_status == 'completed':
|
||
|
|
break
|
||
|
|
if parse_status in ('failed', 'cancelled', 'deleting') or status != 200:
|
||
|
|
raise VerificationError('Temporary file indexing failed')
|
||
|
|
time.sleep(2)
|
||
|
|
else:
|
||
|
|
raise VerificationError('Temporary file indexing exceeded its deadline')
|
||
|
|
report['indexing_seconds'] = round(time.monotonic() - started, 2)
|
||
|
|
for mode in MODES:
|
||
|
|
status, value, latency = client.request('/api/v1/knowledge-bases/' + client.kb_id + '/hybrid-search',
|
||
|
|
'POST', search_payload('资料归档 目录核对', mode, 5))
|
||
|
|
result = inspect_results(value, client.kb_id, 5, mode) if status == 200 else {'passed': False}
|
||
|
|
result.update(mode=mode, status=status, latency_ms=latency)
|
||
|
|
report['retrieval'].append(result)
|
||
|
|
if args.verify_milvus:
|
||
|
|
report['vectors_after_upload'] = {name: client.count_vectors(name) for name in
|
||
|
|
('yawan_acceptance_hybrid_512', 'weknora_embeddings_512')}
|
||
|
|
report['isolated_vector_write'] = (report['vectors_after_upload']['yawan_acceptance_hybrid_512'] > 0
|
||
|
|
and report['vectors_after_upload']['weknora_embeddings_512'] == 0)
|
||
|
|
finally:
|
||
|
|
if client.kb_id:
|
||
|
|
status, _, latency = client.request('/api/v1/knowledge-bases/' + client.kb_id, 'DELETE')
|
||
|
|
report['cleanup'] = {'delete_status': status, 'delete_ms': latency, 'passed': False}
|
||
|
|
if status not in (200, 202, 204):
|
||
|
|
raise VerificationError('Temporary KB cleanup request failed')
|
||
|
|
status, _, _ = client.request('/api/v1/knowledge-bases/' + client.kb_id)
|
||
|
|
report['cleanup']['kb_get_status'] = status
|
||
|
|
report['cleanup']['passed'] = status == 404
|
||
|
|
if args.verify_milvus:
|
||
|
|
started = time.monotonic()
|
||
|
|
while time.monotonic() - started < 90:
|
||
|
|
counts = {name: client.count_vectors(name) for name in
|
||
|
|
('yawan_acceptance_hybrid_512', 'weknora_embeddings_512')}
|
||
|
|
if all(count == 0 for count in counts.values()):
|
||
|
|
break
|
||
|
|
time.sleep(2)
|
||
|
|
report['cleanup'].update(vector_counts=counts,
|
||
|
|
vector_cleanup_seconds=round(time.monotonic() - started, 2))
|
||
|
|
report['cleanup']['passed'] = report['cleanup']['passed'] and all(count == 0 for count in counts.values())
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument('--verify-milvus', action='store_true', help='Read counts for this temporary KB from acceptance and production collections')
|
||
|
|
parser.add_argument('--output', type=Path)
|
||
|
|
parser.add_argument('--index-timeout', type=int, default=180)
|
||
|
|
args = parser.parse_args()
|
||
|
|
if not 10 <= args.index_timeout <= 600:
|
||
|
|
parser.error('index-timeout must be 10-600 seconds')
|
||
|
|
if args.output and args.output.resolve() in {file.resolve() for file in
|
||
|
|
(ACCEPTANCE / 'migration-token.json', ACCEPTANCE / 'legacy-migration-state.json', ACCEPTANCE / 'browser-auth.json')}:
|
||
|
|
parser.error('Output must not overwrite existing authentication or migration state')
|
||
|
|
report = {'acceptance_url': BASE, 'retrieval': [], 'failures': []}
|
||
|
|
try:
|
||
|
|
run(args, report)
|
||
|
|
except VerificationError as error:
|
||
|
|
report['failures'].append(str(error))
|
||
|
|
except Exception:
|
||
|
|
report['failures'].append('Unexpected smoke-test failure; no server error or credentials included')
|
||
|
|
report['passed'] = (not report['failures'] and len(report['retrieval']) == 3
|
||
|
|
and all(row['passed'] for row in report['retrieval'])
|
||
|
|
and report.get('cleanup', {}).get('passed', False)
|
||
|
|
and (not args.verify_milvus or report.get('isolated_vector_write', False)))
|
||
|
|
text = json.dumps(report, ensure_ascii=False, indent=2, allow_nan=False) + '\n'
|
||
|
|
if args.output:
|
||
|
|
args.output.write_text(text, encoding='utf-8')
|
||
|
|
print(text, end='')
|
||
|
|
return 0 if report['passed'] else 1
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
raise SystemExit(main())
|