Enhance logging and error handling in openai module

This commit introduces better logging and error handling in the openai module. The print statements have been replaced with logger calls which provide more insight and control over the output. Additionally, a new unit test file (test_openai.py) has been created to test the 'openai' utilities. Lastly, 'docstring-parser' dependency was added to requirements.txt.

Author: dusvyat

mindsdb_sdk/utils/openai.py

@@ -1,11 +1,31 @@
 import json
+from logging import getLogger
+from typing import List
+
+import openai
+from openai.types.chat import ChatCompletionToolChoiceOptionParam
 
 from mindsdb_sdk.databases import Database
 from tenacity import retry, wait_random_exponential, stop_after_attempt
 
 
-@retry(wait=wait_random_exponential(multiplier=1, max=40), stop=stop_after_attempt(3))
-def chat_completion_request(client, model, messages, tools=None, tool_choice=None):
+DEFAULT_RETRY_MULTIPLIER = 1
+DEFAULT_MAX_WAIT = 40
+DEFAULT_STOP_AFTER_ATTEMPT = 3
+
+logger = getLogger(__name__)
+
+
+@retry(wait=wait_random_exponential(multiplier=DEFAULT_RETRY_MULTIPLIER, max=DEFAULT_MAX_WAIT), stop=stop_after_attempt(
+    DEFAULT_RETRY_MULTIPLIER
+))
+def chat_completion_request(
+        client: openai.OpenAI,
+        model: str,
+        messages: List[dict],
+        tools: List = None,
+        tool_choice: ChatCompletionToolChoiceOptionParam = None
+):
     try:
         response = client.chat.completions.create(
             model=model,
@@ -15,8 +35,8 @@ def chat_completion_request(client, model, messages, tools=None, tool_choice=Non
         )
         return response
     except Exception as e:
-        print("Unable to generate ChatCompletion response")
-        print(f"Exception: {e}")
+        logger.warning("Unable to generate ChatCompletion response")
+        logger.warning(f"Exception: {e}")
         return e
 
 
@@ -153,7 +173,7 @@ def litellm_text2sql_callback_tool(
                         "description":"Data source name",
                     },
                     "connection_args":{
-                        "type":"object",
+                        "type":"string",
                         "description":"Connection arguments for the data source",
                     },
                     "description":{
@@ -246,12 +266,12 @@ def pretty_print_conversation(messages):
 
     for message in messages:
         if message["role"] == "system":
-            print(colored(f"system: {message['content']}\n", role_to_color[message["role"]]))
+            logger.info(colored(f"system: {message['content']}\n", role_to_color[message["role"]]))
         elif message["role"] == "user":
-            print(colored(f"user: {message['content']}\n", role_to_color[message["role"]]))
+            logger.info(colored(f"user: {message['content']}\n", role_to_color[message["role"]]))
         elif message["role"] == "assistant" and message.get("function_call"):
-            print(colored(f"assistant: {message['function_call']}\n", role_to_color[message["role"]]))
+            logger.info(colored(f"assistant: {message['function_call']}\n", role_to_color[message["role"]]))
         elif message["role"] == "assistant" and not message.get("function_call"):
-            print(colored(f"assistant: {message['content']}\n", role_to_color[message["role"]]))
+            logger.info(colored(f"assistant: {message['content']}\n", role_to_color[message["role"]]))
         elif message["role"] == "function":
-            print(colored(f"function ({message['name']}): {message['content']}\n", role_to_color[message["role"]]))
+            logger.info(colored(f"function ({message['name']}): {message['content']}\n", role_to_color[message["role"]]))

requirements.txt

@@ -1,3 +1,4 @@
 requests
 pandas >= 1.3.5
 mindsdb-sql >= 0.13.0, < 0.14.0
+docstring-parser >= 0.7.3

tests/test_openai.py

@@ -0,0 +1,50 @@
+import json
+from unittest.mock import patch, MagicMock
+from mindsdb_sdk.utils import openai
+
+
+def test_chat_completion_request_success():
+    mock_client = MagicMock()
+    mock_client.chat.completions.create.return_value = "Test Response"
+    response = openai.chat_completion_request(mock_client, "text-davinci-002", [{"role": "system", "content": "You are a helpful assistant."}])
+    assert response == "Test Response"
+
+
+def test_make_openai_tool():
+    def test_func(a: int, b: str) -> str:
+        """This is a test function"""
+        return b * a
+    tool = openai.make_openai_tool(test_func)
+    assert tool["function"]["name"] == "test_func"
+    assert tool["function"]["description"] == "This is a test function"
+    assert tool["function"]["parameters"]["properties"]["a"]["type"] == "int"
+    assert tool["function"]["parameters"]["properties"]["b"]["type"] == "str"
+
+
+def test_extract_sql_query():
+    result = "SQLQuery: SELECT * FROM test_table\nSQLResult: [{'column1': 'value1', 'column2': 'value2'}]"
+    query = openai.extract_sql_query(result)
+    assert query == "SELECT * FROM test_table"
+
+
+def test_extract_sql_query_no_query():
+    result = "SQLResult: [{'column1': 'value1', 'column2': 'value2'}]"
+    query = openai.extract_sql_query(result)
+    assert query is None
+
+
+@patch("mindsdb_sdk.utils.openai.query_database")
+def test_execute_function_call_query_database(mock_query_database):
+    mock_query_database.return_value = "Test Result"
+    mock_message = MagicMock()
+    mock_message.tool_calls[0].function.name = "query_database"
+    mock_message.tool_calls[0].function.arguments = json.dumps({"query": "SELECT * FROM test_table"})
+    result = openai.execute_function_call(mock_message, MagicMock())
+    assert result == "Test Result"
+
+
+def test_execute_function_call_no_function():
+    mock_message = MagicMock()
+    mock_message.tool_calls[0].function.name = "non_existent_function"
+    result = openai.execute_function_call(mock_message, MagicMock())
+    assert result == "Error: function non_existent_function does not exist"