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

@@ -9,231 +9,237 @@ from httpx import Timeout
 from typing_extensions import override
 
 if TYPE_CHECKING:
-	from zai.api_resource.agents import Agents
+    from zai.api_resource.agents import Agents
-	from zai.api_resource.assistant import Assistant
+    from zai.api_resource.assistant import Assistant
-	from zai.api_resource.audio import Audio
+    from zai.api_resource.audio import Audio
-	from zai.api_resource.batch import Batches
+    from zai.api_resource.batch import Batches
-	from zai.api_resource.chat import Chat
+    from zai.api_resource.chat import Chat
-	from zai.api_resource.embeddings import Embeddings
+    from zai.api_resource.embeddings import Embeddings
-	from zai.api_resource.files import Files
+    from zai.api_resource.files import Files
-	from zai.api_resource.images import Images
+    from zai.api_resource.images import Images
-	from zai.api_resource.moderations import Moderations
+    from zai.api_resource.moderations import Moderations
-	from zai.api_resource.tools import Tools
+    from zai.api_resource.tools import Tools
-	from zai.api_resource.videos import Videos
+    from zai.api_resource.videos import Videos
-	from zai.api_resource.voice import Voice
+    from zai.api_resource.voice import Voice
-	from zai.api_resource.web_search import WebSearchApi
+    from zai.api_resource.web_search import WebSearchApi
+    from zai.api_resource.file_parser import FileParser
 
 from .core import (
-	NOT_GIVEN,
+    NOT_GIVEN,
-	ZAI_DEFAULT_MAX_RETRIES,
+    ZAI_DEFAULT_MAX_RETRIES,
-	HttpClient,
+    HttpClient,
-	NotGiven,
+    NotGiven,
-	ZaiError,
+    ZaiError,
-	_jwt_token,
+    _jwt_token,
 )
 
 
 class BaseClient(HttpClient):
-	"""
+    """
-	Main client for interacting with the ZAI API
+    Main client for interacting with the ZAI API
 
-	Attributes:
+    Attributes:
-		chat (Chat): Chat completions API resource
+        chat (Chat): Chat completions API resource
-		api_key (str): API key for authentication
+        api_key (str): API key for authentication
-		_disable_token_cache (bool): Whether to disable token caching
+        _disable_token_cache (bool): Whether to disable token caching
-		source_channel (str): Source channel identifier
+        source_channel (str): Source channel identifier
-	"""
+    """
 
-	chat: Chat
+    chat: Chat
-	api_key: str
+    api_key: str
-	base_url: str
+    base_url: str
-	disable_token_cache: bool = True
+    disable_token_cache: bool = True
-	source_channel: str
+    source_channel: str
 
-	def __init__(
+    def __init__(
-		self,
+            self,
-		*,
+            *,
-		api_key: str | None = None,
+            api_key: str | None = None,
-		base_url: str | httpx.URL | None = None,
+            base_url: str | httpx.URL | None = None,
-		timeout: Union[float, Timeout, None, NotGiven] = NOT_GIVEN,
+            timeout: Union[float, Timeout, None, NotGiven] = NOT_GIVEN,
-		max_retries: int = ZAI_DEFAULT_MAX_RETRIES,
+            max_retries: int = ZAI_DEFAULT_MAX_RETRIES,
-		http_client: httpx.Client | None = None,
+            http_client: httpx.Client | None = None,
-		custom_headers: Mapping[str, str] | None = None,
+            custom_headers: Mapping[str, str] | None = None,
-		disable_token_cache: bool = True,
+            disable_token_cache: bool = True,
-		_strict_response_validation: bool = False,
+            _strict_response_validation: bool = False,
-		source_channel: str | None = None,
+            source_channel: str | None = None,
-	) -> None:
+    ) -> None:
-		"""
+        """
-		Initialize the ZAI client
+        Initialize the ZAI client
 
-		Arguments:
+        Arguments:
-			api_key (str | None): API key for authentication.
+            api_key (str | None): API key for authentication.
-									If None, will try to get from ZAI_API_KEY environment variable.
+                                    If None, will try to get from ZAI_API_KEY environment variable.
-			base_url (str | httpx.URL | None): Base URL for the API.
+            base_url (str | httpx.URL | None): Base URL for the API.
-									If None, will try to get from ZAI_BASE_URL environment variable
+                                    If None, will try to get from ZAI_BASE_URL environment variable
-			timeout (Union[float, Timeout, None, NotGiven]): Request timeout configuration
+            timeout (Union[float, Timeout, None, NotGiven]): Request timeout configuration
-			max_retries (int): Maximum number of retries for failed requests
+            max_retries (int): Maximum number of retries for failed requests
-			http_client (httpx.Client | None): Custom HTTP client to use
+            http_client (httpx.Client | None): Custom HTTP client to use
-			custom_headers (Mapping[str, str] | None): Additional headers to include in requests
+            custom_headers (Mapping[str, str] | None): Additional headers to include in requests
-			disable_token_cache (bool): Whether to disable JWT token caching
+            disable_token_cache (bool): Whether to disable JWT token caching
-			_strict_response_validation (bool): Whether to enable strict response validation
+            _strict_response_validation (bool): Whether to enable strict response validation
-			source_channel (str | None): Source channel identifier
+            source_channel (str | None): Source channel identifier
-		"""
+        """
-		if api_key is None:
+        if api_key is None:
-			api_key = os.environ.get('ZAI_API_KEY')
+            api_key = os.environ.get('ZAI_API_KEY')
-		if api_key is None:
+        if api_key is None:
-			raise ZaiError('api_key not provided, please provide it through parameters or environment variables')
+            raise ZaiError('api_key not provided, please provide it through parameters or environment variables')
-		self.api_key = api_key
+        self.api_key = api_key
-		self.source_channel = source_channel
+        self.source_channel = source_channel
-		self.disable_token_cache = disable_token_cache
+        self.disable_token_cache = disable_token_cache
 
-		if base_url is None:
+        if base_url is None:
-			base_url = os.environ.get('ZAI_BASE_URL')
+            base_url = os.environ.get('ZAI_BASE_URL')
-		if base_url is None:
+        if base_url is None:
-			base_url = self.default_base_url
+            base_url = self.default_base_url
-		self.base_url = base_url
+        self.base_url = base_url
 
-		from ._version import __version__
+        from ._version import __version__
 
-		super().__init__(
+        super().__init__(
-			version=__version__,
+            version=__version__,
-			base_url=base_url,
+            base_url=base_url,
-			max_retries=max_retries,
+            max_retries=max_retries,
-			timeout=timeout,
+            timeout=timeout,
-			custom_httpx_client=http_client,
+            custom_httpx_client=http_client,
-			custom_headers=custom_headers,
+            custom_headers=custom_headers,
-			_strict_response_validation=_strict_response_validation,
+            _strict_response_validation=_strict_response_validation,
-		)
+        )
 
-	@property
+    @property
-	def default_base_url(self):
+    def default_base_url(self):
-		raise NotImplementedError('Subclasses must define default_base_url')
+        raise NotImplementedError('Subclasses must define default_base_url')
 
-	@cached_property
+    @cached_property
-	def chat(self) -> Chat:
+    def chat(self) -> Chat:
-		from zai.api_resource.chat import Chat
+        from zai.api_resource.chat import Chat
 
-		return Chat(self)
+        return Chat(self)
 
-	@cached_property
+    @cached_property
-	def assistant(self) -> Assistant:
+    def assistant(self) -> Assistant:
-		from zai.api_resource.assistant import Assistant
+        from zai.api_resource.assistant import Assistant
 
-		return Assistant(self)
+        return Assistant(self)
 
-	@cached_property
+    @cached_property
-	def agents(self) -> Agents:
+    def agents(self) -> Agents:
-		from zai.api_resource.agents import Agents
+        from zai.api_resource.agents import Agents
 
-		return Agents(self)
+        return Agents(self)
 
-	@cached_property
+    @cached_property
-	def embeddings(self) -> Embeddings:
+    def embeddings(self) -> Embeddings:
-		from zai.api_resource.embeddings import Embeddings
+        from zai.api_resource.embeddings import Embeddings
 
-		return Embeddings(self)
+        return Embeddings(self)
 
-	@cached_property
+    @cached_property
-	def batches(self) -> Batches:
+    def batches(self) -> Batches:
-		from zai.api_resource.batch import Batches
+        from zai.api_resource.batch import Batches
 
-		return Batches(self)
+        return Batches(self)
 
-	@cached_property
+    @cached_property
-	def tools(self) -> Tools:
+    def tools(self) -> Tools:
-		from zai.api_resource.tools import Tools
+        from zai.api_resource.tools import Tools
 
-		return Tools(self)
+        return Tools(self)
 
-	@cached_property
+    @cached_property
-	def web_search(self) -> WebSearchApi:
+    def web_search(self) -> WebSearchApi:
-		from zai.api_resource.web_search import WebSearchApi
+        from zai.api_resource.web_search import WebSearchApi
 
-		return WebSearchApi(self)
+        return WebSearchApi(self)
 
-	@cached_property
+    @cached_property
-	def files(self) -> Files:
+    def files(self) -> Files:
-		from zai.api_resource.files import Files
+        from zai.api_resource.files import Files
 
-		return Files(self)
+        return Files(self)
 
-	@cached_property
+    @cached_property
-	def images(self) -> Images:
+    def images(self) -> Images:
-		from zai.api_resource.images import Images
+        from zai.api_resource.images import Images
 
-		return Images(self)
+        return Images(self)
 
-	@cached_property
+    @cached_property
-	def audio(self) -> Audio:
+    def audio(self) -> Audio:
-		from zai.api_resource.audio import Audio
+        from zai.api_resource.audio import Audio
 
-		return Audio(self)
+        return Audio(self)
 
-	@cached_property
+    @cached_property
-	def videos(self) -> Videos:
+    def videos(self) -> Videos:
-		from zai.api_resource.videos import Videos
+        from zai.api_resource.videos import Videos
 
-		return Videos(self)
+        return Videos(self)
 
-	@cached_property
+    @cached_property
-	def moderations(self) -> Moderations:
+    def moderations(self) -> Moderations:
-		from zai.api_resource.moderations import Moderations
+        from zai.api_resource.moderations import Moderations
 
-		return Moderations(self)
+        return Moderations(self)
 
-	@cached_property
+    @cached_property
-	def voice(self) -> Voice:
+    def voice(self) -> Voice:
-		from zai.api_resource.voice import Voice
+        from zai.api_resource.voice import Voice
 
-		return Voice(self)
+        return Voice(self)
 
-	@property
+    @cached_property
-	@override
+    def file_parser(self) -> FileParser:
-	def auth_headers(self) -> dict[str, str]:
+        from zai.api_resource.file_parser import FileParser
-		api_key = self.api_key
+        return FileParser(self)
-		source_channel = self.source_channel or 'python-sdk'
-		if self.disable_token_cache:
-			return {
-				'Authorization': f'Bearer {api_key}',
-				'x-source-channel': source_channel,
-			}
-		else:
-			return {
-				'Authorization': f'Bearer {_jwt_token.generate_token(api_key)}',
-				'x-source-channel': source_channel,
-			}
 
-	def __del__(self) -> None:
+    @property
-		if not hasattr(self, '_has_custom_http_client') or not hasattr(self, 'close') or not hasattr(self, '_client'):
+    @override
-			# if the '__init__' method raised an error, self would not have client attr
+    def auth_headers(self) -> dict[str, str]:
-			return
+        api_key = self.api_key
+        source_channel = self.source_channel or 'python-sdk'
+        if self.disable_token_cache:
+            return {
+                'Authorization': f'Bearer {api_key}',
+                 'x-source-channel': source_channel,
+            }
+        else:
+            return {
+                'Authorization': f'Bearer {_jwt_token.generate_token(api_key)}',
+                'x-source-channel': source_channel,
+            }
 
-		if self._has_custom_http_client:
+    def __del__(self) -> None:
-			return
+        if not hasattr(self, '_has_custom_http_client') or not hasattr(self, 'close') or not hasattr(self, '_client'):
+            # if the '__init__' method raised an error, self would not have client attr
+            return
 
-		try:
+        if self._has_custom_http_client:
-			# Check if client is still valid before closing
+            return
-			if hasattr(self, '_client') and self._client is not None:
+
-				self.close()
+        try:
-		except Exception:
+            # Check if client is still valid before closing
-			# Ignore any exceptions during cleanup to avoid masking the original error
+            if hasattr(self, '_client') and self._client is not None:
-			pass
+                self.close()
+        except Exception:
+            # Ignore any exceptions during cleanup to avoid masking the original error
+            pass
 
 
 class ZaiClient(BaseClient):
-	@property
+    @property
-	def default_base_url(self):
+    def default_base_url(self):
-		return 'https://api.z.ai/api/paas/v4'
+        return 'https://api.z.ai/api/paas/v4'
 
-	@property
+    @property
-	@override
+    @override
-	def auth_headers(self) -> dict[str, str]:
+    def auth_headers(self) -> dict[str, str]:
-		headers = super().auth_headers
+        headers = super().auth_headers
-		headers['Accept-Language'] = 'en-US,en'
+        headers['Accept-Language'] = 'en-US,en'
-		return headers
+        return headers
 
 
 class ZhipuAiClient(BaseClient):
-	@property
+    @property
-	def default_base_url(self):
+    def default_base_url(self):
-		return 'https://open.bigmodel.cn/api/paas/v4'
+        return 'https://open.bigmodel.cn/api/paas/v4'

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