feat: support file parsing function (#42)

Co-authored-by: mengqian <cherish_a_meng@163.com>

examples/file_parsing_example.py

@@ -0,0 +1,81 @@
+from zai import ZaiClient
+import time
+import traceback
+
+client = ZaiClient(
+    base_url="",
+    api_key=""
+)
+
+
+def file_parser_create_example(file_path, tool_type, file_type):
+    """
+    Example: Create a file parsing task
+    """
+    print("=== File Parser Create Example ===")
+    with open(file_path, 'rb') as f:
+        print("Submitting file parsing task ...")
+        response = client.file_parser.create(
+            file=f,
+            file_type=file_type,
+            tool_type=tool_type,
+        )
+        print("Task created successfully. Response:")
+        print(response)
+        # Usually you can get task_id
+        task_id = getattr(response, "task_id", None)
+        return task_id
+
+
+def file_parser_content_example(task_id, format_type="download_link"):
+    """
+    Example: Get file parsing result
+    """
+    print("=== File Parser Content Example ===")
+    try:
+        print(f"Querying parsing result for task_id: {task_id}")
+        response = client.file_parser.content(
+            task_id=task_id,
+            format_type=format_type
+        )
+        return response
+    except Exception as err:
+        print("Failed to get parsing result:", traceback.format_exc())
+        return None
+
+
+def file_parser_complete_example():
+    """
+    Full Example: Submit file for parsing, then poll until result is ready
+    """
+    # 1. Create parsing task
+    # Please modify the local file path
+    file_path = 'your file path'
+    task_id = file_parser_create_example(file_path=file_path, tool_type="lite", file_type="pdf")
+    if not task_id:
+        print("Could not submit file for parsing.")
+        return
+
+    # 2. Poll to get the result
+    max_wait = 60  # Wait up to 1 minute
+    wait_time = 0
+    while wait_time < max_wait:
+        print(f"Waiting {wait_time}/{max_wait} seconds before querying result...")
+        # format_type = text / download_link
+        response = file_parser_content_example(task_id=task_id, format_type="download_link")
+
+        result = response.json()
+        if result.get("status") == "processing":
+            print(result)
+
+            time.sleep(5)
+            wait_time += 5
+        else:
+            print(result)
+            break
+    print("File parser demo completed.")
+
+
+if __name__ == "__main__":
+    print("=== File Parsing Quick Demo ===\n")
+    file_parser_complete_example()

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zai-sdk"
-version = "0.0.4"
+version = "0.0.4.1"
 description = "A SDK library for accessing big model apis from Z.ai"
 authors = ["Z.ai"]
 readme = "README.md"

src/zai/_client.py

@@ -22,6 +22,7 @@ if TYPE_CHECKING:
     from zai.api_resource.videos import Videos
     from zai.api_resource.voice import Voice
     from zai.api_resource.web_search import WebSearchApi
+    from zai.api_resource.file_parser import FileParser
 
 from .core import (
     NOT_GIVEN,
@@ -187,6 +188,11 @@ class BaseClient(HttpClient):
 
         return Voice(self)
 
+    @cached_property
+    def file_parser(self) -> FileParser:
+        from zai.api_resource.file_parser import FileParser
+        return FileParser(self)
+
     @property
     @override
     def auth_headers(self) -> dict[str, str]:

src/zai/_version.py

@@ -1,2 +1,2 @@
 __title__ = 'Z.ai'
-__version__ = '0.0.4'
+__version__ = '0.0.4.1'

src/zai/api_resource/__init__.py

@@ -18,6 +18,8 @@ from .videos import (
 	Videos,
 )
 from .web_search import WebSearchApi
+from .file_parser import FileParser
+
 
 __all__ = [
 	'Videos',
@@ -35,4 +37,5 @@ __all__ = [
 	'Moderations',
 	'WebSearchApi',
 	'Agents',
+	'FileParser',
 ]

src/zai/api_resource/file_parser/__init__.py

@@ -0,0 +1,3 @@
+from .file_parser import FileParser
+
+__all__ = ['FileParser']

src/zai/api_resource/file_parser/file_parser.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+import json
+
+from typing import TYPE_CHECKING, Mapping, cast
+
+import httpx
+from typing_extensions import Literal
+
+from zai.core import (
+    BaseAPI,
+    maybe_transform,
+    NOT_GIVEN,
+    Body,
+    Headers,
+    NotGiven,
+    FileTypes,
+    _legacy_binary_response,
+    _legacy_response,
+    deepcopy_minimal,
+    extract_files,
+    make_request_options
+)
+
+from zai.types.file_parser.file_parser_create_params import FileParserCreateParams
+from zai.types.file_parser.file_parser_resp import FileParserTaskCreateResp
+
+if TYPE_CHECKING:
+    from zai._client import ZaiClient
+
+__all__ = ["FileParser"]
+
+
+class FileParser(BaseAPI):
+
+    def __init__(self, client: "ZaiClient") -> None:
+        super().__init__(client)
+
+    def create(
+            self,
+            *,
+            file: FileTypes = None,
+            file_type: str = None,
+            tool_type: Literal["lite", "expert", "prime"],
+            extra_headers: Headers | None = None,
+            extra_body: Body | None = None,
+            timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+    ) -> FileParserTaskCreateResp:
+
+        if not file:
+            raise ValueError("At least one `file` must be provided.")
+        body = deepcopy_minimal(
+            {
+                "file": file,
+                "file_type": file_type,
+                "tool_type": tool_type,
+            }
+        )
+
+        files = extract_files(cast(Mapping[str, object], body), paths=[["file"]])
+        if files:
+            # It should be noted that the actual Content-Type header that will be
+            # sent to the server will contain a `boundary` parameter, e.g.
+            # multipart/form-data; boundary=---abc--
+            extra_headers = {"Content-Type": "multipart/form-data", **(extra_headers or {})}
+        return self._post(
+            "/files/parser/create",
+            body=maybe_transform(body, FileParserCreateParams),
+            files=files,
+            options=make_request_options(
+                extra_headers=extra_headers, extra_body=extra_body, timeout=timeout
+            ),
+            cast_type=FileParserTaskCreateResp,
+        )
+
+    def content(
+            self,
+            task_id: str,
+            *,
+            format_type: Literal["text", "download_link"],
+            extra_headers: Headers | None = None,
+            extra_body: Body | None = None,
+            timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN,
+    ) -> httpx.Response:
+        """
+        Returns the contents of the specified file.
+
+        Args:
+          extra_headers: Send extra headers
+
+          extra_body: Add additional JSON properties to the request
+
+          timeout: Override the client-level default timeout for this request, in seconds
+        """
+        if not task_id:
+            raise ValueError(f"Expected a non-empty value for `task_id` but received {task_id!r}")
+        extra_headers = {"Accept": "application/binary", **(extra_headers or {})}
+        httpxBinaryResponseContent = self._get(
+            f"/files/parser/result/{task_id}/{format_type}",
+            options=make_request_options(
+                extra_headers=extra_headers, extra_body=extra_body, timeout=timeout
+            ),
+            cast_type=_legacy_binary_response.HttpxBinaryResponseContent,
+        )
+        return httpxBinaryResponseContent.response

src/zai/types/file_parser/file_parser_create_params.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from typing_extensions import Literal, Required, TypedDict
+from zai.core import FileTypes
+
+__all__ = ["FileParserCreateParams", "FileParserDownloadParams"]
+
+
+class FileParserCreateParams(TypedDict):
+    file: FileTypes
+    """Uploaded file"""
+    file_type: str
+    """File type"""
+    tool_type: Literal["lite", "expert", "prime"]
+    """Tool type"""
+
+
+class FileParserDownloadParams(TypedDict):
+    task_id: str
+    """Parsing task id"""
+    format_type: Literal["text", "download_link"]
+    """Result return type"""

src/zai/types/file_parser/file_parser_resp.py

@@ -0,0 +1,16 @@
+from typing import List, Optional
+
+from zai.core import BaseModel
+
+__all__ = [
+    "FileParserTaskCreateResp"
+]
+
+
+class FileParserTaskCreateResp(BaseModel):
+    task_id: str
+    # Task ID
+    message: str
+    # Message
+    success: bool
+    # Whether successful