Waii

WaiiToolSpec

继承自：BaseToolSpec, BaseReader

源代码位于 llama-index-integrations/tools/llama-index-tools-waii/llama_index/tools/waii/base.py

class WaiiToolSpec(BaseToolSpec, BaseReader):
    spec_functions = [
        "get_answer",
        "describe_query",
        "performance_analyze",
        "diff_query",
        "describe_dataset",
        "transcode",
        "get_semantic_contexts",
        "generate_query_only",
        "run_query",
    ]

    def __init__(
        self,
        url: Optional[str] = None,
        api_key: Optional[str] = None,
        database_key: Optional[str] = None,
        verbose: Optional[bool] = False,
    ) -> None:
        from waii_sdk_py import WAII

        WAII.initialize(url=url, api_key=api_key)
        WAII.Database.activate_connection(key=database_key)
        self.verbose = verbose

    def _try_display(self, obj) -> None:
        # only display when verbose is True, we don't want to display too much information by default.
        if self.verbose:
            try:
                from IPython.display import display

                # display df if the function `display` is available (display only available when running with IPYTHON),
                # if it is not available, just ignore the exception.
                display(obj)
            except ImportError:
                # Handle the case where IPython is not available.
                pass

    def _run_query(self, sql: str, return_summary: bool) -> List[Document]:
        from waii_sdk_py import WAII
        from waii_sdk_py.query import RunQueryRequest

        run_result = WAII.Query.run(RunQueryRequest(query=sql))

        self._try_display(run_result.to_pandas_df())

        # create documents based on returned rows
        documents = [Document(text=str(doc)) for doc in run_result.rows]

        if return_summary:
            return self._get_summarization(
                "Summarize the result in text, don't miss any detail.", documents
            )

        return documents

    def load_data(self, ask: str) -> List[Document]:
        """
        Query using natural language and load data from the Database, returning a list of Documents.

        Args:
            ask: a natural language question.

        Returns:
            List[Document]: A list of Document objects.

        """
        query = self.generate_query_only(ask)

        return self._run_query(query, False)

    def _get_summarization(self, original_ask: str, documents) -> Any:
        texts = []

        n_chars = 0
        for i in range(len(documents)):
            t = str(documents[i].text)
            if len(t) + n_chars > 8192:
                texts.append(f"... {len(documents) - i} more results")
                break
            texts.append(t)
            n_chars += len(t)

        summarizer = TreeSummarize()
        return summarizer.get_response(original_ask, texts)

    def get_answer(self, ask: str) -> List[Document]:
        """
        Generate a SQL query and run it against the database, returning the summarization of the answer
        Args:
            ask: a natural language question.

        Returns:
            str: A string containing the summarization of the answer.

        """
        query = self.generate_query_only(ask)

        return self._run_query(query, True)

    def generate_query_only(self, ask: str) -> str:
        """
        Generate a SQL query and NOT run it, returning the query. If you need to get answer, you should use get_answer instead.

        Args:
            ask: a natural language question.

        Returns:
            str: A string containing the query.

        """
        from waii_sdk_py import WAII
        from waii_sdk_py.query import QueryGenerationRequest

        query = WAII.Query.generate(QueryGenerationRequest(ask=ask)).query

        self._try_display(query)

        return query

    def run_query(self, sql: str) -> List[Document]:
        return self._run_query(sql, False)

    def describe_query(self, question: str, query: str) -> str:
        """
        Describe a sql query, returning the summarization of the answer.

        Args:
            question: a natural language question which the people want to ask.
            query: a sql query.

        Returns:
            str: A string containing the summarization of the answer.

        """
        from waii_sdk_py import WAII
        from waii_sdk_py.query import DescribeQueryRequest

        result = WAII.Query.describe(DescribeQueryRequest(query=query))
        result = json.dumps(result.dict(), indent=2)
        self._try_display(result)

        return self._get_summarization(question, [Document(text=result)])

    def performance_analyze(self, query_uuid: str) -> str:
        """
        Analyze the performance of a query, returning the summarization of the answer.

        Args:
            query_uuid: a query uuid, e.g. xxxxxxxxxxxxx...

        Returns:
            str: A string containing the summarization of the answer.

        """
        from waii_sdk_py import WAII
        from waii_sdk_py.query import QueryPerformanceRequest

        result = WAII.Query.analyze_performance(
            QueryPerformanceRequest(query_id=query_uuid)
        )
        return json.dumps(result.dict(), indent=2)

    def diff_query(self, previous_query: str, current_query: str) -> str:
        """
        Diff two sql queries, returning the summarization of the answer.

        Args:
            previous_query: previous sql query.
            current_query: current sql query.

        Returns:
            str: A string containing the summarization of the answer.

        """
        from waii_sdk_py import WAII
        from waii_sdk_py.query import DiffQueryRequest

        result = WAII.Query.diff(
            DiffQueryRequest(query=current_query, previous_query=previous_query)
        )
        result = json.dumps(result.dict(), indent=2)
        return self._get_summarization("get diff summary", [Document(text=result)])

    def describe_dataset(
        self,
        ask: str,
        schema_name: Optional[str] = None,
        table_name: Optional[str] = None,
    ) -> str:
        """
        Describe a dataset (no matter if it is a table or schema), returning the summarization of the answer.
        Example questions like: "describe the dataset", "what the schema is about", "example question for the table xxx", etc.
        When both schema and table are None, describe the whole database.

        Args:
            ask: a natural language question (how you want to describe the dataset).
            schema_name: a schema name (shouldn't include the database name or the table name).
            table_name: a table name. (shouldn't include the database name or the schema name).

        Returns:
            str: A string containing the summarization of the answer.

        """
        from waii_sdk_py import WAII

        catalog = WAII.Database.get_catalogs()

        # filter by schema / table
        schemas = {}
        tables = {}

        for c in catalog.catalogs:
            for s in c.schemas:
                for t in s.tables:
                    if (
                        schema_name is not None
                        and schema_name.lower() != t.name.schema_name.lower()
                    ):
                        continue
                    if table_name is not None:
                        if table_name.lower() != t.name.table_name.lower():
                            continue
                        tables[str(t.name)] = t
                    schemas[str(s.name)] = s

        # remove tables ref from schemas
        for schema in schemas:
            schemas[schema].tables = None

        # generate response
        return self._get_summarization(
            ask + ", use the provided information to get comprehensive summarization",
            [Document(text=str(schemas[schema])) for schema in schemas]
            + [Document(text=str(tables[table])) for table in tables],
        )

    def transcode(
        self,
        instruction: Optional[str] = "",
        source_dialect: Optional[str] = None,
        source_query: Optional[str] = None,
        target_dialect: Optional[str] = None,
    ) -> str:
        """
        Transcode a sql query from one dialect to another, returning generated query.

        Args:
            instruction: instruction in natural language.
            source_dialect: the source dialect of the query.
            source_query: the source query.
            target_dialect: the target dialect of the query.

        Returns:
            str: A string containing the generated query.

        """
        from waii_sdk_py import WAII
        from waii_sdk_py.query import TranscodeQueryRequest

        result = WAII.Query.transcode(
            TranscodeQueryRequest(
                ask=instruction,
                source_dialect=source_dialect,
                source_query=source_query,
                target_dialect=target_dialect,
            )
        )
        return result.query

    def get_semantic_contexts(self) -> Any:
        """Get all pre-defined semantic contexts."""
        from waii_sdk_py import WAII

        return WAII.SemanticContext.get_semantic_context().semantic_context

load_data

load_data(ask: str) -> List[Document]

使用自然语言查询并从数据库加载数据，返回文档列表。

参数

名称	类型	描述	默认值
ask	str	一个自然语言问题。	必需

返回值

类型	描述
List[Document]	List[Document]：Document 对象的列表。

源代码位于 llama-index-integrations/tools/llama-index-tools-waii/llama_index/tools/waii/base.py

def load_data(self, ask: str) -> List[Document]:
    """
    Query using natural language and load data from the Database, returning a list of Documents.

    Args:
        ask: a natural language question.

    Returns:
        List[Document]: A list of Document objects.

    """
    query = self.generate_query_only(ask)

    return self._run_query(query, False)

get_answer

get_answer(ask: str) -> List[Document]

生成 SQL 查询并在数据库上运行，返回答案的摘要。参数：ask：一个自然语言问题。

返回值

名称	类型	描述
str	List[Document]	包含答案摘要的字符串。

源代码位于 llama-index-integrations/tools/llama-index-tools-waii/llama_index/tools/waii/base.py

def get_answer(self, ask: str) -> List[Document]:
    """
    Generate a SQL query and run it against the database, returning the summarization of the answer
    Args:
        ask: a natural language question.

    Returns:
        str: A string containing the summarization of the answer.

    """
    query = self.generate_query_only(ask)

    return self._run_query(query, True)

generate_query_only

generate_query_only(ask: str) -> str

生成 SQL 查询但不运行，返回查询本身。如果需要获取答案，应使用 get_answer。

参数

名称	类型	描述	默认值
ask	str	一个自然语言问题。	必需

返回值

名称	类型	描述
str	str	包含查询的字符串。

源代码位于 llama-index-integrations/tools/llama-index-tools-waii/llama_index/tools/waii/base.py

def generate_query_only(self, ask: str) -> str:
    """
    Generate a SQL query and NOT run it, returning the query. If you need to get answer, you should use get_answer instead.

    Args:
        ask: a natural language question.

    Returns:
        str: A string containing the query.

    """
    from waii_sdk_py import WAII
    from waii_sdk_py.query import QueryGenerationRequest

    query = WAII.Query.generate(QueryGenerationRequest(ask=ask)).query

    self._try_display(query)

    return query

describe_query

describe_query(question: str, query: str) -> str

描述一个 SQL 查询，返回答案的摘要。

参数

名称	类型	描述	默认值
question	str	人们想要问的自然语言问题。	必需
query	str	一个 SQL 查询。	必需

返回值

名称	类型	描述
str	str	包含答案摘要的字符串。

源代码位于 llama-index-integrations/tools/llama-index-tools-waii/llama_index/tools/waii/base.py

def describe_query(self, question: str, query: str) -> str:
    """
    Describe a sql query, returning the summarization of the answer.

    Args:
        question: a natural language question which the people want to ask.
        query: a sql query.

    Returns:
        str: A string containing the summarization of the answer.

    """
    from waii_sdk_py import WAII
    from waii_sdk_py.query import DescribeQueryRequest

    result = WAII.Query.describe(DescribeQueryRequest(query=query))
    result = json.dumps(result.dict(), indent=2)
    self._try_display(result)

    return self._get_summarization(question, [Document(text=result)])

performance_analyze

performance_analyze(query_uuid: str) -> str

分析查询的性能，返回答案的摘要。

参数

名称	类型	描述	默认值
query_uuid	str	一个查询 UUID，例如 xxxxxxxxxxxxx...	必需

返回值

名称	类型	描述
str	str	包含答案摘要的字符串。

源代码位于 llama-index-integrations/tools/llama-index-tools-waii/llama_index/tools/waii/base.py

def performance_analyze(self, query_uuid: str) -> str:
    """
    Analyze the performance of a query, returning the summarization of the answer.

    Args:
        query_uuid: a query uuid, e.g. xxxxxxxxxxxxx...

    Returns:
        str: A string containing the summarization of the answer.

    """
    from waii_sdk_py import WAII
    from waii_sdk_py.query import QueryPerformanceRequest

    result = WAII.Query.analyze_performance(
        QueryPerformanceRequest(query_id=query_uuid)
    )
    return json.dumps(result.dict(), indent=2)

diff_query

diff_query(previous_query: str, current_query: str) -> str

比较两个 SQL 查询的差异，返回答案的摘要。

参数

名称	类型	描述	默认值
previous_query	str	前一个 SQL 查询。	必需
current_query	str	当前 SQL 查询。	必需

返回值

名称	类型	描述
str	str	包含答案摘要的字符串。

源代码位于 llama-index-integrations/tools/llama-index-tools-waii/llama_index/tools/waii/base.py

def diff_query(self, previous_query: str, current_query: str) -> str:
    """
    Diff two sql queries, returning the summarization of the answer.

    Args:
        previous_query: previous sql query.
        current_query: current sql query.

    Returns:
        str: A string containing the summarization of the answer.

    """
    from waii_sdk_py import WAII
    from waii_sdk_py.query import DiffQueryRequest

    result = WAII.Query.diff(
        DiffQueryRequest(query=current_query, previous_query=previous_query)
    )
    result = json.dumps(result.dict(), indent=2)
    return self._get_summarization("get diff summary", [Document(text=result)])

describe_dataset

describe_dataset(ask: str, schema_name: Optional[str] = None, table_name: Optional[str] = None) -> str

描述一个数据集（无论是表还是模式），返回答案的摘要。示例问题如："描述数据集"，"模式是关于什么的"，"表 xxx 的示例问题" 等。当 schema 和 table 都为 None 时，描述整个数据库。

参数

名称	类型	描述	默认值
ask	str	自然语言问题（您希望如何描述数据集）。	必需
schema_name	Optional[str]	模式名称（不应包含数据库名称或表名称）。	无
table_name	Optional[str]	表名称（不应包含数据库名称或模式名称）。	无

返回值

名称	类型	描述
str	str	包含答案摘要的字符串。

源代码位于 llama-index-integrations/tools/llama-index-tools-waii/llama_index/tools/waii/base.py

def describe_dataset(
    self,
    ask: str,
    schema_name: Optional[str] = None,
    table_name: Optional[str] = None,
) -> str:
    """
    Describe a dataset (no matter if it is a table or schema), returning the summarization of the answer.
    Example questions like: "describe the dataset", "what the schema is about", "example question for the table xxx", etc.
    When both schema and table are None, describe the whole database.

    Args:
        ask: a natural language question (how you want to describe the dataset).
        schema_name: a schema name (shouldn't include the database name or the table name).
        table_name: a table name. (shouldn't include the database name or the schema name).

    Returns:
        str: A string containing the summarization of the answer.

    """
    from waii_sdk_py import WAII

    catalog = WAII.Database.get_catalogs()

    # filter by schema / table
    schemas = {}
    tables = {}

    for c in catalog.catalogs:
        for s in c.schemas:
            for t in s.tables:
                if (
                    schema_name is not None
                    and schema_name.lower() != t.name.schema_name.lower()
                ):
                    continue
                if table_name is not None:
                    if table_name.lower() != t.name.table_name.lower():
                        continue
                    tables[str(t.name)] = t
                schemas[str(s.name)] = s

    # remove tables ref from schemas
    for schema in schemas:
        schemas[schema].tables = None

    # generate response
    return self._get_summarization(
        ask + ", use the provided information to get comprehensive summarization",
        [Document(text=str(schemas[schema])) for schema in schemas]
        + [Document(text=str(tables[table])) for table in tables],
    )

transcode

transcode(instruction: Optional[str] = '', source_dialect: Optional[str] = None, source_query: Optional[str] = None, target_dialect: Optional[str] = None) -> str

将 SQL 查询从一种方言转码到另一种，返回生成的查询。

参数

名称	类型	描述	默认值
instruction	Optional[str]	自然语言指令。	''
source_dialect	Optional[str]	查询的源方言。	无
source_query	Optional[str]	源查询。	无
target_dialect	Optional[str]	查询的目标方言。	无

返回值

名称	类型	描述
str	str	包含生成的查询的字符串。

源代码位于 llama-index-integrations/tools/llama-index-tools-waii/llama_index/tools/waii/base.py

def transcode(
    self,
    instruction: Optional[str] = "",
    source_dialect: Optional[str] = None,
    source_query: Optional[str] = None,
    target_dialect: Optional[str] = None,
) -> str:
    """
    Transcode a sql query from one dialect to another, returning generated query.

    Args:
        instruction: instruction in natural language.
        source_dialect: the source dialect of the query.
        source_query: the source query.
        target_dialect: the target dialect of the query.

    Returns:
        str: A string containing the generated query.

    """
    from waii_sdk_py import WAII
    from waii_sdk_py.query import TranscodeQueryRequest

    result = WAII.Query.transcode(
        TranscodeQueryRequest(
            ask=instruction,
            source_dialect=source_dialect,
            source_query=source_query,
            target_dialect=target_dialect,
        )
    )
    return result.query

get_semantic_contexts

get_semantic_contexts() -> Any

获取所有预定义的语义上下文。

源代码位于 llama-index-integrations/tools/llama-index-tools-waii/llama_index/tools/waii/base.py

def get_semantic_contexts(self) -> Any:
    """Get all pre-defined semantic contexts."""
    from waii_sdk_py import WAII

    return WAII.SemanticContext.get_semantic_context().semantic_context