tools/diagnostics.py

#!/usr/bin/env python3
"""tools/diagnostics.py

Consolidated diagnostics for KB, Chroma, FAISS and LLM model verification.

Usage:
  python tools/diagnostics.py check_kb --org ORG_ID
  python tools/diagnostics.py build_faiss --org ORG_ID
  python tools/diagnostics.py verify_models
  python tools/diagnostics.py test_providers
  python tools/diagnostics.py all --org ORG_ID

This file consolidates the checks previously spread across several small
tooling scripts into a single, importable module with subcommands.
"""
import argparse
import json
import os
import sys
import traceback
from pathlib import Path

# Ensure project root is importable
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))


def safe_json(obj):
    try:
        return json.dumps(obj, default=str, indent=2)
    except Exception:
        return str(obj)


def check_kb(org_id: str):
    """Inspect MongoDB kb_documents/kb_chunks and local Chroma for an org."""
    print(f"\n=== Qualora KB Diagnostic for org {org_id} ===\n")
    try:
        from core import get_db, CHROMA_PATH
        from bson import ObjectId
    except Exception as e:
        print("Failed to import core or bson. Ensure you run inside the project venv.")
        print(e)
        return 2

    try:
        db = get_db()
        if db is None:
            print("MongoDB: get_db() returned None (check MONGODB_URI / network)")
        else:
            try:
                safe_oid = ObjectId(str(org_id))
            except Exception:
                print(f"Invalid ObjectId provided: {org_id}")
                safe_oid = None

            if safe_oid:
                docs_count = db.kb_documents.count_documents({"org_id": safe_oid})
                chunks_count = db.kb_chunks.count_documents({"org_id": safe_oid})
                print(f"MongoDB: kb_documents for org {org_id}: {docs_count}")
                print(f"MongoDB: kb_chunks for org {org_id}: {chunks_count}")

                if docs_count:
                    sample_docs = list(db.kb_documents.find({"org_id": safe_oid}).limit(5))
                    for d in sample_docs:
                        d['_id'] = str(d.get('_id'))
                        if 'gridfs_id' in d:
                            d['gridfs_id'] = str(d['gridfs_id'])
                    print("Sample kb_documents:\n", safe_json(sample_docs))

                if chunks_count:
                    sample_chunks = list(db.kb_chunks.find({"org_id": safe_oid}).limit(5))
                    for c in sample_chunks:
                        c['_id'] = str(c.get('_id'))
                        if 'doc_id' in c:
                            c['doc_id'] = str(c.get('doc_id'))
                        if 'text' in c and isinstance(c['text'], str):
                            c['text'] = c['text'][:200].replace('\n', ' ') + ('...' if len(c['text']) > 200 else '')
                    print("Sample kb_chunks:\n", safe_json(sample_chunks))
            else:
                print("MongoDB: Skipping org-specific queries due to invalid ObjectId.")
    except Exception:
        print("MongoDB check failed:")
        traceback.print_exc()

    # --- Chroma checks ---
    try:
        print('\n--- Chroma DB checks ---')
        try:
            import chromadb
            from chromadb.config import Settings
        except Exception as e:
            print('chromadb import failed:', e)
            print('Install chromadb in your venv to run Chroma checks: pip install chromadb')
            return 1

        try:
            client = chromadb.PersistentClient(path=CHROMA_PATH, settings=Settings(anonymized_telemetry=True))
        except Exception:
            client = chromadb.PersistentClient(path=CHROMA_PATH)

        cols = [c.name for c in client.list_collections()]
        print('Chroma collections:', cols)

        target_col = f"org_{str(org_id)}"
        if target_col in cols:
            collection = client.get_collection(target_col)
            try:
                count = collection.count()
            except Exception:
                try:
                    count = len(collection.get(limit=100).get('ids', []))
                except Exception:
                    count = 'unknown'
            print(f"Chroma: collection '{target_col}' exists; count: {count}")
            try:
                sample = collection.get(limit=3)
                print('Sample documents (documents):', sample.get('documents'))
                print('Sample metadatas:', sample.get('metadatas'))
                print('Sample ids:', sample.get('ids'))
            except Exception as e:
                print('Failed to read sample from collection:', e)
        else:
            print(f"Chroma: collection '{target_col}' not found (available: {cols})")

    except Exception:
        print("Chroma check failed:")
        traceback.print_exc()

    print('\n=== Diagnostic complete ===\n')
    return 0


def build_faiss_index(org_id: str, sample_path: str = 'docs/samples/human_chat.txt') -> tuple[int, str]:
    """Build a local FAISS index for L3 testing and save to /tmp/faiss_{org_id}.
    Returns (exit_code, path_or_message).
    """
    try:
        from langchain_community.vectorstores import FAISS
    except Exception as e:
        msg = f'FAISS import failed: {e}'
        print(msg)
        return 2, msg

    try:
        import importlib
        import services.rag as rag
        importlib.reload(rag)
    except Exception as e:
        msg = f'Failed to import services.rag: {e}'
        print(msg)
        return 2, msg

    embed = rag.get_local_embed()
    if not embed:
        msg = 'Local embed model unavailable; cannot build FAISS index.'
        print(msg)
        return 2, msg

    try:
        with open(sample_path, 'r', encoding='utf-8') as f:
            sample = f.read()
    except Exception:
        sample = 'sample text for faiss'

    chunk_texts = [sample[i:i+500] for i in range(0, len(sample), 500)]
    if not chunk_texts:
        chunk_texts = [' ']

    faiss_path = f'/tmp/faiss_{org_id}'
    try:
        import shutil
        if os.path.exists(faiss_path):
            shutil.rmtree(faiss_path)
        os.makedirs(faiss_path, exist_ok=True)
    except Exception:
        pass

    print('Embedding', len(chunk_texts), 'chunks')
    try:
        vs = FAISS.from_texts(chunk_texts, embed)
        vs.save_local(faiss_path)
        print('FAISS index saved to', faiss_path)
        return 0, faiss_path
    except Exception as e:
        msg = f'Failed to build/save FAISS index: {e}'
        print(msg)
        return 3, msg


def verify_llm_models():
    """Run the existing verify logic for OpenRouter, Groq and HuggingFace."""
    # reuse logic adapted from the earlier tools/verify_llm_models.py
    import time
    import httpx
    import traceback
    from core import OPENROUTER_API_KEY, GROQ_API_KEY, HF_SPACE_TOKEN

    OPENROUTER_DEFAULTS = ["openrouter/free", "qwen/qwen3.6-plus:free"]
    GROQ_DEFAULTS = ["llama-3.3-70b-versatile", "llama-3.1-8b-instant"]
    HF_DEFAULTS = ["Qwen/Qwen2.5-7B-Instruct", "mistralai/Mistral-7B-Instruct-v0.3", "meta-llama/Llama-3.1-8B-Instruct"]

    OPENROUTER_MODELS = [m.strip() for m in os.environ.get('OPENROUTER_MODELS', ','.join(OPENROUTER_DEFAULTS)).split(',') if m.strip()]
    GROQ_MODELS = [m.strip() for m in os.environ.get('GROQ_MODELS', ','.join(GROQ_DEFAULTS)).split(',') if m.strip()]
    HF_MODELS = [m.strip() for m in os.environ.get('HF_MODELS', ','.join(HF_DEFAULTS)).split(',') if m.strip()]

    RESULT = {"openrouter": {}, "groq": {}, "hf": {}}

    def verify_openrouter():
        out = {m: {"found": False, "notes": "skipped"} for m in OPENROUTER_MODELS}
        if not OPENROUTER_API_KEY:
            for m in out:
                out[m]["notes"] = "OPENROUTER_API_KEY not set"
            return out

        url = "https://openrouter.ai/api/v1/models"
        headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}"}
        try:
            with httpx.Client(timeout=15.0) as client:
                r = client.get(url, headers=headers)
                if r.status_code == 401:
                    for m in out:
                        out[m]["notes"] = "unauthorized"
                    return out
                if r.status_code == 429:
                    for m in out:
                        out[m]["notes"] = "rate_limited"
                    return out
                r.raise_for_status()
                data = r.json()

                candidates = []
                if isinstance(data, list):
                    for it in data:
                        if isinstance(it, dict):
                            candidates.append(it.get('id') or it.get('model') or it.get('name'))
                        else:
                            candidates.append(str(it))
                elif isinstance(data, dict):
                    if 'models' in data and isinstance(data['models'], list):
                        for it in data['models']:
                            if isinstance(it, dict):
                                candidates.append(it.get('id') or it.get('model') or it.get('name'))
                            else:
                                candidates.append(str(it))
                    else:
                        for k in ('id','model','name'):
                            if k in data:
                                candidates.append(data[k])

                candidates = [c for c in candidates if c]
                for m in OPENROUTER_MODELS:
                    if any(m == c or m in str(c) for c in candidates):
                        out[m]["found"] = True
                        out[m]["notes"] = "found"
                    else:
                        out[m]["notes"] = "not_found_in_list"
                out["_raw_count"] = len(candidates)
                out["_sample"] = candidates[:10]
                if len(candidates) == 0 or any(out[m]["notes"] == "not_found_in_list" for m in OPENROUTER_MODELS):
                    for m in OPENROUTER_MODELS:
                        if out[m]["found"]:
                            continue
                        try:
                            payload = {
                                "model": m,
                                "messages": [
                                    {"role": "system", "content": "You are a lightweight model availability checker."},
                                    {"role": "user",   "content": "Return {\"ok\":1}"}
                                ],
                                "temperature": 0.0,
                                "max_tokens": 8
                            }
                            r2 = client.post("https://openrouter.ai/api/v1/chat/completions", json=payload, headers=headers)
                            if r2.status_code == 429:
                                out[m]["notes"] = "rate_limited"
                                continue
                            if r2.status_code in (401, 403):
                                out[m]["notes"] = "unauthorized"
                                continue
                            r2.raise_for_status()
                            try:
                                dd = r2.json()
                                if isinstance(dd, dict) and dd.get('choices'):
                                    out[m]["found"] = True
                                    out[m]["notes"] = "call_ok"
                                else:
                                    out[m]["notes"] = "no_choices"
                            except Exception as je:
                                out[m]["notes"] = f"invalid_json: {str(je)[:120]}"
                        except Exception as e:
                            out[m]["notes"] = f"error: {str(e)[:200]}"

                return out
        except Exception as e:
            for m in out:
                out[m]["notes"] = f"error: {str(e)[:200]}"
            out["_error"] = traceback.format_exc()
            return out

    def verify_groq():
        out = {m: {"found": False, "notes": "skipped"} for m in GROQ_MODELS}
        if not GROQ_API_KEY:
            for m in out:
                out[m]["notes"] = "GROQ_API_KEY not set"
            return out

        try:
            from services.audit_engine import get_groq_client
            groq_client = get_groq_client()
            if not groq_client:
                for m in out:
                    out[m]["notes"] = "Groq client unavailable"
                return out

            try:
                if hasattr(groq_client, 'models') and hasattr(groq_client.models, 'list'):
                    models_list = groq_client.models.list()
                    candidates = []
                    for it in models_list:
                        try:
                            candidates.append(getattr(it, 'id', None) or getattr(it, 'name', None) or str(it))
                        except Exception:
                            candidates.append(str(it))
                    for m in GROQ_MODELS:
                        if any(m == c or m in str(c) for c in candidates):
                            out[m]["found"] = True
                            out[m]["notes"] = "found"
                        else:
                            out[m]["notes"] = "not_found_in_list"
                    out["_sample"] = candidates[:10]
                    return out
            except Exception:
                pass

            for m in GROQ_MODELS:
                try:
                    resp = groq_client.chat.completions.create(
                        model=m,
                        messages=[{"role":"system","content":"You are a test."},{"role":"user","content":"Respond with {\"ok\":1}"}],
                        temperature=0.0,
                        max_tokens=8
                    )
                    out[m]["found"] = True
                    out[m]["notes"] = "call_ok"
                except Exception as e:
                    se = str(e).lower()
                    if '429' in se or 'rate' in se or 'quota' in se:
                        out[m]["notes"] = "rate_limited"
                    else:
                        out[m]["notes"] = f"error: {str(e)[:200]}"
            return out
        except Exception as e:
            for m in out:
                out[m]["notes"] = f"error: {str(e)[:200]}"
            out["_error"] = traceback.format_exc()
            return out

    def verify_hf():
        out = {m: {"found": False, "notes": "skipped"} for m in HF_MODELS}
        try:
            with httpx.Client(timeout=15.0) as client:
                for m in HF_MODELS:
                    url = f"https://huggingface.co/{m}"
                    try:
                        r = client.get(url)
                        if r.status_code == 200:
                            out[m]["found"] = True
                            out[m]["notes"] = "page_200"
                        elif r.status_code == 404:
                            out[m]["notes"] = "not_found_404"
                        elif r.status_code == 429:
                            out[m]["notes"] = "rate_limited"
                        else:
                            out[m]["notes"] = f"http_{r.status_code}"
                    except Exception as e:
                        out[m]["notes"] = f"error: {str(e)[:200]}"
            return out
        except Exception as e:
            for m in out:
                out[m]["notes"] = f"error: {str(e)[:200]}"
            out["_error"] = traceback.format_exc()
            return out

    print('Verifying LLM cascade model names...')
    RESULT['openrouter'] = verify_openrouter()
    RESULT['groq'] = verify_groq()
    RESULT['hf'] = verify_hf()

    print(json.dumps(RESULT, indent=2))
    bad = False
    for prov, mapping in RESULT.items():
        for k, v in mapping.items():
            if k.startswith('_'):
                continue
            if v.get('notes') in ('not_found_in_list', 'not_found_404') or (v.get('notes', '').startswith('error')):
                bad = True
    if bad:
        return 2
    print('\nAll verified or skipped.')
    return 0


def test_llm_providers():
    """Smoke-test configured LLM providers by invoking audit_engine provider callables."""
    import services.audit_engine as ae
    from core import OPENROUTER_API_KEY, GROQ_API_KEY, HF_SPACE_TOKEN

    PROMPT = (
        "Please return only a compact JSON object matching the audit schema. "
        "Include keys: summary (string), agent_f1_score (number), satisfaction_prediction (Medium), "
        "compliance_risk (Amber), quality_matrix (with numeric fields), compliance_flags (empty list), "
        "behavioral_nudges (empty list), emotions (agent and customer neutral)."
    )

    def safe_call(fn, name, prompt):
        print('\n=== Testing', name, '===')
        try:
            import time
            t0 = time.time()
            res = fn(prompt)
            dt = time.time() - t0
            if isinstance(res, tuple) and len(res) >= 2:
                text, provider = res[0], res[1]
            else:
                text = str(res)
                provider = name
            print('Provider:', provider)
            print('Elapsed: %.2fs' % dt)
            print('Preview:', repr(text)[:400])
            return True
        except Exception as e:
            print('ERROR:', str(e))
            traceback.print_exc()
            return False

    results = {}
    if OPENROUTER_API_KEY:
        print('OPENROUTER_API_KEY found — attempting _call_openrouter')
        results['openrouter'] = safe_call(ae._call_openrouter, 'openrouter', PROMPT)
    else:
        print('OPENROUTER_API_KEY not set — skipping openrouter')
        results['openrouter'] = None

    if getattr(ae, 'GROQ_AVAILABLE', False) and GROQ_API_KEY:
        print('Groq SDK available and GROQ_API_KEY set — attempting _call_groq')
        results['groq'] = safe_call(ae._call_groq, 'groq', PROMPT)
    else:
        print('Groq not configured or SDK missing — skipping groq')
        results['groq'] = None

    if HF_SPACE_TOKEN:
        print('HF_SPACE_TOKEN found — attempting _call_hf_inference')
        results['hf_inference'] = safe_call(ae._call_hf_inference, 'hf_inference', PROMPT)
    else:
        print('HF_SPACE_TOKEN not set — skipping HF Inference')
        results['hf_inference'] = None

    print('\n=== Summary ===')
    for k, v in results.items():
        status = 'SKIPPED' if v is None else ('OK' if v else 'FAILED')
        print(f'{k}: {status}')

    failed = any(v is False for v in results.values() if v is not None)
    if failed:
        return 1
    print('\nAll configured providers either OK or skipped.')

    # verify full cascade
    try:
        print('\n=== Running full LLM cascade to verify fallback behavior ===')
        cascade_out = ae._run_llm_cascade(PROMPT)
        print('Cascade selected provider:', cascade_out[1], 'tier:', cascade_out[2])
        print('Cascade preview:', repr(cascade_out[0])[:400])
    except Exception as e:
        print('Cascade failed:', e)
        return 2

    return 0


def main():
    parser = argparse.ArgumentParser(description='Diagnostics helper: KB/Chroma/FAISS/LLM verifier')
    sub = parser.add_subparsers(dest='cmd')

    p_check = sub.add_parser('check_kb')
    p_check.add_argument('--org', required=True)

    p_faiss = sub.add_parser('build_faiss')
    p_faiss.add_argument('--org', required=True)
    p_faiss.add_argument('--sample', default='docs/samples/human_chat.txt')

    sub.add_parser('verify_models')
    sub.add_parser('test_providers')
    p_all = sub.add_parser('all')
    p_all.add_argument('--org', required=True)

    args = parser.parse_args()
    if args.cmd == 'check_kb':
        return check_kb(args.org)
    if args.cmd == 'build_faiss':
        code, msg = build_faiss_index(args.org, args.sample)
        return code
    if args.cmd == 'verify_models':
        return verify_llm_models()
    if args.cmd == 'test_providers':
        return test_llm_providers()
    if args.cmd == 'all':
        rc = check_kb(args.org)
        if rc != 0:
            print('check_kb returned non-zero:', rc)
        rc2 = verify_llm_models()
        rc3 = test_llm_providers()
        code = 0 if (rc == 0 and rc2 == 0 and rc3 == 0) else 2
        return code


if __name__ == '__main__':
    rc = main()
    if isinstance(rc, int) and rc != 0:
        raise SystemExit(rc)