MindsDB 25.5.4.1__py3-none-any.whl → 25.6.2.0__py3-none-any.whl

mindsdb/integrations/utilities/files/file_reader.py

@@ -1,10 +1,11 @@
+from dataclasses import dataclass, astuple
 import traceback
 import json
 import csv
 from io import BytesIO, StringIO, IOBase
 from pathlib import Path
 import codecs
-from typing import List
+from typing import List, Generator
 
 import filetype
 import pandas as pd
@@ -18,8 +19,27 @@ DEFAULT_CHUNK_SIZE = 500
 DEFAULT_CHUNK_OVERLAP = 250
 
 
-class FileDetectError(Exception):
-    ...
+class FileProcessingError(Exception): ...
+
+
+@dataclass(frozen=True, slots=True)
+class _SINGLE_PAGE_FORMAT:
+    CSV: str = "csv"
+    JSON: str = "json"
+    TXT: str = "txt"
+    PDF: str = "pdf"
+    PARQUET: str = "parquet"
+
+
+SINGLE_PAGE_FORMAT = _SINGLE_PAGE_FORMAT()
+
+
+@dataclass(frozen=True, slots=True)
+class _MULTI_PAGE_FORMAT:
+    XLSX: str = "xlsx"
+
+
+MULTI_PAGE_FORMAT = _MULTI_PAGE_FORMAT()
 
 
 def decode(file_obj: IOBase) -> StringIO:
@@ -56,21 +76,20 @@ def decode(file_obj: IOBase) -> StringIO:
                 data_str = StringIO(byte_str.decode(encoding, errors))
     except Exception as e:
         logger.error(traceback.format_exc())
-        raise FileDetectError("Could not load into string") from e
+        raise FileProcessingError("Could not load into string") from e
 
     return data_str
 
 
 class FormatDetector:
-
-    supported_formats = ['parquet', 'csv', 'xlsx', 'pdf', 'json', 'txt']
-    multipage_formats = ['xlsx']
+    supported_formats = astuple(SINGLE_PAGE_FORMAT) + astuple(MULTI_PAGE_FORMAT)
+    multipage_formats = astuple(MULTI_PAGE_FORMAT)
 
     def __init__(
         self,
-        path: str = None,
-        name: str = None,
-        file: IOBase = None
+        path: str | None = None,
+        name: str | None = None,
+        file: IOBase | None = None,
     ):
         """
         File format detector
@@ -81,16 +100,16 @@ class FormatDetector:
         :param file: file descriptor (via open(...), of BytesIO(...))
         """
         if path is not None:
-            file = open(path, 'rb')
+            file = open(path, "rb")
 
         elif file is not None:
             if name is None:
-                if hasattr(file, 'name'):
+                if hasattr(file, "name"):
                     path = file.name
                 else:
-                    path = 'file'
+                    path = "file"
         else:
-            raise FileDetectError('Wrong arguments: path or file is required')
+            raise FileProcessingError("Wrong arguments: path or file is required")
 
         if name is None:
             name = Path(path).name
@@ -108,14 +127,14 @@ class FormatDetector:
         format = self.get_format_by_name()
         if format is not None:
             if format not in self.supported_formats:
-                raise FileDetectError(f'Not supported format: {format}')
+                raise FileProcessingError(f"Not supported format: {format}")
 
         if format is None and self.file_obj is not None:
             format = self.get_format_by_content()
             self.file_obj.seek(0)
 
         if format is None:
-            raise FileDetectError(f'Unable to detect format: {self.name}')
+            raise FileProcessingError(f"Unable to detect format: {self.name}")
 
         self.format = format
         return format
@@ -124,33 +143,32 @@ class FormatDetector:
         extension = Path(self.name).suffix.strip(".").lower()
         if extension == "tsv":
             extension = "csv"
-            self.parameters['delimiter'] = '\t'
+            self.parameters["delimiter"] = "\t"
 
         return extension or None
 
     def get_format_by_content(self):
         if self.is_parquet(self.file_obj):
-            return "parquet"
+            return SINGLE_PAGE_FORMAT.PARQUET
 
         file_type = filetype.guess(self.file_obj)
         if file_type is not None:
-
             if file_type.mime in {
                 "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
                 "application/vnd.ms-excel",
             }:
-                return 'xlsx'
+                return MULTI_PAGE_FORMAT.XLSX
 
-            if file_type.mime == 'application/pdf':
-                return "pdf"
+            if file_type.mime == "application/pdf":
+                return SINGLE_PAGE_FORMAT.PDF
 
         file_obj = decode(self.file_obj)
 
         if self.is_json(file_obj):
-            return "json"
+            return SINGLE_PAGE_FORMAT.JSON
 
         if self.is_csv(file_obj):
-            return "csv"
+            return SINGLE_PAGE_FORMAT.CSV
 
     @staticmethod
     def is_json(data_obj: StringIO) -> bool:
@@ -198,35 +216,53 @@ class FormatDetector:
         return False
 
 
-class FileReader(FormatDetector):
+def format_column_names(df: pd.DataFrame):
+    df.columns = [column.strip(" \t") for column in df.columns]
+    if len(df.columns) != len(set(df.columns)) or any(len(column_name) == 0 for column_name in df.columns):
+        raise FileProcessingError("Each column should have a unique and non-empty name.")
+
 
+class FileReader(FormatDetector):
     def _get_fnc(self):
         format = self.get_format()
-        func = getattr(self, f'read_{format}', None)
+        func = getattr(self, f"read_{format}", None)
         if func is None:
-            raise FileDetectError(f'Unsupported format: {format}')
-        return func
+            raise FileProcessingError(f"Unsupported format: {format}")
+
+        if format in astuple(MULTI_PAGE_FORMAT):
+
+            def format_multipage(*args, **kwargs):
+                for page_number, df in func(*args, **kwargs):
+                    format_column_names(df)
+                    yield page_number, df
+
+            return format_multipage
+
+        def format_singlepage(*args, **kwargs) -> pd.DataFrame:
+            """Check that the columns have unique not-empty names"""
+            df = func(*args, **kwargs)
+            format_column_names(df)
+            return df
+
+        return format_singlepage
 
     def get_pages(self, **kwargs) -> List[str]:
         """
-            Get list of tables in file
+        Get list of tables in file
         """
         format = self.get_format()
         if format not in self.multipage_formats:
             # only one table
-            return ['main']
+            return ["main"]
 
         func = self._get_fnc()
         self.file_obj.seek(0)
 
-        return [
-            name for name, _ in
-            func(self.file_obj, only_names=True, **kwargs)
-        ]
+        return [name for name, _ in func(self.file_obj, only_names=True, **kwargs)]
 
-    def get_contents(self, **kwargs):
+    def get_contents(self, **kwargs) -> dict[str, pd.DataFrame]:
         """
-            Get all info(pages with content) from file as dict: {tablename, content}
+        Get all info(pages with content) from file as dict: {tablename, content}
         """
         func = self._get_fnc()
         self.file_obj.seek(0)
@@ -234,17 +270,13 @@ class FileReader(FormatDetector):
         format = self.get_format()
         if format not in self.multipage_formats:
             # only one table
-            return {'main': func(self.file_obj, name=self.name, **kwargs)}
+            return {"main": func(self.file_obj, name=self.name, **kwargs)}
 
-        return {
-            name: df
-            for name, df in
-            func(self.file_obj, **kwargs)
-        }
+        return {name: df for name, df in func(self.file_obj, **kwargs)}
 
-    def get_page_content(self, page_name: str = None, **kwargs) -> pd.DataFrame:
+    def get_page_content(self, page_name: str | None = None, **kwargs) -> pd.DataFrame:
         """
-            Get content of a single table
+        Get content of a single table
         """
         func = self._get_fnc()
         self.file_obj.seek(0)
@@ -258,7 +290,7 @@ class FileReader(FormatDetector):
             return df
 
     @staticmethod
-    def _get_csv_dialect(buffer, delimiter=None) -> csv.Dialect:
+    def _get_csv_dialect(buffer, delimiter: str | None = None) -> csv.Dialect | None:
         sample = buffer.readline()  # trying to get dialect from header
         buffer.seek(0)
         try:
@@ -270,42 +302,35 @@ class FileReader(FormatDetector):
             else:
                 accepted_csv_delimiters = [",", "\t", ";"]
             try:
-                dialect = csv.Sniffer().sniff(
-                    sample, delimiters=accepted_csv_delimiters
-                )
-                dialect.doublequote = (
-                    True  # assume that all csvs have " as string escape
-                )
+                dialect = csv.Sniffer().sniff(sample, delimiters=accepted_csv_delimiters)
+                dialect.doublequote = True  # assume that all csvs have " as string escape
             except Exception:
                 dialect = csv.reader(sample).dialect
                 if dialect.delimiter not in accepted_csv_delimiters:
-                    raise Exception(
-                        f"CSV delimeter '{dialect.delimiter}' is not supported"
-                    )
+                    raise FileProcessingError(f"CSV delimeter '{dialect.delimiter}' is not supported")
 
         except csv.Error:
             dialect = None
         return dialect
 
     @classmethod
-    def read_csv(cls, file_obj: BytesIO, delimiter=None, **kwargs):
+    def read_csv(cls, file_obj: BytesIO, delimiter: str | None = None, **kwargs) -> pd.DataFrame:
         file_obj = decode(file_obj)
         dialect = cls._get_csv_dialect(file_obj, delimiter=delimiter)
-
         return pd.read_csv(file_obj, sep=dialect.delimiter, index_col=False)
 
     @staticmethod
-    def read_txt(file_obj: BytesIO, name=None, **kwargs):
+    def read_txt(file_obj: BytesIO, name: str | None = None, **kwargs) -> pd.DataFrame:
         # the lib is heavy, so import it only when needed
         from langchain_text_splitters import RecursiveCharacterTextSplitter
+
         file_obj = decode(file_obj)
 
         try:
             from langchain_core.documents import Document
         except ImportError:
-            raise ImportError(
-                "To import TXT document please install 'langchain-community':\n"
-                "    pip install langchain-community"
+            raise FileProcessingError(
+                "To import TXT document please install 'langchain-community':\n    pip install langchain-community"
             )
         text = file_obj.read()
 
@@ -317,15 +342,10 @@ class FileReader(FormatDetector):
         )
 
         docs = text_splitter.split_documents(documents)
-        return pd.DataFrame(
-            [
-                {"content": doc.page_content, "metadata": doc.metadata}
-                for doc in docs
-            ]
-        )
+        return pd.DataFrame([{"content": doc.page_content, "metadata": doc.metadata} for doc in docs])
 
     @staticmethod
-    def read_pdf(file_obj: BytesIO, name=None, **kwargs):
+    def read_pdf(file_obj: BytesIO, name: str | None = None, **kwargs) -> pd.DataFrame:
         # the libs are heavy, so import it only when needed
         import fitz  # pymupdf
         from langchain_text_splitters import RecursiveCharacterTextSplitter
@@ -340,30 +360,36 @@ class FileReader(FormatDetector):
         split_text = text_splitter.split_text(text)
 
         return pd.DataFrame(
-            {"content": split_text, "metadata": [{"file_format": "pdf", "source_file": name}] * len(split_text)}
+            {
+                "content": split_text,
+                "metadata": [{"file_format": "pdf", "source_file": name}] * len(split_text),
+            }
         )
 
     @staticmethod
-    def read_json(file_obj: BytesIO, **kwargs):
+    def read_json(file_obj: BytesIO, **kwargs) -> pd.DataFrame:
         file_obj = decode(file_obj)
         file_obj.seek(0)
         json_doc = json.loads(file_obj.read())
         return pd.json_normalize(json_doc, max_level=0)
 
     @staticmethod
-    def read_parquet(file_obj: BytesIO, **kwargs):
+    def read_parquet(file_obj: BytesIO, **kwargs) -> pd.DataFrame:
         return pd.read_parquet(file_obj)
 
     @staticmethod
-    def read_xlsx(file_obj: BytesIO, page_name=None, only_names=False, **kwargs):
+    def read_xlsx(
+        file_obj: BytesIO,
+        page_name: str | None = None,
+        only_names: bool = False,
+        **kwargs,
+    ) -> Generator[tuple[str, pd.DataFrame | None], None, None]:
         with pd.ExcelFile(file_obj) as xls:
-
             if page_name is not None:
                 # return specific page
                 yield page_name, pd.read_excel(xls, sheet_name=page_name)
 
             for page_name in xls.sheet_names:
-
                 if only_names:
                     # extract only pages names
                     df = None

mindsdb/integrations/utilities/handler_utils.py

@@ -39,22 +39,36 @@ def get_api_key(
     if "using" in create_args and f"{api_name.lower()}_api_key" in create_args["using"]:
         return create_args["using"][f"{api_name.lower()}_api_key"]
 
+    # 1.5 - Check for generic api_key in using
+    if "using" in create_args and "api_key" in create_args["using"]:
+        return create_args["using"]["api_key"]
+
     # 2
     if f"{api_name.lower()}_api_key" in create_args:
         return create_args[f"{api_name.lower()}_api_key"]
 
-    # 2.5 - Check in params dictionary if it exists (for agents)
+    # 2.5 - Check for generic api_key
+    if "api_key" in create_args:
+        return create_args["api_key"]
+
+    # 3 - Check in params dictionary if it exists (for agents)
     if "params" in create_args and create_args["params"] is not None:
         if f"{api_name.lower()}_api_key" in create_args["params"]:
             return create_args["params"][f"{api_name.lower()}_api_key"]
+        # 3.5 - Check for generic api_key in params
+        if "api_key" in create_args["params"]:
+            return create_args["params"]["api_key"]
 
-    # 3
+    # 4
     if engine_storage is not None:
         connection_args = engine_storage.get_connection_args()
         if f"{api_name.lower()}_api_key" in connection_args:
             return connection_args[f"{api_name.lower()}_api_key"]
+        # 4.5 - Check for generic api_key in connection_args
+        if "api_key" in connection_args:
+            return connection_args["api_key"]
 
-    # 4
+    # 5
     api_key = os.getenv(f"{api_name.lower()}_api_key")
     if api_key is not None:
         return api_key
@@ -62,15 +76,15 @@ def get_api_key(
     if api_key is not None:
         return api_key
 
-    # 5
+    # 6
     config = Config()
     api_cfg = config.get(api_name, {})
     if f"{api_name.lower()}_api_key" in api_cfg:
         return api_cfg[f"{api_name.lower()}_api_key"]
 
-    # 6
-    if 'api_keys' in create_args and api_name in create_args['api_keys']:
-        return create_args['api_keys'][api_name]
+    # 7
+    if "api_keys" in create_args and api_name in create_args["api_keys"]:
+        return create_args["api_keys"][api_name]
 
     if strict:
         provider_upper = api_name.upper()
@@ -79,8 +93,9 @@ def get_api_key(
         error_message = (
             f"API key for {api_name} not found. Please provide it using one of the following methods:\n"
             f"1. Set the {api_key_env_var} environment variable\n"
-            f"2. Provide it as '{api_key_arg}' parameter when creating an agent using the CREATE AGENT syntax\n"
+            f"2. Provide it as '{api_key_arg}' parameter or 'api_key' parameter when creating an agent using the CREATE AGENT syntax\n"
             f"   Example: CREATE AGENT my_agent USING model='gpt-4', provider='{api_name}', {api_key_arg}='your-api-key';\n"
+            f"   Or: CREATE AGENT my_agent USING model='gpt-4', provider='{api_name}', api_key='your-api-key';\n"
         )
         raise Exception(error_message)
     return None

mindsdb/integrations/utilities/sql_utils.py

@@ -46,11 +46,7 @@ class FilterCondition:
 
     def __eq__(self, __value: object) -> bool:
         if isinstance(__value, FilterCondition):
-            return (
-                self.column == __value.column
-                and self.op == __value.op
-                and self.value == __value.value
-            )
+            return self.column == __value.column and self.op == __value.op and self.value == __value.value
         else:
             return False
 
@@ -75,7 +71,7 @@ def make_sql_session():
     from mindsdb.api.executor.controllers.session_controller import SessionController
 
     sql_session = SessionController()
-    sql_session.database = config.get('default_project')
+    sql_session.database = config.get("default_project")
     return sql_session
 
 
@@ -84,44 +80,50 @@ def conditions_to_filter(binary_op: ASTNode):
 
     filters = {}
     for op, arg1, arg2 in conditions:
-        if op != '=':
+        if op != "=":
             raise NotImplementedError
         filters[arg1] = arg2
     return filters
 
 
-def extract_comparison_conditions(binary_op: ASTNode):
-    '''Extracts all simple comparison conditions that must be true from an AST node.
+def extract_comparison_conditions(binary_op: ASTNode, ignore_functions=False):
+    """Extracts all simple comparison conditions that must be true from an AST node.
     Does NOT support 'or' conditions.
-    '''
+    """
     conditions = []
 
     def _extract_comparison_conditions(node: ASTNode, **kwargs):
         if isinstance(node, ast.BinaryOperation):
             op = node.op.lower()
-            if op == 'and':
+            if op == "and":
                 # Want to separate individual conditions, not include 'and' as its own condition.
                 return
-            elif not isinstance(node.args[0], ast.Identifier):
+
+            arg1, arg2 = node.args
+            if ignore_functions and isinstance(arg1, ast.Function):
+                # handle lower/upper
+                if arg1.op.lower() in ("lower", "upper"):
+                    if isinstance(arg1.args[0], ast.Identifier):
+                        arg1 = arg1.args[0]
+
+            if not isinstance(arg1, ast.Identifier):
                 # Only support [identifier] =/</>/>=/<=/etc [constant] comparisons.
-                raise NotImplementedError(f'Not implemented arg1: {node.args[0]}')
+                raise NotImplementedError(f"Not implemented arg1: {arg1}")
 
-            if isinstance(node.args[1], ast.Constant):
-                value = node.args[1].value
-            elif isinstance(node.args[1], ast.Tuple):
-                value = [i.value for i in node.args[1].items]
+            if isinstance(arg2, ast.Constant):
+                value = arg2.value
+            elif isinstance(arg2, ast.Tuple):
+                value = [i.value for i in arg2.items]
             else:
-                raise NotImplementedError(f'Not implemented arg2: {node.args[1]}')
+                raise NotImplementedError(f"Not implemented arg2: {arg2}")
 
-            conditions.append([op, node.args[0].parts[-1], value])
+            conditions.append([op, arg1.parts[-1], value])
         if isinstance(node, ast.BetweenOperation):
             var, up, down = node.args
             if not (
-                isinstance(var, ast.Identifier)
-                and isinstance(up, ast.Constant)
-                and isinstance(down, ast.Constant)
+                isinstance(var, ast.Identifier) and isinstance(up, ast.Constant) and isinstance(down, ast.Constant)
             ):
-                raise NotImplementedError(f'Not implemented: {node}')
+                raise NotImplementedError(f"Not implemented: {node}")
 
             op = node.op.lower()
             conditions.append([op, var.parts[-1], (up.value, down.value)])
@@ -131,16 +133,13 @@ def extract_comparison_conditions(binary_op: ASTNode):
 
 
 def project_dataframe(df, targets, table_columns):
-    '''
-        case-insensitive projection
-        'select A' and 'select a' return different column case but with the same content
-    '''
+    """
+    case-insensitive projection
+    'select A' and 'select a' return different column case but with the same content
+    """
 
     columns = []
-    df_cols_idx = {
-        col.lower(): col
-        for col in df.columns
-    }
+    df_cols_idx = {col.lower(): col for col in df.columns}
     df_col_rename = {}
 
     for target in targets:
@@ -156,10 +155,7 @@ def project_dataframe(df, targets, table_columns):
             col = target.parts[-1]
             col_df = df_cols_idx.get(col.lower())
             if col_df is not None:
-                if (
-                    hasattr(target, 'alias')
-                    and isinstance(target.alias, ast.Identifier)
-                ):
+                if hasattr(target, "alias") and isinstance(target.alias, ast.Identifier):
                     df_col_rename[col_df] = target.alias.parts[0]
                 else:
                     df_col_rename[col_df] = col
@@ -184,14 +180,13 @@ def project_dataframe(df, targets, table_columns):
 
 
 def filter_dataframe(df: pd.DataFrame, conditions: list):
-
     # convert list of conditions to ast.
     # assumes that list was got from extract_comparison_conditions
     where_query = None
     for op, arg1, arg2 in conditions:
         op = op.lower()
 
-        if op == 'between':
+        if op == "between":
             item = ast.BetweenOperation(args=[ast.Identifier(arg1), ast.Constant(arg2[0]), ast.Constant(arg2[1])])
         else:
             if isinstance(arg2, (tuple, list)):
@@ -201,9 +196,9 @@ def filter_dataframe(df: pd.DataFrame, conditions: list):
         if where_query is None:
             where_query = item
         else:
-            where_query = ast.BinaryOperation(op='and', args=[where_query, item])
+            where_query = ast.BinaryOperation(op="and", args=[where_query, item])
 
-    query = ast.Select(targets=[ast.Star()], from_table=ast.Identifier('df'), where=where_query)
+    query = ast.Select(targets=[ast.Star()], from_table=ast.Identifier("df"), where=where_query)
 
     return query_df(df, query)
 
@@ -220,7 +215,7 @@ def sort_dataframe(df, order_by: list):
             continue
 
         cols.append(col)
-        ascending.append(False if order.direction.lower() == 'desc' else True)
+        ascending.append(False if order.direction.lower() == "desc" else True)
     if len(cols) > 0:
         df = df.sort_values(by=cols, ascending=ascending)
     return df