Skip to content

Commit bcdb8fc

Browse files
committed
Refactor function names and improve data fetching
Renamed two functions in openai.py and optimized the data fetching in table_schema.py. Also updated an example script to reflect these changes. The renaming of methods make_mindsdb_tool to make_query_tool and litellm_text2sql_callback_tool to make_data_tool is done for clarity. The simplification of data fetching process by removing redundant steps enhances efficiency. Changes in the example script are made to adhere to these modifications.
1 parent 611d8d7 commit bcdb8fc

File tree

3 files changed

+10
-12
lines changed

3 files changed

+10
-12
lines changed

examples/using_mindsdb_inference_with_text2sql_using_tools.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,10 @@
22

33

44
from mindsdb_sdk.utils.openai import (
5-
make_mindsdb_tool,
6-
execute_function_call,
7-
chat_completion_request,
8-
pretty_print_conversation)
5+
make_query_tool,
6+
execute_function_call,
7+
chat_completion_request,
8+
pretty_print_conversation)
99

1010
import mindsdb_sdk
1111
import os
@@ -33,10 +33,10 @@
3333
api_key=OPENAI_API_KEY
3434
)
3535

36-
database = con.databases.get("example_db")
36+
database = con.databases.get("mindsdb_demo_db")
3737
schema = get_table_schemas(database, included_tables=["airline_passenger_satisfaction"])
3838

39-
tools = [make_mindsdb_tool(schema)]
39+
tools = [make_query_tool(schema)]
4040

4141
SYSTEM_PROMPT = """You are a SQL expert. Given an input question, Answer user questions by generating SQL queries
4242
against the database schema provided in tools

mindsdb_sdk/utils/openai.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ def make_openai_tool(function: callable, description: str = None) -> dict:
9696
return function_dict
9797

9898

99-
def make_mindsdb_tool(schema: dict) -> dict:
99+
def make_query_tool(schema: dict) -> dict:
100100
"""
101101
Make an OpenAI tool for querying a database connection in MindsDB
102102
@@ -128,14 +128,14 @@ def make_mindsdb_tool(schema: dict) -> dict:
128128
}
129129

130130

131-
def litellm_text2sql_callback_tool(
131+
def make_data_tool(
132132
model: str,
133133
data_source: str,
134134
description: str,
135135
connection_args: dict
136136
):
137137
"""
138-
tool passing connection details for datasource to litellm callback
138+
tool passing mindsdb database connection details for datasource to litellm callback
139139
140140
:param model: model name for text to sql completion
141141
:param data_source: data source name

mindsdb_sdk/utils/table_schema.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,9 +48,7 @@ def get_table_schemas(database: Databases, included_tables: List[str] = None, n_
4848

4949
table_schemas = {}
5050
for table in tables:
51-
# Get the first 10 rows of the table
52-
table_df = database.get_table(table).fetch().head(n_rows)
53-
# Convert schema to list of dictionaries
51+
table_df = database.get_table(table).limit(n_rows).fetch()
5452
table_schemas[table] = get_dataframe_schema(table_df)
5553

5654
return table_schemas

0 commit comments

Comments
 (0)