feat: add rerank

Signed-off-by: Abirdcfly <fp544037857@gmail.com>

Dockerfile.rerank-mock

@@ -0,0 +1,23 @@
+# Build the manager binary
+ARG GO_VER=1.21
+FROM golang:${GO_VER} as builder
+ARG GOPROXY=https://goproxy.cn,direct
+WORKDIR /workspace
+# Copy the Go Modules manifests
+COPY . .
+# cache deps before building and copying source so that we don't need to re-download as much
+# and so that source changes don't invalidate our downloaded layer
+RUN go env -w GOPROXY=${GOPROXY}
+RUN go mod download
+
+
+# Build
+RUN CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -a -o rerank-mock tests/rerank-mock/main.go
+# Use distroless as minimal base image to package the manager binary
+# Refer to https://github.com/GoogleContainerTools/distroless for more details
+FROM gcr.io/distroless/static:nonroot
+WORKDIR /
+COPY --from=builder /workspace/rerank-mock .
+USER 65532:65532
+
+ENTRYPOINT ["/rerank-mock"]

api/app-node/retriever/v1alpha1/knowledgebaseretriever_types.go

@@ -38,7 +38,7 @@ type CommonRetrieverConfig struct {
 	// NumDocuments is the max number of documents to return.
 	// +kubebuilder:default=5
 	// +kubebuilder:validation:Minimum=1
-	// +kubebuilder:validation:Maximum=10
+	// +kubebuilder:validation:Maximum=50
 	NumDocuments int `json:"numDocuments,omitempty"`
 	// DocNullReturn is the return statement when the query result is empty from the retriever.
 	// +kubebuilder:default="未找到您询问的内容，请详细描述您的问题"

api/app-node/retriever/v1alpha1/rerankretriever_types.go

@@ -0,0 +1,80 @@
+/*
+Copyright 2024 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 (
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+
+	node "github.com/kubeagi/arcadia/api/app-node"
+	"github.com/kubeagi/arcadia/api/base/v1alpha1"
+)
+
+// RerankRetrieverSpec defines the desired state of RerankRetriever
+type RerankRetrieverSpec struct {
+	v1alpha1.CommonSpec   `json:",inline"`
+	CommonRetrieverConfig `json:",inline"`
+	// the endpoint of the rerank
+	// TODO: should change to model or worker
+	Endpoint string `json:"endpoint,omitempty"`
+}
+
+// RerankRetrieverStatus defines the observed state of RerankRetriever
+type RerankRetrieverStatus struct {
+	// ObservedGeneration is the last observed generation.
+	// +optional
+	ObservedGeneration int64 `json:"observedGeneration,omitempty"`
+
+	// ConditionedStatus is the current status
+	v1alpha1.ConditionedStatus `json:",inline"`
+}
+
+//+kubebuilder:object:root=true
+//+kubebuilder:subresource:status
+
+// RerankRetriever is the Schema for the RerankRetriever API
+type RerankRetriever struct {
+	metav1.TypeMeta   `json:",inline"`
+	metav1.ObjectMeta `json:"metadata,omitempty"`
+
+	Spec   RerankRetrieverSpec   `json:"spec,omitempty"`
+	Status RerankRetrieverStatus `json:"status,omitempty"`
+}
+
+//+kubebuilder:object:root=true
+
+// RerankRetrieverList contains a list of RerankRetriever
+type RerankRetrieverList struct {
+	metav1.TypeMeta `json:",inline"`
+	metav1.ListMeta `json:"metadata,omitempty"`
+	Items           []RerankRetriever `json:"items"`
+}
+
+func init() {
+	SchemeBuilder.Register(&RerankRetriever{}, &RerankRetrieverList{})
+}
+
+var _ node.Node = (*RerankRetriever)(nil)
+
+func (c *RerankRetriever) SetRef() {
+	annotations := node.SetRefAnnotations(c.GetAnnotations(), []node.Ref{node.RetrieverRef.Len(1)}, []node.Ref{node.RetrievalQAChainRef.Len(1)})
+	if c.GetAnnotations() == nil {
+		c.SetAnnotations(annotations)
+	}
+	for k, v := range annotations {
+		c.Annotations[k] = v
+	}
+}

apiserver/docs/docs.go

@@ -1728,6 +1728,11 @@ const docTemplate = `{
                     "type": "string",
                     "example": "q: 旷工最小计算单位为多少天？"
                 },
+                "rerank_score": {
+                    "description": "RerankScore",
+                    "type": "number",
+                    "example": 0.58124
+                },
                 "score": {
                     "description": "vector search score",
                     "type": "number",

apiserver/docs/swagger.json

@@ -1722,6 +1722,11 @@
                     "type": "string",
                     "example": "q: 旷工最小计算单位为多少天？"
                 },
+                "rerank_score": {
+                    "description": "RerankScore",
+                    "type": "number",
+                    "example": 0.58124
+                },
                 "score": {
                     "description": "vector search score",
                     "type": "number",

apiserver/docs/swagger.yaml

@@ -290,6 +290,10 @@ definitions:
         description: Question row
         example: 'q: 旷工最小计算单位为多少天？'
         type: string
+      rerank_score:
+        description: RerankScore
+        example: 0.58124
+        type: number
       score:
         description: vector search score
         example: 0.34

apiserver/pkg/chat/chat_docs.go

@@ -306,10 +306,10 @@ func (cs *ChatServer) GenerateSingleDocSummary(ctx context.Context, req Conversa
 		return "", ErrNoLLMProvidedInApplication
 	}
 	out := map[string]any{
-		"question":       req.Query,
-		"_answer_stream": respStream,
-		"llm":            llm,
-		"documents":      documents,
+		"question": req.Query,
+		runtimebase.OutputAnserStreamChanKeyInArg: respStream,
+		"llm":       llm,
+		"documents": documents,
 	}
 	if req.ResponseMode == "streaming" {
 		out["_need_stream"] = true
@@ -323,7 +323,7 @@ func (cs *ChatServer) GenerateSingleDocSummary(ctx context.Context, req Conversa
 	if err != nil {
 		return "", fmt.Errorf("failed to generate summary due to %s", err.Error())
 	}
-	a, ok := out["_answer"]
+	a, ok := out[runtimebase.OutputAnserKeyInArg]
 	if !ok {
 		return "", errors.New("empty answer")
 	}

config/crd/bases/retriever.arcadia.kubeagi.k8s.com.cn_knowledgebaseretrievers.yaml

@@ -53,7 +53,7 @@ spec:
               numDocuments:
                 default: 5
                 description: NumDocuments is the max number of documents to return.
-                maximum: 10
+                maximum: 50
                 minimum: 1
                 type: integer
               scoreThreshold:

config/crd/bases/retriever.arcadia.kubeagi.k8s.com.cn_rerankretrievers.yaml

@@ -0,0 +1,120 @@
+---
+apiVersion: apiextensions.k8s.io/v1
+kind: CustomResourceDefinition
+metadata:
+  annotations:
+    controller-gen.kubebuilder.io/version: v0.9.2
+  creationTimestamp: null
+  name: rerankretrievers.retriever.arcadia.kubeagi.k8s.com.cn
+spec:
+  group: retriever.arcadia.kubeagi.k8s.com.cn
+  names:
+    kind: RerankRetriever
+    listKind: RerankRetrieverList
+    plural: rerankretrievers
+    singular: rerankretriever
+  scope: Namespaced
+  versions:
+  - name: v1alpha1
+    schema:
+      openAPIV3Schema:
+        description: RerankRetriever is the Schema for the RerankRetriever 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: RerankRetrieverSpec defines the desired state of RerankRetriever
+            properties:
+              creator:
+                description: Creator defines datasource creator (AUTO-FILLED by webhook)
+                type: string
+              description:
+                description: Description defines datasource description
+                type: string
+              displayName:
+                description: DisplayName defines datasource display name
+                type: string
+              docNullReturn:
+                default: 未找到您询问的内容，请详细描述您的问题
+                description: DocNullReturn is the return statement when the query
+                  result is empty from the retriever.
+                type: string
+              endpoint:
+                description: 'the endpoint of the rerank TODO: should change to model
+                  or worker'
+                type: string
+              numDocuments:
+                default: 5
+                description: NumDocuments is the max number of documents to return.
+                maximum: 50
+                minimum: 1
+                type: integer
+              scoreThreshold:
+                default: 0.3
+                description: ScoreThreshold is the cosine distance float score threshold.
+                  Lower score represents more similarity.
+                maximum: 1
+                minimum: 0
+                type: number
+            type: object
+          status:
+            description: RerankRetrieverStatus defines the observed state of RerankRetriever
+            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
+              observedGeneration:
+                description: ObservedGeneration is the last observed generation.
+                format: int64
+                type: integer
+            type: object
+        type: object
+    served: true
+    storage: true
+    subresources:
+      status: {}