Minor fixes.

Author: mgiannakopoulos

singlestoredb/ai/__init__.py

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

singlestoredb/ai/chat.py

@@ -6,7 +6,7 @@
 
 import httpx
 
-from singlestoredb.fusion.handlers.utils import get_workspace_manager
+from singlestoredb import manage_workspaces
 
 try:
     from langchain_openai import ChatOpenAI
@@ -32,7 +32,7 @@
 class SingleStoreChatOpenAI(ChatOpenAI):
     def __init__(self, model_name: str, api_key: Optional[str] = None, **kwargs: Any):
         inference_api_manger = (
-            get_workspace_manager().organizations.current.inference_apis
+            manage_workspaces().organizations.current.inference_apis
         )
         info = inference_api_manger.get(model_name=model_name)
         token = (
@@ -51,7 +51,7 @@ def __init__(self, model_name: str, api_key: Optional[str] = None, **kwargs: Any
 class SingleStoreChat(ChatOpenAI):
     def __init__(self, model_name: str, api_key: Optional[str] = None, **kwargs: Any):
         inference_api_manger = (
-            get_workspace_manager().organizations.current.inference_apis
+            manage_workspaces().organizations.current.inference_apis
         )
         info = inference_api_manger.get(model_name=model_name)
         token = (
@@ -78,7 +78,7 @@ def SingleStoreChatFactory(
     """Return a chat model instance (ChatOpenAI or ChatBedrockConverse).
     """
     inference_api_manager = (
-        get_workspace_manager().organizations.current.inference_apis
+        manage_workspaces().organizations.current.inference_apis
     )
     info = inference_api_manager.get(model_name=model_name)
     token_env = os.environ.get('SINGLESTOREDB_USER_TOKEN')

singlestoredb/ai/embeddings.py

@@ -6,7 +6,7 @@
 
 import httpx
 
-from singlestoredb.fusion.handlers.utils import get_workspace_manager
+from singlestoredb import manage_workspaces
 
 try:
     from langchain_openai import OpenAIEmbeddings
@@ -33,7 +33,7 @@ class SingleStoreEmbeddings(OpenAIEmbeddings):
 
     def __init__(self, model_name: str, **kwargs: Any):
         inference_api_manger = (
-            get_workspace_manager().organizations.current.inference_apis
+            manage_workspaces().organizations.current.inference_apis
         )
         info = inference_api_manger.get(model_name=model_name)
         super().__init__(
@@ -54,7 +54,7 @@ def SingleStoreEmbeddingsFactory(
     """Return an embeddings model instance (OpenAIEmbeddings or BedrockEmbeddings).
     """
     inference_api_manager = (
-        get_workspace_manager().organizations.current.inference_apis
+        manage_workspaces().organizations.current.inference_apis
     )
     info = inference_api_manager.get(model_name=model_name)
     token_env = os.environ.get('SINGLESTOREDB_USER_TOKEN')