MindsDB 25.1.4.0__py3-none-any.whl → 25.1.5.1__py3-none-any.whl

mindsdb/integrations/utilities/files/file_reader.py

@@ -0,0 +1,258 @@
+import traceback
+import json
+import csv
+from io import BytesIO, StringIO
+from pathlib import Path
+import codecs
+
+import filetype
+import pandas as pd
+from charset_normalizer import from_bytes
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+from mindsdb.utilities import log
+
+logger = log.getLogger(__name__)
+
+DEFAULT_CHUNK_SIZE = 500
+DEFAULT_CHUNK_OVERLAP = 250
+
+
+class FileDetectError(Exception):
+    ...
+
+
+def decode(file_obj: BytesIO) -> StringIO:
+    byte_str = file_obj.read()
+    # Move it to StringIO
+    try:
+        # Handle Microsoft's BOM "special" UTF-8 encoding
+        if byte_str.startswith(codecs.BOM_UTF8):
+            data_str = StringIO(byte_str.decode("utf-8-sig"))
+        else:
+            file_encoding_meta = from_bytes(
+                byte_str[: 32 * 1024],
+                steps=32,  # Number of steps/block to extract from my_byte_str
+                chunk_size=1024,  # Set block size of each extraction)
+                explain=False,
+            )
+            best_meta = file_encoding_meta.best()
+            errors = "strict"
+            if best_meta is not None:
+                encoding = file_encoding_meta.best().encoding
+
+                try:
+                    data_str = StringIO(byte_str.decode(encoding, errors))
+                except UnicodeDecodeError:
+                    encoding = "utf-8"
+                    errors = "replace"
+
+                    data_str = StringIO(byte_str.decode(encoding, errors))
+            else:
+                encoding = "utf-8"
+                errors = "replace"
+
+                data_str = StringIO(byte_str.decode(encoding, errors))
+    except Exception as e:
+        logger.error(traceback.format_exc())
+        raise FileDetectError("Could not load into string") from e
+
+    return data_str
+
+
+class FormatDetector:
+
+    def get(self, name, file_obj: BytesIO = None):
+        format = self.get_format_by_name(name)
+        if format is None and file_obj is not None:
+            format = self.get_format_by_content(file_obj)
+
+        if format is not None:
+            return format
+        raise FileDetectError(f'Unable to detect format: {name}')
+
+    def get_format_by_name(self, filename):
+        extension = Path(filename).suffix.strip(".").lower()
+        if extension == "tsv":
+            extension = "csv"
+        return extension or None
+
+    def get_format_by_content(self, file_obj):
+        if self.is_parquet(file_obj):
+            return "parquet"
+
+        file_type = filetype.guess(file_obj)
+        if file_type is None:
+            return
+
+        if file_type.mime in {
+            "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+            "application/vnd.ms-excel",
+        }:
+            return 'xlsx'
+
+        if file_type.mime == 'application/pdf':
+            return "pdf"
+
+        file_obj = decode(file_obj)
+
+        if self.is_json(file_obj):
+            return "json"
+
+        if self.is_csv(file_obj):
+            return "csv"
+
+    def is_json(self, data_obj: StringIO) -> bool:
+        # see if its JSON
+        text = data_obj.read(100).strip()
+        data_obj.seek(0)
+        if len(text) > 0:
+            # it looks like a json, then try to parse it
+            if text.startswith("{") or text.startswith("["):
+                try:
+                    json.loads(data_obj.read())
+                    return True
+                except Exception:
+                    return False
+                finally:
+                    data_obj.seek(0)
+        return False
+
+    def is_csv(self, data_obj: StringIO) -> bool:
+        sample = data_obj.readline()  # trying to get dialect from header
+        data_obj.seek(0)
+        try:
+            csv.Sniffer().sniff(sample)
+
+        except Exception:
+            return False
+
+    def is_parquet(self, data: BytesIO) -> bool:
+        # Check first and last 4 bytes equal to PAR1.
+        # Refer: https://parquet.apache.org/docs/file-format/
+        parquet_sig = b"PAR1"
+        data.seek(0, 0)
+        start_meta = data.read(4)
+        data.seek(-4, 2)
+        end_meta = data.read()
+        data.seek(0)
+        if start_meta == parquet_sig and end_meta == parquet_sig:
+            return True
+        return False
+
+
+class FileReader:
+
+    def _get_csv_dialect(self, buffer) -> csv.Dialect:
+        sample = buffer.readline()  # trying to get dialect from header
+        buffer.seek(0)
+        try:
+            if isinstance(sample, bytes):
+                sample = sample.decode()
+            accepted_csv_delimiters = [",", "\t", ";"]
+            try:
+                dialect = csv.Sniffer().sniff(
+                    sample, delimiters=accepted_csv_delimiters
+                )
+                dialect.doublequote = (
+                    True  # assume that all csvs have " as string escape
+                )
+            except Exception:
+                dialect = csv.reader(sample).dialect
+                if dialect.delimiter not in accepted_csv_delimiters:
+                    raise Exception(
+                        f"CSV delimeter '{dialect.delimiter}' is not supported"
+                    )
+
+        except csv.Error:
+            dialect = None
+        return dialect
+
+    def read(self, format, file_obj: BytesIO, **kwargs) -> pd.DataFrame:
+        func = {
+            'parquet': self.read_parquet,
+            'csv': self.read_csv,
+            'xlsx': self.read_excel,
+            'pdf': self.read_pdf,
+            'json': self.read_json,
+            'txt': self.read_txt,
+        }
+
+        if format not in func:
+            raise FileDetectError(f'Unsupported format: {format}')
+        func = func[format]
+
+        return func(file_obj, **kwargs)
+
+    def read_csv(self, file_obj: BytesIO, **kwargs):
+        file_obj = decode(file_obj)
+        dialect = self._get_csv_dialect(file_obj)
+
+        return pd.read_csv(file_obj, sep=dialect.delimiter, index_col=False)
+
+    def read_txt(self, file_obj: BytesIO, **kwargs):
+        file_obj = decode(file_obj)
+
+        try:
+            from langchain_core.documents import Document
+        except ImportError:
+            raise ImportError(
+                "To import TXT document please install 'langchain-community':\n"
+                "    pip install langchain-community"
+            )
+        text = file_obj.read()
+
+        file_name = None
+        if hasattr(file_obj, "name"):
+            file_name = file_obj.name
+        metadata = {"source": file_name}
+        documents = [Document(page_content=text, metadata=metadata)]
+
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=DEFAULT_CHUNK_SIZE, chunk_overlap=DEFAULT_CHUNK_OVERLAP
+        )
+
+        docs = text_splitter.split_documents(documents)
+        return pd.DataFrame(
+            [
+                {"content": doc.page_content, "metadata": doc.metadata}
+                for doc in docs
+            ]
+        )
+
+    def read_pdf(self, file_obj: BytesIO, **kwargs):
+        import fitz  # pymupdf
+
+        with fitz.open(stream=file_obj) as pdf:  # open pdf
+            text = chr(12).join([page.get_text() for page in pdf])
+
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=DEFAULT_CHUNK_SIZE, chunk_overlap=DEFAULT_CHUNK_OVERLAP
+        )
+
+        split_text = text_splitter.split_text(text)
+
+        return pd.DataFrame(
+            {"content": split_text, "metadata": [{}] * len(split_text)}
+        )
+
+    def read_json(self, file_obj: BytesIO, **kwargs):
+        file_obj = decode(file_obj)
+        file_obj.seek(0)
+        json_doc = json.loads(file_obj.read())
+        return pd.json_normalize(json_doc, max_level=0)
+
+    def read_parquet(self, file_obj: BytesIO, **kwargs):
+        return pd.read_parquet(file_obj)
+
+    def read_excel(self, file_obj: BytesIO, sheet_name=None, **kwargs) -> pd.DataFrame:
+
+        file_obj.seek(0)
+        with pd.ExcelFile(file_obj) as xls:
+            if sheet_name is None:
+                # No sheet specified: Return list of sheets
+                sheet_list = xls.sheet_names
+                return pd.DataFrame(sheet_list, columns=["Sheet_Name"])
+            else:
+                # Specific sheet requested: Load that sheet
+                return pd.read_excel(xls, sheet_name=sheet_name)

mindsdb/integrations/utilities/handlers/api_utilities/microsoft/ms_graph_api_utilities.py

@@ -131,7 +131,8 @@ class MSGraphAPIBaseClient:
         response = self._make_request(api_url, params)
 
         # If the response content is a binary file or a TSV file, return the raw content.
-        if response.headers["Content-Type"] in ("application/octet-stream", "text/tab-separated-values"):
+        if response.headers["Content-Type"] in ("application/octet-stream", "text/plain",
+                                                "text/tab-separated-values", "application/pdf"):
             return response.content
         # Otherwise, return the JSON content.
         else:

mindsdb/integrations/utilities/handlers/auth_utilities/microsoft/ms_graph_api_auth_utilities.py

@@ -43,9 +43,14 @@ class MSGraphAPIDelegatedPermissionsManager:
         # Set the redirect URI based on the request origin.
         # If the request origin is 127.0.0.1 (localhost), replace it with localhost.
         # This is done because the only HTTP origin allowed in Microsoft Entra ID app registration is localhost.
-        request_origin = request.headers.get('ORIGIN') or (request.scheme + '://' + request.host)
-        if not request_origin:
-            raise AuthException('Request origin could not be determined!')
+        try:
+            request_origin = request.headers.get('ORIGIN') or (request.scheme + '://' + request.host)
+            if not request_origin:
+                raise AuthException('Request origin could not be determined!')
+        except RuntimeError:
+            # if it is outside of request context (streaming in agent)
+            request_origin = ''
+
         request_origin = request_origin.replace('127.0.0.1', 'localhost') if 'http://127.0.0.1' in request_origin else request_origin
         self.redirect_uri = request_origin + '/verify-auth'
 

mindsdb/integrations/utilities/rag/chains/map_reduce_summarizer_chain.py

@@ -23,7 +23,7 @@ logger = log.getLogger(__name__)
 Summary = namedtuple('Summary', ['source_id', 'content'])
 
 
-def create_map_reduce_documents_chain(summarization_config: SummarizationConfig, input: str) -> MapReduceDocumentsChain:
+def create_map_reduce_documents_chain(summarization_config: SummarizationConfig, input: str) -> ReduceDocumentsChain:
     '''Creats a chain that map reduces documents into a single consolidated summary
 
     Args:
@@ -43,7 +43,7 @@ def create_map_reduce_documents_chain(summarization_config: SummarizationConfig,
     if 'input' in map_prompt.input_variables:
         map_prompt = map_prompt.partial(input=input)
     # Handles summarization of individual chunks.
-    map_chain = LLMChain(llm=summarization_llm, prompt=map_prompt)
+    # map_chain = LLMChain(llm=summarization_llm, prompt=map_prompt)
 
     reduce_prompt_template = summarization_config.reduce_prompt_template
     reduce_prompt = PromptTemplate.from_template(reduce_prompt_template)
@@ -60,18 +60,12 @@ def create_map_reduce_documents_chain(summarization_config: SummarizationConfig,
     )
 
     # Combines & iteratively reduces mapped documents.
-    reduce_documents_chain = ReduceDocumentsChain(
+    return ReduceDocumentsChain(
         combine_documents_chain=combine_documents_chain,
         collapse_documents_chain=combine_documents_chain,
         # Max number of tokens to group documents into.
         token_max=summarization_config.max_summarization_tokens
     )
-    return MapReduceDocumentsChain(
-        llm_chain=map_chain,
-        reduce_documents_chain=reduce_documents_chain,
-        document_variable_name='docs',
-        return_intermediate_steps=False
-    )
 
 
 class MapReduceSummarizerChain(Chain):
@@ -135,6 +129,8 @@ class MapReduceSummarizerChain(Chain):
         document_chunks = []
         for _, row in all_source_chunks.iterrows():
             metadata = row.get(self.metadata_column_name, {})
+            if row.get('chunk_id', None) is not None:
+                metadata['chunk_index'] = row.get('chunk_id', 0)
             document_chunks.append(Document(page_content=row[self.content_column_name], metadata=metadata))
         # Sort by chunk index if present in metadata so the full document is in its original order.
         document_chunks.sort(key=lambda doc: doc.metadata.get('chunk_index', 0) if doc.metadata else 0)

mindsdb/integrations/utilities/rag/pipelines/rag.py

@@ -298,10 +298,8 @@ class LangChainRAGPipeline:
             examples=retriever_config.examples,
             embeddings_model=embeddings,
             rewrite_prompt_template=retriever_config.rewrite_prompt_template,
-            retry_prompt_template=retriever_config.query_retry_template,
+            metadata_filters_prompt_template=retriever_config.metadata_filters_prompt_template,
             num_retries=retriever_config.num_retries,
-            sql_prompt_template=retriever_config.sql_prompt_template,
-            query_checker_template=retriever_config.query_checker_template,
             embeddings_table=knowledge_base_table._kb.vector_database_table,
             source_table=retriever_config.source_table,
             distance_function=distance_function,

mindsdb/integrations/utilities/rag/retrievers/sql_retriever.py

@@ -1,15 +1,20 @@
 import json
-from typing import List, Optional
+import re
+from pydantic import BaseModel, Field
+from typing import Any, List, Optional
 
 from langchain.chains.llm import LLMChain
 from langchain_core.callbacks.manager import CallbackManagerForRetrieverRun
 from langchain_core.documents.base import Document
 from langchain_core.embeddings import Embeddings
+from langchain_core.exceptions import OutputParserException
 from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.output_parsers import PydanticOutputParser
 from langchain_core.prompts import PromptTemplate
 from langchain_core.retrievers import BaseRetriever
 
 from mindsdb.api.executor.data_types.response_type import RESPONSE_TYPE
+from mindsdb.integrations.libs.response import HandlerResponse
 from mindsdb.integrations.libs.vectordatabase_handler import DistanceFunction, VectorStoreHandler
 from mindsdb.integrations.utilities.rag.settings import LLMExample, MetadataSchema, SearchKwargs
 from mindsdb.utilities import log
@@ -17,6 +22,18 @@ from mindsdb.utilities import log
 logger = log.getLogger(__name__)
 
 
+class MetadataFilter(BaseModel):
+    '''Represents an LLM generated metadata filter to apply to a PostgreSQL query.'''
+    attribute: str = Field(description="Database column to apply filter to")
+    comparator: str = Field(description="PostgreSQL comparator to use to filter database column")
+    value: Any = Field(description="Value to use to filter database column")
+
+
+class MetadataFilters(BaseModel):
+    '''List of LLM generated metadata filters to apply to a PostgreSQL query.'''
+    filters: List[MetadataFilter] = Field(description="List of PostgreSQL metadata filters to apply for user query")
+
+
 class SQLRetriever(BaseRetriever):
     '''Retriever that uses a LLM to generate pgvector queries to do similarity search with metadata filters.
 
@@ -25,10 +42,10 @@ class SQLRetriever(BaseRetriever):
     1. Use a LLM to rewrite the user input to something more suitable for retrieval. For example:
     "Show me documents containing how to finetune a LLM please" --> "how to finetune a LLM"
 
-    2. Use a LLM to generate a pgvector query with metadata filters based on the user input. Provided
-    metadata schemas & examples are used as additional context to generate the query.
+    2. Use a LLM to generate structured metadata filters based on the user input. Provided
+    metadata schemas & examples are used as additional context.
 
-    3. Use a LLM to double check the generated pgvector query is correct.
+    3. Generate a prepared PostgreSQL query from the structured metadata filters.
 
     4. Actually execute the query against our vector database to retrieve documents & return them.
     '''
@@ -37,23 +54,22 @@ class SQLRetriever(BaseRetriever):
     metadata_schemas: Optional[List[MetadataSchema]] = None
     examples: Optional[List[LLMExample]] = None
 
-    embeddings_model: Embeddings
     rewrite_prompt_template: str
-    retry_prompt_template: str
+    metadata_filters_prompt_template: str
+    embeddings_model: Embeddings
     num_retries: int
-    sql_prompt_template: str
-    query_checker_template: str
     embeddings_table: str
     source_table: str
+    source_id_column: str = 'Id'
     distance_function: DistanceFunction
     search_kwargs: SearchKwargs
 
     llm: BaseChatModel
 
-    def _prepare_sql_prompt(self) -> PromptTemplate:
+    def _prepare_metadata_prompt(self) -> PromptTemplate:
         base_prompt_template = PromptTemplate(
-            input_variables=['dialect', 'input', 'embeddings_table', 'source_table', 'embeddings', 'distance_function', 'schema', 'examples'],
-            template=self.sql_prompt_template
+            input_variables=['format_instructions', 'schema', 'examples', 'input', 'embeddings'],
+            template=self.metadata_filters_prompt_template
         )
         schema_prompt_str = ''
         if self.metadata_schemas is not None:
@@ -67,7 +83,7 @@ class SQLRetriever(BaseRetriever):
                     if column.values is not None:
                         column_mapping[column.name]['values'] = column.values
                 column_mapping_json_str = json.dumps(column_mapping, indent=4)
-                schema_str = f'''{i+2}. {schema.table} - {schema.description}
+                schema_str = f'''{i+1}. {schema.table} - {schema.description}
 
 Columns:
 ```json
@@ -86,7 +102,7 @@ Output:
 {example.output}
 
 '''
-            examples_prompt_str += example_str
+                examples_prompt_str += example_str
         return base_prompt_template.partial(
             schema=schema_prompt_str,
             examples=examples_prompt_str
@@ -100,97 +116,89 @@ Output:
         rewrite_chain = LLMChain(llm=self.llm, prompt=rewrite_prompt)
         return rewrite_chain.predict(input=query)
 
-    def _prepare_pgvector_query(self, query: str, run_manager: CallbackManagerForRetrieverRun) -> str:
-        # Incorporate metadata schemas & examples into prompt.
-        sql_prompt = self._prepare_sql_prompt()
-        sql_chain = LLMChain(llm=self.llm, prompt=sql_prompt)
-        # Generate the initial pgvector query.
-        sql_query = sql_chain.predict(
-            # Only pgvector & similarity search is supported.
-            dialect='postgres',
-            input=query,
-            embeddings_table=self.embeddings_table,
-            source_table=self.source_table,
-            distance_function=self.distance_function.value[0],
-            k=self.search_kwargs.k,
-            callbacks=run_manager.get_child() if run_manager else None
-        )
-        query_checker_prompt = PromptTemplate(
-            input_variables=['dialect', 'query'],
-            template=self.query_checker_template
-        )
-        query_checker_chain = LLMChain(llm=self.llm, prompt=query_checker_prompt)
-        # Check the query & return the final result to be executed.
-        return query_checker_chain.predict(
-            dialect='postgres',
-            query=sql_query
-        )
-
-    def _prepare_retry_query(self, query: str, error: str, run_manager: CallbackManagerForRetrieverRun) -> str:
-        sql_prompt = self._prepare_sql_prompt()
-        # Use provided schema as context for retrying failed queries.
-        schema = sql_prompt.partial_variables.get('schema', '')
-        retry_prompt = PromptTemplate(
-            input_variables=['query', 'dialect', 'error', 'embeddings_table', 'schema'],
-            template=self.retry_prompt_template
-        )
-        retry_chain = LLMChain(llm=self.llm, prompt=retry_prompt)
-        # Generate rewritten query.
-        sql_query = retry_chain.predict(
-            query=query,
-            dialect='postgres',
-            error=error,
-            embeddings_table=self.embeddings_table,
-            schema=schema,
-            callbacks=run_manager.get_child() if run_manager else None
-        )
-        query_checker_prompt = PromptTemplate(
-            input_variables=['dialect', 'query'],
-            template=self.query_checker_template
-        )
-        query_checker_chain = LLMChain(llm=self.llm, prompt=query_checker_prompt)
-        # Check the query & return the final result to be executed.
-        return query_checker_chain.predict(
-            dialect='postgres',
-            query=sql_query
+    def _prepare_pgvector_query(self, metadata_filters: List[MetadataFilter]) -> str:
+        # Base select JOINed with document source table.
+        base_query = f'''SELECT * FROM {self.embeddings_table} AS e INNER JOIN {self.source_table} AS s ON (e.metadata->>'original_row_id')::int = s."{self.source_id_column}" '''
+        col_to_schema = {}
+        if not self.metadata_schemas:
+            return ''
+        for schema in self.metadata_schemas:
+            for col in schema.columns:
+                col_to_schema[col.name] = schema
+        joined_schemas = set()
+        for filter in metadata_filters:
+            # Join schemas before filtering.
+            schema = col_to_schema.get(filter.attribute)
+            if schema is None or schema.table in joined_schemas or schema.table == self.source_table:
+                continue
+            joined_schemas.add(schema.table)
+            base_query += schema.join + ' '
+        # Actually construct WHERE conditions from metadata filters.
+        if metadata_filters:
+            base_query += 'WHERE '
+        for i, filter in enumerate(metadata_filters):
+            value = filter.value
+            if isinstance(value, str):
+                value = f"'{value}'"
+            base_query += f'"{filter.attribute}" {filter.comparator} {value}'
+            if i < len(metadata_filters) - 1:
+                base_query += ' AND '
+        base_query += f" ORDER BY e.embeddings {self.distance_function.value[0]} '{{embeddings}}' LIMIT {self.search_kwargs.k};"
+        return base_query
+
+    def _generate_metadata_filters(self, query: str) -> List[MetadataFilter]:
+        parser = PydanticOutputParser(pydantic_object=MetadataFilters)
+        metadata_prompt = self._prepare_metadata_prompt()
+        metadata_filters_chain = LLMChain(llm=self.llm, prompt=metadata_prompt)
+        metadata_filters_output = metadata_filters_chain.predict(
+            format_instructions=parser.get_format_instructions(),
+            input=query
         )
+        # If the LLM outputs raw JSON, use it as-is.
+        # If the LLM outputs anything including a json markdown section, use the last one.
+        json_markdown_output = re.findall(r'```json.*```', metadata_filters_output, re.DOTALL)
+        if json_markdown_output:
+            metadata_filters_output = json_markdown_output[-1]
+            # Clean the json tags.
+            metadata_filters_output = metadata_filters_output[7:]
+            metadata_filters_output = metadata_filters_output[:-3]
+        metadata_filters = parser.invoke(metadata_filters_output)
+        return metadata_filters.filters
+
+    def _prepare_and_execute_query(self, query: str, embeddings_str: str) -> HandlerResponse:
+        try:
+            metadata_filters = self._generate_metadata_filters(query)
+            checked_sql_query = self._prepare_pgvector_query(metadata_filters)
+            checked_sql_query_with_embeddings = checked_sql_query.format(embeddings=embeddings_str)
+            return self.vector_store_handler.native_query(checked_sql_query_with_embeddings)
+        except OutputParserException as e:
+            logger.warning(f'LLM failed to generate structured metadata filters: {str(e)}')
+            return HandlerResponse(RESPONSE_TYPE.ERROR, error_message=str(e))
+        except Exception as e:
+            logger.warning(f'Failed to prepare and execute SQL query from structured metadata: {str(e)}')
+            return HandlerResponse(RESPONSE_TYPE.ERROR, error_message=str(e))
 
     def _get_relevant_documents(
         self, query: str, *, run_manager: CallbackManagerForRetrieverRun
     ) -> List[Document]:
         # Rewrite query to be suitable for retrieval.
         retrieval_query = self._prepare_retrieval_query(query)
-
-        # Generate & check the query to be executed
-        checked_sql_query = self._prepare_pgvector_query(query, run_manager)
-
         # Embed the rewritten retrieval query & include it in the similarity search pgvector query.
         embedded_query = self.embeddings_model.embed_query(retrieval_query)
-        checked_sql_query_with_embeddings = checked_sql_query.format(embeddings=str(embedded_query))
-        # Handle LLM output that has the ```sql delimiter possibly.
-        checked_sql_query_with_embeddings = checked_sql_query_with_embeddings.replace('```sql', '')
-        checked_sql_query_with_embeddings = checked_sql_query_with_embeddings.replace('```', '')
         # Actually execute the similarity search with metadata filters.
-        document_response = self.vector_store_handler.native_query(checked_sql_query_with_embeddings)
+        document_response = self._prepare_and_execute_query(retrieval_query, str(embedded_query))
         num_retries = 0
         while num_retries < self.num_retries:
+            if document_response.resp_type != RESPONSE_TYPE.ERROR and len(document_response.data_frame) > 0:
+                # Successfully retrieved documents.
+                break
             if document_response.resp_type == RESPONSE_TYPE.ERROR:
-                error_msg = document_response.error_message
-                # LLMs won't always generate a working SQL query so we should have a fallback after retrying.
-                logger.info(f'SQL Retriever query {checked_sql_query} failed with error {error_msg}')
-                checked_sql_query = self._prepare_retry_query(checked_sql_query, error_msg, run_manager)
+                # LLMs won't always generate structured metadata so we should have a fallback after retrying.
+                logger.info(f'SQL Retriever query failed with error {document_response.error_message}')
             elif len(document_response.data_frame) == 0:
-                error_msg = "No documents retrieved from query."
-                checked_sql_query = self._prepare_retry_query(checked_sql_query, error_msg, run_manager)
-            else:
-                break
-
-            checked_sql_query_with_embeddings = checked_sql_query.format(embeddings=str(embedded_query))
-            # Handle LLM output that has the ```sql delimiter possibly.
-            checked_sql_query_with_embeddings = checked_sql_query_with_embeddings.replace('```sql', '')
-            checked_sql_query_with_embeddings = checked_sql_query_with_embeddings.replace('```', '')
-            document_response = self.vector_store_handler.native_query(checked_sql_query_with_embeddings)
+                logger.info('No documents retrieved from SQL Retriever query')
 
+            document_response = self._prepare_and_execute_query(retrieval_query, str(embedded_query))
             num_retries += 1
             if num_retries >= self.num_retries:
                 logger.info('Using fallback retriever in SQL retriever.')