68 lines
1.5 KiB
Go
68 lines
1.5 KiB
Go
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
|
|
}
|