MindsDB 1.2.9__py3-none-any.whl → 24.12.4.0__py3-none-any.whl

mindsdb/integrations/handlers/lightwood_handler/tests/test_lightwood_handler.py

@@ -0,0 +1,258 @@
+import os
+import time
+import unittest
+import tempfile
+
+temp_dir = tempfile.mkdtemp(dir='/tmp/', prefix='lightwood_handler_test_')
+os.environ['MINDSDB_STORAGE_DIR'] = os.environ.get('MINDSDB_STORAGE_DIR', temp_dir)
+os.environ['MINDSDB_DB_CON'] = 'sqlite:///' + os.path.join(os.environ['MINDSDB_STORAGE_DIR'], 'mindsdb.sqlite3.db') + '?check_same_thread=False&timeout=30'
+
+from mindsdb.migrations import migrate
+migrate.migrate_to_head()
+
+from mindsdb.utilities.config import Config
+from mindsdb.integrations.utilities.test_utils import HandlerControllerMock, PG_HANDLER_NAME, PG_CONNECTION_DATA
+from mindsdb.interfaces.database.integrations import integration_controller
+from mindsdb.integrations.handlers.lightwood_handler.lightwood_handler.lightwood_handler import LightwoodHandler
+# from mindsdb.integrations.handlers.lightwood_handler.lightwood_handler.utils import load_predictor
+from mindsdb.interfaces.storage.fs import FsStore
+from mindsdb.interfaces.model.model_controller import ModelController
+from mindsdb.integrations.libs.response import RESPONSE_TYPE
+import mindsdb.interfaces.storage.db as db
+
+
+# TODO: drop all models and tables when closing tests
+class LightwoodHandlerTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        # region create permanent integrations
+        for integration_name in ['files', 'lightwood']:
+            integration_record = db.Integration(
+                name=integration_name,
+                data={},
+                engine=integration_name,
+                company_id=None
+            )
+            db.session.add(integration_record)
+            db.session.commit()
+        integration_record = db.Integration(
+            name=PG_HANDLER_NAME,
+            data=PG_CONNECTION_DATA,
+            engine='postgres',
+            company_id=None
+        )
+        db.session.add(integration_record)
+        db.session.commit()
+        # endregion
+
+        handler_controller = integration_controller
+
+        cls.handler = LightwoodHandler(
+            'lightwood',
+            handler_controller=handler_controller,
+            fs_store=FsStore(),
+            model_controller=ModelController()
+        )
+        cls.config = Config()
+
+        cls.target_1 = 'rental_price'
+        cls.data_table_1 = 'demo_data.home_rentals'
+        cls.test_model_1 = 'test_lightwood_home_rentals'
+        cls.test_model_1b = 'test_lightwood_home_rentals_custom'
+
+        cls.target_2 = 'Traffic'
+        cls.data_table_2 = 'demo_data.house_sales'
+        cls.test_model_2 = 'test_lightwood_house_sales'
+
+    def test_00_check_connection(self):
+        conn = self.handler.check_connection()
+        assert conn.success
+
+    def test_01_drop_predictor(self):
+        if self.test_model_1 not in self.handler.get_tables().data_frame.values:
+            # TODO: seems redundant because of test_02
+            query = f"""
+                CREATE PREDICTOR {self.test_model_1}
+                FROM {PG_HANDLER_NAME} (SELECT * FROM {self.data_table_1} limit 50)
+                PREDICT rental_price
+            """
+            self.handler.native_query(query)
+        response = self.handler.native_query(f"DROP PREDICTOR {self.test_model_1}")
+        self.assertTrue(response.type == RESPONSE_TYPE.OK)
+
+    def test_02_train_predictor(self):
+        query = f"""
+            CREATE PREDICTOR {self.test_model_1}
+            FROM {PG_HANDLER_NAME} (SELECT * FROM {self.data_table_1} limit 50)
+            PREDICT rental_price
+        """
+        response = self.handler.native_query(query)
+        time.sleep(5)
+        self.assertTrue(response.type == RESPONSE_TYPE.OK)
+
+    def test_03_retrain_predictor(self):
+        query = f"RETRAIN {self.test_model_1}"
+        response = self.handler.native_query(query)
+        self.assertTrue(response.type == RESPONSE_TYPE.OK)
+
+    def test_04_query_predictor_single_where_condition(self):
+        time.sleep(120) # TODO 
+        query = f"""
+            SELECT target
+            from {self.test_model_1}
+            WHERE sqft=100
+        """
+        response = self.handler.native_query(query)
+        self.assertTrue(response.type == RESPONSE_TYPE.TABLE)
+        self.assertTrue(len(response.data_frame) == 1)
+        self.assertTrue(response.data_frame['sqft'][0] == 100)
+        self.assertTrue(response.data_frame['rental_price'][0] is not None)
+
+    def test_05_query_predictor_multi_where_condition(self):
+        query = f"""
+            SELECT target
+            from {self.test_model_1}
+            WHERE sqft=100
+                  AND number_of_rooms=2
+                  AND number_of_bathrooms=1
+        """
+        response = self.handler.native_query(query)
+        self.assertTrue(response.type == RESPONSE_TYPE.TABLE)
+        self.assertTrue(len(response.data_frame) == 1)
+        self.assertTrue(response.data_frame['number_of_rooms'][0] == 2)
+        self.assertTrue(response.data_frame['number_of_bathrooms'][0] == 1)
+
+    def test_06_train_predictor_custom_jsonai(self):
+        # TODO: turn this into a decorator?
+        if self.test_model_1b in self.handler.get_tables().data_frame.values:  # TODO this accesor feels weird, maybe rethink output format?
+            self.handler.native_query(f"DROP PREDICTOR {self.test_model_1b}")
+
+        using_str = 'model.args={"submodels": [{"module": "LightGBM", "args": {"stop_after": 12, "fit_on_dev": true}}]}'
+        query = f"""
+            CREATE PREDICTOR {self.test_model_1b}
+            FROM {PG_HANDLER_NAME} (SELECT * FROM {self.data_table_1} limit 50)
+            PREDICT rental_price
+            USING {using_str}
+        """
+        response = self.handler.native_query(query)
+        self.assertTrue(response.type == RESPONSE_TYPE.OK)
+        # TODO assert
+        # m = load_predictor(self.handler.storage.get('models')[self.test_model_1b], self.test_model_1b)
+        # assert len(m.ensemble.mixers) == 1
+        # assert type(m.ensemble.mixers[0]).__name__ == 'LightGBM'
+
+    def test_07_list_tables(self):
+        response = self.handler.get_tables()
+        self.assertTrue(response.type == RESPONSE_TYPE.TABLE)
+
+    def test_08_get_columns(self):
+        response = self.handler.get_columns(f'{self.test_model_1}')
+        self.assertTrue(response.type == RESPONSE_TYPE.TABLE)
+
+    # TODO
+    # def test_09_join_predictor_into_table(self):
+    #     into_table = 'test_join_into_lw'
+    #     query = f"SELECT tb.{self.target_1} as predicted, ta.{self.target_1} as truth, ta.sqft from {PG_HANDLER_NAME}.{self.data_table_1} AS ta JOIN {self.test_model_1} AS tb LIMIT 10"
+    #     parsed = self.handler.parser(query, dialect=self.handler.dialect)
+    #     predicted = self.handler.join(parsed, self.data_handler, into=into_table)
+
+    #     # checks whether `into` kwarg does insert into the table or not
+    #     q = f"SELECT * FROM {into_table}"
+    #     qp = self.handler.parser(q, dialect='mysql')
+    #     assert len(self.data_handler.query(qp).data_frame) > 0
+
+    def test_10_train_ts_predictor_multigby_hor4(self):
+        # TODO: handle cap/uncapped column name returned from data handler? Had to rename 'MA' -> 'ma' for test to pass
+        query = f"""
+            CREATE PREDICTOR {self.test_model_2}
+            FROM {PG_HANDLER_NAME} (SELECT * FROM {self.data_table_2})
+            PREDICT ma
+            ORDER BY saledate
+            GROUP BY bedrooms, type
+            WINDOW 8
+            HORIZON 4
+        """
+        if self.test_model_2 not in self.handler.get_tables().data_frame.values:
+            response = self.handler.native_query(query)
+        else:
+            self.handler.native_query(f"DROP PREDICTOR {self.test_model_2}")
+            response = self.handler.native_query(query)
+
+        self.assertTrue(response.type == RESPONSE_TYPE.OK)
+
+        # TODO: reactivate and add to the rest of the TS tests once cache is back on
+        # p = self.handler.storage.get('models')
+        # m = load_predictor(p[self.test_model_2], self.test_model_2)
+        # assert m.problem_definition.timeseries_settings.is_timeseries
+
+    def test_12_train_ts_predictor_multigby_hor1(self):
+        query = f"""
+            CREATE PREDICTOR {self.test_model_2}
+            FROM {PG_HANDLER_NAME} (SELECT * FROM {self.data_table_2})
+            PREDICT ma
+            ORDER BY saledate
+            GROUP BY bedrooms, type
+            WINDOW 8
+            HORIZON 1
+        """
+        if self.test_model_2 not in self.handler.get_tables().data_frame.values:
+            self.handler.native_query(query)
+        else:
+            self.handler.native_query(f"DROP PREDICTOR {self.test_model_2}")
+            self.handler.native_query(query)
+
+    def test_13_train_ts_predictor_no_gby_hor1(self):
+        query = f"""
+            CREATE PREDICTOR {self.test_model_2}
+            FROM {PG_HANDLER_NAME} (SELECT * FROM {self.data_table_2})
+            PREDICT ma
+            ORDER BY saledate
+            WINDOW 8
+            HORIZON 1
+        """
+        if self.test_model_2 not in self.handler.get_tables().data_frame.values:
+            self.handler.native_query(query)
+        else:
+            self.handler.native_query(f"DROP PREDICTOR {self.test_model_2}")
+            self.handler.native_query(query)
+
+    def test_14_train_ts_predictor_no_gby_hor4(self):
+        query = f"""
+            CREATE PREDICTOR {self.test_model_2}
+            FROM {PG_HANDLER_NAME} (SELECT * FROM {self.data_table_2})
+            PREDICT ma
+            ORDER BY saledate
+            WINDOW 8
+            HORIZON 4
+        """
+        if self.test_model_2 not in self.handler.get_tables().data_frame.values:
+            self.handler.native_query(query)
+        else:
+            self.handler.native_query(f"DROP PREDICTOR {self.test_model_2}")
+            self.handler.native_query(query)
+
+    # TODO
+    # def test_15_join_predictor_ts_into(self):
+    #     query = f"""
+    #         SELECT m.saledate as date,
+    #                m.ma as forecast
+    #         FROM mindsdb.{self.test_model_2} m JOIN {PG_HANDLER_NAME}.demo_data.house_sales t
+    #         WHERE t.saledate > LATEST
+    #               AND t.type = 'house'
+    #               AND t.bedrooms = 2
+    #         LIMIT 10
+    #     """
+    #     response = self.handler.native_query(query)
+    #     self.assertTrue(response.type == RESPONSE_TYPE.TABLE)
+
+    # def test_16_join_predictor_ts_model_left(self):
+    #     # TODO: is this one needed?
+    #     target = 'Traffic'
+    #     oby = 'T'
+    #     query = f"SELECT tb.{target} as predicted, ta.{target} as truth, ta.{oby} from mindsdb.{self.test_tsmodel_name_1} AS tb JOIN {self.sql_handler_name}.{self.data_table_name_2} AS ta ON 1=1 WHERE ta.{oby} > LATEST LIMIT 10"
+    #     parsed = self.handler.parser(query, dialect=self.handler.dialect)
+    #     predicted = self.handler.join(parsed, self.data_handler)  # , into=self.model_2_into_table) # TODO: restore when we add support for SQLite and other handlers for `into`
+
+
+if __name__ == "__main__":
+    unittest.main(failfast=True)

mindsdb/integrations/handlers/lightwood_handler/utils.py

@@ -0,0 +1,76 @@
+import gc
+import json
+import sys
+
+import dill
+from lightwood.api.high_level import _module_from_code
+
+
+def unpack_jsonai_old_args(json_ai_override):
+    while '.' in str(list(json_ai_override.keys())):
+        for k in list(json_ai_override.keys()):
+            if '.' in k:
+                nks = k.split('.')
+                obj = json_ai_override
+                for nk in nks[:-1]:
+                    if nk not in obj:
+                        obj[nk] = {}
+                    obj = obj[nk]
+                obj[nks[-1]] = json_ai_override[k]
+                del json_ai_override[k]
+
+
+def load_predictor(predictor_dict, name):
+    try:
+        return dill.loads(predictor_dict['predictor'])
+    except Exception as e:
+        module_name = str(e).lstrip("No module named '").split("'")[0]
+
+        try:
+            del sys.modules[module_name]
+        except Exception:
+            pass
+
+        gc.collect()
+        _module_from_code(predictor_dict['code'], module_name)
+        return dill.loads(predictor_dict['predictor'])
+
+
+def rep_recur(org: dict, ovr: dict):
+    for k in ovr:
+        if k in org:
+            if isinstance(org[k], dict) and isinstance(ovr[k], dict):
+                rep_recur(org[k], ovr[k])
+            else:
+                org[k] = ovr[k]
+        else:
+            org[k] = ovr[k]
+
+
+def brack_to_mod(ovr):
+    if not isinstance(ovr, dict):
+        if isinstance(ovr, list):
+            for i in range(len(ovr)):
+                ovr[i] = brack_to_mod(ovr[i])
+        elif isinstance(ovr, str):
+            if '(' in ovr and ')' in ovr:
+                mod = ovr.split('(')[0]
+                args = {}
+                if '()' not in ovr:
+                    for str_pair in ovr.split('(')[1].split(')')[0].split(','):
+                        k = str_pair.split('=')[0].strip(' ')
+                        v = str_pair.split('=')[1].strip(' ')
+                        args[k] = v
+
+                ovr = {'module': mod, 'args': args}
+            elif '{' in ovr and '}' in ovr:
+                try:
+                    ovr = json.loads(ovr)
+                except Exception:
+                    pass
+        return ovr
+    else:
+        for k in ovr.keys():
+            ovr[k] = brack_to_mod(ovr[k])
+
+    return ovr

mindsdb/integrations/handlers/lindorm_handler/__about__.py

@@ -0,0 +1,9 @@
+__title__ = "MindsDB Lindorm handler"
+__package_name__ = "mindsdb_lindorm_handler"
+__version__ = "0.0.1"
+__description__ = "MindsDB handler for Lindorm"
+__author__ = "Someshfengde"
+__github__ = "https://github.com/someshfengde"
+__pypi__ = "https://pypi.org/project/mindsdb/"
+__license__ = "MIT"
+__copyright__ = "Copyright 2023 - mindsdb"

mindsdb/integrations/handlers/lindorm_handler/__init__.py

@@ -0,0 +1,28 @@
+from mindsdb.integrations.libs.const import HANDLER_TYPE
+
+from .__about__ import __description__ as description
+from .__about__ import __version__ as version
+
+try:
+    from .lindorm_handler import LindormDBHandler as Handler
+
+    import_error = None
+except Exception as e:
+    Handler = None
+    import_error = e
+
+title = "Lindorm"
+name = "lindorm"
+type = HANDLER_TYPE.DATA
+icon_path = 'icon.svg'
+
+__all__ = [
+    "Handler",
+    "version",
+    "name",
+    "type",
+    "title",
+    "description",
+    "import_error",
+    "icon_path"
+]

mindsdb/integrations/handlers/lindorm_handler/icon.svg

@@ -0,0 +1,5 @@
+<svg xmlns="http://www.w3.org/2000/svg" width="40" height="40" viewBox="0 0 40 40" fill="none">
+  <path d="M26.859 19.0325H13.1753V22.1168H26.859V19.0325Z" fill="black"/>
+  <path d="M33.1649 7.83557H24.1111L26.2956 10.9268L32.897 12.9876C33.4797 13.1716 33.9878 13.5379 34.3465 14.0327C34.7052 14.5274 34.8954 15.1243 34.8891 15.7353V25.5859C34.8932 26.1964 34.7021 26.7922 34.3438 27.2865C33.9854 27.7808 33.4785 28.1477 32.897 28.3336L26.2956 30.3532L24.1111 33.4443H33.1649C34.0656 33.4453 34.9576 33.268 35.7896 32.9227C36.6215 32.5773 37.3768 32.0708 38.0122 31.4323C38.6475 30.7938 39.1502 30.0359 39.4913 29.2023C39.8325 28.3687 40.0053 27.4758 39.9999 26.575V14.7049C39.9999 12.889 39.2809 11.1469 38 9.85965C36.7192 8.57239 34.9808 7.84465 33.1649 7.83557Z" fill="black"/>
+  <path d="M6.86931 7.83557H15.8956L13.7386 10.9268L7.13721 12.9876C6.55382 13.1713 6.04479 13.5374 5.68497 14.032C5.32514 14.5266 5.13351 15.1236 5.13824 15.7353V25.5859C5.13567 26.197 5.32814 26.793 5.68766 27.2871C6.04717 27.7813 6.55498 28.1479 7.13721 28.3336L13.7386 30.3532L15.9231 33.4443H6.86931C5.04745 33.4443 3.30022 32.7206 2.01197 31.4324C0.723728 30.1441 0 28.3969 0 26.575L0 14.7049C0 12.883 0.723728 11.1358 2.01197 9.84754C3.30022 8.5593 5.04745 7.83557 6.86931 7.83557Z" fill="black"/>
+</svg>

mindsdb/integrations/handlers/lindorm_handler/lindorm_handler.py

@@ -0,0 +1,220 @@
+from typing import Optional
+import pandas as pd
+import phoenixdb
+from mindsdb_sql_parser import parse_sql
+from mindsdb.utilities.render.sqlalchemy_render import SqlalchemyRender
+from mindsdb.integrations.libs.base import DatabaseHandler
+from pyphoenix.sqlalchemy_phoenix import PhoenixDialect
+from mindsdb_sql_parser.ast.base import ASTNode
+from mindsdb.utilities import log
+from mindsdb.integrations.libs.response import (
+    HandlerStatusResponse as StatusResponse,
+    HandlerResponse as Response,
+    RESPONSE_TYPE
+)
+
+logger = log.getLogger(__name__)
+
+
+class LindormHandler(DatabaseHandler):
+    """
+    This handler handles connection and execution of the Apache Phoenix statements.
+    """
+
+    name = 'lindorm'
+
+    def __init__(self, name: str, connection_data: Optional[dict], **kwargs):
+        """
+        Initialize the handler.
+        Args:
+            name (str): name of particular handler instance
+            connection_data (dict): parameters for connecting to the database
+            **kwargs: arbitrary keyword arguments.
+        """
+        super().__init__(name)
+        self.parser = parse_sql
+        self.dialect = 'phoenix'
+        optional_parameters = ['autocommit', 'lindorm_user', 'lindorm_password']
+        for parameter in optional_parameters:
+            if parameter not in connection_data:
+                connection_data[parameter] = None
+
+        self.connection_data = connection_data
+        self.kwargs = kwargs
+
+        self.connection = None
+        self.is_connected = False
+
+    def __del__(self):
+        if self.is_connected is True:
+            self.disconnect()
+
+    def connect(self) -> StatusResponse:
+        """
+        Set up the connection required by the handler.
+        Returns:
+            HandlerStatusResponse
+        """
+
+        if self.is_connected is True:
+            return self.connection
+
+        lindorm_connection_data = {'lindorm_user': self.connection_data['lindorm_user'], 'lindorm_password': self.connection_data['lindorm_password']}
+
+        self.connection = phoenixdb.connect(
+            url=self.connection_data['url'],
+            autocommit=self.connection_data['autocommit'],
+            **lindorm_connection_data
+        )
+        self.is_connected = True
+
+        return self.connection
+
+    def disconnect(self):
+        """ Close any existing connections
+
+        Should switch self.is_connected.
+        """
+
+        if self.is_connected is False:
+            return
+
+        self.connection.close()
+        self.is_connected = False
+        return self.is_connected
+
+    def check_connection(self) -> StatusResponse:
+        """
+        Check connection to the handler.
+        Returns:
+            HandlerStatusResponse
+        """
+
+        response = StatusResponse(False)
+        need_to_close = self.is_connected is False
+
+        try:
+            self.connect()
+            response.success = True
+        except Exception as e:
+            logger.error(f'Error connecting to the Phoenix Query Server, {e}!')
+            response.error_message = str(e)
+        finally:
+            if response.success is True and need_to_close:
+                self.disconnect()
+            if response.success is False and self.is_connected is True:
+                self.is_connected = False
+
+        return response
+
+    def native_query(self, query: str) -> StatusResponse:
+        """
+        Receive raw query and act upon it somehow.
+        Args:
+            query (str): query in native format
+        Returns:
+            HandlerResponse
+        """
+
+        need_to_close = self.is_connected is False
+
+        connection = self.connect()
+        cursor = connection.cursor()
+
+        try:
+            cursor.execute(query)
+            result = cursor.fetchall()
+            if result:
+                response = Response(
+                    RESPONSE_TYPE.TABLE,
+                    data_frame=pd.DataFrame(
+                        result,
+                        columns=[x[0] for x in cursor.description]
+                    )
+                )
+            else:
+                connection.commit()
+                response = Response(RESPONSE_TYPE.OK)
+        except Exception as e:
+            logger.error(f'Error running query: {query} on the Lindorm Query Server!')
+            response = Response(
+                RESPONSE_TYPE.ERROR,
+                error_message=str(e)
+            )
+
+        cursor.close()
+        if need_to_close is True:
+            self.disconnect()
+
+        return response
+
+    def query(self, query: ASTNode) -> StatusResponse:
+        """
+        Receive query as AST (abstract syntax tree) and act upon it somehow.
+        Args:
+            query (ASTNode): sql query represented as AST. May be any kind
+                of query: SELECT, INTSERT, DELETE, etc
+        Returns:
+            HandlerResponse
+        """
+
+        renderer = SqlalchemyRender(PhoenixDialect)
+        query_str = renderer.get_string(query, with_failback=True)
+        return self.native_query(query_str)
+
+    def get_tables(self) -> StatusResponse:
+        """
+        Return list of entities that will be accessible as tables.
+        Returns:
+            HandlerResponse
+        """
+
+        query = """
+            SELECT DISTINCT TABLE_NAME, TABLE_SCHEM FROM SYSTEM.CATALOG
+        """
+        result = self.native_query(query)
+        df = result.data_frame
+        df = df[df['TABLE_SCHEM'] != 'SYSTEM']
+        df = df.drop('TABLE_SCHEM', axis=1)
+        result.data_frame = df.rename(columns={df.columns[0]: 'table_name'})
+        return result
+
+    def get_columns(self, table_name: str) -> StatusResponse:
+        """
+        Returns a list of entity columns.
+        Args:
+            table_name (str): name of one of tables returned by self.get_tables()
+        Returns:
+            HandlerResponse
+        """
+
+        need_to_close = self.is_connected is False
+
+        connection = self.connect()
+        cursor = connection.cursor()
+
+        try:
+            query = f"SELECT * from {table_name} LIMIT 5"
+            cursor.execute(query)
+            cursor.fetchall()
+
+            response = Response(
+                RESPONSE_TYPE.TABLE,
+                data_frame=pd.DataFrame(
+                    [(x[0], x[1]) for x in cursor.description],
+                    columns=['column_name', 'data_type']
+                )
+            )
+
+        except Exception as e:
+            logger.error(f'Error running query: {query} on the Phoenix Query Server!')
+            response = Response(
+                RESPONSE_TYPE.ERROR,
+                error_message=str(e)
+            )
+
+        cursor.close()
+        if need_to_close is True:
+            self.disconnect()
+
+        return response

mindsdb/integrations/handlers/lindorm_handler/requirements.txt

@@ -0,0 +1,2 @@
+pyphoenix
+phoenixdb

mindsdb/integrations/handlers/lindorm_handler/tests/test_lindorm_handler.py

@@ -0,0 +1,33 @@
+import unittest
+from mindsdb.integrations.handlers.lindorm_handler.lindorm_handler import LindormHandler
+from mindsdb.api.executor.data_types.response_type import RESPONSE_TYPE
+
+
+class LindormHandlerTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.kwargs = {
+            "url": 'http://127.0.0.1:8765',
+            "autocommit": True
+        }
+        cls.handler = LindormHandler('test_phoenix_handler', cls.kwargs)
+
+    def test_0_check_connection(self):
+        assert self.handler.check_connection()
+
+    def test_1_native_query_select(self):
+        query = "SELECT * FROM USERS"
+        result = self.handler.native_query(query)
+        assert result.type is RESPONSE_TYPE.TABLE
+
+    def test_2_get_tables(self):
+        tables = self.handler.get_tables()
+        assert tables.type is not RESPONSE_TYPE.ERROR
+
+    def test_3_get_columns(self):
+        columns = self.handler.get_columns('USERS')
+        assert columns.type is not RESPONSE_TYPE.ERROR
+
+
+if __name__ == '__main__':
+    unittest.main()

mindsdb/integrations/handlers/litellm_handler/__about__.py

@@ -0,0 +1,9 @@
+__title__ = "MindsDB LiteLLM Handler"
+__package_name__ = "mindsdb_litellm_handler"
+__version__ = "0.0.1"
+__description__ = "MindsDB handler for LiteLLM"
+__author__ = "Daniel Usvyat"
+__github__ = "https://github.com/mindsdb/mindsdb"
+__pypi__ = "https://pypi.org/project/mindsdb/"
+__license__ = "MIT"
+__copyright__ = "Copyright 2023- mindsdb"

mindsdb/integrations/handlers/litellm_handler/__init__.py

@@ -0,0 +1,30 @@
+from mindsdb.integrations.libs.const import HANDLER_TYPE
+
+from .__about__ import __description__ as description
+from .__about__ import __version__ as version
+
+try:
+    from .litellm_handler import LiteLLMHandler as Handler
+
+    import_error = None
+except Exception as e:
+    Handler = None
+    import_error = e
+
+title = "LiteLLM"
+name = "litellm"
+type = HANDLER_TYPE.ML
+icon_path = "icon.png"
+permanent = False
+
+
+__all__ = [
+    "Handler",
+    "version",
+    "name",
+    "type",
+    "title",
+    "description",
+    "import_error",
+    "icon_path",
+]