Gitee AI embedding tool (#10903)

This commit is contained in:
方程 2024-11-20 17:40:34 +08:00 committed by GitHub
parent d6ea1e2f12
commit 99b0369f1b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 63 additions and 1 deletions

View File

@ -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"))]

View File

@ -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

View File

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