MindsDB 25.1.2.1__py3-none-any.whl → 25.1.5.0__py3-none-any.whl

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,83 +116,93 @@ 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.
-        return 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
+    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 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}')
+        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:
+                # 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:
+                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.')
-                return self.fallback_retriever._get_relevant_documents(retrieval_query, run_manager)
-            query_to_retry = self._prepare_retry_query(checked_sql_query, error_msg, run_manager)
-            query_to_retry_with_embeddings = query_to_retry.format(embeddings=str(embedded_query))
-            # Handle LLM output that has the ```sql delimiter possibly.
-            query_to_retry_with_embeddings = query_to_retry_with_embeddings.replace('```sql', '')
-            query_to_retry_with_embeddings = query_to_retry_with_embeddings.replace('```', '')
-            document_response = self.vector_store_handler.native_query(query_to_retry_with_embeddings)
-            num_retries += 1
+                return self.fallback_retriever._get_relevant_documents(retrieval_query, run_manager=run_manager)
 
         document_df = document_response.data_frame
         retrieved_documents = []
@@ -185,4 +211,8 @@ Output:
                 document_row.get('content', ''),
                 metadata=document_row.get('metadata', {})
             ))
-        return retrieved_documents
+        if retrieved_documents:
+            return retrieved_documents
+        # If the SQL query constructed did not return any documents, fallback.
+        logger.info('No documents returned from SQL retriever. using fallback retriever.')
+        return self.fallback_retriever._get_relevant_documents(retrieval_query, run_manager=run_manager)

mindsdb/integrations/utilities/rag/settings.py

@@ -3,7 +3,6 @@ from typing import List, Union, Any, Optional, Dict
 
 from langchain_community.vectorstores.chroma import Chroma
 from langchain_community.vectorstores.pgvector import PGVector
-from langchain_community.tools.sql_database.prompt import QUERY_CHECKER as DEFAULT_QUERY_CHECKER_PROMPT_TEMPLATE
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
 from langchain_core.language_models import BaseChatModel
@@ -94,6 +93,25 @@ Output only a single better search query and nothing else like in the example.
 Here is the user input: {input}
 '''
 
+DEFAULT_METADATA_FILTERS_PROMPT_TEMPLATE = '''Construct a list of PostgreSQL metadata filters to filter documents in the database based on the user input.
+
+<< INSTRUCTIONS >>
+{format_instructions}
+
+RETURN ONLY THE FINAL JSON. DO NOT EXPLAIN, JUST RETURN THE FINAL JSON.
+
+<< TABLES YOU HAVE ACCESS TO >>
+
+{schema}
+
+<< EXAMPLES >>
+
+{examples}
+
+Here is the user input:
+{input}
+'''
+
 DEFAULT_SQL_PROMPT_TEMPLATE = '''
 Construct a valid {dialect} SQL query to select documents relevant to the user input.
 Source documents are found in the {source_table} table. You may need to join with other tables to get additional document metadata.
@@ -136,7 +154,6 @@ Columns:
         "description": "Metadata for the document chunk. Always select metadata and always join with the {source_table} table on the string metadata field 'original_row_id'"
     }}
 }}
-```
 
 {schema}
 
@@ -290,6 +307,8 @@ class VectorStoreConfig(BaseModel):
     collection_name: str = DEFAULT_COLLECTION_NAME
     connection_string: str = None
     kb_table: Any = None
+    is_sparse: bool = False
+    vector_size: Optional[int] = None
 
     class Config:
         arbitrary_types_allowed = True
@@ -376,6 +395,13 @@ class MetadataSchema(BaseModel):
     columns: List[ColumnSchema] = Field(
         description="List of column schemas describing the metadata columns available for the table"
     )
+    join: str = Field(
+        description="SQL join string to join this table with source documents table",
+        default=''
+    )
+
+    class Config:
+        frozen = True
 
 
 class LLMExample(BaseModel):
@@ -392,19 +418,9 @@ class SQLRetrieverConfig(BaseModel):
         default_factory=LLMConfig,
         description="LLM configuration to use for generating the final SQL query for retrieval"
     )
-    sql_prompt_template: str = Field(
-        default=DEFAULT_SQL_PROMPT_TEMPLATE,
-        description="""Prompt template to generate the SQL query to execute against the vector database. Currently only pgvector is supported.
-        Has 'dialect', 'input', 'embeddings_table', 'source_table', 'embeddings', 'distance_function', 'schema', and 'examples' input variables.
-        """
-    )
-    query_checker_template: str = Field(
-        default=DEFAULT_QUERY_CHECKER_PROMPT_TEMPLATE,
-        description="Prompt template to use for double checking SQL queries before execution. Has 'query' and 'dialect' input variables."
-    )
-    query_retry_template: str = Field(
-        default=DEFAULT_QUERY_RETRY_PROMPT_TEMPLATE,
-        description="Prompt template to rewrite SQL query that failed. Has 'dialect', 'query', and 'error' input variables."
+    metadata_filters_prompt_template: str = Field(
+        default=DEFAULT_METADATA_FILTERS_PROMPT_TEMPLATE,
+        description="Prompt template to generate PostgreSQL metadata filters. Has 'format_instructions', 'schema', 'examples', and 'input' input variables"
     )
     num_retries: int = Field(
         default=DEFAULT_NUM_QUERY_RETRIES,
@@ -449,8 +465,13 @@ class SummarizationConfig(BaseModel):
 class RerankerConfig(BaseModel):
     model: str = DEFAULT_RERANKING_MODEL
     base_url: str = DEFAULT_LLM_ENDPOINT
-    filtering_threshold: float = 0.99
+    filtering_threshold: float = 0.5
     num_docs_to_keep: Optional[int] = None
+    max_concurrent_requests: int = 20
+    max_retries: int = 3
+    retry_delay: float = 1.0
+    early_stop: bool = True  # Whether to enable early stopping
+    early_stop_threshold: float = 0.8  # Confidence threshold for early stopping
 
 
 class MultiHopRetrieverConfig(BaseModel):

mindsdb/integrations/utilities/sql_utils.py

@@ -178,7 +178,7 @@ def project_dataframe(df, targets, table_columns):
 
     # adapt column names to projection
     if len(df_col_rename) > 0:
-        df = df.rename(columns=df_col_rename)
+        df.rename(columns=df_col_rename, inplace=True)
     return df
 
 

mindsdb/interfaces/agents/agents_controller.py

@@ -1,5 +1,5 @@
 import datetime
-from typing import Dict, Iterator, List, Union, Tuple
+from typing import Dict, Iterator, List, Union, Tuple, Optional
 
 from langchain_core.tools import BaseTool
 from sqlalchemy.orm.attributes import flag_modified
@@ -70,7 +70,7 @@ class AgentsController:
 
         return model, provider
 
-    def get_agent(self, agent_name: str, project_name: str = 'mindsdb') -> db.Agents:
+    def get_agent(self, agent_name: str, project_name: str = 'mindsdb') -> Optional[db.Agents]:
         '''
         Gets an agent by name.
 
@@ -79,7 +79,7 @@ class AgentsController:
             project_name (str): The name of the containing project - must exist
 
         Returns:
-            agent (db.Agents): The database agent object
+            agent (Optional[db.Agents]): The database agent object
         '''
 
         project = self.project_controller.get(name=project_name)
@@ -252,6 +252,16 @@ class AgentsController:
         existing_agent = self.get_agent(agent_name, project_name=project_name)
         if existing_agent is None:
             raise EntityNotExistsError(f'Agent with name not found: {agent_name}')
+        is_demo = (existing_agent.params or {}).get('is_demo', False)
+        if (
+            is_demo and (
+                (name is not None and name != agent_name)
+                or (model_name or provider)
+                or (len(skills_to_add) > 0 or len(skills_to_remove) > 0 or len(skills_to_rewrite) > 0)
+                or (isinstance(params, dict) and len(params) > 1 and 'prompt_template' not in params)
+            )
+        ):
+            raise ValueError("It is forbidden to change properties of the demo object")
 
         if name is not None and name != agent_name:
             # Check to see if updated name already exists
@@ -352,6 +362,8 @@ class AgentsController:
         agent = self.get_agent(agent_name, project_name)
         if agent is None:
             raise ValueError(f'Agent with name does not exist: {agent_name}')
+        if isinstance(agent.params, dict) and agent.params.get('is_demo') is True:
+            raise ValueError('Unable to delete demo object')
         agent.deleted_at = datetime.datetime.now()
         db.session.commit()
 
@@ -362,24 +374,22 @@ class AgentsController:
             project_name: str = 'mindsdb',
             tools: List[BaseTool] = None,
             stream: bool = False) -> Union[Iterator[object], pd.DataFrame]:
-        '''
+        """
         Queries an agent to get a completion.
 
         Parameters:
             agent (db.Agents): Existing agent to get completion from
             messages (List[Dict[str, str]]): Chat history to send to the agent
-            trace_id (str): ID of Langfuse trace to use
-            observation_id (str): ID of parent Langfuse observation to use
             project_name (str): Project the agent belongs to (default mindsdb)
             tools (List[BaseTool]): Tools to use while getting the completion
-            stream (bool): Whether or not to stream the response
+            stream (bool): Whether to stream the response
 
         Returns:
             response (Union[Iterator[object], pd.DataFrame]): Completion as a DataFrame or iterator of completion chunks
 
         Raises:
             ValueError: Agent's model does not exist.
-        '''
+        """
         if stream:
             return self._get_completion_stream(
                 agent,

mindsdb/interfaces/agents/constants.py

@@ -165,6 +165,7 @@ PROVIDER_TO_MODELS = MappingProxyType(
 
 ASSISTANT_COLUMN = "answer"
 CONTEXT_COLUMN = "context"
+TRACE_ID_COLUMN = "trace_id"
 DEFAULT_AGENT_TIMEOUT_SECONDS = 300
 # These should require no additional arguments.
 DEFAULT_AGENT_TOOLS = []