diff --git a/PROJECT b/PROJECT index 5eca591fa..9c5b0a24f 100644 --- a/PROJECT +++ b/PROJECT @@ -49,4 +49,13 @@ resources: kind: Datasource path: github.com/kubeagi/arcadia/api/v1alpha1 version: v1alpha1 +- api: + crdVersion: v1 + namespaced: true + controller: true + domain: kubeagi.k8s.com.cn + group: arcadia + kind: Embeddings + path: github.com/kubeagi/arcadia/api/v1alpha1 + version: v1alpha1 version: "3" diff --git a/api/v1alpha1/embedder.go b/api/v1alpha1/embedder.go new file mode 100644 index 000000000..151308861 --- /dev/null +++ b/api/v1alpha1/embedder.go @@ -0,0 +1,36 @@ +/* +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 v1alpha1 + +import ( + "context" + corev1 "k8s.io/api/core/v1" + "k8s.io/apimachinery/pkg/types" + "sigs.k8s.io/controller-runtime/pkg/client" +) + +func (e Embedder) AuthAPIKey(ctx context.Context, c client.Client) (string, error) { + if e.Spec.Auth == "" { + return "", nil + } + authSecret := &corev1.Secret{} + err := c.Get(ctx, types.NamespacedName{Name: e.Spec.Auth, Namespace: e.Namespace}, authSecret) + if err != nil { + return "", err + } + return string(authSecret.Data["apiKey"]), nil +} diff --git a/api/v1alpha1/embedder_types.go b/api/v1alpha1/embedder_types.go new file mode 100644 index 000000000..28f97c701 --- /dev/null +++ b/api/v1alpha1/embedder_types.go @@ -0,0 +1,77 @@ +/* +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 v1alpha1 + +import ( + "github.com/kubeagi/arcadia/pkg/embeddings" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" +) + +// EDIT THIS FILE! THIS IS SCAFFOLDING FOR YOU TO OWN! +// NOTE: json tags are required. Any new fields you add must have json tags for the fields to be serialized. + +// EmbedderSpec defines the desired state of Embedder +type EmbedderSpec struct { + // INSERT ADDITIONAL SPEC FIELDS - desired state of cluster + // Important: Run "make" to regenerate code after modifying this file + + // Name of the Embedding service + DisplayName string `json:"name,omitempty"` + + // ServiceType indicates the source type of embedding service + ServiceType embeddings.EmbeddingType `json:"serviceType,omitempty"` + + // URL keeps the URL of the embedding service (Must have) + URL string `json:"url,omitempty"` + + // Auth keeps the authentication credentials when access embedding service + // keeps in k8s secret + Auth string `json:"auth,omitempty"` +} + +// EmbeddingsStatus defines the observed state of Embedder +type EmbedderStatus struct { + // INSERT ADDITIONAL STATUS FIELD - define observed state of cluster + // Important: Run "make" to regenerate code after modifying this file + // ConditionedStatus is the current status + ConditionedStatus `json:",inline"` +} + +//+kubebuilder:object:root=true +//+kubebuilder:subresource:status + +// Embedder is the Schema for the embeddings API +type Embedder struct { + metav1.TypeMeta `json:",inline"` + metav1.ObjectMeta `json:"metadata,omitempty"` + + Spec EmbedderSpec `json:"spec,omitempty"` + Status EmbedderStatus `json:"status,omitempty"` +} + +//+kubebuilder:object:root=true + +// EmbedderList contains a list of Embedder +type EmbedderList struct { + metav1.TypeMeta `json:",inline"` + metav1.ListMeta `json:"metadata,omitempty"` + Items []Embedder `json:"items"` +} + +func init() { + SchemeBuilder.Register(&Embedder{}, &EmbedderList{}) +} diff --git a/api/v1alpha1/zz_generated.deepcopy.go b/api/v1alpha1/zz_generated.deepcopy.go index 84f19e72c..cc3e55443 100644 --- a/api/v1alpha1/zz_generated.deepcopy.go +++ b/api/v1alpha1/zz_generated.deepcopy.go @@ -155,6 +155,96 @@ func (in *DatasourceStatus) DeepCopy() *DatasourceStatus { return out } +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *Embedder) DeepCopyInto(out *Embedder) { + *out = *in + out.TypeMeta = in.TypeMeta + in.ObjectMeta.DeepCopyInto(&out.ObjectMeta) + out.Spec = in.Spec + in.Status.DeepCopyInto(&out.Status) +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new Embedder. +func (in *Embedder) DeepCopy() *Embedder { + if in == nil { + return nil + } + out := new(Embedder) + in.DeepCopyInto(out) + return out +} + +// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object. +func (in *Embedder) DeepCopyObject() runtime.Object { + if c := in.DeepCopy(); c != nil { + return c + } + return nil +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *EmbedderList) DeepCopyInto(out *EmbedderList) { + *out = *in + out.TypeMeta = in.TypeMeta + in.ListMeta.DeepCopyInto(&out.ListMeta) + if in.Items != nil { + in, out := &in.Items, &out.Items + *out = make([]Embedder, len(*in)) + for i := range *in { + (*in)[i].DeepCopyInto(&(*out)[i]) + } + } +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new EmbedderList. +func (in *EmbedderList) DeepCopy() *EmbedderList { + if in == nil { + return nil + } + out := new(EmbedderList) + in.DeepCopyInto(out) + return out +} + +// DeepCopyObject is an autogenerated deepcopy function, copying the receiver, creating a new runtime.Object. +func (in *EmbedderList) DeepCopyObject() runtime.Object { + if c := in.DeepCopy(); c != nil { + return c + } + return nil +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *EmbedderSpec) DeepCopyInto(out *EmbedderSpec) { + *out = *in +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new EmbedderSpec. +func (in *EmbedderSpec) DeepCopy() *EmbedderSpec { + if in == nil { + return nil + } + out := new(EmbedderSpec) + in.DeepCopyInto(out) + return out +} + +// DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. +func (in *EmbedderStatus) DeepCopyInto(out *EmbedderStatus) { + *out = *in + in.ConditionedStatus.DeepCopyInto(&out.ConditionedStatus) +} + +// DeepCopy is an autogenerated deepcopy function, copying the receiver, creating a new EmbedderStatus. +func (in *EmbedderStatus) DeepCopy() *EmbedderStatus { + if in == nil { + return nil + } + out := new(EmbedderStatus) + in.DeepCopyInto(out) + return out +} + // DeepCopyInto is an autogenerated deepcopy function, copying the receiver, writing into out. in must be non-nil. func (in *LLM) DeepCopyInto(out *LLM) { *out = *in diff --git a/charts/arcadia/Chart.yaml b/charts/arcadia/Chart.yaml index 375d0adf4..43d52a2b9 100644 --- a/charts/arcadia/Chart.yaml +++ b/charts/arcadia/Chart.yaml @@ -2,7 +2,7 @@ apiVersion: v2 name: arcadia description: A Helm chart(KubeBB Component) for KubeAGI Arcadia type: application -version: 0.1.7 +version: 0.1.8 appVersion: "0.0.0" keywords: - kubeagi diff --git a/charts/arcadia/crds/arcadia.kubeagi.k8s.com.cn_embedders.yaml b/charts/arcadia/crds/arcadia.kubeagi.k8s.com.cn_embedders.yaml new file mode 100644 index 000000000..9abe4115f --- /dev/null +++ b/charts/arcadia/crds/arcadia.kubeagi.k8s.com.cn_embedders.yaml @@ -0,0 +1,99 @@ +--- +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + annotations: + controller-gen.kubebuilder.io/version: v0.9.2 + creationTimestamp: null + name: embedders.arcadia.kubeagi.k8s.com.cn +spec: + group: arcadia.kubeagi.k8s.com.cn + names: + kind: Embedder + listKind: EmbedderList + plural: embedders + singular: embedder + scope: Namespaced + versions: + - name: v1alpha1 + schema: + openAPIV3Schema: + description: Embedder is the Schema for the embedders API + properties: + apiVersion: + description: 'APIVersion defines the versioned schema of this representation + of an object. Servers should convert recognized schemas to the latest + internal value, and may reject unrecognized values. More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#resources' + type: string + kind: + description: 'Kind is a string value representing the REST resource this + object represents. Servers may infer this from the endpoint the client + submits requests to. Cannot be updated. In CamelCase. More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#types-kinds' + type: string + metadata: + type: object + spec: + description: EmbedderSpec defines the desired state of embedder. + properties: + auth: + description: Auth keeps the authentication credentials when access + embedder, keeps in k8s secret. + type: string + displayName: + type: string + type: + description: Type defines the type of embedder. + type: string + url: + description: URL keeps the URL of the embedder service(Must required) + type: string + required: + - type + type: object + status: + description: EmbedderStatus defines the observed state of embedder. + properties: + conditions: + description: Conditions of the resource. + items: + description: A Condition that may apply to a resource. + properties: + lastSuccessfulTime: + description: LastSuccessfulTime is repository Last Successful + Update Time + format: date-time + type: string + lastTransitionTime: + description: LastTransitionTime is the last time this condition + transitioned from one status to another. + format: date-time + type: string + message: + description: A Message containing details about this condition's + last transition from one status to another, if any. + type: string + reason: + description: A Reason for this condition's last transition from + one status to another. + type: string + status: + description: Status of this condition; is it currently True, + False, or Unknown + type: string + type: + description: Type of this condition. At most one of each condition + type may apply to a resource at any point in time. + type: string + required: + - lastTransitionTime + - reason + - status + - type + type: object + type: array + type: object + type: object + served: true + storage: true + subresources: + status: {} diff --git a/config/crd/bases/arcadia.kubeagi.k8s.com.cn_embedders.yaml b/config/crd/bases/arcadia.kubeagi.k8s.com.cn_embedders.yaml new file mode 100644 index 000000000..56a3f370f --- /dev/null +++ b/config/crd/bases/arcadia.kubeagi.k8s.com.cn_embedders.yaml @@ -0,0 +1,98 @@ +--- +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + annotations: + controller-gen.kubebuilder.io/version: v0.9.2 + creationTimestamp: null + name: embedders.arcadia.kubeagi.k8s.com.cn +spec: + group: arcadia.kubeagi.k8s.com.cn + names: + kind: Embedder + listKind: EmbedderList + plural: embedders + singular: embedder + scope: Namespaced + versions: + - name: v1alpha1 + schema: + openAPIV3Schema: + description: Embedder is the Schema for the embeddings API + properties: + apiVersion: + description: 'APIVersion defines the versioned schema of this representation + of an object. Servers should convert recognized schemas to the latest + internal value, and may reject unrecognized values. More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#resources' + type: string + kind: + description: 'Kind is a string value representing the REST resource this + object represents. Servers may infer this from the endpoint the client + submits requests to. Cannot be updated. In CamelCase. More info: https://git.k8s.io/community/contributors/devel/sig-architecture/api-conventions.md#types-kinds' + type: string + metadata: + type: object + spec: + description: EmbedderSpec defines the desired state of Embedder + properties: + auth: + description: Auth keeps the authentication credentials when access + embedding service keeps in k8s secret + type: string + name: + description: Name of the Embedding service + type: string + serviceType: + description: ServiceType indicates the source type of embedding service + type: string + url: + description: URL keeps the URL of the embedding service (Must have) + type: string + type: object + status: + description: EmbeddingsStatus defines the observed state of Embedder + properties: + conditions: + description: Conditions of the resource. + items: + description: A Condition that may apply to a resource. + properties: + lastSuccessfulTime: + description: LastSuccessfulTime is repository Last Successful + Update Time + format: date-time + type: string + lastTransitionTime: + description: LastTransitionTime is the last time this condition + transitioned from one status to another. + format: date-time + type: string + message: + description: A Message containing details about this condition's + last transition from one status to another, if any. + type: string + reason: + description: A Reason for this condition's last transition from + one status to another. + type: string + status: + description: Status of this condition; is it currently True, + False, or Unknown + type: string + type: + description: Type of this condition. At most one of each condition + type may apply to a resource at any point in time. + type: string + required: + - lastTransitionTime + - reason + - status + - type + type: object + type: array + type: object + type: object + served: true + storage: true + subresources: + status: {} diff --git a/config/crd/kustomization.yaml b/config/crd/kustomization.yaml index e543e4197..1e3c33e33 100644 --- a/config/crd/kustomization.yaml +++ b/config/crd/kustomization.yaml @@ -6,6 +6,7 @@ resources: - bases/arcadia.kubeagi.k8s.com.cn_llms.yaml - bases/arcadia.kubeagi.k8s.com.cn_prompts.yaml - bases/arcadia.kubeagi.k8s.com.cn_datasources.yaml +- bases/arcadia.kubeagi.k8s.com.cn_embedders.yaml #+kubebuilder:scaffold:crdkustomizeresource patchesStrategicMerge: @@ -15,6 +16,7 @@ patchesStrategicMerge: #- patches/webhook_in_llms.yaml #- patches/webhook_in_prompts.yaml #- patches/webhook_in_datasources.yaml +#- patches/webhook_in_embedders.yaml #+kubebuilder:scaffold:crdkustomizewebhookpatch # [CERTMANAGER] To enable cert-manager, uncomment all the sections with [CERTMANAGER] prefix. @@ -23,6 +25,7 @@ patchesStrategicMerge: #- patches/cainjection_in_llms.yaml #- patches/cainjection_in_prompts.yaml #- patches/cainjection_in_datasources.yaml +#- patches/cainjection_in_embedders.yaml #+kubebuilder:scaffold:crdkustomizecainjectionpatch # the following config is for teaching kustomize how to do kustomization for CRDs. diff --git a/config/crd/patches/cainjection_in_embedders.yaml b/config/crd/patches/cainjection_in_embedders.yaml new file mode 100644 index 000000000..19bef8ed7 --- /dev/null +++ b/config/crd/patches/cainjection_in_embedders.yaml @@ -0,0 +1,7 @@ +# The following patch adds a directive for certmanager to inject CA into the CRD +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + annotations: + cert-manager.io/inject-ca-from: $(CERTIFICATE_NAMESPACE)/$(CERTIFICATE_NAME) + name: embedders.arcadia.kubeagi.k8s.com.cn diff --git a/config/crd/patches/webhook_in_embedders.yaml b/config/crd/patches/webhook_in_embedders.yaml new file mode 100644 index 000000000..7cd63da9b --- /dev/null +++ b/config/crd/patches/webhook_in_embedders.yaml @@ -0,0 +1,16 @@ +# The following patch enables a conversion webhook for the CRD +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + name: embedders.arcadia.kubeagi.k8s.com.cn +spec: + conversion: + strategy: Webhook + webhook: + clientConfig: + service: + namespace: system + name: webhook-service + path: /convert + conversionReviewVersions: + - v1 diff --git a/config/rbac/embedders_editor_role.yaml b/config/rbac/embedders_editor_role.yaml new file mode 100644 index 000000000..b0c4e893d --- /dev/null +++ b/config/rbac/embedders_editor_role.yaml @@ -0,0 +1,31 @@ +# permissions for end users to edit embeddings. +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + labels: + app.kubernetes.io/name: clusterrole + app.kubernetes.io/instance: embedders-editor-role + app.kubernetes.io/component: rbac + app.kubernetes.io/created-by: arcadia + app.kubernetes.io/part-of: arcadia + app.kubernetes.io/managed-by: kustomize + name: embedders-editor-role +rules: +- apiGroups: + - arcadia.kubeagi.k8s.com.cn + resources: + - embedders + verbs: + - create + - delete + - get + - list + - patch + - update + - watch +- apiGroups: + - arcadia.kubeagi.k8s.com.cn + resources: + - embedders/status + verbs: + - get diff --git a/config/rbac/embedders_viewer_role.yaml b/config/rbac/embedders_viewer_role.yaml new file mode 100644 index 000000000..b64c8a711 --- /dev/null +++ b/config/rbac/embedders_viewer_role.yaml @@ -0,0 +1,27 @@ +# permissions for end users to view embeddings. +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + labels: + app.kubernetes.io/name: clusterrole + app.kubernetes.io/instance: embedders-viewer-role + app.kubernetes.io/component: rbac + app.kubernetes.io/created-by: arcadia + app.kubernetes.io/part-of: arcadia + app.kubernetes.io/managed-by: kustomize + name: embedders-viewer-role +rules: +- apiGroups: + - arcadia.kubeagi.k8s.com.cn + resources: + - embedders + verbs: + - get + - list + - watch +- apiGroups: + - arcadia.kubeagi.k8s.com.cn + resources: + - embedders/status + verbs: + - get diff --git a/config/rbac/role.yaml b/config/rbac/role.yaml index 9e57faf27..f8a56e407 100644 --- a/config/rbac/role.yaml +++ b/config/rbac/role.yaml @@ -31,6 +31,32 @@ rules: - get - patch - update +- apiGroups: + - arcadia.kubeagi.k8s.com.cn + resources: + - embedders + verbs: + - create + - delete + - get + - list + - patch + - update + - watch +- apiGroups: + - arcadia.kubeagi.k8s.com.cn + resources: + - embedders/finalizers + verbs: + - update +- apiGroups: + - arcadia.kubeagi.k8s.com.cn + resources: + - embedders/status + verbs: + - get + - patch + - update - apiGroups: - arcadia.kubeagi.k8s.com.cn resources: diff --git a/config/samples/arcadia_v1alpha1_embedders.yaml b/config/samples/arcadia_v1alpha1_embedders.yaml new file mode 100644 index 000000000..3c465bf87 --- /dev/null +++ b/config/samples/arcadia_v1alpha1_embedders.yaml @@ -0,0 +1,22 @@ +apiVersion: v1 +kind: Secret +metadata: + name: zhipuai +type: Opaque +data: + apiKey: "NGZjY2VjZWIxNjY2Y2QxMTgwOGMyMThkNmQ2MTk5NTAuVENYVXZhUUNXRnlJa3hCMw==" # replace this with your API key +--- +apiVersion: arcadia.kubeagi.k8s.com.cn/v1alpha1 +kind: Embedders +metadata: + labels: + app.kubernetes.io/name: embedders + app.kubernetes.io/instance: embeddings-sample + app.kubernetes.io/part-of: arcadia + app.kubernetes.io/managed-by: kustomize + app.kubernetes.io/created-by: arcadia + name: zhipuai-embedders-sample +spec: + serviceType: "zhipuai" + url: "https://open.bigmodel.cn/api/paas/v3/model-api" # replace this with your LLM URL(Zhipuai use predefined url https://open.bigmodel.cn/api/paas/v3/model-api) + auth: "zhipuai" # replace this with your auth secret diff --git a/config/samples/kustomization.yaml b/config/samples/kustomization.yaml index 50a3b8d80..6163e6797 100644 --- a/config/samples/kustomization.yaml +++ b/config/samples/kustomization.yaml @@ -4,4 +4,5 @@ resources: - arcadia_v1alpha1_llm.yaml - arcadia_v1alpha1_prompt.yaml - arcadia_v1alpha1_datasource.yaml +- arcadia_v1alpha1_embedders.yaml #+kubebuilder:scaffold:manifestskustomizesamples diff --git a/controllers/embedder_controller.go b/controllers/embedder_controller.go new file mode 100644 index 000000000..d5ab1356e --- /dev/null +++ b/controllers/embedder_controller.go @@ -0,0 +1,143 @@ +/* +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 controllers + +import ( + "context" + "fmt" + "github.com/kubeagi/arcadia/pkg/llms/zhipuai" + + "github.com/go-logr/logr" + 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" + ctrl "sigs.k8s.io/controller-runtime" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/log" + + arcadiav1alpha1 "github.com/kubeagi/arcadia/api/v1alpha1" + "github.com/kubeagi/arcadia/pkg/embeddings" +) + +const ( + _StatusNilResponse = "No err replied but response is not string" +) + +// EmbedderReconciler reconciles a Embedder object +type EmbedderReconciler struct { + client.Client + Scheme *runtime.Scheme +} + +//+kubebuilder:rbac:groups=arcadia.kubeagi.k8s.com.cn,resources=embedders,verbs=get;list;watch;create;update;patch;delete +//+kubebuilder:rbac:groups=arcadia.kubeagi.k8s.com.cn,resources=embedders/status,verbs=get;update;patch +//+kubebuilder:rbac:groups=arcadia.kubeagi.k8s.com.cn,resources=embedders/finalizers,verbs=update + +// Reconcile is part of the main kubernetes reconciliation loop which aims to +// move the current state of the cluster closer to the desired state. +// TODO(user): Modify the Reconcile function to compare the state specified by +// the Embedder object against the actual cluster state, and then +// perform operations to make the cluster state reflect the state specified by +// the user. +// +// For more details, check Reconcile and its Result here: +// - https://pkg.go.dev/sigs.k8s.io/controller-runtime@v0.14.4/pkg/reconcile +func (r *EmbedderReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) { + logger := log.FromContext(ctx) + logger.Info("Reconciling embedding resource") + + instance := &arcadiav1alpha1.Embedder{} + err := r.Get(ctx, req.NamespacedName, instance) + if err != nil { + if errors.IsNotFound(err) { + return ctrl.Result{}, nil + } + return ctrl.Result{}, err + } + + err = r.CheckEmbedder(ctx, logger, instance) + if err != nil { + return ctrl.Result{}, err + } + + return ctrl.Result{}, nil +} + +// SetupWithManager sets up the controller with the Manager. +func (r *EmbedderReconciler) SetupWithManager(mgr ctrl.Manager) error { + return ctrl.NewControllerManagedBy(mgr). + For(&arcadiav1alpha1.Embedder{}). + Complete(r) +} + +func (r *EmbedderReconciler) CheckEmbedder(ctx context.Context, logger logr.Logger, instance *arcadiav1alpha1.Embedder) error { + logger.Info("Checking embedding resource") + var err error + var msg string + + // Check Auth availability + apiKey, err := instance.AuthAPIKey(ctx, r.Client) + if err != nil { + return r.UpdateStatus(ctx, instance, nil, err) + } + + switch instance.Spec.ServiceType { + case embeddings.ZhiPuAI: + { + embedClient := zhipuai.NewZhiPuAI(apiKey) + res, err := embedClient.Validate() + if err != nil { + return r.UpdateStatus(ctx, instance, nil, err) + } + msg = res.String() + } + default: + return r.UpdateStatus(ctx, instance, nil, fmt.Errorf("unsupported service type: %s", instance.Spec.ServiceType)) + } + + return r.UpdateStatus(ctx, instance, msg, err) +} + +func (r *EmbedderReconciler) UpdateStatus(ctx context.Context, instance *arcadiav1alpha1.Embedder, t interface{}, err error) error { + instanceCopy := instance.DeepCopy() + if err != nil { + // Set status to unavailable + instanceCopy.Status.SetConditions(arcadiav1alpha1.Condition{ + Type: arcadiav1alpha1.TypeReady, + Status: corev1.ConditionFalse, + Reason: arcadiav1alpha1.ReasonUnavailable, + Message: err.Error(), + LastTransitionTime: metav1.Now(), + }) + } else { + msg, ok := t.(string) + if !ok { + msg = _StatusNilResponse + } + // Set status to available + instanceCopy.Status.SetConditions(arcadiav1alpha1.Condition{ + Type: arcadiav1alpha1.TypeReady, + Status: corev1.ConditionTrue, + Reason: arcadiav1alpha1.ReasonAvailable, + Message: msg, + LastTransitionTime: metav1.Now(), + LastSuccessfulTime: metav1.Now(), + }) + } + return r.Client.Status().Update(ctx, instanceCopy) +} diff --git a/examples/chat_with_document/load.go b/examples/chat_with_document/load.go index 8a50eba0f..7d5915045 100644 --- a/examples/chat_with_document/load.go +++ b/examples/chat_with_document/load.go @@ -19,15 +19,14 @@ package main import ( "context" "fmt" - "github.com/tmc/langchaingo/textsplitter" - "os" - "github.com/spf13/cobra" "github.com/tmc/langchaingo/documentloaders" + "github.com/tmc/langchaingo/textsplitter" + "github.com/tmc/langchaingo/vectorstores/chroma" + "os" zhipuaiembeddings "github.com/kubeagi/arcadia/pkg/embeddings/zhipuai" "github.com/kubeagi/arcadia/pkg/llms/zhipuai" - "github.com/kubeagi/arcadia/pkg/vectorstores/chromadb" ) var ( @@ -93,10 +92,10 @@ func runLoad(ctx context.Context) error { } fmt.Println("Connecting vector database...") - _, err = chromadb.New( - chromadb.WithURL(url), - chromadb.WithEmbedder(embedder), - chromadb.WithNameSpace(namespace), + _, err = chroma.New( + chroma.WithChromaURL(url), + chroma.WithEmbedder(embedder), + chroma.WithNameSpace(namespace), ) if err != nil { return fmt.Errorf("error connecting chroma db: %s", err.Error()) @@ -212,10 +211,10 @@ func loadPDFFiles(ctx context.Context, path string) error { return fmt.Errorf("error creating ZhiPuAI embedder: %s", err.Error()) } - db, err := chromadb.New( - chromadb.WithURL(url), - chromadb.WithEmbedder(embedder), - chromadb.WithNameSpace(namespace), + db, err := chroma.New( + chroma.WithChromaURL(url), + chroma.WithEmbedder(embedder), + chroma.WithNameSpace(namespace), ) if err != nil { return fmt.Errorf("error connecting chroma db: %s", err.Error()) diff --git a/main.go b/main.go index 3939011be..344919b51 100644 --- a/main.go +++ b/main.go @@ -152,6 +152,13 @@ func main() { setupLog.Error(err, "unable to create controller", "controller", "Datasource") os.Exit(1) } + if err = (&controllers.EmbedderReconciler{ + Client: mgr.GetClient(), + Scheme: mgr.GetScheme(), + }).SetupWithManager(mgr); err != nil { + setupLog.Error(err, "unable to create controller", "controller", "Embedder") + os.Exit(1) + } //+kubebuilder:scaffold:builder if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil { diff --git a/pkg/embeddings/embeddings.go b/pkg/embeddings/embeddings.go new file mode 100644 index 000000000..f0727b60e --- /dev/null +++ b/pkg/embeddings/embeddings.go @@ -0,0 +1,32 @@ +/* +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 embeddings + +import langchaingoembeddings "github.com/tmc/langchaingo/embeddings" + +type EmbeddingType string + +const ( + OpenAI EmbeddingType = "openai" + ZhiPuAI EmbeddingType = "zhipuai" + Local EmbeddingType = "local" +) + +type Embedding interface { + Type() EmbeddingType + langchaingoembeddings.Embedder +} diff --git a/pkg/llms/zhipuai/api.go b/pkg/llms/zhipuai/api.go index d46ac46eb..055841d8f 100644 --- a/pkg/llms/zhipuai/api.go +++ b/pkg/llms/zhipuai/api.go @@ -25,6 +25,7 @@ import ( "time" "github.com/r3labs/sse/v2" + "k8s.io/klog/v2" "github.com/kubeagi/arcadia/pkg/llms" ) @@ -32,6 +33,7 @@ import ( const ( ZhipuaiModelAPIURL = "https://open.bigmodel.cn/api/paas/v3/model-api" ZhipuaiModelDefaultTimeout = 300 * time.Second + RetryLimit = 5 ) type Model string @@ -149,7 +151,7 @@ func (z *ZhiPuAI) Validate() (llms.Response, error) { testPrompt := []Prompt{ { Role: "user", - Content: "Hello!", + Content: "Hello", }, } @@ -196,16 +198,12 @@ func (z *ZhiPuAI) CreateEmbedding(ctx context.Context, inputTexts []string) ([][ embeddings := make([][]float32, 0, len(inputTexts)) for _, text := range inputTexts { var retry int - startTime := time.Now() - postResponse := &EmbeddingResponse{} - for retry < 5 && postResponse.Success == false { + for retry < RetryLimit && postResponse.Success == false { time.Sleep(100 * time.Millisecond) retry++ - if retry > 1 { - fmt.Printf("attempt %d... ", retry) - } + klog.Info("retry embedding post quest:", retry) postResponse, err = EmbeddingPost(url, token, EmbeddingText{ Prompt: text, }, ZhipuaiModelDefaultTimeout) @@ -219,9 +217,6 @@ func (z *ZhiPuAI) CreateEmbedding(ctx context.Context, inputTexts []string) ([][ } embeddings = append(embeddings, postResponse.Data.Embedding) - - elapsedTime := time.Since(startTime) / time.Millisecond - fmt.Printf("embedding post time cost: %dms\n", elapsedTime) } return embeddings, nil