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_parsing_content 表中根据 request_id 获取 parsing_data 字段"""
    sql = "SELECT parsing_data FROM knowledge_parsing_content WHERE request_id = %s ORDER BY id DESC LIMIT 10"
    rows = MysqlHelper.get_values(sql, (request_id,)) or []
    
    results = []
    for row in rows:
        parsing_data = row[0]  # 获取 parsing_data 字段
        if parsing_data:
            results.append(parsing_data)
    
    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 _heuristic_expand(input_query: str) -> List[str]:
    """启发式扩展（LLM不可用时的fallback）"""
    base = input_query.strip()
    if not base:
        return []
        
    return [
        base,
        f"{base} 教程",
        f"{base} 实战",
        f"{base} 入门指南",
        f"{base} 高级技巧"
    ]

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)
    
    # 构建prompt
    prompt = _build_prompt(data_samples[0], query)

    # 生成扩展查询
    expanded = _run_llm(prompt)
    if not expanded:
        expanded = _heuristic_expand(query)
    
    # 将扩展查询结果插入到 knowledge_expand_content 表
    try:
        # 先检查是否已存在记录
        check_sql = "SELECT id FROM knowledge_expand_content WHERE request_id = %s LIMIT 1"
        existing_record = MysqlHelper.get_values(check_sql, (requestId,))
        
        # 将 expanded 列表转换为 JSON 字符串
        expand_querys_json = json.dumps(expanded, ensure_ascii=False)
        print(f"expand_querys_json: {expand_querys_json}")
        
        if existing_record:
            # 记录已存在，执行更新
            update_sql = """
            UPDATE knowledge_expand_content 
            SET expand_querys = %s, query = %s, create_time = NOW()
            WHERE request_id = %s
            """
            affected_rows = MysqlHelper.update_values(update_sql, (expand_querys_json, query, requestId))
            logger.info(f"扩展查询结果已更新: requestId={requestId}, affected_rows={affected_rows}")
        else:
            # 记录不存在，执行插入
            insert_sql = """
            INSERT INTO knowledge_expand_content 
            (request_id, create_time, expand_querys, query) 
            VALUES (%s, NOW(), %s, %s)
            """
            insert_result = MysqlHelper.insert_and_get_id(insert_sql, (requestId, expand_querys_json, query))
            logger.info(f"扩展查询结果已插入: requestId={requestId}, insert_id={insert_result}")
        
        # 更新状态为处理完成
        _update_expansion_status(requestId, 2)
            
    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)