import os

import cozeloop
from cozeloop.spec import tracespec
from openai import OpenAI
from openai.types.chat import ChatCompletionSystemMessageParam

def main():
    # 初始化cozeloop client
    client = cozeloop.new_client(prompt_trace=True)

    # 准备用户输入
    input = {"people_num": "2", "destination": "上海", "days_num": "5", "travel_theme": "自由行"}

    # 创建root span
    root = cozeloop.start_span("旅行专家", "chain")
    root.set_input(input)

    try:
        # 拉取Prompt
        prompt = client.get_prompt(prompt_key="CozeLoop_Travel_Master")
        cozeloop_messages = client.prompt_format(prompt, input)
        messages = [
            ChatCompletionSystemMessageParam(
                role="system",
                content=cozeloop_messages[0].content
            )
        ]

        # 调用大模型
        response = call_llm(
            messages=messages,
            prompt=prompt,
        )

        # 处理最终结果
        final_output = response.choices[0].message.content
        root.set_output(final_output)

    finally:
        # root span 上报
        root.finish()

        # 服务退出前调用Close方法，否则可能造成trace数据丢失
        client.close()

def call_llm(messages, prompt):
    # 创建model span
    model_span = cozeloop.start_span("CallLLM", "model")

    ark_client = OpenAI(
        base_url="https://ark.cn-beijing.volces.com/api/v3",
        api_key=os.environ.get("OPENAI_API_KEY")
    )
    temperature = 0.7
    max_tokens = 1000

    try:
        # 设置模型参数
        model_span.set_input({"messages": messages})
        model_span.set_model_provider("volcengine")
        model_span.set_model_name("doubao-1-5-vision-pro-32k-250115")
        model_span.set_model_call_options(tracespec.ModelCallOption(
            temperature=temperature,
            max_tokens=max_tokens
        ))
        model_span.set_prompt(prompt=prompt)

        # 执行模型调用
        response = ark_client.chat.completions.create(
            model="doubao-1-5-vision-pro-32k-250115",
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens
        )

        model_span.set_output(response)
        model_span.set_input_tokens(response.usage.prompt_tokens)
        model_span.set_output_tokens(response.usage.completion_tokens)

        return response
    except Exception as e:
        model_span.set_error(e)
        raise
    finally:
        model_span.finish()

if __name__ == '__main__':
    main()