import asyncio
from autogen_agentchat.teams import SelectorGroupChat
from autogen_agentchat.conditions import TextMentionTermination
from autogen_agentchat.ui import Console
from autogen_agentchat.teams import RoundRobinGroupChat, SelectorGroupChat

from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.conditions import TextMentionTermination
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_core.memory import ListMemory, MemoryContent, MemoryMimeType

def get_model_client() -> OpenAIChatCompletionClient:
    return OpenAIChatCompletionClient(
        model="qwq-plus",
        api_key="",
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        model_info={
            "json_output": False,
            "vision": False,
            "function_calling": True,
            "family": "unknown"
        }
    )

model_client = get_model_client()


async def create_assistant():
    # 确保此函数内所有异步操作都使用await
    # 示例内存初始化（根据文献[6]优化内存管理）
    user_memory = ListMemory()
    await user_memory.add(
        MemoryContent(
            content="The user system is linux",
            mime_type=MemoryMimeType.TEXT
        )
    )

    # 初始化模型客户端（需补充实际配置）
    # model_client = model_client

    return AssistantAgent(
        name="assistant_agent",
        model_client=model_client,
        memory=[user_memory],  # 直接传递内存实例[6](@ref)
        system_message='''...''',  # 保持原有系统消息
        model_client_stream=True
    )

async def create_assistant2():
    # 确保此函数内所有异步操作都使用await
    # 示例内存初始化（根据文献[6]优化内存管理）
    user_memory = ListMemory()
    await user_memory.add(
        MemoryContent(
            content="The user system is linux",
            mime_type=MemoryMimeType.TEXT
        )
    )

    # 初始化模型客户端（需补充实际配置）
    # model_client = model_client

    return AssistantAgent(
        name="assistant_agent2",
        model_client=model_client,
        memory=[user_memory],  # 直接传递内存实例[6](@ref)
        system_message='''...''',  # 保持原有系统消息
        model_client_stream=True
    )


async def main():
    # 异步初始化所有代理（根据文献[3]优化并行初始化）
    assistant= await create_assistant()
    assistant2= await create_assistant2()

    # 构建团队（根据文献[7]优化团队配置）
    team = SelectorGroupChat(
        [assistant, assistant2],
        termination_condition=TextMentionTermination("APPROVE"),
        model_client=model_client,
        selector_prompt="""...""",  # 保持原有提示模板
    )

    # 执行任务流（根据文献[8]优化流处理）
    # try:
    stream = team.run_stream(task=input("Hello! How can I help?\n"))
    #stream = team.run_stream(task=input("Hello! How can I help?\n"))
    await Console(stream, output_stats=True)
    # except asyncio.CancelledError:
    #     print("Operation cancelled by user")
    # finally:
    #     await team.close()  # 确保资源释放

if __name__ == "__main__":
    # 标准事件循环入口（根据文献[4][9]规范）
    asyncio.run(main())

    raise self._make_status_error_from_response(err.response) from None
openai.BadRequestError: Error code: 400 - {'error': {'code': 'invalid_parameter_error', 'param': None, 'message': 'This model only support stream mode, please enable the stream parameter to access the model. ', 'type': 'invalid_request_error'}, 'id': 'chatcmpl-fb835624-c232-90e1-81f6-0b3b19c1e633', 'request_id': 'fb835624-c232-90e1-81f6-0b3b19c1e633'}