MindsDB 25.5.3.0__py3-none-any.whl → 25.5.4.0__py3-none-any.whl

mindsdb/integrations/handlers/stripe_handler/stripe_tables.py

@@ -1,17 +1,11 @@
 import pandas as pd
-
 import stripe
 from typing import Text, List, Dict, Any
-from mindsdb_sql_parser import ast
-from mindsdb.integrations.libs.api_handler import APITable
-from mindsdb.integrations.utilities.handlers.query_utilities import INSERTQueryParser, DELETEQueryParser, UPDATEQueryParser, DELETEQueryExecutor, UPDATEQueryExecutor
 
 from mindsdb_sql_parser import ast
 from mindsdb.integrations.libs.api_handler import APITable
-
+from mindsdb.integrations.utilities.handlers.query_utilities import INSERTQueryParser, DELETEQueryParser, UPDATEQueryParser, DELETEQueryExecutor, UPDATEQueryExecutor
 from mindsdb.integrations.utilities.handlers.query_utilities.select_query_utilities import SELECTQueryParser, SELECTQueryExecutor
-
-
 from mindsdb.utilities import log
 
 logger = log.getLogger(__name__)
@@ -235,7 +229,7 @@ class ProductsTable(APITable):
 
 class PaymentIntentsTable(APITable):
     """The Stripe Payment Intents Table implementation"""
-    
+
     def select(self, query: ast.Select) -> pd.DataFrame:
         """
         Pulls Stripe Payment Intents data.
@@ -273,6 +267,7 @@ class PaymentIntentsTable(APITable):
         payment_intents_df = select_statement_executor.execute_query()
 
         return payment_intents_df
+
     def delete(self, query: ast.Delete) -> None:
         """
         Cancels Stripe Payment Intents and updates the local data.
@@ -307,22 +302,21 @@ class PaymentIntentsTable(APITable):
         payment_intent_ids = canceled_payment_intents_df['id'].tolist()
         self.cancel_payment_intents(payment_intent_ids)
 
-        
         self.payment_intents_df = self.payment_intents_df[~self.payment_intents_df['id'].isin(payment_intent_ids)]
 
     def cancel_payment_intents(self, payment_intent_ids: List[str]) -> None:
         stripe = self.handler.connect()
         for payment_intent_id in payment_intent_ids:
             try:
-              
+
                 payment_intent = stripe.PaymentIntent.retrieve(payment_intent_id)
                 if payment_intent.status in ['requires_payment_method', 'requires_capture', 'requires_confirmation', 'requires_action', 'processing']:
                     stripe.PaymentIntent.cancel(payment_intent_id)
                 else:
-                     logger.warning(f"Payment intent {payment_intent_id} is in status {payment_intent.status} and cannot be canceled.")
+                    logger.warning(f"Payment intent {payment_intent_id} is in status {payment_intent.status} and cannot be canceled.")
             except stripe.error.StripeError as e:
                 logger.error(f"Error cancelling payment intent {payment_intent_id}: {str(e)}")
-                
+
     def update(self, query: 'ast.Update') -> None:
         """
         Updates data in Stripe "POST /v1/payment_intents/:id" API endpoint.
@@ -358,16 +352,6 @@ class PaymentIntentsTable(APITable):
         for payment_intent_id in payment_intent_ids:
             stripe.PaymentIntent.modify(payment_intent_id, **values_to_update)
 
-
-    def get_columns(self) -> List[Text]:
-        return pd.json_normalize(self.get_payment_intents(limit=1)).columns.tolist()
-
-    def get_payment_intents(self, **kwargs) -> List[Dict]:
-        payment_intents = stripe.PaymentIntent.list(**kwargs)
-        return [payment_intent for payment_intent in payment_intents.data]
-
-    
-    
     def insert(self, query: 'ast.Insert') -> None:
         """
         Inserts data into Stripe "POST /v1/payment_intents" API endpoint.
@@ -398,7 +382,7 @@ class PaymentIntentsTable(APITable):
     def create_payment_intent(self, payment_intent_data: list) -> None:
         for data in payment_intent_data:
             stripe.PaymentIntent.create(**data)
-        
+
     def get_columns(self) -> List[Text]:
         return pd.json_normalize(self.get_payment_intents(limit=1)).columns.tolist()
 
@@ -407,7 +391,7 @@ class PaymentIntentsTable(APITable):
         payment_intents = stripe.PaymentIntent.list(**kwargs)
         return [payment_intent.to_dict() for payment_intent in payment_intents]
 
-      
+
 class RefundsTable(APITable):
     """The Stripe Refund Table implementation"""
 
@@ -455,9 +439,9 @@ class RefundsTable(APITable):
     def get_refunds(self, **kwargs) -> List[Dict]:
         stripe = self.handler.connect()
         refunds = stripe.Refund.list(**kwargs)
-        return [refund.to_dict() for refund in refunds ]
-     
-    
+        return [refund.to_dict() for refund in refunds]
+
+
 class PayoutsTable(APITable):
     """The Stripe Payouts Table implementation"""
 

mindsdb/integrations/handlers/supabase_handler/tests/test_supabase_handler.py

@@ -43,4 +43,4 @@ class SupabaseHandlerTest(unittest.TestCase):
 
     def test_7_select_query(self):
         query = "SELECT * FROM dt_test WHERE 'id'='a'"
-        result = self.handler.native_query(query)
+        self.handler.native_query(query)

mindsdb/integrations/handlers/surrealdb_handler/surrealdb_handler.py

@@ -1,4 +1,7 @@
 from typing import Optional
+import pysurrealdb as surreal
+import pandas as pd
+
 from mindsdb_sql_parser.ast.base import ASTNode
 from mindsdb.integrations.libs.base import DatabaseHandler
 from mindsdb.utilities import log
@@ -8,10 +11,7 @@ from mindsdb.integrations.libs.response import (
     HandlerResponse as Response,
     RESPONSE_TYPE
 )
-from .utils.surreal_get_info import *
-
-import pysurrealdb as surreal
-import pandas as pd
+from .utils.surreal_get_info import table_names, column_info
 
 
 logger = log.getLogger(__name__)

mindsdb/integrations/handlers/tdengine_handler/tdengine_handler.py

@@ -18,6 +18,7 @@ from mindsdb.integrations.libs.response import (
 
 logger = log.getLogger(__name__)
 
+
 class TDEngineHandler(DatabaseHandler):
     """
     This handler handles connection and execution of the TDEngine statements.
@@ -27,7 +28,7 @@ class TDEngineHandler(DatabaseHandler):
 
     def __init__(self, name, connection_data: Optional[dict], **kwargs):
         super().__init__(name)
-        
+
         self.parser = parse_sql
         self.dialect = 'tdengine'
         self.kwargs = kwargs
@@ -36,7 +37,6 @@ class TDEngineHandler(DatabaseHandler):
         self.connection = None
         self.is_connected = False
 
-
     def connect(self):
         if self.is_connected is True:
             return self.connection
@@ -44,8 +44,8 @@ class TDEngineHandler(DatabaseHandler):
         config = {
             'url': self.connection_data.get('url', "http://localhost:6041"),
             'token': self.connection_data.get('token'),
-            'user': self.connection_data.get('user','root'),
-            'password': self.connection_data.get('password','taosdata'),
+            'user': self.connection_data.get('user', 'root'),
+            'password': self.connection_data.get('password', 'taosdata'),
             'database': self.connection_data.get('database')
         }
 
@@ -62,7 +62,7 @@ class TDEngineHandler(DatabaseHandler):
         return
 
     def check_connection(self) -> StatusResponse:
-       
+
         result = StatusResponse(False)
         need_to_close = self.is_connected is False
 
@@ -90,29 +90,29 @@ class TDEngineHandler(DatabaseHandler):
         need_to_close = self.is_connected is False
 
         connection = self.connect()
-        cur = connection.cursor() 
+        cur = connection.cursor()
         try:
-                cur.execute(query)
-                
-                if cur.rowcount != 0:
-                    result = cur.fetchall()
-                    response = Response(
-                        RESPONSE_TYPE.TABLE,
-                        pd.DataFrame(
-                            result,
-                            columns=[x[0] for x in cur.description]
-                        )
-                    )
-                else:
-                    response = Response(RESPONSE_TYPE.OK)
-                connection.commit()
-        except Exception as e:
-                logger.error(f'Error running query: {query} on {self.connection_data["database"]}!')
+            cur.execute(query)
+
+            if cur.rowcount != 0:
+                result = cur.fetchall()
                 response = Response(
-                    RESPONSE_TYPE.ERROR,
-                    error_message=str(e)
+                    RESPONSE_TYPE.TABLE,
+                    pd.DataFrame(
+                        result,
+                        columns=[x[0] for x in cur.description]
+                    )
                 )
-                # connection.rollback()
+            else:
+                response = Response(RESPONSE_TYPE.OK)
+            connection.commit()
+        except Exception as e:
+            logger.error(f'Error running query: {query} on {self.connection_data["database"]}!')
+            response = Response(
+                RESPONSE_TYPE.ERROR,
+                error_message=str(e)
+            )
+            # connection.rollback()
         cur.close()
         if need_to_close is True:
             self.disconnect()
@@ -133,7 +133,6 @@ class TDEngineHandler(DatabaseHandler):
         """
         q = 'SHOW TABLES;'
 
-         
         return self.native_query(q)
 
     def get_columns(self, table_name) -> Response:
@@ -142,5 +141,4 @@ class TDEngineHandler(DatabaseHandler):
         """
         q = f'DESCRIBE {table_name};'
 
-        
         return self.native_query(q)

mindsdb/integrations/handlers/tdengine_handler/tests/test_tdengine_handler.py

@@ -7,9 +7,9 @@ class TDEngineHandlerTest(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
         cls.kwargs = {
-            "token":'<token_for_cloud>',
-            "url":"********.cloud.tdengine.com",
-            "database":"temp"
+            "token": '<token_for_cloud>',
+            "url": "********.cloud.tdengine.com",
+            "database": "temp"
         }
         cls.handler = TDEngineHandler('test_tdengine_handler', connection_data=cls.kwargs)
 
@@ -18,26 +18,26 @@ class TDEngineHandlerTest(unittest.TestCase):
 
     def test_1_connect(self):
         assert self.handler.connect()
-    
+
     def test_2_create_table(self):
         query = "CREATE Table `hari` USING `temp` (`id`) TAGS (0);"
         result = self.handler.query(query)
-        assert result.type is not RESPONSE_TYPE.ERROR 
+        assert result.type is not RESPONSE_TYPE.ERROR
         pass
 
     def test_3_insert(self):
         query = "INSERT INTO hari  VALUES (NOW, 12);"
         result = self.handler.query(query)
-        assert result.type is not RESPONSE_TYPE.ERROR 
+        assert result.type is not RESPONSE_TYPE.ERROR
 
     def test_4_native_query_select(self):
         query = "SELECT * FROM hari;"
         result = self.handler.query(query)
-        assert result.type is RESPONSE_TYPE.TABLE 
+        assert result.type is RESPONSE_TYPE.TABLE
 
     def test_5_get_tables(self):
         tables = self.handler.get_tables()
-        assert tables.type is  RESPONSE_TYPE.TABLE
+        assert tables.type is RESPONSE_TYPE.TABLE
 
     def test_6_get_columns(self):
         columns = self.handler.get_columns('hari')

mindsdb/integrations/handlers/tidb_handler/tests/test_tidb_handler.py

@@ -1,5 +1,4 @@
 import unittest
-import pandas as pd
 from mindsdb.integrations.handlers.tidb_handler.tidb_handler import TiDBHandler
 
 
@@ -42,4 +41,4 @@ class TiDBHandlerTest(unittest.TestCase):
 
     def test_7_select_query(self):
         query = "SELECT * FROM dt_test WHERE 'id'='a'"
-        result = self.handler.native_query(query)
+        self.handler.native_query(query)

mindsdb/integrations/handlers/timegpt_handler/timegpt_handler.py

@@ -33,7 +33,7 @@ class TimeGPTHandler(BaseMLEngine):
         if mode == 'forecasting':
             assert time_settings["is_timeseries"], "Specify time series settings in your query"
 
-        timegpt_token =  get_api_key('timegpt', using_args, self.engine_storage, strict=True)
+        timegpt_token = get_api_key('timegpt', using_args, self.engine_storage, strict=True)
         timegpt = NixtlaClient(api_key=timegpt_token)
         assert timegpt.validate_api_key(), "Invalid TimeGPT token provided."
 
@@ -64,7 +64,7 @@ class TimeGPTHandler(BaseMLEngine):
             model_args['add_history'] = True
 
         assert isinstance(model_args["level"], list), "`level` must be a list of integers"
-        assert all([isinstance(l, int) for l in model_args["level"]]), "`level` must be a list of integers"
+        assert all([isinstance(level, int) for level in model_args["level"]]), "`level` must be a list of integers"
 
         self.model_storage.json_set("model_args", model_args)  # persist changes to handler folder
 
@@ -73,7 +73,7 @@ class TimeGPTHandler(BaseMLEngine):
         model_args = self.model_storage.json_get("model_args")
         args = args['predict_params']
         prediction_df = self._transform_to_nixtla_df(df, model_args)
-        
+
         timegpt = NixtlaClient(api_key=model_args['token'])
         assert timegpt.validate_api_key(), "Invalid TimeGPT token provided."
 
@@ -122,7 +122,7 @@ class TimeGPTHandler(BaseMLEngine):
         for i, level in enumerate(levels):
             if i == 0:
                 # NOTE: this should be simplified once we refactor the expected time series output within MindsDB
-                results_df['confidence'] = level/100  # we report the highest level as the overall confidence
+                results_df['confidence'] = level / 100  # we report the highest level as the overall confidence
                 results_df['lower'] = forecast_df[f'TimeGPT-lo-{level}']
                 results_df['upper'] = forecast_df[f'TimeGPT-hi-{level}']
             else:
@@ -171,7 +171,7 @@ class TimeGPTHandler(BaseMLEngine):
                 if mindate > pd.to_datetime('1970-01-01T00:00:00') and maxdate < pd.to_datetime('2050-12-31T23:59:59'):
                     unit = u
             df[date_column] = pd.to_datetime(df[date_column], unit=unit, origin='unix')
-        else: 
+        else:
             df[date_column] = pd.to_datetime(df[date_column])
         df[date_column] = df[date_column].dt.strftime('%Y-%m-%dT%H:%M:%S')  # convert to ISO 8601 format
         df[date_column] = pd.to_datetime(df[date_column])

mindsdb/integrations/handlers/tpot_handler/tpot_handler.py

@@ -1,4 +1,3 @@
-from typing import Optional
 import dill
 import pandas as pd
 from mindsdb.integrations.libs.base import BaseMLEngine
@@ -10,32 +9,30 @@ from sklearn.preprocessing import LabelEncoder
 
 class TPOTHandler(BaseMLEngine):
     name = "TPOT"
+
     def create(self, target: str, df: Optional[pd.DataFrame] = None, args: Optional[Dict] = None) -> None:
         if args is None:
             args = {}
-        type_of_cols=infer_types(df,0).dtypes
-        target_dtype=type_of_cols[target]
-
+        type_of_cols = infer_types(df, 0).dtypes
+        target_dtype = type_of_cols[target]
 
-        if target_dtype in ['binary','categorical','tags']:
+        if target_dtype in ['binary', 'categorical', 'tags']:
             model = TPOTClassifier(generations=args.get('generations', 10),
-                                       population_size=args.get('population_size', 100),
-                                       verbosity=0,
-                                       max_time_mins=args.get('max_time_mins', None),
-                                       n_jobs=args.get('n_jobs', -1))
-            
+                                   population_size=args.get('population_size', 100),
+                                   verbosity=0,
+                                   max_time_mins=args.get('max_time_mins', None),
+                                   n_jobs=args.get('n_jobs', -1))
 
-        elif target_dtype in ['integer','float','quantity'] :
+        elif target_dtype in ['integer', 'float', 'quantity']:
             model = TPOTRegressor(generations=args.get('generations', 10),
-                                      population_size=args.get('population_size', 100),
-                                      verbosity=0,
-                                      max_time_mins=args.get('max_time_mins', None),
-                                      n_jobs=args.get('n_jobs', -1))
-        
+                                  population_size=args.get('population_size', 100),
+                                  verbosity=0,
+                                  max_time_mins=args.get('max_time_mins', None),
+                                  n_jobs=args.get('n_jobs', -1))
 
         if df is not None:
             # Separate out the categorical and non-categorical columns
-            categorical_cols=[col for col, type_col in type_of_cols.items() if type_col in ('categorical', 'binary')]
+            categorical_cols = [col for col, type_col in type_of_cols.items() if type_col in ('categorical', 'binary')]
 
             # Fit a LabelEncoder for each categorical column and store it in a dictionary
             le_dict = {}
@@ -47,30 +44,28 @@ class TPOTHandler(BaseMLEngine):
                 # Encode the categorical column using the fitted LabelEncoder
                 df[col] = le.transform(df[col])
 
-
             model.fit(df.drop(columns=[target]), df[target])
             self.model_storage.json_set('args', args)
             self.model_storage.file_set('le_dict', dill.dumps(le_dict))
             self.model_storage.file_set('model', dill.dumps(model.fitted_pipeline_))
-        else :
+        else:
             raise Exception(
                 "Data is empty!!"
             )
-        
 
     def predict(self, df: pd.DataFrame, args: Optional[Dict] = None) -> pd.DataFrame:
 
-        model=dill.loads(self.model_storage.file_get("model"))
-        le_dict=dill.loads(self.model_storage.file_get("le_dict"))
-        target=self.model_storage.json_get('args').get("target")
+        model = dill.loads(self.model_storage.file_get("model"))
+        le_dict = dill.loads(self.model_storage.file_get("le_dict"))
+        target = self.model_storage.json_get('args').get("target")
 
-         # Encode the categorical columns in the input DataFrame using the saved LabelEncoders
+        # Encode the categorical columns in the input DataFrame using the saved LabelEncoders
         for col, le in le_dict.items():
             if col in df.columns:
                 df[col] = le.transform(df[col])
-                
+
         # Make predictions using the trained TPOT model
-        results=pd.DataFrame(model.predict(df),columns=[target])
+        results = pd.DataFrame(model.predict(df), columns=[target])
 
         # Decode the predicted categorical values back into their original values
         for col, le in le_dict.items():

mindsdb/integrations/handlers/trino_handler/trino_handler.py

@@ -1,9 +1,9 @@
 import re
 from typing import Dict
 import pandas as pd
-from pyhive import (trino, sqlalchemy_trino)
+from pyhive import sqlalchemy_trino
 from mindsdb_sql_parser import parse_sql, ASTNode
-from trino.auth import KerberosAuthentication, BasicAuthentication
+from trino.auth import BasicAuthentication
 from trino.dbapi import connect
 from mindsdb.utilities.render.sqlalchemy_render import SqlalchemyRender
 from mindsdb.integrations.libs.base import DatabaseHandler
@@ -13,10 +13,10 @@ from mindsdb.integrations.libs.response import (
     HandlerResponse as Response,
     RESPONSE_TYPE
 )
-from .trino_config_provider import TrinoConfigProvider
 
 logger = log.getLogger(__name__)
 
+
 class TrinoHandler(DatabaseHandler):
     """
     This handler handles connection and execution of the Trino statements
@@ -59,24 +59,24 @@ class TrinoHandler(DatabaseHandler):
             return self.connection
 
         # option configuration
-        http_scheme='http'
-        auth=None
-        auth_config=None
-        password=None
+        http_scheme = 'http'
+        auth = None
+        auth_config = None
+        password = None
 
         if 'auth' in self.connection_data:
-            auth=self.connection_data['auth']
+            auth = self.connection_data['auth']
         if 'password' in self.connection_data:
-            password=self.connection_data['password']
+            password = self.connection_data['password']
         if 'http_scheme' in self.connection_data:
-           http_scheme=self.connection_data['http_scheme']
+            http_scheme = self.connection_data['http_scheme']
         if 'with' in self.connection_data:
-           self.with_clause=self.connection_data['with']
+            self.with_clause = self.connection_data['with']
 
-        if password and auth=='kerberos':
+        if password and auth == 'kerberos':
             raise Exception("Kerberos authorization doesn't support password.")
         elif password:
-            auth_config=BasicAuthentication(self.connection_data['user'], password)
+            auth_config = BasicAuthentication(self.connection_data['user'], password)
 
         if auth:
             conn = connect(
@@ -159,7 +159,7 @@ class TrinoHandler(DatabaseHandler):
         renderer = SqlalchemyRender(sqlalchemy_trino.TrinoDialect)
         query_str = renderer.get_string(query, with_failback=True)
         modified_query_str = re.sub(
-            "(?is)(CREATE.+TABLE.+\(.*\))",
+            r"(?is)(CREATE.+TABLE.+\(.*\))",
             f"\\1 {self.with_clause}",
             query_str
         )

mindsdb/integrations/handlers/twitter_handler/twitter_handler.py

@@ -1,10 +1,8 @@
 import re
 import os
 import datetime as dt
-import ast
 import time
 from collections import defaultdict
-import pytz
 import io
 import requests
 
@@ -28,6 +26,7 @@ from mindsdb.integrations.libs.response import (
 
 logger = log.getLogger(__name__)
 
+
 class TweetsTable(APITable):
 
     def select(self, query: ast.Select) -> Response:
@@ -39,7 +38,7 @@ class TweetsTable(APITable):
         for op, arg1, arg2 in conditions:
 
             if op == 'or':
-                raise NotImplementedError(f'OR is not supported')
+                raise NotImplementedError('OR is not supported')
             if arg1 == 'created_at':
                 date = parse_utc_date(arg2)
                 if op == '>':
@@ -170,7 +169,6 @@ class TweetsTable(APITable):
 
                 media_ids = [media.media_id]
 
-
             words = re.split('( )', text)
 
             messages = []

mindsdb/integrations/handlers/unify_handler/tests/test_unify_handler.py

@@ -1,20 +1,19 @@
 import os
 import pytest
-import pandas as pd
 
 from mindsdb_sql_parser import parse_sql
 
 from tests.unit.executor_test_base import BaseExecutorTest
 
+
 class TestUnify(BaseExecutorTest):
-    
+
     def run_sql(self, sql):
         ret = self.command_executor.execute_command(parse_sql(sql))
         assert ret.error_code is None
         if ret.data is not None:
             return ret.data.to_df()
 
-
     """
     Integration tests for the Unify handler.
     """
@@ -25,20 +24,20 @@ class TestUnify(BaseExecutorTest):
         """
         super().setup_method()
         self.run_sql("create database proj")
-    
+
     @pytest.mark.skipif(os.environ.get('UNIFY_API_KEY') is None, reason='Missing API key!')
     def test_unify_correct_flow(self):
         self.run_sql(
-        f"""
+            f"""
         CREATE ML_ENGINE unify_engine
         FROM unify
         USING
             unify_api_key = '{os.environ.get('UNIFY_API_KEY')}';
         """
         )
-        
+
         self.run_sql(
-        f"""
+            """
         CREATE MODEL proj.test_unify_correct_flow
         PREDICT output
         USING
@@ -51,7 +50,7 @@ class TestUnify(BaseExecutorTest):
         self.wait_predictor("proj", "test_unify_correct_flow")
 
         result_df = self.run_sql(
-        """
+            """
             SELECT text, output
             FROM proj.test_unify_correct_flow
             WHERE text = 'Hello';

mindsdb/integrations/handlers/unify_handler/unify_handler.py

@@ -12,6 +12,7 @@ from mindsdb.integrations.utilities.handler_utils import get_api_key
 
 logger = log.getLogger(__name__)
 
+
 class UnifyHandler(BaseMLEngine):
     """
     Integration with the Unifyai Python Library
@@ -26,9 +27,8 @@ class UnifyHandler(BaseMLEngine):
         self.model_storage.json_set('args', args)
 
     def predict(self, df: Optional[pd.DataFrame] = None, args: Optional[Dict] = None) -> None:
-        
+
         args = self.model_storage.json_get('args')
-        
 
         input_keys = list(args.keys())
 
@@ -41,7 +41,7 @@ class UnifyHandler(BaseMLEngine):
         if 'provider' not in args['using']:
             raise Exception("Unify requires a provider parameter in the USING clause! Refer to its documentation for more details.")
         provider = args['using']['provider']
-        
+
         self.endpoint = model + '@' + provider
         self.api_key = get_api_key('unify', args["using"], self.engine_storage, strict=False)
         available_endpoints = unify.utils.list_endpoints(api_key=self.api_key)
@@ -51,15 +51,15 @@ class UnifyHandler(BaseMLEngine):
         question_column = args['using']['column']
         if question_column not in df.columns:
             raise RuntimeError(f'Column "{question_column}" not found in input data')
-        
-        result_df = pd.DataFrame() 
-        result_df['predictions'] = df[question_column].apply(self.predict_answer)     
+
+        result_df = pd.DataFrame()
+        result_df['predictions'] = df[question_column].apply(self.predict_answer)
         result_df = result_df.rename(columns={'predictions': args['target']})
-        
+
         return result_df
-    
+
     def predict_answer(self, text):
-        
+
         client = unify.Unify(self.endpoint, api_key=self.api_key)
 
         response = client.generate(text)

mindsdb/integrations/handlers/vertex_handler/vertex_client.py

@@ -2,7 +2,7 @@ from mindsdb.utilities import log
 from google.cloud.aiplatform import init, TabularDataset, Model, Endpoint
 import pandas as pd
 
-from mindsdb.integrations.utilities.handlers.auth_utilities import GoogleServiceAccountOAuth2Manager
+from mindsdb.integrations.utilities.handlers.auth_utilities.google import GoogleServiceAccountOAuth2Manager
 
 logger = log.getLogger(__name__)