## 概念性的类 Python 结构，非可运行代码
from google.adk.agents import Agent
## from google.adk.models.lite_llm import LiteLlm  # 如果使用 ADK 默认智能体不直接支持的模型
## 使用更昂贵的 Gemini Pro 2.5 的智能体
gemini_pro_agent = Agent(
    name="GeminiProAgent",
    model="gemini-2.5-pro",  # 如果实际模型名称不同,这是占位符
    description="一个用于复杂查询的高能力智能体。",
    instruction="您是一个专门解决复杂问题的专家助手。"
)

## 使用更便宜的 Gemini Flash 2.5 的智能体
gemini_flash_agent = Agent(
    name="GeminiFlashAgent",
    model="gemini-2.5-flash",  # 如果实际模型名称不同,这是占位符
    description="一个用于简单查询的快速高效智能体。",
    instruction="您是一个处理简单问题的快速助手。"
)

## 概念性的类 Python 结构,非可运行代码
from google.adk.agents import Agent, BaseAgent
from google.adk.events import Event
from google.adk.agents.invocation_context import InvocationContext
import asyncio

class QueryRouterAgent(BaseAgent):
    name: str = "QueryRouter"
    description: str = "根据复杂性将用户查询路由到适当的LLM智能体。"

    async def _run_async_impl(self, context: InvocationContext) -> AsyncGenerator[Event, None]:
        user_query = context.current_message.text  # 假设文本输入
        query_length = len(user_query.split())  # 简单指标：单词数

        if query_length < 20:  # 示例阈值，用于简单性与复杂性的区分
            print(f"Routing to Gemini Flash Agent for short query (length: {query_length})")
            # 在真实的 ADK 设置中，您会使用 'transfer_to_agent' 或直接调用
            # 为了演示，我们将模拟一个调用并产生其响应
            response = await gemini_flash_agent.run_async(context.current_message)
            yield Event(author=self.name, content=f"Flash Agent processed: {response}")
        else:
            print(f"Routing to Gemini Pro Agent for long query (length: {query_length})")
            response = await gemini_pro_agent.run_async(context.current_message)
            yield Event(author=self.name, content=f"Pro Agent processed: {response}")

CRITIC_SYSTEM_PROMPT = """
您是**批评智能体**,作为我们协作研究助手系统的质量保证部门。您的主要功能是**细致审查和质疑**来自研究智能体的信息,确保**准确性、完整性和无偏见的呈现**。您的职责包括:
* **评估研究发现**的事实正确性、全面性和潜在倾向。
* **识别任何缺失数据**或推理中的不一致。
* **提出关键问题**以改进或扩展当前理解。
* **提供建设性建议**以增强或探索不同角度。
* **验证最终输出是否全面**且平衡。
所有批评必须是建设性的。您的目标是加强研究,而非否定它。清晰组织您的反馈,突出需要修订的具体要点。您的首要目标是确保最终研究产品达到尽可能高的质量标准。
"""

## MIT License
## Copyright (c) 2025 Mahtab Syed
## https://www.linkedin.com/in/mahtabsyed/

import os
import requests
import json
from dotenv import load_dotenv
from openai import OpenAI

## 加载环境变量
load_dotenv()

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
GOOGLE_CUSTOM_SEARCH_API_KEY = os.getenv("GOOGLE_CUSTOM_SEARCH_API_KEY")
GOOGLE_CSE_ID = os.getenv("GOOGLE_CSE_ID")

if not OPENAI_API_KEY or not GOOGLE_CUSTOM_SEARCH_API_KEY or not GOOGLE_CSE_ID:
    raise ValueError(
        "Please set OPENAI_API_KEY, GOOGLE_CUSTOM_SEARCH_API_KEY, and GOOGLE_CSE_ID in your .env file."
    )

client = OpenAI(api_key=OPENAI_API_KEY)

## --- 步骤 1：分类提示词 ---
def classify_prompt(prompt: str) -> dict:
    system_message = {
        "role": "system",
        "content": (
            "You are a classifier that analyzes user prompts and returns one of three categories ONLY:\n\n"
            "- simple\n"
            "- reasoning\n"
            "- internet_search\n\n"
            "Rules:\n"
            "- Use 'simple' for direct factual questions that need no reasoning or current events.\n"
            "- Use 'reasoning' for logic, math, or multi-step inference questions.\n"
            "- Use 'internet_search' if the prompt refers to current events, recent data, or things not in your training data.\n\n"
            "Respond ONLY with JSON like:\n"
            '{ "classification": "simple" }'
        ),
    }
    user_message = {"role": "user", "content": prompt}
    response = client.chat.completions.create(
        model="gpt-4o", messages=[system_message, user_message], temperature=1
    )
    reply = response.choices[0].message.content
    return json.loads(reply)

## --- 步骤 2：Google 搜索 ---
def google_search(query: str, num_results=1) -> list:
    url = "https://www.googleapis.com/customsearch/v1"
    params = {
        "key": GOOGLE_CUSTOM_SEARCH_API_KEY,
        "cx": GOOGLE_CSE_ID,
        "q": query,
        "num": num_results,
    }
    try:
        response = requests.get(url, params=params)
        response.raise_for_status()
        results = response.json()
        if "items" in results and results["items"]:
            return [
                {
                    "title": item.get("title"),
                    "snippet": item.get("snippet"),
                    "link": item.get("link"),
                }
                for item in results["items"]
            ]
        else:
            return []
    except requests.exceptions.RequestException as e:
        return {"error": str(e)}

## --- 步骤 3：生成响应 ---
def generate_response(prompt: str, classification: str, search_results=None) -> str:
    if classification == "simple":
        model = "gpt-4o-mini"
        full_prompt = prompt
    elif classification == "reasoning":
        model = "o4-mini"
        full_prompt = prompt
    elif classification == "internet_search":
        model = "gpt-4o"
        # 将每个搜索结果字典转换为可读字符串
        if search_results:
            search_context = "\n".join(
                [
                    f"Title: {item.get('title')}\nSnippet: {item.get('snippet')}\nLink: {item.get('link')}"
                    for item in search_results
                ]
            )
        else:
            search_context = "未找到搜索结果。"
        full_prompt = f"""使用以下网络结果回答用户查询：{search_context} 查询：{prompt}"""

    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": full_prompt}],
        temperature=1,
    )
    return response.choices[0].message.content, model

## --- 步骤 4：组合路由器 ---
def handle_prompt(prompt: str) -> dict:
    classification_result = classify_prompt(prompt)
    # 删除或注释掉下一行以避免重复打印
    # print("\n🔍 Classification Result:", classification_result)
    classification = classification_result["classification"]

    search_results = None
    if classification == "internet_search":
        search_results = google_search(prompt)
        # print("\n🔍 Search Results:", search_results)

    answer, model = generate_response(prompt, classification, search_results)
    return {"classification": classification, "response": answer, "model": model}

test_prompt = "What is the capital of Australia?"
## test_prompt = "Explain the impact of quantum computing on cryptography."
## test_prompt = "When does the Australian Open 2026 start, give me full date?"
result = handle_prompt(test_prompt)
print("🔍 Classification:", result["classification"])
print("🧠 Model Used:", result["model"])
print("🧠 Response:\n", result["response"])

import requests
import json

response = requests.post(
  url="https://openrouter.ai/api/v1/chat/completions",
  headers={
    "Authorization": "Bearer <OPENROUTER_API_KEY>",
    "HTTP-Referer": "<YOUR_SITE_URL>",  # 可选。用于 openrouter.ai 上排名的网站 URL。
    "X-Title": "<YOUR_SITE_NAME>",      # 可选。用于 openrouter.ai 上排名的网站标题。
  },
  data=json.dumps({
    "model": "openai/gpt-4o",  # 可选
    "messages": [
      {
        "role": "user",
        "content": "What is the meaning of life?"
      }
    ]
  })
)

{
  "model": "openrouter/auto",
  ... // 其他参数
}

{
  "models": ["anthropic/claude-3.5-sonnet", "gryphe/mythomax-l2-13b"],
  ... // 其他参数
}