LQAgentPlatform/data_pipeline/RAG_recall/rag_miluvs/main.py

from __future__ import annotations

import asyncio

from pathlib import Path
from typing import Any, AsyncIterator, Dict, List, Optional

from llm_pipeline.core.config import YamlConfigProvider
from llm_pipeline.core.pipeline import LLMPipeline
from llm_pipeline.interfaces import DataLoader, ResultSaver
from llm_pipeline.entity_extract_v1.dataloaders import EntityExtractV1JsonChunksLoader
from llm_pipeline.entity_extract_v1.prompting import (
    EntityExtractV1JsonResponseParser,
    EntityExtractV1PromptBuilder,
)
from llm_pipeline.entity_extract_v1.factory import build_llm_client as build_extract_llm_client
from llm_pipeline.entity_extract_eval_v1.prompting import (
    EntityEvalV1JsonResponseParser,
    EntityEvalV1PromptBuilder,
)
from llm_pipeline.entity_extract_eval_v1.factory import build_llm_client as build_eval_llm_client
from llm_pipeline.rag_retrieval_eval_v1.factory import (
    build_rag_eval_pipeline_for_entities_items,
    build_rag_eval_pipeline_for_qa_json,
)

from llm_pipeline.entity_extract_v1.factory import build_pipeline_for_csv, build_pipeline_for_json
from llm_pipeline.entity_extract_eval_v1.factory import build_eval_pipeline_for_json


async def run_entity_extract_v1_with_json(
    input_json: str,
    output_json: str = "output_from_json.json",
) -> None:
    """使用 entity_extract_v1 版本：JSON → JSON 处理。"""
    pipeline, _ = build_pipeline_for_json(input_json=input_json, output_json=output_json)
    await pipeline.run()


async def run_entity_extract_v1_with_csv(
    input_csv: str = "input.csv",
    output_csv: str = "output.csv",
) -> None:
    """使用 entity_extract_v1 版本：CSV → CSV 处理。"""
    pipeline, _ = build_pipeline_for_csv(input_csv=input_csv, output_csv=output_csv)
    await pipeline.run()


async def run_entity_eval_v1_with_json(
    input_json: str,
    output_json: str = "output_from_json_eval.json",
) -> None:
    """使用 entity_extract_eval_v1 版本：对抽取结果做专业性评估与过滤。"""
    pipeline, _ = build_eval_pipeline_for_json(
        input_json=input_json,
        output_json=output_json,
    )
    await pipeline.run()


async def run_full_entity_extract_and_eval() -> None:
    """一键运行：先抽取实体，再对结果进行评估过滤。"""
    raw_input = (
        "44_四川公路桥梁建设集团有限公司镇巴（川陕界）至广安高速公路通广段C合同段C4项目经理部_完整结果_20251212_155323.json"
    )
    first_output = "output_from_json.json"
    final_output = "output_from_json_eval.json"

    # 第一步：实体抽取
    await run_entity_extract_v1_with_json(input_json=raw_input, output_json=first_output)
    # 第二步：专业性评估与过滤
    await run_entity_eval_v1_with_json(input_json=first_output, output_json=final_output)


async def run_rag_retrieval_eval_with_qa_json(
    input_json: str,
    output_csv: str = "rag_eval_results.csv",
    collection: str = "first_bfp_collection_test",
    hybrid_top_k: int = 20,
    final_top_k: int = 5,
) -> None:
    """
    使用 rag_retrieval_eval_v1 版本：
    - 输入：单个包含 qa_pairs 的 JSON（与 batch_rag_eval_from_qa.py 兼容）；
    - 过程：对每个实体 name 进行检索召回（multi_stage_recall），并调用 LLM 做命中率评估；
    - 输出：汇总结果写入 CSV，便于统计分析。
    """
    pipeline, _ = build_rag_eval_pipeline_for_qa_json(
        input_json=input_json,
        output_csv=output_csv,
        collection=collection,
        hybrid_top_k=hybrid_top_k,
        final_top_k=final_top_k,
    )
    await pipeline.run()


class InMemoryListSaver(ResultSaver):
    """将流水线结果保存在内存列表中（不落地文件）。"""

    def __init__(self) -> None:
        self.items: List[Dict[str, Any]] = []

    async def save(self, item: Dict[str, Any], result: Dict[str, Any]) -> None:
        self.items.append({**item, **result})


class InMemoryDataLoader(DataLoader):
    """从内存列表提供数据的 DataLoader。"""

    def __init__(self, items: List[Dict[str, Any]]) -> None:
        self._items = items

    async def load_items(self) -> AsyncIterator[Dict[str, Any]]:
        for it in self._items:
            yield it

    def get_total(self) -> Optional[int]:
        return len(self._items)


class InMemoryEntityExtractSaver(ResultSaver):
    """对齐 entity_extract_v1 的 JSON 输出结构，但保存在内存。"""

    def __init__(self) -> None:
        self.items: List[Dict[str, Any]] = []

    async def save(self, item: Dict[str, Any], result: Dict[str, Any]) -> None:
        merged = {**item, **result}
        simplified = {
            "file_name": merged.get("file_name"),
            "chunk_id": merged.get("chunk_id"),
            "section_label": merged.get("section_label"),
            "text": merged.get("text"),
            "entity_extract_result": merged.get("entity_extract_result"),
        }
        self.items.append(simplified)


class InMemoryEvalFilteredSaver(ResultSaver):
    """对齐 entity_extract_eval_v1 的过滤逻辑，但保存在内存。"""

    def __init__(self) -> None:
        self.items: List[Dict[str, Any]] = []

    async def save(self, item: Dict[str, Any], result: Dict[str, Any]) -> None:
        merged = {**item, **result}
        entities_obj = merged.get("entity_extract_result") or {}
        entities = entities_obj.get("entities") if isinstance(entities_obj, dict) else None
        if not entities or not isinstance(entities, list):
            return
        self.items.append(
            {
                "file_name": merged.get("file_name"),
                "chunk_id": merged.get("chunk_id"),
                "section_label": merged.get("section_label"),
                "text": merged.get("text"),
                "entity_extract_result": entities_obj,
            }
        )


async def run_full_extract_eval_and_rag_eval_in_memory(
    input_json: str,
    output_csv: str = "rag_eval_results.csv",
    collection: str = "first_bfp_collection_test",
    hybrid_top_k: int = 20,
    final_top_k: int = 5,
) -> None:
    """
    全流程（不依赖中间文件）：
    1) entity_extract_v1：从 input_json(chunks) 抽取实体概念+背景
    2) entity_extract_eval_v1：专业性评估与过滤
    3) rag_retrieval_eval_v1：用过滤后的实体（name+背景/证据拼 query）做检索召回 + 命中率评估，输出 CSV
    """

    def _iter_input_json_files(path_str: str) -> List[Path]:
        p = Path(path_str)
        if not p.exists():
            raise FileNotFoundError(f"输入路径不存在: {p}")
        if p.is_file():
            return [p]
        if p.is_dir():
            # 目录：递归找 json，固定排序保证可复现
            return sorted(p.rglob("*.json"), key=lambda x: str(x))
        return []

    input_files = _iter_input_json_files(input_json)
    if not input_files:
        print(f"[INFO] 未找到可处理的 JSON 文件: {input_json}")
        return

    all_filtered_items: List[Dict[str, Any]] = []

    # === Stage 1 + 2: per-file extract + eval filter (in-memory) ===
    extract_service = Path(__file__).parent / "llm_pipeline" / "entity_extract_v1" / "service.yaml"
    extract_cfg = YamlConfigProvider(service_path=extract_service)
    extract_client = build_extract_llm_client(extract_cfg)
    extract_prompt = EntityExtractV1PromptBuilder(cfg_provider=extract_cfg)
    extract_parser = EntityExtractV1JsonResponseParser(output_field="entity_extract_result")
    eval_service = Path(__file__).parent / "llm_pipeline" / "entity_extract_eval_v1" / "service.yaml"
    eval_cfg = YamlConfigProvider(service_path=eval_service)
    eval_client = build_eval_llm_client(eval_cfg)
    eval_prompt = EntityEvalV1PromptBuilder(cfg_provider=eval_cfg)
    eval_parser = EntityEvalV1JsonResponseParser(output_field="entity_extract_result")

    for fp in input_files:
        # === Stage 1: entity_extract_v1 (in-memory) ===
        extract_loader = EntityExtractV1JsonChunksLoader(str(fp))
        extract_saver = InMemoryEntityExtractSaver()
        extract_pipeline = LLMPipeline(
            llm_client=extract_client,
            config_provider=extract_cfg,
            data_loader=extract_loader,
            prompt_builder=extract_prompt,
            response_parser=extract_parser,
            result_saver=extract_saver,
        )
        await extract_pipeline.run()
        extracted_items = extract_saver.items

        if not extracted_items:
            print(f"[INFO] 跳过（抽取阶段无输出）: {fp}")
            continue

        # === Stage 2: entity_extract_eval_v1 (in-memory) ===
        eval_loader = InMemoryDataLoader(extracted_items)
        eval_saver = InMemoryEvalFilteredSaver()
        eval_pipeline = LLMPipeline(
            llm_client=eval_client,
            config_provider=eval_cfg,
            data_loader=eval_loader,
            prompt_builder=eval_prompt,
            response_parser=eval_parser,
            result_saver=eval_saver,
        )
        await eval_pipeline.run()
        filtered_items = eval_saver.items

        if not filtered_items:
            print(f"[INFO] 跳过（评估过滤后无有效实体）: {fp}")
            continue

        all_filtered_items.extend(filtered_items)

    # === Stage 3: rag_retrieval_eval_v1 (entities -> retrieval -> hit eval) ===
    if not all_filtered_items:
        print("[INFO] 全部输入处理完成，但未产生任何可用于 RAG 评估的实体。")
        return

    rag_pipeline, _ = build_rag_eval_pipeline_for_entities_items(
        items=all_filtered_items,
        # items=extracted_items,
        output_csv=output_csv,
        collection=collection,
        hybrid_top_k=hybrid_top_k,
        final_top_k=final_top_k,
    )
    await rag_pipeline.run()


if __name__ == "__main__":
    # 默认执行“抽取 → 专业评估过滤 → 检索召回 → 命中率评估（CSV）”全流程（内存承接，不依赖中间文件）
    asyncio.run(
        run_full_extract_eval_and_rag_eval_in_memory(
            input_json="./data",
            output_csv="rag_eval_results.csv",
            collection="first_bfp_collection_test",
            hybrid_top_k=20,
            final_top_k=5,
        )
    )