tools.py

import os
import json
from langchain_community.tools.tavily_search import TavilySearchResults

def get_workdir_root():
    return os.environ.get("WORKDIR_ROOT", './data/llm_result')

WORKDIR_ROOT = get_workdir_root()  # 防止最后一次llm_result已经被改掉了

def read_file(filename):
    filename = os.path.join(WORKDIR_ROOT, filename)
    if not os.path.exists(filename):
        return f"{filename} does not exist, please check the file before reading"
    with open(filename, 'r') as f:
        return "\n".join(f.readlines())

def append_to_file(filename, content):
    filename = os.path.join(WORKDIR_ROOT, filename)
    if not os.path.exists(filename):
        return f"{filename} does not exist, please check the file before writing"
    with open(filename, 'a') as f:  # 使用追加模式 'a'
        f.write(content)
    return 'Append content to file success'

def write_file(filename, content):
    filename = os.path.join(WORKDIR_ROOT, filename)
    if not os.path.exists(WORKDIR_ROOT):
        os.makedirs(WORKDIR_ROOT)

    with open(filename, 'w') as f:  # 使用写入模式 'w'
        f.write(content)
    return 'Write content to file success'

def search(query):
    # 使用 langchain 中的工具
    tavily = TavilySearchResults(max_result=5)
    try:
        ret = tavily.invoke(input=query)
        '''
        ret:
        [{
            "content": "",
            "url": ""
        }]
        '''
        content_list = [obj['content'] for obj in ret]
        return "\n".join(content_list)
    except Exception as err:
        return "Search error: {}".format(err)

# 工具定义
tools_info = [
    {
        "name": "read_file",
        "description": "Read file generated by the agent, should write file before reading",
        "args": [{
            "name": "filename",
            "type": "string",
            "description": "The name of the file to read"
        }]
    },
    {
        "name": "append_to_file",
        "description": "Append LLM content to file, should write file before appending",
        "args": [{
            "name": "filename",
            "type": "string",
            "description": "The name of the file to append"
        },
        {
            "name": "content",
            "type": "string",
            "description": "The content to append"
        }]
    },
    {
        "name": "write_file",
        "description": "Write LLM content to file",
        "args": [{
            "name": "filename",
            "type": "string",
            "description": "The name of the file to write"
        },
        {
            "name": "content",
            "type": "string",
            "description": "The content to write"
        }]
    },
    {
        "name": "search",
        "description": "Search for additional information from the internet when you are unsure about the answer",
        "args": [{
            "name": "query",
            "type": "string",
            "description": "Search query"
        }]
    },
    {
        "name": "finish",
        "description": "完成用户的任务",
        "args": [{
            "name": "answer",
            "type": "string",
            "description": "最后的目标结果"
        }]
    }
]

# 工具的映射，key 是 action_name，value 是函数
tools_map = {
    "read_file": read_file,
    "append_to_file": append_to_file,
    "write_file": write_file,
    "search": search
}

def gen_tools_descri():
    tools_descri = []
    for idx, t in enumerate(tools_info):
        args_descri = []
        for info in t.get("args"):
            args_descri.append({
                "name": info["name"],
                "type": info["type"],
                "description": info["description"]
            })
        args_descri_str = json.dumps(args_descri, ensure_ascii=False)
        # 将工具的描述和参数描述拼接成一个字符串，并直接追加到 tools_descri 列表中
        tools_descri.append(f"{idx + 1}. {t['name']}: {t['description']}, {args_descri_str}")

    # 最后将所有工具描述拼接成一个多行字符串
    tools_prompt = "\n".join(tools_descri)
    return tools_prompt


# 与大模型的交互
# 提示构建: 在 agent_execute 函数中，通过 gen_prompt 将 tools_descri 融入提示中。
# 决策执行: 大模型基于构建的提示内容，决定使用哪些工具（如读取文件、搜索等）。
# 在 response 中，指定的 action 将指向 tools_map 中的相应函数