wehub-resource-sync
119 lines
4.4 KiB
Python
119 lines
4.4 KiB
Python
|
|
""" This example shows an end-to-end recipe for creating a CustomTable, and then creating an Agent process that
|
|
will query the table using natural language.
|
|
|
|
Please note that this example is a 'generalized' and updated version of an earlier example -
|
|
"text2sql-end-to-end-2.py" - now using the more powerful CustomTables class integrated into the LLMfx process
|
|
|
|
The example shows the following steps:
|
|
|
|
1. Creating a custom table resource from a sample CSV file, included in the slim-sql-tool kit, and also
|
|
available in the Examples section with Structured_Tables (customer_table.csv)
|
|
|
|
2 Asking basic natural language questions:
|
|
A. Looks up the table schema
|
|
B. Packages the table schema with query
|
|
C. Runs inference to convert text into SQL
|
|
D. Queries the database with the generated SQL
|
|
E. Returns result
|
|
|
|
3. Using CustomtTable class, this can be run on either Postgres or SQLite DB.
|
|
|
|
Note: as you substitute for your own CSV and JSON, check out the other examples in this section for loading
|
|
configuration ideas and options.
|
|
|
|
"""
|
|
|
|
import os
|
|
|
|
from llmware.agents import LLMfx
|
|
from llmware.resources import CustomTable
|
|
from llmware.configs import LLMWareConfig
|
|
|
|
|
|
def build_table(db=None, table_name=None,load_fp=None,load_file=None):
|
|
|
|
""" Simple example script to take a CSV or JSON/JSONL and create a DB Table. """
|
|
|
|
custom_table = CustomTable(db=db, table_name=table_name)
|
|
analysis = custom_table.validate_csv(load_fp, load_file)
|
|
print("update: analysis from validate_csv: ", analysis)
|
|
|
|
if load_file.endswith(".csv"):
|
|
output = custom_table.load_csv(load_fp, load_file)
|
|
elif load_file.endswith(".jsonl") or load_file.endswith(".json"):
|
|
output = custom_table.load_json(load_fp, load_file)
|
|
else:
|
|
print("file type not supported for db load")
|
|
return -1
|
|
|
|
print("update: output from loading file: ", output)
|
|
|
|
sample_range = min(10, len(custom_table.rows))
|
|
for x in range(0,sample_range):
|
|
print("update: sample rows: ", x, custom_table.rows[x])
|
|
|
|
# stress the schema data type and remediate - use more samples for more accuracy
|
|
updated_schema = custom_table.test_and_remediate_schema(samples=20, auto_remediate=True)
|
|
|
|
print("update: updated schema: ", updated_schema)
|
|
|
|
# insert the rows in the DB
|
|
custom_table.insert_rows()
|
|
|
|
return 1
|
|
|
|
|
|
def agent_natural_language_sql_query(query_list, db=None, table_name=None):
|
|
|
|
""" Query a CustomTable in natural language. """
|
|
|
|
# query starts here
|
|
agent = LLMfx()
|
|
agent.load_tool("sql", sample=False, get_logits=True, temperature=0.0)
|
|
|
|
# Pass direct queries to the DB
|
|
|
|
for i, query in enumerate(query_list):
|
|
|
|
# query_custom_table method is doing all of the work
|
|
# -- looks up the table schema in the db using the table_name
|
|
# -- packages the text-2-sql query prompt
|
|
# -- executes sql method to convert the prompt into a sql query
|
|
# -- attempts to execute the sql query on the db
|
|
# -- returns the db results as 'research' output
|
|
|
|
response = agent.query_custom_table(query,db=db,table=table_name)
|
|
|
|
for x in range(0,len(agent.research_list)):
|
|
print("research: ", x, agent.research_list[x])
|
|
|
|
return agent.research_list
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
# input parameters - db, table_name, path to csv or json/jsonl file
|
|
db = "postgres"
|
|
table_name = "customer_table"
|
|
|
|
# for hello_world, please pull down the customer_table.csv found in the examples repository
|
|
input_fp = "/local/path/to/csv/"
|
|
input_fn = "customer_table.csv"
|
|
|
|
# builds table - only needs to be done once
|
|
build_table(db=db, table_name=table_name, load_fp=input_fp, load_file=input_fn)
|
|
|
|
# list of queries to ask the csv or json/jsonl
|
|
# -- note: these queries are useful but purposefully straightforward with relatively clear alignment
|
|
# to the data schema
|
|
|
|
query_list = ["Which customers have vip customer status of yes?",
|
|
"What is the highest annual spend of any customer?",
|
|
"Which customer has account number 1234953",
|
|
"Which customer has the lowest annual spend?",
|
|
"Is Susan Soinsin a vip customer?"]
|
|
|
|
# agent process to execute the query list
|
|
agent_natural_language_sql_query(query_list, db=db, table_name=table_name)
|