44 lines
975 B
Go
44 lines
975 B
Go
package doubao
|
|
|
|
import (
|
|
"context"
|
|
"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
|
|
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
|
|
"time"
|
|
)
|
|
|
|
type DouBao struct {
|
|
Model string
|
|
Key string
|
|
}
|
|
|
|
func NewDouBao(modelName string, key string) *DouBao {
|
|
return &DouBao{
|
|
Model: modelName,
|
|
Key: key,
|
|
}
|
|
}
|
|
|
|
func (o *DouBao) GetData(ctx context.Context, respHandle func(input string) (string, error), Message []*model.ChatCompletionMessage) (string, error) {
|
|
|
|
client := arkruntime.NewClientWithApiKey(
|
|
o.Key,
|
|
//arkruntime.WithBaseUrl(UrlMap[url]),
|
|
arkruntime.WithRegion("cn-beijing"),
|
|
arkruntime.WithTimeout(2*time.Minute),
|
|
arkruntime.WithRetryTimes(2),
|
|
)
|
|
req := model.CreateChatCompletionRequest{
|
|
Model: o.Model,
|
|
Messages: Message,
|
|
}
|
|
|
|
resp, err := client.CreateChatCompletion(ctx, req)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
result, err := respHandle(*resp.Choices[0].Message.Content.StringValue)
|
|
|
|
return result, err
|
|
}
|