Mutual_Fund/rag/retrieve.py

"""Milvus-only retrieval primitives for the customer-service RAG layer."""
from __future__ import annotations

from inspect import isawaitable
import logging

from rag.embedding import embed_texts


logger = logging.getLogger("rag.retrieve")


BUSINESS_RETRIEVAL = (
    ("fin_faq", "faq", 3, 0.40),
    ("fin_fund_doc", "funddoc", 5, 0.35),
    ("fin_policy", "policy", 5, 0.35),
)


async def _config_value(config_getter, key: str, default):
    value = config_getter(key, str(default))
    if isawaitable(value):
        value = await value
    return type(default)(value)


async def retrieve_candidates(
    query: str,
    customer_id: str | None,
    *,
    milvus_client,
    embedder=embed_texts,
    config_getter,
) -> list[dict]:
    if not query or not query.strip():
        return []
    vectors = await embedder([query])
    vector = vectors[0]
    plans = list(BUSINESS_RETRIEVAL)
    if customer_id:
        plans.append(("customer_memory", "memory", 5, 0.60))

    candidates = []
    for collection_name, key_suffix, default_topk, default_threshold in plans:
        topk = await _config_value(
            config_getter,
            f"agent.customer.rag.topk.{key_suffix}",
            default_topk,
        )
        threshold = await _config_value(
            config_getter,
            f"agent.customer.rag.threshold.{key_suffix}",
            default_threshold,
        )
        expression = ""
        if collection_name == "customer_memory":
            escaped = customer_id.replace("\\", "\\\\").replace('"', '\\"')
            expression = f'customer_id == "{escaped}"'
        result = await milvus_client.search(
            collection_name=collection_name,
            data=[vector],
            limit=topk,
            filter=expression,
            output_fields=["doc_id", "title", "section_title", "text", "strategy"],
        )
        candidates.append(
            {
                "collection_name": collection_name,
                "threshold": threshold,
                "results": result,
            }
        )
    return candidates


def _flatten_hits(results):
    for batch in results or []:
        if isinstance(batch, dict):
            yield batch
        else:
            yield from batch or []


def _format_candidates(candidates) -> list[dict]:
    sources = []
    for candidate in candidates:
        threshold = candidate["threshold"]
        for hit in _flatten_hits(candidate["results"]):
            entity = hit.get("entity") or hit
            distance = hit.get("distance")
            raw_score = hit.get("score")
            if distance is not None:
                score = 1.0 - distance
            else:
                score = raw_score
            if score is None or score < threshold:
                continue
            sources.append(
                {
                    "doc_id": entity.get("doc_id", ""),
                    "title": entity.get("title", ""),
                    "section_title": entity.get("section_title") or None,
                    "chunk_text": entity.get("text", ""),
                    "score": score,
                }
            )
    sources.sort(key=lambda source: source["score"], reverse=True)
    return sources


async def rag_retrieve(
    query: str,
    customer_id: str | None,
    *,
    milvus_client,
    embedder=embed_texts,
    config_getter,
) -> list[dict]:
    """Return the stable source contract consumed by客服 Agent."""
    try:
        candidates = await retrieve_candidates(
            query,
            customer_id,
            milvus_client=milvus_client,
            embedder=embedder,
            config_getter=config_getter,
        )
    except Exception:
        logger.exception("RAG retrieval failed")
        return []
    return _format_candidates(candidates)


async def retrieve_with_status(
    query: str,
    customer_id: str | None,
    *,
    milvus_client,
    embedder=embed_texts,
    config_getter,
) -> dict:
    """Expose operational status while keeping failed source lists empty."""
    try:
        vectors = await embedder([query])
    except Exception:
        logger.exception("RAG embedding failed")
        return {"status": "embedding_failed", "sources": []}

    async def reuse_vector(_texts):
        return vectors

    try:
        candidates = await retrieve_candidates(
            query,
            customer_id,
            milvus_client=milvus_client,
            embedder=reuse_vector,
            config_getter=config_getter,
        )
    except Exception:
        logger.exception("Milvus retrieval failed")
        return {"status": "milvus_unavailable", "sources": []}
    return {"status": "ok", "sources": _format_candidates(candidates)}