# @title Authenticate - Google Colab Only
import sys

from google.colab import auth

if "google.colab.auth" in sys.modules:
    auth.authenticate_user()
    print("Authenticated")

# @title Install AG2, Google Cloud SDK and Dependencies
%pip install -qqq ag2[gemini]==0.10.5 dask[dataframe]==2025.12.0 google-cloud-aiplatform[agent_engines,ag2]==1.135.0 cloudpickle==3.1.2
%pip show ag2 google-cloud-aiplatform dask

# @title Import relevant packages
import pprint

from IPython.display import Markdown
from autogen import AssistantAgent, GroupChat, GroupChatManager, UserProxyAgent
from autogen.cache import Cache
from autogen.code_utils import content_str
import vertexai
from vertexai.agent_engines import AgentEngine, Queryable

print(f"Required libraries successfully imported.")

# @title Set Configuration Parameters
MODEL_NAME = "gemini-2.5-flash"  # @param {'type': 'string'}
PROJECT_ID = "<YOUR_PROJECT_ID>"  # @param {'type': 'string'}
STAGING_BUCKET = "gs://<YOUR_GCS_BUCKET>"  # @param {'type': 'string'}
CACHING_SEED = 42  # @param {'type': 'integer'}
LOCATION = "us-central1"

STAGING_BUCKET = (
    STAGING_BUCKET if STAGING_BUCKET.startswith("gs://") else f"gs://{STAGING_BUCKET}"
)

config_list = [
    {
        "model": MODEL_NAME,
        "project_id": PROJECT_ID,
        "location": LOCATION,
        "api_type": "google",
    }
]
seed = CACHING_SEED
all_config = {
    "gemini_config_list": config_list,
    "caching_seed": seed,
    "staging_bucket": STAGING_BUCKET,
}
print("Configuration:")
pprint.pprint(all_config)

# @title Initialize Vertex AI
vertexai.init(project=PROJECT_ID, location=LOCATION, staging_bucket=STAGING_BUCKET)
print(f"Vertex AI successfully initialized")

research_topic = "Agentic Workflows"  # @param ["Large Language Models", "Real Estate", "Agentic Workflows"]

# @title Define Agents

# Define the Researcher agent. This agent conducts thorough research on the given topic,
# identifying key trends, publications, experts, and relevant data sources.
researcher = AssistantAgent(
    name="Researcher",
    llm_config={"config_list": config_list, "seed": seed},
    description="Conducts thorough research on the given topic.",
    system_message=f"""You are a meticulous researcher. Your goal is to conduct thorough research, identifying key trends, publications, experts, and relevant data sources. You should provide accurate and well-cited information.""",
)

# Define the Literature Reviewer agent. This agent synthesizes research findings into a
# coherent and well-structured literature review, identifying key themes, connections, and gaps in the existing research.
literature_reviewer = AssistantAgent(
    name="Literature_Reviewer",
    llm_config={"config_list": config_list, "seed": seed},
    description="Synthesizes research findings into a coherent literature review.",
    system_message=f"""You are a skilled literature reviewer. Your goal is to synthesize the provided research findings into a coherent and well-structured literature review. You should identify key themes, connections, and gaps in the existing research.""",
)

# Define the Fact Checker agent. This agent verifies the accuracy and validity of information,
# checking sources, identifying potential biases, and ensuring consistency with established knowledge.
fact_checker = AssistantAgent(
    name="Fact_Checker",
    llm_config={"config_list": config_list, "seed": seed},
    description="Verifies the accuracy and validity of information.",
    system_message="""You are a meticulous fact-checker. Your goal is to verify the accuracy and validity of any information presented to you. You should check sources, identify potential biases, and ensure the information is consistent with established knowledge.""",
)

# Define the Data Analyst agent. This agent analyzes data related to the research topic and
# extracts meaningful insights, identifying trends, patterns, and correlations.
data_analyst = AssistantAgent(
    name="Data_Analyst",
    llm_config={"config_list": config_list, "seed": seed},
    description="Analyzes data and extracts meaningful insights.",
    system_message=f"""You are a skilled data analyst. Your goal is to analyze data and extract meaningful insights. You should identify trends, patterns, and correlations in the data and present your findings in a clear and concise manner.""",
)

# Define the User Proxy agent. This agent represents the user and can run code to facilitate tasks.
user_proxy = UserProxyAgent(
    name="User_proxy",
    code_execution_config={
        "last_n_messages": 20,
        "work_dir": "coding",
        "use_docker": False,
    },
    human_input_mode="TERMINATE",
    is_termination_msg=lambda x: content_str(x.get("content")).find("TERMINATE") >= 0,
    description="I represent the user, and can run code.",
)

# Create the GroupChat, a conversable pattern, that facilitates the collaboration on tasks requiring interaction among multiple agents.
# All agents contribute to a single conversation thread and share the same context. We have set the max turns of conversation to 12
groupchat = GroupChat(
    agents=[user_proxy, researcher, literature_reviewer, fact_checker, data_analyst],
    messages=[],
    max_round=12,
)

# Create the GroupChatManager, a specialized agent, to orchestrate the group chat. A GroupChatManager selects an agent to speak,
# broadcasts the message to all other agents, and repeats this process until the conversation stops or maximum number of turns is reached.
# The default selection strategy is "auto", where the manager's LLM, in our case Gemini, chooses the next speaker.
manager = GroupChatManager(
    groupchat=groupchat,
    llm_config={"config_list": config_list, "seed": seed},
    is_termination_msg=lambda x: content_str(x.get("content")).find("TERMINATE") >= 0,
)

# @title Start Multi-agent Conversations
msg = f"""Conduct research on the latest advancements in {research_topic}, analyze relevant data, and produce a comprehensive literature review with verified facts."""
custom_summary_prompt = f"Generate a final report on {research_topic} from the chat history. Just start the report immediately without adding any additional commentary before or after the report."
with Cache.disk() as cache:  # enable caching to reduce inference cost
    result = user_proxy.initiate_chat(
        manager,
        message=msg,
        summary_method="reflection_with_llm",
        summary_args={"summary_prompt": custom_summary_prompt},
        cache=cache,
    )

# @title Inspect Report Generated By Agents
summary = result.summary
final_report = content_str(
    summary.get("content") if isinstance(summary, dict) else summary
)
display(Markdown(final_report))

# @title Define AG2 Agents As a Agent Engine QueryAble

class ResearchApp(Queryable):
    """This class defines a research application using AG2. It facilitates multi-agent conversations for research tasks."""

    base_question = "Conduct research on the latest advancements in {topic}, analyze relevant data, and produce a comprehensive literature review with verified facts."
    custom_summary_prompt = f"Generate a final report on {research_topic} from the chat history. Just start the report immediately without adding any additional commentary before or after the report."

    def __init__(self, project_id: str, location: str):
        """
        Initializes the ResearchApp instance.

        Args:
            project_id (str): The project ID for AG2.
            location (str): The location of the AG2 instance.
        """
        self.project_id = project_id
        self.location = location
        self.set_up_complete = False

    def set_up(self):
        """Sets up the research agents and conversation environment."""
        self.researcher = AssistantAgent(
            name="Researcher",
            llm_config={
                "config_list": all_config["gemini_config_list"],
                "seed": all_config["caching_seed"],
            },
            description="Conducts thorough research on the given topic.",
            system_message="""You are a meticulous researcher. Your goal is to conduct thorough research, identifying key trends, publications, experts, and relevant data sources. You should provide accurate and well-cited information.""",
        )

        self.literature_reviewer = AssistantAgent(
            name="Literature_Reviewer",
            llm_config={
                "config_list": all_config["gemini_config_list"],
                "seed": all_config["caching_seed"],
            },
            description="Synthesizes research findings into a coherent literature review.",
            system_message="""You are a skilled literature reviewer. Your goal is to synthesize the provided research findings into a coherent and well-structured literature review. You should identify key themes, connections, and gaps in the existing research.""",
        )

        self.fact_checker = AssistantAgent(
            name="Fact_Checker",
            llm_config={
                "config_list": all_config["gemini_config_list"],
                "seed": all_config["caching_seed"],
            },
            description="Verifies the accuracy and validity of information.",
            system_message="""You are a meticulous fact-checker. Your goal is to verify the accuracy and validity of any information presented to you. You should check sources, identify potential biases, and ensure the information is consistent with established knowledge.""",
        )

        self.data_analyst = AssistantAgent(
            name="Data_Analyst",
            llm_config={
                "config_list": all_config["gemini_config_list"],
                "seed": all_config["caching_seed"],
            },
            description="Analyzes data and extracts meaningful insights.",
            system_message="""You are a skilled data analyst. Your goal is to analyze data and extract meaningful insights. You should identify trends, patterns, and correlations in the data and present your findings in a clear and concise manner.""",
        )

        self.user_proxy = UserProxyAgent(
            name="User_proxy",
            code_execution_config={
                "last_n_messages": 20,
                "work_dir": "coding",
                "use_docker": False,
            },
            human_input_mode="NEVER",
            is_termination_msg=lambda x: content_str(x.get("content")).find("TERMINATE")
            >= 0,
            description="I represent the user, and can run code.",
        )

        self.groupchat = GroupChat(
            agents=[
                self.user_proxy,
                self.researcher,
                self.literature_reviewer,
                self.fact_checker,
                self.data_analyst,
            ],
            messages=[],
            max_round=12,
        )
        self.manager = GroupChatManager(
            groupchat=self.groupchat,
            llm_config={"config_list": config_list, "seed": seed},
            is_termination_msg=lambda x: content_str(x.get("content")).find("TERMINATE")
            >= 0,
        )
        self.set_up_complete = True

    def query(self, topic: str = "Large Language Models"):
        """
        Queries the multi-agent research system.

        Args:
            topic (str): The research topic. Defaults to "Large Language Models".

        Returns:
            str: A summary of the research conducted by the multi-agent system.
        """
        if not self.set_up_complete:
            self.set_up()
        question = self.base_question.format(topic=topic)
        with Cache.disk() as cache:
            result = self.user_proxy.initiate_chat(
                self.manager,
                message=question,
                summary_method="reflection_with_llm",
                summary_args={"summary_prompt": custom_summary_prompt},
                cache=cache,
            )
        summary = result.summary
        return content_str(
            summary.get("content") if isinstance(summary, dict) else summary
        )


print(f"Research App with multi-agents created.")

# @title Testing - Local Development
local_engine = ResearchApp(project_id=PROJECT_ID, location=LOCATION)
local_engine.set_up()
report = local_engine.query(topic=research_topic)

# @title Inspect Report
display(Markdown(report))

# @title Clean up deployment if needed - If you know the reasoning engine ID please use it in the box

def cleanup_deployments(dry_run: bool = True, engine_id: str = "all"):
    if engine_id.lower() == "all":
        agent_engine_list = AgentEngine.list()
    else:
        try:
            # engine_id can be the numeric ID or the full resource name
            re = vertexai.agent_engines.get(engine_id)
            agent_engine_list = [re]
        except Exception as e:
            print(f"Error retrieving engine '{engine_id}': {e}")
            return

    if agent_engine_list:
        print(f"Found {len(agent_engine_list)} deployment(s) to process:")
        for re in agent_engine_list:
            display_name = re.display_name or "Unnamed"
            resource_name = re.resource_name
            if not dry_run:
                print(f"Deleting {display_name}...")
                print(f"Resource URI: {resource_name}")
                re.delete()
                print(f"  Successfully deleted.")
            else:
                print(f"DRY RUN: Would delete '{display_name}'")
                print(f"Resource URI: '{resource_name}'")
                print("  To execute the deletion, please uncheck 'dry_run_delete_step' and re-run this cell.")
    else:
        print(f"No deployments found matching criteria: {engine_id}")


ENGINE_ID = "all"  # @param {type:"string"}
dry_run_delete_step = True  # @param {type:"boolean"}
cleanup_deployments(dry_run=dry_run_delete_step, engine_id=ENGINE_ID)

# @title Define GCS bucket directory to stage the deployment artifacts
artifacts_gcs_dirname = "ag2_research_app"  # @param {'type': 'string'}

# @title Deploy to Agent Engine - Please Wait
agent_engine = AgentEngine.create(
    ResearchApp(project_id=PROJECT_ID, location=LOCATION),
    display_name="AG2 Research Application",
    gcs_dir_name=artifacts_gcs_dirname,
    description="AG2 example deployed on Vertex AI Agent Engine",
    requirements=[
        "ag2[gemini]==0.10.5",
        "google-cloud-aiplatform[agent_engines,ag2]==1.135.0",
        "google-generativeai==0.8.6",
        "dask[dataframe]==2025.12.0",
        "cloudpickle==3.1.2",
    ],
    extra_packages=[],
)

result = agent_engine.query(topic=research_topic)
display(Markdown(result))

with open("report.md", "w") as f:
    f.write(str(result))

# @title Clean up
ENGINE_ID = "<DEPLOYED_REASONING_ENGINE_ID>"  # @param {type:"string"}
dry_run_delete_step = True  # @param {type:"boolean"}
cleanup_deployments(dry_run=dry_run_delete_step, engine_id=ENGINE_ID)