MindsDB 25.7.3.0__py3-none-any.whl → 25.8.2.0__py3-none-any.whl

mindsdb/interfaces/knowledge_base/preprocessing/text_splitter.py

@@ -0,0 +1,73 @@
+from typing import List
+
+
+class TextSplitter:
+    def __init__(
+        self,
+        chunk_size: int = 1000,
+        chunk_overlap: int = 200,
+        separators: List[str] = None,
+        k_range: float = 0.5,
+        k_ratio: float = 1,
+    ):
+        """
+        Split text into chunks. The logic:
+         - Get a piece of text with chunk_size and try to find the separator at the end of the piece.
+         - The allowed range to find the separator is defined by k_range and k_ratio using formula:
+            k_range * chunk_size / (num * k_ratio + 1)
+            num - is number of a separator from the list
+         - if the separator is not in the rage: switch to the next separator
+         - if the found separator is in the middle of the sentence, use overlapping:
+            - the found text is the current chunk
+            - repeat the search with less strict k_range and k_ratio
+            - the found text will be the beginning of the next chunk
+
+        :param chunk_size: size of the chunk, which must not be exceeded
+        :param separators: list of separators in order of priority
+        :param k_range: defines the range to look for the separator
+        :param k_ratio: defines how much to shrink the range for the next separator
+        """
+        if separators is None:
+            separators = ["\n\n", "\n", ". ", " ", ""]
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+        self.separators = separators
+        self.k_range = k_range
+        self.k_ratio = k_ratio
+
+    def split_text(self, text: str) -> List[str]:
+        chunks = []
+
+        while True:
+            if len(text) < self.chunk_size:
+                chunks.append(text)
+                break
+
+            sep, chunk, shift = self.get_next_chunk(text, self.k_range, self.k_ratio)
+            chunks.append(chunk)
+
+            text = text[shift:]
+        return chunks
+
+    def get_next_chunk(self, text: str, k_range: float, k_ratio: float):
+        # returns chunk with separator and shift for the next search iteration
+
+        chunk = text[: self.chunk_size]
+        # positions = []
+        for i, sep in enumerate(self.separators):
+            pos = chunk.rfind(sep)
+
+            vpos = self.chunk_size - pos
+            if vpos < k_range * self.chunk_size / (i * k_ratio + 1):
+                shift = len(sep) + pos
+                if sep.strip(" ") == "":
+                    # overlapping
+                    sep2, _, shift2 = self.get_next_chunk(text, k_range * 1.5, 0)
+                    if sep2.strip(" ") != "":
+                        # use shift of previous separator
+                        if shift - shift2 < self.chunk_overlap:
+                            shift = shift2
+
+                return sep, chunk[:pos], shift
+
+        raise RuntimeError("Cannot split text")

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,22 +39,19 @@ 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):
                 max_track_value = self.get_max_track_value(df)
                 yield df
-                self.set_progress(df, max_track_value)
+                self.set_progress(max_track_value=max_track_value)
 
         else:
             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:
@@ -64,26 +59,26 @@ class RunningQuery:
 
                 max_track_value = self.get_max_track_value(df)
                 yield df
-                self.set_progress(df, max_track_value)
+                self.set_progress(max_track_value=max_track_value)
 
     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,59 +155,58 @@ 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
             return None
 
-    def set_progress(self, df: pd.DataFrame, max_track_value: int):
+    def set_progress(self, processed_rows: int = None, max_track_value: int = None):
         """
-           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:
-            return
-
-        self.record.processed_rows = self.record.processed_rows + len(df)
-
-        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')
+        if processed_rows is not None and processed_rows > 0:
+            self.record.processed_rows = self.record.processed_rows + processed_rows
+            db.session.commit()
 
-        db.session.commit()
+        if max_track_value is not None:
+            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")
+            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 +214,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 +226,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 +240,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 +254,7 @@ class RunningQuery:
 
 
 class QueryContextController:
-    IGNORE_CONTEXT = '<IGNORE>'
+    IGNORE_CONTEXT = "<IGNORE>"
 
     def handle_db_context_vars(self, query: ASTNode, dn, session) -> tuple:
         """
@@ -300,9 +291,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 +310,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 +342,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 +367,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 +379,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 +393,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 +404,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 +414,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 +431,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 +467,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 +480,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 +492,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 +516,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 +550,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()