Merge branch 'development' into fix/profile-preservation-bug

Author: Krashnicov

README.md

@@ -81,6 +81,7 @@ Agent Zero now supports **Projects** – isolated workspaces with their own prom
 - The framework does not guide or limit the agent in any way. There are no hard-coded rails that agents have to follow.
 - Every prompt, every small message template sent to the agent in its communication loop can be found in the **prompts/** folder and changed.
 - Every default tool can be found in the **python/tools/** folder and changed or copied to create new predefined tools.
+- **Automated configuration** via `A0_SET_` environment variables for deployment automation and easy setup.
 
 ![Prompts](/docs/res/prompts.png)
 

agent.py

@@ -11,8 +11,16 @@
 import uuid
 import models
 
-from python.helpers import extract_tools, files, errors, history, tokens, context as context_helper
-from python.helpers import dirty_json
+from python.helpers import (
+    extract_tools,
+    files,
+    errors,
+    history,
+    tokens,
+    context as context_helper,
+    dirty_json,
+    subagents
+)
 from python.helpers.print_style import PrintStyle
 
 from langchain_core.prompts import (
@@ -69,9 +77,10 @@ def __init__(
         # initialize state
         self.name = name
         self.config = config
+        self.data = data or {}
+        self.output_data = output_data or {}
         self.log = log or Log.Log()
         self.log.context = self
-        self.agent0 = agent0 or Agent(0, self.config, self)
         self.paused = paused
         self.streaming_agent = streaming_agent
         self.task: DeferredTask | None = None
@@ -80,10 +89,9 @@ def __init__(
         AgentContext._counter += 1
         self.no = AgentContext._counter
         self.last_message = last_message or datetime.now(timezone.utc)
-        self.data = data or {}
-        self.output_data = output_data or {}
-
 
+        # initialize agent at last (context is complete now)
+        self.agent0 = agent0 or Agent(0, self.config, self)
 
     @staticmethod
     def get(id: str):
@@ -100,7 +108,7 @@ def use(id: str):
 
     @staticmethod
     def current():
-        ctxid = context_helper.get_context_data("agent_context_id","")
+        ctxid = context_helper.get_context_data("agent_context_id", "")
         if not ctxid:
             return None
         return AgentContext.get(ctxid)
@@ -122,7 +130,8 @@ def all():
     @staticmethod
     def generate_id():
         def generate_short_id():
-            return ''.join(random.choices(string.ascii_letters + string.digits, k=8))
+            return "".join(random.choices(string.ascii_letters + string.digits, k=8))
+
         while True:
             short_id = generate_short_id()
             if short_id not in AgentContext._contexts:
@@ -132,6 +141,7 @@ def generate_short_id():
     def get_notification_manager(cls):
         if cls._notification_manager is None:
             from python.helpers.notification import NotificationManager  # type: ignore
+
             cls._notification_manager = NotificationManager()
         return cls._notification_manager
 
@@ -269,7 +279,6 @@ async def _process_chain(self, agent: "Agent", msg: "UserMessage|str", user=True
             agent.handle_critical_exception(e)
 
 
-
 @dataclass
 class AgentConfig:
     chat_model: models.ModelConfig
@@ -280,7 +289,9 @@ class AgentConfig:
     profile: str = ""
     memory_subdir: str = ""
     knowledge_subdirs: list[str] = field(default_factory=lambda: ["default", "custom"])
-    browser_http_headers: dict[str, str] = field(default_factory=dict)  # Custom HTTP headers for browser requests
+    browser_http_headers: dict[str, str] = field(
+        default_factory=dict
+    )  # Custom HTTP headers for browser requests
     code_exec_ssh_enabled: bool = True
     code_exec_ssh_addr: str = "localhost"
     code_exec_ssh_port: int = 55022
@@ -354,6 +365,7 @@ def __init__(
         asyncio.run(self.call_extensions("agent_init"))
 
     async def monologue(self):
+        error_retries = 0  # counter for critical error retries
         while True:
             try:
                 # loop data dictionary to pass to extensions
@@ -380,7 +392,9 @@ async def monologue(self):
                         prompt = await self.prepare_prompt(loop_data=self.loop_data)
 
                         # call before_main_llm_call extensions
-                        await self.call_extensions("before_main_llm_call", loop_data=self.loop_data)
+                        await self.call_extensions(
+                            "before_main_llm_call", loop_data=self.loop_data
+                        )
 
                         async def reasoning_callback(chunk: str, full: str):
                             await self.handle_intervention()
@@ -389,7 +403,9 @@ async def reasoning_callback(chunk: str, full: str):
                             # Pass chunk and full data to extensions for processing
                             stream_data = {"chunk": chunk, "full": full}
                             await self.call_extensions(
-                                "reasoning_stream_chunk", loop_data=self.loop_data, stream_data=stream_data
+                                "reasoning_stream_chunk",
+                                loop_data=self.loop_data,
+                                stream_data=stream_data,
                             )
                             # Stream masked chunk after extensions processed it
                             if stream_data.get("chunk"):
@@ -405,7 +421,9 @@ async def stream_callback(chunk: str, full: str):
                             # Pass chunk and full data to extensions for processing
                             stream_data = {"chunk": chunk, "full": full}
                             await self.call_extensions(
-                                "response_stream_chunk", loop_data=self.loop_data, stream_data=stream_data
+                                "response_stream_chunk",
+                                loop_data=self.loop_data,
+                                stream_data=stream_data,
                             )
                             # Stream masked chunk after extensions processed it
                             if stream_data.get("chunk"):
@@ -453,6 +471,7 @@ async def stream_callback(chunk: str, full: str):
 
                     # exceptions inside message loop:
                     except InterventionException as e:
+                        error_retries = 0  # reset retry counter on user intervention
                         pass  # intervention message has been handled in handle_intervention(), proceed with conversation loop
                     except RepairableException as e:
                         # Forward repairable errors to the LLM, maybe it can fix them
@@ -462,8 +481,10 @@ async def stream_callback(chunk: str, full: str):
                         PrintStyle(font_color="red", padding=True).print(msg["message"])
                         self.context.log.log(type="error", content=msg["message"])
                     except Exception as e:
-                        # Other exception kill the loop
-                        self.handle_critical_exception(e)
+                        # Retry critical exceptions before failing
+                        error_retries = await self.retry_critical_exception(
+                            e, error_retries
+                        )
 
                     finally:
                         # call message_loop_end extensions
@@ -473,9 +494,13 @@ async def stream_callback(chunk: str, full: str):
 
             # exceptions outside message loop:
             except InterventionException as e:
+                error_retries = 0  # reset retry counter on user intervention
                 pass  # just start over
             except Exception as e:
-                self.handle_critical_exception(e)
+                # Retry critical exceptions before failing
+                error_retries = await self.retry_critical_exception(
+                    e, error_retries
+                )
             finally:
                 self.context.streaming_agent = None  # unset current streamer
                 # call monologue_end extensions
@@ -532,6 +557,30 @@ async def prepare_prompt(self, loop_data: LoopData) -> list[BaseMessage]:
 
         return full_prompt
 
+    async def retry_critical_exception(
+        self, e: Exception, error_retries: int, delay: int = 3, max_retries: int = 1
+    ) -> int:
+        if error_retries >= max_retries:
+            self.handle_critical_exception(e)
+
+        error_message = errors.format_error(e)
+        
+        self.context.log.log(
+            type="warning", content="Critical error occurred, retrying..."
+        )
+        PrintStyle(font_color="orange", padding=True).print(
+            "Critical error occurred, retrying..."
+        )
+        await asyncio.sleep(delay)
+        agent_facing_error = self.read_prompt(
+            "fw.msg_critical_error.md", error_message=error_message
+        )
+        self.hist_add_warning(message=agent_facing_error)
+        PrintStyle(font_color="orange", padding=True).print(
+            agent_facing_error
+        )
+        return error_retries + 1
+
     def handle_critical_exception(self, exception: Exception):
         if isinstance(exception, HandledException):
             raise exception  # Re-raise the exception to kill the loop
@@ -570,27 +619,15 @@ async def get_system_prompt(self, loop_data: LoopData) -> list[str]:
         return system_prompt
 
     def parse_prompt(self, _prompt_file: str, **kwargs):
-        dirs = [files.get_abs_path("prompts")]
-        if (
-            self.config.profile
-        ):  # if agent has custom folder, use it and use default as backup
-            prompt_dir = files.get_abs_path("agents", self.config.profile, "prompts")
-            dirs.insert(0, prompt_dir)
+        dirs = subagents.get_paths(self, "prompts")
         prompt = files.parse_file(
-            _prompt_file, _directories=dirs, **kwargs
+            _prompt_file, _directories=dirs, _agent=self, **kwargs
         )
         return prompt
 
     def read_prompt(self, file: str, **kwargs) -> str:
-        dirs = [files.get_abs_path("prompts")]
-        if (
-            self.config.profile
-        ):  # if agent has custom folder, use it and use default as backup
-            prompt_dir = files.get_abs_path("agents", self.config.profile, "prompts")
-            dirs.insert(0, prompt_dir)
-        prompt = files.read_prompt_file(
-            file, _directories=dirs, **kwargs
-        )
+        dirs = subagents.get_paths(self, "prompts")
+        prompt = files.read_prompt_file(file, _directories=dirs, _agent=self, **kwargs)
         prompt = files.remove_code_fences(prompt)
         return prompt
 
@@ -606,8 +643,12 @@ def hist_add_message(
         self.last_message = datetime.now(timezone.utc)
         # Allow extensions to process content before adding to history
         content_data = {"content": content}
-        asyncio.run(self.call_extensions("hist_add_before", content_data=content_data, ai=ai))
-        return self.history.add_message(ai=ai, content=content_data["content"], tokens=tokens)
+        asyncio.run(
+            self.call_extensions("hist_add_before", content_data=content_data, ai=ai)
+        )
+        return self.history.add_message(
+            ai=ai, content=content_data["content"], tokens=tokens
+        )
 
     def hist_add_user_message(self, message: UserMessage, intervention: bool = False):
         self.history.new_topic()  # user message starts a new topic in history
@@ -720,7 +761,9 @@ async def stream_callback(chunk: str, total: str):
             system_message=call_data["system"],
             user_message=call_data["message"],
             response_callback=stream_callback if call_data["callback"] else None,
-            rate_limiter_callback=self.rate_limiter_callback if not call_data["background"] else None,
+            rate_limiter_callback=(
+                self.rate_limiter_callback if not call_data["background"] else None
+            ),
         )
 
         return response
@@ -742,7 +785,9 @@ async def call_chat_model(
             messages=messages,
             reasoning_callback=reasoning_callback,
             response_callback=response_callback,
-            rate_limiter_callback=self.rate_limiter_callback if not background else None,
+            rate_limiter_callback=(
+                self.rate_limiter_callback if not background else None
+            ),
         )
 
         return response, reasoning
@@ -817,11 +862,15 @@ async def process_tools(self, msg: str):
             # Fallback to local get_tool if MCP tool was not found or MCP lookup failed
             if not tool:
                 tool = self.get_tool(
-                    name=tool_name, method=tool_method, args=tool_args, message=msg, loop_data=self.loop_data
+                    name=tool_name,
+                    method=tool_method,
+                    args=tool_args,
+                    message=msg,
+                    loop_data=self.loop_data,
                 )
 
             if tool:
-                self.loop_data.current_tool = tool # type: ignore
+                self.loop_data.current_tool = tool  # type: ignore
                 try:
                     await self.handle_intervention()
 
@@ -830,14 +879,20 @@ async def process_tools(self, msg: str):
                     await self.handle_intervention()
 
                     # Allow extensions to preprocess tool arguments
-                    await self.call_extensions("tool_execute_before", tool_args=tool_args or {}, tool_name=tool_name)
+                    await self.call_extensions(
+                        "tool_execute_before",
+                        tool_args=tool_args or {},
+                        tool_name=tool_name,
+                    )
 
                     response = await tool.execute(**tool_args)
                     await self.handle_intervention()
 
                     # Allow extensions to postprocess tool response
-                    await self.call_extensions("tool_execute_after", response=response, tool_name=tool_name)
-                    
+                    await self.call_extensions(
+                        "tool_execute_after", response=response, tool_name=tool_name
+                    )
+
                     await tool.after_execution(response)
                     await self.handle_intervention()
 
@@ -889,34 +944,40 @@ async def handle_response_stream(self, stream: str):
             pass
 
     def get_tool(
-        self, name: str, method: str | None, args: dict, message: str, loop_data: LoopData | None, **kwargs
+        self,
+        name: str,
+        method: str | None,
+        args: dict,
+        message: str,
+        loop_data: LoopData | None,
+        **kwargs,
     ):
         from python.tools.unknown import Unknown
         from python.helpers.tool import Tool
 
         classes = []
 
-        # try agent tools first
-        if self.config.profile:
+        # search for tools in agent's folder hierarchy
+        paths = subagents.get_paths(self, "tools", name + ".py", default_root="python")
+        for path in paths:
             try:
-                classes = extract_tools.load_classes_from_file(
-                    "agents/" + self.config.profile + "/tools/" + name + ".py", Tool  # type: ignore[arg-type]
-                )
+                classes = extract_tools.load_classes_from_file(path, Tool)  # type: ignore[arg-type]
+                break
             except Exception:
-                pass
+                continue
 
-        # try default tools
-        if not classes:
-            try:
-                classes = extract_tools.load_classes_from_file(
-                    "python/tools/" + name + ".py", Tool  # type: ignore[arg-type]
-                )
-            except Exception as e:
-                pass
         tool_class = classes[0] if classes else Unknown
         return tool_class(
-            agent=self, name=name, method=method, args=args, message=message, loop_data=loop_data, **kwargs
+            agent=self,
+            name=name,
+            method=method,
+            args=args,
+            message=message,
+            loop_data=loop_data,
+            **kwargs,
         )
 
     async def call_extensions(self, extension_point: str, **kwargs) -> Any:
-        return await call_extensions(extension_point=extension_point, agent=self, **kwargs)
+        return await call_extensions(
+            extension_point=extension_point, agent=self, **kwargs
+        )

agents/agent0/agent.json

@@ -0,0 +1,5 @@
+{
+  "title": "Agent 0",
+  "description": "Main agent of the system communicating directly with the user.",
+  "context": ""
+}

agents/default/agent.json

@@ -0,0 +1,5 @@
+{
+  "title": "Default prompts",
+  "description": "Default prompt file templates. Should be inherited and overriden by specialized prompt profiles.",
+  "context": ""
+}

agents/developer/agent.json

@@ -0,0 +1,5 @@
+{
+  "title": "Developer",
+  "description": "Agent specialized in complex software development.",
+  "context": "Use this agent for software development tasks, including writing code, debugging, refactoring, and architectural design."
+}

agents/hacker/agent.json

@@ -0,0 +1,5 @@
+{
+  "title": "Hacker",
+  "description": "Agent specialized in cyber security and penetration testing.",
+  "context": "Use this agent for cybersecurity tasks such as penetration testing, vulnerability analysis, and security auditing."
+}