feat: remove unecessary code

This commit is contained in:
Eduardo Chiarotti
2024-07-24 21:38:52 -03:00
parent 2df29f3dde
commit 6fd02cdf82

View File

@@ -40,11 +40,11 @@ class NL2SQL(BaseTool):
def _run(self, sql_query: str): def _run(self, sql_query: str):
try: try:
data = self.execute_sql(sql_query) data = self.execute_sql(sql_query)
except Exception: except Exception as exc:
data = ( data = (
f"Based on these tables {self.tables} and columns {self.columns}, " f"Based on these tables {self.tables} and columns {self.columns}, "
"you can create SQL queries to retrieve data from the database." "you can create SQL queries to retrieve data from the database."
f"Get the original request {sql_query} and try to create a SQL query that retrieves the requested data." f"Get the original request {sql_query} and the error {exc} and create the correct SQL query."
) )
return data return data
@@ -53,7 +53,6 @@ class NL2SQL(BaseTool):
engine = create_engine(self.db_uri) engine = create_engine(self.db_uri)
Session = sessionmaker(bind=engine) Session = sessionmaker(bind=engine)
session = Session() session = Session()
try: try:
result = session.execute(text(sql_query)) result = session.execute(text(sql_query))
session.commit() session.commit()
@@ -64,11 +63,8 @@ class NL2SQL(BaseTool):
return data return data
else: else:
return f"Query {sql_query} executed successfully" return f"Query {sql_query} executed successfully"
except Exception as e: except Exception as e:
session.rollback() session.rollback()
print(f"SQL execution error: {e}")
raise e raise e
finally: finally:
session.close() session.close()