#!/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())