Merge pull request #78 from bjwswang/main

feat: fully enable embedding in arcadia with chromadb as the vector s…

README.md

@@ -4,25 +4,17 @@ Our vision is to make it easier for cloud-native applications to integrate with
 
 ## Quick start
 
-### 1. Install arcadia
+1. Install arcadia operator
 
 ```shell
 helm repo add arcadia https://kubeagi.github.io/arcadia
 helm repo update
 helm install arcadia arcadia/arcadia
 ```
 
-### 2. Add a LLM
+2. Add a LLM along with the auth secret
 
-> Take [ZhiPuAI](https://www.zhipuai.cn/) as an example.
-
-1. Prepare auth info
-
-Update apiKey(`Base64 encoded`) in [zhipuai's secret](https://github.com/kubeagi/arcadia/blob/main/config/samples/arcadia_v1alpha1_llm.yaml#L7).
-
-> On how to get apiKey, please refer to [ZhiPuAI](https://open.bigmodel.cn/dev/api#auth)
-
-2. Create a LLM along with the auth secret
+> Update apiKey(`Base64 encoded`) in [secret](https://github.com/kubeagi/arcadia/blob/main/config/samples/arcadia_v1alpha1_llm.yaml#L7).
 
 ```shell
 kubectl apply -f config/samples/arcadia_v1alpha1_llm.yaml
@@ -52,17 +44,26 @@ Output:
 {"code":200,"data":{"choices":[{"content":"\" Kubernetes (also known as K8s) is an open-source container orchestration system for automating the deployment, scaling, and management of containerized applications. It was originally designed by Google, and is now maintained by the Cloud Native Computing Foundation (CNCF).\\n\\nKubernetes provides a platform-as-a-service (PaaS) model, which allows developers to deploy, run, and scale containerized applications with minimal configuration and effort. It does this by abstracting the underlying infrastructure and providing a common set of APIs and tools that can be used to deploy, manage, and scale applications consistently across different environments.\\n\\nKubernetes is widely adopted by organizations of all sizes and has a large, active community of developers contributing to its continued development and improvement. It is available on a variety of platforms, including Linux, Windows, and 移动设备，and can be deployed on-premises, in the cloud, or in a hybrid environment.\"","role":"assistant"}],"request_id":"7865480399259975113","task_id":"7865480399259975113","task_status":"SUCCESS","usage":{"total_tokens":203}},"msg":"操作成功","success":true}
 ```
 
-## Clients
+## Packages
+
+To enhace the AI capability in Golang,we developed some packages.
+
+### LLMs
+
+- ✅ [ZhiPuAI(智谱AI)](https://github.com/kubeagi/arcadia/tree/main/pkg/llms/zhipuai)
+  - [example](https://github.com/kubeagi/arcadia/blob/main/examples/zhipuai/main.go)
+
+### Embeddings
+
+> Fully compatible with [langchain embeddings](https://github.com/tmc/langchaingo/tree/main/embeddings)
 
-We developed some clients in Golang to interfact with AI in Golang.
+- ✅[ZhiPuAI(智谱AI) Embedding](https://github.com/kubeagi/arcadia/tree/main/pkg/embeddings/zhipuai)
 
-### ZhiPuAI(ChatGLM)
+### VectorStores
 
-- [Client](https://github.com/kubeagi/arcadia/tree/main/pkg/llms/zhipuai)
-- [Examples on How to use this Library](https://github.com/kubeagi/arcadia/blob/main/examples/zhipuai/main.go)
+> Fully compatible with [langchain vectorstores](https://github.com/tmc/langchaingo/tree/main/vectorstores)
 
-For ChatGLM,see [here](https://github.com/THUDM/ChatGLM2-6B)
-For `智谱AI`,see [here](https://open.bigmodel.cn/)
+- ✅[ChromaDB](https://github.com/tmc/langchaingo/tree/main/vectorstores)
 
 ## Contribute to Arcadia
 

api/v1alpha1/llm_types.go

@@ -17,8 +17,9 @@ limitations under the License.
 package v1alpha1
 
 import (
-	"github.com/kubeagi/arcadia/pkg/llms"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
+	"github.com/kubeagi/arcadia/pkg/llms"
 )
 
 // LLMSpec defines the desired state of LLM

api/v1alpha1/prompt_types.go

@@ -17,8 +17,9 @@ limitations under the License.
 package v1alpha1
 
 import (
-	llmzhipuai "github.com/kubeagi/arcadia/pkg/llms/zhipuai"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
+	llmzhipuai "github.com/kubeagi/arcadia/pkg/llms/zhipuai"
 )
 
 // PromptSpec defines the desired state of Prompt

api/v1alpha1/prompt_webhook.go

@@ -19,11 +19,12 @@ package v1alpha1
 import (
 	"context"
 
-	llmzhipuai "github.com/kubeagi/arcadia/pkg/llms/zhipuai"
 	"k8s.io/apimachinery/pkg/runtime"
 	ctrl "sigs.k8s.io/controller-runtime"
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/webhook"
+
+	llmzhipuai "github.com/kubeagi/arcadia/pkg/llms/zhipuai"
 )
 
 // log is for logging in this package.
@@ -59,11 +60,11 @@ func (p *Prompt) Default(ctx context.Context, obj runtime.Object) error {
 var _ webhook.CustomValidator = &Prompt{}
 
 // ValidateCreate implements webhook.Validator so a webhook will be registered for the type
-func (r *Prompt) ValidateCreate(ctx context.Context, obj runtime.Object) error {
-	promptlog.Info("validate create", "name", r.Name)
+func (p *Prompt) ValidateCreate(ctx context.Context, obj runtime.Object) error {
+	promptlog.Info("validate create", "name", p.Name)
 
-	if r.Spec.ZhiPuAIParams != nil {
-		if err := llmzhipuai.ValidateModelParams(*r.Spec.ZhiPuAIParams); err != nil {
+	if p.Spec.ZhiPuAIParams != nil {
+		if err := llmzhipuai.ValidateModelParams(*p.Spec.ZhiPuAIParams); err != nil {
 			promptlog.Error(err, "validate model params")
 			return err
 		}
@@ -73,11 +74,11 @@ func (r *Prompt) ValidateCreate(ctx context.Context, obj runtime.Object) error {
 }
 
 // ValidateUpdate implements webhook.Validator so a webhook will be registered for the type
-func (r *Prompt) ValidateUpdate(ctx context.Context, oldObj runtime.Object, newObj runtime.Object) error {
-	promptlog.Info("validate update", "name", r.Name)
+func (p *Prompt) ValidateUpdate(ctx context.Context, oldObj runtime.Object, newObj runtime.Object) error {
+	promptlog.Info("validate update", "name", p.Name)
 
-	if r.Spec.ZhiPuAIParams != nil {
-		if err := llmzhipuai.ValidateModelParams(*r.Spec.ZhiPuAIParams); err != nil {
+	if p.Spec.ZhiPuAIParams != nil {
+		if err := llmzhipuai.ValidateModelParams(*p.Spec.ZhiPuAIParams); err != nil {
 			promptlog.Error(err, "validate model params")
 			return err
 		}
@@ -87,7 +88,7 @@ func (r *Prompt) ValidateUpdate(ctx context.Context, oldObj runtime.Object, newO
 }
 
 // ValidateDelete implements webhook.Validator so a webhook will be registered for the type
-func (r *Prompt) ValidateDelete(ctx context.Context, obj runtime.Object) error {
-	promptlog.Info("validate delete", "name", r.Name)
+func (p *Prompt) ValidateDelete(ctx context.Context, obj runtime.Object) error {
+	promptlog.Info("validate delete", "name", p.Name)
 	return nil
 }

api/v1alpha1/webhook_suite_test.go

@@ -25,9 +25,6 @@ import (
 	"testing"
 	"time"
 
-	. "github.com/onsi/ginkgo"
-	. "github.com/onsi/gomega"
-
 	admissionv1beta1 "k8s.io/api/admission/v1beta1"
 	//+kubebuilder:scaffold:imports
 	"k8s.io/apimachinery/pkg/runtime"
@@ -38,6 +35,9 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/envtest/printer"
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
 )
 
 // These tests use Ginkgo (BDD-style Go testing framework). Refer to

config/samples/arcadia_v1alpha1_llm.yaml

@@ -2,6 +2,7 @@ apiVersion: v1
 kind: Secret
 metadata:
   name: zhipuai
+  namespace: kubebb-system
 type: Opaque
 data:
   apiKey: "ZmVkMDk2NWJjZTAxOTBmZjJiYzY4MWFjMzA2ZDVmM2QuZUlwN3NPWHJueG1XSnhPaw==" # replace this with your API key
@@ -10,6 +11,7 @@ apiVersion: arcadia.kubeagi.k8s.com.cn/v1alpha1
 kind: LLM
 metadata:
   name: zhipuai
+  namespace: kubebb-system
 spec:
   type: "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)

controllers/llm_controller.go

@@ -21,9 +21,6 @@ import (
 	"reflect"
 
 	"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"
@@ -37,6 +34,9 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 
 	arcadiav1alpha1 "github.com/kubeagi/arcadia/api/v1alpha1"
+	"github.com/kubeagi/arcadia/pkg/llms"
+	"github.com/kubeagi/arcadia/pkg/llms/openai"
+	"github.com/kubeagi/arcadia/pkg/llms/zhipuai"
 )
 
 // LLMReconciler reconciles a LLM object

controllers/prompt_controller.go

@@ -21,10 +21,6 @@ import (
 	"reflect"
 
 	"github.com/go-logr/logr"
-	arcadiav1alpha1 "github.com/kubeagi/arcadia/api/v1alpha1"
-	"github.com/kubeagi/arcadia/pkg/llms"
-	"github.com/kubeagi/arcadia/pkg/llms/openai"
-	llmszhipuai "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"
@@ -37,6 +33,11 @@ import (
 	"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"
+	"github.com/kubeagi/arcadia/pkg/llms"
+	"github.com/kubeagi/arcadia/pkg/llms/openai"
+	llmszhipuai "github.com/kubeagi/arcadia/pkg/llms/zhipuai"
 )
 
 // PromptReconciler reconciles a Prompt object

controllers/suite_test.go

@@ -20,9 +20,6 @@ import (
 	"path/filepath"
 	"testing"
 
-	. "github.com/onsi/ginkgo"
-	. "github.com/onsi/gomega"
-
 	"k8s.io/client-go/kubernetes/scheme"
 	"k8s.io/client-go/rest"
 	"sigs.k8s.io/controller-runtime/pkg/client"
@@ -32,7 +29,9 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/log/zap"
 
 	arcadiav1alpha1 "github.com/kubeagi/arcadia/api/v1alpha1"
-	//+kubebuilder:scaffold:imports
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
 )
 
 // These tests use Ginkgo (BDD-style Go testing framework). Refer to

examples/embedding/main.go

@@ -17,47 +17,49 @@ limitations under the License.
 package main
 
 import (
-	llmszhipuai "github.com/kubeagi/arcadia/pkg/llms/zhipuai"
-	"k8s.io/klog/v2"
+	"context"
+	"fmt"
+	"log"
 	"os"
+
+	embedding "github.com/kubeagi/arcadia/pkg/embeddings/zhipuai"
+	"github.com/kubeagi/arcadia/pkg/llms/zhipuai"
+	"github.com/kubeagi/arcadia/pkg/vectorstores/chromadb"
+	"github.com/tmc/langchaingo/schema"
 )
 
 func main() {
-	// usage: go run main.go <apiKey> <testWord>
-	// apiKey 	== The apiKey of user, available on ZhiPuAI dashboard.
-	// testWord == The text to be embedded.
 	if len(os.Args) == 1 {
 		panic("api key is empty")
 	}
 	apiKey := os.Args[1]
 
-	if len(os.Args) == 2 {
-		panic("test word is empty")
+	// init embedder
+	embedder, err := embedding.NewZhiPuAI(
+		embedding.WithClient(*zhipuai.NewZhiPuAI(apiKey)),
+	)
+	if err != nil {
+		panic(fmt.Errorf("error create embedder: %s", err.Error()))
 	}
-	testWord := os.Args[2]
-
-	klog.V(0).Info("try `Embedding`")
-	resp, err := sampleEmbedding(apiKey, testWord)
-
+	// init vector store
+	chroma, err := chromadb.New(context.TODO(), chromadb.WithBasePath("http://localhost:8000"), chromadb.WithEmbedder(embedder))
 	if err != nil {
-		panic(err)
+		panic(fmt.Errorf("error create chroma db: %s", err.Error()))
 	}
-	klog.V(0).Info("Response: \n", resp.Success)
-	klog.V(0).Info("Usage: \n", resp.Data.Usage)
-	klog.V(0).Info("First 5 number of the Result: \n", resp.Data.Embedding[:5])
-	// resp.Data.Embedding is available.
-}
 
-func sampleEmbedding(apiKey, testWord string) (*llmszhipuai.EmbeddingResponse, error) {
-	ai := llmszhipuai.NewZhiPuAI(apiKey)
-	text := llmszhipuai.EmbeddingText{
-		Prompt:    testWord,
-		RequestID: "1",
+	// add documents
+	err = chroma.AddDocuments(context.TODO(), []schema.Document{
+		{PageContent: "This is a document about cats. Cats are great."},
+		{PageContent: "this is a document about dogs. Dogs are great."},
+	})
+	if err != nil {
+		panic(fmt.Errorf("error add documents to chroma db: %s", err.Error()))
 	}
 
-	resp, err := ai.Embedding(text)
+	docs, err := chroma.SimilaritySearch(context.TODO(), "cats", 1)
 	if err != nil {
-		return nil, err
+		log.Fatalf("Error similarity search: %s \n", err.Error())
 	}
-	return resp, nil
+
+	fmt.Printf("SimilaritySearch: %v \n", docs)
 }

examples/zhipuai/main.go

@@ -25,7 +25,7 @@ import (
 )
 
 func main() {
-	if len(os.Args) == 0 {
+	if len(os.Args) == 1 {
 		panic("api key is empty")
 	}
 	apiKey := os.Args[1]

go.mod

@@ -3,7 +3,7 @@ module github.com/kubeagi/arcadia
 go 1.20
 
 require (
-	github.com/amikos-tech/chroma-go v0.0.0-20230809210405-a3fb26bbc4b4
+	github.com/amikos-tech/chroma-go v0.0.0-20230901221218-d0087270239e
 	github.com/go-logr/logr v1.2.0
 	github.com/golang-jwt/jwt v3.2.2+incompatible
 	github.com/onsi/ginkgo v1.16.5

go.sum

@@ -70,8 +70,8 @@ github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuy
 github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
 github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
 github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d/go.mod h1:rBZYJk541a8SKzHPHnH3zbiI+7dagKZ0cgpgrD7Fyho=
-github.com/amikos-tech/chroma-go v0.0.0-20230809210405-a3fb26bbc4b4 h1:ttEvyy/Bo4urvOOedg73ywpq1J3ChoUtGubbka14ptY=
-github.com/amikos-tech/chroma-go v0.0.0-20230809210405-a3fb26bbc4b4/go.mod h1:bPy9xmWK59Ix/nteEqIvPfAI0L07di+uZMb+RiYbles=
+github.com/amikos-tech/chroma-go v0.0.0-20230901221218-d0087270239e h1:RYla80qSHCeP8WL9Hyjfuk9n/KkDcYwzfaYbfc4cqRQ=
+github.com/amikos-tech/chroma-go v0.0.0-20230901221218-d0087270239e/go.mod h1:bPy9xmWK59Ix/nteEqIvPfAI0L07di+uZMb+RiYbles=
 github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY=
 github.com/antlr/antlr4/runtime/Go/antlr v0.0.0-20210826220005-b48c857c3a0e/go.mod h1:F7bn7fEU90QkQ3tnmaTx3LTKLEDqnwWODIYppRQ5hnY=
 github.com/armon/circbuf v0.0.0-20150827004946-bbbad097214e/go.mod h1:3U/XgcO3hCbHZ8TKRvWD2dDTCfh9M9ya+I9JpbB7O8o=

main.go

@@ -24,10 +24,6 @@ import (
 	"path/filepath"
 	"strconv"
 
-	// Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.)
-	// to ensure that exec-entrypoint and run can make use of them.
-	_ "k8s.io/client-go/plugin/pkg/client/auth"
-
 	"k8s.io/apimachinery/pkg/runtime"
 	utilruntime "k8s.io/apimachinery/pkg/util/runtime"
 	clientgoscheme "k8s.io/client-go/kubernetes/scheme"
@@ -37,7 +33,10 @@ import (
 
 	arcadiav1alpha1 "github.com/kubeagi/arcadia/api/v1alpha1"
 	"github.com/kubeagi/arcadia/controllers"
-	//+kubebuilder:scaffold:imports
+
+	// Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.)
+	// to ensure that exec-entrypoint and run can make use of them.
+	_ "k8s.io/client-go/plugin/pkg/client/auth"
 )
 
 var (