Toggle Menu
Prev Up Next

# Source code for langchain_community.agent_toolkits.cassandra_database.toolkit

    """Apache Cassandra Toolkit."""
    from typing import List

    from langchain_core.pydantic_v1 import Field

    from langchain_community.agent_toolkits.base import BaseToolkit
    from langchain_community.tools import BaseTool
    from langchain_community.tools.cassandra_database.tool import (
        GetSchemaCassandraDatabaseTool,
        GetTableDataCassandraDatabaseTool,
        QueryCassandraDatabaseTool,
    )
    from langchain_community.utilities.cassandra_database import CassandraDatabase

    class CassandraDatabaseToolkit(BaseToolkit):
        """Toolkit for interacting with an Apache Cassandra database."""

        db: CassandraDatabase = Field(exclude=True)

        class Config:
            """Configuration for this pydantic object."""

            arbitrary_types_allowed = True

        def get_tools(self) -> List[BaseTool]:
            """Get the tools in the toolkit."""
            return [
                GetSchemaCassandraDatabaseTool(db=self.db),
                QueryCassandraDatabaseTool(db=self.db),
                GetTableDataCassandraDatabaseTool(db=self.db),
            ]