Merge branch 'mindsdb:main' into main

Author: parthiv11

docs/mint.json

@@ -522,7 +522,7 @@
           "pages": [
             "mindsdb_sql/functions/custom_functions",
             "mindsdb_sql/functions/llm_function",
-            "mindsdb_sql/functions/to_markdown"
+            "mindsdb_sql/functions/to_markdown_function"
           ]
         },
         {

mindsdb/integrations/handlers/openai_handler/helpers.py

@@ -4,7 +4,6 @@
 import math
 
 import openai
-from openai import OpenAI
 
 import tiktoken
 
@@ -181,17 +180,16 @@ def count_tokens(messages: List[Dict], encoder: tiktoken.core.Encoding, model_na
         )
 
 
-def get_available_models(api_key: Text, api_base: Text) -> List[Text]:
+def get_available_models(client) -> List[Text]:
     """
     Returns a list of available openai models for the given API key.
 
     Args:
-        api_key (Text): OpenAI API key
-        api_base (Text): OpenAI API base URL
+        client: openai sdk client
 
     Returns:
         List[Text]: List of available models
     """
-    res = OpenAI(api_key=api_key, base_url=api_base).models.list()
+    res = client.models.list()
 
     return [models.id for models in res.data]

mindsdb/integrations/handlers/openai_handler/openai_handler.py

@@ -9,7 +9,7 @@
 import concurrent.futures
 from typing import Text, Tuple, Dict, List, Optional, Any
 import openai
-from openai import OpenAI, NotFoundError, AuthenticationError
+from openai import OpenAI, AzureOpenAI, NotFoundError, AuthenticationError
 import numpy as np
 import pandas as pd
 
@@ -87,7 +87,7 @@ def create_engine(self, connection_args: Dict) -> None:
         if api_key is not None:
             org = connection_args.get('api_organization')
             api_base = connection_args.get('api_base') or os.environ.get('OPENAI_API_BASE', OPENAI_API_BASE)
-            client = self._get_client(api_key=api_key, base_url=api_base, org=org)
+            client = self._get_client(api_key=api_key, base_url=api_base, org=org, args=connection_args)
             OpenAIHandler._check_client_connection(client)
 
     @staticmethod
@@ -188,7 +188,9 @@ def create_validation(target: Text, args: Dict = None, **kwargs: Any) -> None:
                 "temperature",
                 "openai_api_key",
                 "api_organization",
-                "api_base"
+                "api_base",
+                "api_version",
+                "provider",
             }
         )
 
@@ -204,7 +206,7 @@ def create_validation(target: Text, args: Dict = None, **kwargs: Any) -> None:
         api_key = get_api_key('openai', args, engine_storage=engine_storage)
         api_base = args.get('api_base') or connection_args.get('api_base') or os.environ.get('OPENAI_API_BASE', OPENAI_API_BASE)
         org = args.get('api_organization')
-        client = OpenAIHandler._get_client(api_key=api_key, base_url=api_base, org=org)
+        client = OpenAIHandler._get_client(api_key=api_key, base_url=api_base, org=org, args=args)
         OpenAIHandler._check_client_connection(client)
 
     def create(self, target, args: Dict = None, **kwargs: Any) -> None:
@@ -228,7 +230,8 @@ def create(self, target, args: Dict = None, **kwargs: Any) -> None:
             api_key = get_api_key(self.api_key_name, args, self.engine_storage)
             connection_args = self.engine_storage.get_connection_args()
             api_base = args.get('api_base') or connection_args.get('api_base') or os.environ.get('OPENAI_API_BASE') or self.api_base
-            available_models = get_available_models(api_key, api_base)
+            client = self._get_client(api_key=api_key, base_url=api_base, org=args.get('api_organization'), args=args)
+            available_models = get_available_models(client)
 
             if not args.get('mode'):
                 args['mode'] = self.default_mode
@@ -810,6 +813,7 @@ def _tidy(comp: List[openai.types.image.Image]) -> List[Text]:
             api_key=api_key,
             base_url=args.get('api_base'),
             org=args.pop('api_organization') if 'api_organization' in args else None,
+            args=args
         )
 
         try:
@@ -891,7 +895,8 @@ def describe(self, attribute: Optional[Text] = None) -> pd.DataFrame:
                 client = self._get_client(
                     api_key=api_key,
                     base_url=args.get('api_base'),
-                    org=args.get('api_organization')
+                    org=args.get('api_organization'),
+                    args=args,
                 )
                 meta = client.models.retrieve(model_name)
             except Exception as e:
@@ -935,7 +940,7 @@ def finetune(self, df: Optional[pd.DataFrame] = None, args: Optional[Dict] = Non
 
         api_base = using_args.get('api_base', os.environ.get('OPENAI_API_BASE', OPENAI_API_BASE))
         org = using_args.get('api_organization')
-        client = self._get_client(api_key=api_key, base_url=api_base, org=org)
+        client = self._get_client(api_key=api_key, base_url=api_base, org=org, args=args)
 
         args = {**using_args, **args}
         prev_model_name = self.base_model_storage.json_get('args').get('model_name', '')
@@ -1173,7 +1178,7 @@ def _check_ft_status(job_id: Text) -> openai.types.fine_tuning.FineTuningJob:
         return ft_stats, result_file_id
 
     @staticmethod
-    def _get_client(api_key: Text, base_url: Text, org: Optional[Text] = None) -> OpenAI:
+    def _get_client(api_key: Text, base_url: Text, org: Optional[Text] = None, args: dict = None) -> OpenAI:
         """
         Get an OpenAI client with the given API key, base URL, and organization.
 
@@ -1185,4 +1190,11 @@ def _get_client(api_key: Text, base_url: Text, org: Optional[Text] = None) -> Op
         Returns:
             openai.OpenAI: OpenAI client.
         """
+        if args is not None and args.get('provider') == 'azure':
+            return AzureOpenAI(
+                api_key=api_key,
+                azure_endpoint=base_url,
+                api_version=args.get('api_version'),
+                organization=org
+            )
         return OpenAI(api_key=api_key, base_url=base_url, organization=org)