MindsDB 1.3.1__py3-none-any.whl → 25.1.2.0__py3-none-any.whl

mindsdb/integrations/handlers/leonardoai_handler/leonardo_ai_handler.py

@@ -0,0 +1,225 @@
+import os
+import requests
+import time
+from typing import Dict, Optional
+
+import pandas as pd
+
+from mindsdb.integrations.libs.base import BaseMLEngine
+from mindsdb.integrations.libs.llm.utils import get_completed_prompts
+from mindsdb.utilities import log
+from mindsdb.interfaces.storage.model_fs import HandlerStorage
+from mindsdb.utilities.config import Config
+
+logger = log.getLogger(__name__)
+
+LEONARDO_API_BASE = 'https://cloud.leonardo.ai/api/rest/v1'
+
+
+class LeonardoAIHandler(BaseMLEngine):
+    """
+    This integration seamlessly combines MindsDB and Leonardo AI to create a powerful
+    AI-driven solution for creative content generation.
+
+    Content Generation with Leonardo AI: Harness the power of advanced generative
+    models for creative content production. From realistic images to artistic text, Leonardo
+    AI opens up new possibilities for content creators.
+    """
+    name = "leonardo_ai"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.all_models = []
+        self.default_model = '6bef9f1b-29cb-40c7-b9df-32b51c1f67d3'
+        self.base_api = LEONARDO_API_BASE
+        self.rate_limit = 50
+        self.max_batch_size = 5  # default value
+
+    def create(self, target: str, args=None, **kwargs):
+
+        args['target'] = target
+
+        if "using" not in args:
+            raise Exception(
+                "Leornardo Engine requires a USING clause!"
+            )
+
+        self.model_storage.json_set("args", args)
+        api_key = self._get_leonardo_api_key(args, self.engine_storage)  # fetch api key
+
+        # check if API key is valid
+        self.connection = requests.get(
+            "https://cloud.leonardo.ai/api/rest/v1/me",
+            headers={
+                "accept": "application/json",
+                "authorization": f"Bearer {api_key}"
+            }
+        )
+
+        # if valid, check if the model is valid
+        try:
+            if (self.connection.status_code == 200):
+                # get all the available models
+                available_models = self._get_platform_model(args)
+                if not args['using']['model']:
+                    args['using']['model'] = self.default_model
+                elif args['using']['model'] not in available_models:    # if invalid model_id is provided
+                    raise Exception("Invalid Model ID. Please use a valid Model")
+            else:
+                raise Exception("Unable to make connection, please verify the API key.")
+        except Exception:
+            raise Exception("Auth Connection Error, please the modelId or API key")
+
+    def predict(self, df: pd.DataFrame, args: Optional[Dict] = None, **kwargs) -> pd.DataFrame:
+
+        pred_args = args['predict_params'] if args else {}
+        args = self.model_storage.json_get("args")
+
+        prompt_template = pred_args.get('prompt_template', args.get('prompt_template', 'Generate a picture of {{{{text}}}}'))
+
+        # prepare prompts
+        prompts, empty_prompt_id = get_completed_prompts(prompt_template, df)
+        df['__mdb_prompt'] = prompts
+
+        # generate picture based on the given prompt
+        result_df = pd.DataFrame(self.predict_answer(prompts))
+        return result_df
+
+    def _get_leonardo_api_key(self, args, engine_storage: HandlerStorage, strict=True):
+        """
+        API_KEY preference order:
+            1. provided at model creation
+            2. provided at engine creation
+            3. LEONARDO_API_KEY env variable
+            4. leonardo.api_key setting in config.json
+        """
+
+        # 1
+        if "api_key" in args["using"]:
+            return args["using"]["api_key"]
+        # 2
+        connection_args = engine_storage.get_connection_args()
+        if "api_key" in connection_args:
+            return connection_args["api_key"]
+        # 3
+        api_key = os.getenv("LEONARDO_API_KEY")
+        if api_key is not None:
+            return api_key
+        # 4
+        config = Config()
+        leonardo_cfg = config.get("leonardo", {})
+        if "api_key" in leonardo_cfg:
+            return leonardo_cfg["api_key"]
+
+        if strict:
+            raise Exception(
+                'Missing API key "api_key". Either re-create this ML_ENGINE specifying the `api_key` parameter,\
+                 or re-create this model and pass the API key with `USING` syntax.'
+            )
+
+    def _get_platform_model(self, args):
+        """
+        Returns a list of available model based on the API key provided
+        """
+        model_ids = []
+
+        args = self.model_storage.json_get('args')
+        api_key = self._get_leonardo_api_key(args, self.engine_storage)
+
+        self.connection = requests.get(
+            "https://cloud.leonardo.ai/api/rest/v1/platformModels",
+            headers={
+                "accept": "application/json",
+                "authorization": f"Bearer {api_key}"
+            }
+        )
+
+        models = self.connection.json()
+
+        # extract the model ids from the response
+        model_ids = [model['id'] for model in models['custom_models']]
+
+        return model_ids
+
+    def predict_answer(self, prompts, **kwargs):
+        """
+        Generates pictures based on the prompts and returns URLs with few variations.
+
+        Request Flow:
+            - POST request with a prompt is sent
+            - `generation_id` is created for the request
+            - POST request will take couple of seconds to generate the picture, till then the process will be kept busy with a simple math calculation.
+            - New GET request with the `generation_id` will fetch the generated pictures as URLs
+        """
+        args = self.model_storage.json_get('args')
+        tmp_args = args['using']
+        height = tmp_args.get('height', 512)
+        width = tmp_args.get('width', 512)
+        api_key = self._get_leonardo_api_key(args, self.engine_storage)  # fetch API key
+        generation_id = ''
+
+        # Endpoint URL
+        generation_url = "https://cloud.leonardo.ai/api/rest/v1/generations"
+
+        post_headers = {
+            "accept": "application/json",
+            "content-type": "application/json",
+            "authorization": f"Bearer {api_key}"
+        }
+
+        get_headers = {
+            "accept": "application/json",
+            "authorization": f"Bearer {api_key}"
+        }
+
+        # payload
+        generation_payload = {
+            "height": height,
+            "modelId": args['using']['model'],
+            "prompt": f"{prompts}",
+            "width": width,
+        }
+
+        # Make a POST request to generate the image
+        response_generation = requests.post(generation_url, json=generation_payload, headers=post_headers)
+        generation_data = response_generation.json()
+
+        # Wait for 15 seconds
+
+        time.sleep(15)
+
+        # extract generationID from the response
+        generation_id = generation_data['sdGenerationJob']['generationId']
+
+        # ENDPOINT GET URL
+        retrieve_url = f"https://cloud.leonardo.ai/api/rest/v1/generations/{generation_id}"
+
+        # GET request to retrieve image URLs
+        response_retrieve = requests.get(retrieve_url, headers=get_headers)
+        retrieve_data = response_retrieve.json()
+
+        # extract URLs from the response
+        generated_images = retrieve_data["generations_by_pk"]["generated_images"]
+        image_urls = [image["url"] for image in generated_images]
+
+        url_dicts = []
+
+        for url in image_urls:
+            url_dicts.append({'url': url})
+
+        img_urls = pd.DataFrame(url_dicts, columns=['url'])
+        return img_urls
+
+    def describe(self, attribute: Optional[str] = None) -> pd.DataFrame:
+        args = self.model_storage.json_get('args')
+        model, target = args['using']['model'], args['target']
+        prompt_template = args.get('prompt_template', 'Generate a picture of {{{{text}}}}')
+
+        if attribute == "features":
+            return pd.DataFrame([[target, prompt_template]], columns=['target_column', 'mindsdb_prompt_template'])
+        elif attribute == "metadata":
+            api_key = self._get_leonardo_api_key(args, self.engine_storage)
+            return pd.DataFrame([[target, api_key, model]], columns=['target', 'api_key', 'model_name'])
+        else:
+            tables = ['args', 'api_key']
+            return pd.DataFrame(tables, columns=['tables'])

mindsdb/integrations/handlers/libsql_handler/__about__.py

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

mindsdb/integrations/handlers/libsql_handler/__init__.py

@@ -0,0 +1,28 @@
+from mindsdb.integrations.libs.const import HANDLER_TYPE
+
+from .__about__ import __version__ as version, __description__ as description
+from .connection_args import connection_args, connection_args_example
+try:
+    from .libsql_handler import LibSQLHandler as Handler
+    import_error = None
+except Exception as e:
+    Handler = None
+    import_error = e
+
+title = "LibSQL"
+name = "libsql"
+type = HANDLER_TYPE.DATA
+icon_path = "icon.svg"
+
+__all__ = [
+    "Handler",
+    "version",
+    "name",
+    "type",
+    "title",
+    "description",
+    "connection_args",
+    "connection_args_example",
+    "import_error",
+    "icon_path",
+]

mindsdb/integrations/handlers/libsql_handler/connection_args.py

@@ -0,0 +1,22 @@
+from collections import OrderedDict
+
+from mindsdb.integrations.libs.const import HANDLER_CONNECTION_ARG_TYPE as ARG_TYPE
+
+
+connection_args = OrderedDict(
+    database={
+        "type": ARG_TYPE.STR,
+        "description": "The database file where the data will be stored.",
+    },
+    sync_url={
+        "type": ARG_TYPE.STR,
+        "description": "The database URL where the data is synced with.",
+    },
+    auth_token={
+        "type": ARG_TYPE.STR,
+        "description": "The JWT auth token to authenticate with the sync database.",
+        "secret": True
+    },
+)
+
+connection_args_example = OrderedDict(database="chinook.db")

mindsdb/integrations/handlers/libsql_handler/icon.svg

@@ -0,0 +1,3 @@
+<svg xmlns="http://www.w3.org/2000/svg" width="40" height="40" viewBox="0 0 40 40" fill="none">
+  <path fill-rule="evenodd" clip-rule="evenodd" d="M22.8386 1.74509C21.0821 0.751636 18.9179 0.751639 17.1613 1.74509L4.83866 8.71437C3.08209 9.70781 2 11.5438 2 13.5307V27.4693C2 29.456 3.08209 31.2921 4.83866 32.2856L17.1613 39.2549C18.9179 40.2484 21.0821 40.2484 22.8386 39.2549L35.1614 32.2856C36.9179 31.2921 38 29.456 38 27.4693V13.5307C38 11.5438 36.9179 9.70781 35.1614 8.71437L22.8386 1.74509ZM29.5863 10.4846C29.8386 10.3419 29.8385 9.98518 29.5862 9.8425L20.9462 4.95598C20.3607 4.62483 19.6393 4.62483 19.0537 4.95598L6.73108 11.9252C6.14556 12.2564 5.78486 12.8684 5.78486 13.5307V27.4693C5.78486 28.1315 6.14556 28.7435 6.73108 29.0747L19.0537 36.044C19.6393 36.375 20.3607 36.375 20.9462 36.044L33.2689 29.0747C33.8544 28.7435 34.2151 28.1315 34.2151 27.4693V23.3043C34.2151 23.0194 33.9007 22.841 33.6484 22.9827L20.8172 30.1904C20.6911 30.2611 20.5339 30.1719 20.5339 30.0296L20.5339 27.4095C20.5339 26.6148 20.9665 25.8805 21.6689 25.4831L34.026 18.4909C34.143 18.4247 34.2151 18.3023 34.2151 18.1699V13.5307C34.2151 13.2147 34.133 12.9101 33.9834 12.642C33.8938 12.4812 33.6841 12.4442 33.5217 12.5354L20.8172 19.6719C20.6911 19.7427 20.5339 19.6535 20.5339 19.5111L20.5339 16.8909C20.5339 16.0963 20.9665 15.3621 21.6689 14.9646L29.5863 10.4846Z" fill="#88FFE4"/>
+</svg>

mindsdb/integrations/handlers/libsql_handler/libsql_handler.py

@@ -0,0 +1,192 @@
+from typing import Optional
+
+import pandas as pd
+import libsql_experimental as libsql
+
+from mindsdb_sql_parser import parse_sql
+from mindsdb.integrations.libs.base import DatabaseHandler
+
+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 LibSQLHandler(DatabaseHandler):
+    """
+    This handler handles connection and execution of the LibSQL statements.
+    """
+
+    name = "libsql"
+
+    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 = "libsql"
+        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
+
+        args = self.connection_data
+        # sync_url and auth_token are optional
+        # sync_url is used to sync the local database from the remote database
+        # auth_token is used as the authentication token for the remote database
+        if args.get("sync_url"):
+            self.connection = libsql.connect(
+                database=args["database"],
+                sync_url=args["sync_url"],
+                auth_token=args["auth_token"],
+            )
+        else:
+            self.connection = libsql.connect(database=args["database"])
+
+        self.is_connected = True
+
+        return self.connection
+
+    def disconnect(self):
+        """
+        Close any existing connections.
+        """
+
+        if self.is_connected is False:
+            return
+
+        self.connection = None
+        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 SQLite {self.connection_data["database"]}, {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 {self.connection_data["database"]}!'
+            )
+            response = Response(RESPONSE_TYPE.ERROR, error_message=str(e))
+
+        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
+        """
+        return self.native_query(query)
+
+    def get_tables(self) -> StatusResponse:
+        """
+        Return list of entities that will be accessible as tables.
+        Returns:
+            HandlerResponse
+        """
+
+        query = "SELECT name from sqlite_master where type= 'table';"
+        result = self.native_query(query)
+        df = result.data_frame
+        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
+        """
+
+        query = f"PRAGMA table_info([{table_name}]);"
+        result = self.native_query(query)
+        df = result.data_frame
+        result.data_frame = df.rename(
+            columns={"name": "column_name", "type": "data_type"}
+        )
+        return result

mindsdb/integrations/handlers/libsql_handler/requirements.txt

@@ -0,0 +1 @@
+libsql-experimental

mindsdb/integrations/handlers/libsql_handler/tests/test_libsql_handler.py

@@ -0,0 +1,45 @@
+import unittest
+from mindsdb.integrations.handlers.libsql_handler.libsql_handler import LibSQLHandler
+from mindsdb.api.executor.data_types.response_type import RESPONSE_TYPE
+
+
+class LibSQLHandlerTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.kwargs = {
+            "database": "tests/test.db",
+        }
+        cls.handler = LibSQLHandler("test_libsql_handler", cls.kwargs)
+
+    def test_0_check_connection(self):
+        assert self.handler.check_connection()
+
+    def test_1_create_table(self):
+        query = "CREATE TABLE IF NOT EXISTS user (id INTEGER PRIMARY KEY, name TEXT, age INTEGER)"
+        result = self.handler.native_query(query)
+        assert result.type is RESPONSE_TYPE.OK
+
+    def test_2_insert_rows(self):
+        query = (
+            "INSERT OR IGNORE INTO user (name, age) VALUES ('Alice', 30), ('Bob', 25)"
+        )
+        result = self.handler.native_query(query)
+        print(result)
+        assert result.type is RESPONSE_TYPE.OK
+
+    def test_3_native_query_select(self):
+        query = "SELECT * FROM user"
+        result = self.handler.native_query(query)
+        assert result.type is RESPONSE_TYPE.TABLE
+
+    def test_4_get_tables(self):
+        tables = self.handler.get_tables()
+        assert tables.type is not RESPONSE_TYPE.ERROR
+
+    def test_5_get_columns(self):
+        columns = self.handler.get_columns("user")
+        assert columns.type is not RESPONSE_TYPE.ERROR
+
+
+if __name__ == "__main__":
+    unittest.main()

mindsdb/integrations/handlers/lightdash_handler/__about__.py

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

mindsdb/integrations/handlers/lightdash_handler/__init__.py

@@ -0,0 +1,27 @@
+from mindsdb.integrations.libs.const import HANDLER_TYPE
+
+from .__about__ import __description__ as description
+from .__about__ import __version__ as version
+
+try:
+    from .lightdash_handler import LightdashHandler as Handler
+    import_error = None
+except Exception as e:
+    Handler = None
+    import_error = e
+
+title = "Lightdash"
+name = "lightdash"
+type = HANDLER_TYPE.DATA
+icon_path = "icon.png"
+
+__all__ = [
+    "Handler",
+    "version",
+    "name",
+    "type",
+    "title",
+    "description",
+    "import_error",
+    "icon_path",
+]