diff --git a/singlestoredb/ai/__init__.py b/singlestoredb/ai/__init__.py index 67f50b8da..6f97c99f4 100644 --- a/singlestoredb/ai/__init__.py +++ b/singlestoredb/ai/__init__.py @@ -1 +1,2 @@ +from .chat import SingleStoreChatOpenAI # noqa: F401 from .embeddings import SingleStoreEmbeddings # noqa: F401 diff --git a/singlestoredb/ai/chat.py b/singlestoredb/ai/chat.py new file mode 100644 index 000000000..100bcb454 --- /dev/null +++ b/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, + ) diff --git a/singlestoredb/ai/embeddings.py b/singlestoredb/ai/embeddings.py index c9fb1195d..106b66328 100644 --- a/singlestoredb/ai/embeddings.py +++ b/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, + ) diff --git a/singlestoredb/management/inference_api.py b/singlestoredb/management/inference_api.py new file mode 100644 index 000000000..f017339d4 --- /dev/null +++ b/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) diff --git a/singlestoredb/management/organization.py b/singlestoredb/management/organization.py index d6415e046..2c0f917df 100644 --- a/singlestoredb/management/organization.py +++ b/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)