Created InferenceClient class, use that in Deployment

tamirse · tamirse · commit 68275a2a3994 · 2025-04-01T10:57:11.000+03:00
diff --git a/datacrunch/InferenceClient/__init__.py b/datacrunch/InferenceClient/__init__.py
diff --git a/datacrunch/InferenceClient/inference_client.py b/datacrunch/InferenceClient/inference_client.py
@@ -0,0 +1,198 @@
+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
+from urllib.parse import urlparse
+
+
+class InferenceClientError(Exception):
+    """Base exception for InferenceClient errors."""
+    pass
+
+
+@dataclass_json(undefined=Undefined.EXCLUDE)
+@dataclass
+class InferenceResponse:
+    body: Any
+    headers: CaseInsensitiveDict[str]
+    status_code: int
+    status_text: str
+
+
+class InferenceClient:
+    def __init__(self, inference_key: str, endpoint_base_url: str, timeout_seconds: int = 300) -> 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.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):
+        response = self.post(
+            path, json=data, timeout_seconds=timeout_seconds, headers=headers)
+
+        return InferenceResponse(
+            body=response.json(),
+            headers=response.headers,
+            status_code=response.status_code,
+            status_text=response.reason
+        )
+
+    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) -> dict:
+        """
+        Check the health status of the API.
+
+        Returns:
+            dict: Health status information
+
+        Raises:
+            InferenceClientError: If the health check fails
+        """
+        try:
+            response = self.get('/health')
+            return response.json()
+        except InferenceClientError as e:
+            raise InferenceClientError(f"Health check failed: {str(e)}")
diff --git a/datacrunch/containers/containers.py b/datacrunch/containers/containers.py
@@ -1,11 +1,11 @@
 import requests
-from requests.structures import CaseInsensitiveDict
 from dataclasses import dataclass
 from dataclasses_json import dataclass_json, Undefined  # type: ignore
 from typing import List, Optional, Dict, Any
 from enum import Enum
 
 from datacrunch.http_client.http_client import HTTPClient
+from datacrunch.InferenceClient.inference_client import InferenceClient, InferenceResponse
 
 
 # API endpoints
@@ -238,7 +238,6 @@ class Deployment:
     :param endpoint_base_url: Optional base URL for the deployment endpoint
     :param scaling: Optional scaling configuration
     :param created_at: Optional timestamp when the deployment was created
-    :param inference_key: Optional inference key for the deployment
     """
     name: str
     container_registry_settings: ContainerRegistrySettings
@@ -249,13 +248,13 @@ class Deployment:
     scaling: Optional[ScalingOptions] = None
     created_at: Optional[str] = None
 
-    _inference_key: Optional[str] = None
+    _inference_client: Optional[InferenceClient] = None
 
     def __str__(self):
         """String representation of the deployment, excluding sensitive information."""
-        # Get all attributes except _inference_key
+        # Get all attributes except _inference_client
         attrs = {k: v for k, v in self.__dict__.items() if k !=
-                 '_inference_key'}
+                 '_inference_client'}
         # Format each attribute
         attr_strs = [f"{k}={repr(v)}" for k, v in attrs.items()]
         return f"Deployment({', '.join(attr_strs)})"
@@ -265,38 +264,47 @@ def __repr__(self):
         return self.__str__()
 
     @classmethod
-    def from_dict_with_inference_key(cls, data: Dict[str, Any], inference_key: str = None, **kwargs) -> 'Deployment':
+    def from_dict_with_inference_key(cls, data: Dict[str, Any], inference_key: str = None) -> 'Deployment':
         """Create a Deployment instance from a dictionary with an inference key.
 
         :param data: Dictionary containing deployment data
         :param inference_key: inference key to set on the deployment
-        :param **kwargs: Additional arguments to pass to from_dict
         :return: Deployment instance
         """
         deployment = Deployment.from_dict(data, infer_missing=True)
-        deployment._inference_key = inference_key
+        if inference_key and deployment.endpoint_base_url:
+            deployment._inference_client = InferenceClient(
+                inference_key=inference_key,
+                endpoint_base_url=deployment.endpoint_base_url
+            )
         return deployment
 
-    def run_sync(self, data: Dict[str, Any], path: str = "", timeout_seconds: int = 60 * 5):
-        if self._inference_key is None:
-            # TODO: do something better
-            raise ValueError("Inference key is not set")
+    def set_inference_client(self, inference_key: str) -> None:
+        """Set the inference client for this deployment.
 
-        response = requests.post(
-            url=f"{self.endpoint_base_url}{path}",
-            json=data,
-            headers={"Authorization": f"Bearer {self._inference_key}"},
-            timeout=timeout_seconds
+        :param inference_key: The inference key to use for authentication
+        :type inference_key: str
+        :raises ValueError: If endpoint_base_url is not set
+        """
+        if self.endpoint_base_url is None:
+            raise ValueError(
+                "Endpoint base URL must be set to use inference client")
+        self._inference_client = InferenceClient(
+            inference_key=inference_key,
+            endpoint_base_url=self.endpoint_base_url
         )
 
-        return InferenceResponse(
-            body=response.json(),
-            headers=response.headers,
-            status_code=response.status_code,
-            status_text=response.reason
-        )
+    def run_sync(self, data: Dict[str, Any], path: str = "", timeout_seconds: int = 60 * 5) -> InferenceResponse:
+        if self._inference_client is None:
+            if self.endpoint_base_url is None:
+                raise ValueError(
+                    "Endpoint base URL must be set to use run_sync")
+            raise ValueError(
+                "Inference client not initialized. Use from_dict_with_inference_key or set_inference_client to initialize inference capabilities.")
+        return self._inference_client.run_sync(data, path, timeout_seconds)
 
     def health(self):
+        # TODO: use inference client?
         healthcheck_path = "health"
         if self.containers and self.containers[0].healthcheck and self.containers[0].healthcheck.path:
             healthcheck_path = self.containers[0].healthcheck.path.lstrip('/')
@@ -310,15 +318,6 @@ def health(self):
     healthcheck = health
 
 
-@dataclass_json(undefined=Undefined.EXCLUDE)
-@dataclass
-class InferenceResponse:
-    body: Any
-    headers: CaseInsensitiveDict[str]
-    status_code: int
-    status_text: str
-
-
 @dataclass_json
 @dataclass
 class ReplicaInfo:
