feat 存档

backend/doubao/ai_client.go

@@ -0,0 +1,54 @@
+package doubao
+
+import (
+	"fmt"
+)
+
+type Client struct {
+	enableSearch   bool    // 是否开启搜索
+	searchStrategy string  // 搜索策略（开了搜索才生效）
+	temperature    float32 // 采样温度
+}
+
+func NewAIClient(enableSearch bool, searchStrategy string, temperature float32) *Client {
+	return &Client{
+		enableSearch:   enableSearch,
+		searchStrategy: searchStrategy,
+		temperature:    temperature,
+	}
+}
+
+func (c *Client) RequestAI(model, key string, pre []string, prompt []string, cb func(content string, i int) error) ([]Message, int32, int32, error) {
+	inToken, outToken := int32(0), int32(0)
+	messages := make([]Message, 0)
+	for _, s := range pre {
+		messages = append(messages, Message{
+			Role:    "user",
+			Content: s,
+		})
+	}
+	for i, step := range prompt {
+		messages = append(messages, Message{
+			Role:    "user",
+			Content: step,
+		})
+		r, err := request(model, key, messages)
+		if err != nil {
+			return nil, 0, 0, err
+		}
+		inToken += r.Usage.InputTokens
+		outToken += r.Usage.OutputTokens
+		if len(r.Choices) == 0 {
+			fmt.Printf("AI response err: %+v\n", r)
+			return nil, 0, 0, err
+		} else {
+			messages = append(messages, r.Choices[0].Message)
+		}
+		fmt.Printf("AI 回答 %v: %v\n", i, r.Choices[0].Message.Content)
+		if err = cb(r.Choices[0].Message.Content, i); err != nil {
+			return nil, 0, 0, err
+		}
+	}
+	fmt.Printf("AI finished, use token: in-%v out-%v\n", inToken, outToken)
+	return messages, inToken, outToken, nil
+}

backend/doubao/ai_request.go

@@ -0,0 +1,67 @@
+package doubao
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+	"net/http"
+)
+
+type RequestBody struct {
+	Model               string    `json:"model"`
+	Messages            []Message `json:"messages"`
+	MaxCompletionTokens int32     `json:"max_completion_tokens"`
+}
+
+type Message struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
+type SearchOptions struct {
+	SearchStrategy string `json:"search_strategy"`
+}
+
+type ResponseBody struct {
+	Choices []struct {
+		FinishReason string  `json:"finish_reason"`
+		Message      Message `json:"message"`
+	} `json:"choices"`
+	Usage struct {
+		InputTokens  int32 `json:"prompt_tokens"`
+		OutputTokens int32 `json:"completion_tokens"`
+		TotalTokens  int32 `json:"total_tokens"`
+	} `json:"usage"`
+}
+
+func request(model, key string, messages []Message) (*ResponseBody, error) {
+	requestBody := RequestBody{
+		Model:               model,
+		Messages:            messages,
+		MaxCompletionTokens: 65535,
+	}
+	jsonData, _ := json.Marshal(requestBody)
+
+	req, err := http.NewRequest("POST", "https://ark.cn-beijing.volces.com/api/v3/chat/completions", bytes.NewBuffer(jsonData))
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("Authorization", "Bearer "+key)
+	req.Header.Set("Content-Type", "application/json")
+
+	resp, err := (&http.Client{}).Do(req)
+	if err != nil {
+		return nil, err
+	}
+	defer resp.Body.Close()
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+	var res ResponseBody
+	if err = json.Unmarshal(body, &res); err != nil {
+		return nil, err
+	}
+	return &res, nil
+}