Sql retriever

Defines an abstract base class to create an SQL database retriever.

This module provides the BaseSQLRetriever class, which serves as a foundation for implementing database retrieval systems in Gen AI applications.

Authors

Berty C L Tobing (berty.c.l.tobing@gdplabs.id) Immanuel Rhesa (immanuel.rhesa@gdplabs.id)

References

NONE

`BaseSQLRetriever(sql_data_store, lm_request_processor=None, extract_func=None, preprocess_query_func=None)`

Bases: Component, ABC

An abstract base class for the SQL database retriever used in Gen AI applications.

This class defines the interface for SQL database retriever, which are responsible for retrieving data or information based on a given query.

Attributes:

Name	Type	Description
`sql_data_store`	`BaseSQLDataStore`	The SQL database data store to be used.
`lm_request_processor`	`LMRequestProcessor \| None`	The LMRequestProcessor instance to be used for modifying a failed query. If not None, will modify the query upon failure.
`extract_func`	`Callable[[str \| list[str] \| dict[str, str \| list[str]]], str \| list[str]]`	A function to extract the modified query from the LM output.
`preprocess_query_func`	`Callable[[str], str] \| None`	A function to preprocess the query before execution.

Initializes the BaseSQLRetriever object.

Parameters:

Name	Type	Description	Default
`sql_data_store`	`BaseSQLDataStore`	The SQL database data store to be used.	required
`lm_request_processor`	`LMRequestProcessor \| None`	The LMRequestProcessor instance to be used for modifying a failed query. If not None, will modify the query upon failure.	`None`
`extract_func`	`Callable[[str \| list[str] \| dict[str, str \| list[str]]], str \| list[str]]`	A function to extract the transformed query from the output. Defaults to None, in which case a default extractor will be used.	`None`
`preprocess_query_func`	`Callable[[str], str] \| None`	A function to preprocess the query before. Defaults to None.	`None`

Retrieve data based on the query.

This method should be implemented by subclasses to define the specific retrieval logic.

Parameters:

Name	Type	Description	Default
`query`	`str`	The query string to retrieve data.	required
`event_emitter`	`EventEmitter \| None`	The event emitter to emit events. Defaults to None.	`None`
`prompt_kwargs`	`dict[str, Any] \| None`	Additional keyword arguments for the prompt. Defaults to None.	`None`
`return_query`	`bool`	If True, returns a tuple of the executed query and the result. If False, returns only the result. Defaults to False.	`False`

Returns:

Type	Description
`DataFrame \| tuple[str, DataFrame]`	pd.DataFrame \| tuple[str, pd.DataFrame]: The result of the retrieval process. If return_query is True, returns a tuple of the executed query and the result. If return_query is False, returns only the result.

Raises:

Type	Description
`NotImplementedError`	If the method is not implemented.