MindsDB 25.7.2.0__py3-none-any.whl → 25.7.4.0__py3-none-any.whl

mindsdb/interfaces/knowledge_base/executor.py

@@ -217,6 +217,17 @@ class KnowledgeBaseQueryExecutor:
                 f'Operator "{content_condition.op}" is not supported for condition: {content_condition}'
             )
 
+    @staticmethod
+    def to_include_content(content_condition: BinaryOperation) -> List[str]:
+        """
+        Handles positive conditions for content. Returns list of content values
+        """
+        if content_condition.op == "IN":
+            return [item.value for item in content_condition.args[1].items]
+
+        elif content_condition.op in ("=", "LIKE"):
+            return [content_condition.args[1].value]
+
     def to_excluded_ids(
         self, content_condition: BinaryOperation, other_conditions: List[BinaryOperation]
     ) -> Optional[List[str]]:
@@ -290,11 +301,17 @@ class KnowledgeBaseQueryExecutor:
             if len(content_filters) > 0:
                 content_filters2 = []
                 exclude_ids = set()
+                include_contents = set()
                 # exclude content conditions
                 for condition in content_filters:
                     ids = self.to_excluded_ids(condition, other_filters)
                     if ids is not None:
                         exclude_ids.update(ids)
+                        continue
+                    contents = self.to_include_content(condition)
+                    if contents is not None:
+                        include_contents.update(contents)
+                        continue
                     else:
                         # keep origin content filter
                         content_filters2.append(condition)
@@ -305,6 +322,13 @@ class KnowledgeBaseQueryExecutor:
                     condition = BinaryOperation(op="NOT IN", args=[Identifier(self.id_column), Tuple(values)])
                     other_filters.append(condition)
                 # execute content filters
+                if include_contents:
+                    content = " AND ".join(include_contents)
+                    result = self.execute_content_condition(
+                        BinaryOperation(op="=", args=[Identifier(self.content_column), Constant(content)]),
+                        other_filters,
+                    )
+                    results.append(result)
                 for condition in content_filters2:
                     result = self.execute_content_condition(condition, other_filters)
                     results.append(result)

mindsdb/interfaces/knowledge_base/llm_client.py

@@ -54,12 +54,12 @@ class LLMClient:
 
             self.client = module.Handler
 
-    def completion(self, messages: List[dict]) -> str:
+    def completion(self, messages: List[dict], json_output: bool = False) -> str:
         """
         Call LLM completion and get response
         """
         params = self.params
-
+        params["json_output"] = json_output
         if self.provider in ("azure_openai", "openai"):
             response = self.client.chat.completions.create(
                 model=params["model_name"],
@@ -69,6 +69,6 @@ class LLMClient:
         else:
             kwargs = params.copy()
             model = kwargs.pop("model_name")
-
+            kwargs.pop("provider", None)
             response = self.client.completion(self.provider, model=model, messages=messages, args=kwargs)
             return response.choices[0].message.content

mindsdb/interfaces/knowledge_base/preprocessing/document_preprocessor.py

@@ -1,16 +1,17 @@
+import re
+import html
+import asyncio
 from typing import List, Dict, Optional, Any
+
 import pandas as pd
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-import asyncio
-
+from langchain_core.documents import Document as LangchainDocument
 
 from mindsdb.integrations.utilities.rag.splitters.file_splitter import (
     FileSplitter,
     FileSplitterConfig,
 )
-
 from mindsdb.interfaces.agents.langchain_agent import create_chat_model
-
 from mindsdb.interfaces.knowledge_base.preprocessing.models import (
     PreprocessingConfig,
     ProcessedChunk,
@@ -21,7 +22,6 @@ from mindsdb.interfaces.knowledge_base.preprocessing.models import (
 )
 from mindsdb.utilities import log
 
-from langchain_core.documents import Document as LangchainDocument
 
 logger = log.getLogger(__name__)
 
@@ -123,11 +123,11 @@ class ContextualPreprocessor(DocumentPreprocessor):
 
     DEFAULT_CONTEXT_TEMPLATE = """
 <document>
-{{WHOLE_DOCUMENT}}
+{WHOLE_DOCUMENT}
 </document>
 Here is the chunk we want to situate within the whole document
 <chunk>
-{{CHUNK_CONTENT}}
+{CHUNK_CONTENT}
 </chunk>
 Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk. Answer only with the succinct context and nothing else."""
 
@@ -149,12 +149,20 @@ Please give a short succinct context to situate this chunk within the overall do
         self.summarize = self.config.summarize
 
     def _prepare_prompts(self, chunk_contents: list[str], full_documents: list[str]) -> list[str]:
-        prompts = [
-            self.context_template.replace("{{WHOLE_DOCUMENT}}", full_document) for full_document in full_documents
-        ]
-        prompts = [
-            prompt.replace("{{CHUNK_CONTENT}}", chunk_content) for prompt, chunk_content in zip(prompts, chunk_contents)
-        ]
+        def tag_replacer(match):
+            tag = match.group(0)
+            if tag.lower() not in ["<document>", "</document>", "<chunk>", "</chunk>"]:
+                return tag
+            return html.escape(tag)
+
+        tag_pattern = r"</?document>|</?chunk>"
+        prompts = []
+        for chunk_content, full_document in zip(chunk_contents, full_documents):
+            chunk_content = re.sub(tag_pattern, tag_replacer, chunk_content, flags=re.IGNORECASE)
+            full_document = re.sub(tag_pattern, tag_replacer, full_document, flags=re.IGNORECASE)
+            prompts.append(
+                self.DEFAULT_CONTEXT_TEMPLATE.format(WHOLE_DOCUMENT=full_document, CHUNK_CONTENT=chunk_content)
+            )
 
         return prompts
 

mindsdb/interfaces/query_context/context_controller.py

@@ -7,9 +7,7 @@ import pandas as pd
 
 from mindsdb_sql_parser import Select, Star, OrderBy
 
-from mindsdb_sql_parser.ast import (
-    Identifier, BinaryOperation, Last, Constant, ASTNode
-)
+from mindsdb_sql_parser.ast import Identifier, BinaryOperation, Last, Constant, ASTNode
 from mindsdb.integrations.utilities.query_traversal import query_traversal
 from mindsdb.utilities.cache import get_cache
 
@@ -22,15 +20,15 @@ from .last_query import LastQuery
 
 class RunningQuery:
     """
-      Query in progres
+    Query in progres
     """
 
-    OBJECT_TYPE = 'query'
+    OBJECT_TYPE = "query"
 
     def __init__(self, record: db.Queries):
         self.record = record
         self.sql = record.sql
-        self.database = record.database or config.get('default_project')
+        self.database = record.database or config.get("default_project")
 
     def get_partitions(self, dn, step_call, query: Select) -> Iterable:
         """
@@ -41,7 +39,7 @@ class RunningQuery:
         :param query: AST query to execute
         :return: generator with query results
         """
-        if dn.has_support_stream():
+        if hasattr(dn, "has_support_stream") and dn.has_support_stream():
             query2 = self.get_partition_query(step_call.current_step_num, query, stream=True)
 
             for df in dn.query_stream(query2, fetch_size=self.batch_size):
@@ -53,10 +51,7 @@ class RunningQuery:
             while True:
                 query2 = self.get_partition_query(step_call.current_step_num, query, stream=False)
 
-                response = dn.query(
-                    query=query2,
-                    session=step_call.session
-                )
+                response = dn.query(query=query2, session=step_call.session)
                 df = response.data_frame
 
                 if df is None or len(df) == 0:
@@ -68,22 +63,22 @@ class RunningQuery:
 
     def get_partition_query(self, step_num: int, query: Select, stream=False) -> Select:
         """
-           Generate query for fetching the next partition
-           It wraps query to
-              select * from ({query})
-              where {track_column} > {previous_value}
-              order by track_column
-              limit size {batch_size}
-           And fill track_column, previous_value, batch_size
+        Generate query for fetching the next partition
+        It wraps query to
+           select * from ({query})
+           where {track_column} > {previous_value}
+           order by track_column
+           limit size {batch_size}
+        And fill track_column, previous_value, batch_size
 
-           If steam is true:
-             - if track_column is defined:
-                - don't add limit
-             - else:
-                - return user query without modifications
+        If steam is true:
+          - if track_column is defined:
+             - don't add limit
+          - else:
+             - return user query without modifications
         """
 
-        track_column = self.record.parameters.get('track_column')
+        track_column = self.record.parameters.get("track_column")
         if track_column is None and stream:
             # if no track column for stream fetching: it is not resumable query, execute original query
 
@@ -93,31 +88,30 @@ class RunningQuery:
             return query
 
         if not stream and track_column is None:
-            raise ValueError('Track column is not defined')
+            raise ValueError("Track column is not defined")
 
         query = Select(
             targets=[Star()],
             from_table=query,
             order_by=[OrderBy(Identifier(track_column))],
-
         )
         if not stream:
             query.limit = Constant(self.batch_size)
 
-        track_value = self.record.context.get('track_value')
+        track_value = self.record.context.get("track_value")
         # is it different step?
-        cur_step_num = self.record.context.get('step_num')
+        cur_step_num = self.record.context.get("step_num")
         if cur_step_num is not None and cur_step_num != step_num:
             # reset track_value
             track_value = None
-            self.record.context['track_value'] = None
-            self.record.context['step_num'] = step_num
-            flag_modified(self.record, 'context')
+            self.record.context["track_value"] = None
+            self.record.context["step_num"] = step_num
+            flag_modified(self.record, "context")
             db.session.commit()
 
         if track_value is not None:
             query.where = BinaryOperation(
-                op='>',
+                op=">",
                 args=[Identifier(track_column), Constant(track_value)],
             )
 
@@ -126,24 +120,22 @@ class RunningQuery:
     def get_info(self):
         record = self.record
         return {
-            'id': record.id,
-            'sql': record.sql,
-            'database': record.database,
-            'started_at': record.started_at,
-            'finished_at': record.finished_at,
-            'parameters': record.parameters,
-            'context': record.context,
-            'processed_rows': record.processed_rows,
-            'error': record.error,
-            'updated_at': record.updated_at,
+            "id": record.id,
+            "sql": record.sql,
+            "database": record.database,
+            "started_at": record.started_at,
+            "finished_at": record.finished_at,
+            "parameters": record.parameters,
+            "context": record.context,
+            "processed_rows": record.processed_rows,
+            "error": record.error,
+            "updated_at": record.updated_at,
         }
 
     def add_to_task(self):
-
         task_record = db.Tasks(
             company_id=ctx.company_id,
             user_class=ctx.user_class,
-
             object_type=self.OBJECT_TYPE,
             object_id=self.record.id,
         )
@@ -163,24 +155,24 @@ class RunningQuery:
 
     def set_params(self, params: dict):
         """
-            Store parameters of the step which is about to be split into partitions
+        Store parameters of the step which is about to be split into partitions
         """
 
-        if 'batch_size' not in params:
-            params['batch_size'] = 1000
+        if "batch_size" not in params:
+            params["batch_size"] = 1000
 
         self.record.parameters = params
-        self.batch_size = self.record.parameters['batch_size']
+        self.batch_size = self.record.parameters["batch_size"]
         db.session.commit()
 
     def get_max_track_value(self, df: pd.DataFrame) -> Optional[pd.DataFrame]:
         """
-            return max value to use in `set_progress`.
-            this function is called before execution substeps,
-             `set_progress` function - after
+        return max value to use in `set_progress`.
+        this function is called before execution substeps,
+         `set_progress` function - after
         """
-        if 'track_column' in self.record.parameters:
-            track_column = self.record.parameters['track_column']
+        if "track_column" in self.record.parameters:
+            track_column = self.record.parameters["track_column"]
             return df[track_column].max()
         else:
             # stream mode
@@ -188,7 +180,7 @@ class RunningQuery:
 
     def set_progress(self, df: pd.DataFrame, max_track_value: int):
         """
-           Store progres of the query, it is called after processing of batch
+        Store progres of the query, it is called after processing of batch
         """
 
         if len(df) == 0:
@@ -196,26 +188,26 @@ class RunningQuery:
 
         self.record.processed_rows = self.record.processed_rows + len(df)
 
-        cur_value = self.record.context.get('track_value')
+        cur_value = self.record.context.get("track_value")
         new_value = max_track_value
         if new_value is not None:
             if cur_value is None or new_value > cur_value:
-                self.record.context['track_value'] = new_value
-                flag_modified(self.record, 'context')
+                self.record.context["track_value"] = new_value
+                flag_modified(self.record, "context")
 
         db.session.commit()
 
     def on_error(self, error: Exception, step_num: int, steps_data: dict):
         """
-            Saves error of the query in database
-            Also saves step data and current step num to be able to resume query
+        Saves error of the query in database
+        Also saves step data and current step num to be able to resume query
         """
         self.record.error = str(error)
-        self.record.context['step_num'] = step_num
-        flag_modified(self.record, 'context')
+        self.record.context["step_num"] = step_num
+        flag_modified(self.record, "context")
 
         # save steps_data
-        cache = get_cache('steps_data')
+        cache = get_cache("steps_data")
         data = pickle.dumps(steps_data, protocol=5)
         cache.set(str(self.record.id), data)
 
@@ -223,10 +215,10 @@ class RunningQuery:
 
     def mark_as_run(self):
         """
-            Mark query as running and reset error of the query
+        Mark query as running and reset error of the query
         """
         if self.record.finished_at is not None:
-            raise RuntimeError('The query already finished')
+            raise RuntimeError("The query already finished")
 
         if self.record.started_at is None:
             self.record.started_at = dt.datetime.now()
@@ -235,13 +227,13 @@ class RunningQuery:
             self.record.error = None
             db.session.commit()
         else:
-            raise RuntimeError('The query might be running already')
+            raise RuntimeError("The query might be running already")
 
     def get_state(self) -> dict:
         """
-            Returns stored state for resuming the query
+        Returns stored state for resuming the query
         """
-        cache = get_cache('steps_data')
+        cache = get_cache("steps_data")
         key = self.record.id
         data = cache.get(key)
         cache.delete(key)
@@ -249,13 +241,13 @@ class RunningQuery:
         steps_data = pickle.loads(data)
 
         return {
-            'step_num': self.record.context.get('step_num'),
-            'steps_data': steps_data,
+            "step_num": self.record.context.get("step_num"),
+            "steps_data": steps_data,
         }
 
     def finish(self):
         """
-            Mark query as finished
+        Mark query as finished
         """
 
         self.record.finished_at = dt.datetime.now()
@@ -263,7 +255,7 @@ class RunningQuery:
 
 
 class QueryContextController:
-    IGNORE_CONTEXT = '<IGNORE>'
+    IGNORE_CONTEXT = "<IGNORE>"
 
     def handle_db_context_vars(self, query: ASTNode, dn, session) -> tuple:
         """
@@ -300,9 +292,9 @@ class QueryContextController:
             values = self._get_init_last_values(l_query, dn, session)
             if rec is None:
                 self.__add_context_record(context_name, query_str, values)
-                if context_name.startswith('job-if-'):
+                if context_name.startswith("job-if-"):
                     # add context for job also
-                    self.__add_context_record(context_name.replace('job-if', 'job'), query_str, values)
+                    self.__add_context_record(context_name.replace("job-if", "job"), query_str, values)
             else:
                 rec.values = values
         else:
@@ -319,20 +311,19 @@ class QueryContextController:
 
     def remove_lasts(self, query):
         def replace_lasts(node, **kwargs):
-
             # find last in where
             if isinstance(node, BinaryOperation):
                 if isinstance(node.args[0], Identifier) and isinstance(node.args[1], Last):
                     node.args = [Constant(0), Constant(0)]
-                    node.op = '='
+                    node.op = "="
 
         # find lasts
         query_traversal(query, replace_lasts)
         return query
 
-    def _result_callback(self, l_query: LastQuery,
-                         context_name: str, query_str: str,
-                         df: pd.DataFrame, columns_info: list):
+    def _result_callback(
+        self, l_query: LastQuery, context_name: str, query_str: str, df: pd.DataFrame, columns_info: list
+    ):
         """
         This function handlers result from executed query and updates context variables with new values
 
@@ -352,12 +343,12 @@ class QueryContextController:
         values = {}
         # get max values
         for info in l_query.get_last_columns():
-            target_idx = info['target_idx']
+            target_idx = info["target_idx"]
             if target_idx is not None:
                 # get by index
-                col_name = columns_info[target_idx]['name']
+                col_name = columns_info[target_idx]["name"]
             else:
-                col_name = info['column_name']
+                col_name = info["column_name"]
                 # get by name
             if col_name not in df:
                 continue
@@ -377,7 +368,7 @@ class QueryContextController:
                         continue
 
             if value is not None:
-                values[info['table_name']] = {info['column_name']: value}
+                values[info["table_name"]] = {info["column_name"]: value}
 
         self.__update_context_record(context_name, query_str, values)
 
@@ -389,10 +380,9 @@ class QueryContextController:
         """
 
         context_name = self.gen_context_name(object_type, object_id)
-        for rec in db.session.query(db.QueryContext).filter_by(
-            context_name=context_name,
-            company_id=ctx.company_id
-        ).all():
+        for rec in (
+            db.session.query(db.QueryContext).filter_by(context_name=context_name, company_id=ctx.company_id).all()
+        ):
             db.session.delete(rec)
         db.session.commit()
 
@@ -404,11 +394,7 @@ class QueryContextController:
         """
         last_values = {}
         for query, info in l_query.get_init_queries():
-
-            response = dn.query(
-                query=query,
-                session=session
-            )
+            response = dn.query(query=query, session=session)
             data = response.data_frame
             columns_info = response.columns
 
@@ -419,7 +405,7 @@ class QueryContextController:
 
                 idx = None
                 for i, col in enumerate(columns_info):
-                    if col['name'].upper() == info['column_name'].upper():
+                    if col["name"].upper() == info["column_name"].upper():
                         idx = i
                         break
 
@@ -429,7 +415,7 @@ class QueryContextController:
                     value = row[idx]
 
             if value is not None:
-                last_values[info['table_name']] = {info['column_name']: value}
+                last_values[info["table_name"]] = {info["column_name"]: value}
 
         return last_values
 
@@ -446,7 +432,7 @@ class QueryContextController:
         if len(context_stack) > 0:
             return context_stack[-1]
         else:
-            return ''
+            return ""
 
     def set_context(self, object_type: str = None, object_id: int = None):
         """
@@ -482,9 +468,9 @@ class QueryContextController:
         """
 
         if object_type is None:
-            return ''
+            return ""
         if object_id is not None:
-            object_type += '-' + str(object_id)
+            object_type += "-" + str(object_id)
         return object_type
 
     def get_context_vars(self, object_type: str, object_id: int) -> List[dict]:
@@ -495,10 +481,7 @@ class QueryContextController:
         """
         context_name = self.gen_context_name(object_type, object_id)
         vars = []
-        for rec in db.session.query(db.QueryContext).filter_by(
-            context_name=context_name,
-            company_id=ctx.company_id
-        ):
+        for rec in db.session.query(db.QueryContext).filter_by(context_name=context_name, company_id=ctx.company_id):
             if rec.values is not None:
                 vars.append(rec.values)
 
@@ -510,21 +493,17 @@ class QueryContextController:
         Find and return record for context and query string
         """
 
-        return db.session.query(db.QueryContext).filter_by(
-            query=query_str,
-            context_name=context_name,
-            company_id=ctx.company_id
-        ).first()
+        return (
+            db.session.query(db.QueryContext)
+            .filter_by(query=query_str, context_name=context_name, company_id=ctx.company_id)
+            .first()
+        )
 
     def __add_context_record(self, context_name: str, query_str: str, values: dict) -> db.QueryContext:
         """
         Creates record (for context and query string) with values and returns it
         """
-        rec = db.QueryContext(
-            query=query_str,
-            context_name=context_name,
-            company_id=ctx.company_id,
-            values=values)
+        rec = db.QueryContext(query=query_str, context_name=context_name, company_id=ctx.company_id, values=values)
         db.session.add(rec)
         return rec
 
@@ -538,27 +517,23 @@ class QueryContextController:
 
     def get_query(self, query_id: int) -> RunningQuery:
         """
-           Get running query by id
+        Get running query by id
         """
 
-        rec = db.Queries.query.filter(
-            db.Queries.id == query_id,
-            db.Queries.company_id == ctx.company_id
-        ).first()
+        rec = db.Queries.query.filter(db.Queries.id == query_id, db.Queries.company_id == ctx.company_id).first()
 
         if rec is None:
-            raise RuntimeError(f'Query not found: {query_id}')
+            raise RuntimeError(f"Query not found: {query_id}")
         return RunningQuery(rec)
 
     def create_query(self, query: ASTNode, database: str = None) -> RunningQuery:
         """
-           Create a new running query from AST query
+        Create a new running query from AST query
         """
 
         # remove old queries
         remove_query = db.session.query(db.Queries).filter(
-            db.Queries.company_id == ctx.company_id,
-            db.Queries.finished_at < (dt.datetime.now() - dt.timedelta(days=1))
+            db.Queries.company_id == ctx.company_id, db.Queries.finished_at < (dt.datetime.now() - dt.timedelta(days=1))
         )
         for rec in remove_query.all():
             self.get_query(rec.id).remove_from_task()
@@ -576,27 +551,19 @@ class QueryContextController:
 
     def list_queries(self) -> List[dict]:
         """
-           Get list of all running queries with metadata
+        Get list of all running queries with metadata
         """
 
-        query = db.session.query(db.Queries).filter(
-            db.Queries.company_id == ctx.company_id
-        )
-        return [
-            RunningQuery(record).get_info()
-            for record in query
-        ]
+        query = db.session.query(db.Queries).filter(db.Queries.company_id == ctx.company_id)
+        return [RunningQuery(record).get_info() for record in query]
 
     def cancel_query(self, query_id: int):
         """
-           Cancels running query by id
+        Cancels running query by id
         """
-        rec = db.Queries.query.filter(
-            db.Queries.id == query_id,
-            db.Queries.company_id == ctx.company_id
-        ).first()
+        rec = db.Queries.query.filter(db.Queries.id == query_id, db.Queries.company_id == ctx.company_id).first()
         if rec is None:
-            raise RuntimeError(f'Query not found: {query_id}')
+            raise RuntimeError(f"Query not found: {query_id}")
 
         self.get_query(rec.id).remove_from_task()