feat: add cogVideo tool (#10456)
parent
22dee4f6f3
commit
4fe5297e35
@ -0,0 +1,24 @@
|
|||||||
|
from typing import Any, Union
|
||||||
|
|
||||||
|
from zhipuai import ZhipuAI
|
||||||
|
|
||||||
|
from core.tools.entities.tool_entities import ToolInvokeMessage
|
||||||
|
from core.tools.tool.builtin_tool import BuiltinTool
|
||||||
|
|
||||||
|
|
||||||
|
class CogVideoTool(BuiltinTool):
|
||||||
|
def _invoke(
|
||||||
|
self, user_id: str, tool_parameters: dict[str, Any]
|
||||||
|
) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
|
||||||
|
client = ZhipuAI(
|
||||||
|
base_url=self.runtime.credentials["zhipuai_base_url"],
|
||||||
|
api_key=self.runtime.credentials["zhipuai_api_key"],
|
||||||
|
)
|
||||||
|
if not tool_parameters.get("prompt") and not tool_parameters.get("image_url"):
|
||||||
|
return self.create_text_message("require at least one of prompt and image_url")
|
||||||
|
|
||||||
|
response = client.videos.generations(
|
||||||
|
model="cogvideox", prompt=tool_parameters.get("prompt"), image_url=tool_parameters.get("image_url")
|
||||||
|
)
|
||||||
|
|
||||||
|
return self.create_json_message(response.dict())
|
||||||
@ -0,0 +1,30 @@
|
|||||||
|
from typing import Any, Union
|
||||||
|
|
||||||
|
import httpx
|
||||||
|
from zhipuai import ZhipuAI
|
||||||
|
|
||||||
|
from core.tools.entities.tool_entities import ToolInvokeMessage
|
||||||
|
from core.tools.tool.builtin_tool import BuiltinTool
|
||||||
|
|
||||||
|
|
||||||
|
class CogVideoJobTool(BuiltinTool):
|
||||||
|
def _invoke(
|
||||||
|
self, user_id: str, tool_parameters: dict[str, Any]
|
||||||
|
) -> Union[ToolInvokeMessage, list[ToolInvokeMessage]]:
|
||||||
|
client = ZhipuAI(
|
||||||
|
api_key=self.runtime.credentials["zhipuai_api_key"],
|
||||||
|
base_url=self.runtime.credentials["zhipuai_base_url"],
|
||||||
|
)
|
||||||
|
|
||||||
|
response = client.videos.retrieve_videos_result(id=tool_parameters.get("id"))
|
||||||
|
result = [self.create_json_message(response.dict())]
|
||||||
|
if response.task_status == "SUCCESS":
|
||||||
|
for item in response.video_result:
|
||||||
|
video_cover_image = self.create_image_message(item.cover_image_url)
|
||||||
|
result.append(video_cover_image)
|
||||||
|
video = self.create_blob_message(
|
||||||
|
blob=httpx.get(item.url).content, meta={"mime_type": "video/mp4"}, save_as=self.VariableKey.VIDEO
|
||||||
|
)
|
||||||
|
result.append(video)
|
||||||
|
|
||||||
|
return result
|
||||||
@ -0,0 +1,21 @@
|
|||||||
|
identity:
|
||||||
|
name: cogvideo_job
|
||||||
|
author: hjlarry
|
||||||
|
label:
|
||||||
|
en_US: CogVideo Result
|
||||||
|
zh_Hans: CogVideo 结果获取
|
||||||
|
description:
|
||||||
|
human:
|
||||||
|
en_US: Get the result of CogVideo tool generation.
|
||||||
|
zh_Hans: 根据 CogVideo 工具返回的 id 获取视频生成结果。
|
||||||
|
llm: Get the result of CogVideo tool generation. The input is the id which is returned by the CogVideo tool. The output is the url of video and video cover image.
|
||||||
|
parameters:
|
||||||
|
- name: id
|
||||||
|
type: string
|
||||||
|
label:
|
||||||
|
en_US: id
|
||||||
|
human_description:
|
||||||
|
en_US: The id returned by the CogVideo.
|
||||||
|
zh_Hans: CogVideo 工具返回的 id。
|
||||||
|
llm_description: The id returned by the cogvideo.
|
||||||
|
form: llm
|
||||||
Loading…
Reference in New Issue