MindsDB 25.4.1.0__py3-none-any.whl → 25.4.2.1__py3-none-any.whl

mindsdb/integrations/handlers/chromadb_handler/chromadb_handler.py

@@ -1,5 +1,6 @@
 import ast
 import sys
+import os
 from typing import Dict, List, Optional, Union
 import hashlib
 
@@ -67,6 +68,8 @@ class ChromaDBHandler(VectorStoreHandler):
             "persist_directory": self.persist_directory,
         }
 
+        self._use_handler_storage = False
+
         self.connect()
 
     def validate_connection_parameters(self, name, **kwargs):
@@ -79,11 +82,15 @@ class ChromaDBHandler(VectorStoreHandler):
 
         config = ChromaHandlerConfig(**_config)
 
-        if config.persist_directory and not self.handler_storage.is_temporal:
-            # get full persistence directory from handler storage
-            self.persist_directory = self.handler_storage.folder_get(
-                config.persist_directory
-            )
+        if config.persist_directory:
+            if os.path.isabs(config.persist_directory):
+                self.persist_directory = config.persist_directory
+            elif not self.handler_storage.is_temporal:
+                # get full persistence directory from handler storage
+                self.persist_directory = self.handler_storage.folder_get(
+                    config.persist_directory
+                )
+                self._use_handler_storage = True
 
         return config
 
@@ -105,7 +112,7 @@ class ChromaDBHandler(VectorStoreHandler):
 
     def _sync(self):
         """Sync the database to disk if using persistent storage"""
-        if self.persist_directory:
+        if self.persist_directory and self._use_handler_storage:
             self.handler_storage.folder_sync(self.persist_directory)
 
     def __del__(self):
@@ -162,6 +169,8 @@ class ChromaDBHandler(VectorStoreHandler):
             FilterOperator.LESS_THAN_OR_EQUAL: "$lte",
             FilterOperator.GREATER_THAN: "$gt",
             FilterOperator.GREATER_THAN_OR_EQUAL: "$gte",
+            FilterOperator.IN: "$in",
+            FilterOperator.NOT_IN: "$nin",
         }
 
         if operator not in mapping:
@@ -308,7 +317,7 @@ class ChromaDBHandler(VectorStoreHandler):
         }
 
         if columns is not None:
-            payload = {column: payload[column] for column in columns}
+            payload = {column: payload[column] for column in columns if column != TableField.DISTANCE.value}
 
         # always include distance
         distance_filter = None
@@ -316,10 +325,11 @@ class ChromaDBHandler(VectorStoreHandler):
         if distances is not None:
             payload[distance_col] = distances
 
-            for cond in conditions:
-                if cond.column == distance_col:
-                    distance_filter = cond
-                    break
+            if conditions is not None:
+                for cond in conditions:
+                    if cond.column == distance_col:
+                        distance_filter = cond
+                        break
 
         df = pd.DataFrame(payload)
         if distance_filter is not None:
@@ -413,8 +423,8 @@ class ChromaDBHandler(VectorStoreHandler):
             collection.upsert(
                 ids=data_dict[TableField.ID.value],
                 documents=data_dict[TableField.CONTENT.value],
-                embeddings=data_dict.get(TableField.EMBEDDINGS.value),
-                metadatas=data_dict.get(TableField.METADATA.value)
+                embeddings=data_dict.get(TableField.EMBEDDINGS.value, None),
+                metadatas=data_dict.get(TableField.METADATA.value, None)
             )
             self._sync()
         except Exception as e:

mindsdb/integrations/handlers/langchain_embedding_handler/langchain_embedding_handler.py

@@ -104,6 +104,22 @@ def construct_model_from_args(args: Dict) -> Embeddings:
     return model
 
 
+def row_to_document(row: pd.Series) -> str:
+    """
+    Convert a row in the input dataframe into a document
+
+    Default implementation is to concatenate all the columns
+    in the form of
+    field1: value1\nfield2: value2\n...
+    """
+    fields = row.index.tolist()
+    values = row.values.tolist()
+    document = "\n".join(
+        [f"{field}: {value}" for field, value in zip(fields, values)]
+    )
+    return document
+
+
 class LangchainEmbeddingHandler(BaseMLEngine):
     """
     Bridge class to connect langchain.embeddings module to mindsDB
@@ -180,7 +196,7 @@ class LangchainEmbeddingHandler(BaseMLEngine):
             )
 
         # convert each row into a document
-        df_texts = df[input_columns].apply(self.row_to_document, axis=1)
+        df_texts = df[input_columns].apply(row_to_document, axis=1)
         embeddings = model.embed_documents(df_texts.tolist())
 
         # create a new dataframe with the embeddings
@@ -188,21 +204,6 @@ class LangchainEmbeddingHandler(BaseMLEngine):
 
         return df_embeddings
 
-    def row_to_document(self, row: pd.Series) -> str:
-        """
-        Convert a row in the input dataframe into a document
-
-        Default implementation is to concatenate all the columns
-        in the form of
-        field1: value1\nfield2: value2\n...
-        """
-        fields = row.index.tolist()
-        values = row.values.tolist()
-        document = "\n".join(
-            [f"{field}: {value}" for field, value in zip(fields, values)]
-        )
-        return document
-
     def finetune(
         self, df: Union[DataFrame, None] = None, args: Union[Dict, None] = None
     ) -> None:

mindsdb/integrations/handlers/langchain_handler/langchain_handler.py

@@ -50,6 +50,7 @@ class LangChainHandler(BaseMLEngine):
         - OpenAI
         - Anthropic
         - Anyscale
+        - Google
         - LiteLLM
         - Ollama
 

mindsdb/integrations/handlers/mssql_handler/mssql_handler.py

@@ -177,7 +177,7 @@ class SqlServerHandler(DatabaseHandler):
                         )
                     )
                 else:
-                    response = Response(RESPONSE_TYPE.OK)
+                    response = Response(RESPONSE_TYPE.OK, affected_rows=cur.rowcount)
                 connection.commit()
             except Exception as e:
                 logger.error(f'Error running query: {query} on {self.database}, {e}!')

mindsdb/integrations/handlers/mysql_handler/mysql_handler.py

@@ -178,10 +178,11 @@ class MySQLHandler(DatabaseHandler):
                         pd.DataFrame(
                             result,
                             columns=[x[0] for x in cur.description]
-                        )
+                        ),
+                        affected_rows=cur.rowcount
                     )
                 else:
-                    response = Response(RESPONSE_TYPE.OK)
+                    response = Response(RESPONSE_TYPE.OK, affected_rows=cur.rowcount)
         except mysql.connector.Error as e:
             logger.error(f'Error running query: {query} on {self.connection_data["database"]}!')
             response = Response(

mindsdb/integrations/handlers/oracle_handler/oracle_handler.py

@@ -205,8 +205,10 @@ class OracleHandler(DatabaseHandler):
         with connection.cursor() as cur:
             try:
                 cur.execute(query)
-                result = cur.fetchall()
-                if result:
+                if cur.description is None:
+                    response = Response(RESPONSE_TYPE.OK, affected_rows=cur.rowcount)
+                else:
+                    result = cur.fetchall()
                     response = Response(
                         RESPONSE_TYPE.TABLE,
                         data_frame=pd.DataFrame(
@@ -214,8 +216,6 @@ class OracleHandler(DatabaseHandler):
                             columns=[row[0] for row in cur.description],
                         ),
                     )
-                else:
-                    response = Response(RESPONSE_TYPE.OK)
 
                 connection.commit()
             except DatabaseError as database_error:

mindsdb/integrations/handlers/pgvector_handler/pgvector_handler.py

@@ -46,7 +46,8 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
 
     def _make_connection_args(self):
         cloud_pgvector_url = os.environ.get('KB_PGVECTOR_URL')
-        if cloud_pgvector_url is not None:
+        # if no connection args and shared pg vector defined - use it
+        if len(self.connection_args) == 0 and cloud_pgvector_url is not None:
             result = urlparse(cloud_pgvector_url)
             self.connection_args = {
                 'host': result.hostname,
@@ -149,7 +150,7 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
         for key, value in filter_conditions.items():
             if key == "embeddings":
                 continue
-            if value['op'].lower() == 'in':
+            if value['op'].lower() in ('in', 'not in'):
                 values = list(repr(i) for i in value['value'])
                 value['value'] = '({})'.format(', '.join(values))
             else:
@@ -157,7 +158,7 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
             where_clauses.append(f'{key} {value["op"]} {value["value"]}')
 
         if len(where_clauses) > 1:
-            return f"WHERE{' AND '.join(where_clauses)}"
+            return f"WHERE {' AND '.join(where_clauses)}"
         elif len(where_clauses) == 1:
             return f"WHERE {where_clauses[0]}"
         else:
@@ -165,9 +166,9 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
 
     @staticmethod
     def _construct_full_after_from_clause(
+        where_clause: str,
         offset_clause: str,
         limit_clause: str,
-        where_clause: str,
     ) -> str:
 
         return f"{where_clause} {offset_clause} {limit_clause}"
@@ -195,21 +196,26 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
         # given filter conditions, construct where clause
         where_clause = self._construct_where_clause(filter_conditions)
 
-        # construct full after from clause, where clause + offset clause + limit clause
-        after_from_clause = self._construct_full_after_from_clause(
-            where_clause, offset_clause, limit_clause
-        )
-
-        if columns is None:
-            targets = '*'
+        # Handle distance column specially since it's calculated, not stored
+        modified_columns = []
+        has_distance = False
+        if columns is not None:
+            for col in columns:
+                if col == TableField.DISTANCE.value:
+                    has_distance = True
+                else:
+                    modified_columns.append(col)
         else:
-            targets = ', '.join(columns)
+            modified_columns = ['id', 'content', 'embeddings', 'metadata']
+            has_distance = True
+
+        targets = ', '.join(modified_columns)
 
 
         if filter_conditions:
 
             if embedding_search:
-                search_vector = filter_conditions["embeddings"]["value"][0]
+                search_vector = filter_conditions["embeddings"]["value"]
                 filter_conditions.pop("embeddings")
 
                 if self._is_sparse:
@@ -227,15 +233,19 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
                     # Use cosine similarity for dense vectors
                     distance_op = "<=>"
 
-                return f"SELECT {targets} FROM {table_name} ORDER BY embeddings {distance_op} '{search_vector}' ASC {after_from_clause}"
+                # Calculate distance as part of the query if needed
+                if has_distance:
+                    targets = f"{targets}, (embeddings {distance_op} '{search_vector}') as distance"
+                
+                return f"SELECT {targets} FROM {table_name} {where_clause} ORDER BY embeddings {distance_op} '{search_vector}' ASC {limit_clause} {offset_clause} "
 
             else:
                 # if filter conditions, return rows that satisfy the conditions
-                return f"SELECT {targets} FROM {table_name} {after_from_clause}"
+                return f"SELECT {targets} FROM {table_name} {where_clause} {limit_clause} {offset_clause}"
 
         else:
             # if no filter conditions, return all rows
-            return f"SELECT {targets} FROM {table_name} {after_from_clause}"
+            return f"SELECT {targets} FROM {table_name} {limit_clause} {offset_clause}"
 
     def _check_table(self, table_name: str):
         # Apply namespace for a user

mindsdb/integrations/handlers/postgres_handler/postgres_handler.py

@@ -1,6 +1,7 @@
 import time
 import json
 from typing import Optional
+import threading
 
 import pandas as pd
 import psycopg
@@ -77,6 +78,8 @@ class PostgresHandler(DatabaseHandler):
         self.is_connected = False
         self.thread_safe = True
 
+        self._insert_lock = threading.Lock()
+
     def __del__(self):
         if self.is_connected:
             self.disconnect()
@@ -228,7 +231,7 @@ class PostgresHandler(DatabaseHandler):
                 else:
                     cur.execute(query)
                 if cur.pgresult is None or ExecStatus(cur.pgresult.status) == ExecStatus.COMMAND_OK:
-                    response = Response(RESPONSE_TYPE.OK)
+                    response = Response(RESPONSE_TYPE.OK, affected_rows=cur.rowcount)
                 else:
                     result = cur.fetchall()
                     df = DataFrame(
@@ -238,7 +241,8 @@ class PostgresHandler(DatabaseHandler):
                     self._cast_dtypes(df, cur.description)
                     response = Response(
                         RESPONSE_TYPE.TABLE,
-                        df
+                        data_frame=df,
+                        affected_rows=cur.rowcount
                     )
                 connection.commit()
             except Exception as e:
@@ -255,26 +259,51 @@ class PostgresHandler(DatabaseHandler):
 
         return response
 
-    def insert(self, table_name: str, df: pd.DataFrame):
+    def insert(self, table_name: str, df: pd.DataFrame) -> Response:
         need_to_close = not self.is_connected
 
         connection = self.connect()
 
-        columns = [f'"{c}"' for c in df.columns]
+        columns = df.columns
+
+        # postgres 'copy' is not thread safe. use lock to prevent concurrent execution
+        with self._insert_lock:
+            resp = self.get_columns(table_name)
+
+        # copy requires precise cases of names: get current column names from table and adapt input dataframe columns
+        if resp.data_frame is not None and not resp.data_frame.empty:
+            db_columns = {
+                c.lower(): c
+                for c in resp.data_frame['Field']
+            }
+
+            # try to get case of existing column
+            columns = [
+                db_columns.get(c.lower(), c)
+                for c in columns
+            ]
+
+        columns = [f'"{c}"' for c in columns]
+        rowcount = None
+
         with connection.cursor() as cur:
             try:
-                with cur.copy(f'copy "{table_name}" ({",".join(columns)}) from STDIN  WITH CSV') as copy:
-                    df.to_csv(copy, index=False, header=False)
+                with self._insert_lock:
+                    with cur.copy(f'copy "{table_name}" ({",".join(columns)}) from STDIN WITH CSV') as copy:
+                        df.to_csv(copy, index=False, header=False)
 
-                connection.commit()
+                    connection.commit()
             except Exception as e:
                 logger.error(f'Error running insert to {table_name} on {self.database}, {e}!')
                 connection.rollback()
                 raise e
+            rowcount = cur.rowcount
 
         if need_to_close:
             self.disconnect()
 
+        return Response(RESPONSE_TYPE.OK, affected_rows=rowcount)
+
     @profiler.profile()
     def query(self, query: ASTNode) -> Response:
         """

mindsdb/integrations/handlers/redshift_handler/redshift_handler.py

@@ -52,7 +52,7 @@ class RedshiftHandler(PostgresHandler):
         with connection.cursor() as cur:
             try:
                 cur.executemany(query, df.values.tolist())
-                response = Response(RESPONSE_TYPE.OK)
+                response = Response(RESPONSE_TYPE.OK, affected_rows=cur.rowcount)
 
                 connection.commit()
             except Exception as e:

mindsdb/integrations/handlers/snowflake_handler/snowflake_handler.py

@@ -230,18 +230,25 @@ class SnowflakeHandler(DatabaseHandler):
                     # Fallback for CREATE/DELETE/UPDATE. These commands returns table with single column,
                     # but it cannot be retrieved as pandas DataFrame.
                     result = cur.fetchall()
-                    if result:
-                        response = Response(
-                            RESPONSE_TYPE.TABLE,
-                            DataFrame(
-                                result,
-                                columns=[x[0] for x in cur.description]
+                    match result:
+                        case (
+                            [{'number of rows inserted': affected_rows}]
+                            | [{'number of rows deleted': affected_rows}]
+                            | [{'number of rows updated': affected_rows, 'number of multi-joined rows updated': _}]
+                        ):
+                            response = Response(RESPONSE_TYPE.OK, affected_rows=affected_rows)
+                        case list():
+                            response = Response(
+                                RESPONSE_TYPE.TABLE,
+                                DataFrame(
+                                    result,
+                                    columns=[x[0] for x in cur.description]
+                                )
                             )
-                        )
-                    else:
-                        # Looks like SnowFlake always returns something in response, so this is suspicious
-                        logger.warning('Snowflake did not return any data in response.')
-                        response = Response(RESPONSE_TYPE.OK)
+                        case _:
+                            # Looks like SnowFlake always returns something in response, so this is suspicious
+                            logger.warning('Snowflake did not return any data in response.')
+                            response = Response(RESPONSE_TYPE.OK)
             except Exception as e:
                 logger.error(f"Error running query: {query} on {self.connection_data.get('database')}, {e}!")
                 response = Response(

mindsdb/integrations/libs/llm/config.py

@@ -1,6 +1,6 @@
 from typing import Any, Dict, List, Optional
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, Field
 
 
 class BaseLLMConfig(BaseModel):
@@ -104,3 +104,13 @@ class NvidiaNIMConfig(BaseLLMConfig):
 class MindsdbConfig(BaseLLMConfig):
     model_name: str
     project_name: str
+
+
+# See https://python.langchain.com/api_reference/google_genai/chat_models/langchain_google_genai.chat_models.ChatGoogleGenerativeAI.html
+class GoogleConfig(BaseLLMConfig):
+    model: str = Field(description="Gemini model name to use (e.g., 'gemini-1.5-pro')")
+    temperature: Optional[float] = Field(default=None, description="Controls randomness in responses")
+    top_p: Optional[float] = Field(default=None, description="Nucleus sampling parameter")
+    top_k: Optional[int] = Field(default=None, description="Number of highest probability tokens to consider")
+    max_output_tokens: Optional[int] = Field(default=None, description="Maximum number of tokens to generate")
+    google_api_key: Optional[str] = Field(default=None, description="API key for Google Generative AI")

mindsdb/integrations/libs/llm/utils.py

@@ -10,6 +10,7 @@ from mindsdb.integrations.libs.llm.config import (
     AnthropicConfig,
     AnyscaleConfig,
     BaseLLMConfig,
+    GoogleConfig,
     LiteLLMConfig,
     OllamaConfig,
     OpenAIConfig,
@@ -31,6 +32,8 @@ DEFAULT_ANTHROPIC_MODEL = "claude-3-haiku-20240307"
 DEFAULT_ANYSCALE_MODEL = "meta-llama/Llama-2-7b-chat-hf"
 DEFAULT_ANYSCALE_BASE_URL = "https://api.endpoints.anyscale.com/v1"
 
+DEFAULT_GOOGLE_MODEL = "gemini-2.5-pro-preview-03-25"
+
 DEFAULT_LITELLM_MODEL = "gpt-3.5-turbo"
 DEFAULT_LITELLM_PROVIDER = "openai"
 DEFAULT_LITELLM_BASE_URL = "https://ai.dev.mindsdb.com"
@@ -225,6 +228,15 @@ def get_llm_config(provider: str, args: Dict) -> BaseLLMConfig:
             openai_organization=args.get("api_organization", None),
             request_timeout=args.get("request_timeout", None),
         )
+    if provider == "google":
+        return GoogleConfig(
+            model=args.get("model_name", DEFAULT_GOOGLE_MODEL),
+            temperature=temperature,
+            top_p=args.get("top_p", None),
+            top_k=args.get("top_k", None),
+            max_output_tokens=args.get("max_tokens", None),
+            google_api_key=args["api_keys"].get("google", None),
+        )
 
     raise ValueError(f"Provider {provider} is not supported.")
 

mindsdb/integrations/libs/ml_handler_process/learn_process.py

@@ -78,8 +78,7 @@ def learn_process(data_integration_ref: dict, problem_definition: dict, fetch_da
                     query_ast = parse_sql(fetch_data_query)
                     sqlquery = SQLQuery(query_ast, session=sql_session)
 
-                result = sqlquery.fetch(view='dataframe')
-                training_data_df = result['result']
+                training_data_df = sqlquery.fetched_data.to_df()
 
             training_data_columns_count, training_data_rows_count = 0, 0
             if training_data_df is not None:

mindsdb/integrations/libs/response.py

@@ -1,3 +1,4 @@
+from typing import Optional
 from pandas import DataFrame
 
 from mindsdb.utilities import log
@@ -8,13 +9,16 @@ from mindsdb_sql_parser.ast import ASTNode
 logger = log.getLogger(__name__)
 
 class HandlerResponse:
-    def __init__(self, resp_type: RESPONSE_TYPE, data_frame: DataFrame = None,
-                 query: ASTNode = 0, error_code: int = 0, error_message: str = None) -> None:
+    def __init__(self, resp_type: RESPONSE_TYPE, data_frame: DataFrame = None, query: ASTNode = 0, error_code: int = 0,
+                 error_message: Optional[str] = None, affected_rows: Optional[int] = None) -> None:
         self.resp_type = resp_type
         self.query = query
         self.data_frame = data_frame
         self.error_code = error_code
         self.error_message = error_message
+        self.affected_rows = affected_rows
+        if isinstance(self.affected_rows, int) is False or self.affected_rows < 0:
+            self.affected_rows = 0
 
     @property
     def type(self):
@@ -35,13 +39,14 @@ class HandlerResponse:
                  "error": self.error_message}
 
     def __repr__(self):
-        return "%s: resp_type=%s, query=%s, data_frame=%s, err_code=%s, error=%s" % (
+        return "%s: resp_type=%s, query=%s, data_frame=%s, err_code=%s, error=%s, affected_rows=%s" % (
                 self.__class__.__name__,
                 self.resp_type,
                 self.query,
                 self.data_frame,
                 self.error_code,
-                self.error_message
+                self.error_message,
+                self.affected_rows
             )
 
 class HandlerStatusResponse:

mindsdb/integrations/libs/vectordatabase_handler.py

@@ -20,7 +20,7 @@ from mindsdb_sql_parser.ast.base import ASTNode
 
 from mindsdb.integrations.libs.response import RESPONSE_TYPE, HandlerResponse
 from mindsdb.utilities import log
-from mindsdb.integrations.utilities.sql_utils import conditions_to_filter, FilterCondition, FilterOperator
+from mindsdb.integrations.utilities.sql_utils import FilterCondition, FilterOperator
 
 from mindsdb.integrations.utilities.query_traversal import query_traversal
 from .base import BaseHandler
@@ -39,6 +39,7 @@ class TableField(Enum):
     METADATA = "metadata"
     SEARCH_VECTOR = "search_vector"
     DISTANCE = "distance"
+    RELEVANCE = "relevance"
 
 
 class DistanceFunction(Enum):
@@ -69,6 +70,10 @@ class VectorStoreHandler(BaseHandler):
             "name": TableField.METADATA.value,
             "data_type": "json",
         },
+        {
+            "name": TableField.DISTANCE.value,
+            "data_type": "float",
+        },
     ]
 
     def validate_connection_parameters(self, name, **kwargs):
@@ -231,7 +236,7 @@ class VectorStoreHandler(BaseHandler):
 
         return self.do_upsert(table_name, pd.DataFrame(data))
 
-    def _dispatch_update(self, query: Update):
+    def dispatch_update(self, query: Update, conditions: List[FilterCondition] = None):
         """
         Dispatch update query to the appropriate method.
         """
@@ -250,8 +255,15 @@ class VectorStoreHandler(BaseHandler):
                     pass
             row[k] = v
 
-        filters = conditions_to_filter(query.where)
-        row.update(filters)
+        if conditions is None:
+            where_statement = query.where
+            conditions = self.extract_conditions(where_statement)
+
+        for condition in conditions:
+            if condition.op != FilterOperator.EQUAL:
+                raise NotImplementedError
+
+            row[condition.column] = condition.value
 
         # checks
         if TableField.EMBEDDINGS.value not in row:
@@ -381,7 +393,7 @@ class VectorStoreHandler(BaseHandler):
             CreateTable: self._dispatch_create_table,
             DropTables: self._dispatch_drop_table,
             Insert: self._dispatch_insert,
-            Update: self._dispatch_update,
+            Update: self.dispatch_update,
             Delete: self.dispatch_delete,
             Select: self.dispatch_select,
         }