xinmi/scripts/milvus_verify_acceptance.py

#!/usr/bin/env python3
"""Verify retrieval in the isolated acceptance API without changing its data.

Only login and read-only search POSTs are allowed. Credentials and response
prose are never included in the report. No files are written unless --output
is supplied, and credential files are never refreshed or overwritten.
"""

import argparse
import datetime
import json
import math
import re
import time
import urllib.error
import urllib.request
from pathlib import Path


BASE = 'http://127.0.0.1:18080'
ACCEPTANCE = Path(__file__).resolve().parents[2] / 'acceptance'
UUID = re.compile(r'^[0-9a-fA-F]{8}(?:-[0-9a-fA-F]{4}){3}-[0-9a-fA-F]{12}$')
MODES = {
    'keyword-only': (False, True),
    'vector-only': (True, False),
    'hybrid': (False, False),
}
QUERIES = {
    '1': '产品 使用 说明',
    '2': '练习 方法 步骤',
    '3': '沟通 情绪 理解',
    '4': '关系 感受 表达',
    '5': '技巧 步骤 学习',
    '6': '需求 价值 选择',
}


class VerificationError(Exception):
    """A fixed, credential-free error suitable for the report."""


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


class Client:
    def __init__(self, timeout):
        self.timeout = timeout
        self.opener = urllib.request.build_opener(
            urllib.request.ProxyHandler({}), NoRedirect())

    def request(self, route, token=None, payload=None):
        search_route = re.fullmatch(
            r'/api/v1/knowledge-bases/[0-9a-fA-F-]+/hybrid-search', route)
        allowed_read = route in ('/health', '/api/v1/auth/me')
        allowed_login = route == '/api/v1/auth/login'
        if not ((payload is None and allowed_read) or
                (payload is not None and (search_route or allowed_login))):
            raise VerificationError('Request outside the read-only acceptance allowlist')
        headers = {'Accept': 'application/json'}
        if token:
            headers['Authorization'] = 'Bearer ' + token
        body = None
        if payload is not None:
            headers['Content-Type'] = 'application/json'
            body = json.dumps(payload).encode('utf-8')
        request = urllib.request.Request(BASE + route, data=body, headers=headers)
        started = time.monotonic()
        try:
            try:
                response = self.opener.open(request, timeout=self.timeout)
            except urllib.error.HTTPError as error:
                response = error
            with response:
                status = response.code
                raw = response.read(8 * 1024 * 1024 + 1)
            if len(raw) > 8 * 1024 * 1024:
                raise VerificationError('Response exceeded the verification size limit')
            # Denied requests need only a status; never retain their error prose.
            value = json.loads(raw) if status == 200 else None
            return status, value, round((time.monotonic() - started) * 1000, 2)
        except VerificationError:
            raise
        except (OSError, urllib.error.URLError, ValueError):
            raise VerificationError('Request failed, timed out, or returned invalid JSON') from None


def read_object(file, label):
    try:
        value = json.loads(file.read_text(encoding='utf-8'))
    except (OSError, ValueError):
        raise VerificationError(label + ' could not be read as JSON') from None
    if not isinstance(value, dict):
        raise VerificationError(label + ' must contain a JSON object')
    return value


def search_payload(query, mode, count):
    disable_keywords, disable_vector = MODES[mode]
    return {
        'query_text': query,
        'match_count': count,
        'vector_threshold': 0,
        'keyword_threshold': 0,
        'disable_keywords_match': disable_keywords,
        'disable_vector_match': disable_vector,
        'skip_context_enrichment': True,
    }


def inspect_results(value, kb_id, count, mode):
    if not isinstance(value, dict) or value.get('success') is not True:
        return {'passed': False, 'failure': 'Invalid search response envelope'}
    rows = value.get('data') or []
    if not isinstance(rows, list) or any(not isinstance(row, dict) for row in rows):
        return {'passed': False, 'failure': 'Invalid search results shape'}
    scores = [row.get('score') for row in rows]
    finite = all(isinstance(score, (int, float)) and not isinstance(score, bool)
                 and math.isfinite(score) for score in scores)
    descending = finite and all(left + 1e-9 >= right
                                for left, right in zip(scores, scores[1:]))
    in_scope = all(row.get('knowledge_base_id') == kb_id for row in rows)
    ids = [row.get('id') for row in rows]
    unique = all(isinstance(item, str) and item for item in ids) and len(set(ids)) == len(ids)
    # Match types 0 and 1 are vector and keyword hits respectively.
    expected = {'keyword-only': {1}, 'vector-only': {0}, 'hybrid': {0, 1}}[mode]
    correct_mode = all(type(row.get('match_type')) is int and row['match_type'] in expected for row in rows)
    return {
        'passed': bool(rows) and len(rows) <= count and finite and descending and in_scope and unique and correct_mode,
        'count': len(rows),
        'nonempty': bool(rows),
        'within_requested_count': len(rows) <= count,
        'scores': [round(score, 8) for score in scores] if finite else [],
        'scores_finite': finite,
        'scores_descending': descending,
        'kb_scope_valid': in_scope,
        'unique_chunk_ids': unique,
        'match_types_valid': correct_mode,
    }


def verify(args, report):
    credentials = read_object(args.credentials, 'Migration credentials')
    migration = read_object(args.state, 'Migration state')
    foreign = read_object(args.foreign_auth, 'Foreign tenant authentication')
    entries = migration.get('knowledge_bases', {})
    if not isinstance(entries, dict) or set(entries) != set(QUERIES):
        raise VerificationError('Migration state must describe exactly the six expected knowledge bases')
    kb_ids = []
    for legacy_id, row in entries.items():
        if not isinstance(row, dict) or not UUID.fullmatch(str(row.get('id', ''))):
            raise VerificationError('Migration state contains an invalid knowledge base ID')
        kb_ids.append(row['id'])
    if len(set(kb_ids)) != 6:
        raise VerificationError('Migration state contains duplicate knowledge base IDs')
    if not all(isinstance(credentials.get(key), str) and credentials[key] for key in ('email', 'password')):
        raise VerificationError('Migration login credentials are incomplete')
    if not isinstance(foreign.get('token'), str) or not foreign['token']:
        raise VerificationError('Existing foreign tenant token is required')

    client = Client(args.timeout)
    status, _, latency = client.request('/health')
    report['checks'].append({'check': 'acceptance_health', 'status': status,
                             'latency_ms': latency, 'passed': status == 200})
    if status != 200:
        raise VerificationError('Acceptance API is not healthy')

    status, login, latency = client.request('/api/v1/auth/login', payload={
        'email': credentials['email'], 'password': credentials['password'],
    })
    if status != 200 or not isinstance(login, dict) or not isinstance(login.get('token'), str):
        raise VerificationError('Migration login failed')
    token = login['token']
    tenant_id = login.get('active_tenant', {}).get('id')
    correct_tenant = tenant_id is not None and tenant_id == migration.get('tenant_id') == credentials.get('tenant_id')
    report['checks'].append({'check': 'migration_login_scope', 'status': status,
                             'latency_ms': latency, 'passed': correct_tenant})
    if not correct_tenant:
        raise VerificationError('Migration login is outside the expected tenant')

    status, identity, latency = client.request('/api/v1/auth/me', foreign['token'])
    details = identity.get('data', {}) if isinstance(identity, dict) else {}
    foreign_tenant = (details.get('tenant') or {}).get('id')
    user = details.get('user') or {}
    memberships = details.get('memberships') or []
    isolated = (status == 200 and foreign_tenant is not None and foreign_tenant != tenant_id
                and not user.get('is_system_admin') and not user.get('can_access_all_tenants')
                and not any(member.get('tenant_id') == tenant_id for member in memberships if isinstance(member, dict)))
    report['checks'].append({'check': 'foreign_identity_valid_and_isolated', 'status': status,
                             'latency_ms': latency, 'passed': isolated})
    if not isolated:
        report['failures'].append('Foreign authentication is expired, invalid, or not isolated; denial tests cannot prove tenant isolation')

    for legacy_id in sorted(entries):
        kb_id = entries[legacy_id]['id']
        route = '/api/v1/knowledge-bases/' + kb_id + '/hybrid-search'
        payload = search_payload(QUERIES[legacy_id], 'keyword-only', args.top_k)
        for check, auth, enabled, statuses in (
            ('unauthenticated_denied', None, True, (401, 403)),
            ('foreign_tenant_denied', foreign['token'], isolated, (403, 404)),
        ):
            if not enabled:
                continue
            result = {'check': check, 'legacy_kb': legacy_id, 'kb_id': kb_id}
            try:
                status, _, latency = client.request(route, auth, payload)
                result.update(status=status, latency_ms=latency, passed=status in statuses)
            except VerificationError as error:
                result.update(passed=False, failure=str(error))
            report['checks'].append(result)

        for mode in MODES:
            result = {'legacy_kb': legacy_id, 'kb_id': kb_id, 'mode': mode,
                      'query': QUERIES[legacy_id], 'top_k': args.top_k}
            try:
                status, value, latency = client.request(route, token, search_payload(QUERIES[legacy_id], mode, args.top_k))
                result.update(status=status, latency_ms=latency)
                result.update(inspect_results(value, kb_id, args.top_k, mode) if status == 200
                              else {'passed': False, 'failure': 'Search HTTP status was not 200'})
            except VerificationError as error:
                result.update(passed=False, failure=str(error))
            report['retrieval'].append(result)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--credentials', type=Path, default=ACCEPTANCE / 'migration-token.json')
    parser.add_argument('--state', type=Path, default=ACCEPTANCE / 'legacy-migration-state.json')
    parser.add_argument('--foreign-auth', type=Path, default=ACCEPTANCE / 'browser-auth.json')
    parser.add_argument('--output', type=Path, help='Write the same JSON report to this file only')
    parser.add_argument('--timeout', type=float, default=45, help='Timeout in seconds for every HTTP request (1-180)')
    parser.add_argument('--top-k', type=int, default=10, help='Requested primary matches (1-50)')
    args = parser.parse_args()
    if not 1 <= args.timeout <= 180 or not 1 <= args.top_k <= 50:
        parser.error('timeout must be 1-180 seconds and top-k must be 1-50')
    if args.output and args.output.resolve() in {file.resolve() for file in (args.credentials, args.state, args.foreign_auth)}:
        parser.error('output must not overwrite authentication or migration state')
    started = time.monotonic()
    report = {
        'acceptance_url': BASE,
        'started_at': datetime.datetime.now(datetime.timezone.utc).isoformat(),
        'checks': [], 'retrieval': [], 'failures': [],
        'limitations': ['API validation does not independently attest the configured backend or collection',
                        'Nonempty and ordered matches are smoke checks, not a relevance-quality benchmark'],
    }
    try:
        verify(args, report)
    except VerificationError as error:
        report['failures'].append(str(error))
    except Exception:
        # Do not expose response contents, credentials, or traceback locals.
        report['failures'].append('Unexpected verification failure; inspect the verifier and response shapes locally')
    all_checks = report['checks'] + report['retrieval']
    report['passed'] = (not report['failures'] and len(report['retrieval']) == 18
                        and all(item.get('passed') for item in all_checks))
    report['duration_seconds'] = round(time.monotonic() - started, 2)
    report['summary'] = {'checks': len(all_checks),
                         'passed': sum(bool(item.get('passed')) for item in all_checks),
                         'failed': sum(not item.get('passed') for item in all_checks)}
    encoded = json.dumps(report, ensure_ascii=False, indent=2, allow_nan=False) + '\n'
    if args.output:
        try:
            args.output.write_text(encoded, encoding='utf-8')
        except OSError:
            report['passed'] = False
            report['failures'].append('Report output could not be written')
            encoded = json.dumps(report, ensure_ascii=False, indent=2, allow_nan=False) + '\n'
    print(encoded, end='')
    return 0 if report['passed'] else 1


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