Spaces:
Runtime error
Runtime error
File size: 3,463 Bytes
37a6ee1 | 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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | import os
import json
import asyncio
from backend.services.schema_service import schema_service
from backend.services.llm_service import llm_service
from backend.services.sql_service import sql_service
# Clean up db first to prevent already exists
db_path = "databases/schema/empdept.db"
if os.path.exists(db_path):
os.remove(db_path)
sql = """
CREATE TABLE DEPT (
DEPTNO INT PRIMARY KEY,
DNAME VARCHAR(14),
LOC VARCHAR(13)
);
CREATE TABLE EMP (
EMPNO INT PRIMARY KEY,
ENAME VARCHAR(10),
JOB VARCHAR(9),
MGR INT,
HIREDATE DATETIME,
SAL DECIMAL(7,2),
COMM DECIMAL(7,2),
DEPTNO INT,
CONSTRAINT FK_DEPTNO FOREIGN KEY (DEPTNO) REFERENCES DEPT(DEPTNO),
CONSTRAINT FK_MGR FOREIGN KEY (MGR) REFERENCES EMP(EMPNO)
);
CREATE TABLE SALGRADE (
GRADE INT PRIMARY KEY,
LOSAL INT,
HISAL INT
);
INSERT INTO DEPT VALUES (10, 'ACCOUNTING', 'NEW YORK');
INSERT INTO DEPT VALUES (20, 'RESEARCH', 'DALLAS');
INSERT INTO DEPT VALUES (30, 'SALES', 'CHICAGO');
INSERT INTO DEPT VALUES (40, 'OPERATIONS', 'BOSTON');
INSERT INTO EMP VALUES (7839, 'KING', 'PRESIDENT', NULL, '1981-11-17', 5000, NULL, 10);
INSERT INTO EMP VALUES (7698, 'BLAKE', 'MANAGER', 7839, '1981-05-01', 2850, NULL, 30);
INSERT INTO EMP VALUES (7782, 'CLARK', 'MANAGER', 7839, '1981-06-09', 2450, NULL, 10);
INSERT INTO EMP VALUES (7566, 'JONES', 'MANAGER', 7839, '1981-04-02', 2975, NULL, 20);
INSERT INTO EMP VALUES (7654, 'MARTIN', 'SALESMAN', 7698, '1981-09-28', 1250, 1400, 30);
INSERT INTO EMP VALUES (7499, 'ALLEN', 'SALESMAN', 7698, '1981-02-20', 1600, 300, 30);
INSERT INTO EMP VALUES (7844, 'TURNER', 'SALESMAN', 7698, '1981-09-08', 1500, 0, 30);
INSERT INTO EMP VALUES (7900, 'JAMES', 'CLERK', 7698, '1981-12-03', 950, NULL, 30);
INSERT INTO EMP VALUES (7521, 'WARD', 'SALESMAN', 7698, '1981-02-22', 1250, 500, 30);
INSERT INTO EMP VALUES (7902, 'FORD', 'ANALYST', 7566, '1981-12-03', 3000, NULL, 20);
INSERT INTO EMP VALUES (7369, 'SMITH', 'CLERK', 7902, '1980-12-17', 800, NULL, 20);
INSERT INTO EMP VALUES (7788, 'SCOTT', 'ANALYST', 7566, '1982-12-09', 3000, NULL, 20);
INSERT INTO EMP VALUES (7876, 'ADAMS', 'CLERK', 7788, '1983-01-12', 1100, NULL, 20);
INSERT INTO EMP VALUES (7934, 'MILLER', 'CLERK', 7782, '1982-01-23', 1300, NULL, 10);
INSERT INTO SALGRADE VALUES (1, 700, 1200);
INSERT INTO SALGRADE VALUES (2, 1201, 1400);
INSERT INTO SALGRADE VALUES (3, 1401, 2000);
INSERT INTO SALGRADE VALUES (4, 2001, 3000);
INSERT INTO SALGRADE VALUES (5, 3001, 9999);
"""
def run_tests():
print("Testing register_schema_db...")
info = schema_service.register_schema_db("empdept", schema_sql_content=sql)
print("Relationships:")
for rel in info['relationships']:
print(f" - {rel['from_table']}.{rel['from_column']} -> {rel['to_table']}.{rel['to_column']}")
questions = [
"What are the names of employees in the SALES department?",
"Which employees earn a grade 4 salary?",
"Who is SMITH's manager?"
]
for q in questions:
print(f"\\nQ: {q}")
res_llm = llm_service.generate_sql_schema(q, info)
sql_query = res_llm["sql"]
print(f"SQL: {sql_query}")
try:
res = sql_service.execute_query("empdept", sql_query, mode="schema")
print(f"Results: {res['rows']}")
except Exception as e:
print(f"Error executing: {e}")
if __name__ == '__main__':
run_tests()
|