CakeCN/hsa
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

init

Browse Source
This commit is contained in:
CakeCN
2024-12-03 16:21:19 +08:00
commit 4198ca63b1
975 changed files with 333413 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
AlgoriAgent/examples/conversation_nl2sql/configs/model_configs.json Normal file
View File

@@ -0,0 +1,22 @@
[
{
"config_name": "gpt-3.5-turbo",
"model_type": "openai-chat",
"model_name": "gpt-3.5-turbo",
"api_key": "xxx",
"organization": "xxx",
"generate_args": {
"temperature": 0.0
}
},
{
"config_name": "gpt-4",
"model_type": "openai-chat",
"model_name": "gpt-4",
"api_key": "xxx",
"organization": "xxx",
"generate_args": {
"temperature": 0.5
}
}
]

85
AlgoriAgent/examples/conversation_nl2sql/database/concert_singer/schema.sql Normal file
View File

@@ -0,0 +1,85 @@
PRAGMA foreign_keys = ON;
CREATE TABLE "stadium" (
"Stadium_ID" int,
"Location" text,
"Name" text,
"Capacity" int,
"Highest" int,
"Lowest" int,
"Average" int,
PRIMARY KEY ("Stadium_ID")
);
INSERT INTO "stadium" VALUES (1,"Raith Rovers","Stark's Park","10104","4812","1294","2106");
INSERT INTO "stadium" VALUES (2,"Ayr United","Somerset Park","11998","2363","1057","1477");
INSERT INTO "stadium" VALUES (3,"East Fife","Bayview Stadium","2000","1980","533","864");
INSERT INTO "stadium" VALUES (4,"Queen's Park","Hampden Park","52500","1763","466","730");
INSERT INTO "stadium" VALUES (5,"Stirling Albion","Forthbank Stadium","3808","1125","404","642");
INSERT INTO "stadium" VALUES (6,"Arbroath","Gayfield Park","4125","921","411","638");
INSERT INTO "stadium" VALUES (7,"Alloa Athletic","Recreation Park","3100","1057","331","637");
INSERT INTO "stadium" VALUES (9,"Peterhead","Balmoor","4000","837","400","615");
INSERT INTO "stadium" VALUES (10,"Brechin City","Glebe Park","3960","780","315","552");
CREATE TABLE "singer" (
"Singer_ID" int,
"Name" text,
"Country" text,
"Song_Name" text,
"Song_release_year" text,
"Age" int,
"Is_male" bool,
PRIMARY KEY ("Singer_ID")
);
INSERT INTO "singer" VALUES (1,"Joe Sharp","Netherlands","You","1992",52,"F");
INSERT INTO "singer" VALUES (2,"Timbaland","United States","Dangerous","2008",32,"T");
INSERT INTO "singer" VALUES (3,"Justin Brown","France","Hey Oh","2013",29,"T");
INSERT INTO "singer" VALUES (4,"Rose White","France","Sun","2003",41,"F");
INSERT INTO "singer" VALUES (5,"John Nizinik","France","Gentleman","2014",43,"T");
INSERT INTO "singer" VALUES (6,"Tribal King","France","Love","2016",25,"T");
CREATE TABLE "concert" (
"concert_ID" int,
"concert_Name" text,
"Theme" text,
"Stadium_ID" text,
"Year" text,
PRIMARY KEY ("concert_ID"),
FOREIGN KEY ("Stadium_ID") REFERENCES "stadium"("Stadium_ID")
);
INSERT INTO "concert" VALUES (1,"Auditions","Free choice",1,2014);
INSERT INTO "concert" VALUES (2,"Super bootcamp","Free choice 2",2,2014);
INSERT INTO "concert" VALUES (3,"Home Visits","Bleeding Love",2,2015);
INSERT INTO "concert" VALUES (4,"Week 1","Wide Awake",10,2014);
INSERT INTO "concert" VALUES (5,"Week 1","Happy Tonight",9,2015);
INSERT INTO "concert" VALUES (6,"Week 2","Party All Night",7,2015);
CREATE TABLE "singer_in_concert" (
"concert_ID" int,
"Singer_ID" text,
PRIMARY KEY ("concert_ID","Singer_ID"),
FOREIGN KEY ("concert_ID") REFERENCES "concert"("concert_ID"),
FOREIGN KEY ("Singer_ID") REFERENCES "singer"("Singer_ID")
);
INSERT INTO "singer_in_concert" VALUES (1, 2);
INSERT INTO "singer_in_concert" VALUES (1, 3);
INSERT INTO "singer_in_concert" VALUES (1, 5);
INSERT INTO "singer_in_concert" VALUES (2, 3);
INSERT INTO "singer_in_concert" VALUES (2, 6);
INSERT INTO "singer_in_concert" VALUES (3, 5);
INSERT INTO "singer_in_concert" VALUES (4, 4);
INSERT INTO "singer_in_concert" VALUES (5, 6);
INSERT INTO "singer_in_concert" VALUES (5, 3);
INSERT INTO "singer_in_concert" VALUES (6, 2);

1
AlgoriAgent/examples/conversation_nl2sql/database/train.json Normal file
View File

File diff suppressed because one or more lines are too long

498
AlgoriAgent/examples/conversation_nl2sql/react_nl2sql.ipynb Normal file
View File

@@ -0,0 +1,498 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Convert Natural Language to SQL Query with ReAct Agent\n",
"\n",
"This notebook will show \n",
"- how to use ReAct agent to convert natural language to SQL query\n",
"- how to use a customized NL2SQL algorithm, DAIL-SQL, in AgentScope \n",
"\n",
"## Background\n",
"\n",
"NL2SQl is a challenging task in both database and natural language processing communities. It aims to convert natural language question into SQL query. In research community, there are many works exploring the potential of LLMs in NL2SQL, and it would be very interesting to explore this task with LLM-empowered agents. \n",
"\n",
"In this notebook, we try to equip the ReAct agent with a SOTA NL2SQL algorithm, DAIL-SQL, to see if the agent-level reasoning could bring any insights to this task. The implementation of DAIL-SQL refer to its [GitHub repository](https://github.com/BeachWang/DAIL-SQL) and [paper](https://arxiv.org/abs/2308.15363).\n",
"\n",
"## Prerequisites\n",
"\n",
"- Follow [READMD.md](https://github.com/modelscope/agentscope) to install AgentScope.\n",
"- Install the third-party libraries used in [DAIL-SQL](https://github.com/BeachWang/DAIL-SQL/blob/main/requirements.txt).\n",
"- Prepare a model configuration. AgentScope supports both local deployed model services (CPU or GPU) and third-party services. More details and example model configurations please refer to our [tutorial](https://modelscope.github.io/agentscope/en/tutorial/203-model.html).\n",
"- Get familiar with service functions, ServiceToolkit module and build-in ReAct agent. The corresponding example can be found in [ReAct agent example](../conversation_with_react_agent/main.ipynb).\n",
"\n",
"## Note\n",
"\n",
"- The example is tested with the following models. For other models, you may need to adjust the prompt.\n",
" - gpt-4 (The model that [DAIL-SQL](https://github.com/BeachWang/DAIL-SQL) is built on)\n",
" - gpt-3.5-turbo "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 1: Initializing model config and SQL database\n",
"First, we initalize the agentscope config and load the corresponding model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from typing import Callable\n",
"import agentscope\n",
"from agentscope.models import load_model_by_config_name\n",
"agentscope.init(\n",
" model_configs=\"./configs/model_configs.json\",\n",
" project=\"Conversation with NL2SQL\",\n",
")\n",
"loaded_model = load_model_by_config_name('gpt-4')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Then, we init and provide the corresponding database path in sqlite file format.\n",
"Here we generate the sqlite file using the provided sql commands [here](./database/concert_singer/schema.sql). \n",
"You can also use the .sqlite format file directly."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from sql_utils import create_sqlite_db_from_schema\n",
"example_db_id = \"concert_singer\"\n",
"db_schema_path = \"./database/concert_singer/schema.sql\"\n",
"db_sqlite_path = \"./database/concert_singer/concert_singer.sqlite\"\n",
"create_sqlite_db_from_schema(db_schema_path, db_sqlite_path)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can show a general description of the database schema using the code below."
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"The database consists of four tables: \n",
"\n",
"1. Stadium: This table includes the following fields: Stadium_ID (which is the primary key), Location, Name, Capacity, Highest, Lowest, and Average.\n",
"\n",
"2. Singer: This table includes the following fields: Singer_ID (which is the primary key), Name, Country, Song_Name, Song_release_year, Age, and Is_male (a Boolean field that indicates the singer's gender).\n",
"\n",
"3. Concert: This table is linked to the Stadium table via the Stadium_ID field. The table includes the fields concert_ID (which is the primary key), concert_Name, Theme, Stadium_ID, and Year.\n",
"\n",
"4. Singer_in_Concert: This table acts as a junction table to establish a many-to-many relationship between the Singer and Concert tables. It includes the fields concert_ID and Singer_ID, which both make up the primary key for the table. These are also reference as foreign keys to their respective basetable: the Concert table via the concert_ID field, and the Singer table via the Singer_ID field.\n",
"\n",
"In summary, this database seems designed to track concerts, including details about singers, the concert's location (stadium), and which singer performed at which concert.\n"
]
}
],
"source": [
"from loguru import logger\n",
"from sql_utils import SQLPrompt\n",
"describe_prompt = SQLPrompt().describe_schema(db_sqlite_path)\n",
"sql_description = loaded_model([{\"role\": \"assistant\", \"content\": describe_prompt}]).text\n",
"logger.chat(sql_description)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 2: Setting up tool functions\n",
"Then, we define the tools needed for React Agent to perform SQL query. Namely, our agent should be able to generate sql query given natural language and perform sql query to get the result.\n",
"\n",
"- We referenced a third-party Text-to-SQL [DAIL-SQL](https://github.com/BeachWang/DAIL-SQL) to generate Text-to-SQL prompt.\n",
"- We use the query_sqlite service function in the agentscope.service module. "
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"from agentscope.service import (\n",
" ServiceResponse, \n",
" ServiceExecStatus, \n",
" ServiceToolkit, \n",
" query_sqlite,\n",
")\n",
"from sql_utils import DailSQLPromptGenerator\n",
"\n",
"\n",
"def generate_sql_query(question: str, db_path: str, model: Callable) -> ServiceResponse:\n",
" \"\"\"\n",
" Generate the sql query based on the natural language question from user.\n",
" \n",
" Args:\n",
" question (`str`):\n",
" The question asked by user in natural language.\n",
" db_path (`str`):\n",
" The path to the .sqlite file.\n",
" model (`Callable`):\n",
" The LLM model used to generate the SQL query.\n",
" \"\"\"\n",
" prompt_helper = DailSQLPromptGenerator(db_path)\n",
" prepared_prompt = prompt_helper.generate_prompt({\"content\": question})\n",
" \n",
" def get_response_from_prompt(prompt: dict, model: Callable) -> str:\n",
" \"\"\"\n",
" Generate response from prompt using LLM\n",
" \"\"\"\n",
" messages = [{\"role\": \"assistant\", \"content\": prompt}]\n",
" sql = model(messages).text\n",
" sql = \" \".join(sql.replace(\"\\n\", \" \").split())\n",
" sql = sql.strip().split(\"/*\")[0]\n",
" if sql.startswith(\"SELECT\"):\n",
" response = sql + \"\\n\"\n",
" elif sql.startswith(\" \"):\n",
" response = \"SELECT\" + sql + \"\\n\"\n",
" else:\n",
" response = \"SELECT \" + sql + \"\\n\"\n",
" return response\n",
"\n",
" sql_response = get_response_from_prompt(\n",
" prepared_prompt[\"prompt\"], model=model\n",
" )\n",
"\n",
" return ServiceResponse(\n",
" ServiceExecStatus.SUCCESS,\n",
" sql_response\n",
" )\n",
"\n",
"\n",
"# Use Service Toolkit to set up tool functions for LLMs\n",
"service_toolkit = ServiceToolkit()\n",
"service_toolkit.add(generate_sql_query, db_path=db_sqlite_path, model=loaded_model)\n",
"service_toolkit.add(query_sqlite, database=db_sqlite_path)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 3: Create the ReAct agent to help perform SQL queries\n",
"\n",
"Now, we can init the ReAct Agent using the defined tools and interact with the agent."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from agentscope.agents import ReActAgent\n",
"agent = ReActAgent(\n",
" name=\"assistant\",\n",
" model_config_name='gpt-4',\n",
" service_toolkit=service_toolkit,\n",
" sys_prompt=\"You are a helpful agent that preform SQL queries base on natual language instructions.\",\n",
" verbose=True, # set verbose to True to show the reasoning process\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 4: Interact with the SQL ReAct Agent using natural lanugage\n",
"\n",
"Now let's see how the agent have response to our questions:"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"user: How many singers do we have?\n",
"##################### ITER 1, STEP 1: REASONING ######################\n",
"assistant: {\n",
" \"thought\": \"To answer the user's question, I need to find out how many unique singers are listed in the database. I will use the generate_sql_query function to create a SQL query to count distinct singers.\",\n",
" \"speak\": \"Let me find this information for you.\",\n",
" \"function\": [\n",
" {\n",
" \"name\": \"generate_sql_query\",\n",
" \"arguments\": {\n",
" \"question\": \"How many unique singers are we have in the database?\"\n",
" }\n",
" }\n",
" ]\n",
"}\n",
"####################### ITER 1, STEP 2: ACTION #######################\n",
">>> Executing function generate_sql_query ...\n",
">>> END \n",
"system: Execution Results:\n",
"1. generate_sql_query:\n",
" [EXECUTE STATUS]: SUCCESS\n",
" [EXECUTE RESULT]: SELECT To find out how many unique singers are in the database, you would need to count the distinct singer IDs in the \"singer\" table. The appropriate SQL query for this would be: ```sql SELECT COUNT(DISTINCT Singer_ID) FROM singer; ```\n",
"\n",
"\n",
"##################### ITER 2, STEP 1: REASONING ######################\n",
"assistant: {\n",
" \"thought\": \"The generate_sql_query function has provided the necessary SQL query. Now I will use the query_sqlite function to execute this query and retrieve the count of unique singers.\",\n",
" \"speak\": \"I'm now running the query to count the number of unique singers in the database.\",\n",
" \"function\": [\n",
" {\n",
" \"name\": \"query_sqlite\",\n",
" \"arguments\": {\n",
" \"allow_change_data\": false,\n",
" \"query\": \"SELECT COUNT(DISTINCT Singer_ID) FROM singer;\"\n",
" }\n",
" }\n",
" ]\n",
"}\n",
"####################### ITER 2, STEP 2: ACTION #######################\n",
">>> Executing function query_sqlite ...\n",
">>> END \n",
"system: Execution Results:\n",
"1. query_sqlite:\n",
" [EXECUTE STATUS]: SUCCESS\n",
" [EXECUTE RESULT]: [(6,)]\n",
"\n",
"##################### ITER 3, STEP 1: REASONING ######################\n",
"assistant: {\n",
" \"thought\": \"I successfully executed the query to count the distinct singers in the database and received the result.\",\n",
" \"speak\": \"We have 6 unique singers in the database.\",\n",
" \"function\": []\n",
"}\n"
]
}
],
"source": [
"from agentscope.message import Msg\n",
"\n",
"mss = Msg(\n",
" name=\"user\", \n",
" content=\"How many singers do we have?\", \n",
" role=\"user\"\n",
")\n",
"logger.chat(mss)\n",
"\n",
"sql_query_mss1 = agent(mss)"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"user: Show all countries and the number of singers in each country.\n",
"##################### ITER 1, STEP 1: REASONING ######################\n",
"assistant: {\n",
" \"thought\": \"To show all countries and the number of singers in each, I'll need to create and run an SQL query that counts singers by their country. This involves grouping the data by country and then counting the number of singers in each group.\",\n",
" \"speak\": \"Let me find out how many singers there are in each country for you.\",\n",
" \"function\": [\n",
" {\n",
" \"name\": \"generate_sql_query\",\n",
" \"arguments\": {\n",
" \"question\": \"Show all countries and the number of singers in each country.\"\n",
" }\n",
" }\n",
" ]\n",
"}\n",
"####################### ITER 1, STEP 2: ACTION #######################\n",
">>> Executing function generate_sql_query ...\n",
">>> END \n",
"system: Execution Results:\n",
"1. generate_sql_query:\n",
" [EXECUTE STATUS]: SUCCESS\n",
" [EXECUTE RESULT]: SELECT To show all countries and the number of singers in each country based on the database schema provided, you can use the SQL query below. This query groups singers by their country and counts the number of singers in each country: ```sql SELECT Country, COUNT(*) AS NumberOfSingers FROM singer GROUP BY Country ``` This SQL query selects the `Country` column from the `singer` table and uses the `COUNT(*)` function to count the total number of singers for each country. The `GROUP BY` clause groups the results by the `Country`, so you get the count of singers for each distinct country listed in the `singer` table.\n",
"\n",
"\n",
"##################### ITER 2, STEP 1: REASONING ######################\n",
"assistant: {\n",
" \"thought\": \"The generate_sql_query function has provided the correct SQL query for showing all countries and the number of singers in each. Now, I need to execute this query using the query_sqlite function to get the desired data.\",\n",
" \"speak\": \"I'm now running the query to find out the number of singers in each country.\",\n",
" \"function\": [\n",
" {\n",
" \"name\": \"query_sqlite\",\n",
" \"arguments\": {\n",
" \"allow_change_data\": false,\n",
" \"query\": \"SELECT Country, COUNT(*) AS NumberOfSingers FROM singer GROUP BY Country\"\n",
" }\n",
" }\n",
" ]\n",
"}\n",
"####################### ITER 2, STEP 2: ACTION #######################\n",
">>> Executing function query_sqlite ...\n",
">>> END \n",
"system: Execution Results:\n",
"1. query_sqlite:\n",
" [EXECUTE STATUS]: SUCCESS\n",
" [EXECUTE RESULT]: [('France', 4), ('Netherlands', 1), ('United States', 1)]\n",
"\n",
"##################### ITER 3, STEP 1: REASONING ######################\n",
"assistant: {\n",
" \"thought\": \"The query has been executed successfully and I got the result indicating the number of singers in each country. Now I will present this data to the user.\",\n",
" \"speak\": \"We have 4 singers from France, 1 singer from Netherlands, and 1 singer from the United States.\",\n",
" \"function\": []\n",
"}\n"
]
}
],
"source": [
"mss = Msg(\n",
" name=\"user\", \n",
" content=\"Show all countries and the number of singers in each country.\", \n",
" role=\"user\"\n",
")\n",
"logger.chat(mss)\n",
"\n",
"sql_query_mss2 = agent(mss)"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"user: How many singers have held concert in multiple stadiums? Give me the name of these singers.\n",
"##################### ITER 1, STEP 1: REASONING ######################\n",
"assistant: {\n",
" \"thought\": \"To find out how many singers have held concerts in multiple stadiums and get their names, I need to create a SQL query that joins relevant tables to count the number of distinct stadiums for each singer. Singers with counts greater than one have held concerts in multiple stadiums. I will start by generating the SQL query.\",\n",
" \"speak\": \"Let me find out which singers have held concerts in multiple stadiums for you.\",\n",
" \"function\": [\n",
" {\n",
" \"name\": \"generate_sql_query\",\n",
" \"arguments\": {\n",
" \"question\": \"How many singers have held concerts in multiple stadiums and what are their names?\"\n",
" }\n",
" }\n",
" ]\n",
"}\n",
"####################### ITER 1, STEP 2: ACTION #######################\n",
">>> Executing function generate_sql_query ...\n",
">>> END \n",
"system: Execution Results:\n",
"1. generate_sql_query:\n",
" [EXECUTE STATUS]: SUCCESS\n",
" [EXECUTE RESULT]: SELECT s.Name, COUNT(DISTINCT c.Stadium_ID) FROM singer AS s JOIN singer_in_concert AS sic ON s.Singer_ID = sic.Singer_ID JOIN concert AS c ON sic.concert_ID = c.concert_ID GROUP BY s.Name HAVING COUNT(DISTINCT c.Stadium_ID) > 1\n",
"\n",
"\n",
"##################### ITER 2, STEP 1: REASONING ######################\n",
"assistant: {\n",
" \"thought\": \"The generate_sql_query function successfully generated the necessary SQL query to find singers who have held concerts in multiple stadiums along with their names. I will now use the query_sqlite function to execute this SQL query and retrieve the required information.\",\n",
" \"speak\": \"I'm now running the query to find out which singers have performed in multiple stadiums.\",\n",
" \"function\": [\n",
" {\n",
" \"name\": \"query_sqlite\",\n",
" \"arguments\": {\n",
" \"allow_change_data\": false,\n",
" \"query\": \"SELECT s.Name, COUNT(DISTINCT c.Stadium_ID) FROM singer AS s JOIN singer_in_concert AS sic ON s.Singer_ID = sic.Singer_ID JOIN concert AS c ON sic.concert_ID = c.concert_ID GROUP BY s.Name HAVING COUNT(DISTINCT c.Stadium_ID) > 1\"\n",
" }\n",
" }\n",
" ]\n",
"}\n",
"####################### ITER 2, STEP 2: ACTION #######################\n",
">>> Executing function query_sqlite ...\n",
">>> END \n",
"system: Execution Results:\n",
"1. query_sqlite:\n",
" [EXECUTE STATUS]: SUCCESS\n",
" [EXECUTE RESULT]: [('John Nizinik', 2), ('Justin Brown', 3), ('Timbaland', 2), ('Tribal King', 2)]\n",
"\n",
"##################### ITER 3, STEP 1: REASONING ######################\n",
"assistant: {\n",
" \"thought\": \"The SQL query has been executed successfully and the resulting dataset includes all singers who have held concerts in multiple stadiums.\",\n",
" \"speak\": \"The singers who have held concerts in multiple stadiums are John Nizinik, Justin Brown, Timbaland, and Tribal King.\",\n",
" \"function\": []\n",
"}\n"
]
}
],
"source": [
"mss = Msg(\n",
" name=\"user\", \n",
" content=\"How many singers have held concert in multiple stadiums? Give me the name of these singers.\", \n",
" role=\"user\"\n",
")\n",
"logger.chat(mss)\n",
"sql_query_mss3 = agent(mss)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"You can use the following code to build a conversation with the ReAct agent:\n",
"\n",
"``` python\n",
"from agentscope.agents import UserAgent\n",
"\n",
"user = UserAgent(name=\"User\")\n",
"\n",
"x = None\n",
"while True:\n",
" x = user(x)\n",
" if x.content == \"exit\":\n",
" break\n",
" x = agent(x)\n",
"```"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The capabilities of the ReAct agent are indeed very powerful. By modifying prompts and expanding the corresponding usage tools, the ReAct agent can leverage the combination of LLM and tools to accomplish more complex tasks.\n",
"\n",
"We encourage users to delve deeper into exploring the Agent's functionalities and to try out the capabilities of the LLM agent on their own."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.18"
}
},
"nbformat": 4,
"nbformat_minor": 4
}

304
AlgoriAgent/examples/conversation_nl2sql/sql_utils.py Normal file
View File

@@ -0,0 +1,304 @@
# -*- coding: utf-8 -*-
"""
Utils and helpers for performing sql querys.
Referenced from https://github.com/BeachWang/DAIL-SQL.
"""
import sqlite3
import json
from typing import Union
from sqlite3 import Connection
import os
import numpy as np
def query_sqlite(
queries: Union[list[str], str],
path_db: str = None,
cur: str = None,
) -> Union[list, str]:
"""Execute queries and return results."""
assert not (
path_db is None and cur is None
), "path_db and cur cannot be NoneType at the same time"
con: Connection
close_in_func = False
if cur is None:
con = sqlite3.connect(path_db)
cur = con.cursor()
close_in_func = True
if isinstance(queries, str):
results = cur.execute(queries).fetchall()
elif isinstance(queries, list):
results = []
for query in queries:
res = cur.execute(query).fetchall()
results.append(res)
else:
raise TypeError(f"queries cannot be {type(queries)}")
# close the connection if needed
if close_in_func:
con.close()
return results
def create_sqlite_db_from_schema(schema_path: str, db_path: str) -> None:
"""Create a SQLite database file from a schema SQL file.
Args:
schema_path: The file path to the schema SQL file.
db_path: The file path for the SQLite database to be created.
"""
if os.path.exists(db_path):
print(
f"Database file '{db_path}' already exists. ",
)
return
conn = sqlite3.connect(db_path)
with open(schema_path, "r", encoding="utf-8") as schema_file:
schema_sql = schema_file.read()
cursor = conn.cursor()
cursor.executescript(schema_sql)
conn.commit()
conn.close()
def get_table_names(path_db: str = None, cur: str = None) -> list:
"""Get names of all tables within the database,
and reuse cur if it's not None"""
table_names = query_sqlite(
queries="SELECT name FROM sqlite_master WHERE type='table'",
path_db=path_db,
cur=cur,
)
table_names = [_[0] for _ in table_names]
return table_names
def get_sql_for_database(path_db: str = None, cur: str = None) -> list:
"""
Get sql table from database
"""
close_in_func = False
if cur is None:
con = sqlite3.connect(path_db)
cur = con.cursor()
close_in_func = True
table_names = get_table_names(path_db, cur)
queries = [
f"SELECT sql FROM sqlite_master WHERE tbl_name='{name}'"
for name in table_names
]
sqls = query_sqlite(queries, path_db, cur)
if close_in_func:
cur.close()
return [_[0][0] for _ in sqls]
class SQLPrompt:
"""SQL prompt helper"""
def __init__(self) -> None:
self.template_info = "/* Given the following database schema: */\n{}"
self.template_question = "/* Answer the following: {} */"
self.template_agent_prompt = (
"You are a helpful agent that preform"
"SQL querys base on natual language instructions."
"Please describe the database schema provided"
"in a simple and understandable manner. "
)
self.is_sql_prompt = (
"Please read the user's question below and "
"determine whether the question is an appropriate "
"query for the given SQL schema. \n"
"If the question is indeed a query pertaining to the SQL schema, "
'respond with "YES". '
"If the question is not a query related to the SQL schema, "
"provide a brief explanation to the user explaining why their "
"question does not correspond to a SQL query within the "
"context of the schema. "
"The user's question is: "
)
def format_target(self, example: dict) -> str:
"""Format sql prompt"""
return self.format_question(example) + "\nSELECT "
def format_question(self, example: dict) -> str:
"""Format question"""
sqls = get_sql_for_database(example["path_db"])
prompt_info = self.template_info.format("\n\n".join(sqls))
prompt_question = self.template_question.format(example["question"])
prompt_components = [prompt_info, prompt_question]
prompt = "\n\n".join(prompt_components)
return prompt
def describe_schema(self, db_path: str) -> str:
"""Describe SQL schema"""
sqls = get_sql_for_database(db_path)
prompt_info = self.template_info.format("\n\n".join(sqls))
prompt_components = [
self.template_agent_prompt,
"DB schema info: ",
prompt_info,
]
prompt = "\n\n".join(prompt_components)
return prompt
def is_sql_question(self, example: dict) -> str:
"""whether the input is a sql question or not"""
sqls = get_sql_for_database(example["path_db"])
prompt_info = self.template_info.format("\n\n".join(sqls))
prompt_components = [
prompt_info,
self.is_sql_prompt,
example["question"],
]
prompt = "\n".join(prompt_components)
return prompt
class QuestionSqlExampleStyle:
"""Provide QA pair as examples"""
def get_example_prefix(self) -> str:
"""get example prefix"""
return (
"/* Some SQL examples are provided based on similar problems: */\n"
)
def format_example(self, example: dict) -> str:
"""format example"""
template_qa = "/* Answer the following: {} */\n{}"
return template_qa.format(example["question"], example["query"])
class EuclideanDistanceExampleSelector:
"""Select similar sample question"""
def __init__(self) -> None:
self.train_json_path = "./database/train.json"
with open(self.train_json_path, "r", encoding="utf-8") as file:
self.train_json = json.load(file)
self.train_questions = [_["question"] for _ in self.train_json]
self.SELECT_MODEL = "sentence-transformers/all-mpnet-base-v2"
from sentence_transformers import SentenceTransformer
self.bert_model = SentenceTransformer(self.SELECT_MODEL, device="cpu")
try:
self.train_embeddings = np.load("./.cache/train_embeddings.npy")
except FileNotFoundError:
self.train_embeddings = self.bert_model.encode(
self.train_questions,
)
np.save("./.cache/train_embeddings.npy", self.train_embeddings)
def get_examples(
self,
target: dict,
num_example: int,
) -> list:
"""Get similar question examples for few shot"""
target_embedding = self.bert_model.encode([target["question"]])
# find the most similar question in train dataset
from sklearn.metrics.pairwise import euclidean_distances
distances = np.squeeze(
euclidean_distances(target_embedding, self.train_embeddings),
).tolist()
pairs = list(zip(distances, range(len(distances))))
train_json = self.train_json
pairs_sorted = sorted(pairs, key=lambda x: x[0])
top_pairs = []
for d, index in pairs_sorted:
top_pairs.append((index, d))
if len(top_pairs) >= num_example:
break
return [train_json[index] for (index, d) in top_pairs]
class DailSQLPromptGenerator:
"""Generate prompt given the dataset and question"""
def __init__(
self,
db_path: str,
) -> None:
self.db_path = db_path
self.sql_prompt = SQLPrompt()
self.question_selector = EuclideanDistanceExampleSelector()
self.question_style = QuestionSqlExampleStyle()
self.SEP_EXAMPLE = "\n\n"
self.scope_factor = 100
self.NUM_EXAMPLE = 9
self.cross_domain = False
def describe_schema(self) -> str:
"""Describe the sql"""
return self.sql_prompt.describe_schema(self.db_path)
def is_sql_question_prompt(self, question: str) -> str:
"""
prompt for LLM to judge whether the question is appropriate
"""
target = {
"path_db": self.db_path,
"question": question,
}
return self.sql_prompt.is_sql_question(target)
def generate_prompt(self, x: dict = None) -> dict:
"""
Generate prompt given input question
"""
question = x["content"]
target = {
"path_db": self.db_path,
"question": question,
}
prompt_target = self.sql_prompt.format_target(target)
if self.NUM_EXAMPLE != 0:
examples = self.question_selector.get_examples(
target,
self.NUM_EXAMPLE * self.scope_factor,
)
prompt_example = []
question = target["question"]
example_prefix = self.question_style.get_example_prefix()
for example in examples:
example_format = self.question_style.format_example(example)
prompt_example.append(example_format)
if len(prompt_example) >= self.NUM_EXAMPLE:
break
n_valid_example = len(prompt_example)
if len(prompt_example) > 0:
prompt = example_prefix + self.SEP_EXAMPLE.join(
prompt_example + [prompt_target],
)
else:
prompt = self.SEP_EXAMPLE.join(
prompt_example + [prompt_target],
)
else:
n_valid_example = 0
prompt = prompt_target
return {
"prompt": prompt,
"n_examples": n_valid_example,
}