Context Store

The Context Store module is responsible for storing and retrieving context about the connected Databases and the business logic of the stored data. This information is subsequently accessed and used by the Text-to-SQL module for prompt engineering to generate accurate SQL to answer the natural language question.

There is currently a single implementation of the Context Store which accesses both the Vector store and Application Storage componenets to store and find closest validated SQL queries and information about the DB tables and rows.

Abstract Context Store Class

All implementations of the context store must inherit from the abstract ContextStore class and implement the following methods.

ContextStore

This abstract class provides a consistent interface for working with context retrieval and management.

doc_store_collection

The name of the collection in the database where context information is stored.

__init__(self, system: System)

Initializes the ContextStore instance.

Parameters:

system (System) – The system object.

retrieve_context_for_question(self, nl_question: str, number_of_samples: int = 3) → Tuple[List[dict] | None, List[dict] | None]

Given a natural language question, this method retrieves a single string containing information about relevant data stores, tables, and columns necessary for building the SQL query. This information includes example questions, corresponding SQL queries, and metadata about the tables (e.g., categorical columns). The retrieved string is then passed to the text-to-SQL generator.

Parameters:

• nl_question (str) – The natural language question.

• number_of_samples (int) – The number of context samples to retrieve.

Returns:

A list of dictionaries containing context information for generating SQL (contain few-shot samples and instructions).

Return type:

Tuple[List[dict] | None, List[dict] | None]

add_golden_records(self, golden_records: List[GoldenRecordRequest]) → List[GoldenRecord]

This method adds NL <> SQL pairs to the context store. These pairs serve as examples and will be included in prompts to the Language Model.

Parameters:

golden_records (List[GoldenRecordRequest]) – A list of natural language questions and the SQL that can be used to answer them.

Returns:

A list of GoldenRecord objects that were added to the context store.

Return type:

List[GoldenRecord]

remove_golden_records(self, ids: List) → bool

This method removes NL <> SQL pairs that were stored in the context store so they are not used in few-shot prompting.

Parameters:

ids (List) – A list of ids of the golden records to be removed.

Returns:

True if the records were successfully removed.

Return type:

bool