xinmi/scripts/milvus_audit_acceptance.py

#!/usr/bin/env python3
"""Read-only vector copy audit and distinct-query acceptance latency sample.

Reports only counts, checksums, and timings; never knowledge prose or credentials.
"""
import hashlib
import json
import math
import statistics
import time
from pathlib import Path

from milvus_copy_acceptance import FIELDS, post
from milvus_verify_acceptance import (ACCEPTANCE, Client, VerificationError,
                                     inspect_results, read_object, search_payload)

SOURCE = 'weknora_embeddings_512'
TARGET = 'yawan_acceptance_hybrid_512'
MILVUS = 'http://127.0.0.1:19530'
QUERIES = {
    '1': ['产品资料应如何整理和解释', '介绍产品时需要说明哪些使用方法'],
    '2': ['如何安排由浅入深的练习流程', '基础练习中怎样判断动作是否正确'],
    '3': ['怎样通过沟通理解对方的需求', '遇到紧张情绪时如何调整表达方式'],
    '4': ['如何描述亲密关系中的感受和期待', '怎样建立双方舒适的互动氛围'],
    '5': ['学习新技巧时如何分解具体步骤', '练习过程有哪些常见错误需要纠正'],
    '6': ['如何理解客户需求并解释产品价值', '怎样回应疑问并帮助客户作出选择'],
}


def query(collection, filter_text, fields, limit=100):
    return post(MILVUS, 'entities/query', {
        'collectionName': collection, 'filter': filter_text,
        'outputFields': fields, 'limit': limit, 'consistencyLevel': 'Strong',
    })


def digest(row):
    return hashlib.sha256(json.dumps({field: row[field] for field in FIELDS},
                                     sort_keys=True, ensure_ascii=False,
                                     separators=(',', ':')).encode()).hexdigest()


def run():
    state = read_object(ACCEPTANCE / 'legacy-migration-state.json', 'Migration state')
    copy = read_object(ACCEPTANCE / 'hybrid-copy-state.json', 'Copy state')
    manifest = json.loads((ACCEPTANCE / 'hybrid-chunk-manifest.json').read_text())
    report = {'source': SOURCE, 'target': TARGET, 'storage': [], 'queries': [],
              'scope': 'Six migrated business knowledge bases only',
              'latency_scope': 'Sequential distinct queries, warmed service; cache not cleared or instrumented',
              'started_at': time.time()}
    business_ids = {kb['id'] for kb in state['knowledge_bases'].values()}
    missing_ids = set(copy['missing'])
    business_missing = [row['id'] for row in manifest
                        if row['id'] in missing_ids and row['knowledge_base_id'] in business_ids]
    report['missing_business_chunks'] = len(business_missing)
    report['missing_historical_test_chunks'] = len(missing_ids) - len(business_missing)
    for legacy_id, kb in sorted(state['knowledge_bases'].items()):
        entries = [row for row in manifest if row['knowledge_base_id'] == kb['id']]
        filter_text = 'knowledge_base_id == ' + json.dumps(kb['id'])
        counts = {name: int(query(name, filter_text, ['count(*)'])[0]['count(*)'])
                  for name in (SOURCE, TARGET)}
        positions = sorted({0, len(entries) // 2, len(entries) - 1})
        ids = [entries[index]['id'] for index in positions]
        samples = {}
        for name in (SOURCE, TARGET):
            rows = query(name, 'chunk_id in ' + json.dumps(ids), FIELDS, len(ids))
            samples[name] = {row['chunk_id']: digest(row) for row in rows}
        equal = (set(samples[SOURCE]) == set(ids) and samples[SOURCE] == samples[TARGET])
        report['storage'].append({
            'legacy_kb': legacy_id, 'kb_id': kb['id'], 'expected_chunks': len(entries),
            'source_count': counts[SOURCE], 'target_count': counts[TARGET],
            'sample_count': len(ids), 'sample_fields_equal': equal,
            'passed': equal and all(count == len(entries) for count in counts.values()),
        })
    credentials = read_object(ACCEPTANCE / 'migration-token.json', 'Migration credentials')
    client = Client(45)
    status, login, _ = client.request('/api/v1/auth/login', payload={
        'email': credentials['email'], 'password': credentials['password'],
    })
    if status != 200 or login.get('active_tenant', {}).get('id') != state['tenant_id']:
        raise VerificationError('Unexpected login scope')
    for legacy_id, queries in sorted(QUERIES.items()):
        kb_id = state['knowledge_bases'][legacy_id]['id']
        for index, text in enumerate(queries):
            status, result, latency = client.request(
                '/api/v1/knowledge-bases/' + kb_id + '/hybrid-search', login['token'],
                search_payload(text, 'hybrid', 10))
            check = inspect_results(result, kb_id, 10, 'hybrid') if status == 200 else {'passed': False}
            report['queries'].append({'legacy_kb': legacy_id, 'query_index': index,
                                      'query_sha256': hashlib.sha256(text.encode()).hexdigest(),
                                      'status': status, 'latency_ms': latency, **check})
    latencies = sorted(row['latency_ms'] for row in report['queries'])
    report['latency_ms'] = {'min': min(latencies), 'median': statistics.median(latencies),
                            'p95_nearest_rank': latencies[math.ceil(len(latencies) * .95) - 1],
                            'max': max(latencies), 'count': len(latencies)}
    report['passed'] = (not business_missing and all(row['passed'] for row in report['storage'])
                        and all(row['passed'] for row in report['queries']))
    report['finished_at'] = time.time()
    output = ACCEPTANCE / 'milvus-storage-and-latency.json'
    output.write_text(json.dumps(report, ensure_ascii=False, indent=2) + '\n')
    print(json.dumps({key: report[key] for key in (
        'passed', 'missing_business_chunks', 'missing_historical_test_chunks', 'latency_ms')}, indent=2))
    return 0 if report['passed'] else 1


if __name__ == '__main__':
    raise SystemExit(run())