xinmi/scripts/milvus_upload_smoke_acceptance.py

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