SQLRetriever output unrelated information

[JonOnEarth]

Hi, I used SQLRetriever to retriever the Database, it returned something like:

[NodeWithScore(node=TextNode(id_='41595ebd-ee90-44c0-b99d-ab27e78fef3e', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, text="[('descrip...

It will be the output to my next module in the query pipeline with synthesis response, and the answer will use unrelated information such as the node information.
such that:

Based on the data, it appears that Tik is working better. The top-performing type is TextNode,...

Sometimes it even outputs the node id...

I checked the original code of SQLRetriever, it outputs the format like the example above, may I ask the reason, is there someone who met this issue?

[dosubot[bot]]

Hey @JonOnEarth! I'm here to help you with any bugs, questions, or contributions you have. Let's tackle this issue together!

The SQLRetriever is returning output that includes unrelated information such as node details because the retrieved nodes are not being filtered or processed to exclude irrelevant data before being passed to the next module. To ensure that only relevant data is passed to the next module, you can use node postprocessors to filter and process the nodes.

Here is how you can ensure that only relevant data is passed to the next module:

1. Define a custom node postprocessor that filters out unrelated information.
2. Add this custom postprocessor to the NLSQLRetriever.

Example of a custom node postprocessor:

from llama_index.legacy.schema import NodeWithScore, QueryBundle
from typing import List

class CustomNodePostprocessor:
    def postprocess_nodes(self, nodes: List[NodeWithScore], query_bundle: QueryBundle) -> List[NodeWithScore]:
        # Implement your filtering logic here
        filtered_nodes = [node for node in nodes if self.is_relevant(node)]
        return filtered_nodes

    def is_relevant(self, node: NodeWithScore) -> bool:
        # Define your criteria for relevance
        return "relevant_keyword" in node.node.text

Adding the custom postprocessor to the NLSQLRetriever:

custom_postprocessor = CustomNodePostprocessor()

class CustomNLSQLRetriever(NLSQLRetriever):
    def retrieve_with_metadata(self, str_or_query_bundle: QueryType) -> Tuple[List[NodeWithScore], Dict]:
        retrieved_nodes, metadata = super().retrieve_with_metadata(str_or_query_bundle)
        filtered_nodes = custom_postprocessor.postprocess_nodes(retrieved_nodes, str_or_query_bundle)
        return filtered_nodes, metadata

By using a custom node postprocessor, you can filter out unrelated information and ensure that only relevant data is passed to the next module ^[1][2][3].

To continue talking to Dosu, mention @dosu.

[JonOnEarth]

Thanks @dosu, I just want the SQLRetriever to output the text data. When I say irrelevant data mean it outputs the whole structure of the data. How can I just get the text data from the original output:

[NodeWithScore(node=TextNode(id_='41595ebd-ee90-44c0-b99d-ab27e78fef3e', embedding=None, metadata={}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={}, text="[('descrip...

to:

"descrip..."

[dosubot[bot]]

To modify the SQLRetriever to output only the text data instead of the entire structure, you need to set the return_raw parameter to False when initializing the SQLRetriever. This will ensure that the results are formatted to include only the text data.

Here is the relevant part of the code where you can make this change:

class NLSQLRetriever(BaseRetriever, PromptMixin):
    """Text-to-SQL Retriever.

    Args:
        sql_database (SQLDatabase): SQL database.
        text_to_sql_prompt (BasePromptTemplate): Prompt template for text-to-sql.
        context_query_kwargs (dict): Mapping from table name to context query.
        tables (Union[List[str], List[Table]]): List of table names or Table objects.
        table_retriever (ObjectRetriever[SQLTableSchema]): Object retriever for
            SQLTableSchema objects.
        context_str_prefix (str): Prefix for context string.
        service_context (ServiceContext): Service context.
        return_raw (bool): Whether to return plain-text dump of SQL results, or parsed into Nodes.
        handle_sql_errors (bool): Whether to handle SQL errors.
        sql_only (bool): Whether to get only sql and not the sql query result.
        llm (Optional[LLM]): Language model to use.

    """

    def __init__(
        self,
        sql_database: SQLDatabase,
        text_to_sql_prompt: Optional[BasePromptTemplate] = None,
        context_query_kwargs: Optional[dict] = None,
        tables: Optional[Union[List[str], List[Table]]] = None,
        table_retriever: Optional[ObjectRetriever[SQLTableSchema]] = None,
        context_str_prefix: Optional[str] = None,
        sql_parser_mode: SQLParserMode = SQLParserMode.DEFAULT,
        llm: Optional[LLMType] = "default",
        service_context: Optional[ServiceContext] = None,
        return_raw: bool = False,  # Set return_raw to False
        handle_sql_errors: bool = True,
        sql_only: bool = False,
        callback_manager: Optional[CallbackManager] = None,
        verbose: bool = False,
        **kwargs: Any,
    ) -> None:
        """Initialize params."""
        self._sql_retriever = SQLRetriever(sql_database, return_raw=return_raw)
        self._sql_database = sql_database
        self._get_tables = self._load_get_tables_fn(
            sql_database, tables, context_query_kwargs, table_retriever
        )
        self._context_str_prefix = context_str_prefix
        self._service_context = service_context or ServiceContext.from_defaults(llm=llm)
        self._text_to_sql_prompt = text_to_sql_prompt or DEFAULT_TEXT_TO_SQL_PROMPT
        self._sql_parser_mode = sql_parser_mode
        self._sql_parser = self._load_sql_parser(sql_parser_mode, self._service_context)
        self._handle_sql_errors = handle_sql_errors
        self._sql_only = sql_only
        self._verbose = verbose
        super().__init__(callback_manager)

By setting return_raw to False, the SQLRetriever will format the results to include only the text data ^[1][2].

To continue talking to Dosu, mention @dosu.

[JonOnEarth]

It still gives me the output like:

[NodeWithScore(node=TextNode(id_='f3496669-fd8c-43be-9d83-33dbaa51b5fc', embedding=None, metadata={'winning': 0, 'testing': 0, 'losing': 0, 'fatiguing': 0}, excluded_embed_metadata_keys=[], excluded_l...

Can you @dosu look at the code inside the SQLRetriever:

def retrieve_with_metadata(
        self, str_or_query_bundle: QueryType
    ) -> Tuple[List[NodeWithScore], Dict]:
        """Retrieve with metadata."""
        if isinstance(str_or_query_bundle, str):
            query_bundle = QueryBundle(str_or_query_bundle)
        else:
            query_bundle = str_or_query_bundle
        raw_response_str, metadata = self._sql_database.run_sql(query_bundle.query_str)
        if self._return_raw:
            return [NodeWithScore(node=TextNode(text=raw_response_str))], metadata
        else:
            # return formatted
            results = metadata["result"]
            col_keys = metadata["col_keys"]
            return self._format_node_results(results, col_keys), metadata

there is no option to output only the text data, always include the

[NodeWithScore(node=TextNode(id_='f3496669-fd8c-43be-9d83-33dbaa51b5fc', embedding=None, metadata={...

[dosubot[bot]]

To ensure that the SQLRetriever outputs only the text data, you can modify the retrieve_with_metadata method to extract and return just the text content from the TextNode objects. Here is how you can adjust the method:

1. Add a utility function to extract text from NodeWithScore objects.
2. Modify the retrieve_with_metadata method to use this utility function when return_raw is set to False.

Here is the updated code:

def retrieve_with_metadata(
    self, str_or_query_bundle: QueryType
) -> Tuple[List[str], Dict]:
    """Retrieve with metadata."""
    if isinstance(str_or_query_bundle, str):
        query_bundle = QueryBundle(str_or_query_bundle)
    else:
        query_bundle = str_or_query_bundle
    raw_response_str, metadata = self._sql_database.run_sql(query_bundle.query_str)
    if self._return_raw:
        return [raw_response_str], metadata
    else:
        # return formatted
        results = metadata["result"]
        col_keys = metadata["col_keys"]
        nodes_with_scores = self._format_node_results(results, col_keys)
        # Extract text from nodes
        text_data = [node.node.text for node in nodes_with_scores]
        return text_data, metadata

This modification ensures that when return_raw is set to False, the method will return a list of text data extracted from the TextNode objects instead of the entire node structure. This should give you the desired output format.

To continue talking to Dosu, mention @dosu.