feat: able to configure base url for openai kind of embedders/llms

Signed-off-by: bjwswang <[email protected]>
kubeagi · Nov 28, 2023 · 9de7083 · 9de7083
1 parent 35f464c
commit 9de7083
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 6 deletions.
diff --git a/controllers/embedder_controller.go b/controllers/embedder_controller.go
@@ -151,7 +151,7 @@ func (r *EmbedderReconciler) check3rdPartyEmbedder(ctx context.Context, logger l
 		}
 		msg = res.String()
 	case embeddings.OpenAI:
-		embedClient := openai.NewOpenAI(apiKey)
+		embedClient := openai.NewOpenAI(apiKey, instance.Spec.Enpoint.URL)
 		res, err := embedClient.Validate()
 		if err != nil {
 			return r.UpdateStatus(ctx, instance, nil, err)

diff --git a/controllers/llm_controller.go b/controllers/llm_controller.go
@@ -156,7 +156,7 @@ func (r *LLMReconciler) check3rdPartyLLM(ctx context.Context, logger logr.Logger
 		}
 		msg = res.String()
 	case llms.OpenAI:
-		embedClient := openai.NewOpenAI(apiKey)
+		embedClient := openai.NewOpenAI(apiKey, instance.Spec.Enpoint.URL)
 		res, err := embedClient.Validate()
 		if err != nil {
 			return r.UpdateStatus(ctx, instance, nil, err)

diff --git a/controllers/prompt_controller.go b/controllers/prompt_controller.go
@@ -130,7 +130,7 @@ func (r *PromptReconciler) CallLLM(ctx context.Context, logger logr.Logger, prom
 		llmClient = llmszhipuai.NewZhiPuAI(apiKey)
 		callData = prompt.Spec.ZhiPuAIParams.Marshal()
 	case llms.OpenAI:
-		llmClient = openai.NewOpenAI(apiKey)
+		llmClient = openai.NewOpenAI(apiKey, llm.Spec.Enpoint.URL)
 	default:
 		llmClient = llms.NewUnknowLLM()
 	}

diff --git a/pkg/llms/openai/api.go b/pkg/llms/openai/api.go
@@ -33,12 +33,17 @@ const (
 var _ llms.LLM = (*OpenAI)(nil)
 
 type OpenAI struct {
-	apiKey string
+	apiKey  string
+	baseURL string
 }
 
-func NewOpenAI(auth string) *OpenAI {
+func NewOpenAI(apiKey string, baseURL string) *OpenAI {
+	if baseURL == "" {
+		baseURL = OpenaiModelAPIURL
+	}
 	return &OpenAI{
-		apiKey: auth,
+		apiKey:  apiKey,
+		baseURL: baseURL,
 	}
 }