"""
测试流式接口的完整脚本
API流程：
1. 登录获取token
2. 创建thread（可选，也可以使用无状态接口）
3. 调用流式接口

两种方式：
- 方式1: POST /api/runs/stream (无状态，自动创建thread)
- 方式2: POST /api/threads/{thread_id}/runs/stream (需要先创建thread)
"""
import requests
import json
import asyncio
import httpx  # 异步 HTTP 客户端
# ============ 配置区域 ============
BASE_URL = "http://127.0.0.1:8001"
LOGIN_URL = f"{BASE_URL}/api/v1/auth/login/username"
THREADS_URL = f"{BASE_URL}/api/threads"
RUNS_STREAM_URL = f"{BASE_URL}/api/runs/stream"  # 无状态流式接口

# ============ 登录获取token ============
def get_token(username: str) -> str:
    """登录获取token"""
    response = requests.post(
        LOGIN_URL,
        json={"username": username},
        headers={"Content-Type": "application/json"}
    )
    # print(f"登录状态码: {response.status_code}")
    data = response.json()
    # print(f"登录响应: {json.dumps(data, indent=2, ensure_ascii=False)}")
    
    # 根据实际返回结构调整
    token = data.get("token") or data.get("access_token") or data.get("data", {}).get("token")
    return token
# ============ 方式2: 先创建thread，再调用流式接口 ============
def create_thread(token: str) -> str:
    """创建thread，返回thread_id"""
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {token}"
    }
    
    response = requests.post(
        THREADS_URL,
        json={},  # 空body，使用默认值
        headers=headers
    )
    
    print(f"创建thread状态码: {response.status_code}")
    data = response.json()
    print(f"创建thread响应: {json.dumps(data, indent=2, ensure_ascii=False)}")
    
    return data.get("thread_id")
def get_agent_name_description(agent_name="Taiwan-intelligence-gathering-agent"):
    url = f"{BASE_URL}/api/agents/{agent_name}"
    token=get_token("ji")
    headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {token}"
        }
    response = requests.get(
            url,
            headers=headers,
        )
    print(response.json())
    return response.json()["name"].lower(),"agent_name:"+response.json()["name"].lower()+","+response.json()["description"]

# ============ 创建智能体 ============
def create_agent(
    name: str,
    description: str = "",
    model: str | None = None,
    tool_groups: list[str] | None = None,
    skills: list[str] | None = None,
    soul: str = "",
):
    """
    创建一个新的智能体
    
    Args:
        name: 智能体名称（必须是 hyphen-case，如 my-agent）
        description: 智能体描述
        model: 可选的模型覆盖（如 gpt-4, claude-3-opus 等）
        tool_groups: 工具组白名单（如 ["web", "file"]）
        skills: 技能白名单（如 ["skill1", "skill2"]）
        soul: SOUL.md 内容（智能体人格/身份描述）
    
    Returns:
        创建结果
    """
    TOKEN =get_token("ji")
    url = f"{BASE_URL}/api/agents"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {TOKEN}"
    }
    
    payload = {
        "name": name,
        "description": description,
        "soul": soul
    }
    
    # 添加可选字段
    if model is not None:
        payload["model"] = model
    if tool_groups is not None:
        payload["tool_groups"] = tool_groups
    if skills is not None:
        payload["skills"] = skills
    
    print(f"请求: POST {url}")
    print(f"请求体: {json.dumps(payload, indent=2, ensure_ascii=False)}")
    
    response = requests.post(url, json=payload, headers=headers)
    
    print(f"状态码: {response.status_code}")
    
    if response.status_code == 201:
        result = response.json()
        print(f"✓ 智能体创建成功!")
        print(f"响应: {json.dumps(result, indent=2, ensure_ascii=False)}")
        return result
    else:
        print(f"✗ 创建失败: {response.text}")
        return f"✗ 创建失败: {response.text}"

def create_main_agent(zhihui_name,agent_names):
    d={}
    token=get_token("ji")
    l=[]
    for i in agent_names:
        # print(i)
        name,description=get_agent_name_description(i)
        thread_id_tmp=create_thread(token)
        d[name]=thread_id_tmp
        l.append(description)
    name="你是总控控制协调智能体,负责协调其他智能体干活"
    miaoshu=f"""#核心定位：
    用途：跟用用户的任务目标,调用其他智能体干活,只负责协调汇总,当你觉得已有信息支持回答用户问题时候,停止调度,并输出最后的结果
    #调度方法：
    通过使用agent_orchestration技能调用其他智能体干活,使用技能时传入agent_name和需要做的活
    #可调度智能体：
    {l}"""
    thread_id_tmp=create_thread(token)
    engname=zhihui_name
    d[engname]=thread_id_tmp
    tmp_skills=["agent_orchestration"]
    tmpres=create_agent(
            name=engname,
            description=name,
            model="glm-5",  # 指定模型
            skills=tmp_skills,  # 关联技能
            soul=miaoshu)
    # print("创建智能体响应:",tmpres)
    if "失败" in tmpres and "exists" not in tmpres:
        return tmpres
    print(f"创建协调智能体成功,可以协调列表:{l}")
    return d
import re
def extract_name_task(text="python3 scripts/agent_orchestration.py \"Us-intelligence-gathering-agent\" \"收集和分析美国人民对特朗普的看法，包括支持率和反对率、不同群体的观点差异、主要支持理由和反对理由等\""):
    try:
        pattern = r'\.py\s+"([^"]*)"\s+"([^"]*)"'
        match = re.search(pattern, text)
        if match:
            agent_name = match.group(1)
            task_content = match.group(2)
        return agent_name,task_content
    except:
        return "",""
def update_agent_message(thread_id,new_content):
    token=get_token("ji")
    url = f"{BASE_URL}/api/threads/{thread_id}/state"
    headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {token}"
        }
    tmp = requests.get(
            url,
            headers=headers
        ).json()
    # print(tmp)
    if "messages" in tmp["values"]:
         tmp["values"]["messages"].append({"type": "human", "content": new_content})
    else:
         tmp["values"]["messages"]=[]
         tmp["values"]["messages"].append({"type": "human", "content": new_content})
    url = f"{BASE_URL}/api/threads/{thread_id}/state"
    headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {token}"
        }
    response = requests.post(
            url,
            json=tmp,
            headers=headers,
            stream=True
        )
    del tmp
    print(f"线程{thread_id}更新状态码: {response.status_code}")
def Broadcast_message(message,d_agent_thread_id,pingbi_agent_name):
    for key in d_agent_thread_id:
            if key not in pingbi_agent_name:
                update_agent_message(d_agent_thread_id[key],message)
                print("广播消息成功:",key,message)
null=None
true=True
false=False
async def leader_agent(d_agent_thread_id, agent_name, new_message, skill_stop_names):
    """使用 httpx 异步请求的 leader_agent"""
    # 使用 to_thread 异步运行同步的 Broadcast_message
    await asyncio.to_thread(Broadcast_message, "总控智能体收到用户消息:" + new_message, d_agent_thread_id, [agent_name])
    thread_id = d_agent_thread_id[agent_name]
    l = []
    token = get_token("ji")
    url = f"{BASE_URL}/api/threads/{thread_id}/runs/stream"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {token}"
    }
    payload = {
        "input": {
            "messages": [
                {
                    "type": "human",
                    "content": [
                        {
                            "type": "text",
                            "text": new_message
                        }
                    ],
                    "additional_kwargs": {}
                }
            ]
        },
        "config": {
            "recursion_limit": 1000
        },
        "context": {
            "agent_name": agent_name,
            "model_name": "deepseek-chat",
            "mode": "pro",
            "reasoning_effort": "medium",
            "thinking_enabled": true,
            "is_plan_mode": false,
            "subagent_enabled": false,
            "thread_id": thread_id
        },
        "stream_mode": [
            "messages-tuple",
            "values",
        ],
        "stream_subgraphs": true,
        "stream_resumable": true,
        "assistant_id": "lead_agent",
        "on_disconnect": "continue",
        "excluded_tools": ["web_search"],
        "skill_stop_names": skill_stop_names,
    }
    
    async with httpx.AsyncClient(timeout=None) as client:
        async with client.stream("POST", url, json=payload, headers=headers) as response:
            print(f"状态码: {response.status_code}")
            if response.status_code != 200:
                body = await response.aread()
                raise RuntimeError(
                    f"upstream {url} returned HTTP {response.status_code}: "
                    f"{body.decode('utf-8', 'replace')[:1000]}"
                )
            async for line in response.aiter_lines():
                if line:
                    yield line
                    if line[:5] == "data:":
                        print(line)
                        l.append(line)
    if len(l) < 2:
        raise RuntimeError(
            f"upstream stream produced no usable data lines (got {len(l)} line(s)). "
            f"Check the server log on {url} for the real failure."
        )
    
    null = None
    tmpres = eval(l[-2][5:])
    tmpres = tmpres["messages"][-1]
    tmp_task = []
    if tmpres.get("tool_calls", []) == []:
        yield {"status": []}
    else:
        for tool_call in tmpres.get("tool_calls", []):
            try:
                agentname, task_text = extract_name_task(str(tool_call["args"]))
                print("agentname", agentname)
                print("task_text:", task_text)
                # 使用 to_thread 异步运行同步的 Broadcast_message
                await asyncio.to_thread(Broadcast_message, f"总控智能体给{agentname}派活,内容为:" + task_text, d_agent_thread_id, [agent_name, agentname])
                tmp_task.append([agentname, task_text])
            except:
                pass
        yield {"status": tmp_task}


async def special_agent(d_agent_thread_id, agent_name, new_message, skill_stop_names):
    """使用 httpx 异步请求的 special_agent"""
    thread_id = d_agent_thread_id[agent_name]
    l = []
    token = get_token("ji")
    url = f"{BASE_URL}/api/threads/{thread_id}/runs/stream"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {token}"
    }
    payload = {
        "input": {
            "messages": [
                {
                    "type": "human",
                    "content": [
                        {
                            "type": "text",
                            "text": new_message
                        }
                    ],
                    "additional_kwargs": {}
                }
            ]
        },
        "config": {
            "recursion_limit": 1000
        },
        "context": {
            "agent_name": agent_name,
            "model_name": "deepseek-chat",
            "mode": "pro",
            "reasoning_effort": "medium",
            "thinking_enabled": true,
            "is_plan_mode": false,
            "subagent_enabled": false,
            "thread_id": thread_id
        },
        "stream_mode": [
            "messages-tuple",
            "values",
        ],
        "stream_subgraphs": true,
        "stream_resumable": true,
        "assistant_id": "lead_agent",
        "on_disconnect": "continue",
        "excluded_tools": ["web_search", "ask_clarification", "present_files", "view_image"],
        "skill_stop_names": skill_stop_names,
    }
    
    async with httpx.AsyncClient(timeout=None) as client:
        async with client.stream("POST", url, json=payload, headers=headers) as response:
            print(f"状态码: {response.status_code}")
            if response.status_code != 200:
                body = await response.aread()
                raise RuntimeError(
                    f"upstream {url} returned HTTP {response.status_code}: "
                    f"{body.decode('utf-8', 'replace')[:1000]}"
                )
            async for line in response.aiter_lines():
                if line:
                    yield line
                    if line[:5] == "data:":
                        print(line)
                        l.append(line)
    if len(l) < 2:
        raise RuntimeError(
            f"upstream stream produced no usable data lines (got {len(l)} line(s)). "
            f"Check the server log on {url} for the real failure."
        )
    
    null = None
    tmpres = eval(l[-2][5:])
    tmpres = tmpres["messages"][-1]
    # 使用 to_thread 异步运行同步的 Broadcast_message
    await asyncio.to_thread(Broadcast_message, f"子智能体{agent_name}完成{new_message}工作,交付内容为:" + tmpres["content"], d_agent_thread_id, [agent_name])
    yield {"status": f"子智能体{agent_name}完成{new_message}工作"}

async def agent_init(agent_names, main_agent_name):
    """异步的 agent_init，同步函数调用使用 to_thread"""
    agent_names = [i.lower() for i in agent_names]
    # 使用 to_thread 异步运行同步的 create_main_agent
    tmpd = await asyncio.to_thread(create_main_agent, main_agent_name, agent_names)
    message = """
你的职责:你只是若干个智能体的其中一员,你们的共同目标是完成一个庞大的任务,总控控制协调智能体会分配特定的活给你,你只需要用你自身的能力来完成这个特定的活
注意:你只是一个子智能体,你不必把你的分析结果存放到文件里面去,你只关心于完成总控控制协调智能体分配的活即可
"""
    # 使用 to_thread 异步运行同步的 Broadcast_message
    await asyncio.to_thread(Broadcast_message, message, tmpd, [main_agent_name])
    return main_agent_name, tmpd


from fastapi import FastAPI, HTTPException,Body
from fastapi.responses  import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
import asyncio
import uvicorn

app = FastAPI()
 # CORS 配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

import time
@app.post("/generate_mul_agent")
async def generate_stream(payload: dict = Body(...)):
    """
    根据 payload 中的 agent_type 字段判断调用 leader_agent 还是 special_agent
    
    payload 结构示例:
    {
        "agent_type": "leader",  # 或 "special"
        "d_agent_thread_id": {...},
        "agent_name": "xxx",
        "new_message": "xxx",
        "skill_stop_names": [...]
    }
    """
    # 从 payload 中提取参数
    agent_type = payload.get("agent_type", "leader")  # 默认为 leader
    d_agent_thread_id = payload.get("d_agent_thread_id", {})
    agent_name = payload.get("agent_name", "")
    new_message = payload.get("new_message", "")
    skill_stop_names = payload.get("skill_stop_names", [])
    
    # 流式响应生成器
    async def stream_generator():
        # 根据 agent_type 选择调用的函数
        if agent_type == "leader":
            async_generator = leader_agent(d_agent_thread_id, agent_name, new_message, skill_stop_names)
        else:  # special
            async_generator = special_agent(d_agent_thread_id, agent_name, new_message, skill_stop_names)
        
        # 异步迭代生成器
        async for chunk in async_generator:
            if isinstance(chunk, dict):
                yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n"
            else:
                yield f"{chunk}"
    
    # 使用 SSE 流式响应
    return StreamingResponse(
        stream_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
            "Connection": "keep-alive",
        }
    )


@app.post("/agent_init")
async def init_agents(payload: dict = Body(...)):
    """
    初始化智能体接口
    payload 结构示例:
    {
        "agent_names": ["agent1", "agent2"],
        "main_agent_name": "main-agent"
    }
    """
    agent_names = payload.get("agent_names", [])
    main_agent_name = payload.get("main_agent_name", "main-agent")
    
    # 调用异步的 agent_init
    name, thread_id_dict = await agent_init(agent_names, main_agent_name)
    
    return {
        "status": "success",
        "main_agent_name": name,
        "thread_ids": thread_id_dict
    }
if __name__ == "__main__": 
    uvicorn.run("testapi:app",  host="0.0.0.0", port=8095, workers=1,limit_concurrency=1000)  # 多进程提升并发能力