ZennyKenny commited on
Commit
7306c07
ยท
verified ยท
1 Parent(s): 4943ed5

add plain english support

Browse files
Files changed (1) hide show
  1. app.py +33 -10
app.py CHANGED
@@ -5,24 +5,32 @@ from sqlalchemy import text
5
  from smolagents import tool, CodeAgent, HfApiModel
6
  import spaces
7
 
 
8
  from database import engine, receipts
9
 
 
10
  load_dotenv(override=True)
11
  hf_token = os.getenv("HF_TOKEN")
12
 
 
 
 
13
  @tool
14
  def sql_engine(query: str) -> str:
15
  """
16
- Allows you to perform SQL queries on the table. Returns a string representation of the result.
17
- The table is named 'receipts'. Its description is as follows:
18
- Columns:
19
  - receipt_id: INTEGER
20
  - customer_name: VARCHAR(16)
21
  - price: FLOAT
22
  - tip: FLOAT
23
 
24
  Args:
25
- query: The query to perform. This should be correct SQL.
 
 
 
26
  """
27
  output = ""
28
  try:
@@ -34,21 +42,36 @@ def sql_engine(query: str) -> str:
34
  output = f"Error: {str(e)}"
35
  return output.strip()
36
 
 
37
  agent = CodeAgent(
38
- tools=[sql_engine],
39
  model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct", token=hf_token),
40
  )
41
 
42
- def query_sql(user_query):
43
- return sql_engine(user_query)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
44
 
45
  # Define Gradio interface
46
  demo = gr.Interface(
47
  fn=query_sql,
48
- inputs=gr.Textbox(label="Enter your SQL Query"),
49
  outputs=gr.Textbox(label="Query Result"),
50
- title="SQL Query Executor",
51
- description="Enter SQL queries to interact with an in-memory SQLite database.",
52
  flagging_mode="never",
53
  )
54
 
 
5
  from smolagents import tool, CodeAgent, HfApiModel
6
  import spaces
7
 
8
+ # Import database
9
  from database import engine, receipts
10
 
11
+ # Load Hugging Face API token
12
  load_dotenv(override=True)
13
  hf_token = os.getenv("HF_TOKEN")
14
 
15
+ if not hf_token:
16
+ raise ValueError("HF_TOKEN environment variable is not set.")
17
+
18
  @tool
19
  def sql_engine(query: str) -> str:
20
  """
21
+ Executes an SQL query on the 'receipts' table and returns results.
22
+
23
+ Table Schema:
24
  - receipt_id: INTEGER
25
  - customer_name: VARCHAR(16)
26
  - price: FLOAT
27
  - tip: FLOAT
28
 
29
  Args:
30
+ query: The SQL query to execute.
31
+
32
+ Returns:
33
+ Query result as a string.
34
  """
35
  output = ""
36
  try:
 
42
  output = f"Error: {str(e)}"
43
  return output.strip()
44
 
45
+ # Initialize CodeAgent to generate SQL queries from natural language
46
  agent = CodeAgent(
47
+ tools=[sql_engine], # Ensure sql_engine is properly registered
48
  model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct", token=hf_token),
49
  )
50
 
51
+ def query_sql(user_query: str) -> str:
52
+ """
53
+ Converts natural language input to an SQL query using CodeAgent
54
+ and returns the execution results.
55
+
56
+ Args:
57
+ user_query: The user's request in natural language.
58
+
59
+ Returns:
60
+ The query result from the database.
61
+ """
62
+ # Generate SQL from natural language
63
+ generated_sql = agent.run(f"Convert this request into SQL: {user_query}")
64
+
65
+ # Execute the SQL query and return the result
66
+ return sql_engine(generated_sql)
67
 
68
  # Define Gradio interface
69
  demo = gr.Interface(
70
  fn=query_sql,
71
+ inputs=gr.Textbox(label="Enter your query in plain English"),
72
  outputs=gr.Textbox(label="Query Result"),
73
+ title="Natural Language to SQL Executor",
74
+ description="Enter a plain English request, and the AI will generate an SQL query and return the results.",
75
  flagging_mode="never",
76
  )
77