#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   server.py
@Time    :   2025/10/27 15:43:49
@Author  :   Anderson Huang 
@Version :   1.0
@Contact :   huanganfa66@163.com
@Desc    :   None
'''

# here put the import lib
import os
import asyncio
import argparse
from typing import List
from agents import Runner
from openai import AsyncOpenAI
from mcp.server.fastmcp import FastMCP
from agents import set_default_openai_client, set_tracing_disabled

# 从.env文件加载环境变量
from dotenv import load_dotenv
load_dotenv()  

from deepresearch_agents.planner_agent import planner_agent, WebSearchItem, WebSearchPlan
from deepresearch_agents.search_agent import search_agent
from deepresearch_agents.writer_agent import writer_agent, ReportData

# 初始化MCP服务器
mcp = FastMCP(name="DeepResearch")
USER_AGENT = "Deepresearch-app/1.0"

# ################# 辅助函数 #########################

async def _plan_searches(query: str) -> WebSearchPlan:
    """
    用于进行某个搜索主题的搜索计划。
    args:
        query(str): 搜索主题
    returns:
        WebSearchPlan: 搜索计划
    """
    result = await Runner.run(
        starting_agent=planner_agent, 
        input=f"Query: {query}"
        )
    return result.final_output_as(WebSearchPlan)

async def _search(item: WebSearchItem) -> str:
    """
    用于进行某个搜索主题的搜索。
    args:
        item(WebSearchItem): 搜索项
    returns:
        str: 搜索结果
    """
    try:
        result = await Runner.run(
            starting_agent=search_agent,
            input=f"Search term: {item.query}\nReason for searching: {item.reason}"
        )
        return result.final_output_as(str)
    except Exception as e:
        return None

async def _execute_searches(plan: WebSearchPlan) -> List[str]:
    """
    执行搜索计划，返回搜索结果。

    args:
        plan(WebSearchPlan): 搜索计划
    returns:
        List[str]: 搜索结果
    """
    tasks = [
        asyncio.create_task(
            _search(item)
        )
        for item in plan.items
    ]
    results = []
    for task in asyncio.as_completed(tasks):
        result = await task
        if result is not None:
            results.append(result)
    return results

async def _write_report(query: str, search_results: List[str]) -> ReportData:
    """
    用于根据搜索结果编写报告。

    args:
        query(str): 搜索主题
        search_results(List[str]): 搜索结果

    returns:
        ReportData: 报告数据
    """
    result = await Runner.run(
        starting_agent=writer_agent,
        input=f"Original query: {query}\nSummaried search results: {search_results}"
    )
    return result.final_output_as(ReportData)

# ################# MCP服务器主函数 #########################

@mcp.tool()
async def deep_research(query: str) -> ReportData:
    """
    用于进行某个搜索主题的深度研究。
    
    args:
        query(str): 搜索主题
        
    returns:
        ReportData: 报告数据
    """
    search_plan = await _plan_searches(query)
    search_results = await _execute_searches(search_plan)
    report = await _write_report(query, search_results)
    return report

def main():
    # 从环境变量获取默认值，优先级：环境变量 > 代码默认值
    # default_api_key = os.getenv("OPENAI_API_KEY", None)
    default_base_url = os.getenv("BASE_URL", "https://api.openai.com/v1")
    default_tracing_disabled = os.getenv("TRACING_DISABLED", "true").lower() == "true"
    
    parser = argparse.ArgumentParser(description="DeepResearch Server")
    parser.add_argument("--openai_api_key", type=str, required=True, help="Your Openai API Key")
    parser.add_argument("--base_url", type=str, required=False, default=default_base_url, help="Openai API base url")
    parser.add_argument("--tracing_disabled", type=bool, required=False, default=default_tracing_disabled, help="是否禁用OpenAI客户端跟踪")
    args = parser.parse_args()

    # OpenAI客户端
    openai_client = AsyncOpenAI(
        base_url = args.base_url,
        api_key = args.openai_api_key
    )
    # 设置默认OpenAI客户端
    set_default_openai_client(openai_client)
    # 禁用OpenAI客户端的跟踪
    set_tracing_disabled(args.tracing_disabled)

    # 以标准I/O方式启动MCP服务器
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
    












