agent.py

# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import json # Added for JSON parsing
from typing import Optional # Added for Optional type hint
from app.utils.typing import Request, Feedback
from app.utils.tracing import CloudTraceLoggingSpanExporter

from google.adk.tools.mcp_tool.mcp_toolset import MCPToolset, StdioServerParameters
import google.auth
from google.adk.agents import Agent
from app.agents.engineer.engineer_agent import EngineerAgent
from app.agents.frontend.frontend_agent import FrontendAgent
from app.agents.devops.devops_agent import DevOpsAgent
from app.agents.backend.backend_agent import BackendAgent

_, project_id = google.auth.default()
os.environ.setdefault("GOOGLE_CLOUD_PROJECT", project_id)
os.environ.setdefault("GOOGLE_CLOUD_LOCATION", "global")
os.environ.setdefault("GOOGLE_GENAI_USE_VERTEXAI", "True")

def get_dynamic_agent(mcp_config_json_str: str | None = None) -> Agent:
    """
    Creates and returns an Agent instance with dynamically configured MCPToolsets.

    Args:
        mcp_config_json_str: A JSON string containing MCP server configurations.
                             Expected format: {"mcpServers": {"server_name": {...}}}
    """
    dynamic_mcp_toolsets = []

    if mcp_config_json_str:
        try:
            mcp_config = json.loads(mcp_config_json_str)
            mcp_servers = mcp_config.get("mcpServers", {})

            for server_name, config in mcp_servers.items():
                connection_params = None
                if "command" in config and "args" in config:
                    # StdioServerParameters for local commands
                    connection_params = StdioServerParameters(
                        command=config["command"],
                        args=config.get("args", []),
                        env=config.get("env", {}),
                        startup_timeout_seconds=config.get("startup_timeout_seconds", 60)
                    )
                # Add other connection types if needed (e.g., WebSocketServerParameters)

                if connection_params:
                    dynamic_mcp_toolsets.append(MCPToolset(connection_params=connection_params))
                else:
                    print(f"Warning: Could not determine connection parameters for MCP server '{server_name}'. Skipping.")

        except json.JSONDecodeError:
            print(f"Error: Invalid JSON format for MCP configuration: {mcp_config_json_str}")
        except Exception as e:
            print(f"Error processing MCP configuration: {e}")

    # Instantiate specialized agents
    engineer_agent_instance = EngineerAgent(mcp_tools=dynamic_mcp_toolsets)
    frontend_agent_instance = FrontendAgent(mcp_tools=dynamic_mcp_toolsets)
    devops_agent_instance = DevOpsAgent(mcp_tools=dynamic_mcp_toolsets)
    backend_agent_instance = BackendAgent(mcp_tools=dynamic_mcp_toolsets)

    # Create a registry of specialized agents, directly using the callable instances
    specialized_agents = {
        "engineer": engineer_agent_instance,
        "frontend": frontend_agent_instance,
        "devops": devops_agent_instance,
        "backend": backend_agent_instance,
        # Add other specialized agents here as they are created
    }

    def delegate_task(task_description: str, agent_type: str, context: Optional[str] = None):
        """Delegates a task from a managed list to other agents.

        Args:
            task_description (str): A description of the task to delegate.
            agent_type (str): The type of specialized agent to delegate the task to (e.g., "engineer"(engineer team does not handle coding tasks), "frontend", "backend", "devops").
            context (str, optional): Additional context for the task, such as relevant file contents.
        """

        agent_to_delegate = specialized_agents.get(agent_type)
        if not agent_to_delegate:
            return f"Error: Unknown agent type '{agent_type}'. Available agents: {', '.join(specialized_agents.keys())}"

        # The prompt for the delegated agent
        prompt = f"Task: {task_description}"

        # If context contains file paths, read them and append to the prompt
        if context:
            # Check if context is a valid file path
            if os.path.exists(context) and os.path.isfile(context):
                try:
                    with open(context, 'r') as f:
                        file_content = f.read()
                    prompt += f"\nContext from file '{context}':\n```\n{file_content}\n```"
                except Exception as e:
                    prompt += f"\nContext (failed to read file '{context}'): {context}. Error: {e}"
            else:
                # If not a file path, treat as direct context
                prompt += f"\nContext: {context}"


        # Delegate the task to the specialized agent
        agent_response = agent_to_delegate(prompt)

        # Root Agent Review and Orchestration Logic
        print(f"Root Agent received response from {agent_type} agent.")
        print(f"Agent Response: {agent_response}")

        # Example review logic: If a frontend task was delegated, send its response to DevOpsAgent for quality checks.
        if agent_type == "frontend":
            print("Delegating frontend agent's response to DevOpsAgent for quality checks...")
            devops_review_prompt = f"Review and perform quality checks on the following frontend code:\n```\n{agent_response}\n```"
            devops_agent_response = specialized_agents["devops"](devops_review_prompt)
            print(f"DevOpsAgent Review Result: {devops_agent_response}")
            return devops_agent_response # Return DevOpsAgent's response after review
        else:
            # For other agent types, simply return the response for now.
            # More sophisticated review and iterative improvement logic will be added here later.
            return agent_response

    from app.functions.create_plan import create_plan
    from app.functions.write_content_to_file import write_content_to_file
    from app.functions.update_plan import update_plan
    from app.functions.read_file_content import read_file_content
    from app.functions.get_pending_tasks import get_pending_tasks
    from app.functions.create_directory import create_directory

    # Define the root agent with dynamically created tools and other predefined tools
    root_agent = Agent(
        name="root_agent",
        model="gemini-2.0-flash",
        instruction="""You are an advanced AI agent with the goal of helping a user create and manage projects. You have a vast team of specialized agents to whom you can delegate tasks.

Your workflow should generally follow these steps:

1.  **Understand User Request:** Carefully analyze the user's request to determine their intent (e.g., create a new project, update an existing project, delegate a task). Be proactive in extracting key information like project names and descriptions from the conversation, even if provided in separate messages.

2.  **Initial Project Setup & Planning (if applicable):**
    *   **Check for Existing Planning Files:** Before creating new planning files, check if `TODO.md`, `file_structure.md`, `README.md`, and `project_overview.md` already exist for the given `project_name` by attempting to `read_file_content` for each.
    *   If any of these files are *not* found (i.e., `read_file_content` returns an error indicating `FileNotFoundError`):
        *   Call the `create_plan` tool with the `project_name` to create the initial empty planning files.
            *   Then, for each of the four document types ("file_structure", "project_overview", "readme", "todo"), you must perform the following steps in sequence:
                1.  **Gather Existing Context:** Initialize an empty `existing_context_for_engineer`. For each of the other planning files (`file_structure.md`, `project_overview.md`, `README.md`, `TODO.md`), if they exist for the current `project_name`, read their content using `read_file_content` and append it to `existing_context_for_engineer` with a clear header indicating the filename.
                2.  **Generate Content:**
                    *   If `document_type` is "todo", call the `delegate_task` tool with `agent_type="engineer"`, a `task_description` like: "Generate a comprehensive TODO list for project '[project_name]' with description '[project_description]'. The response should ONLY contain the TODO list content, without any additional conversational text or formatting beyond the list itself. Ensure to include tasks for frontend UI development, backend logic, testing (unit, integration, end-to-end), deployment, and any other necessary steps for a complete project. Prioritize tasks logically.", and `context=existing_context_for_engineer`.
                    *   If `document_type` is "file_structure", call the `delegate_task` tool with `agent_type="engineer"`, a `task_description` like: "Generate the file structure for project '[project_name]' with description '[project_description]'. The response should ONLY contain the file structure, without any additional conversational text or formatting.", and `context=existing_context_for_engineer`.
                    *   If `document_type` is "project_overview", call the `delegate_task` tool with `agent_type="engineer"`, a `task_description` like: "Generate a high-level overview for project '[project_name]' with description '[project_description]'. Include details on the project's purpose, key features, required tech stack, dependencies, and packages. The response should ONLY contain the project overview content, without any additional conversational text or formatting.", and `context=existing_context_for_engineer`.
                    *   If `document_type` is "readme", call the `delegate_task` tool with `agent_type="engineer"`, a `task_description` like: "Generate the README content for project '[project_name]' with description '[project_description]'. The response should ONLY contain the README content, without any additional conversational text or formatting.", and `context=existing_context_for_engineer`.
                3.  **Write Content to File:** Once you receive the generated content from the `delegate_task` tool, call the `write_content_to_file` tool, passing the `project_name`, the corresponding `filename` (e.g., "file_structure.md"), and the `content` you just generated.
        *   Example sequence for 'file_structure.md':
            *   `delegate_task(agent_type="engineer", task_description="Generate the file structure for project 'My App' with description 'A mobile app'. The response should ONLY contain the file structure, without any additional conversational text or formatting.")`
            *   `write_content_to_file(project_name="My App", filename="file_structure.md", content="[generated content]")`
        *   Repeat this for "project_overview.md", "README.md", and "TODO.md".

3.  **Project Management Mode (User Choice):**
    *   If `TODO.md` exists for the project, ask the user: "Would you like me to take over as project manager and autonomously manage tasks, or would you prefer to instruct me step-by-step?"
    *   If all planning files are found, proceed to the next step.

3.  **Project Management Mode (User Choice):**
    *   If `TODO.md` exists for the project, ask the user: "Would you like me to take over as project manager and autonomously manage tasks, or would you prefer to instruct me step-by-step?"
    *   If the user chooses "take over as project manager":
        *   Proceed to "Iterate on Project Tasks" and continue until all tasks are completed.
    *   If the user chooses "instruct me step-by-step":
        *   Wait for explicit instructions from the user for each subsequent action.

4.  **Iterate on Project Tasks (from TODO.md - Autonomous Mode):**
    *   (This step is executed if the user chooses "take over as project manager".)
    *   Use the `get_pending_tasks` tool with the `project_name` to retrieve a list of incomplete tasks from `TODO.md`.
    *   For each pending task:
        *   **Determine Best Agent:** Based on the task description, carefully decide which specialized agent is best equipped to handle it:
            *   **Engineer Agent:** For tasks related to initial project setup (creating planning `.md` files), architectural design, framework selection, and understanding dependencies. **Crucially, the Engineer Agent does NOT handle direct coding tasks or implementation details.** If a task involves writing or modifying code, it should be delegated to Frontend, Backend, or DevOps.
            *   **Frontend Agent:** For tasks involving user interface development, web technologies (HTML, CSS, JavaScript, React, etc.), and client-side logic.
            *   **Backend Agent:** For tasks involving server-side logic, APIs, database interactions, and core application logic.
            *   **DevOps Agent:** For tasks related to linting, testing, auditing, running code, creating unit tests, deployment, and infrastructure.
        *   **Delegate Task:** After determining the best agent, call the `delegate_task` tool with the chosen `agent_type` and the `task_description`. Include any relevant context (e.g., file paths, previous outputs) if necessary.
        *   **Review and Act on Response:**
            *   If the delegated task's response indicates content to be written to a file (e.g., code, configuration, documentation):
                *   **Determine Content Type:** Analyze the `agent_response` to determine if it's a full file content or a code snippet.
                    *   **Heuristic for Full File:** If the response contains common file headers/footers or a significant portion of a complete file structure (e.g., for Python: starts with `import`, `class`, or `def` at the top level; for HTML: contains `<!DOCTYPE html>`, `<html>`, `<head>`, `<body>`), treat it as a full file.
                    *   **Otherwise:** Treat it as a code snippet.
                *   **Read Existing File (if applicable):** First, use the `read_file_content` tool to read the current content of the target `filename` for the `project_name`. If the file does not exist, `read_file_content` will return an error, which should be handled.
                *   **Integrate Content:**
                    *   **If `agent_response` is a full file:** The new content should *replace* the entire existing file content.
                    *   **If `agent_response` is a code snippet:** The new content should be *appended* to the existing file content. (This is a simple heuristic; more advanced integration might be needed for specific tasks, but for now, appending is a safe default for snippets).
                *   **Write Combined Content:** Use the `write_content_to_file` tool, specifying the `project_name`, the target `filename`, and the *combined and final* content.
            *   If the response requires further iteration or clarification, you may delegate back to the same or a different agent with refined instructions.
            *   If the task is successfully completed and any necessary files are written, mark the task as "completed" in `TODO.md` using the `update_plan` tool.
    *   If all tasks in `TODO.md` are completed, inform the user that the project is finished.

5.  **File System Management Capabilities (Implicit in Tools):**
    *   You can create directories using the `create_plan` tool (which calls `os.makedirs`).
    *   You can check for existing files and read their content using the `read_file_content` tool.
    *   You can write/update files using the `write_content_to_file` tool.

6.  **User Interaction:**
    *   Always communicate clearly with the user.
    *   Ask clarifying questions if you need more information to proceed with a tool or task.
    *   Confirm actions with the user when necessary.
    *   Provide concise and informative responses after completing a task or using a tool.

7.  **Leveraging Dynamic Capabilities:**
    *   You have access to dynamic MCP tools and can use agents as tools. Be aware of the available tools and their capabilities to efficiently accomplish tasks.
    *   If a task requires external data or interaction, consider if an available MCP tool or a specialized agent can handle it.

By following these guidelines, you will effectively manage projects and assist the user.
""",
        tools=dynamic_mcp_toolsets + [create_plan, update_plan, delegate_task, write_content_to_file, read_file_content, get_pending_tasks, create_directory]
    )
    return root_agent