Drake's
0%

使用大型語言模型(LLM)完成Text-to-SQL任務 - 使用Function Calling

發表於 更新於 分類於 閱讀次數： Disqus：

上一篇說明Langchain的SQL Agent範例如何透過ReAct Prompting實作，而目前LLM其實都已經具有Function Calling的能力，比起要解析LLM回傳的Action和Action Input再另外執行，直接透過Function Calling的效果其實更好。甚至目前的LLM能力已經可以在不強迫產生Thought的推論動作下模型就能順利解決問題。

Function Calling可以允許LLM使用外部的API或Function，以克服無法取得最新知識或是外部資料的限制。只要將prompt和可以允許執行的Function提供給LLM，模型就能在處理prompt時選擇能夠處理當下任務的Function並得到回傳資訊。在取得Function所提供的外部回傳資訊後，可以再提供給模型作最後的回答。

下面使用Langchain的SQLDatabase提供執行的Function，與Gemini的Function Calling完成Text-to-SQL任務。

首先建立Chinook資料庫，並使用langchain的SQLDatabase建立資料庫的查詢引擎：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import sqlite3
import requests

# 下載 SQL 腳本
url = "https://raw.githubusercontent.com/lerocha/chinook-database/master/ChinookDatabase/DataSources/Chinook_Sqlite.sql"
response = requests.get(url)

# 儲存 SQL 腳本到本地檔案
with open('Chinook_Sqlite.sql', 'w') as f:
    f.write(response.text)

# 連接到 SQLite 資料庫 (會自動建立資料庫檔案)
conn = sqlite3.connect('Chinook.db')
cursor = conn.cursor()

# 執行 SQL 腳本來建立資料庫
with open('Chinook_Sqlite.sql', 'r') as f:
    cursor.executescript(f.read())

# 提交變更並關閉連接
conn.commit()
conn.close()
1
2
3
from langchain_community.utilities import SQLDatabase
# 準備sqlite db
db = SQLDatabase.from_uri("sqlite:///Chinook.db")

Gemini會使用function name, docstring, parameters和parameter type annotations來決定是否要使用特定的函式來回答問題。這裡借用Langchain的SQLDatabase的三個Function並寫好docstring和parameter。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
def sql_db_list_tables():
    """Tool for getting tables names.

    Args:
        Input is an empty string
    """
    return ", ".join(db.get_usable_table_names());

def sql_db_schema(table_names: str):
    """Tool for getting metadata about a SQL database. Get the schema and sample rows for the specified SQL tables.

    Args:
        table_names: A comma-separated list of the table names for which to return the schema.
    """
    return db.get_table_info_no_throw(
            [t.strip() for t in table_names.split(",")]
        )

def sql_db_query(query: str):
    """Tool for querying a SQL database. Execute a SQL query against the database and get back the result..

    Args:
        query: A detailed and correct SQL query.
    """
    return db.run_no_throw(query)

接著設定function calling config，這裡設定成ANY會強迫讓模型使用提供的Function，並提供以上三個Function。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
from google.generativeai.types import content_types
from collections.abc import Iterable

def tool_config_from_mode(mode: str, fns: Iterable[str] = ()):
    """Create a tool config with the specified function calling mode."""
    return content_types.to_tool_config(
        {"function_calling_config": {"mode": mode, "allowed_function_names": fns}}
    )

sql_tool_dict = {
    "sql_db_query": sql_db_query,
    "sql_db_schema": sql_db_schema,
    "sql_db_list_tables": sql_db_list_tables
}

tool_config = tool_config_from_mode(mode="any", fns=list(sql_tool_dict.keys()))

以下參考Langchain的prompt內容並擷取prefix當作instruction還有輸入的問題為input。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
instruction = """
You are an agent designed to interact with a SQL database.
Given an input question, create a syntactically correct sqlite query to run, then look at the results of the query and return the answer.
Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 10 results.
You can order the results by a relevant column to return the most interesting examples in the database.
Never query for all the columns from a specific table, only ask for the relevant columns given the question.
You have access to tools for interacting with the database.
Only use the below tools. Only use the information returned by the below tools to construct your final answer.
You MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.

DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.

If the question does not seem related to the database, just return "I don't know" as the answer.

First, You should look at the tables in the database to see what I can query.  Then I should query the schema of the most relevant tables.
"""

text2sql_input = """
How many employees are there?
"""

如果使用generate_content會需要手動管理與模型的互動，所以這裡會把將instruction、prompt和tools提供給模型並與模型建立chat session，其中將enable_automatic_function_calling設成True時會允許模型自己呼叫我們所提供的Function。

1
2
3
4
5
6
7
8
9
10
# Configure Gemini
genai.configure(api_key=GEMINI_API_KEY)
model = genai.GenerativeModel(model_name='gemini-1.5-flash', \
                system_instruction=instruction, \
                tools=list(sql_tool_dict.values()))

chat = model.start_chat(enable_automatic_function_calling=True)

response = chat.send_message(text2sql_input)
response.text

最後可以從模型得到正確回答「There are 8 employees.」。

1
There are 8 employees.

可以從以下的chat history看到模型會從input的內容開始，先呼叫查詢所有的資料表後，接著選出Employee資料表查詢schema資訊，最後下一段模型生成出來的SQL查詢後回答最終的答案。

1
2
3
for content in chat.history:
    print(content.role, "->", [type(part).to_dict(part) for part in content.parts])
    print("-" * 80)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
user -> [{'text': '\nHow many employees are there?\n'}]
--------------------------------------------------------------------------------
model -> [{'function_call': {'name': 'sql_db_list_tables', 'args': {}, 'id': ''}}]
--------------------------------------------------------------------------------
user -> [{'function_response': {'name': 'sql_db_list_tables', 'response': {'result': 'Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track'}, 'id': ''}}]
--------------------------------------------------------------------------------
model -> [{'function_call': {'name': 'sql_db_schema', 'args': {'table_names': 'Employee'}, 'id': ''}}]
--------------------------------------------------------------------------------
user -> [{'function_response': {'name': 'sql_db_schema', 'response': {'result': '\nCREATE TABLE "Employee" (\n\t"EmployeeId" INTEGER NOT NULL, \n\t"LastName" NVARCHAR(20) NOT NULL, \n\t"FirstName" NVARCHAR(20) NOT NULL, \n\t"Title" NVARCHAR(30), \n\t"ReportsTo" INTEGER, \n\t"BirthDate" DATETIME, \n\t"HireDate" DATETIME, \n\t"Address" NVARCHAR(70), \n\t"City" NVARCHAR(40), \n\t"State" NVARCHAR(40), \n\t"Country" NVARCHAR(40), \n\t"PostalCode" NVARCHAR(10), \n\t"Phone" NVARCHAR(24), \n\t"Fax" NVARCHAR(24), \n\t"Email" NVARCHAR(60), \n\tPRIMARY KEY ("EmployeeId"), \n\tFOREIGN KEY("ReportsTo") REFERENCES "Employee" ("EmployeeId")\n)\n\n/*\n3 rows from Employee table:\nEmployeeId\tLastName\tFirstName\tTitle\tReportsTo\tBirthDate\tHireDate\tAddress\tCity\tState\tCountry\tPostalCode\tPhone\tFax\tEmail\n1\tAdams\tAndrew\tGeneral Manager\tNone\t1962-02-18 00:00:00\t2002-08-14 00:00:00\t11120 Jasper Ave NW\tEdmonton\tAB\tCanada\tT5K 2N1\t+1 (780) 428-9482\t+1 (780) 428-3457\tandrew@chinookcorp.com\n2\tEdwards\tNancy\tSales Manager\t1\t1958-12-08 00:00:00\t2002-05-01 00:00:00\t825 8 Ave SW\tCalgary\tAB\tCanada\tT2P 2T3\t+1 (403) 262-3443\t+1 (403) 262-3322\tnancy@chinookcorp.com\n3\tPeacock\tJane\tSales Support Agent\t2\t1973-08-29 00:00:00\t2002-04-01 00:00:00\t1111 6 Ave SW\tCalgary\tAB\tCanada\tT2P 5M5\t+1 (403) 262-3443\t+1 (403) 262-6712\tjane@chinookcorp.com\n*/'}, 'id': ''}}]
--------------------------------------------------------------------------------
model -> [{'function_call': {'name': 'sql_db_query', 'args': {'query': 'SELECT COUNT(*) FROM Employee'}, 'id': ''}}]
--------------------------------------------------------------------------------
user -> [{'function_response': {'name': 'sql_db_query', 'response': {'result': '[(8,)]'}, 'id': ''}}]
--------------------------------------------------------------------------------
model -> [{'text': 'There are 8 employees.'}]
--------------------------------------------------------------------------------

參考資料:
Generative-AI function-calling
Gemini cookbook Function_Calling
Gemini cookbook Function_Calling_config