Source code for langchain_community.tools.cassandra_database.tool

"""与Apache Cassandra数据库交互的工具。"""
from __future__ import annotations

import traceback
from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Type, Union

from langchain_core.callbacks import CallbackManagerForToolRun
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.tools import BaseTool

from langchain_community.utilities.cassandra_database import CassandraDatabase

if TYPE_CHECKING:
    from cassandra.cluster import ResultSet


class BaseCassandraDatabaseTool(BaseModel):
    """与Apache Cassandra数据库交互的基本工具。"""

    db: CassandraDatabase = Field(exclude=True)

    class Config(BaseTool.Config):
        pass


class _QueryCassandraDatabaseToolInput(BaseModel):
    query: str = Field(..., description="A detailed and correct CQL query.")


class QueryCassandraDatabaseTool(BaseCassandraDatabaseTool, BaseTool):
    """工具用于使用提供的CQL查询Apache Cassandra数据库。"""

    name: str = "cassandra_db_query"
    description: str = """
    Execute a CQL query against the database and get back the result.
    If the query is not correct, an error message will be returned.
    If an error is returned, rewrite the query, check the query, and try again.
    """
    args_schema: Type[BaseModel] = _QueryCassandraDatabaseToolInput

    def _run(
        self,
        query: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> Union[str, Sequence[Dict[str, Any]], ResultSet]:
        """执行查询，返回结果或错误消息。"""
        try:
            return self.db.run(query)
        except Exception as e:
            """Format the error message"""
            return f"Error: {e}\n{traceback.format_exc()}"


class _GetSchemaCassandraDatabaseToolInput(BaseModel):
    keyspace: str = Field(
        ...,
        description=("The name of the keyspace for which to return the schema."),
    )


class GetSchemaCassandraDatabaseTool(BaseCassandraDatabaseTool, BaseTool):
    """用于获取Apache Cassandra数据库中keyspace的模式的工具。"""

    name: str = "cassandra_db_schema"
    description: str = """
    Input to this tool is a keyspace name, output is a table description 
    of Apache Cassandra tables.
    If the query is not correct, an error message will be returned.
    If an error is returned, report back to the user that the keyspace 
    doesn't exist and stop.
    """

    args_schema: Type[BaseModel] = _GetSchemaCassandraDatabaseToolInput

    def _run(
        self,
        keyspace: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        """获取keyspace的模式。"""
        try:
            tables = self.db.get_keyspace_tables(keyspace)
            return "".join([table.as_markdown() + "\n\n" for table in tables])
        except Exception as e:
            """Format the error message"""
            return f"Error: {e}\n{traceback.format_exc()}"


class _GetTableDataCassandraDatabaseToolInput(BaseModel):
    keyspace: str = Field(
        ...,
        description=("The name of the keyspace containing the table."),
    )
    table: str = Field(
        ...,
        description=("The name of the table for which to return data."),
    )
    predicate: str = Field(
        ...,
        description=("The predicate for the query that uses the primary key."),
    )
    limit: int = Field(
        ...,
        description=("The maximum number of rows to return."),
    )


class GetTableDataCassandraDatabaseTool(BaseCassandraDatabaseTool, BaseTool):
    """```python
# 用于从Apache Cassandra数据库中的表中获取数据的工具。
# 使用WHERE子句指定查询的谓词，该谓词使用主键。空谓词将返回所有行。尽量避免这种情况。
# 使用limit指定要返回的行数。空限制将返回所有行。
```"""

    name: str = "cassandra_db_select_table_data"
    description: str = """
    Tool for getting data from a table in an Apache Cassandra database. 
    Use the WHERE clause to specify the predicate for the query that uses the 
    primary key. A blank predicate will return all rows. Avoid this if possible. 
    Use the limit to specify the number of rows to return. A blank limit will 
    return all rows.
    """
    args_schema: Type[BaseModel] = _GetTableDataCassandraDatabaseToolInput

    def _run(
        self,
        keyspace: str,
        table: str,
        predicate: str,
        limit: int,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        """从keyspace中的表中获取数据。"""
        try:
            return self.db.get_table_data(keyspace, table, predicate, limit)
        except Exception as e:
            """Format the error message"""
            return f"Error: {e}\n{traceback.format_exc()}"