sdk_generate/internal/call/callLLM.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
}