Talk2DatabaseAgent / app /utils /database_connection.py
Kshitijk20's picture
added modular code
34a5262
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit
class DatabaseConnection:
def __init__(self,connection_string: str):
self.db = None
self.setup_database_connection(connection_string)
def setup_database_connection(self, connection_string: str):
"""Set up database connection and initialize tools"""
try:
# Initialize database connection
self.db = SQLDatabase.from_uri(connection_string)
print("Database connection successful!")
return self.db
except ImportError as e:
print(f"Database driver import error: {str(e)}")
raise ValueError(f"Missing database driver or invalid database type: {str(e)}")
except ValueError as e:
print(f"Invalid connection string or configuration: {str(e)}")
raise
except Exception as e:
print(f"Unexpected error during database connection: {str(e)}")
raise ValueError(f"Failed to establish database connection: {str(e)}")