MindsDB 25.6.3.1__py3-none-any.whl → 25.7.1.0__py3-none-any.whl

mindsdb/interfaces/agents/mcp_client_agent.py

@@ -71,11 +71,11 @@ class MCPLangchainAgent(LangchainAgent):
         self,
         agent: db.Agents,
         model: dict = None,
-        params: dict = None,
+        llm_params: dict = None,
         mcp_host: str = "127.0.0.1",
         mcp_port: int = 47337,
     ):
-        super().__init__(agent, model, params)
+        super().__init__(agent, model, llm_params)
         self.mcp_host = mcp_host
         self.mcp_port = mcp_port
         self.exit_stack = AsyncExitStack()
@@ -251,10 +251,10 @@ def create_mcp_agent(
         raise ValueError(f"Agent {agent_name} not found in project {project_name}")
 
     # Get merged parameters (defaults + agent params)
-    merged_params = agent_controller.get_agent_llm_params(agent_db.params)
+    llm_params = agent_controller.get_agent_llm_params(agent_db.params)
 
     # Create MCP agent with merged parameters
-    mcp_agent = MCPLangchainAgent(agent_db, params=merged_params, mcp_host=mcp_host, mcp_port=mcp_port)
+    mcp_agent = MCPLangchainAgent(agent_db, llm_params=llm_params, mcp_host=mcp_host, mcp_port=mcp_port)
 
     # Wrap for LiteLLM compatibility
     return LiteLLMAgentWrapper(mcp_agent)

mindsdb/interfaces/agents/mindsdb_database_agent.py

@@ -96,27 +96,7 @@ class MindsDBSQL(SQLDatabase):
             # Log the query for debugging
             logger.info(f"Executing SQL query: {command}")
 
-            # Removing backticks causes in query execution.
-            # remove backticks
-            # command = command.replace('`', '')
-
-            # Parse the SQL string to an AST object first
-            from mindsdb_sql_parser import parse_sql
-
-            ast_query = parse_sql(command)
-
-            # Now execute the parsed query
-            result = self._sql_agent.skill_tool.get_command_executor().execute_command(
-                ast_query, database_name="mindsdb"
-            )
-
-            # Convert ExecuteAnswer to a DataFrame for easier manipulation
-            if result.data is not None:
-                df = result.data.to_df()
-                return df.to_string(index=False)
-
-            else:
-                return "Query executed successfully, but returned no data."
+            return self._sql_agent.query(command)
 
         except Exception as e:
             logger.error(f"Error executing SQL command: {str(e)}\n{traceback.format_exc()}")
@@ -127,28 +107,6 @@ class MindsDBSQL(SQLDatabase):
                 return f"Error executing knowledge base query: {str(e)}. Please check that the knowledge base exists and your query syntax is correct."
             return f"Error: {str(e)}"
 
-    # def run_no_throw(self, command: str, fetch: str = "all") -> str:
-    #     """Execute a SQL command and return the result as a string.
-    #
-    #     This method catches any exceptions and returns an error message instead of raising an exception.
-    #
-    #     Args:
-    #         command: The SQL command to execute
-    #         fetch: Whether to fetch 'all' results or just 'one'
-    #
-    #     Returns:
-    #         A string representation of the result or an error message
-    #     """
-    #     command = extract_essential(command)
-    #     try:
-    #         return self._sql_agent.query_safe(command)
-    #     except Exception as e:
-    #         logger.error(f"Error executing SQL command: {str(e)}")
-    #         # If this is a knowledge base query, provide a more helpful error message
-    #         if "knowledge_base" in command.lower() or any(kb in command for kb in self._sql_agent.get_usable_knowledge_base_names()):
-    #             return f"Error executing knowledge base query: {str(e)}. Please check that the knowledge base exists and your query syntax is correct."
-    #         return f"Error: {str(e)}"
-
     def get_usable_knowledge_base_names(self) -> List[str]:
         """Get a list of usable knowledge base names.
 
@@ -160,3 +118,12 @@ class MindsDBSQL(SQLDatabase):
         except Exception as e:
             logger.error(f"Error getting usable knowledge base names: {str(e)}")
             return []
+
+    def check_knowledge_base_permission(self, name):
+        """Get a list of usable knowledge base names.
+
+        Returns:
+            A list of knowledge base names that can be used in queries
+        """
+
+        return self._sql_agent.check_knowledge_base_permission(name)

mindsdb/interfaces/data_catalog/data_catalog_reader.py

@@ -11,8 +11,6 @@ class DataCatalogReader(BaseDataCatalog):
         """
         Read the metadata from the data catalog and return it as a string.
         """
-        if not self.is_data_catalog_supported():
-            return f"Data catalog is not supported for database '{self.database_name}'."
         tables = self._read_metadata()
         if not tables:
             self.logger.warning(f"No metadata found for database '{self.database_name}'")
@@ -20,16 +18,37 @@ class DataCatalogReader(BaseDataCatalog):
 
         metadata_str = "Data Catalog: \n"
         if hasattr(self.data_handler, "meta_get_handler_info"):
-            metadata_str += self.data_handler.meta_get_handler_info() + "\n\n"
+            info = self.data_handler.meta_get_handler_info()
+            if info:
+                metadata_str += info + "\n\n"
 
         for table in tables:
             metadata_str += table.as_string() + "\n\n"
         return metadata_str
 
+    def read_metadata_as_records(self) -> list:
+        """
+        Read the metadata from the data catalog and return it as a list of database records.
+        """
+        tables = self._read_metadata()
+        if not tables:
+            self.logger.warning(f"No metadata found for database '{self.database_name}'")
+            return []
+        return tables
+
+    def get_handler_info(self) -> str:
+        """
+        Get the handler info for the database.
+        """
+        return self.data_handler.meta_get_handler_info()
+
     def _read_metadata(self) -> list:
         """
         Read the metadata from the data catalog and return it in a structured format.
         """
+        if not self.is_data_catalog_supported():
+            return f"Data catalog is not supported for database '{self.database_name}'."
+
         query = db.session.query(db.MetaTables).filter_by(integration_id=self.integration_id)
         if self.table_names:
             cleaned_table_names = [name.strip("`").split(".")[-1] for name in self.table_names]

mindsdb/interfaces/knowledge_base/controller.py

@@ -6,6 +6,7 @@ import decimal
 
 import pandas as pd
 import numpy as np
+from sqlalchemy.orm.attributes import flag_modified
 
 from mindsdb_sql_parser.ast import BinaryOperation, Constant, Identifier, Select, Update, Delete, Star
 from mindsdb_sql_parser.ast.mindsdb import CreatePredictor
@@ -33,6 +34,7 @@ from mindsdb.interfaces.variables.variables_controller import variables_controll
 from mindsdb.interfaces.knowledge_base.preprocessing.models import PreprocessingConfig, Document
 from mindsdb.interfaces.knowledge_base.preprocessing.document_preprocessor import PreprocessorFactory
 from mindsdb.interfaces.knowledge_base.evaluate import EvaluateBase
+from mindsdb.interfaces.knowledge_base.executor import KnowledgeBaseQueryExecutor
 from mindsdb.interfaces.model.functions import PredictorRecordNotFound
 from mindsdb.utilities.exception import EntityExistsError, EntityNotExistsError
 from mindsdb.integrations.utilities.sql_utils import FilterCondition, FilterOperator
@@ -46,25 +48,21 @@ from mindsdb.integrations.utilities.rag.rerankers.base_reranker import BaseLLMRe
 
 logger = log.getLogger(__name__)
 
-KB_TO_VECTORDB_COLUMNS = {"id": "original_doc_id", "chunk_id": "id", "chunk_content": "content"}
-
 
 def get_model_params(model_params: dict, default_config_key: str):
     """
     Get model parameters by combining default config with user provided parameters.
     """
-    # If the default config key is for reranking and the switch to use the default LLM is enabled,
-    # switch to the default LLM model.
-    if default_config_key == "default_reranking_model" and config.get("default_reranking_model").get(
-        "use_default_llm", False
-    ):
-        default_config_key = "default_llm_model"
-
     combined_model_params = copy.deepcopy(config.get(default_config_key, {}))
 
     if model_params:
+        if not isinstance(model_params, dict):
+            raise ValueError("Model parameters must be passed as a JSON object")
+
         combined_model_params.update(model_params)
 
+    combined_model_params.pop("use_default_llm", None)
+
     return combined_model_params
 
 
@@ -105,8 +103,6 @@ def get_reranking_model_from_params(reranking_model_params: dict):
         params_copy["api_key"] = get_api_key(provider, params_copy, strict=False)
     params_copy["model"] = params_copy.pop("model_name", None)
 
-    params_copy.pop("use_default_llm", None)
-
     return BaseLLMReranker(**params_copy)
 
 
@@ -144,23 +140,29 @@ class KnowledgeBaseTable:
         self.document_loader = None
         self.model_params = None
 
+        self.kb_to_vector_columns = {"id": "_original_doc_id", "chunk_id": "id", "chunk_content": "content"}
+        if self._kb.params.get("version", 0) < 2:
+            self.kb_to_vector_columns["id"] = "original_doc_id"
+
     def configure_preprocessing(self, config: Optional[dict] = None):
         """Configure preprocessing for the knowledge base table"""
         logger.debug(f"Configuring preprocessing with config: {config}")
         self.document_preprocessor = None  # Reset existing preprocessor
-        if config is not None:
-            # Ensure content_column is set for JSON chunking if not already specified
-            if config.get("type") == "json_chunking" and config.get("json_chunking_config"):
-                if "content_column" not in config["json_chunking_config"]:
-                    config["json_chunking_config"]["content_column"] = "content"
-
-            preprocessing_config = PreprocessingConfig(**config)
-            self.document_preprocessor = PreprocessorFactory.create_preprocessor(preprocessing_config)
-            logger.debug(f"Created preprocessor of type: {type(self.document_preprocessor)}")
-        else:
-            # Always create a default preprocessor if none specified
-            self.document_preprocessor = PreprocessorFactory.create_preprocessor()
-            logger.debug("Created default preprocessor")
+        if config is None:
+            config = {}
+
+        # Ensure content_column is set for JSON chunking if not already specified
+        if config.get("type") == "json_chunking" and config.get("json_chunking_config"):
+            if "content_column" not in config["json_chunking_config"]:
+                config["json_chunking_config"]["content_column"] = "content"
+
+        preprocessing_config = PreprocessingConfig(**config)
+        self.document_preprocessor = PreprocessorFactory.create_preprocessor(preprocessing_config)
+
+        # set doc_id column name
+        self.document_preprocessor.config.doc_id_column_name = self.kb_to_vector_columns["id"]
+
+        logger.debug(f"Created preprocessor of type: {type(self.document_preprocessor)}")
 
     def select_query(self, query: Select) -> pd.DataFrame:
         """
@@ -169,6 +171,30 @@ class KnowledgeBaseTable:
         :param query: query to KB table
         :return: dataframe with the result table
         """
+
+        # Copy query for complex execution via DuckDB: DISTINCT, GROUP BY etc.
+        query_copy = copy.deepcopy(query)
+
+        executor = KnowledgeBaseQueryExecutor(self)
+        df = executor.run(query)
+
+        if (
+            query.group_by is not None
+            or query.order_by is not None
+            or query.having is not None
+            or query.distinct is True
+            or len(query.targets) != 1
+            or not isinstance(query.targets[0], Star)
+        ):
+            query_copy.where = None
+            if "metadata" in df.columns:
+                df["metadata"] = df["metadata"].apply(to_json)
+
+            df = query_df(df, query_copy, session=self.session)
+
+        return df
+
+    def select(self, query, disable_reranking=False):
         logger.debug(f"Processing select query: {query}")
 
         # Extract the content query text for potential reranking
@@ -180,9 +206,6 @@ class KnowledgeBaseTable:
         query.from_table = Identifier(parts=[self._kb.vector_database_table])
         logger.debug(f"Set table name to: {self._kb.vector_database_table}")
 
-        # Copy query for complex execution via DuckDB: DISTINCT, GROUP BY etc.
-        query_copy = copy.deepcopy(query)
-
         query.targets = [
             Identifier(TableField.ID.value),
             Identifier(TableField.CONTENT.value),
@@ -197,7 +220,6 @@ class KnowledgeBaseTable:
         conditions = []
         query_text = None
         relevance_threshold = None
-        reranking_enabled_flag = True
         query_conditions = db_handler.extract_conditions(query.where)
         if query_conditions is not None:
             for item in query_conditions:
@@ -213,10 +235,9 @@ class KnowledgeBaseTable:
                         logger.error(error_msg)
                         raise ValueError(error_msg)
                 elif item.column == "reranking":
-                    reranking_enabled_flag = item.value
-                    # cast to boolean
-                    if isinstance(reranking_enabled_flag, str):
-                        reranking_enabled_flag = reranking_enabled_flag.lower() not in ("false")
+                    if item.value is False or (isinstance(item.value, str) and item.value.lower() == "false"):
+                        disable_reranking = True
+
                 elif item.column == "relevance" and item.op.value != FilterOperator.GREATER_THAN_OR_EQUAL.value:
                     raise ValueError(
                         f"Invalid operator for relevance: {item.op.value}. Only GREATER_THAN_OR_EQUAL is allowed."
@@ -248,66 +269,59 @@ class KnowledgeBaseTable:
                 limit = 100
             query.limit = Constant(limit)
 
-        df = db_handler.dispatch_select(query, conditions)
+        allowed_metadata_columns = self._get_allowed_metadata_columns()
+        df = db_handler.dispatch_select(query, conditions, allowed_metadata_columns=allowed_metadata_columns)
         df = self.addapt_result_columns(df)
 
         logger.debug(f"Query returned {len(df)} rows")
         logger.debug(f"Columns in response: {df.columns.tolist()}")
         # Check if we have a rerank_model configured in KB params
-        df = self.add_relevance(df, query_text, relevance_threshold, reranking_enabled_flag)
+        df = self.add_relevance(df, query_text, relevance_threshold, disable_reranking)
 
-        if (
-            query.group_by is not None
-            or query.order_by is not None
-            or query.having is not None
-            or query.distinct is True
-            or len(query.targets) != 1
-            or not isinstance(query.targets[0], Star)
-        ):
-            query_copy.where = None
-            if "metadata" in df.columns:
-                df["metadata"] = df["metadata"].apply(to_json)
+        return df
 
-            df = query_df(df, query_copy, session=self.session)
+    def _get_allowed_metadata_columns(self) -> List[str] | None:
+        # Return list of KB columns to restrict querying, if None: no restrictions
 
-        return df
+        if self._kb.params.get("version", 0) < 2:
+            # disable for old version KBs
+            return None
+
+        user_columns = self._kb.params.get("metadata_columns", [])
+        dynamic_columns = self._kb.params.get("inserted_metadata", [])
+
+        columns = set(user_columns) | set(dynamic_columns)
+        return [col.lower() for col in columns]
 
     def score_documents(self, query_text, documents, reranking_model_params):
         reranker = get_reranking_model_from_params(reranking_model_params)
         return reranker.get_scores(query_text, documents)
 
-    def add_relevance(self, df, query_text, relevance_threshold=None, reranking_enabled_flag=True):
+    def add_relevance(self, df, query_text, relevance_threshold=None, disable_reranking=False):
         relevance_column = TableField.RELEVANCE.value
 
         reranking_model_params = get_model_params(self._kb.params.get("reranking_model"), "default_reranking_model")
-        if reranking_model_params and query_text and len(df) > 0 and reranking_enabled_flag:
+        if reranking_model_params and query_text and len(df) > 0 and not disable_reranking:
             # Use reranker for relevance score
-            try:
-                logger.info(f"Using knowledge reranking model from params: {reranking_model_params}")
-                # Apply custom filtering threshold if provided
-                if relevance_threshold is not None:
-                    reranking_model_params["filtering_threshold"] = relevance_threshold
-                    logger.info(f"Using custom filtering threshold: {relevance_threshold}")
-
-                reranker = get_reranking_model_from_params(reranking_model_params)
-                # Get documents to rerank
-                documents = df["chunk_content"].tolist()
-                # Use the get_scores method with disable_events=True
-                scores = reranker.get_scores(query_text, documents)
-                # Add scores as the relevance column
-                df[relevance_column] = scores
-
-                # Filter by threshold
-                scores_array = np.array(scores)
-                df = df[scores_array > reranker.filtering_threshold]
-                logger.debug(f"Applied reranking with params: {reranking_model_params}")
-            except Exception as e:
-                logger.error(f"Error during reranking: {str(e)}")
-                # Fallback to distance-based relevance
-                if "distance" in df.columns:
-                    df[relevance_column] = 1 / (1 + df["distance"])
-                else:
-                    logger.info("No distance or reranker available")
+
+            logger.info(f"Using knowledge reranking model from params: {reranking_model_params}")
+            # Apply custom filtering threshold if provided
+            if relevance_threshold is not None:
+                reranking_model_params["filtering_threshold"] = relevance_threshold
+                logger.info(f"Using custom filtering threshold: {relevance_threshold}")
+
+            reranker = get_reranking_model_from_params(reranking_model_params)
+            # Get documents to rerank
+            documents = df["chunk_content"].tolist()
+            # Use the get_scores method with disable_events=True
+            scores = reranker.get_scores(query_text, documents)
+            # Add scores as the relevance column
+            df[relevance_column] = scores
+
+            # Filter by threshold
+            scores_array = np.array(scores)
+            df = df[scores_array > reranker.filtering_threshold]
+            logger.debug(f"Applied reranking with params: {reranking_model_params}")
 
         elif "distance" in df.columns:
             # Calculate relevance from distance
@@ -327,12 +341,12 @@ class KnowledgeBaseTable:
         if conditions is None:
             return
         for condition in conditions:
-            if condition.column in KB_TO_VECTORDB_COLUMNS:
-                condition.column = KB_TO_VECTORDB_COLUMNS[condition.column]
+            if condition.column in self.kb_to_vector_columns:
+                condition.column = self.kb_to_vector_columns[condition.column]
 
     def addapt_result_columns(self, df):
         col_update = {}
-        for kb_col, vec_col in KB_TO_VECTORDB_COLUMNS.items():
+        for kb_col, vec_col in self.kb_to_vector_columns.items():
             if vec_col in df.columns:
                 col_update[vec_col] = kb_col
 
@@ -341,7 +355,7 @@ class KnowledgeBaseTable:
         columns = list(df.columns)
         # update id, get from metadata
         df[TableField.ID.value] = df[TableField.METADATA.value].apply(
-            lambda m: None if m is None else m.get("original_doc_id")
+            lambda m: None if m is None else m.get(self.kb_to_vector_columns["id"])
         )
 
         # id on first place
@@ -528,8 +542,8 @@ class KnowledgeBaseTable:
 
                     metadata = {
                         **base_metadata,
-                        "original_row_index": str(idx),  # provide link to original row index
-                        "content_column": col,
+                        "_original_row_index": str(idx),  # provide link to original row index
+                        "_content_column": col,
                     }
 
                     raw_documents.append(Document(content=content_str, id=doc_id, metadata=metadata))
@@ -624,16 +638,22 @@ class KnowledgeBaseTable:
             metadata_columns = [column_map.get(col.lower(), col) for col in metadata_columns]
             logger.debug(f"Mapped metadata columns: {metadata_columns}")
 
-        if content_columns is not None:
-            content_columns = list(set(content_columns).intersection(columns))
-            if len(content_columns) == 0:
-                raise ValueError(f"Content columns {params.get('content_columns')} not found in dataset: {columns}")
+        content_columns = list(set(content_columns).intersection(columns))
+        if len(content_columns) == 0:
+            raise ValueError(f"Content columns {params.get('content_columns')} not found in dataset: {columns}")
 
-            if metadata_columns is not None:
-                metadata_columns = list(set(metadata_columns).intersection(columns))
-            else:
-                # all the rest columns
-                metadata_columns = list(set(columns).difference(content_columns))
+        if metadata_columns is not None:
+            metadata_columns = list(set(metadata_columns).intersection(columns))
+        else:
+            # all the rest columns
+            metadata_columns = list(set(columns).difference(content_columns))
+
+            # update list of used columns
+            inserted_metadata = set(self._kb.params.get("inserted_metadata", []))
+            inserted_metadata.update(metadata_columns)
+            self._kb.params["inserted_metadata"] = list(inserted_metadata)
+            flag_modified(self._kb, "params")
+            db.session.commit()
 
         # Add content columns directly (don't combine them)
         for col in content_columns:
@@ -659,7 +679,7 @@ class KnowledgeBaseTable:
                     elif isinstance(value, dict):
                         metadata.update(value)
                         continue
-                    else:
+                    elif value is not None:
                         value = str(value)
                     metadata[col] = value
                 return metadata
@@ -766,15 +786,10 @@ class KnowledgeBaseTable:
         llm_model = args.pop("model_name")
         engine = args.pop("provider")
 
-        llm_model = f"{engine}/{llm_model}"
-
-        if "base_url" in args:
-            args["api_base"] = args.pop("base_url")
-
         module = session.integration_controller.get_handler_module("litellm")
         if module is None or module.Handler is None:
             raise ValueError(f'Unable to use "{engine}" provider. Litellm handler is not installed')
-        return module.Handler.embeddings(llm_model, messages, args)
+        return module.Handler.embeddings(engine, llm_model, messages, args)
 
     def build_rag_pipeline(self, retrieval_config: dict):
         """
@@ -896,6 +911,8 @@ class KnowledgeBaseController:
     manages knowledge bases
     """
 
+    KB_VERSION = 2
+
     def __init__(self, session) -> None:
         self.session = session
 
@@ -907,6 +924,7 @@ class KnowledgeBaseController:
         params: dict,
         preprocessing_config: Optional[dict] = None,
         if_not_exists: bool = False,
+        keyword_search_enabled: bool = False,
         # embedding_model: Identifier = None, # Legacy: Allow MindsDB models to be passed as embedding_model.
     ) -> db.KnowledgeBase:
         """
@@ -961,10 +979,7 @@ class KnowledgeBaseController:
         #         # it is params for model
         #         embedding_params.update(params["embedding_model"])
 
-        if "embedding_model" in params:
-            if not isinstance(params["embedding_model"], dict):
-                raise ValueError("embedding_model should be JSON object with model parameters.")
-            embedding_params.update(params["embedding_model"])
+        embedding_params = get_model_params(params.get("embedding_model", {}), "default_embedding_model")
 
         # if model_name is None:  # Legacy
         model_name = self._create_embedding_model(
@@ -1023,7 +1038,10 @@ class KnowledgeBaseController:
             vector_db_name, vector_table_name = storage.parts
 
         # create table in vectordb before creating KB
-        self.session.datahub.get(vector_db_name).integration_handler.create_table(vector_table_name)
+        vector_store_handler = self.session.datahub.get(vector_db_name).integration_handler
+        vector_store_handler.create_table(vector_table_name)
+        if keyword_search_enabled:
+            vector_store_handler.add_full_text_index(vector_table_name, TableField.CONTENT.value)
         vector_database_id = self.session.integration_controller.get(vector_db_name)["id"]
 
         # Store sparse vector settings in params if specified
@@ -1033,6 +1051,7 @@ class KnowledgeBaseController:
             if vector_size is not None:
                 params["vector_config"]["vector_size"] = vector_size
 
+        params["version"] = self.KB_VERSION
         kb = db.KnowledgeBase(
             name=name,
             project_id=project_id,

mindsdb/interfaces/knowledge_base/evaluate.py

@@ -7,7 +7,7 @@ import pandas as pd
 import datetime as dt
 
 from mindsdb.api.executor.sql_query.result_set import ResultSet
-from mindsdb_sql_parser import Identifier, Select, Constant, Star, parse_sql
+from mindsdb_sql_parser import Identifier, Select, Constant, Star, parse_sql, BinaryOperation
 from mindsdb.utilities import log
 
 from mindsdb.interfaces.knowledge_base.llm_client import LLMClient
@@ -130,6 +130,8 @@ class EvaluateBase:
         integration_name = table_name.parts[0]
         table_name = Identifier(parts=table_name.parts[1:])
         dn = self.session.datahub.get(integration_name)
+        if dn is None:
+            raise ValueError(f"Can't find database: {integration_name}")
         return dn, table_name
 
     def save_to_table(self, table_name: Identifier, df: pd.DataFrame, is_replace=False):
@@ -168,13 +170,13 @@ class EvaluateBase:
             test_data = self.generate_test_data(gen_params)
 
             self.save_to_table(test_table, test_data, is_replace=True)
-        else:
-            test_data = self.read_from_table(test_table)
 
         if params.get("evaluate", True) is False:
             # no evaluate is required
             return pd.DataFrame()
 
+        test_data = self.read_from_table(test_table)
+
         scores = self.evaluate(test_data)
         scores["name"] = self.name
         scores["created_at"] = dt.datetime.now()
@@ -256,7 +258,13 @@ class EvaluateRerank(EvaluateBase):
 
             start_time = time.time()
             logger.debug(f"Querying [{i + 1}/{len(questions)}]: {question}")
-            df_answers = self.kb.select_query(Select(targets=[Identifier("chunk_content")], limit=Constant(self.TOP_K)))
+            df_answers = self.kb.select_query(
+                Select(
+                    targets=[Identifier("chunk_content")],
+                    where=BinaryOperation(op="=", args=[Identifier("content"), Constant(question)]),
+                    limit=Constant(self.TOP_K),
+                )
+            )
             query_time = time.time() - start_time
 
             proposed_responses = list(df_answers["chunk_content"])
@@ -410,7 +418,7 @@ class EvaluateDocID(EvaluateBase):
     Checks if ID in response from KB is matched with doc ID in test dataset
     """
 
-    TOP_K = 100
+    TOP_K = 20
 
     def generate(self, sampled_df: pd.DataFrame) -> pd.DataFrame:
         if "id" not in sampled_df.columns:
@@ -462,7 +470,11 @@ class EvaluateDocID(EvaluateBase):
             start_time = time.time()
             logger.debug(f"Querying [{i + 1}/{len(questions)}]: {question}")
             df_answers = self.kb.select_query(
-                Select(targets=[Identifier("chunk_content"), Identifier("id")], limit=Constant(self.TOP_K))
+                Select(
+                    targets=[Identifier("chunk_content"), Identifier("id")],
+                    where=BinaryOperation(op="=", args=[Identifier("content"), Constant(question)]),
+                    limit=Constant(self.TOP_K),
+                )
             )
             query_time = time.time() - start_time
 
@@ -511,6 +523,6 @@ class EvaluateDocID(EvaluateBase):
             "total": total_questions,
             "total_found": total_found,
             "retrieved_in_top_10": accurate_in_top_10,
-            "cumulative_recall": cumulative_recall,
+            "cumulative_recall": json.dumps(cumulative_recall),
             "avg_query_time": avg_query_time,
         }