Final fixes.

Author: mgiannakopoulos

singlestoredb/ai/chat.py

@@ -32,35 +32,16 @@
 
 def SingleStoreChatFactory(
     model_name: str,
-    api_key: Optional[Union[Optional[str], Callable[[], Optional[str]]]] = None,
+    api_key: Optional[str] = None,
     streaming: bool = True,
     http_client: Optional[httpx.Client] = None,
-    obo_token: Optional[Union[Optional[str], Callable[[], Optional[str]]]] = None,
     obo_token_getter: Optional[Callable[[], Optional[str]]] = None,
     base_url: Optional[str] = None,
     hosting_platform: Optional[str] = None,
     **kwargs: Any,
 ) -> Union[ChatOpenAI, ChatBedrockConverse]:
     """Return a chat model instance (ChatOpenAI or ChatBedrockConverse).
     """
-    # Handle api_key and obo_token as callable functions
-    if callable(api_key):
-        api_key_getter_fn = api_key
-    else:
-        def api_key_getter_fn() -> Optional[str]:
-            if api_key is None:
-                return os.environ.get('SINGLESTOREDB_USER_TOKEN')
-            return api_key
-
-    if obo_token_getter is not None:
-        obo_token_getter_fn = obo_token_getter
-    else:
-        if callable(obo_token):
-            obo_token_getter_fn = obo_token
-        else:
-            def obo_token_getter_fn() -> Optional[str]:
-                return obo_token
-
     # handle model info
     if base_url is None:
         base_url = os.environ.get('SINGLESTOREDB_INFERENCE_API_BASE_URL')
@@ -126,12 +107,12 @@ def obo_token_getter_fn() -> Optional[str]:
 
         def _inject_headers(request: Any, **_ignored: Any) -> None:
             """Inject dynamic auth/OBO headers prior to Bedrock sending."""
-            if api_key_getter_fn is not None:
-                token_val = api_key_getter_fn()
-                if token_val:
-                    request.headers['Authorization'] = f'Bearer {token_val}'
-            if obo_token_getter_fn is not None:
-                obo_val = obo_token_getter_fn()
+            token_env_val = os.environ.get('SINGLESTOREDB_USER_TOKEN')
+            token_val = api_key if api_key is not None else token_env_val
+            if token_val:
+                request.headers['Authorization'] = f'Bearer {token_val}'
+            if obo_token_getter is not None:
+                obo_val = obo_token_getter()
                 if obo_val:
                     request.headers['X-S2-OBO'] = obo_val
             request.headers.pop('X-Amz-Date', None)
@@ -167,7 +148,8 @@ def _inject_headers(request: Any, **_ignored: Any) -> None:
         )
 
     # OpenAI / Azure OpenAI path
-    token = api_key_getter_fn()
+    token_env = os.environ.get('SINGLESTOREDB_USER_TOKEN')
+    token = api_key if api_key is not None else token_env
 
     openai_kwargs = dict(
         base_url=info.connection_url,

singlestoredb/ai/embeddings.py

@@ -32,34 +32,15 @@
 
 def SingleStoreEmbeddingsFactory(
     model_name: str,
-    api_key: Optional[Union[Optional[str], Callable[[], Optional[str]]]] = None,
+    api_key: Optional[str] = None,
     http_client: Optional[httpx.Client] = None,
-    obo_token: Optional[Union[Optional[str], Callable[[], Optional[str]]]] = None,
     obo_token_getter: Optional[Callable[[], Optional[str]]] = None,
     base_url: Optional[str] = None,
     hosting_platform: Optional[str] = None,
     **kwargs: Any,
 ) -> Union[OpenAIEmbeddings, BedrockEmbeddings]:
     """Return an embeddings model instance (OpenAIEmbeddings or BedrockEmbeddings).
     """
-    # Handle api_key and obo_token as callable functions
-    if callable(api_key):
-        api_key_getter_fn = api_key
-    else:
-        def api_key_getter_fn() -> Optional[str]:
-            if api_key is None:
-                return os.environ.get('SINGLESTOREDB_USER_TOKEN')
-            return api_key
-
-    if obo_token_getter is not None:
-        obo_token_getter_fn = obo_token_getter
-    else:
-        if callable(obo_token):
-            obo_token_getter_fn = obo_token
-        else:
-            def obo_token_getter_fn() -> Optional[str]:
-                return obo_token
-
     # handle model info
     if base_url is None:
         base_url = os.environ.get('SINGLESTOREDB_INFERENCE_API_BASE_URL')
@@ -125,12 +106,12 @@ def obo_token_getter_fn() -> Optional[str]:
 
         def _inject_headers(request: Any, **_ignored: Any) -> None:
             """Inject dynamic auth/OBO headers prior to Bedrock sending."""
-            if api_key_getter_fn is not None:
-                token_val = api_key_getter_fn()
-                if token_val:
-                    request.headers['Authorization'] = f'Bearer {token_val}'
-            if obo_token_getter_fn is not None:
-                obo_val = obo_token_getter_fn()
+            token_env_val = os.environ.get('SINGLESTOREDB_USER_TOKEN')
+            token_val = api_key if api_key is not None else token_env_val
+            if token_val:
+                request.headers['Authorization'] = f'Bearer {token_val}'
+            if obo_token_getter is not None:
+                obo_val = obo_token_getter()
                 if obo_val:
                     request.headers['X-S2-OBO'] = obo_val
             request.headers.pop('X-Amz-Date', None)
@@ -157,7 +138,8 @@ def _inject_headers(request: Any, **_ignored: Any) -> None:
         )
 
     # OpenAI / Azure OpenAI path
-    token = api_key_getter_fn()
+    token_env = os.environ.get('SINGLESTOREDB_USER_TOKEN')
+    token = api_key if api_key is not None else token_env
 
     openai_kwargs = dict(
         base_url=info.connection_url,