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()