Source code for langchain_community.agent_toolkits.spark_sql.toolkit

"""Toolkit for interacting with Spark SQL."""
from typing import List

from langchain_core.language_models import BaseLanguageModel
from langchain_core.pydantic_v1 import Field
from langchain_core.tools import BaseToolkit

from langchain_community.tools import BaseTool
from langchain_community.tools.spark_sql.tool import (
    InfoSparkSQLTool,
    ListSparkSQLTool,
    QueryCheckerTool,
    QuerySparkSQLTool,
)
from langchain_community.utilities.spark_sql import SparkSQL


[docs]class SparkSQLToolkit(BaseToolkit): """Toolkit for interacting with Spark SQL.""" db: SparkSQL = Field(exclude=True) llm: BaseLanguageModel = Field(exclude=True) class Config: """Configuration for this pydantic object.""" arbitrary_types_allowed = True
[docs] def get_tools(self) -> List[BaseTool]: """Get the tools in the toolkit.""" return [ QuerySparkSQLTool(db=self.db), InfoSparkSQLTool(db=self.db), ListSparkSQLTool(db=self.db), QueryCheckerTool(db=self.db, llm=self.llm), ]