use openai api for openai types of requests, making it compatible with some other providers

Author: Chenglong-MS

py-src/data_formulator/agents/client_utils.py

@@ -1,5 +1,6 @@
 import os
-from litellm import completion
+import litellm
+import openai
 from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 
 class Client(object):
@@ -14,11 +15,17 @@ def __init__(self, endpoint, model, api_key=None,  api_base=None, api_version=No
 
         # other params, including temperature, max_completion_tokens, api_base, api_version
         self.params = {
-            "api_key": api_key,
             "temperature": 0.7,
             "max_completion_tokens": 1200,
         }
 
+        if api_key is not None and api_key != "":
+            self.params["api_key"] = api_key
+        if api_base is not None and api_base != "":
+            self.params["api_base"] = api_base
+        if api_version is not None and api_version != "":
+            self.params["api_version"] = api_version
+
         if self.endpoint == "gemini":
             if model.startswith("gemini/"):
                 self.model = model
@@ -52,9 +59,19 @@ def get_completion(self, messages):
         Supports OpenAI, Azure, Ollama, and other providers via LiteLLM.
         """
         # Configure LiteLLM 
-        return completion(
-            model=self.model,
-            messages=messages,
-            drop_params=True,
-            **self.params
-        )
+
+        if self.endpoint == "openai":
+            client = openai.OpenAI(api_key=self.params["api_key"], base_url=self.params["api_base"] if "api_base" in self.params else None)
+
+            return client.chat.completions.create(
+                model=self.model,
+                messages=messages,
+                **{k: v for k, v in self.params.items() if k != "api_key"}
+            )
+        else:
+            return litellm.completion(
+                model=self.model,
+                messages=messages,
+                drop_params=True,
+                **self.params
+            )

src/views/ModelSelectionDialog.tsx

@@ -181,7 +181,7 @@ export const ModelSelectionButton: React.FC<{}> = ({ }) => {
                 PaperComponent={({ children }) => (
                     <Paper>
                         <Typography sx={{ p: 1, color: 'gray', fontStyle: 'italic', fontSize: '0.75rem' }}>
-                            suggestions
+                            examples
                         </Typography>
                         {children}
                     </Paper>
@@ -226,7 +226,7 @@ export const ModelSelectionButton: React.FC<{}> = ({ }) => {
                 PaperComponent={({ children }) => (
                     <Paper>
                         <Typography sx={{ p: 1, color: 'gray', fontStyle: 'italic', fontSize: 'small' }}>
-                            suggestions
+                            examples
                         </Typography>
                         {children}
                     </Paper>
@@ -267,23 +267,17 @@ export const ModelSelectionButton: React.FC<{}> = ({ }) => {
 
                         let model = {endpoint, model: newModel, api_key: newApiKey, api_base: newApiBase, api_version: newApiVersion, id: id};
 
-                        console.log("checkpont 2")
-
                         dispatch(dfActions.addModel(model));
                         dispatch(dfActions.selectModel(id));
                         setTempSelectedModeId(id);
 
-                        console.log("checkpont 3")
-
                         testModel(model); 
 
                         setNewEndpoint("");
                         setNewModel("");
                         setNewApiKey(undefined);
                         setNewApiBase(undefined);
                         setNewApiVersion(undefined);
-
-                        console.log("checkpont 4")
                     }}>
                     <AddCircleIcon />
                 </IconButton>