package main

import (
    "context"
    "os"

    "github.com/coze-dev/cozeloop-go"
    "github.com/coze-dev/cozeloop-go/entity"
    "github.com/coze-dev/cozeloop-go/spec/tracespec"
    "github.com/volcengine/volcengine-go-sdk/service/arkruntime"
    "github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
    "github.com/volcengine/volcengine-go-sdk/volcengine"
)

var (
    // 初始化cozeloop client
    client, _ = cozeloop.NewClient(cozeloop.WithPromptTrace(true))
)

func main() {
    ctx := context.Background()
    // 准备用户输入
    input := map[string]any{"people_num": "2", "destination": "上海", "days_num": "5", "travel_theme": "自由行"}

    // 创建root span
    ctx, root := client.StartSpan(ctx, "旅行专家", "chain")
    root.SetInput(ctx, input)

    // 拉取Prompt
    prompt, err := client.GetPrompt(ctx, cozeloop.GetPromptParam{PromptKey: "CozeLoop_Travel_Master"})
    if err != nil {
       panic(err)
    }
    messages, err := client.PromptFormat(ctx, prompt, input)
    if err != nil {
       panic(err)
    }

    // 调用模型
    resp, err := callLLM(ctx, messages, prompt)
    if err != nil {
       panic(err)
    }

    root.SetOutput(ctx, *resp.Choices[0].Message.Content.StringValue)
    // root span上报
    root.Finish(ctx)

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

func callLLM(ctx context.Context, messages []*entity.Message, prompt *entity.Prompt) (*model.ChatCompletionResponse, error) {
    // 创建model span
    ctx, span := client.StartSpan(ctx, "CallLLM", "model")
    defer span.Finish(ctx)

    temperature := float32(0.7)
    maxTokens := 1000
    arkClient := arkruntime.NewClientWithApiKey(
       os.Getenv("ARK_API_KEY"),
    )
    req := model.CreateChatCompletionRequest{
       Model: "doubao-1-5-vision-pro-32k-250115",
       Messages: []*model.ChatCompletionMessage{
          {
             Role: model.ChatMessageRoleSystem,
             Content: &model.ChatCompletionMessageContent{
                StringValue: volcengine.String(*messages[0].Content),
             },
          },
       },
       Temperature: &temperature,
       MaxTokens: &maxTokens,
    }
    span.SetInput(ctx, map[string]interface{}{"messages": req.Messages})
    span.SetModelProvider(ctx, "volcengine")
    span.SetModelName(ctx, req.Model)
    span.SetModelCallOptions(ctx, tracespec.ModelCallOption{
       Temperature: temperature,  
       MaxTokens: int64(maxTokens),
    })
    span.SetPrompt(ctx, *prompt)

    // 调用模型
    resp, err := arkClient.CreateChatCompletion(ctx, req)
    if err != nil {
       span.SetError(ctx, err)
       return nil, err
    }

    span.SetOutput(ctx, resp)
    span.SetInputTokens(ctx, resp.Usage.PromptTokens)
    span.SetOutputTokens(ctx, resp.Usage.CompletionTokens)

    return &resp, nil
}