from julep import AsyncJulep

client = AsyncJulep(api_key="your_api_key")

agent = await client.agents.create(
    name="讲故事代理",
    model="gpt-4-turbo",
    about="您是一个创意讲故事代理，能够根据想法创作引人入胜的故事并生成漫画面板。",
)

session = await client.sessions.create(agent_id=agent.id)

# 向代理发送消息
response = await client.sessions.chat(
    session_id=session.id,
    message="你好，能给我讲个故事吗？",
)

print(response)

# 🛠️ 为代理添加图像生成工具（DALL·E）
await client.agents.tools.create(
    agent_id=agent.id,
    name="image_generator",
    description="使用此工具根据描述生成图像。",
    integration={
        "provider": "dalle",
        "method": "generate_image",
        "setup": {
            "api_key": "your_dalle_api_key",
        },
    },
)

# 📋 任务
# 创建一个任务，接受一个想法并创建故事和 4 格漫画
task = await client.tasks.create(
    agent_id=agent.id,
    name="故事和漫画创作器",
    description="根据一个想法创作故事并生成 4 格漫画来说明故事。",
    main=[
        # 步骤 1：生成故事并将其概括为 4 个面板
        {
            "prompt": [
                {
                    "role": "system",
                    "content": "您是 {{agent.name}}。{{agent.about}}"
                },
                {
                    "role": "user",
                    "content": (
                        "基于想法 '{{_.idea}}'，写一个适合 4 格漫画的短故事。"
                        "提供故事和一个编号列表，包含 4 个简短描述，每个描述对应一个面板，说明故事中的关键时刻。"
                    ),
                },
            ],
            "unwrap": True,
        },
        # 步骤 2：提取面板描述和故事
        {
            "evaluate": {
                "story": "_.split('1. ')[0].strip()",
                "panels": "re.findall(r'\\d+\\.\\s*(.*?)(?=\\d+\\.\\s*|$)', _)",
            }
        },
        # 步骤 3：使用图像生成器工具为每个面板生成图像
        {
            "foreach": {
                "in": "_.panels",
                "do": {
                    "tool": "image_generator",
                    "arguments": {
                        "description": "_",
                    },
                },
            },
        },
        # 步骤 4：为故事生成一个吸引人的标题
        {
            "prompt": [
                {
                    "role": "system",
                    "content": "您是 {{agent.name}}。{{agent.about}}"
                },
                {
                    "role": "user",
                    "content": "根据以下故事，生成一个吸引人的标题。\n\n故事：{{outputs[1].story}}",
                },
            ],
            "unwrap": True,
        },
        # 步骤 5：返回故事、生成的图像和标题
        {
            "return": {
                "title": "outputs[3]",
                "story": "outputs[1].story",
                "comic_panels": "[output.image.url for output in outputs[2]]",
            }
        },
    ],
)

# 🚀 执行任务，输入一个想法
execution = await client.executions.create(
    task_id=task.id,
    input={"idea": "一只学会飞翔的猫"}
)

# 🎉 观看故事和漫画面板的生成过程
await client.executions.stream(execution_id=execution.id)

client.agents.tools.create(
    agent_id=agent.id,
    name="web_search",
    description="搜索网络以获取信息。",
    integration={
        "provider": "google",
        "method": "search",
        "setup": {"api_key": "your_google_api_key"},
    },
)

session = client.sessions.create(
    agent_id=agent.id,
    user_id="user123",
    context_overflow="adaptive"
)

# 在同一会话中继续对话
response = client.sessions.chat(
    session_id=session.id,
    message="继续我们之前的对话。"
)

# 上传文档
document = client.documents.create(
    file="path/to/document.pdf",
    metadata={"category": "research_paper"}
)

# 搜索文档
results = client.documents.search(
    query="AI 进展",
    filter={"category": "research_paper"}
)