LongjoeAgent/core/agent.py

import os
import dotenv
import datetime
from pathlib import Path

from langchain.agents import create_agent, AgentState
from langchain_openai import ChatOpenAI
from langchain_core.messages import (
    SystemMessage,
    HumanMessage,
    BaseMessage,
    trim_messages,
)
from tools.tool_factory import get_all_tools
from langchain_core.runnables import RunnableConfig
from langchain.agents.middleware import before_model
from langgraph.runtime import Runtime
from typing import Any, List, Sequence
from langchain.messages import RemoveMessage
from langgraph.graph.message import REMOVE_ALL_MESSAGES
import sqlite3
from config.settings import settings
from langchain_core.messages.utils import count_tokens_approximately

dotenv.load_dotenv()


def create_system_prompt(
    backend_url: str = "", token: str = "", username: str = "default"
) -> str:
    auth_status = "已认证" if token else "未认证"
    backend_available = "API可用" if backend_url and token else "仅数据查询"
    knowledge_status = (
        "知识库可用" if settings.KNOWLEDGE_BASE_ENABLED else "知识库已禁用"
    )
    echart_status = "图表可用" if settings.ECHARTS_ENABLED else "图表已禁用"

    if settings.KNOWLEDGE_BASE_ENABLED:
        # 知识库启用时的提示词
        system_prompt = f"""龙嘉软件助手- 用户:{username} 认证:{auth_status} 服务:{backend_available} 知识库:{knowledge_status} 图表:{echart_status}
    职责：ERP数据查询和问题解答，按用户语言回答。

    **核心安全指令 (必遵)**：
    1.  **当前凭据 (每次工具调用必须使用)**：
        - 后端地址: {backend_url if backend_url else '无'}
        - API令牌: {token if token else '无'}
    2.  **禁止沿用历史**：**严禁**从对话历史中复制、沿用任何旧的后端地址、令牌或工具参数。历史记录仅用于理解背景，其中的工具详情**不能**作为本次调用的参数来源。
    3.  **调用规范**：调用查询工具时，**必须且只能**使用上方提供的当前凭据。
    工作流：
    1. 分析问题意图，提取模块关键词
    2. 如果是数据查询类问题,直接调用相关工具查询数据
    3. 如果是其他问题,则通过工具搜索知识库,知识库工具使用流程:a.通过关键字获取相关文章列表,b.判断哪些文章最符合,c.再通过工具获取文章内容.严格按文章内容回复,不能编造答案.
    4. 关键词要精准，避免无意义词
    工具调用规格:
    - 如果连续3次调用相同工具相同参数,自动停止
    - 工具返回相同结果但仍在重复调用时,自动停止
    回答规则：
    - 知识库找不到时提示"正在学习该问题"
    - {"需要个人数据时验证认证状态" if backend_url else "仅提供数据查询和知识库支持"}
    - 保护隐私，专业准确,精炼简要
    时间：{datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
    数据查询结果尽量以 Markdown 表格格式输出，格式如下：
    | 列名1 | 列名2 | 列名3 |
    | :--- | :--- | :--- |
    | 数据1 | 数据2 | 数据3 |
    | 数据4 | 数据5 | 数据6 |
    """
    else:
        # 知识库禁用时的提示词 - 灵活处理工具返回结果
        system_prompt = f"""龙嘉软件助手- 用户:{username} 认证:{auth_status} 服务:{backend_available} 知识库:{knowledge_status} 图表:{echart_status}
    职责：处理ERP数据查询类问题，按用户语言回答。
    **核心安全指令 (必遵)**：
    1.  **当前凭据 (每次工具调用必须使用)**：
        - 后端地址: {backend_url if backend_url else '无'}
        - API令牌: {token if token else '无'}
    2.  **禁止沿用历史**：**严禁**从对话历史中复制、沿用任何旧的后端地址、令牌或工具参数。历史记录仅用于理解背景，其中的工具详情**不能**作为本次调用的参数来源。
    3.  **调用规范**：调用查询工具时，**必须且只能**使用上方提供的当前凭据。
    工作流：
    1. 分析问题意图，判断是否为数据查询类问题
    2. 如果是数据查询类问题，直接调用相关工具查询数据
    3. 根据工具返回的结果进行回答：
    - 如果工具返回了具体数据，按数据内容回答
    - 如果工具返回了错误信息（如"API返回错误","查询失败","没有权限"等），如实告知用户错误信息
    - 如果工具返回空数据或"未找到数据"，如实告知用户
    4. 如果是非数据查询类问题（如疑问、流程、操作等），回复："知识库正在完善，无法回答该问题"
    工具调用规格:
    - 禁止连续调用相同工具相同参数
    - 工具返回相同结果但仍在重复调用时,自动停止
    回答规则：
    - 如用户提出非ERP范围的问题(例如："你好"等闲聊),明确告知用户自己的职责：主要处理ERP数据查询类问题
    - 工具提示没有权限时,明确回复用户没有权限
    - 严格按工具返回的内容回答，不能编造答案,可对结果进行简单总结
    - 当工具返回错误信息时，如实转达给用户，不要添加额外解释
    - 保持专业、准确、简洁的回答风格
    {"- 需要个人数据时验证认证状态" if backend_url else "- 仅提供数据查询支持"}
    当前时间：{datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
    数据查询结果尽量以 Markdown 表格格式输出，格式如下：
    | 列名1 | 列名2 | 列名3 |
    | :--- | :--- | :--- |
    | 数据1 | 数据2 | 数据3 |
    | 数据4 | 数据5 | 数据6 |
    """
    if settings.ECHARTS_ENABLED:
        system_prompt = (
            system_prompt
            + """
    并且根据数据的格式,主动选择合适的图表输出,你可以输出柱状图、折线图、饼图。

    饼图格式范例如下:
    ```echarts
    {{
    "title": {{
        "text": "浏览器份额", "left": "center" }}
    "tooltip": {{
        "trigger": "item" }},
    "legend": {{
        "orient": "vertical", "left": "left" }},
    "series": [
        {{
        "name": "Share",
        "type": "pie",
        "radius": "55%",
        "center": ["50%", "60%"],
        "data": [
            {{"value": 1048, "name": "Chrome" }}
            {{"value": 735, "name": "Firefox" }}
            {{"value": 580, "name": "Edge" }}
        ]
        }}
    ]
    }}
    ```

    柱状图格式范例如下:
    ```echarts
    {{
    "title": {{"text": "每周销量" }}
    "tooltip": {{}},
    "xAxis": {{"type": "category", "data": ["Mon","Tue","Wed","Thu","Fri","Sat","Sun"] }}
    "yAxis": {{"type": "value" }}
    "series": [
        {{"type": "bar", "data": [120, 200, 150, 80, 70, 110, 130] }}
    ]
    }}
    ```

    折线图,格式范例如下:
    ```echarts
    {{
    "title": {{ "text": "温度趋势" }},
    "tooltip": {{ "trigger": "axis" }},
    "legend": {{ "data": ["最高", "最低"] }},
    "xAxis": {{
        "type": "category",
        "boundaryGap": false,
        "data": ["Mon","Tue","Wed","Thu","Fri","Sat","Sun"]
    }},
    "yAxis": {{"type": "value" }},
    "series": [
        {{"name": "最高", "type": "line", "data": [11, 11, 15, 13, 12, 13, 10], "smooth": true }}
        {{"name": "最低", "type": "line", "data": [1, -2, 2, 5, 3, 2, 0], "smooth": true }}
    ]
    }}
    ```
    """
        )

    return system_prompt


def get_day_number(date=None):
    """获取日期编号 (YYYYMMDD 格式)"""
    if date is None:
        date = datetime.datetime.now()
    return date.strftime("%Y%m%d")  # 格式: 20251229


def get_sqlite_checkpointer():
    """创建按天分割的SQLite检查点保存器"""
    try:
        from langgraph.checkpoint.sqlite import SqliteSaver

        # 获取当前日期编号
        current_day = get_day_number()

        # 数据库文件存放目录
        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        base_dir = os.path.join(project_root, "data", "checkpoints")
        os.makedirs(base_dir, exist_ok=True)

        # 数据库文件名格式: checkpoints_20251229.db
        db_filename = f"checkpoints_{current_day}.db"
        db_path = os.path.join(base_dir, db_filename)
        # checkpointer = SqliteSaver.from_conn_string(db_path)

        conn = sqlite3.connect(db_path, check_same_thread=False)
        conn.execute("PRAGMA wal_autocheckpoint=500")  # 2MB 就提交
        conn.execute("PRAGMA journal_size_limit=52428800")  # 最大 50MB
        checkpointer = SqliteSaver(conn)

        return checkpointer

    except Exception as e:
        print(f"[ERROR]创建 SQLite 检查器失败: {e}")
        import traceback

        traceback.print_exc()

        # 回退到内存保存器
        from langgraph.checkpoint.memory import InMemorySaver

        print("[WARN]使用 InMemorySaver 作为回退")
        return InMemorySaver()


def cleanup_old_checkpoints(max_days=7):
    """清理超过指定天数的旧检查点文件（可选功能）"""
    try:
        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        base_dir = os.path.join(project_root, "data", "checkpoints")
        if not os.path.exists(base_dir):
            return

        # 获取当前日期
        current_date = datetime.datetime.now()

        # 遍历目录中的所有.db文件
        for filename in os.listdir(base_dir):
            if filename.startswith("checkpoints_") and filename.endswith(".db"):
                try:
                    print(f"检查旧检查点文件: {filename}")
                    # 提取日期 (checkpoints_day_20251229.db -> 20251229)
                    date_str = filename.replace("checkpoints_day_", "").replace(
                        ".db", ""
                    )
                    file_date = datetime.datetime.strptime(date_str, "%Y%m%d")

                    # 计算天数差
                    days_diff = (current_date - file_date).days

                    # 删除超过 max_days 天的旧数据
                    if days_diff > max_days:
                        file_path = os.path.join(base_dir, filename)
                        os.remove(file_path)
                        print(
                            f"[CLEAN]清理旧检查点文件: {filename} (超过 {max_days} 天)"
                        )

                except (ValueError, IndexError):
                    # 文件名不符合预期，跳过
                    continue

    except Exception as e:
        print(f"[WARN]清理旧检查点失败: {e}")


# 创建agent
def create_langchain_agent(
    backend_url: str = "",
    token: str = "",
    username: str = "default",
    thread_id: str = "default",
):
    llm = ChatOpenAI(
        model=settings.LLM_MODEL,
        temperature=settings.LLM_TEMPERATURE,
        api_key=settings.DEEPSEEK_API_KEY,
        base_url=settings.DEEPSEEK_BASE_URL,
        max_tokens=settings.LLM_MAX_TOKENS,
    )

    tools = get_all_tools()
    # 添加调试信息
    print(f"[DEBUG]Agent 创建调试信息:")
    print(f"  - 用户: {username}")
    print(f"  - Thread ID: {thread_id}")
    print(f"  - 后端地址: {backend_url}")
    print(f"  - Token: {'已提供' if token else '未提供'}")
    print(f"  - 工具数量: {len(tools)}")

    for i, tool in enumerate(tools):
        print(f"  - 工具 {i+1}: {tool.name}")

    # 获取动态的system_prompt
    system_prompt = create_system_prompt(backend_url, token, username)
    print(system_prompt)
    # def simple_turn_based_trim(
    #     messages: Sequence[BaseMessage],
    #     keep_turns: int = 3,
    #     system_message: BaseMessage = None,
    # ) -> List[BaseMessage]:
    #     """
    #     修正版：按完整对话轮次修剪消息
    #     每轮对话从Human开始，到下一个Human之前结束
    #     """
    #     if not messages:
    #         return []
    #     # 分离系统消息（始终保留）
    #     system_messages = []
    #     other_messages = []
    #     for msg in messages:
    #         if (
    #             isinstance(msg, SystemMessage)
    #             or getattr(msg, "type", None) == "system"
    #             or getattr(msg, "role", None) == "system"
    #             or msg.__class__.__name__ == "SystemMessage"
    #         ):
    #             system_messages.append(msg)
    #         else:
    #             other_messages.append(msg)

    #     if len(other_messages) <= 1:
    #         return system_messages + other_messages

    #     # 找出所有Human消息的位置
    #     human_indices = []
    #     for i, msg in enumerate(other_messages):
    #         if (
    #             isinstance(msg, HumanMessage)
    #             or getattr(msg, "type", None) == "human"
    #             or getattr(msg, "role", None) == "user"
    #         ):
    #             human_indices.append(i)

    #     # 如果Human消息不足keep_turns轮，返回所有
    #     if not human_indices or len(human_indices) <= keep_turns:
    #         return system_messages + other_messages

    #     # 计算起始索引
    #     start_idx = human_indices[-keep_turns]

    #     # 获取要保留的消息
    #     preserved_messages = other_messages[start_idx:]

    #     # 4. 返回从该索引开始的所有消息
    #     result = system_messages + preserved_messages
    #     # print(f"修剪后消息数: {len(result)}")

    #     return result

    # @before_model
    # def trim_messages(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    #     """Keep only the last few messages to fit context window."""
    #     messages = state["messages"]

    #     if len(messages) <= 3:
    #         return None  # No changes needed

    #     # 保留最后4轮对话
    #     trimmed_messages = simple_turn_based_trim(messages, keep_turns=4)

    #     return {"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES)] + trimmed_messages}

    @before_model
    def trim_messages_middleware(
        state: AgentState, runtime: Runtime
    ) -> dict[str, Any] | None:
        """使用官方trim_messages函数修剪消息"""
        messages = state.get("messages", [])
        print(f"trim_messages_middleware[DEBUG]原始消息数: {len(messages)}")
        # if len(messages) <= 3:
        #     return None  # 不需要修剪

        trimmed_messages = trim_messages(
            messages,
            max_tokens=1000,
            strategy="last",  # 保留最近的对话
            token_counter=count_tokens_approximately,  # token计数器
            start_on="human",  # 从human消息开始计算轮次
            include_system=True,  # 包含系统消息
        )

        # 添加调试信息
        original_count = len(messages)
        trimmed_count = len(trimmed_messages)
        print(f"trim_messages_middleware[DEBUG]修剪后消息数: {trimmed_count}")

        if trimmed_count < original_count:
            print(f"[INFO]消息修剪: {original_count} -> {trimmed_count} 条消息")

        return {"messages": trimmed_messages}

    # 使用SQLiteSaver（按天分割）
    checkpointer = get_sqlite_checkpointer()
    # print(f"打印检查点保存器: {checkpointer}")

    # 可选：清理旧检查点（可配置为定期执行）
    if os.getenv("AUTO_CLEANUP", "false").lower() == "true":
        cleanup_old_checkpoints(max_days=7)  # 保留最近7天数据

    agent = create_agent(
        llm,
        tools,
        checkpointer=checkpointer,
        system_prompt=system_prompt,
        middleware=[trim_messages_middleware],
    )

    return agent