|
| 1 | +// Copyright (c) 2025 Bytedance Ltd. and/or its affiliates |
| 2 | +// SPDX-License-Identifier: MIT |
| 3 | + |
| 4 | +package main |
| 5 | + |
| 6 | +import ( |
| 7 | + "context" |
| 8 | + "encoding/json" |
| 9 | + "fmt" |
| 10 | + "net/http" |
| 11 | + "time" |
| 12 | + |
| 13 | + "github.com/coze-dev/cozeloop-go" |
| 14 | + "github.com/coze-dev/cozeloop-go/entity" |
| 15 | + "github.com/coze-dev/cozeloop-go/internal/util" |
| 16 | + "github.com/coze-dev/cozeloop-go/spec/tracespec" |
| 17 | +) |
| 18 | + |
| 19 | +// If you want to use the jinja templates in prompts, you can refer to the following. |
| 20 | +func main() { |
| 21 | + // 1.Create a prompt on the platform |
| 22 | + // You can create a Prompt on the platform's Prompt development page (set Prompt Key to 'prompt_hub_demo'), add the following messages to the template, and submit a version. |
| 23 | + // System: You are a helpful bot, the conversation topic is {{var1}}. |
| 24 | + // Placeholder: placeholder1 |
| 25 | + // User: My question is {{var2}} |
| 26 | + // Placeholder: placeholder2 |
| 27 | + |
| 28 | + ctx := context.Background() |
| 29 | + |
| 30 | + // Set the following environment variables first. |
| 31 | + // COZELOOP_WORKSPACE_ID=your workspace id |
| 32 | + // COZELOOP_API_TOKEN=your token |
| 33 | + // 2.New loop client |
| 34 | + client, err := cozeloop.NewClient( |
| 35 | + // Set whether to report a trace span when get or format prompt. |
| 36 | + // Default value is false. |
| 37 | + cozeloop.WithPromptTrace(true)) |
| 38 | + if err != nil { |
| 39 | + panic(err) |
| 40 | + } |
| 41 | + |
| 42 | + llmRunner := llmRunner{ |
| 43 | + client: client, |
| 44 | + } |
| 45 | + |
| 46 | + // 1. start root span |
| 47 | + ctx, span := llmRunner.client.StartSpan(ctx, "root_span", "main_span", nil) |
| 48 | + |
| 49 | + // 3.Get the prompt |
| 50 | + prompt, err := llmRunner.client.GetPrompt(ctx, cozeloop.GetPromptParam{ |
| 51 | + PromptKey: "prompt_hub_demo", |
| 52 | + // If version is not specified, the latest version of the corresponding prompt will be obtained |
| 53 | + Version: "0.0.1", |
| 54 | + }) |
| 55 | + if err != nil { |
| 56 | + fmt.Printf("get prompt failed: %v\n", err) |
| 57 | + return |
| 58 | + } |
| 59 | + if prompt != nil { |
| 60 | + // Get messages of the prompt |
| 61 | + if prompt.PromptTemplate != nil { |
| 62 | + messages, err := json.Marshal(prompt.PromptTemplate.Messages) |
| 63 | + if err != nil { |
| 64 | + fmt.Printf("json marshal failed: %v\n", err) |
| 65 | + return |
| 66 | + } |
| 67 | + fmt.Printf("prompt messages=%s\n", string(messages)) |
| 68 | + } |
| 69 | + // Get llm config of the prompt |
| 70 | + if prompt.LLMConfig != nil { |
| 71 | + llmConfig, err := json.Marshal(prompt.LLMConfig) |
| 72 | + if err != nil { |
| 73 | + fmt.Printf("json marshal failed: %v\n", err) |
| 74 | + } |
| 75 | + fmt.Printf("prompt llm config=%s\n", llmConfig) |
| 76 | + } |
| 77 | + |
| 78 | + // 4.Format messages of the prompt |
| 79 | + userMessageContent := "Hello!" |
| 80 | + assistantMessageContent := "Hello!" |
| 81 | + |
| 82 | + messages, err := llmRunner.client.PromptFormat(ctx, prompt, map[string]any{ |
| 83 | + "var_string": "hi", |
| 84 | + "var_int": 5, |
| 85 | + "var_bool": true, |
| 86 | + "var_float": 1.0, |
| 87 | + "var_object": map[string]interface{}{ |
| 88 | + "name": "John", |
| 89 | + "age": 30, |
| 90 | + "hobbies": []string{"reading", "coding"}, |
| 91 | + "address": map[string]interface{}{ |
| 92 | + "city": "bejing", |
| 93 | + "street": "123 Main", |
| 94 | + }, |
| 95 | + }, |
| 96 | + "var_array_string": []string{ |
| 97 | + "hello", "nihao", |
| 98 | + }, |
| 99 | + "var_array_boolean": []bool{ |
| 100 | + true, false, true, |
| 101 | + }, |
| 102 | + "var_array_int": []int64{ |
| 103 | + 1, 2, 3, 4, |
| 104 | + }, |
| 105 | + "var_array_float": []float64{ |
| 106 | + 1.0, 2.0, |
| 107 | + }, |
| 108 | + "var_array_object": []interface{}{ |
| 109 | + map[string]interface{}{"key": "123"}, |
| 110 | + map[string]interface{}{"value": 100}, |
| 111 | + }, |
| 112 | + // Placeholder variable type should be entity.Message/*entity.Message/[]entity.Message/[]*entity.Message |
| 113 | + "placeholder1": []*entity.Message{ |
| 114 | + { |
| 115 | + Role: entity.RoleUser, |
| 116 | + Content: &userMessageContent, |
| 117 | + }, |
| 118 | + { |
| 119 | + Role: entity.RoleAssistant, |
| 120 | + Content: &assistantMessageContent, |
| 121 | + }, |
| 122 | + }, |
| 123 | + // Other variables in the prompt template that are not provided with corresponding values will be considered as empty values |
| 124 | + }) |
| 125 | + if err != nil { |
| 126 | + fmt.Printf("prompt format failed: %v\n", err) |
| 127 | + return |
| 128 | + } |
| 129 | + data, err := json.Marshal(messages) |
| 130 | + if err != nil { |
| 131 | + fmt.Printf("json marshal failed: %v\n", err) |
| 132 | + return |
| 133 | + } |
| 134 | + fmt.Printf("formatted messages=%s\n", string(data)) |
| 135 | + |
| 136 | + // 5. llm call |
| 137 | + err = llmRunner.llmCall(ctx, messages) |
| 138 | + if err != nil { |
| 139 | + return |
| 140 | + } |
| 141 | + } |
| 142 | + |
| 143 | + // 6. span finish |
| 144 | + span.Finish(ctx) |
| 145 | + |
| 146 | + // 7. (optional) flush or close |
| 147 | + // -- force flush, report all traces in the queue |
| 148 | + // Warning! In general, this method is not needed to be call, as spans will be automatically reported in batches. |
| 149 | + // Note that flush will block and wait for the report to complete, and it may cause frequent reporting, |
| 150 | + // affecting performance. |
| 151 | + llmRunner.client.Flush(ctx) |
| 152 | +} |
| 153 | + |
| 154 | +type llmRunner struct { |
| 155 | + client cozeloop.Client |
| 156 | +} |
| 157 | + |
| 158 | +func (r *llmRunner) llmCall(ctx context.Context, messages []*entity.Message) (err error) { |
| 159 | + ctx, span := r.client.StartSpan(ctx, "llmCall", tracespec.VModelSpanType, nil) |
| 160 | + defer span.Finish(ctx) |
| 161 | + |
| 162 | + // llm is processing |
| 163 | + //baseURL := "https://xxx" |
| 164 | + //ak := "****" |
| 165 | + modelName := "gpt-4o-2024-05-13" |
| 166 | + maxTokens := 1000 // range: [0, 4096] |
| 167 | + //transport := &MyTransport{ |
| 168 | + // DefaultTransport: &http.Transport{}, |
| 169 | + //} |
| 170 | + //config := openai.DefaultAzureConfig(ak, baseURL) |
| 171 | + //config.HTTPClient = &http.Client{ |
| 172 | + // Transport: transport, |
| 173 | + //} |
| 174 | + //client := openai.NewClientWithConfig(config) |
| 175 | + //marshal, err := json.Marshal(messages) |
| 176 | + //if err != nil { |
| 177 | + // return err |
| 178 | + //} |
| 179 | + //chatCompletionMessage := make([]openai.ChatCompletionMessage, 0) |
| 180 | + //err = json.Unmarshal(marshal, &chatCompletionMessage) |
| 181 | + //if err != nil { |
| 182 | + // return err |
| 183 | + //} |
| 184 | + //resp, err := client.CreateChatCompletion( |
| 185 | + // ctx, |
| 186 | + // openai.ChatCompletionRequest{ |
| 187 | + // Model: modelName, |
| 188 | + // Messages: chatCompletionMessage, |
| 189 | + // MaxTokens: maxTokens, |
| 190 | + // TopP: 0.95, |
| 191 | + // N: 1, |
| 192 | + // PresencePenalty: 1.0, |
| 193 | + // FrequencyPenalty: 1.0, |
| 194 | + // Temperature: 0.6, |
| 195 | + // }, |
| 196 | + //) |
| 197 | + |
| 198 | + // mock resp |
| 199 | + time.Sleep(1 * time.Second) |
| 200 | + respChoices := []string{ |
| 201 | + "Hello! Can I help you?", |
| 202 | + } |
| 203 | + respPromptTokens := 11 |
| 204 | + respCompletionTokens := 52 |
| 205 | + |
| 206 | + // set tag key: `input` |
| 207 | + span.SetInput(ctx, convertModelInput(messages)) |
| 208 | + // set tag key: `output` |
| 209 | + span.SetOutput(ctx, respChoices) |
| 210 | + // set tag key: `model_provider`, e.g., openai, etc. |
| 211 | + span.SetModelProvider(ctx, "openai") |
| 212 | + // set tag key: `start_time_first_resp` |
| 213 | + // Timestamp of the first packet return from LLM, unit: microseconds. |
| 214 | + // When `start_time_first_resp` is set, a tag named `latency_first_resp` calculated |
| 215 | + // based on the span's StartTime will be added, meaning the latency for the first packet. |
| 216 | + span.SetStartTimeFirstResp(ctx, time.Now().UnixMicro()) |
| 217 | + // set tag key: `input_tokens`. The amount of input tokens. |
| 218 | + // when the `input_tokens` value is set, it will automatically sum with the `output_tokens` to calculate the `tokens` tag. |
| 219 | + span.SetInputTokens(ctx, respPromptTokens) |
| 220 | + // set tag key: `output_tokens`. The amount of output tokens. |
| 221 | + // when the `output_tokens` value is set, it will automatically sum with the `input_tokens` to calculate the `tokens` tag. |
| 222 | + span.SetOutputTokens(ctx, respCompletionTokens) |
| 223 | + // set tag key: `model_name`, e.g., gpt-4-1106-preview, etc. |
| 224 | + span.SetModelName(ctx, modelName) |
| 225 | + span.SetTags(ctx, map[string]interface{}{ |
| 226 | + tracespec.CallOptions: tracespec.ModelCallOption{ |
| 227 | + Temperature: 0.6, |
| 228 | + MaxTokens: int64(maxTokens), |
| 229 | + TopP: 0.95, |
| 230 | + N: 1, |
| 231 | + PresencePenalty: util.Ptr(float32(1.0)), |
| 232 | + FrequencyPenalty: util.Ptr(float32(1.0)), |
| 233 | + }, |
| 234 | + }) |
| 235 | + |
| 236 | + return nil |
| 237 | +} |
| 238 | + |
| 239 | +type MyTransport struct { |
| 240 | + Header http.Header |
| 241 | + DefaultTransport http.RoundTripper |
| 242 | +} |
| 243 | + |
| 244 | +func (transport *MyTransport) RoundTrip(req *http.Request) (*http.Response, error) { |
| 245 | + for key, values := range transport.Header { |
| 246 | + for _, value := range values { |
| 247 | + req.Header.Add(key, value) |
| 248 | + } |
| 249 | + } |
| 250 | + return transport.DefaultTransport.RoundTrip(req) |
| 251 | +} |
| 252 | + |
| 253 | +func convertModelInput(messages []*entity.Message) *tracespec.ModelInput { |
| 254 | + modelMessages := make([]*tracespec.ModelMessage, 0) |
| 255 | + for _, message := range messages { |
| 256 | + modelMessages = append(modelMessages, &tracespec.ModelMessage{ |
| 257 | + Role: string(message.Role), |
| 258 | + Content: util.PtrValue(message.Content), |
| 259 | + }) |
| 260 | + } |
| 261 | + |
| 262 | + return &tracespec.ModelInput{ |
| 263 | + Messages: modelMessages, |
| 264 | + } |
| 265 | +} |
0 commit comments