From 62d4c8fa583abca81312a7b3a16c75e98adbacc3 Mon Sep 17 00:00:00 2001
From: glide-the <2533736852@qq.com>
Date: Fri, 12 Jul 2024 17:53:43 +0800
Subject: [PATCH] =?UTF-8?q?cogvideo=20=E4=B8=9A=E5=8A=A1=E9=9B=86=E6=88=90?=
 =?UTF-8?q?=20(#44)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 tests/integration_tests/test_videos.py     | 43 ++++++++++++
 zhipuai/_client.py                         |  1 +
 zhipuai/api_resource/__init__.py           |  4 ++
 zhipuai/api_resource/videos/__init__.py    |  7 ++
 zhipuai/api_resource/videos/videos.py      | 77 ++++++++++++++++++++++
 zhipuai/types/video/__init__.py            |  7 ++
 zhipuai/types/video/video_create_params.py | 16 +++++
 zhipuai/types/video/video_object.py        | 28 ++++++++
 8 files changed, 183 insertions(+)
 create mode 100644 tests/integration_tests/test_videos.py
 create mode 100644 zhipuai/api_resource/videos/__init__.py
 create mode 100644 zhipuai/api_resource/videos/videos.py
 create mode 100644 zhipuai/types/video/__init__.py
 create mode 100644 zhipuai/types/video/video_create_params.py
 create mode 100644 zhipuai/types/video/video_object.py

diff --git a/tests/integration_tests/test_videos.py b/tests/integration_tests/test_videos.py
new file mode 100644
index 0000000..1b6d553
--- /dev/null
+++ b/tests/integration_tests/test_videos.py
@@ -0,0 +1,43 @@
+from zhipuai import ZhipuAI
+import zhipuai
+
+import logging
+import logging.config
+
+
+def test_videos(logging_conf):
+    logging.config.dictConfig(logging_conf)  # type: ignore
+    client = ZhipuAI()  # 填写您自己的APIKey
+    try:
+        response = client.videos.generations(
+            model="cogvideo",
+            prompt="一些相信光的人，举着奥特曼"
+        )
+        print(response)
+
+
+
+    except zhipuai.core._errors.APIRequestFailedError as err:
+        print(err)
+    except zhipuai.core._errors.APIInternalError as err:
+        print(err)
+    except zhipuai.core._errors.APIStatusError as err:
+        print(err)
+
+
+def test_retrieve_videos_result(logging_conf):
+    logging.config.dictConfig(logging_conf)  # type: ignore
+    client = ZhipuAI()  # 填写您自己的APIKey
+    try:
+        response = client.videos.retrieve_videos_result(
+            id="1014908828144164571910908"
+        )
+        print(response)
+
+
+    except zhipuai.core._errors.APIRequestFailedError as err:
+        print(err)
+    except zhipuai.core._errors.APIInternalError as err:
+        print(err)
+    except zhipuai.core._errors.APIStatusError as err:
+        print(err)
diff --git a/zhipuai/_client.py b/zhipuai/_client.py
index f3b00c4..903469c 100644
--- a/zhipuai/_client.py
+++ b/zhipuai/_client.py
@@ -60,6 +60,7 @@ def __init__(
         self.batches = api_resource.Batches(self)
         self.knowledge = api_resource.Knowledge(self)
         self.tools = api_resource.Tools(self)
+        self.videos = api_resource.Videos(self)
 
     @property
     @override
diff --git a/zhipuai/api_resource/__init__.py b/zhipuai/api_resource/__init__.py
index 34b0bfb..b1692ba 100644
--- a/zhipuai/api_resource/__init__.py
+++ b/zhipuai/api_resource/__init__.py
@@ -27,7 +27,11 @@
 from .tools import (
     Tools
 )
+from .videos import (
+    Videos,
+)
 __all__ = [
+    'Videos',
     'AsyncCompletions',
     'Chat',
     'Completions',
diff --git a/zhipuai/api_resource/videos/__init__.py b/zhipuai/api_resource/videos/__init__.py
new file mode 100644
index 0000000..36303a7
--- /dev/null
+++ b/zhipuai/api_resource/videos/__init__.py
@@ -0,0 +1,7 @@
+from .videos import (
+    Videos,
+)
+__all__ = [
+    'Videos',
+
+]
\ No newline at end of file
diff --git a/zhipuai/api_resource/videos/videos.py b/zhipuai/api_resource/videos/videos.py
new file mode 100644
index 0000000..263b216
--- /dev/null
+++ b/zhipuai/api_resource/videos/videos.py
@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+
+from typing import TYPE_CHECKING, List, Mapping, cast, Optional, Dict
+from typing_extensions import Literal
+
+from ...types.video import video_create_params
+from ...types.video import VideoObject
+from ...core import BaseAPI, maybe_transform
+from ...core import NOT_GIVEN, Body, Headers, NotGiven
+
+import httpx
+
+from ...core import (
+    make_request_options,
+)
+from ...core import deepcopy_minimal, extract_files
+
+if TYPE_CHECKING:
+    from ..._client import ZhipuAI
+
+__all__ = ["Videos"]
+
+
+class Videos(BaseAPI):
+
+    def __init__(self, client: "ZhipuAI") -> None:
+        super().__init__(client)
+
+    def generations(
+            self,
+            model: str,
+            prompt: str,
+            *,
+            request_id: str = None,
+            extra_headers: Headers | None = None,
+            extra_body: Body | None = None,
+            timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+    ) -> VideoObject:
+
+        if not model and not model:
+            raise ValueError("At least one of `model` and `prompt` must be provided.")
+        body = deepcopy_minimal(
+            {
+                "model": model,
+                "prompt": prompt,
+                "request_id": request_id,
+            }
+        )
+        return self._post(
+            "/videos/generations",
+            body=maybe_transform(body, video_create_params.VideoCreateParams),
+            options=make_request_options(
+                extra_headers=extra_headers, extra_body=extra_body, timeout=timeout
+            ),
+            cast_type=VideoObject,
+        )
+
+    def retrieve_videos_result(
+            self,
+            id: str,
+            *,
+            extra_headers: Headers | None = None,
+            extra_body: Body | None = None,
+            timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+    ) -> VideoObject:
+
+        if not id:
+            raise ValueError("At least one of `id` must be provided.")
+
+        return self._get(
+            f"/async-result/{id}",
+            options=make_request_options(
+                extra_headers=extra_headers, extra_body=extra_body, timeout=timeout
+            ),
+            cast_type=VideoObject,
+        )
\ No newline at end of file
diff --git a/zhipuai/types/video/__init__.py b/zhipuai/types/video/__init__.py
new file mode 100644
index 0000000..de085d4
--- /dev/null
+++ b/zhipuai/types/video/__init__.py
@@ -0,0 +1,7 @@
+
+from .video_object import (
+    VideoObject,
+    VideoResult
+)
+
+__all__ = ["VideoObject", "VideoResult"]
diff --git a/zhipuai/types/video/video_create_params.py b/zhipuai/types/video/video_create_params.py
new file mode 100644
index 0000000..ee0eb6a
--- /dev/null
+++ b/zhipuai/types/video/video_create_params.py
@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from typing import List, Optional
+
+from typing_extensions import Literal, Required, TypedDict
+
+__all__ = ["VideoCreateParams"]
+
+
+class VideoCreateParams(TypedDict, total=False):
+    model: str
+    """模型编码"""
+    prompt: str
+    """所需视频的文本描述"""
+    request_id: str
+    """由用户端传参，需保证唯一性；用于区分每次请求的唯一标识，用户端不传时平台会默认生成。"""
diff --git a/zhipuai/types/video/video_object.py b/zhipuai/types/video/video_object.py
new file mode 100644
index 0000000..5c7006d
--- /dev/null
+++ b/zhipuai/types/video/video_object.py
@@ -0,0 +1,28 @@
+from typing import Optional, List
+
+from ...core import BaseModel
+
+__all__ = ["VideoObject", "VideoResult"]
+
+
+class VideoResult(BaseModel):
+    url: str
+    """视频url"""
+
+
+class VideoObject(BaseModel):
+    id: Optional[str] = None
+    """智谱 AI 开放平台生成的任务订单号，调用请求结果接口时请使用此订单号"""
+
+    model: str
+    """模型名称"""
+
+    video_result: List[VideoResult]
+    """视频生成结果"""
+
+    task_status: str
+    """处理状态，PROCESSING（处理中），SUCCESS（成功），FAIL（失败）
+    注：处理中状态需通过查询获取结果"""
+
+    request_id: str
+    """用户在客户端请求时提交的任务编号或者平台生成的任务编号"""