AIyaodianzhushou/diagnose_online.py

# -*- coding: utf-8 -*-
"""
Production Diagnostic Script
Deploy to server and run to trace the full RAG pipeline for any query.

Usage:
  python diagnose_online.py "青霉素过敏处理"
  python diagnose_online.py "青霉素过敏处理" --top-k 30
  python diagnose_online.py "青霉素过敏处理" --call-api  # also call the API
"""
import sys
import os
import json
import re
import asyncio
import argparse
import math
from pathlib import Path

# ============================================================
# Environment Setup
# ============================================================
def load_env():
    """Load .env file from project root."""
    candidates = [
        Path(__file__).resolve().parent / ".env",
        Path.cwd() / ".env",
        Path("/opt/pharmacopoeia-ai/.env"),
    ]
    for env_file in candidates:
        if env_file.exists():
            with open(env_file, encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if line and not line.startswith("#") and "=" in line:
                        key, _, val = line.partition("=")
                        os.environ.setdefault(key.strip(), val.strip())
            print(f"[ENV] Loaded from {env_file}")
            return True
    print("[ENV] WARNING: No .env file found, using env vars directly")
    return False


load_env()

# ============================================================
# Intent Classification (copy from retriever.py)
# ============================================================
def classify_intent(query: str) -> str:
    q = query.strip()
    negation_patterns = [
        r"不是", r"没有", r"并非", r"不算", r"不属于",
        r"这不是", r"我没有", r"不包含", r"不涉及",
    ]
    if any(re.search(pat, q) for pat in negation_patterns):
        return "drug_query"

    usage_keywords = [
        "怎么吃", "吃多少", "怎么用", "一天几次", "多长时间",
        "能一起吃", "孕妇能用", "儿童用量", "哺乳期",
        "饭前还是饭后", "空腹", "过量", "漏服", "停药",
        "副作用多大", "伤肝吗", "伤肾吗", "安全吗",
    ]
    safety_sections = [
        "副作用", "不良反应", "禁忌", "注意事项",
        "能不能", "可以吗", "会不会",
    ]
    regulation_keywords = [
        "凡例", "通则规定", "制剂通则",
        "一般规定", "通用技术要求", "检验方法通则",
    ]
    exam_keywords = [
        "执业药师考试", "考点", "历年真题", "考试大纲",
        "高频考点", "报名时间",
    ]
    symptom_keywords = [
        "吃了什么药", "吃什么药", "该吃", "推荐用药", "推荐下用药",
        "买什么药", "推荐什么药", "用什么药", "用药建议",
        "发烧", "咳嗽", "感冒", "腹泻", "头疼", "头痛",
        "嗓子疼", "流鼻涕", "鼻塞", "肚子疼", "胃疼",
        "过敏", "皮肤痒", "失眠", "便秘", "牙疼",
        "体温", "多少度", "退烧", "止痛", "止泻",
    ]

    if any(kw in q for kw in usage_keywords):
        return "usage_guide"
    if any(kw in q for kw in safety_sections):
        return "usage_guide"
    if any(kw in q for kw in symptom_keywords):
        return "symptom_advice"
    if any(kw in q for kw in regulation_keywords):
        return "regulation"
    if any(kw in q for kw in exam_keywords):
        return "exam_tutor"
    return "drug_query"


# ============================================================
# Tokenization (copy from reranker.py)
# ============================================================
def tokenize_query(query: str) -> set:
    tokens = set()
    for n in range(2, 5):
        for i in range(len(query) - n + 1):
            seg = query[i:i + n]
            if all('一' <= c <= '鿿' for c in seg):
                tokens.add(seg)
    for word in re.findall(r'[a-zA-Z0-9]+', query):
        tokens.add(word.lower())
    return tokens


def keyword_overlap_score(query_terms: set, content: str) -> float:
    if not query_terms:
        return 0.0
    matched = sum(1 for t in query_terms if t in content)
    return matched / len(query_terms)


# ============================================================
# Database Connection & Vector Search
# ============================================================
def get_db_url() -> str:
    host = os.environ.get("POSTGRES_HOST", "localhost")
    port = os.environ.get("POSTGRES_PORT", "5432")
    db = os.environ.get("POSTGRES_DB", "pharmacopoeia")
    user = os.environ.get("POSTGRES_USER", "postgres")
    password = os.environ.get("POSTGRES_PASSWORD", "postgres")
    return f"postgresql+asyncpg://{user}:{password}@{host}:{port}/{db}"


async def get_query_embedding(query: str) -> list:
    """Call DashScope embedding API."""
    import httpx
    api_key = os.environ.get("QWEN_API_KEY", "")
    if not api_key:
        raise RuntimeError("QWEN_API_KEY not set")

    url = "https://dashscope.aliyuncs.com/api/v1/services/embeddings/text-embedding/text-embedding"
    async with httpx.AsyncClient(timeout=30) as client:
        resp = await client.post(
            url,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {api_key}",
            },
            json={
                "model": "text-embedding-v3",
                "input": {"texts": [query]},
                "parameters": {"text_type": "query"},
            },
        )
        data = resp.json()
    if data.get("code") and data.get("code") != "":
        raise RuntimeError(f"Embedding API error: {data.get('message')}")
    return data["output"]["embeddings"][0]["embedding"]


async def vector_search(query_vec: list, top_k: int = 20) -> list:
    """Run pgvector cosine similarity search."""
    from sqlalchemy.ext.asyncio import create_async_engine
    from sqlalchemy import text

    db_url = get_db_url()
    vec_str = "[" + ",".join(str(v) for v in query_vec) + "]"

    engine = create_async_engine(db_url, pool_size=5, max_overflow=10, pool_pre_ping=True)
    try:
        async with engine.connect() as conn:
            result = await conn.execute(
                text("""
                    SELECT content, source, drug_id, section,
                           1 - (vec <=> CAST(:qv AS vector)) AS similarity
                    FROM drug_chunks
                    WHERE vec IS NOT NULL
                    ORDER BY vec <=> CAST(:qv AS vector)
                    LIMIT :k
                """),
                {"qv": vec_str, "k": top_k},
            )
            rows = result.fetchall()

        docs = []
        for row in rows:
            docs.append({
                "content": row[0],
                "source": row[1],
                "drug_id": row[2],
                "section": row[3],
                "score": float(row[4]),
            })
        return docs
    finally:
        await engine.dispose()


async def get_db_stats() -> dict:
    """Get statistics about the drug_chunks table."""
    from sqlalchemy.ext.asyncio import create_async_engine
    from sqlalchemy import text

    db_url = get_db_url()
    engine = create_async_engine(db_url)
    try:
        async with engine.connect() as conn:
            # Total chunks
            r = await conn.execute(text("SELECT COUNT(*) FROM drug_chunks WHERE vec IS NOT NULL"))
            total_chunks = r.fetchone()[0]

            # By source
            r = await conn.execute(text("""
                SELECT source, COUNT(*) as cnt
                FROM drug_chunks
                WHERE vec IS NOT NULL
                GROUP BY source
                ORDER BY cnt DESC
                LIMIT 20
            """))
            by_source = [(row[0], row[1]) for row in r.fetchall()]

            # By section
            r = await conn.execute(text("""
                SELECT section, COUNT(*) as cnt
                FROM drug_chunks
                WHERE vec IS NOT NULL
                GROUP BY section
                ORDER BY cnt DESC
                LIMIT 20
            """))
            by_section = [(row[0], row[1]) for row in r.fetchall()]

            # Drugs with penicillin in name
            r = await conn.execute(text("""
                SELECT DISTINCT d.name, d.category, d.source_volume
                FROM drugs d
                WHERE d.name LIKE '%青霉素%'
            """))
            penicillin_drugs = [(row[0], row[1], row[2]) for row in r.fetchall()]

            # Drugs with 清开灵 in name
            r = await conn.execute(text("""
                SELECT DISTINCT d.name, d.category, d.source_volume
                FROM drugs d
                WHERE d.name LIKE '%清开灵%'
            """))
            qkl_drugs = [(row[0], row[1], row[2]) for row in r.fetchall()]

            # Count drugs that have allergy in their sections
            r = await conn.execute(text("""
                SELECT COUNT(DISTINCT drug_id)
                FROM drug_chunks
                WHERE content LIKE '%过敏%' AND vec IS NOT NULL
            """))
            allergy_chunk_count = r.fetchone()[0]

            # Count total unique drugs
            r = await conn.execute(text("SELECT COUNT(*) FROM drugs WHERE is_active = TRUE"))
            total_drugs = r.fetchone()[0]

            return {
                "total_chunks": total_chunks,
                "total_drugs": total_drugs,
                "allergy_chunk_count": allergy_chunk_count,
                "by_source": by_source,
                "by_section": by_section,
                "penicillin_drugs": penicillin_drugs,
                "qkl_drugs": qkl_drugs,
            }
    finally:
        await engine.dispose()


# ============================================================
# Reranker Simulation
# ============================================================
def simulate_reranker(query: str, documents: list, top_k: int = 5) -> list:
    """Simulate the reranker behavior on retrieval results."""
    query_terms = tokenize_query(query)
    MIN_THRESHOLD = 0.3

    # Step 1: Filter low similarity
    filtered = [d for d in documents if d.get("score", 0) >= MIN_THRESHOLD]

    # Step 2: Keyword bonus
    for doc in filtered:
        content = doc.get("content", "")
        bonus = keyword_overlap_score(query_terms, content)
        doc["keyword_score"] = bonus
        doc["fused_score"] = doc.get("score", 0) + bonus * 0.3

    # Step 3: Sort by fused score
    sorted_docs = sorted(filtered, key=lambda d: d.get("fused_score", 0), reverse=True)

    # Step 4: Dedup
    deduped = []
    seen_texts = []
    for doc in sorted_docs:
        content = doc.get("content", "")
        # Simple dedup: check first 200 chars
        sample = set(content[:200])
        is_dup = False
        for seen in seen_texts[-5:]:
            seen_sample = set(seen[:200])
            intersection = len(sample & seen_sample)
            union = len(sample | seen_sample)
            if union > 0 and intersection / union > 0.8:
                is_dup = True
                break
        if not is_dup:
            deduped.append(doc)
            seen_texts.append(content)

    return deduped[:top_k]


# ============================================================
# Extract drug name from chunk content
# ============================================================
def extract_drug_name(content: str) -> str:
    m = re.match(r'【(.+?) - ', content)
    return m.group(1) if m else "unknown"


# ============================================================
# API Caller
# ============================================================
async def call_production_api(query: str, token: str, api_base: str = "https://pharmacopoeia.kailin.com.cn"):
    """Call the production chat API and capture full SSE response."""
    import httpx

    if not token:
        print("\n[API] ERROR: No token provided (use --token).")
        return None

    print(f"\n[API] Calling {api_base}/api/v1/chat/stream ...")
    print(f"  Query: {query}")

    events = {"intent": None, "status": [], "content": [], "meta": None}

    async with httpx.AsyncClient(timeout=120) as client:
        async with client.stream(
            "POST",
            f"{api_base}/api/v1/chat/stream",
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {token}",
            },
            json={"message": query, "conversation_id": f"diag-{os.urandom(4).hex()}"},
        ) as stream:
            current_event = None
            async for line in stream.aiter_lines():
                if line.startswith("event: "):
                    current_event = line[7:].strip()
                elif line.startswith("data: "):
                    data = line[6:]
                    if current_event == "intent":
                        events["intent"] = data
                        print(f"  [SSE] intent: {data}")
                    elif current_event == "status":
                        events["status"].append(data)
                        print(f"  [SSE] status: {data}")
                    elif current_event == "content":
                        if data != "[DONE]":
                            events["content"].append(data)
                    elif current_event == "meta":
                        try:
                            events["meta"] = json.loads(data)
                            print(f"  [SSE] meta: sources={len(events['meta'].get('sources', []))} items")
                        except json.JSONDecodeError:
                            pass

    answer = "".join(events["content"])
    print(f"\n  Answer length: {len(answer)} chars")

    return {
        "intent": events["intent"],
        "sources": events["meta"].get("sources", []) if events["meta"] else [],
        "answer_preview": answer[:500] + ("..." if len(answer) > 500 else ""),
        "status": events["status"],
    }


# ============================================================
# Main Diagnostic
# ============================================================
async def run_diagnosis(query: str, top_k: int = 20, call_api: bool = False,
                        api_base: str = "https://pharmacopoeia.kailin.com.cn",
                        token: str = ""):
    print("=" * 80)
    print(f"  RAG Pipeline Diagnosis: '{query}'")
    print("=" * 80)

    # ===== Phase 1: Intent Classification =====
    print(f"\n{'='*80}")
    print(f"  PHASE 1: Intent Classification")
    print(f"{'='*80}")
    intent = classify_intent(query)
    print(f"  Query: '{query}'")
    print(f"  Intent: '{intent}'")

    # Show why
    symptom_keywords = [
        "发烧", "咳嗽", "感冒", "腹泻", "头疼", "头痛",
        "嗓子疼", "流鼻涕", "鼻塞", "肚子疼", "胃疼",
        "过敏", "皮肤痒", "失眠", "便秘", "牙疼",
    ]
    matched = [kw for kw in symptom_keywords if kw in query]
    if matched:
        print(f"  Matched symptom keywords: {matched}")
        print(f"  >> This triggers 'symptom_advice' template -> LLM is told to recommend drugs")

    print(f"\n  PROMPT TEMPLATE (intent={intent}):")
    # Inline prompt templates
    PROMPT_MAP = {
        "symptom_advice": "SYMPTOM_ADVICE: 病情评估 -> 用药方案(推荐药品) -> 注意事项 -> 就医指征",
        "drug_query": "DRUG_QUERY: 结论 -> 详细说明 -> 注意事项 -> 来源明细",
        "usage_guide": "USAGE_GUIDE: 结论 -> 用法用量 -> 禁忌 -> 不良反应 -> 注意事项",
        "regulation": "REGULATION: 摘要 -> 原文引用 -> 条款出处 -> 关联条款",
        "exam_tutor": "EXAM_TUTOR: 考点定位 -> 知识要点 -> 记忆技巧 -> 考试频率",
        "no_docs": "NO_DOCS: 结论 -> 详细说明(通用知识) -> 来源说明",
    }
    template = PROMPT_MAP.get(intent, "UNKNOWN")
    print(f"    {template}")

    # ===== Phase 2: Database Statistics =====
    print(f"\n{'='*80}")
    print(f"  PHASE 2: Database Overview")
    print(f"{'='*80}")
    try:
        stats = await get_db_stats()
        print(f"  Total active drugs: {stats['total_drugs']}")
        print(f"  Total chunks (with vec): {stats['total_chunks']}")
        print(f"  Chunks containing '过敏': {stats['allergy_chunk_count']} "
              f"({stats['allergy_chunk_count']/max(stats['total_chunks'],1)*100:.1f}%)")

        print(f"\n  Penicillin-related drugs in DB:")
        for name, cat, vol in stats['penicillin_drugs']:
            print(f"    - {name} ({cat}, {vol})")

        print(f"\n  Qingkailing-related drugs in DB:")
        for name, cat, vol in stats['qkl_drugs']:
            print(f"    - {name} ({cat}, {vol})")

        print(f"\n  Chunks by source (top 10):")
        for src, cnt in stats['by_source'][:10]:
            print(f"    {src}: {cnt} chunks")

        print(f"\n  Chunks by section (top 15):")
        for sec, cnt in stats['by_section'][:15]:
            print(f"    {sec}: {cnt} chunks")
    except Exception as e:
        print(f"  DB STATS FAILED: {e}")
        stats = None

    # ===== Phase 3: Vector Search =====
    print(f"\n{'='*80}")
    print(f"  PHASE 3: Vector Search (pgvector cosine similarity, top-{top_k})")
    print(f"{'='*80}")

    try:
        print(f"  Getting embedding for query...")
        query_vec = await get_query_embedding(query)
        print(f"  Embedding dim: {len(query_vec)}")

        print(f"  Searching drug_chunks...")
        docs = await vector_search(query_vec, top_k=top_k)
        print(f"  Retrieved: {len(docs)} documents")

        if docs:
            print(f"\n  Similarity score range: {docs[0]['score']:.4f} ~ {docs[-1]['score']:.4f}")
            above_threshold = sum(1 for d in docs if d['score'] >= 0.3)
            print(f"  Above reranker threshold (0.3): {above_threshold}/{len(docs)}")

            # Show all results
            print(f"\n  --- VECTOR SEARCH RESULTS ---")
            for i, doc in enumerate(docs):
                drug_name = extract_drug_name(doc['content'])
                content_preview = doc['content'][:120].replace('\n', ' ')
                flag = ""
                if drug_name and "青霉素" in drug_name:
                    flag = " <<< PENICILLIN MATCH"
                elif drug_name and "清开灵" in drug_name:
                    flag = " <<< QINGKAILING MATCH"
                print(f"\n  [{i+1}] score={doc['score']:.4f} | drug='{drug_name}' | section='{doc['section']}'{flag}")
                print(f"      {content_preview}...")
    except Exception as e:
        print(f"  VECTOR SEARCH FAILED: {e}")
        import traceback
        traceback.print_exc()
        docs = []

    # ===== Phase 4: Reranker Simulation =====
    print(f"\n{'='*80}")
    print(f"  PHASE 4: Reranker Simulation (keyword overlap + dedup)")
    print(f"{'='*80}")

    query_terms = tokenize_query(query)
    print(f"  Query tokens ({len(query_terms)}): {sorted(query_terms)}")

    if docs:
        reranked = simulate_reranker(query, docs, top_k=5)
        print(f"\n  After filtering (score>=0.3), keyword bonus, dedup: {len(reranked)} docs")

        print(f"\n  --- RERANKED TOP-5 ---")
        for i, doc in enumerate(reranked):
            drug_name = extract_drug_name(doc['content'])
            content_preview = doc['content'][:150].replace('\n', ' ')
            print(f"\n  [{i+1}] fused_score={doc['fused_score']:.4f} "
                  f"(vec={doc['score']:.4f} + kw_bonus={doc['keyword_score']:.2f}*0.3)")
            print(f"      drug='{drug_name}' | section='{doc['section']}' | source='{doc.get('source','')}'")
            print(f"      {content_preview}")

        # Find Qingkailing results
        qkl_in_results = [d for d in docs if "清开灵" in d.get("content", "")]
        qkl_in_reranked = [d for d in reranked if "清开灵" in d.get("content", "")]
        print(f"\n  Qingkailing in top-{top_k} vector results: {len(qkl_in_results)}")
        print(f"  Qingkailing in top-5 reranked: {len(qkl_in_reranked)}")

        if qkl_in_reranked:
            print(f"\n  *** QINGKAILING MADE IT TO TOP-5! ***")
            for d in qkl_in_reranked:
                print(f"  drug='{extract_drug_name(d['content'])}' "
                      f"score={d['score']:.4f} kw_bonus={d['keyword_score']:.2f}")

        # Show all Qingkailing in vector results
        if qkl_in_results:
            print(f"\n  All Qingkailing entries in vector results:")
            for d in qkl_in_results:
                drug_name = extract_drug_name(d['content'])
                preview = d['content'][:100].replace('\n', ' ')
                print(f"    [{drug_name}/{d['section']}] score={d['score']:.4f}: {preview}...")

    # ===== Phase 5: API Call (optional) =====
    if call_api:
        print(f"\n{'='*80}")
        print(f"  PHASE 5: Production API Call")
        print(f"{'='*80}")
        try:
            api_result = await call_production_api(query, token, api_base)
            if api_result:
                print(f"\n  API Intent: {api_result['intent']}")
                print(f"  API Sources ({len(api_result['sources'])}):")
                for s in api_result['sources']:
                    print(f"    - {s.get('name', '?')} | section={s.get('section', '?')} "
                          f"| score={s.get('score', 0):.4f} | source={s.get('source', '?')}")
                print(f"\n  Answer preview:")
                print(f"  {api_result['answer_preview']}")
        except Exception as e:
            print(f"  API CALL FAILED: {e}")

    # ===== Phase 6: Summary =====
    print(f"\n{'='*80}")
    print(f"  DIAGNOSIS SUMMARY")
    print(f"{'='*80}")

    issues = []

    # Issue 1: Intent
    if intent == "symptom_advice" and "过敏" in query:
        issues.append({
            "severity": "CRITICAL",
            "category": "Intent Classification",
            "detail": f"Query '{query}' classified as 'symptom_advice'. "
                      f"The keyword '过敏' triggers symptom-based drug recommendation, "
                      f"but this query is about managing a known drug allergy, not seeking "
                      f"drugs for allergy symptoms.",
            "fix": "Add known-allergy detection: r'.+过敏(处理|怎么办|急救|替代)' -> usage_guide"
        })

    # Issue 2: Vector search data imbalance
    if stats:
        allergy_pct = stats['allergy_chunk_count'] / max(stats['total_chunks'], 1) * 100
        if allergy_pct > 30:
            issues.append({
                "severity": "HIGH",
                "category": "Data Imbalance",
                "detail": f"{allergy_pct:.0f}% of chunks contain '过敏'. "
                          f"Any allergy-related query will match a huge number of "
                          f"unrelated drug entries.",
                "fix": "Consider section-specific indexing or adding drug-name boost "
                       "in vector search scoring."
            })

        if stats['qkl_drugs'] and not stats['penicillin_drugs']:
            issues.append({
                "severity": "HIGH",
                "category": "Missing Data",
                "detail": f"Qingkailing entries found ({len(stats['qkl_drugs'])}), "
                          f"but no penicillin entries in database. "
                          f"System cannot return relevant results for penicillin queries.",
                "fix": "Ensure chemical drug data (volume 2) is imported."
            })

    # Issue 3: Reranker weakness
    if docs:
        top_scores = [d['score'] for d in docs[:5]]
        score_spread = top_scores[0] - top_scores[-1] if len(top_scores) > 1 else 0
        if score_spread < 0.1:
            issues.append({
                "severity": "MEDIUM",
                "category": "Weak Reranker",
                "detail": f"Top-5 vector scores too close (spread={score_spread:.4f}). "
                          f"Keyword-based reranker cannot effectively distinguish relevance.",
                "fix": "Implement Cross-Encoder reranker (BGE-Reranker-v2-m3)."
            })

    # Issue 4: Always present
    issues.append({
        "severity": "MEDIUM",
        "category": "Prompt Template",
        "detail": "SYMPTOM_ADVICE template forces LLM to recommend drugs "
                  "('用药方案: 方案一, 方案二...'). When query is about "
                  "allergy management, this produces incorrect responses.",
        "fix": "Add constraint in SYMPTOM_ADVICE: if user asks about known drug "
               "allergy, provide management advice, not drug recommendations."
    })

    for issue in issues:
        print(f"\n  [{issue['severity']}] {issue['category']}")
        print(f"  {issue['detail']}")
        print(f"  Fix: {issue['fix']}")

    print(f"\n{'='*80}")
    print(f"  Diagnosis complete. {len(issues)} issues found.")
    print(f"{'='*80}")


# ============================================================
# CLI Entry Point
# ============================================================
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Production RAG pipeline diagnostics"
    )
    parser.add_argument(
        "query", nargs="?", default="青霉素过敏处理",
        help="Query to diagnose (default: 青霉素过敏处理)"
    )
    parser.add_argument(
        "--top-k", type=int, default=20,
        help="Number of vector search results to retrieve (default: 20)"
    )
    parser.add_argument(
        "--call-api", action="store_true",
        help="Also call the production API endpoint"
    )
    parser.add_argument(
        "--api-base", default="https://pharmacopoeia.kailin.com.cn",
        help="API base URL (default: https://pharmacopoeia.kailin.com.cn)"
    )
    parser.add_argument(
        "--token", default="",
        help="JWT token for API authentication (required when --call-api)"
    )
    args = parser.parse_args()

    if args.call_api and not args.token:
        print("ERROR: --token is required when using --call-api")
        import sys
        sys.exit(1)

    asyncio.run(run_diagnosis(
        query=args.query,
        top_k=args.top_k,
        call_api=args.call_api,
        api_base=args.api_base,
        token=args.token,
    ))