Spaces:
Runtime error
Runtime error
"""Toolkit for interacting with a SQL database.""" | |
from typing import List | |
from pydantic import Field | |
from langchain.agents.agent_toolkits.base import BaseToolkit | |
from langchain.sql_database import SQLDatabase | |
from langchain.tools import BaseTool | |
from langchain.tools.sql_database.tool import ( | |
InfoSQLDatabaseTool, | |
ListSQLDatabaseTool, | |
QueryCheckerTool, | |
QuerySQLDataBaseTool, ClarifyTool, | |
) | |
class SQLDatabaseToolkit(BaseToolkit): | |
"""Toolkit for interacting with SQL databases.""" | |
db: SQLDatabase = Field(exclude=True) | |
def dialect(self) -> str: | |
"""Return string representation of dialect to use.""" | |
return self.db.dialect | |
class Config: | |
"""Configuration for this pydantic object.""" | |
arbitrary_types_allowed = True | |
def get_tools(self) -> List[BaseTool]: | |
"""Get the tools in the toolkit.""" | |
return [ | |
QuerySQLDataBaseTool(db=self.db), | |
InfoSQLDatabaseTool(db=self.db), | |
ListSQLDatabaseTool(db=self.db), | |
QueryCheckerTool(db=self.db), | |
] | |
class SimpleSQLDatabaseToolkit(BaseToolkit): | |
"""Toolkit for interacting with SQL databases.""" | |
db: SQLDatabase = Field(exclude=True) | |
def dialect(self) -> str: | |
"""Return string representation of dialect to use.""" | |
return self.db.dialect | |
class Config: | |
"""Configuration for this pydantic object.""" | |
arbitrary_types_allowed = True | |
def get_tools(self) -> List[BaseTool]: | |
"""Get the tools in the toolkit.""" | |
return [ | |
QuerySQLDataBaseTool(db=self.db), | |
] | |