"""
Google ADK LLM 智能体完整示例
展示 LlmAgent 的各种配置和用法

对应教程：第03章 - LLM 智能体详解
"""

# 导入 ADK 核心模块
from google.adk.agents import Agent, LlmAgent  # Agent 和 LlmAgent
from google.adk.agents import SequentialAgent   # 顺序工作流
from google.adk.models.lite_llm import LiteLlm  # LiteLLM 适配器
from google.adk.tools import google_search      # Google 搜索工具
from google.adk.tools import code_execution     # 代码执行工具

# 导入运行相关模块
from google.adk.runners import Runner            # 运行器
from google.adk.sessions import InMemorySessionService  # 内存会话服务
from google.genai import types                   # Google GenAI 类型定义

# 导入异步库
import asyncio                                  # 异步编程库


# ========================================
# 示例一：基础 LlmAgent 配置
# ========================================

# 使用 Gemini 模型（默认）
gemini_agent = Agent(
    name="gemini_agent",                    # Agent 名称
    model="gemini-2.0-flash",               # 使用 Gemini 2.0 Flash
    instruction="你是一个专业的编程助手。",    # 系统指令
    description="擅长编程和代码审查。",       # 描述（用于多 Agent 路由）
)

# 使用 LiteLLM 适配 DeepSeek 模型
deepseek_agent = Agent(
    name="deepseek_agent",                  # Agent 名称
    model=LiteLlm(                          # 使用 LiteLLM 适配器
        model="deepseek/deepseek-chat",     # 格式：provider/model_name
        api_key="你的_DEEPSEEK_API_KEY",    # API Key
    ),
    instruction="你是一个中文写作助手。",     # 系统指令
)

# 使用 LiteLLM 适配 OpenAI 模型
openai_agent = Agent(
    name="openai_agent",                    # Agent 名称
    model=LiteLlm(                          # 使用 LiteLLM 适配器
        model="openai/gpt-4o",              # OpenAI GPT-4o
        api_key="你的_OPENAI_API_KEY",      # API Key
    ),
    instruction="你是一个数据分析助手。",     # 系统指令
)


# ========================================
# 示例二：使用 Google 内置工具
# ========================================

# 带搜索能力的 Agent
search_agent = Agent(
    name="search_agent",                    # Agent 名称
    model="gemini-2.0-flash",               # 使用 Gemini 模型
    instruction=(                           # 系统指令
        "你是一个信息查询助手。\n"
        "使用 Google Search 工具搜索最新信息。\n"
        "根据搜索结果提供准确、全面的回答。"
    ),
    tools=[google_search],                  # 注册 Google 搜索工具
)

# 带代码执行能力的 Agent
code_agent = Agent(
    name="code_agent",                      # Agent 名称
    model="gemini-2.0-flash",               # 使用 Gemini 模型
    instruction=(                           # 系统指令
        "你是一个数据分析助手。\n"
        "使用 Code Execution 工具执行 Python 代码。\n"
        "确保代码正确且高效。"
    ),
    tools=[code_execution],                 # 注册代码执行工具
)

# 多功能 Agent（搜索 + 代码执行）
super_agent = Agent(
    name="super_agent",                     # Agent 名称
    model="gemini-2.0-flash",               # 使用 Gemini 模型
    instruction="你是一个全能助手，可以搜索信息和执行代码。",  # 指令
    tools=[google_search, code_execution],  # 注册多个工具
)


# ========================================
# 示例三：output_key 用法
# ========================================

# 第一步：提取信息的 Agent
extractor = Agent(
    name="Extractor",                       # 信息提取 Agent
    model="gemini-2.0-flash",               # 模型
    instruction="从用户输入中提取关键信息。只输出提取的信息。",  # 指令
    output_key="extracted_info",            # 将输出保存到 state['extracted_info']
)

# 第二步：分析信息的 Agent（引用第一步的输出）
analyzer = Agent(
    name="Analyzer",                        # 信息分析 Agent
    model="gemini-2.0-flash",               # 模型
    instruction=(                           # 指令：通过 {key} 引用 state 中的值
        "分析以下提取的信息，给出深度分析。\n"
        "提取的信息：{extracted_info}"
    ),
    output_key="analysis_result",           # 将分析结果保存到 state
)

# 创建顺序流水线
pipeline = SequentialAgent(
    name="InfoPipeline",                    # 流水线名称
    sub_agents=[extractor, analyzer],       # 按顺序执行
)


# ========================================
# 示例四：Agent Transfer（动态委派）
# ========================================

# 定义专门的子 Agent
booking_agent = LlmAgent(
    name="BookingAgent",                    # 预订 Agent
    model="gemini-2.0-flash",               # 模型
    description="负责处理航班和酒店的预订请求。",  # 描述
    instruction="你是一个预订助手，帮助用户预订航班和酒店。",  # 指令
)

info_agent = LlmAgent(
    name="InfoAgent",                       # 信息查询 Agent
    model="gemini-2.0-flash",               # 模型
    description="负责回答一般性问题和提供信息。",  # 描述
    instruction="你是一个信息助手，回答用户的各种问题。",  # 指令
)

# 定义协调器 Agent
coordinator = LlmAgent(
    name="Coordinator",                     # 协调器
    model="gemini-2.0-flash",               # 模型
    description="主协调器，负责分配任务。",  # 描述
    instruction=(                           # 指令
        "你是一个协调器。\n"
        "预订任务交给 BookingAgent。\n"
        "信息查询交给 InfoAgent。\n"
        "使用 transfer_to_agent 工具委派任务。"
    ),
    sub_agents=[booking_agent, info_agent], # 注册子 Agent
)


# ========================================
# 示例五：通过代码运行 Agent
# ========================================

# 定义应用信息
APP_NAME = "llm_demo_app"                   # 应用名称
USER_ID = "user_001"                        # 用户 ID
SESSION_ID = "session_001"                  # 会话 ID


async def run_demo():
    """运行 LlmAgent 演示"""

    # 创建会话服务
    session_service = InMemorySessionService()  # 内存会话服务

    # 创建会话
    session = await session_service.create_session(
        app_name=APP_NAME,                   # 应用名称
        user_id=USER_ID,                     # 用户 ID
        session_id=SESSION_ID,               # 会话 ID
    )

    # 创建运行器
    runner = Runner(
        agent=gemini_agent,                  # 使用 Gemini Agent
        app_name=APP_NAME,                   # 应用名称
        session_service=session_service,     # 会话服务
    )

    # 构造用户消息
    content = types.Content(
        role='user',                         # 角色：用户
        parts=[types.Part(text='用 Python 写一个冒泡排序算法')],  # 消息
    )

    # 运行 Agent
    events = runner.run_async(
        user_id=USER_ID,                     # 用户 ID
        session_id=SESSION_ID,               # 会话 ID
        new_message=content,                 # 用户消息
    )

    # 处理事件流
    async for event in events:               # 遍历所有事件
        if event.is_final_response():        # 如果是最终响应
            response = event.content.parts[0].text  # 提取文本
            print(f"Agent 回复:\n{response}")  # 打印回复


# 运行演示
if __name__ == "__main__":                  # 当脚本被直接运行时
    asyncio.run(run_demo())                  # 执行异步函数