BaseText2SQLView

dbally.views.freeform.text2sql.view.BaseText2SQLView

BaseText2SQLView(engine: Engine)

Bases: BaseView, ABC

Text2SQLFreeformView is a class designed to interact with the database using text2sql queries.

Source code in src/dbally/views/freeform/text2sql/view.py

def __init__(
    self,
    engine: Engine,
) -> None:
    super().__init__()
    self._engine = engine
    self._table_index = {table.name: table for table in self.get_tables()}

list_few_shots

list_few_shots() -> List[FewShotExample]

List all examples to be injected into few-shot prompt.

RETURNS	DESCRIPTION
List[FewShotExample]	List of few-shot examples

Source code in src/dbally/views/base.py

def list_few_shots(self) -> List[FewShotExample]:
    """
    List all examples to be injected into few-shot prompt.

    Returns:
        List of few-shot examples
    """
    return []

get_tables abstractmethod

get_tables() -> List[TableConfig]

Get the tables used by the view.

RETURNS	DESCRIPTION
List[TableConfig]	The list of tables used by the view.

Source code in src/dbally/views/freeform/text2sql/view.py

@abstractmethod
def get_tables(self) -> List[TableConfig]:
    """
    Get the tables used by the view.

    Returns:
        The list of tables used by the view.
    """

ask async

ask(query: str, llm: LLM, event_tracker: Optional[EventTracker] = None, n_retries: int = 3, dry_run: bool = False, llm_options: Optional[LLMOptions] = None) -> ViewExecutionResult

Executes the query and returns the result. It generates the SQL query from the natural language query and executes it against the database. It retries the process in case of errors.

PARAMETER	DESCRIPTION
query	The natural language query to execute. TYPE: str
llm	The LLM used to execute the query. TYPE: LLM
event_tracker	The event tracker used to audit the query execution. TYPE: Optional[EventTracker] DEFAULT: None
n_retries	The number of retries to execute the query in case of errors. TYPE: int DEFAULT: 3
dry_run	If True, the query will not be used to fetch data from the datasource. TYPE: bool DEFAULT: False
llm_options	Options to use for the LLM. TYPE: Optional[LLMOptions] DEFAULT: None

RETURNS	DESCRIPTION
ViewExecutionResult	The result of the query.

RAISES	DESCRIPTION
Text2SQLError	If the text2sql query generation fails after n_retries.

Source code in src/dbally/views/freeform/text2sql/view.py

async def ask(
    self,
    query: str,
    llm: LLM,
    event_tracker: Optional[EventTracker] = None,
    n_retries: int = 3,
    dry_run: bool = False,
    llm_options: Optional[LLMOptions] = None,
) -> ViewExecutionResult:
    """
    Executes the query and returns the result. It generates the SQL query from the natural language query and
    executes it against the database. It retries the process in case of errors.

    Args:
        query: The natural language query to execute.
        llm: The LLM used to execute the query.
        event_tracker: The event tracker used to audit the query execution.
        n_retries: The number of retries to execute the query in case of errors.
        dry_run: If True, the query will not be used to fetch data from the datasource.
        llm_options: Options to use for the LLM.

    Returns:
        The result of the query.

    Raises:
        Text2SQLError: If the text2sql query generation fails after n_retries.
    """
    sql, rows = None, None
    exceptions = []

    tables = self.get_tables()
    examples = self.list_few_shots()

    prompt_format = SQLGenerationPromptFormat(
        question=query,
        dialect=self._engine.dialect.name,
        tables=tables,
        examples=examples,
    )
    formatted_prompt = SQL_GENERATION_TEMPLATE.format_prompt(prompt_format)

    for _ in range(n_retries + 1):
        # We want to catch all exceptions to retry the process.
        # pylint: disable=broad-except
        try:
            sql, parameters, formatted_prompt = await self._generate_sql(
                conversation=formatted_prompt,
                llm=llm,
                event_tracker=event_tracker,
                llm_options=llm_options,
            )

            if dry_run:
                return ViewExecutionResult(results=[], context={"sql": sql})

            rows = await self._execute_sql(sql, parameters, event_tracker=event_tracker)
            break
        except Exception as e:
            formatted_prompt = formatted_prompt.add_user_message(f"Response is invalid! Error: {e}")
            exceptions.append(e)
            continue

    if rows is None:
        raise Text2SQLError("Text2SQL query generation failed", exceptions=exceptions) from exceptions[-1]

    # The underscore is used by sqlalchemy to avoid conflicts with column names
    # pylint: disable=protected-access
    return ViewExecutionResult(
        results=[dict(row._mapping) for row in rows],
        context={
            "sql": sql,
        },
    )

list_similarity_indexes

list_similarity_indexes() -> Dict[AbstractSimilarityIndex, List[IndexLocation]]

List all similarity indexes used by the view.

RETURNS	DESCRIPTION
Dict[AbstractSimilarityIndex, List[IndexLocation]]	Mapping of similarity indexes to their locations in the (view_name, table_name, column_name) format.

Source code in src/dbally/views/freeform/text2sql/view.py

def list_similarity_indexes(self) -> Dict[AbstractSimilarityIndex, List[IndexLocation]]:
    """
    List all similarity indexes used by the view.

    Returns:
        Mapping of similarity indexes to their locations in the (view_name, table_name, column_name) format.
    """
    indexes = defaultdict(list)
    for table in self.get_tables():
        for column in table.columns:
            if column.similarity_index:
                indexes[column.similarity_index].append((self.__class__.__name__, table.name, column.name))
    return indexes

dbally.views.freeform.text2sql.config.TableConfig dataclass

TableConfig(name: str, columns: List[ColumnConfig], description: Optional[str] = None)

Configuration of a table used in the Text2SQL view.

name instance-attribute

name: str

columns instance-attribute

columns: List[ColumnConfig]

description class-attribute instance-attribute

description: Optional[str] = None

ddl property

ddl: str

Returns the DDL for the table which can be provided to the LLM as a context.

RETURNS	DESCRIPTION
str	The DDL for the table.

dbally.views.freeform.text2sql.config.ColumnConfig dataclass

ColumnConfig(name: str, data_type: str, description: Optional[str] = None, similarity_index: Optional[SimilarityIndex] = None)

Configuration of a column used in the Text2SQL view.

name instance-attribute

name: str

data_type instance-attribute

data_type: str

description class-attribute instance-attribute

description: Optional[str] = None

similarity_index class-attribute instance-attribute

similarity_index: Optional[SimilarityIndex] = None