# 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.

from google.adk import Agent
from google.genai import types

research_plan_agent = Agent(
    model="gemini-1.5-flash",
    name="research_plan_agent",
    description="I can help generate research plan.",
    instruction="""\
Your task is to create a research plan according to the user's query.

# Here are the instructions for creating the research plan:

+ Focus on finding specific things, e.g. products, data, etc.
+ Have the personality of a work colleague that is very helpful and explains things very nicely.
+ Don't mention your name unless you are asked.
+ Think about the most common things that you would need to research.
+ Think about possible answers when creating the plan.
+ Your task is to create the sections that should be researched. You will output high level headers, preceded by ##
+ Underneath each header, write a short sentence on what we want to find there.
+ The headers will follow the logical analysis pattern, as well as logical exploration pattern.
+ The headers should be a statement, not be in the form of questions.
+ The header will not include roman numerals or anything of the sort, e.g. ":", etc
+ Do not include things that you cannot possibly know about from using Google Search: e.g. sales forecasting, competitors, profitability analysis, etc.
+ Do not have an executive summary
+ In each section describe specifically what will be researched.
+ Never use "we will", but rather "I will".
+ Don't ask for clarifications from the user.
+ Do not ask the user for clarifications or if they have any other questions.
+ All headers should be bolded.
+ If you have steps in the plan that depend on other information, make sure they are 2 diferent sections in the plan.
+ At the end mention that you will start researching.

# Instruction on replying format

+ Start with your name as "[research_plan_agent]: ".
+ Output the content you want to say.

Output summary:
""",
    flow="single",
    sub_agents=[],
    generate_content_config=types.GenerateContentConfig(
        temperature=0.1,
    ),
)


question_generation_agent = Agent(
    model="gemini-1.5-flash",
    name="question_generation_agent",
    description="I can help generate questions related to user's question.",
    instruction="""\
Generate questions related to the research plan generated by research_plan_agent.

# Instruction on replying format

Your reply should be a numbered lsit.

For each question, reply in the following format: "[question_generation_agent]: [generated questions]"

Here is an example of the generated question list:

1. [question_generation_agent]: which state is San Jose in?
2. [question_generation_agent]: how google website is designed?
""",
    flow="single",
    sub_agents=[],
    generate_content_config=types.GenerateContentConfig(
        temperature=0.1,
    ),
)

information_retrieval_agent = Agent(
    model="gemini-1.5-flash",
    name="information_retrieval_agent",
    description=(
        "I can help retrieve information related to question_generation_agent's"
        " question."
    ),
    instruction="""\
Inspect all the questions after "[question_generation_agent]: " and asnwer them.

# Instruction on replying format

Always start with "[information_retrieval_agent]: "

For the answer of one question:

- Start with a title with one line summary of the reply.
- The title line should be bolded and starts with No.x of the corresponding question.
- Have a paragraph of detailed explain.

# Instruction on exiting the loop

- If you see there are less than 20 questions by "question_generation_agent", do not say "[exit]".
- If you see there are already great or equal to 20 questions asked by "question_generation_agent", say "[exit]" at last to exit the loop.
""",
    flow="single",
    sub_agents=[],
    generate_content_config=types.GenerateContentConfig(
        temperature=0.1,
    ),
)

question_sources_generation_agent = Agent(
    model="gemini-1.5-flash",
    name="question_sources_generation_agent",
    description=(
        "I can help generate questions and retrieve related information."
    ),
    instruction="Generate questions and retrieve information.",
    flow="loop",
    sub_agents=[
        question_generation_agent,
        information_retrieval_agent,
    ],
    generate_content_config=types.GenerateContentConfig(
        temperature=0.1,
    ),
)

summary_agent = Agent(
    model="gemini-1.5-flash",
    name="summary_agent",
    description="I can help summarize information of previous content.",
    instruction="""\
Summarize information in all historical messages that were replied by "question_generation_agent" and "information_retrieval_agent".

# Instruction on replying format

- The output should be like an essay that has a title, an abstract, multiple paragraphs for each topic and a conclusion.
- Each paragraph should maps to one or more question in historical content.
""",
    flow="single",
    generate_content_config=types.GenerateContentConfig(
        temperature=0.8,
    ),
)

research_assistant = Agent(
    model="gemini-1.5-flash",
    name="research_assistant",
    description="I can help with research question.",
    instruction="Help customers with their need.",
    flow="sequential",
    sub_agents=[
        research_plan_agent,
        question_sources_generation_agent,
        summary_agent,
    ],
    generate_content_config=types.GenerateContentConfig(
        temperature=0.1,
    ),
)

spark_agent = Agent(
    model="gemini-1.5-flash",
    name="spark_assistant",
    description="I can help with non-research question.",
    instruction="Help customers with their need.",
    flow="auto",
    sub_agents=[research_assistant],
    generate_content_config=types.GenerateContentConfig(
        temperature=0.1,
    ),
)

root_agent = spark_agent