business-assistant/scripts/deploy_check.py

"""Read-only deployment diagnostics; never print credentials or provider payloads."""
import json
from collections import Counter
from pathlib import Path


def result(ok, code, **details):
    return {'ok': bool(ok), 'code': code, **details}


def milvus(root, client_factory=None):
    from src.knowledge_profiles import for_agent, prepared_path
    from src.step2_knowledge_ingest.ingest import config

    settings = config(root / '.env')
    try:
        business = _milvus_snapshot(prepared_path(root, 'business'), settings, client_factory)
    except Exception:
        business = result(False, 'business_check_failed')
    agents = {'business': business}
    talent = prepared_path(root, 'talent')
    if talent.exists():
        try:
            agents['talent'] = _milvus_snapshot(talent, for_agent(settings, 'talent'), client_factory)
        except Exception:
            # Optional-agent errors must not hide the existing business report.
            agents['talent'] = result(False, 'talent_check_failed')
    healthy = all(row['ok'] for row in agents.values())
    code = business['code'] if len(agents) == 1 else 'matched' if healthy else 'knowledge_difference'
    return {**business, 'ok': healthy, 'code': code,
            'agents': agents,
            'expected_total_all_agents': sum(row.get('expected_total', 0) for row in agents.values()),
            'indexed_total_all_agents': sum(row.get('indexed_total', 0) for row in agents.values())}


def _milvus_snapshot(prepared, settings, client_factory):
    from src.knowledge_profiles import check_records
    from src.step2_knowledge_ingest.ingest import content_parts, load, route

    manifest_path = prepared / 'ingestion/manifest.json'
    if not manifest_path.is_file():
        return result(False, 'manifest_missing')
    manifest = json.loads(manifest_path.read_text(encoding='utf-8'))
    if manifest.get('complete') is not True or not isinstance(manifest.get('databases'), dict):
        return result(False, 'manifest_incomplete')
    expected_segments = Counter()
    local_records = Counter()
    records = load(prepared)
    check_records(records, settings)
    for record in records:
        kind = route(record)
        expected_segments[kind] += len(content_parts(record))
        local_records[kind] += 1
    if dict(local_records) != manifest['counts'] or set(manifest['databases']) != set(local_records):
        return result(False, 'manifest_local_mismatch')
    if settings.get('KNOWLEDGE_AGENT') == 'talent' and set(local_records) != {'article', 'qa', 'policy'}:
        return result(False, 'talent_libraries_missing')
    if client_factory is None:
        from pymilvus import MilvusClient
        client_factory = MilvusClient
    client = client_factory(uri=settings['MILVUS_URI'], timeout=settings['MILVUS_TIMEOUT_SECONDS'])
    try:
        return _milvus_counts(client, settings, manifest, expected_segments)
    finally:
        close = getattr(client, 'close', None)
        if close is not None:
            close()


def _milvus_counts(client, settings, manifest, expected_segments):
    timeout = settings['MILVUS_TIMEOUT_SECONDS']
    available = set(client.list_databases(timeout=timeout))
    libraries = {}
    for kind, spec in manifest['databases'].items():
        expected = manifest['counts'].get(kind)
        database = spec['database']
        if database != settings['MILVUS_DB_PREFIX'] + '_' + kind or database not in available:
            libraries[kind] = result(False, 'database_missing_or_mismatched', database=database, expected_records=expected)
            continue
        client.use_database(database, timeout=timeout)
        names = set(client.list_collections(timeout=timeout))
        index_name, content_name = spec['index_collection'], spec['content_collection']
        if index_name not in names or content_name not in names:
            libraries[kind] = result(False, 'collection_missing', database=database, expected_records=expected)
            continue
        def count(name):
            return client.query(name, filter='', output_fields=['count(*)'],
                                consistency_level='Strong', timeout=timeout)[0]['count(*)']
        indexed, segments = count(index_name), count(content_name)
        expected_parts = expected_segments[kind]
        ok = indexed == expected and segments == expected_parts
        libraries[kind] = result(ok, 'matched' if ok else 'count_mismatch',
                                 database=database,
                                 expected_records=expected, indexed_records=indexed,
                                 expected_content_segments=expected_parts, content_segments=segments,
                                 index_delta=indexed - expected, content_delta=segments - expected_parts)
    return result(all(row['ok'] for row in libraries.values()) and bool(libraries),
                  'matched' if all(row['ok'] for row in libraries.values()) and libraries else 'knowledge_difference',
                  snapshot=manifest.get('snapshot'), libraries=libraries,
                  expected_total=sum(manifest['counts'].values()),
                  indexed_total=sum(row.get('indexed_records', 0) for row in libraries.values()))


def redis_check(root, client_factory=None):
    from src.step3_question_answer.context import load_context_config

    settings = load_context_config(root / '.env')
    if client_factory is None:
        from redis import Redis
        client_factory = Redis.from_url
    client = client_factory(settings.url, socket_connect_timeout=5, socket_timeout=5,
                            decode_responses=False, protocol=2)
    try:
        connected = bool(client.ping())
        return result(connected, 'connected' if connected else 'ping_failed')
    finally:
        client.close()


def dms(root, session_factory=None, company_factory=None):
    """Verify the five DMS columns used by the app through read-only metadata calls."""
    from src.step5_session_analysis.dms import DmsClient, load_dms_config
    from src.step3_question_answer.dms_company import DmsCompanyClient
    from src.dms_schema import names as schema_names

    session_factory = session_factory or DmsClient
    company_factory = company_factory or DmsCompanyClient
    schema_path = root / 'config/dms_schema.json'
    checks = {}
    try:
        session = session_factory(load_dms_config(root / '.env', schema_path=schema_path))
        session.connect()
        checks['sessions'] = result(True, 'usable', columns=3)
    except Exception as exc:
        checks['sessions'] = result(False, safe_code(exc))
    try:
        from src.dms_schema import load_schema
        schema = load_schema(schema_path)
        config = load_dms_config(root / '.env', scope_names=schema_names(schema, ('basic', 'honor')),
                                 schema_path=schema_path)
        company = company_factory(config)
        company.connect(config.schema['columns']['basic']['id'])
        company.connect(config.schema['columns']['honor']['id'])
        checks['company'] = result(True, 'usable', columns=2)
    except Exception as exc:
        checks['company'] = result(False, safe_code(exc))
    healthy = all(row['ok'] for row in checks.values())
    return result(healthy, 'usable' if healthy else 'api_failure', scopes=checks)


def deepseek(root, send=None):
    from src.step1_knowledge_prepare.summarize import load_config, transport

    config = load_config(root / '.env')
    payload = {'model': config.model, 'messages': [{'role': 'user', 'content': 'Reply with OK.'}],
               'max_tokens': 128, 'stream': False}
    response = (send or transport)(payload, config)
    choices = response.get('choices') if isinstance(response, dict) else None
    content = choices[0].get('message', {}).get('content') if isinstance(choices, list) and choices and isinstance(choices[0], dict) else None
    return result(isinstance(content, str) and bool(content.strip()),
                  'usable' if isinstance(content, str) and content.strip() else 'empty_model_response')


def qcc(root, company_name, client=None):
    from src.step3_question_answer.qcc import QccClient, load_config

    client = client or QccClient(load_config(root / '.env'))
    checks = {}
    try:
        search = client.search(company_name)
        checks['886'] = result(True, 'usable', candidates=len(search['candidates']))
    except Exception as exc:
        checks['886'] = result(False, safe_code(exc))
        search = None
    exact = [row for row in search['candidates'] if row['Name'] == company_name] if search else []
    if len(exact) != 1:
        code = 'exact_company_not_found' if search else 'search_failed'
        checks['410'] = result(False, code)
        checks['1001'] = result(False, code)
    else:
        company = exact[0]
        try:
            client.basic(company)
            checks['410'] = result(True, 'usable')
        except Exception as exc:
            checks['410'] = result(False, safe_code(exc))
        try:
            honors = client.honors(company)
            checks['1001'] = result(honors['complete'], 'usable' if honors['complete'] else safe_code(Exception(honors['error'])),
                                    records=len(honors['records']), pages=len(honors['sources']))
        except Exception as exc:
            checks['1001'] = result(False, safe_code(exc))
    return result(all(row['ok'] for row in checks.values()), 'usable' if all(row['ok'] for row in checks.values()) else 'api_failure',
                  interfaces=checks)


def safe_code(exc):
    code = str(exc)
    allowed = {'provider_failure', 'qcc_auth_failed', 'invalid_qcc_configuration', 'invalid_result', 'invalid_search_result', 'invalid_candidate',
               'duplicate_candidate', 'company_identity_mismatch', 'invalid_honor_result',
               'invalid_honor_data', 'inconsistent_honor_result', 'inconsistent_honor_paging',
               'page_limit', 'network_failure', 'invalid_json', 'response_too_large',
               'invalid_dms_configuration', 'invalid_dms_login_url', 'dms_scope_denied', 'dms_redirect_refused',
               'dms_transport_failed', 'dms_invalid_envelope', 'dms_permission_denied',
               'dms_token_invalid', 'dms_service_failed', 'dms_auth_retry_later',
               'dms_login_failed', 'dms_direct_api_requires_token', 'dms_column_mismatch',
               'dms_parent_mismatch', 'dms_schema_mismatch', 'dms_write_scope_denied',
               'dms_operation_denied', 'invalid_deepseek_base_url', 'missing_deepseek_configuration',
               'request_timeout', 'response_incomplete', 'invalid_response_json'}
    return code if code in allowed or (code.startswith('http_') and code[5:].isdigit()) else 'request_failed'


def run(root, company_name=None, action='check'):
    root = Path(root)
    checks = {}
    probes = {'milvus': lambda: milvus(root), 'redis': lambda: redis_check(root),
              'dms': lambda: dms(root),
              'deepseek': lambda: deepseek(root), 'qcc': lambda: qcc(root, company_name)}
    names = tuple(probes) if action == 'check' else (action.removeprefix('check-'),)
    if not names or any(name not in probes for name in names):
        raise ValueError('invalid_check_action')
    for name in names:
        try:
            checks[name] = probes[name]()
        except Exception as exc:
            checks[name] = result(False, safe_code(exc))
    report = {'ok': all(row['ok'] for row in checks.values()), 'checks': checks}
    print(json.dumps(report, ensure_ascii=False), flush=True)
    return 0 if report['ok'] else 1