fix: get more accurate emberder segmentation

Signed-off-by: Abirdcfly <fp544037857@gmail.com>
Abirdcfly · Dec 11, 2023 · 2b3daec · 2b3daec
1 parent b175755
commit 2b3daec
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 9 deletions.
diff --git a/controllers/knowledgebase_controller.go b/controllers/knowledgebase_controller.go
@@ -30,6 +30,7 @@ import (
 	"github.com/tmc/langchaingo/documentloaders"
 	langchainembeddings "github.com/tmc/langchaingo/embeddings"
 	"github.com/tmc/langchaingo/llms/openai"
+	"github.com/tmc/langchaingo/schema"
 	"github.com/tmc/langchaingo/textsplitter"
 	"github.com/tmc/langchaingo/vectorstores/chroma"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
@@ -435,17 +436,22 @@ func (r *KnowledgeBaseReconciler) handleFile(ctx context.Context, log logr.Logge
 		return err
 	}
 	dataReader := bytes.NewReader(data)
+	var documents []schema.Document
 	var loader documentloaders.Loader
 	switch filepath.Ext(fileName) {
-	case "txt":
+	case ".txt":
 		loader = documentloaders.NewText(dataReader)
-	case "csv":
+	case ".csv":
 		if v == arcadiav1alpha1.ObjectTypeQA {
 			loader = pkgdocumentloaders.NewQACSV(dataReader, fileName, "q", "a")
+			documents, err = loader.Load(ctx)
+			if err != nil {
+				return err
+			}
 		} else {
 			loader = documentloaders.NewCSV(dataReader)
 		}
-	case "html", "htm":
+	case ".html", ".htm":
 		loader = documentloaders.NewHTML(dataReader)
 	default:
 		loader = documentloaders.NewText(dataReader)
@@ -475,11 +481,15 @@ func (r *KnowledgeBaseReconciler) handleFile(ctx context.Context, log logr.Logge
 	//	)
 	//}
 
-	documents, err := loader.LoadAndSplit(ctx, split)
-	if err != nil {
-		return err
+	if len(documents) == 0 {
+		documents, err = loader.LoadAndSplit(ctx, split)
+		if err != nil {
+			return err
+		}
+	}
+	for i, doc := range documents {
+		log.V(5).Info(fmt.Sprintf("document[%d]: embedding:%s, metadata:%v", i, doc.PageContent, doc.Metadata))
 	}
-
 	switch store.Spec.Type() { // nolint: gocritic
 	case arcadiav1alpha1.VectorStoreTypeChroma:
 		s, err := chroma.New(
@@ -511,6 +521,7 @@ func (r *KnowledgeBaseReconciler) reconcileDelete(ctx context.Context, log logr.
 			chroma.WithChromaURL(vectorStore.Spec.Enpoint.URL),
 			chroma.WithDistanceFunction(vectorStore.Spec.Chroma.DistanceFunction),
 			chroma.WithNameSpace(kb.VectorStoreCollectionName()),
+			chroma.WithOpenAiAPIKey("fake"),
 		)
 		if err != nil {
 			log.Error(err, "reconcile delete: init vector store error, may leave garbage data")

diff --git a/deploy/charts/arcadia/Chart.yaml b/deploy/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.38
+version: 0.1.39
 appVersion: "0.0.1"
 
 keywords:

diff --git a/deploy/charts/arcadia/templates/controller.yaml b/deploy/charts/arcadia/templates/controller.yaml
@@ -27,6 +27,7 @@ spec:
       containers:
       - command:
         - /manager
+        - --v={{ .Values.controller.logLevel }}
         env:
           - name: POD_NAME
             valueFrom:

diff --git a/deploy/charts/arcadia/values.yaml b/deploy/charts/arcadia/values.yaml
@@ -15,6 +15,7 @@ controller:
     requests:
       cpu: 10m
       memory: 64Mi
+  logLevel: 2
 # graphql and bff server
 apiserver:
   image: kubeagi/arcadia:v0.1.0-20231207-11c8738

diff --git a/tests/example-test.sh b/tests/example-test.sh
@@ -180,7 +180,7 @@ kind load docker-image controller:example-e2e --name=$KindName
 
 info "3. install arcadia"
 kubectl create namespace arcadia
-helm install -narcadia arcadia deploy/charts/arcadia -f tests/deploy-values.yaml --set controller.image=controller:example-e2e --set apiserver.image=controller:example-e2e --wait --timeout $HelmTimeout
+helm install -narcadia arcadia deploy/charts/arcadia -f tests/deploy-values.yaml --set controller.image=controller:example-e2e --set apiserver.image=controller:example-e2e --wait --timeout $HelmTimeout --set controller.logLevel=5
 
 info "4. check system datasource arcadia-minio(system datasource)"
 waitCRDStatusReady "Datasource" "arcadia" "arcadia-minio"
@@ -264,4 +264,5 @@ waitCRDStatusReady "Application" "arcadia" "base-chat-with-knowledgebase"
 sleep 3
 curl -XPOST http://127.0.0.1:8081/chat --data '{"query":"旷工最小计算单位为多少天？","response_mode":"blocking","conversion_id":"","app_name":"base-chat-with-knowledgebase", "app_namespace":"arcadia"}' | jq -e '.message'
 
+exit 1
 info "all finished! ✅"