supporting multiple LLM backend (require a little more testing)

Author: Chenglong-MS

py-src/data_formulator/agents/agent_code_explanation.py

@@ -66,9 +66,8 @@ def transform_data(df_0):
 
 class CodeExplanationAgent(object):
 
-    def __init__(self, client, model):
+    def __init__(self, client):
         self.client = client
-        self.model = model
 
     def run(self, input_tables, code):
 
@@ -82,9 +81,7 @@ def run(self, input_tables, code):
                     {"role":"user","content": user_query}]
 
         ###### the part that calls open_ai
-        response = self.client.chat.completions.create(
-            model=self.model, messages = messages, temperature=0.7, max_tokens=1200,
-            top_p=0.95, n=1, frequency_penalty=0, presence_penalty=0, stop=None)
+        response = self.client.get_completion(messages = messages)
 
         logger.info('\n=== explanation output ===>\n')
         logger.info(response.choices[0].message.content)

py-src/data_formulator/agents/agent_concept_derive.py

@@ -167,9 +167,8 @@
 
 class ConceptDeriveAgent(object):
 
-    def __init__(self, client, model):
+    def __init__(self, client):
         self.client = client
-        self.model = model
 
     def run(self, input_table, input_fields, output_field, description, n=1):
         """derive a new concept based on input table, input fields, and output field name, (and description)
@@ -190,9 +189,7 @@ def run(self, input_table, input_fields, output_field, description, n=1):
                     {"role":"user","content": user_query}]
 
         ###### the part that calls open_ai
-        response = self.client.chat.completions.create(
-            model=self.model, messages = messages, temperature=0.7, max_tokens=1200,
-            top_p=0.95, n=n, frequency_penalty=0, presence_penalty=0, stop=None)
+        response = self.client.get_completion(messages = messages)
 
         #log = {'messages': messages, 'response': response.model_dump(mode='json')}
 

py-src/data_formulator/agents/agent_data_clean.py

@@ -78,8 +78,7 @@
 
 class DataCleanAgent(object):
 
-    def __init__(self, client, model):
-        self.model = model
+    def __init__(self, client):
         self.client = client
 
     def run(self, content_type, raw_data, image_cleaning_instruction):
@@ -129,9 +128,7 @@ def run(self, content_type, raw_data, image_cleaning_instruction):
         messages = [system_message, user_prompt]
 
         ###### the part that calls open_ai
-        response = self.client.chat.completions.create(
-            model=self.model, messages = messages, temperature=0.7, max_tokens=1200,
-            top_p=0.95, n=1, frequency_penalty=0, presence_penalty=0, stop=None)
+        response = self.client.get_completion(messages = messages)
 
         candidates = []
         for choice in response.choices:

py-src/data_formulator/agents/agent_data_filter.py

@@ -125,9 +125,8 @@ def filter_row(row, df):
 
 class DataFilterAgent(object):
 
-    def __init__(self, client, model):
+    def __init__(self, client):
         self.client = client
-        self.model = model
 
     def process_gpt_result(self, input_table, response, messages):
         #log = {'messages': messages, 'response': response.model_dump(mode='json')}
@@ -177,9 +176,7 @@ def run(self, input_table, description):
                     {"role":"user","content": user_query}]
 
         ###### the part that calls open_ai
-        response = self.client.chat.completions.create(
-            model=self.model, messages = messages, temperature=0.7, max_tokens=1200,
-            top_p=0.95, n=1, frequency_penalty=0, presence_penalty=0, stop=None)
+        response = self.client.get_completion(messages = messages)
 
         return self.process_gpt_result(input_table, response, messages)
 
@@ -190,8 +187,6 @@ def followup(self, input_table, dialog, new_instruction: str, n=1):
                               "content": new_instruction + '\nupdate the filter function accordingly'}]
 
         ##### the part that calls open_ai
-        response = self.client.chat.completions.create(
-            model=self.model, messages=messages, temperature=0.7, max_tokens=1200,
-            top_p=0.95, n=n, frequency_penalty=0, presence_penalty=0, stop=None)
+        response = self.client.get_completion(messages = messages)
 
         return self.process_gpt_result(input_table, response, messages)

py-src/data_formulator/agents/agent_data_load.py

@@ -124,9 +124,8 @@
 
 class DataLoadAgent(object):
 
-    def __init__(self, client, model):
+    def __init__(self, client):
         self.client = client
-        self.model = model
 
     def run(self, input_data, n=1):
 
@@ -140,9 +139,7 @@ def run(self, input_data, n=1):
                     {"role":"user","content": user_query}]
 
         ###### the part that calls open_ai
-        response = self.client.chat.completions.create(
-            model=self.model, messages=messages, temperature=0.2, max_tokens=4096,
-            top_p=0.95, n=n, frequency_penalty=0, presence_penalty=0, stop=None)
+        response = self.client.get_completion(messages = messages)
 
         #log = {'messages': messages, 'response': response.model_dump(mode='json')}
 

py-src/data_formulator/agents/agent_data_rec.py

@@ -126,9 +126,8 @@ def transform_data(df):
 
 class DataRecAgent(object):
 
-    def __init__(self, client, model, system_prompt=None):
+    def __init__(self, client, system_prompt=None):
         self.client = client
-        self.model = model
         self.system_prompt = system_prompt if system_prompt is not None else SYSTEM_PROMPT
 
     def process_gpt_response(self, input_tables, messages, response):
@@ -192,7 +191,7 @@ def run(self, input_tables, description, n=1):
         messages = [{"role":"system", "content": self.system_prompt},
                     {"role":"user","content": user_query}]
 
-        response = completion_response_wrapper(self.client, self.model, messages, n)
+        response = completion_response_wrapper(self.client, messages, n)
 
         return self.process_gpt_response(input_tables, messages, response)
 
@@ -204,7 +203,6 @@ def followup(self, input_tables, dialog, new_instruction: str, n=1):
 
         messages = [*dialog, {"role":"user", "content": f"Update: \n\n{new_instruction}"}]
 
-        ##### the part that calls open_ai
-        response = completion_response_wrapper(self.client, self.model, messages, n)
+        response = completion_response_wrapper(self.client, messages, n)
 
         return self.process_gpt_response(input_tables, messages, response)

py-src/data_formulator/agents/agent_data_transform_v2.py

@@ -178,12 +178,10 @@ def transform_data(df):
 ```
 '''
 
-def completion_response_wrapper(client, model, messages, n):
+def completion_response_wrapper(client, messages, n):
     ### wrapper for completion response, especially handling errors
     try:
-        response = client.chat.completions.create(
-                model=model, messages=messages, temperature=0.7, max_tokens=1200,
-                top_p=0.95, n=n, frequency_penalty=0, presence_penalty=0, stop=None)
+        response = client.get_completion(messages = messages)
     except Exception as e:
         response = e
 
@@ -192,9 +190,8 @@ def completion_response_wrapper(client, model, messages, n):
 
 class DataTransformationAgentV2(object):
 
-    def __init__(self, client, model, system_prompt=None):
+    def __init__(self, client, system_prompt=None):
         self.client = client
-        self.model = model
         self.system_prompt = system_prompt if system_prompt is not None else SYSTEM_PROMPT
 
     def process_gpt_response(self, input_tables, messages, response):
@@ -265,7 +262,7 @@ def run(self, input_tables, description, expected_fields: list[str], n=1):
         messages = [{"role":"system", "content": self.system_prompt},
                     {"role":"user","content": user_query}]
 
-        response = completion_response_wrapper(self.client, self.model, messages, n)
+        response = completion_response_wrapper(self.client, messages, n)
 
         return self.process_gpt_response(input_tables, messages, response)
 
@@ -287,6 +284,6 @@ def followup(self, input_tables, dialog, output_fields: list[str], new_instructi
         messages = [*updated_dialog, {"role":"user", 
                               "content": f"Update the code above based on the following instruction:\n\n{json.dumps(goal, indent=4)}"}]
 
-        response = completion_response_wrapper(self.client, self.model, messages, n)
+        response = completion_response_wrapper(self.client, messages, n)
 
         return self.process_gpt_response(input_tables, messages, response)

py-src/data_formulator/agents/agent_data_transformation.py

@@ -122,9 +122,8 @@ def transform_data(df_0):
 
 class DataTransformationAgent(object):
 
-    def __init__(self, client, model):
+    def __init__(self, client):
         self.client = client
-        self.model = model
 
     def process_gpt_response(self, input_tables, messages, response):
         """process gpt response to handle execution"""
@@ -185,9 +184,7 @@ def run(self, input_tables, description, expected_fields: list[str], n=1, enrich
                     {"role":"user","content": user_query}]
 
         ###### the part that calls open_ai
-        response = self.client.chat.completions.create(
-            model=self.model, messages = messages, temperature=0.7, max_tokens=1200,
-            top_p=0.95, n=n, frequency_penalty=0, presence_penalty=0, stop=None)
+        response = self.client.get_completion(messages = messages)
 
         return self.process_gpt_response(input_tables, messages, response)
 
@@ -207,9 +204,7 @@ def followup(self, input_tables, dialog, output_fields: list[str], new_instructi
                               "content": "Update the code above based on the following instruction:\n\n" + new_instruction + output_fields_instr}]
 
         ##### the part that calls open_ai
-        response = self.client.chat.completions.create(
-            model=self.model, messages=messages, temperature=0.7, max_tokens=1200,
-            top_p=0.95, n=n, frequency_penalty=0, presence_penalty=0, stop=None)
+        response = self.client.get_completion(messages = messages)
 
         logger.info(response)
 

py-src/data_formulator/agents/agent_generic_py_concept.py

@@ -157,9 +157,8 @@ def derive(row, df):
 
 class GenericPyConceptDeriveAgent(object):
 
-    def __init__(self, client, model_version):
+    def __init__(self, client):
         self.client = client
-        self.model_version = model_version
 
     def process_gpt_response(self, input_table, output_field, response, messages):
         #log = {'messages': messages, 'response': response.model_dump(mode='json')}
@@ -220,10 +219,7 @@ def run(self, input_table, output_field, description):
                     {"role":"user","content": user_query}]
 
         ###### the part that calls open_ai
-        response = self.client.chat.completions.create(
-            model=self.model_version, 
-            messages = messages, temperature=0.7, max_tokens=1200,
-            top_p=0.95, n=1, frequency_penalty=0, presence_penalty=0, stop=None)
+        response = self.client.get_completion(messages = messages)
 
         return self.process_gpt_response(input_table, output_field, response, messages)
 
@@ -234,9 +230,7 @@ def followup(self, input_table, dialog, output_field: str, new_instruction: str,
                               "content": new_instruction + '\n update the function accordingly'}]
 
         ##### the part that calls open_ai
-        response = self.client.chat.completions.create(
-            model=self.model, messages=messages, temperature=0.7, max_tokens=1200,
-            top_p=0.95, n=n, frequency_penalty=0, presence_penalty=0, stop=None)
+        response = self.client.get_completion(messages = messages)
 
         candidates = self.process_gpt_response(input_table, output_field, response, messages)
 

py-src/data_formulator/agents/agent_py_concept_derive.py

@@ -131,8 +131,7 @@ def derive(writing, reading, math):
 
 class PyConceptDeriveAgent(object):
 
-    def __init__(self, client, model):
-        self.model = model
+    def __init__(self, client):
         self.client = client
 
     def run(self, input_table, input_fields, output_field, description):
@@ -163,9 +162,7 @@ def derive({arg_string}):
                     {"role":"user","content": user_query}]
 
         ###### the part that calls open_ai
-        response = self.client.chat.completions.create(
-            model=self.model, messages = messages, temperature=0.7, max_tokens=1200,
-            top_p=0.95, n=1, frequency_penalty=0, presence_penalty=0, stop=None)
+        response = self.client.get_completion(messages = messages)
 
         #log = {'messages': messages, 'response': response.model_dump(mode='json')}