51 lines
1.1 KiB
Go
51 lines
1.1 KiB
Go
package third_party
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
|
|
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
|
|
)
|
|
|
|
type Hsyq struct {
|
|
mapClient map[string]*arkruntime.Client
|
|
}
|
|
|
|
func NewHsyq() *Hsyq {
|
|
return &Hsyq{
|
|
mapClient: make(map[string]*arkruntime.Client),
|
|
}
|
|
}
|
|
|
|
func (h *Hsyq) getClient(key string) *arkruntime.Client {
|
|
var client *arkruntime.Client
|
|
if _, ok := h.mapClient[key]; !ok {
|
|
client = h.mapClient[key]
|
|
} else {
|
|
client = arkruntime.NewClientWithApiKey(
|
|
key,
|
|
arkruntime.WithRegion("cn-beijing"),
|
|
arkruntime.WithTimeout(2*time.Minute),
|
|
arkruntime.WithRetryTimes(2),
|
|
)
|
|
h.mapClient[key] = client
|
|
}
|
|
return client
|
|
}
|
|
|
|
// 火山引擎
|
|
func (h *Hsyq) RequestHsyq(ctx context.Context, key string, modelName string, prompt []*model.ChatCompletionMessage) (model.ChatCompletionResponse, error) {
|
|
req := model.CreateChatCompletionRequest{
|
|
Model: modelName,
|
|
Messages: prompt,
|
|
}
|
|
|
|
resp, err := h.getClient(key).CreateChatCompletion(ctx, req)
|
|
if err != nil {
|
|
return model.ChatCompletionResponse{ID: ""}, err
|
|
}
|
|
|
|
return resp, err
|
|
}
|