Spaces:
Runtime error
Runtime error
File size: 1,063 Bytes
ed4d993 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 |
"""Tool for the Dataherald Hosted API"""
from typing import Optional, Type
from langchain_core.callbacks import CallbackManagerForToolRun
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain_core.tools import BaseTool
from langchain_community.utilities.dataherald import DataheraldAPIWrapper
class DataheraldTextToSQLInput(BaseModel):
prompt: str = Field(
description="Natural language query to be translated to a SQL query."
)
class DataheraldTextToSQL(BaseTool):
"""Tool that queries using the Dataherald SDK."""
name: str = "dataherald"
description: str = (
"A wrapper around Dataherald. "
"Text to SQL. "
"Input should be a prompt and an existing db_connection_id"
)
api_wrapper: DataheraldAPIWrapper
args_schema: Type[BaseModel] = DataheraldTextToSQLInput
def _run(
self,
prompt: str,
run_manager: Optional[CallbackManagerForToolRun] = None,
) -> str:
"""Use the Dataherald tool."""
return self.api_wrapper.run(prompt)
|