大活
This commit is contained in:
407
art-agent/backend/venv/Lib/site-packages/replicate/client.py
Normal file
407
art-agent/backend/venv/Lib/site-packages/replicate/client.py
Normal file
@@ -0,0 +1,407 @@
|
||||
import asyncio
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
from datetime import datetime
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
AsyncIterator,
|
||||
Dict,
|
||||
Iterable,
|
||||
Iterator,
|
||||
Mapping,
|
||||
Optional,
|
||||
Type,
|
||||
Union,
|
||||
)
|
||||
|
||||
import httpx
|
||||
from typing_extensions import Unpack
|
||||
|
||||
from replicate.__about__ import __version__
|
||||
from replicate.account import Accounts
|
||||
from replicate.collection import Collections
|
||||
from replicate.deployment import Deployments
|
||||
from replicate.exceptions import ReplicateError
|
||||
from replicate.file import Files
|
||||
from replicate.hardware import HardwareNamespace as Hardware
|
||||
from replicate.model import Models
|
||||
from replicate.prediction import Predictions
|
||||
from replicate.run import async_run, run
|
||||
from replicate.stream import async_stream, stream
|
||||
from replicate.training import Trainings
|
||||
from replicate.webhook import Webhooks
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from replicate.stream import ServerSentEvent
|
||||
|
||||
|
||||
class Client:
|
||||
"""A Replicate API client library"""
|
||||
|
||||
__client: Optional[httpx.Client] = None
|
||||
__async_client: Optional[httpx.AsyncClient] = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
api_token: Optional[str] = None,
|
||||
*,
|
||||
base_url: Optional[str] = None,
|
||||
timeout: Optional[httpx.Timeout] = None,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self._api_token = api_token
|
||||
self._base_url = base_url
|
||||
self._timeout = timeout
|
||||
self._client_kwargs = kwargs
|
||||
|
||||
self.poll_interval = float(os.environ.get("REPLICATE_POLL_INTERVAL", "0.5"))
|
||||
|
||||
@property
|
||||
def _client(self) -> httpx.Client:
|
||||
if not self.__client:
|
||||
self.__client = _build_httpx_client(
|
||||
httpx.Client,
|
||||
self._api_token,
|
||||
self._base_url,
|
||||
self._timeout,
|
||||
**self._client_kwargs,
|
||||
) # type: ignore[assignment]
|
||||
return self.__client # type: ignore[return-value]
|
||||
|
||||
@property
|
||||
def _async_client(self) -> httpx.AsyncClient:
|
||||
if not self.__async_client:
|
||||
self.__async_client = _build_httpx_client(
|
||||
httpx.AsyncClient,
|
||||
self._api_token,
|
||||
self._base_url,
|
||||
self._timeout,
|
||||
**self._client_kwargs,
|
||||
) # type: ignore[assignment]
|
||||
return self.__async_client # type: ignore[return-value]
|
||||
|
||||
def _request(self, method: str, path: str, **kwargs) -> httpx.Response:
|
||||
resp = self._client.request(method, path, **kwargs)
|
||||
_raise_for_status(resp)
|
||||
|
||||
return resp
|
||||
|
||||
async def _async_request(self, method: str, path: str, **kwargs) -> httpx.Response:
|
||||
resp = await self._async_client.request(method, path, **kwargs)
|
||||
_raise_for_status(resp)
|
||||
|
||||
return resp
|
||||
|
||||
@property
|
||||
def accounts(self) -> Accounts:
|
||||
"""
|
||||
Namespace for operations related to accounts.
|
||||
"""
|
||||
|
||||
return Accounts(client=self)
|
||||
|
||||
@property
|
||||
def collections(self) -> Collections:
|
||||
"""
|
||||
Namespace for operations related to collections of models.
|
||||
"""
|
||||
return Collections(client=self)
|
||||
|
||||
@property
|
||||
def deployments(self) -> Deployments:
|
||||
"""
|
||||
Namespace for operations related to deployments.
|
||||
"""
|
||||
return Deployments(client=self)
|
||||
|
||||
@property
|
||||
def files(self) -> Files:
|
||||
"""
|
||||
Namespace for operations related to files.
|
||||
"""
|
||||
return Files(client=self)
|
||||
|
||||
@property
|
||||
def hardware(self) -> Hardware:
|
||||
"""
|
||||
Namespace for operations related to hardware.
|
||||
"""
|
||||
return Hardware(client=self)
|
||||
|
||||
@property
|
||||
def models(self) -> Models:
|
||||
"""
|
||||
Namespace for operations related to models.
|
||||
"""
|
||||
return Models(client=self)
|
||||
|
||||
@property
|
||||
def predictions(self) -> Predictions:
|
||||
"""
|
||||
Namespace for operations related to predictions.
|
||||
"""
|
||||
return Predictions(client=self)
|
||||
|
||||
@property
|
||||
def trainings(self) -> Trainings:
|
||||
"""
|
||||
Namespace for operations related to trainings.
|
||||
"""
|
||||
return Trainings(client=self)
|
||||
|
||||
@property
|
||||
def webhooks(self) -> Webhooks:
|
||||
"""
|
||||
Namespace for operations related to webhooks.
|
||||
"""
|
||||
return Webhooks(client=self)
|
||||
|
||||
def run(
|
||||
self,
|
||||
ref: str,
|
||||
input: Optional[Dict[str, Any]] = None,
|
||||
*,
|
||||
use_file_output: Optional[bool] = True,
|
||||
**params: Unpack["Predictions.CreatePredictionParams"],
|
||||
) -> Union[Any, Iterator[Any]]: # noqa: ANN401
|
||||
"""
|
||||
Run a model and wait for its output.
|
||||
"""
|
||||
|
||||
return run(self, ref, input, use_file_output=use_file_output, **params)
|
||||
|
||||
async def async_run(
|
||||
self,
|
||||
ref: str,
|
||||
input: Optional[Dict[str, Any]] = None,
|
||||
*,
|
||||
use_file_output: Optional[bool] = True,
|
||||
**params: Unpack["Predictions.CreatePredictionParams"],
|
||||
) -> Union[Any, AsyncIterator[Any]]: # noqa: ANN401
|
||||
"""
|
||||
Run a model and wait for its output asynchronously.
|
||||
"""
|
||||
|
||||
return await async_run(
|
||||
self, ref, input, use_file_output=use_file_output, **params
|
||||
)
|
||||
|
||||
def stream(
|
||||
self,
|
||||
ref: str,
|
||||
*,
|
||||
input: Optional[Dict[str, Any]] = None,
|
||||
use_file_output: Optional[bool] = True,
|
||||
**params: Unpack["Predictions.CreatePredictionParams"],
|
||||
) -> Iterator["ServerSentEvent"]:
|
||||
"""
|
||||
Stream a model's output.
|
||||
"""
|
||||
|
||||
return stream(self, ref, input, use_file_output=use_file_output, **params)
|
||||
|
||||
async def async_stream(
|
||||
self,
|
||||
ref: str,
|
||||
input: Optional[Dict[str, Any]] = None,
|
||||
*,
|
||||
use_file_output: Optional[bool] = True,
|
||||
**params: Unpack["Predictions.CreatePredictionParams"],
|
||||
) -> AsyncIterator["ServerSentEvent"]:
|
||||
"""
|
||||
Stream a model's output asynchronously.
|
||||
"""
|
||||
|
||||
return async_stream(self, ref, input, use_file_output=use_file_output, **params)
|
||||
|
||||
|
||||
# Adapted from https://github.com/encode/httpx/issues/108#issuecomment-1132753155
|
||||
class RetryTransport(httpx.AsyncBaseTransport, httpx.BaseTransport):
|
||||
"""A custom HTTP transport that automatically retries requests using an exponential backoff strategy
|
||||
for specific HTTP status codes and request methods.
|
||||
"""
|
||||
|
||||
RETRYABLE_METHODS = frozenset(["HEAD", "GET", "PUT", "DELETE", "OPTIONS", "TRACE"])
|
||||
RETRYABLE_STATUS_CODES = frozenset(
|
||||
[
|
||||
429, # Too Many Requests
|
||||
503, # Service Unavailable
|
||||
504, # Gateway Timeout
|
||||
]
|
||||
)
|
||||
MAX_BACKOFF_WAIT = 60
|
||||
|
||||
def __init__( # pylint: disable=too-many-arguments
|
||||
self,
|
||||
wrapped_transport: Union[httpx.BaseTransport, httpx.AsyncBaseTransport],
|
||||
*,
|
||||
max_attempts: int = 10,
|
||||
max_backoff_wait: float = MAX_BACKOFF_WAIT,
|
||||
backoff_factor: float = 0.1,
|
||||
jitter_ratio: float = 0.1,
|
||||
retryable_methods: Optional[Iterable[str]] = None,
|
||||
retry_status_codes: Optional[Iterable[int]] = None,
|
||||
) -> None:
|
||||
self._wrapped_transport = wrapped_transport
|
||||
|
||||
if jitter_ratio < 0 or jitter_ratio > 0.5:
|
||||
raise ValueError(
|
||||
f"jitter ratio should be between 0 and 0.5, actual {jitter_ratio}"
|
||||
)
|
||||
|
||||
self.max_attempts = max_attempts
|
||||
self.backoff_factor = backoff_factor
|
||||
self.retryable_methods = (
|
||||
frozenset(retryable_methods)
|
||||
if retryable_methods
|
||||
else self.RETRYABLE_METHODS
|
||||
)
|
||||
self.retry_status_codes = (
|
||||
frozenset(retry_status_codes)
|
||||
if retry_status_codes
|
||||
else self.RETRYABLE_STATUS_CODES
|
||||
)
|
||||
self.jitter_ratio = jitter_ratio
|
||||
self.max_backoff_wait = max_backoff_wait
|
||||
|
||||
def _calculate_sleep(
|
||||
self, attempts_made: int, headers: Union[httpx.Headers, Mapping[str, str]]
|
||||
) -> float:
|
||||
retry_after_header = (headers.get("Retry-After") or "").strip()
|
||||
if retry_after_header:
|
||||
if retry_after_header.isdigit():
|
||||
return float(retry_after_header)
|
||||
|
||||
try:
|
||||
parsed_date = datetime.fromisoformat(retry_after_header).astimezone()
|
||||
diff = (parsed_date - datetime.now().astimezone()).total_seconds()
|
||||
if diff > 0:
|
||||
return min(diff, self.max_backoff_wait)
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
backoff = self.backoff_factor * (2 ** (attempts_made - 1))
|
||||
jitter = (backoff * self.jitter_ratio) * random.choice([1, -1]) # noqa: S311
|
||||
total_backoff = backoff + jitter
|
||||
return min(total_backoff, self.max_backoff_wait)
|
||||
|
||||
def handle_request(self, request: httpx.Request) -> httpx.Response:
|
||||
response = self._wrapped_transport.handle_request(request) # type: ignore
|
||||
|
||||
if request.method not in self.retryable_methods:
|
||||
return response
|
||||
|
||||
remaining_attempts = self.max_attempts - 1
|
||||
attempts_made = 1
|
||||
|
||||
while True:
|
||||
if (
|
||||
remaining_attempts < 1
|
||||
or response.status_code not in self.retry_status_codes
|
||||
):
|
||||
return response
|
||||
|
||||
response.close()
|
||||
|
||||
sleep_for = self._calculate_sleep(attempts_made, response.headers)
|
||||
time.sleep(sleep_for)
|
||||
|
||||
response = self._wrapped_transport.handle_request(request) # type: ignore
|
||||
|
||||
attempts_made += 1
|
||||
remaining_attempts -= 1
|
||||
|
||||
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
|
||||
response = await self._wrapped_transport.handle_async_request(request) # type: ignore
|
||||
|
||||
if request.method not in self.retryable_methods:
|
||||
return response
|
||||
|
||||
remaining_attempts = self.max_attempts - 1
|
||||
attempts_made = 1
|
||||
|
||||
while True:
|
||||
if (
|
||||
remaining_attempts < 1
|
||||
or response.status_code not in self.retry_status_codes
|
||||
):
|
||||
return response
|
||||
|
||||
await response.aclose()
|
||||
|
||||
sleep_for = self._calculate_sleep(attempts_made, response.headers)
|
||||
await asyncio.sleep(sleep_for)
|
||||
|
||||
response = await self._wrapped_transport.handle_async_request(request) # type: ignore
|
||||
|
||||
attempts_made += 1
|
||||
remaining_attempts -= 1
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self._wrapped_transport.aclose() # type: ignore
|
||||
|
||||
def close(self) -> None:
|
||||
self._wrapped_transport.close() # type: ignore
|
||||
|
||||
|
||||
def _get_api_token_from_environment() -> Optional[str]:
|
||||
"""Get API token from cog current scope if available, otherwise from environment."""
|
||||
try:
|
||||
import cog # noqa: I001 # pyright: ignore [reportMissingImports]
|
||||
|
||||
for key, value in cog.current_scope().context.items():
|
||||
if key.upper() == "REPLICATE_API_TOKEN":
|
||||
return value
|
||||
except: # noqa: S110,E722,BLE001 we don't want this code to cause clients to fail
|
||||
pass
|
||||
|
||||
return os.environ.get("REPLICATE_API_TOKEN")
|
||||
|
||||
|
||||
def _build_httpx_client(
|
||||
client_type: Type[Union[httpx.Client, httpx.AsyncClient]],
|
||||
api_token: Optional[str] = None,
|
||||
base_url: Optional[str] = None,
|
||||
timeout: Optional[httpx.Timeout] = None,
|
||||
**kwargs,
|
||||
) -> Union[httpx.Client, httpx.AsyncClient]:
|
||||
headers = kwargs.pop("headers", {})
|
||||
if "User-Agent" not in headers:
|
||||
headers["User-Agent"] = f"replicate-python/{__version__}"
|
||||
if "Authorization" not in headers and (
|
||||
api_token := api_token or _get_api_token_from_environment()
|
||||
):
|
||||
headers["Authorization"] = f"Bearer {api_token}"
|
||||
|
||||
base_url = (
|
||||
base_url or os.environ.get("REPLICATE_BASE_URL") or "https://api.replicate.com"
|
||||
)
|
||||
if base_url == "":
|
||||
base_url = "https://api.replicate.com"
|
||||
|
||||
timeout = timeout or httpx.Timeout(
|
||||
5.0, read=30.0, write=30.0, connect=5.0, pool=10.0
|
||||
)
|
||||
|
||||
transport = kwargs.pop("transport", None) or (
|
||||
httpx.HTTPTransport()
|
||||
if client_type is httpx.Client
|
||||
else httpx.AsyncHTTPTransport()
|
||||
)
|
||||
|
||||
return client_type(
|
||||
base_url=base_url,
|
||||
headers=headers,
|
||||
timeout=timeout,
|
||||
transport=RetryTransport(wrapped_transport=transport), # type: ignore[arg-type]
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _raise_for_status(resp: httpx.Response) -> None:
|
||||
if 400 <= resp.status_code < 600:
|
||||
raise ReplicateError.from_response(resp)
|
||||
Reference in New Issue
Block a user