LQAgentPlatform/foundation/ai/rag/retrieval/query_rewrite.py

import uuid
import asyncio
from foundation.infrastructure.config.config import config_handler
from foundation.observability.logger.loggering import server_logger
from foundation.ai.agent.generate.model_generate import generate_model_client
from foundation.database.base.vector.milvus_vector import MilvusVectorManager
from core.construction_review.component.reviewers.utils import prompt_loader

class QueryRewriteManager():
    """
    召回管理器，实现多路召回功能
    """

    def __init__(self):
        # 获取部署的模型列表
        self.generate_model_client = generate_model_client
        self.prompt_loader = prompt_loader

    def query_extract(self, review_content):
        """
        从审查条文中提取query
        return:
            query: str
            background: str
            parameters: str
        """
        try:
            # 获取提示词模板并组装
            task_prompt = self.prompt_loader.get_prompt_template(
                reviewer_type="query_extract",  # 审查器类型
                review_content=review_content   # 传入审查内容作为参数
            )

            # 构建任务提示信息 - 参考标准模式
            task_prompt_info = {
                "task_prompt": task_prompt,  # 使用组装好的提示词
                "task_name": "query_extract"
            }

            # 生成唯一的trace_id用于追踪
            trace_id = str(uuid.uuid4())

            # 调用模型生成接口（使用异步运行）
            model_response = asyncio.run(self.generate_model_client.get_model_generate_invoke(
                trace_id=trace_id,
                task_prompt_info=task_prompt_info
            ))

            # 记录日志
            server_logger.info(f"Query 提取完成长度: {len(review_content)}")
            return model_response

        except Exception as e:
            server_logger.error(f"Query 提取失败: {str(e)}")
            return None