Source code for langchain_community.utilities.dataherald

"""调用Dataherald的工具。"""
from typing import Any, Dict, Optional

from langchain_core.pydantic_v1 import BaseModel, Extra, root_validator
from langchain_core.utils import get_from_dict_or_env


class DataheraldAPIWrapper(BaseModel):
    """数据传送包装器。

    使用文档：

    1. 前往Dataherald并注册
    2. 创建一个API密钥
    3. 将您的API密钥保存到DATAHERALD_API_KEY环境变量中
    4. pip install dataherald

    """

    dataherald_client: Any  #: :meta private:
    db_connection_id: str
    dataherald_api_key: Optional[str] = None

    class Config:
        """此pydantic对象的配置。"""

        extra = Extra.forbid

    @root_validator()
    def validate_environment(cls, values: Dict) -> Dict:
        """验证环境中是否存在API密钥和Python包。"""
        dataherald_api_key = get_from_dict_or_env(
            values, "dataherald_api_key", "DATAHERALD_API_KEY"
        )
        values["dataherald_api_key"] = dataherald_api_key

        try:
            import dataherald

        except ImportError:
            raise ImportError(
                "dataherald is not installed. "
                "Please install it with `pip install dataherald`"
            )

        client = dataherald.Dataherald(api_key=dataherald_api_key)
        values["dataherald_client"] = client

        return values

    def run(self, prompt: str) -> str:
        """通过Dataherald生成一个SQL查询并解析结果。"""
        from dataherald.types.sql_generation_create_params import Prompt

        prompt_obj = Prompt(text=prompt, db_connection_id=self.db_connection_id)
        res = self.dataherald_client.sql_generations.create(prompt=prompt_obj)

        try:
            answer = res.sql
            if not answer:
                # We don't want to return the assumption alone if answer is empty
                return "No answer"
            else:
                return f"Answer: {answer}"

        except StopIteration:
            return "Dataherald wasn't able to answer it"