singlestoredb 1.16.1__py3-none-any.whl

singlestoredb/management/inference_api.py

@@ -0,0 +1,105 @@
+#!/usr/bin/env python
+"""SingleStoreDB Cloud Inference API."""
+import os
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+from .utils import vars_to_str
+from singlestoredb.exceptions import ManagementError
+from singlestoredb.management.manager import Manager
+
+
+class InferenceAPIInfo(object):
+    """
+    Inference API definition.
+
+    This object is not directly instantiated. It is used in results
+    of API calls on the :class:`InferenceAPIManager`. See :meth:`InferenceAPIManager.get`.
+    """
+
+    service_id: str
+    model_name: str
+    name: str
+    connection_url: str
+    project_id: str
+    hosting_platform: str
+
+    def __init__(
+        self,
+        service_id: str,
+        model_name: str,
+        name: str,
+        connection_url: str,
+        project_id: str,
+        hosting_platform: str,
+    ):
+        self.service_id = service_id
+        self.connection_url = connection_url
+        self.model_name = model_name
+        self.name = name
+        self.project_id = project_id
+        self.hosting_platform = hosting_platform
+
+    @classmethod
+    def from_dict(
+        cls,
+        obj: Dict[str, Any],
+    ) -> 'InferenceAPIInfo':
+        """
+        Construct a Inference API from a dictionary of values.
+
+        Parameters
+        ----------
+        obj : dict
+            Dictionary of values
+
+        Returns
+        -------
+        :class:`Job`
+
+        """
+        out = cls(
+            service_id=obj['serviceID'],
+            project_id=obj['projectID'],
+            model_name=obj['modelName'],
+            name=obj['name'],
+            connection_url=obj['connectionURL'],
+            hosting_platform=obj['hostingPlatform'],
+        )
+        return out
+
+    def __str__(self) -> str:
+        """Return string representation."""
+        return vars_to_str(self)
+
+    def __repr__(self) -> str:
+        """Return string representation."""
+        return str(self)
+
+
+class InferenceAPIManager(object):
+    """
+    SingleStoreDB Inference APIs manager.
+
+    This class should be instantiated using :attr:`Organization.inference_apis`.
+
+    Parameters
+    ----------
+    manager : InferenceAPIManager, optional
+        The InferenceAPIManager the InferenceAPIManager belongs to
+
+    See Also
+    --------
+    :attr:`InferenceAPI`
+    """
+
+    def __init__(self, manager: Optional[Manager]):
+        self._manager = manager
+        self.project_id = os.environ.get('SINGLESTOREDB_PROJECT')
+
+    def get(self, model_name: str) -> InferenceAPIInfo:
+        if self._manager is None:
+            raise ManagementError(msg='Manager not initialized')
+        res = self._manager._get(f'inferenceapis/{self.project_id}/{model_name}').json()
+        return InferenceAPIInfo.from_dict(res)