import os
import asyncio

from agentscope.agent import ReActAgent
from agentscope.formatter import OpenAIChatFormatter
from agentscope.memory import InMemoryMemory
from agentscope.model import OpenAIChatModel
from agentscope.tool import Toolkit, execute_python_code
from agentscope.message import ImageBlock, Msg, TextBlock, URLSource

LLAMA_CPP_BASE_URL = os.getenv("LLAMA_CPP_BASE_URL", "http://yyplab.site:8033/v1")
LLAMA_CPP_MODEL_NAME = os.getenv("LLAMA_CPP_MODEL_NAME", "Qwen3.5-35B-A3B-UD-Q8_K_XL.gguf")
LLAMA_CPP_API_KEY = os.getenv("LLAMA_CPP_API_KEY", "EMPTY")


async def main():
    toolkit = Toolkit()
    toolkit.register_tool_function(execute_python_code)

    yyp_bot = ReActAgent(
        name="yyp_bot",
        sys_prompt="你是一个名为 yyp_bot 的助手",
        model=OpenAIChatModel(
            model_name=LLAMA_CPP_MODEL_NAME,
            api_key=LLAMA_CPP_API_KEY,
            client_kwargs={"base_url": LLAMA_CPP_BASE_URL},
            stream=True,
        ),
        memory=InMemoryMemory(),
        formatter=OpenAIChatFormatter(),
        toolkit=toolkit,
    )

    while True:
        user_text = input("请输入问题：").strip()
        if user_text == "exit":
            break

        image_path = input("请输入图片路径，留空则不发送图片：").strip()

        if image_path:
            msg = Msg(
                name="user",
                role="user",
                content=[
                    TextBlock(
                        type="text",
                        text=user_text,
                    ),
                    ImageBlock(
                        type="image",
                        source=URLSource(
                            type="url",
                            url=image_path,
                        ),
                    ),
                ],
            )
        else:
            msg = Msg(
                name="user",
                role="user",
                content=user_text,
            )

        await yyp_bot(msg)

asyncio.run(main())