Bladeren bron

Gitee AI embedding tool (#10903)

方程 5 maanden geleden
bovenliggende
commit
99b0369f1b

+ 25 - 0
api/core/tools/provider/builtin/gitee_ai/tools/embedding.py

@@ -0,0 +1,25 @@
+from typing import Any, Union
+
+import requests
+
+from core.tools.entities.tool_entities import ToolInvokeMessage
+from core.tools.tool.builtin_tool import BuiltinTool
+
+
+class GiteeAIToolEmbedding(BuiltinTool):
+    def _invoke(
+        self, user_id: str, tool_parameters: dict[str, Any]
+    ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
+        headers = {
+            "content-type": "application/json",
+            "authorization": f"Bearer {self.runtime.credentials['api_key']}",
+        }
+
+        payload = {"inputs": tool_parameters.get("inputs")}
+        model = tool_parameters.get("model", "bge-m3")
+        url = f"https://ai.gitee.com/api/serverless/{model}/embeddings"
+        response = requests.post(url, json=payload, headers=headers)
+        if response.status_code != 200:
+            return self.create_text_message(f"Got Error Response:{response.text}")
+
+        return [self.create_text_message(response.content.decode("utf-8"))]

+ 37 - 0
api/core/tools/provider/builtin/gitee_ai/tools/embedding.yaml

@@ -0,0 +1,37 @@
+identity:
+  name: embedding
+  author: gitee_ai
+  label:
+    en_US: embedding
+  icon: icon.svg
+description:
+  human:
+    en_US: Generate word embeddings using Serverless-supported models (compatible with OpenAI)
+  llm: This tool is used to generate word embeddings from text input.
+parameters:
+  - name: model
+    type: string
+    required: true
+    in: path
+    description:
+      en_US: Supported Embedding (compatible with OpenAI) interface models
+    enum:
+      - bge-m3
+      - bge-large-zh-v1.5
+      - bge-small-zh-v1.5
+    label:
+      en_US: Service Model
+      zh_Hans: 服务模型
+    default: bge-m3
+    form: form
+  - name: inputs
+    type: string
+    required: true
+    label:
+      en_US: Input Text
+      zh_Hans: 输入文本
+    human_description:
+      en_US: The text input used to generate embeddings.
+      zh_Hans: 用于生成词向量的输入文本。
+    llm_description: This text input will be used to generate embeddings.
+    form: llm

+ 1 - 1
api/core/tools/provider/builtin/gitee_ai/tools/text-to-image.py

@@ -6,7 +6,7 @@ from core.tools.entities.tool_entities import ToolInvokeMessage
 from core.tools.tool.builtin_tool import BuiltinTool
 
 
-class GiteeAITool(BuiltinTool):
+class GiteeAIToolText2Image(BuiltinTool):
     def _invoke(
         self, user_id: str, tool_parameters: dict[str, Any]
     ) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]: