universal-mcp-agents 0.1.14__py3-none-any.whl → 0.1.16__py3-none-any.whl

universal_mcp/agents/builder/prompts.py

@@ -1,54 +1,107 @@
-NEW_AGENT_PROMPT = r"""
-# ROLE & GOAL
-You are a specialized Agent Generation AI. Your primary function is to create a complete, high-quality AI agent profile based on the information provided.
+import json
 
-# INPUTS
-1.  **User Task (Optional):** A brief, initial request from the user. This might be vague or specific.
-2.  **Conversation History (Optional):** A transcript of a conversation. This is the **primary source of truth**. If the conversation history is provided, it should be prioritized over the User Task to understand the user's full, potentially multi-step, objective.
+from universal_mcp.types import ToolConfig
 
-# INSTRUCTIONS
-Analyze the available inputs to fully understand the user's intent. Synthesize this understanding into a complete agent profile according to the specified JSON schema.
 
-- The first line of the `instructions` field in your output MUST be a single, complete sentence that serves as the definitive task for the agent. This sentence should be synthesized from your analysis of the inputs.
-- The rest of the `instructions` should provide clear, actionable commands for the agent, covering its role, responsibilities, interaction style, and output formatting.
+def _build_prompt(
+    user_task: str | None = None,
+    conversation_history: list[dict] | None = None,
+    existing_instructions: str | None = None,
+    modification_request: str | None = None,
+    tool_config: ToolConfig | None = None,
+) -> str:
+    """Dynamically builds a cohesive and effective prompt for the LLM based on the provided inputs."""
 
-# TASK
-Based on the following inputs, generate a complete agent profile.
+    core_prompt = r"""
+You are a master AI Agent Architect. Your purpose is to design and define highly effective AI agents by interpreting user requests and generating a precise agent profile in JSON format.
 
-**User Task:**
-{user_task}
+Your process is systematic and thorough. You will analyze all provided information to construct a complete and coherent agent definition.
+"""
 
-**Conversation History:**
-{conversation_history}
+    analysis_sections = ["\n# I. Analysis of Provided Inputs\n"]
+    analysis_sections.append("You are to analyze the following information to understand the user's requirements:\n")
 
-**YOUR JSON OUTPUT:**
-"""
+    if user_task:
+        analysis_sections.append(f"## Primary User Task:\n```\n{user_task}\n```\n")
+
+    if conversation_history:
+        analysis_sections.append(
+            "## Conversation History:\n"
+            "Pay special attention to the messages from the 'human' user. These are direct expressions of their needs and expectations for the agent's behavior. Include the user specific personal information like email-id or anything else which is personal in the agent's instruction.\n"
+            f"```json\n{json.dumps(conversation_history, indent=2)}\n```\n"
+        )
 
-MODIFY_AGENT_PROMPT = r"""
-# ROLE & GOAL
-You are an expert Agent Modification AI. Your task is to intelligently update an existing AI agent's profile by integrating a user's modification request.
+    if existing_instructions:
+        analysis_sections.append(
+            "## Existing Agent Instructions:\n"
+            "This is the baseline definition for the current agent. Your task will be to modify this based on the user's new requests.\n"
+            f"```\n{existing_instructions}\n```\n"
+        )
 
-# CORE PRINCIPLES
-1.  **Synthesize, Don't Just Add:** Do not simply append the modification request. You must seamlessly integrate the user's feedback into the agent's existing instructions, creating a new, coherent set of commands.
-2.  **Holistic Update:** The modification may require changes to not only the instructions but also the agent's name, description, and expertise. You must re-evaluate the entire profile.
-3.  **Prioritize New Information:** The user's latest modification request and any new conversation history are the primary drivers for the changes.
+    if modification_request:
+        analysis_sections.append(
+            "## Modification Request:\n"
+            "The user wants to change the existing agent. You must incorporate these changes into the new agent definition.\n"
+            f"```\n{modification_request}\n```\n"
+        )
 
-# INPUTS
-1.  **Existing Agent Instructions:** The original instructions that define the agent's current behavior.
-2.  **Modification Request:** The user's new input, specifying the desired changes.
-3.  **Conversation History (Optional):** A transcript of the conversation that may provide additional context for the modification.
+    if tool_config:
+        analysis_sections.append(
+            "## Tool Configuration:\n"
+            "The agent has access to the following tools. The agent's instructions should reflect the appropriate use of these tools.\n"
+            f"```json\n{json.dumps(tool_config, indent=2)}\n```\n"
+        )
 
-# TASK
-Update the agent profile based on the user's feedback.
+    framework_prompt = r"""
+# II. Agent Definition Framework
 
-**Existing Agent Instructions:**
-{existing_instructions}
+Based on your analysis, you will now define the agent's profile.
 
-**Modification Request:**
-{modification_request}
+## 1. Intent Synthesis
+- **Primary Goal:** In a single sentence, what is the core objective of this agent?
+- **Key Requirements & Constraints:** List any specific requirements, rules, or limitations the agent must adhere to.
+
+## 2. Agent Profile Generation
+You will now construct the complete agent profile.
+
+- **Name (2-4 words):** A concise and memorable name that reflects the agent's core function.
+- **Description (1-2 sentences):** A clear and compelling summary of the agent's purpose and value.
+- **Expertise:** A specific, well-defined area of expertise (e.g., "Python Code Generation and Debugging," not "Programming").
+- **Instructions:**
+    - This is the most critical part of your output. Write a comprehensive set of system instructions for the agent.
+    - The instructions should contain all the necessary details for the agent to call the tools , use the information from the conversation history, and fulfill the user's primary task.
+    - The instructions should be written in markdown and be direct, actionable commands.
+    - Start with the user's primary task.
+    - Clearly define the agent's role and responsibilities.
+    - Provide explicit rules for its behavior and interaction style.
+    - If tools are provided, explain how and when the agent should use them.
+    - Specify the desired output format (e.g., JSON, markdown, plain text).
+- **Schedule:**
+    - If the user specifies a schedule, provide a cron expression for when the agent should run.
+    - The output for the schedule should only be the cron expression itself (e.g., "0 9 * * *"). Do not add any explanatory text.
+"""
+
+    final_task_prompt = r"""
+# III. Your Task
+
+Generate a single JSON object that represents the complete agent profile. The JSON object should have the following structure:
+
+{
+  "name": "...",
+  "description": "...",
+  "expertise": "...",
+  "instructions": "...",
+  "schedule": "..."
+}
+
+**YOUR JSON OUTPUT:**
+"""
 
-**Conversation History:**
-{conversation_history}
+    full_prompt = [
+        core_prompt,
+        "".join(analysis_sections),
+        framework_prompt,
+        final_task_prompt,
+    ]
 
-**YOUR UPDATED JSON OUTPUT:**
-"""
+    return "\n".join(full_prompt)

universal_mcp/agents/builder/state.py

@@ -21,4 +21,4 @@ class BuilderState(TypedDict):
     user_task: str | None
     generated_agent: Agent | None
     tool_config: ToolConfig | None
-    messages: Annotated[Sequence[BaseMessage], add_messages]
+    messages: Annotated[Sequence[BaseMessage], add_messages]

universal_mcp/agents/codeact0/__init__.py

@@ -1,3 +1,4 @@
 from .agent import CodeActAgent
+from .playbook_agent import CodeActPlaybookAgent
 
-__all__ = ["CodeActAgent"]
+__all__ = ["CodeActAgent", "CodeActPlaybookAgent"]

universal_mcp/agents/codeact0/agent.py

@@ -18,9 +18,9 @@ from universal_mcp.agents.codeact0.prompts import (
 )
 from universal_mcp.agents.codeact0.sandbox import eval_unsafe, execute_ipython_cell
 from universal_mcp.agents.codeact0.state import CodeActState
-from universal_mcp.agents.utils import filter_retry_on
-from universal_mcp.agents.codeact0.utils import inject_context
+from universal_mcp.agents.codeact0.utils import inject_context, smart_truncate
 from universal_mcp.agents.llm import load_chat_model
+from universal_mcp.agents.utils import filter_retry_on
 
 
 class CodeActAgent(BaseAgent):
@@ -64,9 +64,11 @@ class CodeActAgent(BaseAgent):
                 raise ValueError("Tools are configured but no registry is provided")
             # Langchain tools are fine
             exported_tools = await self.registry.export_tools(self.tools_config, ToolFormat.LANGCHAIN)
-        additional_tools= [smart_print, data_extractor, ai_classify, call_llm]
+        additional_tools = [smart_print, data_extractor, ai_classify, call_llm]
         additional_tools = [t if isinstance(t, StructuredTool) else create_tool(t) for t in additional_tools]
-        self.instructions, self.tools_context = create_default_prompt(exported_tools, additional_tools, self.instructions)
+        self.instructions, self.tools_context = create_default_prompt(
+            exported_tools, additional_tools, self.instructions
+        )
 
         def call_model(state: CodeActState) -> Command[Literal["sandbox"]]:
             messages = [{"role": "system", "content": self.instructions}] + state["messages"]
@@ -114,7 +116,12 @@ class CodeActAgent(BaseAgent):
                 existing_context = state.get("context", {})
                 context = {**existing_context, **add_context}
                 # Execute the script in the sandbox
-                output, new_context, new_add_context = self.eval_fn(code, context, previous_add_context)
+
+                output, new_context, new_add_context = self.eval_fn(
+                    code, context, previous_add_context, 180
+                )  # default timeout 3 min
+                output = smart_truncate(output)
+
                 return Command(
                     goto="call_model",
                     update={

universal_mcp/agents/codeact0/langgraph_agent.py

@@ -1,17 +1,14 @@
-import asyncio
-
-from langgraph.checkpoint.memory import MemorySaver
-from rich import print
 from universal_mcp.agentr.registry import AgentrRegistry
 
-from universal_mcp.agents.codeact0.agent import CodeActAgent
-from universal_mcp.agents.utils import messages_to_list
+from universal_mcp.agents.codeact0.playbook_agent import CodeActPlaybookAgent
+
+
 async def agent():
-    agent_obj = CodeActAgent(
-            name="CodeAct Agent",
-            instructions="Be very concise in your answers.",
-            model="anthropic:claude-4-sonnet-20250514",
-            tools={"google_calendar": ["get_upcoming_events"], "exa" : ["search_with_filters"]},
-            registry=AgentrRegistry()
-        )
-    return await agent_obj._build_graph()
+    agent_obj = CodeActPlaybookAgent(
+        name="CodeAct Agent",
+        instructions="Be very concise in your answers.",
+        model="anthropic:claude-4-sonnet-20250514",
+        tools=[],
+        registry=AgentrRegistry(),
+    )
+    return await agent_obj._build_graph()

universal_mcp/agents/codeact0/llm_tool.py

@@ -5,7 +5,7 @@ from typing import Any, Literal, cast
 from langchain.chat_models import init_chat_model
 from langchain_openai import AzureChatOpenAI
 
-from universal_mcp.agents.codeact0.utils import get_message_text, light_copy
+from universal_mcp.agents.codeact0.utils import get_message_text
 
 MAX_RETRIES = 3
 
@@ -27,7 +27,7 @@ def smart_print(data: Any) -> None:
     Args:
         data: Either a dictionary with string keys, or a list of such dictionaries
     """
-    print(light_copy(data))
+    print(light_copy(data)) # noqa
 
 
 def creative_writer(

universal_mcp/agents/codeact0/playbook_agent.py

@@ -0,0 +1,364 @@
+import inspect
+import json
+import re
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal, cast
+
+from langchain_core.messages import AIMessage, RemoveMessage, ToolMessage
+from langchain_core.tools import StructuredTool
+from langchain_core.tools import tool as create_tool
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.graph import START, StateGraph
+from langgraph.types import Command, RetryPolicy
+from universal_mcp.tools.registry import ToolRegistry
+from universal_mcp.types import ToolFormat, ToolConfig
+
+from universal_mcp.agents.base import BaseAgent
+from universal_mcp.agents.codeact0.llm_tool import ai_classify, call_llm, data_extractor, smart_print
+from universal_mcp.agents.codeact0.prompts import (
+    create_default_prompt,
+)
+from universal_mcp.agents.codeact0.sandbox import eval_unsafe, execute_ipython_cell
+from universal_mcp.agents.codeact0.state import CodeActState
+from universal_mcp.agents.codeact0.tools import create_meta_tools, enter_playbook_mode, exit_playbook_mode, get_valid_tools
+from universal_mcp.agents.codeact0.utils import inject_context, smart_truncate
+from universal_mcp.agents.llm import load_chat_model
+from universal_mcp.agents.utils import filter_retry_on, get_message_text, convert_tool_ids_to_dict
+
+PLAYBOOK_PLANNING_PROMPT = """Now, you are tasked with creating a reusable playbook from the user's previous workflow.
+
+TASK: Analyze the conversation history and code execution to create a step-by-step plan for a reusable function. Do not include the searching and loading of tools. Assume that the tools have already been loaded.
+
+Your plan should:
+1. Identify the key steps in the workflow
+2. Mark user-specific variables that should become the main playbook function parameters using `variable_name` syntax. Intermediate variables should not be highlighted using ``
+3. Keep the logic generic and reusable
+4. Be clear and concise
+
+Example:
+```
+1. Connect to database using `db_connection_string`
+2. Query user data for `user_id`
+3. Process results and calculate `metric_name`
+4. Send notification to `email_address`
+```
+
+Now create a plan based on the conversation history. Enclose it between ``` and ```. Ask the user if the plan is okay."""
+
+
+
+PLAYBOOK_CONFIRMING_PROMPT = """Now, you are tasked with confirming the playbook plan. Return True if the user is happy with the plan, False otherwise. Do not say anything else in your response. The user response will be the last message in the chain.
+"""
+
+PLAYBOOK_GENERATING_PROMPT = """Now, you are tasked with generating the playbook function. Return the function in Python code.
+Do not include any other text in your response.
+The function should be a single, complete piece of code that can be executed independently, based on previously executed code snippets that executed correctly.
+The parameters of the function should be the same as the final confirmed playbook plan.
+Do not include anything other than python code in your response
+"""
+
+
+class CodeActPlaybookAgent(BaseAgent):
+    def __init__(
+        self,
+        name: str,
+        instructions: str,
+        model: str,
+        memory: BaseCheckpointSaver | None = None,
+        tools: ToolConfig | None = None,
+        registry: ToolRegistry | None = None,
+        playbook_registry: object | None = None,
+        sandbox_timeout: int = 20,
+        **kwargs,
+    ):
+        super().__init__(
+            name=name,
+            instructions=instructions,
+            model=model,
+            memory=memory,
+            **kwargs,
+        )
+        self.model_instance = load_chat_model(model, thinking=True)
+        self.tools_config = tools or []
+        self.registry = registry
+        self.playbook_registry = playbook_registry
+        self.eval_fn = eval_unsafe
+        self.sandbox_timeout = sandbox_timeout
+        self.processed_tools: list[StructuredTool | Callable] = []
+
+    async def _build_graph(self):
+        meta_tools = create_meta_tools(self.registry)
+        additional_tools = [smart_print, data_extractor, ai_classify, call_llm, meta_tools["web_search"]]
+        self.additional_tools = [t if isinstance(t, StructuredTool) else create_tool(t) for t in additional_tools]
+        async def call_model(state: CodeActState) -> Command[Literal["sandbox", "execute_tools"]]:
+            self.exported_tools = []
+            if self.tools_config:
+                # Convert dict format to list format if needed
+                if isinstance(self.tools_config, dict):
+                    self.tools_config = [
+                        f"{provider}__{tool}" 
+                        for provider, tools in self.tools_config.items() 
+                        for tool in tools
+                    ]
+                if not self.registry:
+                    raise ValueError("Tools are configured but no registry is provided")
+                # Langchain tools are fine
+            self.tools_config.extend(state.get('selected_tool_ids',[]))
+            self.exported_tools = await self.registry.export_tools(self.tools_config, ToolFormat.LANGCHAIN)
+            self.final_instructions, self.tools_context = create_default_prompt(
+                self.exported_tools, self.additional_tools, self.instructions
+            )
+            messages = [{"role": "system", "content": self.final_instructions}] + state["messages"]
+
+            # Run the model and potentially loop for reflection
+            model_with_tools = self.model_instance.bind_tools(
+                tools=[
+                    execute_ipython_cell,
+                    enter_playbook_mode,
+                    meta_tools["search_functions"],
+                    meta_tools["load_functions"],
+                ],
+                tool_choice="auto",
+            )
+            response = cast(AIMessage, model_with_tools.invoke(messages))
+            if response.tool_calls:
+                return Command(goto="execute_tools", update={"messages": [response]})
+            else:
+                return Command(update={"messages": [response], "model_with_tools": model_with_tools})
+
+            # if response.tool_calls:
+            #     if len(response.tool_calls) > 1:
+            #         raise Exception("Not possible in Claude with llm.bind_tools(tools=tools, tool_choice='auto')")
+            #     if response.tool_calls[0]["name"] == "enter_playbook_mode":
+            #         return Command(goto="playbook", update = {"playbook_mode": "planning"})
+            #     if response.tool_calls[0]["name"] != "execute_ipython_cell":
+            #         raise Exception(
+            #             f"Unexpected tool call: {response.tool_calls[0]['name']}. Expected 'execute_ipython_cell'."
+            #         )
+            #     if (
+            #         response.tool_calls[0]["args"].get("snippet") is None
+            #         or not response.tool_calls[0]["args"]["snippet"].strip()
+            #     ):
+            #         raise Exception("Tool call 'execute_ipython_cell' requires a non-empty 'snippet' argument.")
+            #     return Command(goto="sandbox", update={"messages": [response]})
+            # else:
+            #     return Command(update={"messages": [response]})
+
+        async def execute_tools(state: CodeActState) -> Command[Literal["call_model", "playbook", "sandbox"]]:
+            """Execute tool calls"""
+            last_message = state["messages"][-1]
+            tool_calls = last_message.tool_calls if isinstance(last_message, AIMessage) else []
+
+            tool_messages = []
+            new_tool_ids = []
+            ask_user = False
+            ai_msg = ""
+            tool_result = ""
+
+            for tool_call in tool_calls:
+                try:
+                    if tool_call["name"] == "enter_playbook_mode":
+                        tool_message = ToolMessage(
+                            content=json.dumps("Entered Playbook Mode."),
+                            name=tool_call["name"],
+                            tool_call_id=tool_call["id"],
+                        )
+                        return Command(
+                            goto="playbook",
+                            update={"playbook_mode": "planning", "messages": [tool_message]}, #Entered Playbook mode
+                        )
+                    elif tool_call["name"] == "execute_ipython_cell":
+                        return Command(goto="sandbox")
+                    elif tool_call["name"] == "load_functions":  # Handle load_functions separately
+                        valid_tools, unconnected_links = await get_valid_tools(
+                            tool_ids=tool_call["args"]["tool_ids"], registry=self.registry
+                        )
+                        new_tool_ids.extend(valid_tools)
+                        # Create tool message response
+                        tool_result = f"Successfully loaded {len(valid_tools)} tools: {valid_tools}"
+                        links = "\n".join(unconnected_links)
+                        if links:
+                            ask_user = True
+                            ai_msg = f"Please login to the following app(s) using the following links and let me know in order to proceed:\n {links} "
+                    elif tool_call["name"] == "search_functions":
+                        tool_result = await meta_tools["search_functions"].ainvoke(tool_call["args"])
+                except Exception as e:
+                    tool_result = f"Error during {tool_call}: {e}"
+
+                tool_message = ToolMessage(
+                    content=json.dumps(tool_result),
+                    name=tool_call["name"],
+                    tool_call_id=tool_call["id"],
+                )
+                tool_messages.append(tool_message)
+
+            if new_tool_ids:
+                self.tools_config.extend(new_tool_ids)
+                self.exported_tools = await self.registry.export_tools(self.tools_config, ToolFormat.LANGCHAIN)
+                self.final_instructions, self.tools_context = create_default_prompt(
+                    self.exported_tools, self.additional_tools, self.instructions
+                )
+
+            if ask_user:
+                tool_messages.append(AIMessage(content=ai_msg))
+                return Command(update={"messages": tool_messages, "selected_tool_ids": new_tool_ids})
+
+            return Command(goto="call_model", update={"messages": tool_messages, "selected_tool_ids": new_tool_ids})
+
+        # If eval_fn is a async, we define async node function.
+        if inspect.iscoroutinefunction(self.eval_fn):
+            raise ValueError("eval_fn must be a synchronous function, not a coroutine.")
+            # async def sandbox(state: StateSchema):
+            #     existing_context = state.get("context", {})
+            #     context = {**existing_context, **tools_context}
+            #     # Execute the script in the sandbox
+            #     output, new_vars = await eval_fn(state["script"], context)
+            #     new_context = {**existing_context, **new_vars}
+            #     return {
+            #         "messages": [{"role": "user", "content": output}],
+            #         "context": new_context,
+            #     }
+        else:
+
+            def sandbox(state: CodeActState) -> Command[Literal["call_model"]]:
+                tool_call = state["messages"][-1].tool_calls[0]  # type: ignore
+                code = tool_call["args"]["snippet"]
+                previous_add_context = state.get("add_context", {})
+                add_context = inject_context(previous_add_context, self.tools_context)
+                existing_context = state.get("context", {})
+                context = {**existing_context, **add_context}
+                # Execute the script in the sandbox
+
+                output, new_context, new_add_context = self.eval_fn(
+                    code, context, previous_add_context, 180
+                )  # default timeout 3 min
+                output = smart_truncate(output)
+
+                return Command(
+                    goto="call_model",
+                    update={
+                        "messages": [
+                            ToolMessage(
+                                content=output,
+                                name=tool_call["name"],
+                                tool_call_id=tool_call["id"],
+                            )
+                        ],
+                        "context": new_context,
+                        "add_context": new_add_context,
+                    },
+                )
+
+        def playbook(state: CodeActState) -> Command[Literal["call_model"]]:
+            playbook_mode = state.get("playbook_mode")
+            if playbook_mode == "planning":
+                planning_instructions = self.instructions + PLAYBOOK_PLANNING_PROMPT
+                messages = [{"role": "system", "content": planning_instructions}] + state["messages"]
+
+                response = self.model_instance.invoke(messages)
+                response = cast(AIMessage, response)
+                response_text = get_message_text(response)
+                # Extract plan from response text between triple backticks
+                plan_match = re.search(r'```(.*?)```', response_text, re.DOTALL)
+                if plan_match:
+                    plan = plan_match.group(1).strip()
+                else:
+                    plan = response_text.strip()
+                return Command(update={"messages": [response], "playbook_mode": "confirming", "plan": plan})
+                
+
+            elif playbook_mode == "confirming":
+                confirmation_instructions = self.instructions + PLAYBOOK_CONFIRMING_PROMPT
+                messages = [{"role": "system", "content": confirmation_instructions}] + state["messages"]
+                response = self.model_instance.invoke(messages, stream=False)
+                response = get_message_text(response)
+                if "true" in response.lower():
+                    return Command(goto="playbook", update={"playbook_mode": "generating"})
+                else:
+                    return Command(goto="playbook", update={"playbook_mode": "planning"})
+
+                    
+
+            elif playbook_mode == "generating":
+                generating_instructions = self.instructions + PLAYBOOK_GENERATING_PROMPT
+                messages = [{"role": "system", "content": generating_instructions}] + state["messages"]
+                response = cast(AIMessage, self.model_instance.invoke(messages))
+                raw_content = get_message_text(response)
+                func_code = raw_content.strip()
+                func_code = func_code.replace("```python", "").replace("```", "")
+                func_code = func_code.strip()
+
+                # Extract function name (handle both regular and async functions)
+                match = re.search(r"^\s*(?:async\s+)?def\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*\(", func_code, re.MULTILINE)
+                if match:
+                    function_name = match.group(1)
+                else:
+                    function_name = "generated_playbook"
+
+                # Save or update an Agent using the helper registry
+                saved_note = ""
+                try:
+                    if not self.playbook_registry:
+                        raise ValueError("Playbook registry is not configured")
+
+                    # Build instructions payload embedding the plan and function code
+                    instructions_payload = {
+                        "playbookPlan": state["plan"],
+                        "playbookScript": {
+                            "name": function_name,
+                            "code": func_code,
+                        },
+                    }
+
+                    # Convert tool ids list to dict
+                    tool_dict = convert_tool_ids_to_dict(state["selected_tool_ids"])
+
+                    res = self.playbook_registry.create_agent(
+                        name=function_name,
+                        description=f"Generated playbook: {function_name}",
+                        instructions=instructions_payload,
+                        tools=tool_dict,
+                        visibility="private",
+                    )
+                    saved_note = f"Successfully created your playbook! Check it out here: [View Playbook](https://wingmen.info/agents/{res.id})"
+                except Exception as e:
+                    saved_note = f"Failed to save generated playbook as Agent '{function_name}': {e}"
+
+                # Mock tool call for exit_playbook_mode (for testing/demonstration)
+                mock_exit_tool_call = {
+                    "name": "exit_playbook_mode",
+                    "args": {},
+                    "id": "mock_exit_playbook_123"
+                }
+                mock_assistant_message = AIMessage(
+                    content=saved_note, 
+                    tool_calls=[mock_exit_tool_call]
+                )
+
+                
+                # Mock tool response for exit_playbook_mode
+                mock_exit_tool_response = ToolMessage(
+                    content=json.dumps(f"Exited Playbook Mode.{saved_note}"),
+                    name="exit_playbook_mode",
+                    tool_call_id="mock_exit_playbook_123"
+                )
+
+                return Command(update={"messages": [mock_assistant_message, mock_exit_tool_response], "playbook_mode": "normal"})
+
+        def route_entry(state: CodeActState) -> Literal["call_model", "playbook"]:
+            """Route to either normal mode or playbook creation"""
+            if state.get("playbook_mode") in ["planning", "confirming", "generating"]:
+                return "playbook"
+
+            return "call_model"
+
+        agent = StateGraph(state_schema=CodeActState)
+        agent.add_node(call_model, retry_policy=RetryPolicy(max_attempts=3, retry_on=filter_retry_on))
+        agent.add_node(sandbox)
+        agent.add_node(playbook)
+        agent.add_node(execute_tools)
+        agent.add_conditional_edges(START, route_entry)
+        # agent.add_edge(START, "call_model")
+        return agent.compile(checkpointer=self.memory)