run.py

import time

from tools import tools
from llm import ZhipuLLM
from state import AgentState
from tool_executor import executor

FILE_NAME = "checkpoint.json"

class AgentRuntime:

    def __init__(self,state):
        self.llm = ZhipuLLM()
        self.state = state

    def run(self,task,resume=False):
        if resume:
            self.state.load(FILE_NAME)
        messages,step_count,max_steps=self.state.state()
        tokens = self.llm.estimate_tokens(messages)
        if tokens >= 4096:
            messages = self.llm.compress(messages)
            print(f"触发压缩，压缩前：{tokens} tokens，压缩后：{self.llm.estimate_tokens(messages)} tokens")
        messages.append({"role": "user", "content": task})
        while True:
            step_count += 1
            if step_count >= max_steps:
                print("共执行：",step_count,"轮，超过最大step数")
                self.state.reset(FILE_NAME)
                return None
            # print(messages)
            for attempt in range(3):
                try:
                    response = self.llm.generate(messages, tools)
                    break
                except Exception as e:
                    print(f"LLM调用失败，正在进行第{attempt+1}/3次重试 \t {e}")
                    time.sleep(2 ** attempt)
            else:
                print("重试全部失败，放弃本任务")
                return None

            if not response.tool_calls:
                messages.append({"role": "assistant", "content": response.content})
                self.state.update(messages, step_count,FILE_NAME)
                print("共执行：",step_count,"轮，任务已完成")
                return response
            else:
                tool_calls_list = [
                    {
                        "id":tc.id,
                        "type":"function",
                        "function":{
                            "name": tc.function.name,
                            "arguments": tc.function.arguments
                        }
                    }for tc in response.tool_calls
                ]
                # print(tool_calls_list)
                messages.append({"role": "assistant", "content": response.content, "tool_calls":tool_calls_list})
                for tool_call in response.tool_calls:
                    # print(tool_call.function.name, tool_call.function.arguments)
                    try:
                        tool_output = executor(tool_call.function.name, tool_call.function.arguments)
                    except Exception as e:
                        tool_output = f"工具调用失败:{e}"
                    # print(tool_output)
                    messages.append({"role": "tool", "tool_call_id": tool_call.id,"content": tool_output})
                self.state.update(messages, step_count, FILE_NAME)


if __name__ == '__main__':
    state = AgentState()
    resume = input("Resume:")
    if resume == "Y" or resume == "y":
        print("Resume:yes")
        resume = True
    else:
        resume = False
    while True:
        agent = AgentRuntime(state)
        task = input("User:")
        response = agent.run(task, resume)
        if response==None:
            print("失败！")
            continue
        print("Agent:",response.content)
        # print("raw response:",response)
        print("State:",state.state())