Add Wrapper for ChatOpenAI and OpenAIEmbeddings (#62)

Author: hbindu-ss

singlestoredb/ai/__init__.py

@@ -1 +1,2 @@
+from .chat import SingleStoreChatOpenAI  # noqa: F401
 from .embeddings import SingleStoreEmbeddings  # noqa: F401

singlestoredb/ai/chat.py

@@ -0,0 +1,26 @@
+import os
+from typing import Any
+
+from singlestoredb.fusion.handlers.utils import get_workspace_manager
+
+try:
+    from langchain_openai import ChatOpenAI
+except ImportError:
+    raise ImportError(
+        'Could not import langchain_openai python package. '
+        'Please install it with `pip install langchain_openai`.',
+    )
+
+
+class SingleStoreChatOpenAI(ChatOpenAI):
+    def __init__(self, model_name: str, **kwargs: Any):
+        inference_api_manger = (
+            get_workspace_manager().organizations.current.inference_apis
+        )
+        info = inference_api_manger.get(model_name=model_name)
+        super().__init__(
+            base_url=info.connection_url,
+            api_key=os.environ.get('SINGLESTOREDB_USER_TOKEN'),
+            model=model_name,
+            **kwargs,
+        )

singlestoredb/ai/embeddings.py

@@ -1,24 +1,27 @@
-import os as _os
+import os
 from typing import Any
 
+from singlestoredb.fusion.handlers.utils import get_workspace_manager
+
 try:
-    from langchain_community.embeddings.ollama import OllamaEmbeddings
+    from langchain_openai import OpenAIEmbeddings
 except ImportError:
     raise ImportError(
-        'Could not import langchain_community python package. '
-        'Please install it with `pip install langchain_community`.',
+        'Could not import langchain_openai python package. '
+        'Please install it with `pip install langchain_openai`.',
     )
 
 
-class SingleStoreEmbeddings(OllamaEmbeddings):
-
-    def __init__(self, **kwargs: Any):
-        url = _os.getenv('SINGLESTORE_AI_EXPERIMENTAL_URL')
-        if not url:
-            raise ValueError(
-                "Environment variable 'SINGLESTORE_AI_EXPERIMENTAL_URL' must be set",
-            )
+class SingleStoreEmbeddings(OpenAIEmbeddings):
 
-        base_url = url.strip('/v1')
-        kwargs = {'model': 'nomic-embed-text', **kwargs}
-        super().__init__(base_url=base_url, **kwargs)
+    def __init__(self, model_name: str, **kwargs: Any):
+        inference_api_manger = (
+            get_workspace_manager().organizations.current.inference_apis
+        )
+        info = inference_api_manger.get(model_name=model_name)
+        super().__init__(
+            base_url=info.connection_url,
+            api_key=os.environ.get('SINGLESTOREDB_USER_TOKEN'),
+            model=model_name,
+            **kwargs,
+        )

singlestoredb/management/inference_api.py

@@ -0,0 +1,101 @@
+#!/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
+
+    def __init__(
+        self,
+        service_id: str,
+        model_name: str,
+        name: str,
+        connection_url: str,
+        project_id: str,
+    ):
+        self.service_id = service_id
+        self.connection_url = connection_url
+        self.model_name = model_name
+        self.name = name
+        self.project_id = project_id
+
+    @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'],
+        )
+        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)

singlestoredb/management/organization.py

@@ -7,6 +7,7 @@
 from typing import Union
 
 from ..exceptions import ManagementError
+from .inference_api import InferenceAPIManager
 from .job import JobsManager
 from .manager import Manager
 from .utils import vars_to_str
@@ -207,3 +208,19 @@ def jobs(self) -> JobsManager:
         :class:`JobsManager`
         """
         return JobsManager(self._manager)
+
+    @property
+    def inference_apis(self) -> InferenceAPIManager:
+        """
+        Retrieve a SingleStoreDB inference api manager.
+
+        Parameters
+        ----------
+        manager : WorkspaceManager, optional
+            The WorkspaceManager the InferenceAPIManager belongs to
+
+        Returns
+        -------
+        :class:`InferenceAPIManager`
+        """
+        return InferenceAPIManager(self._manager)