knowledge-agent/agents/expand_agent/agent.py

import sys
import os
from typing import Any, Dict, List
import json

# 添加项目根目录到 Python 路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from utils.logging_config import get_logger
from utils.mysql_db import MysqlHelper

logger = get_logger('ExpandAgent')

try:
    from gemini import GeminiProcessor
    HAS_GEMINI = True
except ImportError:
    HAS_GEMINI = False


def _fetch_parsing_data_by_request(request_id: str) -> List[str]:
    """从 knowledge_extraction_content 表中根据 request_id 获取 data 字段"""
    sql = "SELECT data, content_id, parsing_id FROM knowledge_extraction_content WHERE request_id = %s"
    rows = MysqlHelper.get_values(sql, (request_id,)) or []
    
    results = []
    for row in rows:
        data = row[0]  # 获取 data 字段
        content_id = row[1]  # 获取 content_id 字段
        parsing_id = row[2]  # 获取 parsing_id 字段
        if data:
            results.append({"data": data, "content_id": content_id, "parsing_id": parsing_id})
    
    print(f"Final results: {len(results)} items")
    return results


def _build_prompt(data_samples: str, input_query: str) -> str:
    """构建用于扩展查询的 Prompt，使用 expansion.md 模板"""
    
    # 读取 expansion.md 模板文件
    try:
        template_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), "prompt", "expansion.md")
        with open(template_path, 'r', encoding='utf-8') as f:
            template = f.read()
        
        # 替换模板中的占位符
        prompt = template.replace("{Original_Query}", input_query or "无")
        prompt = prompt.replace("{Content_For_Expansion}", data_samples)
        
        return prompt
        
    except Exception as e:
        logger.error(f"读取 expansion.md 模板失败: {e}")
        # 如果模板读取失败，使用备用 prompt
        return f"""你是一位顶级的知识库专家，精通语义分析、信息检索和搜索优化策略。

根据以下原始查询和内容，生成扩展的查询词列表：

**原始查询:** {input_query or "无"}

**扩展要基于的内容:** {data_samples}

请生成一个JSON数组形式的扩展查询词列表，确保每个查询词都具体、可操作、信息量丰富。

输出格式：
```json
[
    "查询词1",
    "查询词2",
    ...
]
```"""


def _run_llm(prompt: str) -> List[str]:
    """调用LLM生成扩展查询"""
    if not HAS_GEMINI:
        return []
        
    try:
        processor = GeminiProcessor()
        result = processor.process(content=prompt, system_prompt="你是专业的查询扩展助手")
        print(f"result: {result}")
        
        # 处理返回结果
        if isinstance(result, dict):
            # 如果有错误，直接返回空列表
            if "error" in result:
                logger.error(f"Gemini API 返回错误: {result['error']}")
                return []
            # 如果结果在 result 字段中
            text = result.get("result", "") or result.get("raw_response", "")
        else:
            text = str(result)
        
        # 清理文本，移除 markdown 代码块标记
        if "```json" in text:
            # 提取 ```json 和 ``` 之间的内容
            start = text.find("```json") + 7
            end = text.find("```", start)
            if end != -1:
                text = text[start:end].strip()
        elif "```" in text:
            # 提取 ``` 之间的内容
            start = text.find("```") + 3
            end = text.find("```", start)
            if end != -1:
                text = text[start:end].strip()
        
        # 尝试解析 JSON
        try:
            queries = json.loads(text)
            if isinstance(queries, list):
                # 确保所有元素都是字符串
                return [str(q) for q in queries if q]
            else:
                logger.warning(f"Gemini 返回的不是列表格式: {type(queries)}")
                return []
        except json.JSONDecodeError as e:
            logger.error(f"JSON 解析失败: {e}, 原始文本: {text}")
            return []
            
    except Exception as e:
        logger.error(f"LLM调用失败: {e}")
        return []

def _update_expansion_status(requestId: str, status: int):
    sql = "UPDATE knowledge_request SET expansion_status = %s WHERE request_id = %s"
    MysqlHelper.update_values(sql, (status, requestId))

def execute_expand_agent_with_api(requestId: str, query: str = "") -> Dict[str, Any]:
    """对外暴露的API：根据requestId查询数据，生成扩展查询"""
    # 获取数据（可能多条）
    data_samples = _fetch_parsing_data_by_request(requestId)

    try:
        total = 0
        success = 0
        if not data_samples:
            # 即使没有数据，也基于 query 生成一次兜底扩展
            # prompt = _build_prompt("", query)
            # expanded = _run_llm(prompt)
            # if not expanded:
            #     expanded = _heuristic_expand(query)
            # expand_querys_json = json.dumps(expanded, ensure_ascii=False)
            # insert_sql = """
            # INSERT INTO knowledge_expand_content 
            # (request_id, create_time, expand_querys, query) 
            # VALUES (%s, NOW(), %s, %s)
            # """
            # MysqlHelper.insert_and_get_id(insert_sql, (requestId, expand_querys_json, query))
            # total = 1
            # success = 1 if expanded else 0
            logger.info(f"没有数据，不进行扩展，直接返回: requestId={requestId}, total={total}, success={success}")
        else:
            # 针对每条 parsing_data 分别生成与入库
            insert_sql = """
            INSERT INTO knowledge_expand_content 
            (request_id, create_time, expand_querys, query, content_id, parsing_id) 
            VALUES (%s, NOW(), %s, %s, %s, %s)
            """
            for sample in data_samples:
                total += 1
                if not sample:
                    continue
                prompt = _build_prompt(sample["data"], query)
                expanded = _run_llm(prompt)
                try:
                    expand_querys_json = json.dumps(expanded, ensure_ascii=False)
                    MysqlHelper.insert_and_get_id(insert_sql, (requestId, expand_querys_json, query, sample["content_id"], sample["parsing_id"]))
                    success += 1
                except Exception as ie:
                    logger.error(f"单条扩展结果入库失败: requestId={requestId}, error={ie}")

        # 更新状态为处理完成
        _update_expansion_status(requestId, 2)
        logger.info(f"扩展完成: requestId={requestId}, total={total}, success={success}")
    except Exception as e:
        logger.error(f"保存扩展查询结果到数据库时出错: requestId={requestId}, error={e}")
        _update_expansion_status(requestId, 3)
    
    return {"status": 1, "requestId": requestId}


if __name__ == "__main__":
    queries = execute_expand_agent_with_api("REQUEST_001")
    print(queries)