Skip to content

Commit 6718cb0

Browse files
committed
Merge branch 'support-openai-completion-with-toools' of https://github.com/mindsdb/mindsdb_python_sdk into staging
2 parents 49ea39a + 28658f7 commit 6718cb0

File tree

4 files changed

+479
-0
lines changed

4 files changed

+479
-0
lines changed
Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
from openai import OpenAI
2+
from mindsdb_sdk.utils.openai import extract_sql_query, query_database, chat_completion_request, \
3+
pretty_print_conversation
4+
5+
import mindsdb_sdk
6+
import os
7+
8+
from mindsdb_sdk.utils.table_schema import get_table_schemas
9+
10+
# generate the key at https://llm.mdb.ai
11+
MINDSDB_API_KEY = os.environ.get("MINDSDB_API_KEY")
12+
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
13+
14+
MODEL = "gpt-3.5-turbo"
15+
16+
# the prompt should be a question that can be answered by the database
17+
SYSTEM_PROMPT = """You are a SQL expert. Given an input question, first create a syntactically correct SQL query to run, then look at the results of the query and return the answer to the input question.
18+
Unless the user specifies in the question a specific number of examples to obtain, query for at most 5 results using the LIMIT clause as per SQL standards. You can order the results to return the most informative data in the database.
19+
Never query for all columns from a table. You must query only the columns that are needed to answer the question. Wrap each column name in backticks (`) to denote them as identifiers.
20+
Pay attention to use only the column names you can see in the tables below. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.
21+
Pay attention to use CURRENT_DATE function to get the current date, if the question involves "today".
22+
23+
Use the following format:
24+
25+
Question: <Question here>
26+
SQLQuery: <SQL Query to run>
27+
SQLResult: <Result of the SQLQuery>
28+
Answer: <Final answer here>
29+
30+
Only use the following tables:
31+
32+
{schema}
33+
"""
34+
PROMPT = "what was the average delay on arrivals?"
35+
36+
37+
def generate_system_prompt(system_prompt: str, schema: dict) -> dict:
38+
prompt = {
39+
"role":"system",
40+
"content":system_prompt.format(schema=schema)
41+
}
42+
return prompt
43+
44+
45+
def generate_user_prompt(query: str) -> dict:
46+
prompt = {
47+
"role":"user",
48+
"content":query
49+
}
50+
return prompt
51+
52+
53+
con = mindsdb_sdk.connect()
54+
55+
# given database name, returns schema and database object
56+
# using example_db from mindsdb
57+
58+
database = con.databases.get("example_db")
59+
schema = get_table_schemas(database, included_tables=["airline_passenger_satisfaction"])
60+
61+
# client_mindsdb_serve = OpenAI(
62+
# api_key=MINDSDB_API_KEY,
63+
# base_url="https://llm.mdb.ai"
64+
# )
65+
66+
client_mindsdb_serve = OpenAI(
67+
api_key=OPENAI_API_KEY
68+
)
69+
70+
messages = [
71+
generate_system_prompt(SYSTEM_PROMPT, schema),
72+
generate_user_prompt(PROMPT)
73+
]
74+
75+
chat_response = chat_completion_request(client=client_mindsdb_serve, model=MODEL, messages=messages, tools=None,
76+
tool_choice=None)
77+
78+
# extract the SQL query from the response
79+
query = extract_sql_query(chat_response.choices[0].message.content)
80+
81+
result = query_database(database, query)
82+
83+
# generate the user prompt with the query result, this will be used to generate the final response
84+
query = generate_user_prompt(f"Given this SQLResult: {str(result)} provide Answer: ")
85+
86+
# add the query to the messages list
87+
messages.append(query)
88+
89+
# generate the final response
90+
chat_completion_gpt = chat_completion_request(client=client_mindsdb_serve, model=MODEL, messages=messages, tools=None,
91+
tool_choice=None)
92+
response = chat_completion_gpt.choices[0].message.content
93+
94+
pretty_print_conversation(messages)
95+
Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
from openai import OpenAI
2+
3+
4+
from mindsdb_sdk.utils.openai import (
5+
make_mindsdb_tool,
6+
execute_function_call,
7+
chat_completion_request,
8+
pretty_print_conversation)
9+
10+
import mindsdb_sdk
11+
import os
12+
13+
from mindsdb_sdk.utils.table_schema import get_table_schemas
14+
15+
# generate the key at https://llm.mdb.ai
16+
MINDSDB_API_KEY = os.environ.get("MINDSDB_API_KEY")
17+
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
18+
19+
MODEL = "gpt-3.5-turbo"
20+
21+
22+
con = mindsdb_sdk.connect()
23+
24+
# given database name, returns schema and database object
25+
# using example_db from mindsdb
26+
27+
# client_mindsdb_serve = OpenAI(
28+
# api_key=MINDSDB_API_KEY,
29+
# base_url="https://llm.mdb.ai"
30+
# )
31+
32+
client_mindsdb_serve = OpenAI(
33+
api_key=OPENAI_API_KEY
34+
)
35+
36+
database = con.databases.get("example_db")
37+
schema = get_table_schemas(database, included_tables=["airline_passenger_satisfaction"])
38+
39+
tools = [make_mindsdb_tool(schema)]
40+
41+
SYSTEM_PROMPT = """You are a SQL expert. Given an input question, Answer user questions by generating SQL queries
42+
against the database schema provided in tools
43+
Unless the user specifies in the question a specific number of examples to obtain, query for at most 5 results using the
44+
LIMIT clause as per SQL standards. You can order the results to return the most informative data in the database.
45+
Never query for all columns from a table. You must query only the columns that are needed to answer the question.
46+
Wrap each column name in backticks (`) to denote them as identifiers.
47+
Pay attention to use only the column names you can see in the tables below.
48+
Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.
49+
Pay attention to use CURRENT_DATE function to get the current date, if the question involves "today"."""
50+
51+
messages = [{
52+
"role":"system", "content":SYSTEM_PROMPT
53+
}, {"role":"user", "content":"what was the average delay on arrivals?"}]
54+
55+
chat_response = chat_completion_request(client=client_mindsdb_serve, model=MODEL, messages=messages, tools=tools, tool_choice=None)
56+
57+
assistant_message = chat_response.choices[0].message
58+
59+
assistant_message.content = str(assistant_message.tool_calls[0].function)
60+
61+
messages.append({"role": assistant_message.role, "content": assistant_message.content})
62+
63+
if assistant_message.tool_calls:
64+
results = execute_function_call(message=assistant_message, database=database)
65+
messages.append({
66+
"role": "function", "tool_call_id": assistant_message.tool_calls[0].id,
67+
"name": assistant_message.tool_calls[0].function.name,
68+
"content": results
69+
})
70+
71+
pretty_print_conversation(messages)

0 commit comments

Comments
 (0)