fix: SQLDatabase has no attribute set_event_manager (#4192)

This commit is contained in:
Marcelo Nunes Alves 2024-11-03 21:20:33 -03:00 committed by GitHub
commit dd4a9f908c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -2,26 +2,33 @@ from langchain_community.utilities.sql_database import SQLDatabase
from sqlalchemy import create_engine
from sqlalchemy.pool import StaticPool
from langflow.custom import CustomComponent
from langflow.custom import Component
from langflow.io import (
Output,
StrInput,
)
class SQLDatabaseComponent(CustomComponent):
class SQLDatabaseComponent(Component):
display_name = "SQLDatabase"
description = "SQL Database"
name = "SQLDatabase"
def build_config(self):
return {
"uri": {"display_name": "URI", "info": "URI to the database."},
}
inputs = [
StrInput(name="uri", display_name="URI", info="URI to the database.", required=True),
]
outputs = [
Output(display_name="SQLDatabase", name="SQLDatabase", method="build_sqldatabase"),
]
def clean_up_uri(self, uri: str) -> str:
if uri.startswith("postgres://"):
uri = uri.replace("postgres://", "postgresql://")
return uri.strip()
def build(self, uri: str) -> SQLDatabase:
uri = self.clean_up_uri(uri)
def build_sqldatabase(self) -> SQLDatabase:
uri = self.clean_up_uri(self.uri)
# Create an engine using SQLAlchemy with StaticPool
engine = create_engine(uri, poolclass=StaticPool)
return SQLDatabase(engine)