AIyaodianzhushou/backend-python/app/rag/retriever.py

"""
混合检索器：pgvector 向量检索 + BM25 关键词检索
"""
import os
import json
from typing import Optional

import httpx
from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine

from app.core.config import get_settings

settings = get_settings()

EMBEDDING_URL = "https://dashscope.aliyuncs.com/api/v1/services/embeddings/text-embedding/text-embedding"
EMBEDDING_MODEL = "text-embedding-v3"
DB_URL = settings.database_url


def classify_intent(query: str) -> str:
    q = query.strip()

    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"


async def _get_query_embedding(query: str) -> list[float]:
    api_key = os.environ.get("QWEN_API_KEY") or settings.qwen_api_key
    async with httpx.AsyncClient(timeout=30) as client:
        resp = await client.post(
            EMBEDDING_URL,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {api_key}",
            },
            json={
                "model": EMBEDDING_MODEL,
                "input": {"texts": [query]},
                "parameters": {"text_type": "query"},
            },
        )
        data = resp.json()
    if data.get("code") and data.get("code") != "":
        raise RuntimeError(f"Embedding error: {data.get('message')}")
    return data["output"]["embeddings"][0]["embedding"]


class MixedRetriever:
    def __init__(self):
        self.table_name = "drug_chunks"
        self.vector_dim = settings.embedding_dim

    async def search(
        self,
        query: str,
        intent: str = "drug_query",
        top_k: int = 20,
        filters: Optional[dict] = None,
    ) -> list[dict]:
        query_vec = await _get_query_embedding(query)
        vec_str = "[" + ",".join(str(v) for v in query_vec) + "]"

        engine = create_async_engine(DB_URL)
        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()