feat: able to set/update models in LLM and Embedder

Signed-off-by: bjwswang <[email protected]>
kubeagi · Jan 24, 2024 · 3d5751f · 3d5751f
1 parent 2626e7b
commit 3d5751f
Show file tree

Hide file tree

Showing 7 changed files with 134 additions and 36 deletions.
diff --git a/apiserver/graph/generated/generated.go b/apiserver/graph/generated/generated.go
diff --git a/apiserver/graph/generated/models_gen.go b/apiserver/graph/generated/models_gen.go
diff --git a/apiserver/graph/schema/embedder.graphqls b/apiserver/graph/schema/embedder.graphqls
@@ -54,6 +54,11 @@ input CreateEmbedderInput {
     规则:  目前支持 zhipuai,openai两种接口类型
     """
     type: String
+
+    """
+    此LLM支持调用的模型列表
+    """
+    models: [String!]
 }
 
 input UpdateEmbedderInput {
@@ -80,6 +85,11 @@ input UpdateEmbedderInput {
     规则:  目前支持 zhipuai,openai两种接口类型
     """
     type: String
+
+    """
+    此LLM支持调用的模型列表
+    """
+    models: [String!]
 }
 
 type EmbedderQuery {

diff --git a/apiserver/graph/schema/llm.graphqls b/apiserver/graph/schema/llm.graphqls
@@ -54,6 +54,11 @@ input CreateLLMInput {
     规则:  目前支持 zhipuai,openai两种接口类型
     """
     type: String
+
+        """
+    此LLM支持调用的模型列表
+    """
+    models: [String!]
 }
 
 input UpdateLLMInput {
@@ -79,6 +84,11 @@ input UpdateLLMInput {
     规则:  目前支持 zhipuai,openai两种接口类型
     """
     type: String
+
+    """
+    此LLM支持调用的模型列表
+    """
+    models: [String!]
 }
 
 type LLMQuery {

diff --git a/apiserver/pkg/embedder/embedder.go b/apiserver/pkg/embedder/embedder.go
@@ -130,7 +130,8 @@ func CreateEmbedder(ctx context.Context, c dynamic.Interface, input generated.Cr
 					URL: input.Endpointinput.URL,
 				},
 			},
-			Type: embeddings.EmbeddingType(servicetype),
+			Type:   embeddings.EmbeddingType(servicetype),
+			Models: input.Models,
 		},
 	}
 
@@ -210,6 +211,11 @@ func UpdateEmbedder(ctx context.Context, c dynamic.Interface, input *generated.U
 		updatedEmbedder.Spec.Type = embeddings.EmbeddingType(*input.Type)
 	}
 
+	// update LLM's models if specified
+	if input.Models != nil {
+		updatedEmbedder.Spec.Models = input.Models
+	}
+
 	// Update endpoint
 	if input.Endpointinput != nil {
 		endpoint, err := common.MakeEndpoint(ctx, c, generated.TypedObjectReferenceInput{

diff --git a/apiserver/pkg/llm/llm.go b/apiserver/pkg/llm/llm.go
@@ -185,7 +185,8 @@ func CreateLLM(ctx context.Context, c dynamic.Interface, input generated.CreateL
 					URL: input.Endpointinput.URL,
 				},
 			},
-			Type: llms.LLMType(APIType),
+			Type:   llms.LLMType(APIType),
+			Models: input.Models,
 		},
 	}
 	common.SetCreator(ctx, &llm.Spec.CommonSpec)
@@ -266,6 +267,11 @@ func UpdateLLM(ctx context.Context, c dynamic.Interface, input *generated.Update
 		updatedLLM.Spec.Type = llms.LLMType(*input.Type)
 	}
 
+	// update LLM's models if specified
+	if input.Models != nil {
+		updatedLLM.Spec.Models = input.Models
+	}
+
 	// Update endpoint
 	if input.Endpointinput != nil {
 		endpoint, err := common.MakeEndpoint(ctx, c, generated.TypedObjectReferenceInput{