from openai import OpenAI

client = OpenAI()

# 创建助手
assistant = client.beta.assistants.create(
    name="数据分析助手",
    instructions="""你是一个专业的数据分析师。
    帮助用户分析数据、生成图表和报告。
    使用代码解释器处理数据文件。
    输出时使用中文，并提供清晰的分析结论。""",
    model="gpt-4o",
    tools=[
        {"type": "code_interpreter"},
        {"type": "file_search"},
    ],
)

print(f"助手 ID: {assistant.id}")

# 列出所有助手
assistants = client.beta.assistants.list(limit=10)
for a in assistants.data:
    print(f"  {a.id}: {a.name}")

# 更新助手
client.beta.assistants.update(
    assistant_id=assistant.id,
    instructions="更新后的指令...",
    model="gpt-4o",
)

# 删除助手
client.beta.assistants.delete(assistant_id=assistant.id)

# 创建线程
thread = client.beta.threads.create()
print(f"线程 ID: {thread.id}")

# 添加消息
message = client.beta.threads.messages.create(
    thread_id=thread.id,
    role="user",
    content="请分析这份销售数据，找出趋势和异常。",
)

# 上传文件
file = client.files.create(
    file=open("sales_data.csv", "rb"),
    purpose="assistants",
)

# 带文件的消息
message = client.beta.threads.messages.create(
    thread_id=thread.id,
    role="user",
    content="分析这份销售数据",
    attachments=[{
        "file_id": file.id,
        "tools": [{"type": "file_search"}, {"type": "code_interpreter"}]
    }]
)

import time

# 创建运行
run = client.beta.threads.runs.create(
    thread_id=thread.id,
    assistant_id=assistant.id,
)

# 轮询等待完成
while run.status in ["queued", "in_progress", "cancelling"]:
    time.sleep(1)
    run = client.beta.threads.runs.retrieve(
        thread_id=thread.id,
        run_id=run.id,
    )
    print(f"状态: {run.status}")

if run.status == "completed":
    # 获取回复
    messages = client.beta.threads.messages.list(
        thread_id=thread.id,
        order="desc",
        limit=1,
    )
    for msg in messages.data:
        print(msg.content[0].text.value)
elif run.status == "failed":
    print(f"运行失败: {run.last_error}")

# 一步完成：创建运行并等待结果
run = client.beta.threads.runs.create_and_poll(
    thread_id=thread.id,
    assistant_id=assistant.id,
)

if run.status == "completed":
    messages = client.beta.threads.messages.list(thread_id=thread.id)
    for msg in messages.data:
        if msg.role == "assistant":
            print(msg.content[0].text.value)

with client.beta.threads.runs.stream(
    thread_id=thread.id,
    assistant_id=assistant.id,
) as stream:
    for event in stream:
        if event.event == "thread.message.delta":
            delta = event.data.delta
            if delta.content:
                for block in delta.content:
                    if hasattr(block, 'text') and block.text:
                        print(block.text.value, end="", flush=True)
        elif event.event == "thread.run.failed":
            print(f"\n运行失败: {event.data.last_error}")

print()

# 上传文件用于代码解释器
file = client.files.create(
    file=open("data.csv", "rb"),
    purpose="assistants",
)

# 创建带代码解释器的助手
analyst = client.beta.assistants.create(
    name="数据分析专家",
    instructions="用 Python 分析数据并生成可视化图表。",
    model="gpt-4o",
    tools=[{"type": "code_interpreter"}],
    tool_resources={
        "code_interpreter": {
            "file_ids": [file.id]
        }
    }
)

# 运行分析
thread = client.beta.threads.create()
client.beta.threads.messages.create(
    thread_id=thread.id,
    role="user",
    content="请分析数据中的趋势，生成折线图和柱状图。",
)

run = client.beta.threads.runs.create_and_poll(
    thread_id=thread.id,
    assistant_id=analyst.id,
)

# 获取生成的文件（图表等）
if run.status == "completed":
    messages = client.beta.threads.messages.list(thread_id=thread.id)
    for msg in messages.data:
        for content in msg.content:
            if content.type == "image":
                print(f"生成图表: {content.image.file_id}")
            elif content.type == "text":
                print(content.text.value)

# 创建向量存储
vector_store = client.beta.vector_stores.create(
    name="产品文档",
)

# 上传文件到向量存储
files = ["manual.pdf", "faq.txt", "specs.docx"]
file_streams = [open(f, "rb") for f in files]

file_batch = client.beta.vector_stores.file_batches.upload_and_poll(
    vector_store_id=vector_store.id,
    files=file_streams,
)

print(f"文件状态: {file_batch.status}")
print(f"文件数量: {file_batch.file_counts}")

# 创建带检索功能的助手
doc_assistant = client.beta.assistants.create(
    name="文档问答助手",
    instructions="基于上传的文档回答用户问题，引用具体来源。",
    model="gpt-4o",
    tools=[{"type": "file_search"}],
    tool_resources={
        "file_search": {
            "vector_store_ids": [vector_store.id]
        }
    }
)

run = client.beta.threads.runs.create(
    thread_id=thread.id,
    assistant_id=doc_assistant.id,
    tool_choice={"type": "file_search"},
)

import json

# 定义工具
tools = [
    {
        "type": "function",
        "function": {
            "name": "get_user_info",
            "description": "获取用户信息",
            "parameters": {
                "type": "object",
                "properties": {
                    "user_id": {"type": "string", "description": "用户ID"}
                },
                "required": ["user_id"]
            }
        }
    }
]

assistant = client.beta.assistants.create(
    name="客服助手",
    instructions="帮助用户查询订单和账户信息。",
    model="gpt-4o",
    tools=tools,
)

# 运行并处理工具调用
def run_with_tools(thread_id, assistant_id):
    run = client.beta.threads.runs.create_and_poll(
        thread_id=thread_id,
        assistant_id=assistant_id,
    )

    # 处理工具调用
    while run.status == "requires_action":
        tool_outputs = []
        for call in run.required_action.submit_tool_outputs.tool_calls:
            if call.function.name == "get_user_info":
                args = json.loads(call.function.arguments)
                result = json.dumps({"name": "张三", "tier": "VIP"})
                tool_outputs.append({
                    "tool_call_id": call.id,
                    "output": result,
                })

        run = client.beta.threads.runs.submit_tool_outputs_and_poll(
            thread_id=thread_id,
            run_id=run.id,
            tool_outputs=tool_outputs,
        )

    return run

class AssistantApp:
    """基于 Assistants API 的对话应用"""

    def __init__(self, assistant_id: str):
        self.client = OpenAI()
        self.assistant_id = assistant_id
        self.thread = self.client.beta.threads.create()

    def chat(self, user_input: str) -> str:
        """发送消息并获取回复"""
        # 添加用户消息
        self.client.beta.threads.messages.create(
            thread_id=self.thread.id,
            role="user",
            content=user_input,
        )

        # 运行
        run = self.client.beta.threads.runs.create_and_poll(
            thread_id=self.thread.id,
            assistant_id=self.assistant_id,
        )

        if run.status == "completed":
            messages = self.client.beta.threads.messages.list(
                thread_id=self.thread.id,
                order="desc",
                limit=1,
            )
            return messages.data[0].content[0].text.value
        else:
            return f"运行状态: {run.status}"

    def get_history(self) -> list[dict]:
        """获取对话历史"""
        messages = self.client.beta.threads.messages.list(
            thread_id=self.thread.id,
            order="asc",
        )
        return [
            {"role": msg.role, "content": msg.content[0].text.value}
            for msg in messages.data
            if msg.content[0].type == "text"
        ]


# 使用
app = AssistantApp(assistant_id="asst_xxx")
print(app.chat("你好，请帮我分析数据"))
print(app.chat("按月份统计趋势"))