diff --git a/examples/file_parsing_example.py b/examples/file_parsing_example.py new file mode 100644 index 0000000..bb4285f --- /dev/null +++ b/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() diff --git a/pyproject.toml b/pyproject.toml index cfdac2b..7bdde89 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/zai/_client.py b/src/zai/_client.py index 93e3c8c..fa5c51a 100644 --- a/src/zai/_client.py +++ b/src/zai/_client.py @@ -9,231 +9,237 @@ from typing_extensions import override if TYPE_CHECKING: - from zai.api_resource.agents import Agents - from zai.api_resource.assistant import Assistant - from zai.api_resource.audio import Audio - from zai.api_resource.batch import Batches - from zai.api_resource.chat import Chat - from zai.api_resource.embeddings import Embeddings - from zai.api_resource.files import Files - from zai.api_resource.images import Images - from zai.api_resource.moderations import Moderations - from zai.api_resource.tools import Tools - 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.agents import Agents + from zai.api_resource.assistant import Assistant + from zai.api_resource.audio import Audio + from zai.api_resource.batch import Batches + from zai.api_resource.chat import Chat + from zai.api_resource.embeddings import Embeddings + from zai.api_resource.files import Files + from zai.api_resource.images import Images + from zai.api_resource.moderations import Moderations + from zai.api_resource.tools import Tools + 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, - ZAI_DEFAULT_MAX_RETRIES, - HttpClient, - NotGiven, - ZaiError, - _jwt_token, + NOT_GIVEN, + ZAI_DEFAULT_MAX_RETRIES, + HttpClient, + NotGiven, + ZaiError, + _jwt_token, ) class BaseClient(HttpClient): - """ - Main client for interacting with the ZAI API - - Attributes: - chat (Chat): Chat completions API resource - api_key (str): API key for authentication - _disable_token_cache (bool): Whether to disable token caching - source_channel (str): Source channel identifier - """ - - chat: Chat - api_key: str - base_url: str - disable_token_cache: bool = True - source_channel: str - - def __init__( - self, - *, - api_key: str | None = None, - base_url: str | httpx.URL | None = None, - timeout: Union[float, Timeout, None, NotGiven] = NOT_GIVEN, - max_retries: int = ZAI_DEFAULT_MAX_RETRIES, - http_client: httpx.Client | None = None, - custom_headers: Mapping[str, str] | None = None, - disable_token_cache: bool = True, - _strict_response_validation: bool = False, - source_channel: str | None = None, - ) -> None: - """ - Initialize the ZAI client - - Arguments: - api_key (str | None): API key for authentication. - If None, will try to get from ZAI_API_KEY environment variable. - base_url (str | httpx.URL | None): Base URL for the API. - If None, will try to get from ZAI_BASE_URL environment variable - timeout (Union[float, Timeout, None, NotGiven]): Request timeout configuration - max_retries (int): Maximum number of retries for failed requests - http_client (httpx.Client | None): Custom HTTP client to use - custom_headers (Mapping[str, str] | None): Additional headers to include in requests - disable_token_cache (bool): Whether to disable JWT token caching - _strict_response_validation (bool): Whether to enable strict response validation - source_channel (str | None): Source channel identifier - """ - if api_key is None: - api_key = os.environ.get('ZAI_API_KEY') - if api_key is None: - raise ZaiError('api_key not provided, please provide it through parameters or environment variables') - self.api_key = api_key - self.source_channel = source_channel - self.disable_token_cache = disable_token_cache - - if base_url is None: - base_url = os.environ.get('ZAI_BASE_URL') - if base_url is None: - base_url = self.default_base_url - self.base_url = base_url - - from ._version import __version__ - - super().__init__( - version=__version__, - base_url=base_url, - max_retries=max_retries, - timeout=timeout, - custom_httpx_client=http_client, - custom_headers=custom_headers, - _strict_response_validation=_strict_response_validation, - ) - - @property - def default_base_url(self): - raise NotImplementedError('Subclasses must define default_base_url') - - @cached_property - def chat(self) -> Chat: - from zai.api_resource.chat import Chat - - return Chat(self) - - @cached_property - def assistant(self) -> Assistant: - from zai.api_resource.assistant import Assistant - - return Assistant(self) - - @cached_property - def agents(self) -> Agents: - from zai.api_resource.agents import Agents - - return Agents(self) - - @cached_property - def embeddings(self) -> Embeddings: - from zai.api_resource.embeddings import Embeddings - - return Embeddings(self) - - @cached_property - def batches(self) -> Batches: - from zai.api_resource.batch import Batches - - return Batches(self) - - @cached_property - def tools(self) -> Tools: - from zai.api_resource.tools import Tools - - return Tools(self) - - @cached_property - def web_search(self) -> WebSearchApi: - from zai.api_resource.web_search import WebSearchApi - - return WebSearchApi(self) - - @cached_property - def files(self) -> Files: - from zai.api_resource.files import Files - - return Files(self) - - @cached_property - def images(self) -> Images: - from zai.api_resource.images import Images - - return Images(self) - - @cached_property - def audio(self) -> Audio: - from zai.api_resource.audio import Audio - - return Audio(self) - - @cached_property - def videos(self) -> Videos: - from zai.api_resource.videos import Videos - - return Videos(self) + """ + Main client for interacting with the ZAI API + + Attributes: + chat (Chat): Chat completions API resource + api_key (str): API key for authentication + _disable_token_cache (bool): Whether to disable token caching + source_channel (str): Source channel identifier + """ + + chat: Chat + api_key: str + base_url: str + disable_token_cache: bool = True + source_channel: str + + def __init__( + self, + *, + api_key: str | None = None, + base_url: str | httpx.URL | None = None, + timeout: Union[float, Timeout, None, NotGiven] = NOT_GIVEN, + max_retries: int = ZAI_DEFAULT_MAX_RETRIES, + http_client: httpx.Client | None = None, + custom_headers: Mapping[str, str] | None = None, + disable_token_cache: bool = True, + _strict_response_validation: bool = False, + source_channel: str | None = None, + ) -> None: + """ + Initialize the ZAI client + + Arguments: + api_key (str | None): API key for authentication. + If None, will try to get from ZAI_API_KEY environment variable. + base_url (str | httpx.URL | None): Base URL for the API. + If None, will try to get from ZAI_BASE_URL environment variable + timeout (Union[float, Timeout, None, NotGiven]): Request timeout configuration + max_retries (int): Maximum number of retries for failed requests + http_client (httpx.Client | None): Custom HTTP client to use + custom_headers (Mapping[str, str] | None): Additional headers to include in requests + disable_token_cache (bool): Whether to disable JWT token caching + _strict_response_validation (bool): Whether to enable strict response validation + source_channel (str | None): Source channel identifier + """ + if api_key is None: + api_key = os.environ.get('ZAI_API_KEY') + if api_key is None: + raise ZaiError('api_key not provided, please provide it through parameters or environment variables') + self.api_key = api_key + self.source_channel = source_channel + self.disable_token_cache = disable_token_cache + + if base_url is None: + base_url = os.environ.get('ZAI_BASE_URL') + if base_url is None: + base_url = self.default_base_url + self.base_url = base_url + + from ._version import __version__ + + super().__init__( + version=__version__, + base_url=base_url, + max_retries=max_retries, + timeout=timeout, + custom_httpx_client=http_client, + custom_headers=custom_headers, + _strict_response_validation=_strict_response_validation, + ) + + @property + def default_base_url(self): + raise NotImplementedError('Subclasses must define default_base_url') + + @cached_property + def chat(self) -> Chat: + from zai.api_resource.chat import Chat + + return Chat(self) + + @cached_property + def assistant(self) -> Assistant: + from zai.api_resource.assistant import Assistant + + return Assistant(self) + + @cached_property + def agents(self) -> Agents: + from zai.api_resource.agents import Agents + + return Agents(self) + + @cached_property + def embeddings(self) -> Embeddings: + from zai.api_resource.embeddings import Embeddings + + return Embeddings(self) + + @cached_property + def batches(self) -> Batches: + from zai.api_resource.batch import Batches + + return Batches(self) + + @cached_property + def tools(self) -> Tools: + from zai.api_resource.tools import Tools + + return Tools(self) + + @cached_property + def web_search(self) -> WebSearchApi: + from zai.api_resource.web_search import WebSearchApi + + return WebSearchApi(self) + + @cached_property + def files(self) -> Files: + from zai.api_resource.files import Files + + return Files(self) + + @cached_property + def images(self) -> Images: + from zai.api_resource.images import Images + + return Images(self) + + @cached_property + def audio(self) -> Audio: + from zai.api_resource.audio import Audio + + return Audio(self) + + @cached_property + def videos(self) -> Videos: + from zai.api_resource.videos import Videos + + return Videos(self) - @cached_property - def moderations(self) -> Moderations: - from zai.api_resource.moderations import Moderations + @cached_property + def moderations(self) -> Moderations: + from zai.api_resource.moderations import Moderations - return Moderations(self) + return Moderations(self) - @cached_property - def voice(self) -> Voice: - from zai.api_resource.voice import Voice + @cached_property + def voice(self) -> Voice: + from zai.api_resource.voice import Voice - return Voice(self) + return Voice(self) - @property - @override - def auth_headers(self) -> dict[str, str]: - 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, - } + @cached_property + def file_parser(self) -> FileParser: + from zai.api_resource.file_parser import FileParser + return FileParser(self) - def __del__(self) -> None: - 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 + @property + @override + def auth_headers(self) -> dict[str, str]: + 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: - return + def __del__(self) -> None: + 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: - # Check if client is still valid before closing - if hasattr(self, '_client') and self._client is not None: - self.close() - except Exception: - # Ignore any exceptions during cleanup to avoid masking the original error - pass + if self._has_custom_http_client: + return + + try: + # Check if client is still valid before closing + if hasattr(self, '_client') and self._client is not None: + self.close() + except Exception: + # Ignore any exceptions during cleanup to avoid masking the original error + pass class ZaiClient(BaseClient): - @property - def default_base_url(self): - return 'https://api.z.ai/api/paas/v4' + @property + def default_base_url(self): + return 'https://api.z.ai/api/paas/v4' - @property - @override - def auth_headers(self) -> dict[str, str]: - headers = super().auth_headers - headers['Accept-Language'] = 'en-US,en' - return headers + @property + @override + def auth_headers(self) -> dict[str, str]: + headers = super().auth_headers + headers['Accept-Language'] = 'en-US,en' + return headers class ZhipuAiClient(BaseClient): - @property - def default_base_url(self): - return 'https://open.bigmodel.cn/api/paas/v4' + @property + def default_base_url(self): + return 'https://open.bigmodel.cn/api/paas/v4' diff --git a/src/zai/_version.py b/src/zai/_version.py index 6f1fe34..55b66c6 100644 --- a/src/zai/_version.py +++ b/src/zai/_version.py @@ -1,2 +1,2 @@ __title__ = 'Z.ai' -__version__ = '0.0.4' +__version__ = '0.0.4.1' diff --git a/src/zai/api_resource/__init__.py b/src/zai/api_resource/__init__.py index a49a808..616b46a 100644 --- a/src/zai/api_resource/__init__.py +++ b/src/zai/api_resource/__init__.py @@ -18,6 +18,8 @@ Videos, ) from .web_search import WebSearchApi +from .file_parser import FileParser + __all__ = [ 'Videos', @@ -35,4 +37,5 @@ 'Moderations', 'WebSearchApi', 'Agents', + 'FileParser', ] diff --git a/src/zai/api_resource/file_parser/__init__.py b/src/zai/api_resource/file_parser/__init__.py new file mode 100644 index 0000000..b263267 --- /dev/null +++ b/src/zai/api_resource/file_parser/__init__.py @@ -0,0 +1,3 @@ +from .file_parser import FileParser + +__all__ = ['FileParser'] \ No newline at end of file diff --git a/src/zai/api_resource/file_parser/file_parser.py b/src/zai/api_resource/file_parser/file_parser.py new file mode 100644 index 0000000..2238e72 --- /dev/null +++ b/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 diff --git a/src/zai/types/file_parser/__init__.py b/src/zai/types/file_parser/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/zai/types/file_parser/file_parser_create_params.py b/src/zai/types/file_parser/file_parser_create_params.py new file mode 100644 index 0000000..73da78b --- /dev/null +++ b/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""" diff --git a/src/zai/types/file_parser/file_parser_resp.py b/src/zai/types/file_parser/file_parser_resp.py new file mode 100644 index 0000000..6a753d5 --- /dev/null +++ b/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