Spaces:
Running
Running
| # Copyright (c) 2026 PaddlePaddle Authors. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import json | |
| import os | |
| from typing import Any, Dict, Optional | |
| import aiohttp | |
| from ._core import ( | |
| extract_api_message_from_payload, | |
| extract_job_id, | |
| raise_for_status, | |
| unwrap_api_response, | |
| ) | |
| from .errors import ( | |
| NetworkError, | |
| RequestTimeoutError, | |
| ResponseFormatError, | |
| ResultParseError, | |
| ) | |
| from ._http import API_PATH, DEFAULT_BASE_URL | |
| class AsyncHTTPClient: | |
| def __init__( | |
| self, | |
| token: str, | |
| base_url: str, | |
| timeout: float, | |
| client_platform: Optional[str] = None, | |
| ): | |
| self._token = token | |
| self._base_url = base_url.rstrip("/") | |
| self._jobs_url = f"{self._base_url}{API_PATH}" | |
| self._timeout = timeout | |
| self._client_platform = client_platform | |
| self._session = None | |
| def timeout(self) -> float: | |
| return self._timeout | |
| async def __aenter__(self): | |
| await self._ensure_session() | |
| return self | |
| async def __aexit__(self, *args): | |
| await self.close() | |
| async def _ensure_session(self): | |
| if self._session is None: | |
| self._session = aiohttp.ClientSession( | |
| headers=self._api_headers(), | |
| timeout=aiohttp.ClientTimeout(total=self._timeout), | |
| ) | |
| async def close(self): | |
| if self._session: | |
| await self._session.close() | |
| self._session = None | |
| def _api_headers(self) -> dict: | |
| headers = {"Authorization": f"Bearer {self._token}"} | |
| if self._client_platform: | |
| headers["Client-Platform"] = self._client_platform | |
| return headers | |
| async def submit_url( | |
| self, | |
| model: str, | |
| file_url: str, | |
| optional_payload: dict, | |
| page_ranges: Optional[str] = None, | |
| batch_id: Optional[str] = None, | |
| ) -> str: | |
| body = { | |
| "fileUrl": file_url, | |
| "model": model, | |
| "optionalPayload": optional_payload, | |
| } | |
| if page_ranges is not None: | |
| body["pageRanges"] = page_ranges | |
| if batch_id is not None: | |
| body["batchId"] = batch_id | |
| await self._ensure_session() | |
| async with self._session.post( | |
| self._jobs_url, | |
| json=body, | |
| headers={"Content-Type": "application/json"}, | |
| ) as resp: | |
| await self._raise_for_response(resp) | |
| data = await self._response_data(resp) | |
| return extract_job_id(data) | |
| async def submit_file( | |
| self, | |
| model: str, | |
| file_path: str, | |
| optional_payload: dict, | |
| page_ranges: Optional[str] = None, | |
| batch_id: Optional[str] = None, | |
| ) -> str: | |
| if not os.path.exists(file_path): | |
| raise FileNotFoundError(file_path) | |
| form = aiohttp.FormData() | |
| form.add_field("model", model) | |
| form.add_field("optionalPayload", json.dumps(optional_payload)) | |
| if page_ranges is not None: | |
| form.add_field("pageRanges", page_ranges) | |
| if batch_id is not None: | |
| form.add_field("batchId", batch_id) | |
| with open(file_path, "rb") as f: | |
| file_data = f.read() | |
| form.add_field( | |
| "file", | |
| file_data, | |
| filename=os.path.basename(file_path), | |
| ) | |
| await self._ensure_session() | |
| async with self._session.post(self._jobs_url, data=form) as resp: | |
| await self._raise_for_response(resp) | |
| data = await self._response_data(resp) | |
| return extract_job_id(data) | |
| async def get_job_status(self, job_id: str) -> Dict[str, Any]: | |
| await self._ensure_session() | |
| async with self._session.get(f"{self._jobs_url}/{job_id}") as resp: | |
| await self._raise_for_response(resp) | |
| return await self._response_data(resp) | |
| async def get_batch_status(self, batch_id: str) -> Dict[str, Any]: | |
| await self._ensure_session() | |
| async with self._session.get(f"{self._jobs_url}/batch/{batch_id}") as resp: | |
| await self._raise_for_response(resp) | |
| return await self._response_data(resp) | |
| async def fetch_jsonl(self, url: str) -> list: | |
| timeout = aiohttp.ClientTimeout(total=self._timeout) | |
| async with aiohttp.ClientSession(timeout=timeout) as bare_session: | |
| async with bare_session.get(url) as resp: | |
| await self._raise_for_response(resp) | |
| text = await resp.text() | |
| try: | |
| lines = text.strip().split("\n") | |
| return [json.loads(line) for line in lines if line.strip()] | |
| except json.JSONDecodeError as e: | |
| raise ResultParseError( | |
| f"Malformed JSONL result payload: {e}" | |
| ) from e | |
| async def _raise_for_response(self, resp) -> None: | |
| if 200 <= resp.status < 300: | |
| return | |
| try: | |
| body = await resp.json() | |
| msg = ( | |
| extract_api_message_from_payload(body) | |
| if isinstance(body, dict) | |
| else None | |
| ) | |
| if not msg: | |
| msg = await resp.text() | |
| except Exception: | |
| msg = await resp.text() | |
| raise_for_status(resp.status, msg) | |
| async def _response_data(self, resp) -> Dict[str, Any]: | |
| try: | |
| payload = await resp.json() | |
| except Exception as e: | |
| raise ResponseFormatError(f"Response body is not valid JSON: {e}") from e | |
| return unwrap_api_response(payload, resp.status) | |