universal-mcp-agents 0.1.13__py3-none-any.whl → 0.1.15__py3-none-any.whl

universal_mcp/agents/builder/builder.py

@@ -1,44 +1,26 @@
+import asyncio
 import json
-from collections import defaultdict
 
-from langchain_core.language_models import BaseChatModel
-from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.messages import HumanMessage
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.graph import END, START, StateGraph
+from langgraph.types import Command
 from loguru import logger
 from universal_mcp.tools.registry import ToolRegistry
 from universal_mcp.types import ToolConfig
 
 from universal_mcp.agents.base import BaseAgent
-from universal_mcp.agents.builder.prompts import (
-    AGENT_BUILDER_INSTRUCTIONS,
-    AGENT_FROM_CONVERSATION_PROMPT,
-    TASK_SYNTHESIS_PROMPT,
+from universal_mcp.agents.builder.helper import (
+    _clean_conversation_history,
+    _extract_tools_from_history,
+    _merge_tool_configs,
 )
+from universal_mcp.agents.builder.prompts import _build_prompt
 from universal_mcp.agents.builder.state import Agent, BuilderState
 from universal_mcp.agents.llm import load_chat_model
 from universal_mcp.agents.shared.tool_node import build_tool_node_graph
 
 
-async def generate_agent(llm: BaseChatModel, task: str, old_agent: Agent | None = None) -> Agent:
-    """Generates an agent from a task, optionally modifying an existing one."""
-    prompt_parts = [AGENT_BUILDER_INSTRUCTIONS]
-    if old_agent:
-        prompt_parts.append(
-            "\nThe user wants to modify the following agent design. "
-            "Incorporate their feedback into a new design.\n\n"
-            f"**User Feedback:** {task}\n\n"
-            f"{old_agent.model_dump_json(indent=2)}"
-        )
-    else:
-        prompt_parts.append(f"\n\n**Task:** {task}")
-
-    prompt = "\n".join(prompt_parts)
-    structured_llm = llm.with_structured_output(Agent)
-    agent = await structured_llm.ainvoke(prompt)
-    return agent
-
-
 class BuilderAgent(BaseAgent):
     def __init__(
         self,
@@ -57,169 +39,175 @@ class BuilderAgent(BaseAgent):
             **kwargs,
         )
         self.registry = registry
-        self.llm: BaseChatModel = load_chat_model(model, thinking=False)
+        self.llm = load_chat_model(model, thinking=False)
 
-    def _entry_point_router(self, state: BuilderState):
+    async def invoke(
+        self,
+        thread_id: str,
+        user_input: dict,
+    ):
         """
-        Determines the entry point of the graph based on the input format and conversation history.
-        - If input is a JSON with 'conversation_history', it builds from the conversation.
-        - If an agent has already been generated, it assumes a modification is requested.
-        - Otherwise, it starts a fresh build from a text prompt.
+        Overrides BaseAgent.invoke to build or modify an agent.
+        This is the primary entry point for the Builder Agent.
         """
-        last_message_content = state["messages"][-1].content
-        try:
-            # Case 1: Input is a JSON for building from a conversation
-            payload = json.loads(last_message_content)
-            if isinstance(payload, dict) and "conversation_history" in payload and "tool_config" in payload:
-                logger.info("Routing to: build from conversation history.")
-                return "synthesize_from_conversation"
-        except (json.JSONDecodeError, TypeError):
-            # Input is not a valid JSON, proceed as an interactive build
-            pass
-
-        # Case 2: It's an interactive build, check for modification vs. new
-        if state.get("generated_agent"):
-            logger.info("Routing to: modify existing agent.")
-            return "synthesize_new_task"
-        else:
-            logger.info("Routing to: new agent build.")
-            return "prepare_for_build"
-
-    async def _prepare_for_build(self, state: BuilderState):
-        """Sets the initial user task to begin the build process."""
-        last_message = state["messages"][-1]
-        task = last_message.content
-        yield {
-            "user_task": task,
-        }
+        keys = ("userInput", "agent", "tools", "messages")
+        userInput, agent_data, tools, messages = (user_input.get(k) for k in keys)
+        agent = Agent(**agent_data) if agent_data else None
+
+        await self.ainit()
+        graph = self._graph
+
+        initial_state = BuilderState(
+            user_task=userInput,
+            generated_agent=agent,
+            tool_config=tools,
+            messages=[],
+        )
 
-    async def _create_agent(self, state: BuilderState):
-        """Creates or updates the agent definition from a user_task."""
-        task = state["user_task"]
-        agent = state.get("generated_agent")
+        if messages:
+            initial_state["messages"] = [HumanMessage(content=json.dumps(messages))]
+        elif not userInput and not agent:
+            raise ValueError("Either 'user_input' or 'messages' must be provided for a new agent.")
 
-        generated_agent = await generate_agent(self.llm, task, agent)
-        yield {
-            "generated_agent": generated_agent,
-        }
+        run_metadata = {"agent_name": self.name, "is_background_run": False}
 
-    async def _get_tool_config_for_task(self, task: str) -> ToolConfig:
-        """Helper method to find and configure tools for a given task string."""
-        tool_finder_graph = build_tool_node_graph(self.llm, self.registry)
-        initial_state = {
-            "original_task": task,
-            "decomposition_attempts": 0,
+        config = {
+            "configurable": {"thread_id": thread_id},
+            "metadata": run_metadata,
+            "run_id": thread_id,
+            "run_name": self.name,
         }
-        final_state = await tool_finder_graph.ainvoke(initial_state)
-        execution_plan = final_state.get("execution_plan")
 
-        if not execution_plan:
-            return {}
+        final_state = await graph.ainvoke(initial_state, config=config)
+        return final_state
 
-        apps_with_tools = defaultdict(list)
-        for step in execution_plan:
-            app_id = step.get("app_id")
-            tool_ids = step.get("tool_ids")
-            if app_id and tool_ids:
-                apps_with_tools[app_id].extend(tool_ids)
-
-        return {app_id: list(set(tools)) for app_id, tools in apps_with_tools.items()}
-
-    async def _create_tool_config(self, state: BuilderState):
-        """Creates the tool configuration for the agent."""
-        task = state["user_task"]
-        tool_config = await self._get_tool_config_for_task(task)
-        yield {
-            "tool_config": tool_config,
-        }
+    def _entry_point_router(self, state: BuilderState):
+        """
+        Determines the entry point of the graph based on the initial state.
+        """
+        has_agent = state.get("generated_agent") is not None
+        has_messages = bool(state.get("messages"))
+        has_user_task = bool(state.get("user_task"))
+
+        if has_agent:
+            logger.info("Routing to: modify_agent.")
+            return "modify_agent"
+        elif has_messages:
+            logger.info("Routing to: create_agent_from_history.")
+            return "create_agent_from_history"
+        elif has_user_task:
+            logger.info("Routing to: create_agent_from_input.")
+            return "create_agent_from_input"
+        else:
+            raise ValueError("Invalid initial state. Cannot determine route.")
+
+    async def _create_agent_from_input(self, state: BuilderState) -> Command:
+        """SCENARIO 1: Generates a new agent from a single user_input, running agent and tool creation in parallel."""
+        user_task = state["user_task"]
+        logger.info(f"Creating new agent from input: '{user_task}'")
 
-    async def _synthesize_new_task_from_feedback(self, state: BuilderState):
-        """Synthesizes a new user_task from the original task and subsequent user feedback."""
-        original_task = next((msg.content for msg in state["messages"] if isinstance(msg, HumanMessage)), None)
-        modification_request = state["messages"][-1].content
+        structured_llm = self.llm.with_structured_output(Agent)
 
-        if not original_task:
-            raise ValueError("Could not find the original task in the conversation history.")
+        async def _task_generate_agent():
+            prompt = _build_prompt(user_task=user_task)
+            return await structured_llm.ainvoke(prompt)
 
-        synthesis_prompt = TASK_SYNTHESIS_PROMPT.format(
-            original_task=original_task,
-            modification_request=modification_request,
+        async def _task_find_tools():
+            return await self._get_tool_config_for_task(user_task)
+
+        # Run agent creation and tool finding concurrently for max efficiency
+        agent_profile, tool_config = await asyncio.gather(_task_generate_agent(), _task_find_tools())
+
+        logger.info(f"Successfully created agent '{agent_profile.name}' with tools: {tool_config}")
+
+        return Command(
+            update={"generated_agent": agent_profile, "tool_config": tool_config},
+            goto=END,
         )
 
-        response = await self.llm.ainvoke(synthesis_prompt)
-        new_synthesized_task = response.content.strip()
-        logger.info(f"The new synthesized task is: {new_synthesized_task}")
-        yield {
-            "user_task": new_synthesized_task,
-        }
+    async def _create_agent_from_history(self, state: BuilderState) -> Command:
+        """SCENARIO 2: Generates an agent by synthesizing a conversation history."""
+        user_task = state.get("user_task")
 
-    async def _synthesize_from_conversation(self, state: BuilderState):
-        """
-        Takes conversation history and used tools from input to synthesize a complete agent profile.
-        This is a one-shot generation.
-        """
         content_str = state["messages"][-1].content
-        initial_input = json.loads(content_str)
+        raw_history = json.loads(content_str)
+        conversation_history = _clean_conversation_history(raw_history)
+
+        logger.info(f"Creating new agent from conversation history (length: {len(conversation_history)}).")
 
-        conversation_history = initial_input.get("conversation_history")
-        tool_config = initial_input.get("tool_config")
+        # 1. Generate the agent profile first to get the definitive instructions
+        tools_from_history = _extract_tools_from_history(raw_history)
+        prompt = _build_prompt(
+            user_task=user_task,
+            conversation_history=conversation_history,
+            tool_config=tools_from_history,
+        )
+        structured_llm = self.llm.with_structured_output(Agent)
+        generated_agent = await structured_llm.ainvoke(prompt)
+        logger.info(f"Successfully generated agent profile for '{generated_agent.name}'.")
+
+        # 2. Synthesize tool configuration based on the new instructions and history
+        tools_from_instructions = await self._get_tool_config_for_task(generated_agent.instructions)
 
-        if not conversation_history or not tool_config:
-            raise ValueError("Input must be a dictionary containing 'conversation_history' and 'tool_config'.")
+        final_tool_config = _merge_tool_configs(tools_from_history, tools_from_instructions)
+        logger.info(f"Final synthesized tool configuration: {final_tool_config}")
 
-        prompt = AGENT_FROM_CONVERSATION_PROMPT.format(
-            conversation_history=json.dumps(conversation_history, indent=2),
-            tool_config=json.dumps(tool_config, indent=2),
+        return Command(
+            update={
+                "generated_agent": generated_agent,
+                "tool_config": final_tool_config,
+            },
+            goto=END,
         )
 
-        structured_llm = self.llm.with_structured_output(Agent)
-        generated_agent_profile = await structured_llm.ainvoke(prompt)
-
-        yield {
-            "generated_agent": generated_agent_profile,
-            "tool_config": tool_config,
-            "messages": [
-                AIMessage(
-                    content=f"Successfully generated agent '{generated_agent_profile.name}' from the conversation history."
-                )
-            ],
-        }
+    async def _modify_agent(self, state: BuilderState) -> Command:
+        """SCENARIO 3: Modifies an existing agent and re-evaluates its tool configuration."""
+        existing_agent = state["generated_agent"]
+        modification_request = state["user_task"]
+        existing_tools = state["tool_config"]
 
-    async def _build_graph(self):
-        """Builds the conversational agent graph."""
-        builder = StateGraph(BuilderState)
+        logger.info(f"Modifying existing agent '{existing_agent.name}' with request: '{modification_request}'")
 
-        # Add nodes
-        builder.add_node("prepare_for_build", self._prepare_for_build)
-        builder.add_node("create_agent", self._create_agent)
-        builder.add_node("create_tool_config", self._create_tool_config)
-        builder.add_node("synthesize_new_task", self._synthesize_new_task_from_feedback)
-        builder.add_node("synthesize_from_conversation", self._synthesize_from_conversation)
-
-        # The conditional entry point decides the workflow
-        builder.add_conditional_edges(
-            START,
-            self._entry_point_router,
-            {
-                "prepare_for_build": "prepare_for_build",
-                "synthesize_new_task": "synthesize_new_task",
-                "synthesize_from_conversation": "synthesize_from_conversation",
+        # 1. Generate the modified agent profile to get the new definitive instructions
+        prompt = _build_prompt(
+            existing_instructions=existing_agent.instructions,
+            modification_request=modification_request,
+        )
+        structured_llm = self.llm.with_structured_output(Agent)
+        modified_agent = await structured_llm.ainvoke(prompt)
+        logger.info(f"Successfully generated modified agent profile for '{modified_agent.name}'.")
+
+        # 2. Update tool configuration based on the NEW instructions, preserving existing tools
+        tools_from_new_instructions = await self._get_tool_config_for_task(modified_agent.instructions)
+        final_tool_config = _merge_tool_configs(existing_tools, tools_from_new_instructions)
+        logger.info(f"Final updated tool configuration: {final_tool_config}")
+
+        return Command(
+            update={
+                "generated_agent": modified_agent,
+                "tool_config": final_tool_config,
             },
+            goto=END,
         )
 
-        # Path for a fresh interactive build
-        builder.add_edge("prepare_for_build", "create_agent")
-        builder.add_edge("prepare_for_build", "create_tool_config")
+    async def _get_tool_config_for_task(self, task: str) -> ToolConfig:
+        """Helper method to find and configure tools for a given task string."""
+        if not task:
+            return {}
+        tool_finder_graph = build_tool_node_graph(self.llm, self.registry)
+        final_state = await tool_finder_graph.ainvoke({"original_task": task})
+        return final_state.get("execution_plan") or {}
 
-        # Path for modifying an existing build
-        builder.add_edge("synthesize_new_task", "create_agent")
-        builder.add_edge("synthesize_new_task", "create_tool_config")
+    async def _build_graph(self):
+        """Builds the conversational agent graph with the new, scenario-based structure."""
+        builder = StateGraph(BuilderState)
 
-        # Path for building from conversation ends after its single step
-        builder.add_edge("synthesize_from_conversation", END)
+        # Add the three self-contained nodes for each scenario
+        builder.add_node("create_agent_from_input", self._create_agent_from_input)
+        builder.add_node("create_agent_from_history", self._create_agent_from_history)
+        builder.add_node("modify_agent", self._modify_agent)
 
-        # Interactive creation nodes lead to the end of the run
-        builder.add_edge("create_agent", END)
-        builder.add_edge("create_tool_config", END)
+        # The entry point router directs to one of the three nodes, and they all go to END
+        builder.add_conditional_edges(START, self._entry_point_router)
 
         return builder.compile(checkpointer=self.memory)

universal_mcp/agents/builder/helper.py

@@ -0,0 +1,71 @@
+import collections
+from collections import defaultdict
+
+from loguru import logger
+from universal_mcp.types import ToolConfig
+
+
+def _extract_tools_from_history(history: list[dict]) -> ToolConfig:
+    """
+    Parses a conversation history to find and extract all tool names,
+    returning them in a structured ToolConfig format.
+
+    This function identifies messages with a "type" of "tool", extracts the
+    tool's name from the "name" key, and filters out a predefined list of
+    excluded tools. The remaining tool names are expected to be in an
+    "app_id__tool_id" format. These are then organized into a dictionary
+    mapping each app_id to a sorted list of its associated tool_ids.
+    """
+    apps_with_tools = collections.defaultdict(set)
+    excluded_tools = {"search_tools", "load_tools"}
+
+    for message in history:
+        if message.get("type") == "tool":
+            full_tool_name = message.get("name")
+            if not full_tool_name or full_tool_name in excluded_tools:
+                continue
+
+            if "__" in full_tool_name:
+                app_id, tool_id = full_tool_name.split("__", 1)
+                apps_with_tools[app_id].add(tool_id)
+
+    return {app_id: sorted(list(tools)) for app_id, tools in apps_with_tools.items()}
+
+
+def _clean_conversation_history(history: list[dict]) -> list[dict]:
+    """
+    Filters a raw conversation history, keeping only messages relevant for
+    agent synthesis (human, ai, and tool messages with a name containing double underscores).
+    """
+    cleaned_history = []
+    for message in history:
+        msg_type = message.get("type")
+
+        if msg_type in ["human", "ai"]:
+            cleaned_history.append(message)
+        elif msg_type == "tool" and isinstance(message.get("name"), str) and "__" in message["name"]:
+            cleaned_history.append(message)
+
+    return cleaned_history
+
+
+def _merge_tool_configs(old_config: ToolConfig, new_config: ToolConfig) -> ToolConfig:
+    """Merges two tool configurations, taking the union of tools for each app."""
+    if not old_config:
+        return new_config
+    if not new_config:
+        return old_config
+
+    # Start with a copy of the old configuration
+    merged_config = defaultdict(set)
+    for app, tools in old_config.items():
+        merged_config[app].update(tools)
+
+    # Add the new tools, ensuring uniqueness
+    for app, tools in new_config.items():
+        merged_config[app].update(tools)
+
+    # Convert the sets back to sorted lists for consistent output
+    final_config = {app: sorted(list(tool_set)) for app, tool_set in merged_config.items()}
+    logger.info(f"Merged tool configuration: {final_config}")
+    return final_config