From 6d2292addb6d6c5c9fb3fe7b54154db70d4e3cab Mon Sep 17 00:00:00 2001 From: Lanture1064 Date: Wed, 23 Aug 2023 16:23:07 +0800 Subject: [PATCH] refactor: Extract LLM validate logic Signed-off-by: Lanture1064 --- controllers/llm_controller.go | 87 ++++++++++++----------------------- pkg/llms/llms.go | 4 ++ pkg/llms/openai/api.go | 72 +++++++++++++++++++++++++++++ pkg/llms/zhipuai/api.go | 31 +++++++++++++ 4 files changed, 137 insertions(+), 57 deletions(-) create mode 100644 pkg/llms/openai/api.go diff --git a/controllers/llm_controller.go b/controllers/llm_controller.go index 3b71f2817..ef8f8fd80 100644 --- a/controllers/llm_controller.go +++ b/controllers/llm_controller.go @@ -19,17 +19,20 @@ package controllers import ( "context" "fmt" - "net/http" - "github.com/go-logr/logr" + "github.com/kubeagi/arcadia/pkg/llms" + "github.com/kubeagi/arcadia/pkg/llms/openai" + "github.com/kubeagi/arcadia/pkg/llms/zhipuai" corev1 "k8s.io/api/core/v1" "k8s.io/apimachinery/pkg/api/errors" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" "k8s.io/apimachinery/pkg/runtime" "k8s.io/apimachinery/pkg/types" ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/builder" "sigs.k8s.io/controller-runtime/pkg/client" "sigs.k8s.io/controller-runtime/pkg/log" + "sigs.k8s.io/controller-runtime/pkg/predicate" "sigs.k8s.io/controller-runtime/pkg/reconcile" arcadiav1alpha1 "github.com/kubeagi/arcadia/api/v1alpha1" @@ -81,7 +84,7 @@ func (r *LLMReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.R // SetupWithManager sets up the controller with the Manager. func (r *LLMReconciler) SetupWithManager(mgr ctrl.Manager) error { return ctrl.NewControllerManagedBy(mgr). - For(&arcadiav1alpha1.LLM{}). + For(&arcadiav1alpha1.LLM{}, builder.WithPredicates(LLMPredicates{})). Complete(r) } @@ -89,7 +92,26 @@ func (r *LLMReconciler) SetupWithManager(mgr ctrl.Manager) error { func (r *LLMReconciler) CheckLLM(ctx context.Context, logger logr.Logger, instance *arcadiav1alpha1.LLM) error { logger.Info("Checking LLM instance") // Check new URL/Auth availability - err := r.TestLLMAvailability(ctx, instance, logger) + var err error + var validator llms.Validator + + secret := &corev1.Secret{} + err = r.Get(ctx, types.NamespacedName{Name: instance.Spec.Auth, Namespace: instance.Namespace}, secret) + if err != nil { + return err + } + apiKey := string(secret.Data["apiKey"]) + + switch instance.Spec.Type { + case llms.OpenAI: + validator = openai.NewOpenAI(apiKey) + case llms.ZhiPuAI: + validator = zhipuai.NewZhiPuAI(apiKey) + default: + return fmt.Errorf("unknown LLM type: %s", instance.Spec.Type) + } + + res, err := validator.Validate() if err != nil { // Set status to unavailable instance.Status.SetConditions(arcadiav1alpha1.Condition{ @@ -105,64 +127,15 @@ func (r *LLMReconciler) CheckLLM(ctx context.Context, logger logr.Logger, instan Type: arcadiav1alpha1.TypeReady, Status: corev1.ConditionTrue, Reason: arcadiav1alpha1.ReasonAvailable, - Message: "Available", + Message: res, LastTransitionTime: metav1.Now(), LastSuccessfulTime: metav1.Now(), }) } - return r.Client.Status().Update(ctx, instance) -} - -// TestLLMAvailability tests LLM availability. -func (r *LLMReconciler) TestLLMAvailability(ctx context.Context, instance *arcadiav1alpha1.LLM, logger logr.Logger) error { - logger.Info("Testing LLM availability") - - //TODO: change URL & request for different types of LLM instance - // For openai instance, we use the "GET model" api. - // For Zhipuai instance, we send a standard async request. - testURL := instance.Spec.URL + "/v1/models" - - if instance.Spec.Auth == "" { - return fmt.Errorf("auth is empty") - } - - // get auth by secret name - var auth string - secret := &corev1.Secret{} - err := r.Get(ctx, types.NamespacedName{Name: instance.Spec.Auth, Namespace: instance.Namespace}, secret) - if err != nil { - return err - } - - auth = "Bearer " + string(secret.Data["apiKey"]) - err = SendTestRequest("GET", testURL, auth) - if err != nil { - return err - } - - return nil + return r.Client.Status().Update(ctx, instance) } -func SendTestRequest(method string, url string, auth string) error { - req, err := http.NewRequest(method, url, nil) - if err != nil { - return err - } - - req.Header.Set("Authorization", auth) - req.Header.Set("Content-Type", "application/json") - - cli := &http.Client{} - resp, err := cli.Do(req) - if err != nil { - return err - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - return fmt.Errorf("returns unexpected status code: %d", resp.StatusCode) - } - - return nil +type LLMPredicates struct { + predicate.Funcs } diff --git a/pkg/llms/llms.go b/pkg/llms/llms.go index cc3f891c2..533639d7c 100644 --- a/pkg/llms/llms.go +++ b/pkg/llms/llms.go @@ -28,3 +28,7 @@ type Response interface { String() string Bytes() []byte } + +type Validator interface { + Validate() (string, error) +} diff --git a/pkg/llms/openai/api.go b/pkg/llms/openai/api.go new file mode 100644 index 000000000..d2db8115a --- /dev/null +++ b/pkg/llms/openai/api.go @@ -0,0 +1,72 @@ +/* +Copyright 2023 KubeAGI. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package openai + +import ( + "fmt" + "net/http" + "time" +) + +const ( + OPENAI_MODEL_API_URL = "https://api.openai.com/v1" + OPENAI_MODEL_Default_Timeout = 300 * time.Second +) + +type OpenAI struct { + apiKey string +} + +func NewOpenAI(auth string) *OpenAI { + return &OpenAI{ + apiKey: auth, + } +} + +func (o *OpenAI) Validate() (string, error) { + // Validate OpenAI type CRD LLM Instance + // instance.Spec.URL should be like "https://api.openai.com/" + + if o.apiKey == "" { + // TODO: maybe we should consider local pseudo-openAI LLM worker that doesn't require an apiKey? + return "", fmt.Errorf("auth is empty") + } + + testURL := OPENAI_MODEL_API_URL + "/models" + testAuth := "Bearer " + o.apiKey // openAI official requirement + + req, err := http.NewRequest("GET", testURL, nil) + if err != nil { + return "", err + } + + req.Header.Set("Authorization", testAuth) + req.Header.Set("Content-Type", "application/json") + + cli := &http.Client{} + resp, err := cli.Do(req) + if err != nil { + return "", err + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + return "", fmt.Errorf("returns unexpected status code: %d", resp.StatusCode) + } + + return "GET model request returns Status OK", nil +} diff --git a/pkg/llms/zhipuai/api.go b/pkg/llms/zhipuai/api.go index 307c541cd..2256986fa 100644 --- a/pkg/llms/zhipuai/api.go +++ b/pkg/llms/zhipuai/api.go @@ -124,3 +124,34 @@ func (z *ZhiPuAI) SSEInvoke(params ModelParams, handler func(*sse.Event)) error } return Stream(url, token, params, ZHIPUAI_MODEL_Default_Timeout, nil) } + +func (z *ZhiPuAI) Validate() (string, error) { + url := BuildAPIURL(ZhiPuAILite, ZhiPuAIAsyncInvoke) + token, err := GenerateToken(z.apiKey, API_TOKEN_TTL_SECONDS) + if err != nil { + return "", err + } + + testPrompt := []Prompt{ + { + Role: "user", + Content: "Hello!", + }, + } + + testParam := ModelParams{ + Method: ZhiPuAIAsyncInvoke, + Model: ZhiPuAILite, + Temperature: 0.95, + TopP: 0.7, + Prompt: testPrompt, + } + + response, err := Post(url, token, testParam, ZHIPUAI_MODEL_Default_Timeout) + if err != nil { + return "", err + } + responseMsg := "TaskID:" + response.Data.TaskID + + return responseMsg, nil +}