Merge pull request #116 from mindsdb/support-openai-completion-with-toools

Refactor function names and improve data fetching

Author: dusvyat

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

@@ -96,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
 
@@ -128,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

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