Spaces:
Runtime error
Runtime error
| import os | |
| import json | |
| import traceback | |
| import sys | |
| from pathlib import Path | |
| cerebrum_path = Path.home() / "repos" / "terra" / "src" / "cerebrum" | |
| sys.path.insert(0, str(cerebrum_path)) | |
| REMOTE = os.getenv("HF_SPACE") == "true" | |
| CEREBRUM_AVAILABLE = False | |
| try: | |
| from aios.llm.apis import llm_chat, llm_call_tool, llm_chat_with_json_output | |
| from aios.interface import AutoTool | |
| from aios.config.config_manager import config as aios_config | |
| CEREBRUM_AVAILABLE = True | |
| print("✅ Loaded aios-agent-sdk (Cerebrum) successfully.") | |
| except ImportError: | |
| print("⚠️ cerebrum (aios-agent-sdk) not found. Falling back to mock mode.") | |
| def llm_chat(*args, **kwargs): | |
| return {"response": {"response_message": "[MOCK] Chat reply"}} | |
| def llm_call_tool(*args, **kwargs): | |
| return {"response": {"response_message": "[MOCK] Tool reply"}} | |
| def llm_chat_with_json_output(*args, **kwargs): | |
| return {"response": {"response_message": '[{"action_type": "chat", "action": "Mock Plan", "tool_use": []}]'}} | |
| # class AutoTool: | |
| # @staticmethod | |
| # def from_batch_preloaded(tools): | |
| # return [] | |
| from cerebrum.llm.apis import llm_chat, llm_call_tool, llm_chat_with_json_output | |
| from cerebrum.interface import AutoTool | |
| import os | |
| import json | |
| def get_config(): | |
| from cerebrum.config.config_manager import config | |
| return config | |
| config = get_config() | |
| aios_kernel_url = config.get_kernel_url() | |
| class DemoAgent: | |
| def __init__(self, agent_name): | |
| self.agent_name = agent_name | |
| self.config = self.load_config() | |
| self.tools = [ | |
| tool.get_tool_call_format() | |
| for tool in AutoTool.from_batch_preloaded(self.config["tools"]) | |
| ] | |
| self.plan_max_fail_times = 3 | |
| self.tool_call_max_fail_times = 3 | |
| self.start_time = None | |
| self.end_time = None | |
| self.request_waiting_times: list = [] | |
| self.request_turnaround_times: list = [] | |
| self.messages = [] | |
| self.workflow_mode = "manual" # (manual, automatic) | |
| self.rounds = 0 | |
| def load_config(self): | |
| script_path = os.path.abspath(__file__) | |
| script_dir = os.path.dirname(script_path) | |
| config_file = os.path.join(script_dir, "config.json") | |
| with open(config_file, "r") as f: | |
| config = json.load(f) | |
| return config | |
| def pre_select_tools(self, tool_names): | |
| pre_selected_tools = [] | |
| for tool_name in tool_names: | |
| for tool in self.tools: | |
| if tool["function"]["name"] == tool_name: | |
| pre_selected_tools.append(tool) | |
| break | |
| return pre_selected_tools | |
| def build_system_instruction(self): | |
| prefix = "".join(["".join(self.config["description"])]) | |
| plan_instruction = "".join( | |
| [ | |
| f"You are given the available tools from the tool list: {json.dumps(self.tools)} to help you solve problems. ", | |
| "Generate a plan with comprehensive yet minimal steps to fulfill the task. ", | |
| "The plan must follow the json format as below: ", | |
| "[", | |
| '{"action_type": "action_type_value", "action": "action_value","tool_use": [tool_name1, tool_name2,...]}', | |
| '{"action_type": "action_type_value", "action": "action_value", "tool_use": [tool_name1, tool_name2,...]}', | |
| "...", | |
| "]", | |
| "In each step of the planned plan, identify tools to use and recognize no tool is necessary. ", | |
| "Followings are some plan examples. ", | |
| "[" "[", | |
| '{"action_type": "tool_use", "action": "gather information from arxiv. ", "tool_use": ["arxiv"]},', | |
| '{"action_type": "chat", "action": "write a summarization based on the gathered information. ", "tool_use": []}', | |
| "];", | |
| "[", | |
| '{"action_type": "tool_use", "action": "gather information from arxiv. ", "tool_use": ["arxiv"]},', | |
| '{"action_type": "chat", "action": "understand the current methods and propose ideas that can improve ", "tool_use": []}', | |
| "]", | |
| "]", | |
| ] | |
| ) | |
| if self.workflow_mode == "manual": | |
| self.messages.append({"role": "system", "content": prefix}) | |
| else: | |
| assert self.workflow_mode == "automatic" | |
| self.messages.append({"role": "system", "content": prefix}) | |
| self.messages.append({"role": "user", "content": plan_instruction}) | |
| def automatic_workflow(self): | |
| for i in range(self.plan_max_fail_times): | |
| response = llm_chat_with_json_output( | |
| messages=self.messages, | |
| message_return_type="json" | |
| )["response"]["response_message"] | |
| try: | |
| workflow = json.loads(response) | |
| except: | |
| workflow = None | |
| self.rounds += 1 | |
| if workflow: | |
| return workflow | |
| else: | |
| self.messages.append( | |
| { | |
| "role": "assistant", | |
| "content": f"Fail {i+1} times to generate a valid plan. I need to regenerate a plan", | |
| } | |
| ) | |
| return None | |
| def manual_workflow(self): | |
| workflow = [ | |
| { | |
| "action_type": "call_tool", | |
| "action": "Search for relevant papers", | |
| "tool_use": ["demo_author/arxiv"], | |
| }, | |
| { | |
| "action_type": "chat", | |
| "action": "Provide responses based on the user's query", | |
| "tool_use": [], | |
| }, | |
| ] | |
| return workflow | |
| def run(self, task_input): | |
| self.build_system_instruction() | |
| self.messages.append({"role": "user", "content": task_input}) | |
| workflow = None | |
| if self.workflow_mode == "automatic": | |
| workflow = self.automatic_workflow() | |
| self.messages = self.messages[:1] # clear long context | |
| else: | |
| assert self.workflow_mode == "manual" | |
| workflow = self.manual_workflow() | |
| self.messages.append( | |
| { | |
| "role": "user", | |
| "content": f"[Thinking]: The workflow generated for the problem is {json.dumps(workflow)}. Follow the workflow to solve the problem step by step. ", | |
| } | |
| ) | |
| try: | |
| if workflow: | |
| final_result = "" | |
| for i, step in enumerate(workflow): | |
| action_type = step["action_type"] | |
| action = step["action"] | |
| tool_use = step["tool_use"] | |
| prompt = f"At step {i + 1}, you need to: {action}. " | |
| self.messages.append({"role": "user", "content": prompt}) | |
| if tool_use: | |
| selected_tools = self.pre_select_tools(tool_use) | |
| else: | |
| selected_tools = None | |
| if action_type == "call_tool": | |
| response = llm_call_tool( | |
| agent_name=self.agent_name, | |
| messages=self.messages, | |
| tools=selected_tools, | |
| base_url=aios_kernel_url | |
| )["response"] | |
| else: | |
| response = llm_chat( | |
| agent_name=self.agent_name, | |
| messages=self.messages, | |
| base_url=aios_kernel_url | |
| )["response"] | |
| self.messages.append({"role": "assistant", "content": response["response_message"]}) | |
| self.rounds += 1 | |
| final_result = self.messages[-1]["content"] | |
| return { | |
| "agent_name": self.agent_name, | |
| "result": final_result, | |
| "rounds": self.rounds, | |
| } | |
| else: | |
| return { | |
| "agent_name": self.agent_name, | |
| "result": "Failed to generate a valid workflow in the given times.", | |
| "rounds": self.rounds, | |
| } | |
| except Exception as e: | |
| return {} |