51 lines
1.1 KiB
Go
51 lines
1.1 KiB
Go
package call
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
|
|
"github.com/sashabaranov/go-openai"
|
|
)
|
|
|
|
type LlmCallSet struct {
|
|
ApiKey string
|
|
BaseUrl string
|
|
ModelName string
|
|
}
|
|
|
|
type CallLLM struct {
|
|
ModelName string
|
|
openaiClient *openai.Client
|
|
}
|
|
|
|
func NewCallLLM(LlmCallSet *LlmCallSet) *CallLLM {
|
|
clientConfig := openai.DefaultConfig(LlmCallSet.ApiKey)
|
|
clientConfig.BaseURL = LlmCallSet.BaseUrl
|
|
client := openai.NewClientWithConfig(clientConfig)
|
|
return &CallLLM{
|
|
openaiClient: client,
|
|
ModelName: LlmCallSet.ModelName,
|
|
}
|
|
}
|
|
|
|
func (s *CallLLM) Do(ctx context.Context, RequestOption openai.ChatCompletionRequest, modelName ...string) (string, error) {
|
|
if modelName == nil {
|
|
modelName = append(modelName, s.ModelName)
|
|
}
|
|
if modelName[0] == "" {
|
|
return "", fmt.Errorf("model name is nil")
|
|
}
|
|
RequestOption.Model = modelName[0]
|
|
resp, err := s.openaiClient.CreateChatCompletion(ctx, RequestOption)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
|
|
if len(resp.Choices) == 0 {
|
|
return "", fmt.Errorf("模型未返回任何内容")
|
|
}
|
|
content := resp.Choices[0].Message.Content
|
|
//content := test.Res
|
|
return content, nil
|
|
}
|