Merge branch 'staging' of https://github.com/mindsdb/mindsdb_python_sdk into staging

Author: tmichaeldb

examples/using_mindsdb_inference_with_text2sql_using_tools.py

@@ -2,10 +2,10 @@
 
 
 from mindsdb_sdk.utils.openai import (
-                                      make_mindsdb_tool,
-                                      execute_function_call,
-                                      chat_completion_request,
-                                      pretty_print_conversation)
+    make_query_tool,
+    execute_function_call,
+    chat_completion_request,
+    pretty_print_conversation)
 
 import mindsdb_sdk
 import os
@@ -33,10 +33,10 @@
     api_key=OPENAI_API_KEY
 )
 
-database = con.databases.get("example_db")
+database = con.databases.get("mindsdb_demo_db")
 schema = get_table_schemas(database, included_tables=["airline_passenger_satisfaction"])
 
-tools = [make_mindsdb_tool(schema)]
+tools = [make_query_tool(schema)]
 
 SYSTEM_PROMPT = """You are a SQL expert. Given an input question, Answer user questions by generating SQL queries 
 against the database schema provided in tools 

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
 
 
@@ -29,7 +49,6 @@ def make_openai_tool(function: callable, description: str = None) -> dict:
 
     :return: dictionary containing function metadata
     """
-    # You will need to pip install docstring-parser to use this function
 
     import inspect
     import docstring_parser
@@ -77,7 +96,7 @@ def make_openai_tool(function: callable, description: str = None) -> dict:
     return function_dict
 
 
-def make_mindsdb_tool(schema: dict) -> dict:
+def make_query_tool(schema: dict) -> dict:
     """
     Make an OpenAI tool for querying a database connection in MindsDB
 
@@ -109,14 +128,14 @@ def make_mindsdb_tool(schema: dict) -> dict:
     }
 
 
-def litellm_text2sql_callback_tool(
+def make_data_tool(
     model: str,
     data_source: str,
     description: str,
     connection_args: dict
 ):
     """
-    tool passing connection details for datasource to litellm callback
+    tool passing mindsdb database connection details for datasource to litellm callback
 
     :param model: model name for text to sql completion
     :param data_source: data source name
@@ -153,7 +172,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 +265,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"]]))

mindsdb_sdk/utils/table_schema.py

@@ -48,9 +48,7 @@ def get_table_schemas(database: Databases, included_tables: List[str] = None, n_
 
     table_schemas = {}
     for table in tables:
-        # Get the first 10 rows of the table
-        table_df = database.get_table(table).fetch().head(n_rows)
-        # Convert schema to list of dictionaries
+        table_df = database.get_table(table).limit(n_rows).fetch()
         table_schemas[table] = get_dataframe_schema(table_df)
 
     return table_schemas

requirements.txt

@@ -1,3 +1,6 @@
 requests
 pandas >= 1.3.5
 mindsdb-sql >= 0.13.0, < 0.14.0
+docstring-parser >= 0.7.3
+tenacity >= 8.0.1
+openai >= 1.15.0

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"