singlestoredb 1.12.4__py3-none-any.whl → 1.13.1__py3-none-any.whl

singlestoredb/__init__.py

@@ -13,7 +13,7 @@ Examples
 
 """
 
-__version__ = '1.12.4'
+__version__ = '1.13.1'
 
 from typing import Any
 

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/apps/__init__.py

@@ -1,2 +1,3 @@
 from ._cloud_functions import run_function_app  # noqa: F401
 from ._dashboards import run_dashboard_app  # noqa: F401
+from ._python_udfs import run_udf_app  # noqa: F401

singlestoredb/apps/_config.py

@@ -8,10 +8,12 @@ class AppConfig:
     listen_port: int
     base_url: str
     base_path: str
+    notebook_server_id: str
     app_token: Optional[str]
     user_token: Optional[str]
     running_interactively: bool
     is_gateway_enabled: bool
+    is_local_dev: bool
 
     @staticmethod
     def _read_variable(name: str) -> str:
@@ -28,6 +30,8 @@ class AppConfig:
         port = cls._read_variable('SINGLESTOREDB_APP_LISTEN_PORT')
         base_url = cls._read_variable('SINGLESTOREDB_APP_BASE_URL')
         base_path = cls._read_variable('SINGLESTOREDB_APP_BASE_PATH')
+        notebook_server_id = cls._read_variable('SINGLESTOREDB_NOTEBOOK_SERVER_ID')
+        is_local_dev_env_var = cls._read_variable('SINGLESTOREDB_IS_LOCAL_DEV')
 
         workload_type = os.environ.get('SINGLESTOREDB_WORKLOAD_TYPE')
         running_interactively = workload_type == 'InteractiveNotebook'
@@ -49,10 +53,12 @@ class AppConfig:
             listen_port=int(port),
             base_url=base_url,
             base_path=base_path,
+            notebook_server_id=notebook_server_id,
             app_token=app_token,
             user_token=user_token,
             running_interactively=running_interactively,
             is_gateway_enabled=is_gateway_enabled,
+            is_local_dev=is_local_dev_env_var == 'true',
         )
 
     @property

singlestoredb/apps/_connection_info.py

@@ -1,4 +1,6 @@
 from dataclasses import dataclass
+from typing import Any
+from typing import Dict
 from typing import Optional
 
 
@@ -8,3 +10,9 @@ class ConnectionInfo:
 
     # Only present in interactive mode
     token: Optional[str]
+
+
+@dataclass
+class UdfConnectionInfo:
+    url: str
+    functions: Dict[str, Any]

singlestoredb/apps/_python_udfs.py

@@ -0,0 +1,85 @@
+import asyncio
+import os
+import typing
+
+from ..functions.ext.asgi import Application
+from ._config import AppConfig
+from ._connection_info import UdfConnectionInfo
+from ._process import kill_process_by_port
+
+if typing.TYPE_CHECKING:
+    from ._uvicorn_util import AwaitableUvicornServer
+
+# Keep track of currently running server
+_running_server: 'typing.Optional[AwaitableUvicornServer]' = None
+
+
+async def run_udf_app(
+    replace_existing: bool,
+    log_level: str = 'error',
+    kill_existing_app_server: bool = True,
+) -> UdfConnectionInfo:
+    global _running_server
+    from ._uvicorn_util import AwaitableUvicornServer
+
+    try:
+        import uvicorn
+    except ImportError:
+        raise ImportError('package uvicorn is required to run python udfs')
+
+    app_config = AppConfig.from_env()
+
+    if kill_existing_app_server:
+        # Shutdown the server gracefully if it was started by us.
+        # Since the uvicorn server doesn't start a new subprocess
+        # killing the process would result in kernel dying.
+        if _running_server is not None:
+            await _running_server.shutdown()
+            _running_server = None
+
+        # Kill if any other process is occupying the port
+        kill_process_by_port(app_config.listen_port)
+
+    base_url = generate_base_url(app_config)
+
+    udf_suffix = ''
+    if app_config.running_interactively:
+        udf_suffix = '_test'
+    app = Application(url=base_url, app_mode='managed', name_suffix=udf_suffix)
+
+    config = uvicorn.Config(
+        app,
+        host='0.0.0.0',
+        port=app_config.listen_port,
+        log_level=log_level,
+    )
+    _running_server = AwaitableUvicornServer(config)
+
+    # Register the functions
+    app.register_functions(replace=replace_existing)
+
+    asyncio.create_task(_running_server.serve())
+    await _running_server.wait_for_startup()
+
+    print(f'Python UDF registered at {base_url}')
+
+    return UdfConnectionInfo(base_url, app.get_function_info())
+
+
+def generate_base_url(app_config: AppConfig) -> str:
+    if not app_config.is_gateway_enabled:
+        raise RuntimeError('Python UDFs are not available if Nova Gateway is not enabled')
+
+    if not app_config.running_interactively:
+        return app_config.base_url
+
+    # generate python udf endpoint for interactive notebooks
+    gateway_url = os.environ.get('SINGLESTOREDB_NOVA_GATEWAY_ENDPOINT')
+    if app_config.is_local_dev:
+        gateway_url = os.environ.get('SINGLESTOREDB_NOVA_GATEWAY_DEV_ENDPOINT')
+        if gateway_url is None:
+            raise RuntimeError(
+                'Missing SINGLESTOREDB_NOVA_GATEWAY_DEV_ENDPOINT environment variable.',
+            )
+
+    return f'{gateway_url}/pythonudfs/{app_config.notebook_server_id}/interactive/'

singlestoredb/config.py

@@ -317,7 +317,7 @@ register_option(
     'external_function.app_mode', 'string',
     functools.partial(
         check_str,
-        valid_values=['remote', 'collocated'],
+        valid_values=['remote', 'collocated', 'managed'],
     ),
     'remote',
     'Specifies the mode of operation of the external function application.',
@@ -407,6 +407,18 @@ register_option(
     environ=['SINGLESTOREDB_EXT_FUNC_LOG_LEVEL'],
 )
 
+register_option(
+    'external_function.name_prefix', 'string', check_str, '',
+    'Prefix to add to external function names.',
+    environ=['SINGLESTOREDB_EXT_FUNC_NAME_PREFIX'],
+)
+
+register_option(
+    'external_function.name_suffix', 'string', check_str, '',
+    'Suffix to add to external function names.',
+    environ=['SINGLESTOREDB_EXT_FUNC_NAME_SUFFIX'],
+)
+
 register_option(
     'external_function.connection', 'string', check_str,
     os.environ.get('SINGLESTOREDB_URL') or None,
@@ -415,7 +427,7 @@ register_option(
 )
 
 register_option(
-    'external_function.host', 'string', check_str, '127.0.0.1',
+    'external_function.host', 'string', check_str, 'localhost',
     'Specifies the host to bind the server to.',
     environ=['SINGLESTOREDB_EXT_FUNC_HOST'],
 )

singlestoredb/functions/__init__.py

@@ -1,2 +1,16 @@
-from .decorator import tvf  # noqa: F401
 from .decorator import udf  # noqa: F401
+from .typing import Masked  # noqa: F401
+from .typing import Table  # noqa: F401
+from .utils import pack_vector  # noqa: F401
+from .utils import pack_vectors  # noqa: F401
+from .utils import unpack_vector  # noqa: F401
+from .utils import unpack_vectors  # noqa: F401
+from .utils import VectorTypes
+
+
+F32 = VectorTypes.F32
+F64 = VectorTypes.F64
+I8 = VectorTypes.I8
+I16 = VectorTypes.I16
+I32 = VectorTypes.I32
+I64 = VectorTypes.I64