datacrunch 1.16.0__py3-none-any.whl → 1.17.1__py3-none-any.whl

datacrunch/InferenceClient/inference_client.py

@@ -1,514 +0,0 @@
-from dataclasses import dataclass
-from dataclasses_json import dataclass_json, Undefined  # type: ignore
-import requests
-from requests.structures import CaseInsensitiveDict
-from typing import Optional, Dict, Any, Union, Generator
-from urllib.parse import urlparse
-from enum import Enum
-
-
-class InferenceClientError(Exception):
-    """Base exception for InferenceClient errors."""
-
-    pass
-
-
-class AsyncStatus(str, Enum):
-    Initialized = 'Initialized'
-    Queue = 'Queue'
-    Inference = 'Inference'
-    Completed = 'Completed'
-
-
-@dataclass_json(undefined=Undefined.EXCLUDE)
-@dataclass
-class InferenceResponse:
-    headers: CaseInsensitiveDict[str]
-    status_code: int
-    status_text: str
-    _original_response: requests.Response
-    _stream: bool = False
-
-    def _is_stream_response(self, headers: CaseInsensitiveDict[str]) -> bool:
-        """Check if the response headers indicate a streaming response.
-
-        Args:
-            headers: The response headers to check
-
-        Returns:
-            bool: True if the response is likely a stream, False otherwise
-        """
-        # Standard chunked transfer encoding
-        is_chunked_transfer = headers.get('Transfer-Encoding', '').lower() == 'chunked'
-        # Server-Sent Events content type
-        is_event_stream = headers.get('Content-Type', '').lower() == 'text/event-stream'
-        # NDJSON
-        is_ndjson = headers.get('Content-Type', '').lower() == 'application/x-ndjson'
-        # Stream JSON
-        is_stream_json = headers.get('Content-Type', '').lower() == 'application/stream+json'
-        # Keep-alive
-        is_keep_alive = headers.get('Connection', '').lower() == 'keep-alive'
-        # No content length
-        has_no_content_length = 'Content-Length' not in headers
-
-        # No Content-Length with keep-alive often suggests streaming (though not definitive)
-        is_keep_alive_and_no_content_length = is_keep_alive and has_no_content_length
-
-        return (
-            self._stream
-            or is_chunked_transfer
-            or is_event_stream
-            or is_ndjson
-            or is_stream_json
-            or is_keep_alive_and_no_content_length
-        )
-
-    def output(self, is_text: bool = False) -> Any:
-        try:
-            if is_text:
-                return self._original_response.text
-            return self._original_response.json()
-        except Exception as e:
-            # if the response is a stream (check headers), raise relevant error
-            if self._is_stream_response(self._original_response.headers):
-                raise InferenceClientError(
-                    'Response might be a stream, use the stream method instead'
-                )
-            raise InferenceClientError(f'Failed to parse response as JSON: {str(e)}')
-
-    def stream(self, chunk_size: int = 512, as_text: bool = True) -> Generator[Any, None, None]:
-        """Stream the response content.
-
-        Args:
-            chunk_size: Size of chunks to stream, in bytes
-            as_text: If True, stream as text using iter_lines. If False, stream as binary using iter_content.
-
-        Returns:
-            Generator yielding chunks of the response
-        """
-        if as_text:
-            for chunk in self._original_response.iter_lines(chunk_size=chunk_size):
-                if chunk:
-                    yield chunk
-        else:
-            for chunk in self._original_response.iter_content(chunk_size=chunk_size):
-                if chunk:
-                    yield chunk
-
-
-class InferenceClient:
-    def __init__(
-        self, inference_key: str, endpoint_base_url: str, timeout_seconds: int = 60 * 5
-    ) -> None:
-        """
-        Initialize the InferenceClient.
-
-        Args:
-            inference_key: The authentication key for the API
-            endpoint_base_url: The base URL for the API
-            timeout_seconds: Request timeout in seconds
-
-        Raises:
-            InferenceClientError: If the parameters are invalid
-        """
-        if not inference_key:
-            raise InferenceClientError('inference_key cannot be empty')
-
-        parsed_url = urlparse(endpoint_base_url)
-        if not parsed_url.scheme or not parsed_url.netloc:
-            raise InferenceClientError('endpoint_base_url must be a valid URL')
-
-        self.inference_key = inference_key
-        self.endpoint_base_url = endpoint_base_url.rstrip('/')
-        self.base_domain = self.endpoint_base_url[: self.endpoint_base_url.rindex('/')]
-        self.deployment_name = self.endpoint_base_url[self.endpoint_base_url.rindex('/') + 1 :]
-        self.timeout_seconds = timeout_seconds
-        self._session = requests.Session()
-        self._global_headers = {
-            'Authorization': f'Bearer {inference_key}',
-            'Content-Type': 'application/json',
-        }
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self._session.close()
-
-    @property
-    def global_headers(self) -> Dict[str, str]:
-        """
-        Get the current global headers that will be used for all requests.
-
-        Returns:
-            Dictionary of current global headers
-        """
-        return self._global_headers.copy()
-
-    def set_global_header(self, key: str, value: str) -> None:
-        """
-        Set or update a global header that will be used for all requests.
-
-        Args:
-            key: Header name
-            value: Header value
-        """
-        self._global_headers[key] = value
-
-    def set_global_headers(self, headers: Dict[str, str]) -> None:
-        """
-        Set multiple global headers at once that will be used for all requests.
-
-        Args:
-            headers: Dictionary of headers to set globally
-        """
-        self._global_headers.update(headers)
-
-    def remove_global_header(self, key: str) -> None:
-        """
-        Remove a global header.
-
-        Args:
-            key: Header name to remove from global headers
-        """
-        if key in self._global_headers:
-            del self._global_headers[key]
-
-    def _build_url(self, path: str) -> str:
-        """Construct the full URL by joining the base URL with the path."""
-        return f'{self.endpoint_base_url}/{path.lstrip("/")}'
-
-    def _build_request_headers(
-        self, request_headers: Optional[Dict[str, str]] = None
-    ) -> Dict[str, str]:
-        """
-        Build the final headers by merging global headers with request-specific headers.
-
-        Args:
-            request_headers: Optional headers specific to this request
-
-        Returns:
-            Merged headers dictionary
-        """
-        headers = self._global_headers.copy()
-        if request_headers:
-            headers.update(request_headers)
-        return headers
-
-    def _make_request(self, method: str, path: str, **kwargs) -> requests.Response:
-        """
-        Make an HTTP request with error handling.
-
-        Args:
-            method: HTTP method to use
-            path: API endpoint path
-            **kwargs: Additional arguments to pass to the request
-
-        Returns:
-            Response object from the request
-
-        Raises:
-            InferenceClientError: If the request fails
-        """
-        timeout = kwargs.pop('timeout_seconds', self.timeout_seconds)
-        try:
-            response = self._session.request(
-                method=method,
-                url=self._build_url(path),
-                headers=self._build_request_headers(kwargs.pop('headers', None)),
-                timeout=timeout,
-                **kwargs,
-            )
-            response.raise_for_status()
-            return response
-        except requests.exceptions.Timeout:
-            raise InferenceClientError(f'Request to {path} timed out after {timeout} seconds')
-        except requests.exceptions.RequestException as e:
-            raise InferenceClientError(f'Request to {path} failed: {str(e)}')
-
-    def run_sync(
-        self,
-        data: Dict[str, Any],
-        path: str = '',
-        timeout_seconds: int = 60 * 5,
-        headers: Optional[Dict[str, str]] = None,
-        http_method: str = 'POST',
-        stream: bool = False,
-    ):
-        """Make a synchronous request to the inference endpoint.
-
-        Args:
-            data: The data payload to send with the request
-            path: API endpoint path. Defaults to empty string.
-            timeout_seconds: Request timeout in seconds. Defaults to 5 minutes.
-            headers: Optional headers to include in the request
-            http_method: HTTP method to use. Defaults to "POST".
-            stream: Whether to stream the response. Defaults to False.
-
-        Returns:
-            InferenceResponse: Object containing the response data.
-
-        Raises:
-            InferenceClientError: If the request fails
-        """
-        response = self._make_request(
-            http_method,
-            path,
-            json=data,
-            timeout_seconds=timeout_seconds,
-            headers=headers,
-            stream=stream,
-        )
-
-        return InferenceResponse(
-            headers=response.headers,
-            status_code=response.status_code,
-            status_text=response.reason,
-            _original_response=response,
-        )
-
-    def run(
-        self,
-        data: Dict[str, Any],
-        path: str = '',
-        timeout_seconds: int = 60 * 5,
-        headers: Optional[Dict[str, str]] = None,
-        http_method: str = 'POST',
-        no_response: bool = False,
-    ):
-        """Make an asynchronous request to the inference endpoint.
-
-        Args:
-            data: The data payload to send with the request
-            path: API endpoint path. Defaults to empty string.
-            timeout_seconds: Request timeout in seconds. Defaults to 5 minutes.
-            headers: Optional headers to include in the request
-            http_method: HTTP method to use. Defaults to "POST".
-            no_response: If True, don't wait for response. Defaults to False.
-
-        Returns:
-            AsyncInferenceExecution: Object to track the async execution status.
-            If no_response is True, returns None.
-
-        Raises:
-            InferenceClientError: If the request fails
-        """
-        # Add relevant headers to the request, to indicate that the request is async
-        headers = headers or {}
-        if no_response:
-            # If no_response is True, use the "Prefer: respond-async-proxy" header to run async and don't wait for the response
-            headers['Prefer'] = 'respond-async-proxy'
-            self._make_request(
-                http_method,
-                path,
-                json=data,
-                timeout_seconds=timeout_seconds,
-                headers=headers,
-            )
-            return
-        # Add the "Prefer: respond-async" header to the request, to run async and wait for the response
-        headers['Prefer'] = 'respond-async'
-
-        response = self._make_request(
-            http_method,
-            path,
-            json=data,
-            timeout_seconds=timeout_seconds,
-            headers=headers,
-        )
-
-        result = response.json()
-        execution_id = result['Id']
-
-        return AsyncInferenceExecution(self, execution_id, AsyncStatus.Initialized)
-
-    def get(
-        self,
-        path: str,
-        params: Optional[Dict[str, Any]] = None,
-        headers: Optional[Dict[str, str]] = None,
-        timeout_seconds: Optional[int] = None,
-    ) -> requests.Response:
-        return self._make_request(
-            'GET', path, params=params, headers=headers, timeout_seconds=timeout_seconds
-        )
-
-    def post(
-        self,
-        path: str,
-        json: Optional[Dict[str, Any]] = None,
-        data: Optional[Union[str, Dict[str, Any]]] = None,
-        params: Optional[Dict[str, Any]] = None,
-        headers: Optional[Dict[str, str]] = None,
-        timeout_seconds: Optional[int] = None,
-    ) -> requests.Response:
-        return self._make_request(
-            'POST',
-            path,
-            json=json,
-            data=data,
-            params=params,
-            headers=headers,
-            timeout_seconds=timeout_seconds,
-        )
-
-    def put(
-        self,
-        path: str,
-        json: Optional[Dict[str, Any]] = None,
-        data: Optional[Union[str, Dict[str, Any]]] = None,
-        params: Optional[Dict[str, Any]] = None,
-        headers: Optional[Dict[str, str]] = None,
-        timeout_seconds: Optional[int] = None,
-    ) -> requests.Response:
-        return self._make_request(
-            'PUT',
-            path,
-            json=json,
-            data=data,
-            params=params,
-            headers=headers,
-            timeout_seconds=timeout_seconds,
-        )
-
-    def delete(
-        self,
-        path: str,
-        params: Optional[Dict[str, Any]] = None,
-        headers: Optional[Dict[str, str]] = None,
-        timeout_seconds: Optional[int] = None,
-    ) -> requests.Response:
-        return self._make_request(
-            'DELETE',
-            path,
-            params=params,
-            headers=headers,
-            timeout_seconds=timeout_seconds,
-        )
-
-    def patch(
-        self,
-        path: str,
-        json: Optional[Dict[str, Any]] = None,
-        data: Optional[Union[str, Dict[str, Any]]] = None,
-        params: Optional[Dict[str, Any]] = None,
-        headers: Optional[Dict[str, str]] = None,
-        timeout_seconds: Optional[int] = None,
-    ) -> requests.Response:
-        return self._make_request(
-            'PATCH',
-            path,
-            json=json,
-            data=data,
-            params=params,
-            headers=headers,
-            timeout_seconds=timeout_seconds,
-        )
-
-    def head(
-        self,
-        path: str,
-        params: Optional[Dict[str, Any]] = None,
-        headers: Optional[Dict[str, str]] = None,
-        timeout_seconds: Optional[int] = None,
-    ) -> requests.Response:
-        return self._make_request(
-            'HEAD',
-            path,
-            params=params,
-            headers=headers,
-            timeout_seconds=timeout_seconds,
-        )
-
-    def options(
-        self,
-        path: str,
-        params: Optional[Dict[str, Any]] = None,
-        headers: Optional[Dict[str, str]] = None,
-        timeout_seconds: Optional[int] = None,
-    ) -> requests.Response:
-        return self._make_request(
-            'OPTIONS',
-            path,
-            params=params,
-            headers=headers,
-            timeout_seconds=timeout_seconds,
-        )
-
-    def health(self, healthcheck_path: str = '/health') -> requests.Response:
-        """
-        Check the health status of the API.
-
-        Returns:
-            requests.Response: The response from the health check
-
-        Raises:
-            InferenceClientError: If the health check fails
-        """
-        try:
-            return self.get(healthcheck_path)
-        except InferenceClientError as e:
-            raise InferenceClientError(f'Health check failed: {str(e)}')
-
-
-@dataclass_json(undefined=Undefined.EXCLUDE)
-@dataclass
-class AsyncInferenceExecution:
-    _inference_client: 'InferenceClient'
-    id: str
-    _status: AsyncStatus
-    INFERENCE_ID_HEADER = 'X-Inference-Id'
-
-    def status(self) -> AsyncStatus:
-        """Get the current stored status of the async inference execution. Only the status value type
-
-        Returns:
-            AsyncStatus: The status object
-        """
-
-        return self._status
-
-    def status_json(self) -> Dict[str, Any]:
-        """Get the current status of the async inference execution. Return the status json
-
-        Returns:
-            Dict[str, Any]: The status response containing the execution status and other metadata
-        """
-        url = (
-            f'{self._inference_client.base_domain}/status/{self._inference_client.deployment_name}'
-        )
-        response = self._inference_client._session.get(
-            url,
-            headers=self._inference_client._build_request_headers(
-                {self.INFERENCE_ID_HEADER: self.id}
-            ),
-        )
-
-        response_json = response.json()
-        self._status = AsyncStatus(response_json['Status'])
-
-        return response_json
-
-    def result(self) -> Dict[str, Any]:
-        """Get the results of the async inference execution.
-
-        Returns:
-            Dict[str, Any]: The results of the inference execution
-        """
-        url = (
-            f'{self._inference_client.base_domain}/result/{self._inference_client.deployment_name}'
-        )
-        response = self._inference_client._session.get(
-            url,
-            headers=self._inference_client._build_request_headers(
-                {self.INFERENCE_ID_HEADER: self.id}
-            ),
-        )
-
-        if response.headers['Content-Type'] == 'application/json':
-            return response.json()
-        else:
-            return {'result': response.text}
-
-    # alias for get_results
-    output = result

datacrunch/_version.py

@@ -1,6 +0,0 @@
-try:
-    from importlib.metadata import version
-
-    __version__ = version('datacrunch')
-except Exception:
-    __version__ = '0.0.0+dev'  # fallback for development

datacrunch/authentication/authentication.py

@@ -1,105 +0,0 @@
-import requests
-import time
-
-from datacrunch.http_client.http_client import handle_error
-
-TOKEN_ENDPOINT = '/oauth2/token'
-
-CLIENT_CREDENTIALS = 'client_credentials'
-REFRESH_TOKEN = 'refresh_token'
-
-
-class AuthenticationService:
-    """A service for client authentication"""
-
-    def __init__(self, client_id: str, client_secret: str, base_url: str) -> None:
-        self._base_url = base_url
-        self._client_id = client_id
-        self._client_secret = client_secret
-
-    def authenticate(self) -> dict:
-        """Authenticate the client and store the access & refresh tokens
-
-        returns an authentication data dictionary with the following schema:
-        {
-            "access_token": token str,
-            "refresh_token": token str,
-            "scope": scope str,
-            "token_type": token type str,
-            "expires_in": duration until expires in seconds
-        }
-
-        :return: authentication data (tokens, scope, token type, expires in)
-        :rtype: dict
-        """
-        url = self._base_url + TOKEN_ENDPOINT
-        payload = {
-            'grant_type': CLIENT_CREDENTIALS,
-            'client_id': self._client_id,
-            'client_secret': self._client_secret,
-        }
-
-        response = requests.post(url, json=payload, headers=self._generate_headers())
-        handle_error(response)
-
-        auth_data = response.json()
-
-        self._access_token = auth_data['access_token']
-        self._refresh_token = auth_data['refresh_token']
-        self._scope = auth_data['scope']
-        self._token_type = auth_data['token_type']
-        self._expires_at = time.time() + auth_data['expires_in']
-
-        return auth_data
-
-    def refresh(self) -> dict:
-        """Authenticate the client using the refresh token - refresh the access token.
-
-        updates the object's tokens, and:
-        returns an authentication data dictionary with the following schema:
-        {
-            "access_token": token str,
-            "refresh_token": token str,
-            "scope": scope str,
-            "token_type": token type str,
-            "expires_in": duration until expires in seconds
-        }
-
-        :return: authentication data (tokens, scope, token type, expires in)
-        :rtype: dict
-        """
-        url = self._base_url + TOKEN_ENDPOINT
-
-        payload = {'grant_type': REFRESH_TOKEN, 'refresh_token': self._refresh_token}
-
-        response = requests.post(url, json=payload, headers=self._generate_headers())
-
-        # if refresh token is also expired, authenticate again:
-        if response.status_code == 401 or response.status_code == 400:
-            return self.authenticate()
-        else:
-            handle_error(response)
-
-        auth_data = response.json()
-
-        self._access_token = auth_data['access_token']
-        self._refresh_token = auth_data['refresh_token']
-        self._scope = auth_data['scope']
-        self._token_type = auth_data['token_type']
-        self._expires_at = time.time() + auth_data['expires_in']
-
-        return auth_data
-
-    def _generate_headers(self):
-        # get the first 10 chars of the client id
-        client_id_truncated = self._client_id[:10]
-        headers = {'User-Agent': 'datacrunch-python-' + client_id_truncated}
-        return headers
-
-    def is_expired(self) -> bool:
-        """Returns true if the access token is expired.
-
-        :return: True if the access token is expired, otherwise False.
-        :rtype: bool
-        """
-        return time.time() >= self._expires_at

datacrunch/balance/balance.py

@@ -1,50 +0,0 @@
-BALANCE_ENDPOINT = '/balance'
-
-
-class Balance:
-    """A balance model class"""
-
-    def __init__(self, amount: float, currency: str) -> None:
-        """Initialize a new Balance object
-
-        :param amount: Balance amount
-        :type amount: float
-        :param currency: currency code
-        :type currency: str
-        """
-        self._amount = amount
-        self._currency = currency
-
-    @property
-    def amount(self) -> float:
-        """Get the balance amount
-
-        :return: amount
-        :rtype: float
-        """
-        return self._amount
-
-    @property
-    def currency(self) -> str:
-        """Get the currency code
-
-        :return: currency code
-        :rtype: str
-        """
-        return self._currency
-
-
-class BalanceService:
-    """A service for interacting with the balance endpoint"""
-
-    def __init__(self, http_client) -> None:
-        self._http_client = http_client
-
-    def get(self) -> Balance:
-        """Get the client's current balance
-
-        :return: Balance object containing the amount and currency.
-        :rtype: Balance
-        """
-        balance = self._http_client.get(BALANCE_ENDPOINT).json()
-        return Balance(balance['amount'], balance['currency'])