import json
from datetime import datetime
from langchain_ollama import ChatOllama
from langchain_core.tools import tool
from langchain.agents import create_agent

# ------ 1. 定义工具
OLLAMA_URL = "http://localhost:11434"
@tool
def gen_map(prompt: str) -> str:
    """生成2D游戏地图"""
    print('gen_map: ', prompt)
    # 调用第三方文生图大模型
    return "http://127.0.0.1/map.jpg"

@tool
def gen_role(prompt: str) -> str:
    """生成2D游戏角色"""
    print('gen_role: ', prompt)
    # 调用第三方文生图大模型
    return "http://127.0.0.1/role.png"
    
@tool
def gen_sound(prompt: str) -> str:
    """生成2D游戏背景音乐"""
    print('gen_sound: ', prompt)
    # 调用第三方文生音乐大模型
    return "http://127.0.0.1/sound.mp3"
    
@tool
def gen_config(prompt: str) -> str:
    """生成2D游戏关卡配置文件"""
    print('gen_config: ', prompt)
    # 调用第三方大模型(deepseek、qwen、hunyuan,...)
    return "http://127.0.0.1/config.json"

tools = [gen_map, gen_role, gen_sound, gen_config]

# ------ 2. 初始化模型
llm = ChatOllama(
    model="qwen2.5:7b",
    base_url=OLLAMA_URL,
    temperature=0,
)

# ------ 3. 创建并运行 Agent
agent = create_agent(model=llm, tools=tools)

# 非流式输出
result = agent.invoke({"messages": [("user", """
任务：
生成一张卡通风格的横版游戏地图。生成一张卡通风格的古装英雄角色。生成一段紧张的游戏背景音乐。生成第3关配置文件。
输出json格式：
{
 "map": url,
 "role": url,
 "sound": url,
 "config": url
}
""")]})

print()
print("输出结果：")
print(result['messages'][-1].content)

# 将所有生成的资源汇总成一个zip

运行测试