AI Agent for Natural Language Queries to Generate Charts from a Database.
from datagpt.tool.llm_openai import OpenAI
from datagpt.log import logger
import json


# Plan 1: Stable output of config configuration
DATA2CHART_PROMPT_CONFIG = """

## Description
You are a JS engineer proficient in the latest "EChart" framework, responsible for converting data into EChart charts. If the requirement cannot produce content in the appropriate format, please return directly: "ChartError".

My original question was to your "Requirement"
SQL Result comes from the SQL query and is given to you in the task to be solved at the end of this paragraph.

In order to correctly render the results into the EChart chart, please carefully analyze the core query target and SQL Result of the original question.

First, based on the requirements, original question, and the resulting data, consider the following 6 points to analyze which chart type is more suitable:

The range of chart types includes (in order of priority):

1.Gauge chart (if there is only one data point in the array, prioritize using a gauge chart).
2.Pie chart (if the original question involves querying proportions or percentages, prioritize using a pie chart).
3.Radar chart (if the data has a single dimension and the number of data points is <= 6, prioritize generating a radar chart. In this case, the series' data array should contain only one object that includes all the values, and generate a matching indicator array based on the data in the series).
4.Line chart.
5.Scatter plot.
6.Multi-series bar chart (if there are multiple series in the array, prioritize using a bar chart. Additionally, if none of the above 1-5 situations apply, use a bar chart).
If the conditions stated in the parentheses above are met, the corresponding type of graph should be chosen as a priority.

Return all numerical values with a maximum precision of two decimal places.
Next, you can optimize the format of the "result" field, adjusting it to the data structure that you believe is most suitable for rendering the chart of the previously determined TYPE. You can replace the fields as needed.

Then, based on the optimized "result" field, provide me with the complete JS code for the options parameter required by ECharts in the setOption function, ensuring that all dimension axes and titles meet the requirements of the original problem.

You should strictly return three items to me in the following order, adhering to the specific requirements for each field as mentioned in the comments of the JSON format.
The optimized "result" field.
The correct parameters required in the options used for rendering, including: x (corresponding to the field name of the x-axis), y (corresponding to the field name of the y-axis), names of the x-axis and y-axis, title, tooltip, etc. Please do not disclose any other code or comments to me.
The final data format you return should strictly adhere to the following requirements for JSON format. The number of elements in the array should match the number returned by the API, and there should be no additional descriptive text included.
{{
     "result" :{{
           "a":["value1"],
           "b":["value2"]
      }},
     "option":{{
                "type":TYPE,
                "color":["#95a5fd", "#fd7f82", "#fec077", "#fee77d", "#95e38f", "#37bbff", "#74b6ff", "#d09dff","#444547"],
                "grid": {{
                    "top": "10%"
                }},
                "title": {{
                    "text": "XXX",//Replace the original question with more appropriate and concise wording.
                    "left": "center"
                }},
                "legend": {{
                    "orient": "horizontal",
                    "top": "bottom"
                }}
                "xAxis": //only exists when TYPE is bar or lines
                {{
                    "data": [...],//you should realize a function to put all item's name into the array when type is bar or lines
                    "axisLabel": {{
                        "rotate": XXX,//If the number of data points is less than 5, set the value to 0. If it is less than 10, set the value to 30. If it is greater than 15, set the value to 45.
                        "interval": 0
                    }},
                    "show":bool//set true when TYPE is bar or lines.others set false
                }},
                "yAxis": {{}},//only exists when TYPE is bar or lines.
                "series": [{{
                    "name": "Title Based on Your Understanding",
                    "type": TYPE,
                    "radius": "40%",//(gauge:80%,other:40%)                   
                   
                    "data(only exists when TYPE is radar)":
                    [
                        //Only this ONE object exists if TYPE is radar.
                        {{
                            "name":"",
                            "value":number1,number2,...,numberN  // Put all values in the array if TYPE is radar.
                        }}
                        //No more objects! if TYPE is radar
                    ]

                }},{{}}],
                "radar"://only exists when TYPE is radar 
                {{
                    "indicator":
                    [
                        {{
                            "name":"series[0].name",
                            "max":MAX//Retrieve the maximum value from the 'value' array inside the 'series.data' automatically.
                        }},{{
                            "name":"series[1].name",
                            "max":MAX//Ensure that the MAX value equals the previous MAX value
                        }}……
                    ]
                }}
        }}
}}

## Examples
Requirement: Proportion of the area occupied by the top 6 parking spaces in terms of total area.
SQL Result：[["space_name", "space_area"], ["MUJI无印良品", 200.0], ["优衣库", 140.0], ["105", 50.0], ["奈雪の茶", 40.0], ["103", 35.0], ["机房", 35.0]]
Result: {{}}

Requirement: Empty data or data with excessively high dimensions
SQL Result:empty or error format
Result: ChartError

## Task to solve
Requirement: {requirement}
SQL Result：{result}
Result:
"""


# plan2: More aggressive dynamic function configuration
DATA2CHART_PROMPT_FUNCTION = """

## Description
You are a JS engineer proficient in the latest "EChart" framework, responsible for converting data into EChart charts. If the requirement cannot produce content in the appropriate format, please return directly: "ChartError".

My original question was to your "Requirement"
SQL Result comes from the SQL query and is given to you in the task to be solved at the end of this paragraph.
In order to correctly render the results into the EChart chart, please carefully analyze the core query target and SQL Result of the original question.

First, based on the requirements, original question, and the resulting data, consider the following 6 points to analyze which chart type is more suitable:

The range of chart types includes (in order of priority):

1.Gauge chart (if there is only one data point in the array, prioritize using a gauge chart).
2.Pie chart (if the original question involves querying proportions or percentages, prioritize using a pie chart).
3.Radar chart (if the data has a single dimension and the number of data points is <= 6, prioritize generating a radar chart. In this case, the series' data array should contain only one object that includes all the values, and generate a matching indicator array based on the data in the series).
4.Line chart.
5.Scatter plot.
6.Multi-series bar chart (if there are multiple series in the array, prioritize using a bar chart. Additionally, if none of the above 1-5 situations apply, use a bar chart).
If the conditions stated in the parentheses above are met, the corresponding type of graph should be chosen as a priority.

Then, you can optimize the format of the 'result' field and adjust it to the data structure that you consider most suitable for rendering the chart of the type determined in the previous step. You can use '...' to indicate the omission of a large amount of data content. You can replace the fields as needed.
Next, based on the optimized 'result' field, please provide me with the complete dynamic JavaScript function that is required to generate the 'option' JSON string parameter for the setOption method in ECharts. In other words, I need you to return dynamic JS code. The requirements for this code are as follows:

1.The first parameter is SQL Result, and in the provided code, this result needs to be dynamically transformed into key data required for various parameters in the options.
2.The type field, which represents the table type, should be defined at the beginning of the function. Determine the specific type based on the above requirements, and directly use this type within the following JSON.
3.The algorithm supports multiple series to ensure that the chart can display the complete data when it comes in.
4.You only return the function code. Once the code is returned, please immediately stop answering.
5.The returned result must adhere to the JSON format required by ECharts options to ensure the complete rendering of the chart.
6.The XXX characters in the JSON example below need to be replaced with the most appropriate content based on your understanding of the requirements.
7.For specific requirements of each field in the returned option, please refer to the comments in the JSON format requirements below.
8.The data format returned by this function must strictly adhere to the following JSON format requirements, and the number of elements in the array should be the same as the data quantity returned by the API (i.e., consistent with the data volume of the first parameter of this function):
{{
                "type":TYPE,
                "color":["#95a5fd", "#fd7f82", "#fec077", "#fee77d", "#95e38f", "#37bbff", "#74b6ff", "#d09dff","#444547"],
                "grid": {{
                    "top": "10%"
                }},
                "title": {{
                    "text": "XXX",//Replace the original question with more appropriate and concise wording.
                    "left": "center"
                }},
                "legend": {{
                    "orient": "horizontal",
                    "top": "bottom"
                }}
                "xAxis": //only exists when TYPE is bar or lines
                {{
                    "data": [...],//you should realize a function to put all item's name into the array when type is bar or lines
                    "axisLabel": {{
                        "rotate": XXX,//If the number of data points is less than 5, set the value to 0. If it is less than 10, set the value to 30. If it is greater than 15, set the value to 45.
                        "interval": 0
                    }},
                    "show":bool//set true when TYPE is bar or lines.others set false
                }},
                "yAxis": {{}},//only exists when TYPE is bar or lines.
                "series": [{{
                    "name": "Title Based on Your Understanding",
                    "type": TYPE,
                    "radius": "40%",//(gauge:80%,other:40%)                   
                   
                    "data(only exists when TYPE is radar)":
                    [
                        //Only this ONE object exists if TYPE is radar.
                        {{
                            "name":"",
                            "value":[number1,number2,...,numberN]//Put all values in the array if TYPE is radar.
                        }}
                        //No more objects! if TYPE is radar
                    ]

                }},{{}}],
                "radar"://only exists when TYPE is radar 
                {{
                    "indicator":
                    [
                        {{
                            "name":"series[0].name",
                            "max":MAX//Retrieve the maximum value from the 'value' array inside the 'series.data' automatically.
                        }},{{
                            "name":"series[1].name",
                            "max":MAX//Ensure that the MAX value equals the previous MAX value
                        }}……
                    ]
                }}
}}


## Examples
Requirement: Proportion of the area occupied by the top 6 parking spaces in terms of total area.
Sql Result:[["space_name", "space_area"], ["nameAA", 200.0], ……, ["nameYY", 35.0], ["nameZZ", 35.0]]
Result: 
function generateOption(apiResultList){{
    var option= {{}};// Initialization of the entire options content, it must be in JSON format.
    var type = "pie";//Here, the comment explains the reasons for choosing the chart type.
    //This is the specific algorithm to break down the apiResultList data.
    //Here is the complete content that meets the requirements of the option:
    //1、build option basic structure by template before
    //2、filling all the titles and data by a dynamic function
    //3、check the option again and promise it works currectly
    return option;
}}
//NO Note!

Requirement: An illogical question
Sql Result:Empty data or data with excessively high dimensions
Result: ChartError
//NO Note!

## Task to solve
Requirement: {requirement}
Sql Result:{result}
Result:
"""


class Data2Chart:
    def __init__(self):
        self.llm = OpenAI()

    def gen_chart(self, requirement: str, result: list, render_mode: str) -> str:
        if render_mode == "Config":
            result = json.dumps(result)
            prompt = DATA2CHART_PROMPT_CONFIG.format(
                requirement=requirement, result=result
            )
        elif render_mode == "Function":
            # TODO:Provide fewer data samples and allow AI to know the format based on limited data
            # if len(result) >= 3:
            #     result = [result[0],result[1],result[-1]]
            result = json.dumps(result)
            prompt = DATA2CHART_PROMPT_FUNCTION.format(
                requirement=requirement, result=result
            )

        return self.llm.ask(prompt)

datagpt/action/data2chart.py
digai-co/DataGPT
from datagpt.config import config
from datagpt.tool.db import db


class FetchData:
    def __init__(self):
        self.db = db

    def fetch(self, sql: str):
        return self.db.run(sql)

datagpt/action/fetch_data.py
digai-co/DataGPT
from datagpt.config import config
from datagpt.util import root
from datagpt.memory.faiss import FaissStore
from datagpt.tool.llm_openai import OpenAI
from datagpt.log import logger
from datagpt.tool.db import db

TEXT2SQL_PROMPT = """
## Description
You are a data engineer responsible for translating the requirements of management or business 
personnel into SQL statements that can be executed in {database_type}. Information about each table 
in the database will be provided to you, including table names, table descriptions, field names, 
types, and descriptions for each field. You need to select the relevant tables based on this 
information and write the corresponding SQL statements. If the requirements go beyond the information 
provided by these tables, please be sure to return: "BeyondError."

Note: ONLY ONE sql statement for each requirement.


## Examples
Requirement: Find the name, employee ID, and age of the oldest employee.
SQL: SELECT employee_id, name FROM staff ORDER BY age DESC LIMIT 1;

Requirement: Who is my best friend?
SQL: BeyondError

## Database Information
{database_schema}

## Task to solve
Requirement: {requirement}
SQL:
"""


class Text2SQL:
    def __init__(self):
        self.llm = OpenAI()

    def gen_sql(self, requirement: str) -> str:
        # Search top k tables related to this requirement from long-term memory
        memory_dir = config.get("memory.dir")
        schema_file = config.get("memory.schema_file")
        if not (root / memory_dir / f"{schema_file}.pkl").exists():
            self._cache_schema(memory_dir, schema_file)

        schema_store = FaissStore(memory_dir, schema_file)
        schema = schema_store.search(requirement)
        
        # Generate sql for this requirement
        prompt = TEXT2SQL_PROMPT.format(
            database_type=db.get_type(), database_schema=schema, requirement=requirement
        )
        sql = self.llm.ask(prompt)
        if sql == "BeyondError":
            logger.error(f"Error on generating sql for [{requirement}]")

        return sql

    def _cache_schema(self, memory_dir: str, schema_file: str):
        """Cache schemas of all tables to long-term memory"""
        schemas = {}

        tables = db.get_tables()
        for table in tables:
            fields = db.get_fields(table.name)
            schema = "\n".join(
                [
                    f"col:{f.col_name}, type:{f.col_type}, description:{f.col_comment}"
                    for f in fields
                ]
            )
            schemas[
                f"{table.name}:{table.comment}"
            ] = f"table: {table.name}\ntable description: {table.comment}\ncolumns information:\n{schema}\n"

        schema_store = FaissStore(memory_dir, schema_file)
        schema_store.write(schemas)

datagpt/action/text2sql.py
digai-co/DataGPT
import os

import yaml

from datagpt.util import root


class Config:
    """Usage：
    config = Config("config.yaml")
    secret_key = config.get_key("MY_SECRET_KEY")
    print("Secret key:", secret_key)
    """

    default_yaml_file = root / "config/config.yaml"

    def __init__(self, yaml_file=default_yaml_file):
        self._config = {}
        self.__init_with_config_files_and_env(self._config, yaml_file)

    def __init_with_config_files_and_env(self, configs: dict, yaml_file):
        for _yaml_file in [yaml_file]:
            if not _yaml_file.exists():
                continue

            with open(_yaml_file, "r", encoding="utf-8") as file:
                yaml_data = yaml.safe_load(file)
                if not yaml_data:
                    continue

                configs.update(yaml_data)

        configs.update(os.environ)

    def __lookup(self, dct, key):
        if "." in key:
            key, node = key.split(".", 1)
            return self.__lookup(dct[key], node)
        elif key in dct:
            return dct[key]
        else:
            return None

    def get(self, *keys: str, default=None):
        for key in keys:
            value = self._config.get(key)
            if value is not None:
                return value
            value = self.__lookup(self._config, key)
            if value is not None:
                return value
        return default


config = Config()

datagpt/config.py
digai-co/DataGPT
from datagpt.action import text2sql, fetch_data, data2chart
from datagpt.log import logger


class DataGPT:
    def __init__(self):
        self.text2sql = text2sql.Text2SQL()
        self.fetchdata = fetch_data.FetchData()
        self.data2chart = data2chart.Data2Chart()

    def run(self, text: str, render_mode: str):
        """from text to chart"""

        logger.debug(f"input text: {text}")
        sql = self.text2sql.gen_sql(text)
        if sql == "BeyondError":
            return sql, None, None

        logger.debug(f"generated sql: {sql}")
        return self.__excute_sql(sql, text, render_mode)

    def run_sql(self, text: str, sql: str, render_mode: str):
        """from sql to chart"""

        logger.debug(f"input sql: {sql}")
        return self.__excute_sql(sql, text, render_mode)

    def __excute_sql(self, sql, text, render_mode):
        data = self.fetchdata.fetch(sql)
        if data == "SqlError":
            return sql, data, None

        logger.debug(f"fetched data: {data}")
        chart = self.data2chart.gen_chart(text, data, render_mode)
        logger.debug(f"rendered chart: {chart}")

        return sql, data, chart

datagpt/datagpt.py
digai-co/DataGPT
import sys

from loguru import logger as _logger

from datagpt.util import root


def get_logger(print_level="INFO", logfile_level="DEBUG"):
    _logger.remove()
    _logger.add(sys.stderr, level=print_level)

    _logger.add(root / "logs/log.txt", level=logfile_level)
    return _logger


logger = get_logger()

datagpt/log.py
digai-co/DataGPT
import os
import pickle

import faiss
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import FAISS

from datagpt.config import config
from datagpt.util import root


class FaissStore:
    def __init__(self, dir: str, name: str):
        self.name = name
        if not os.path.exists(root / dir):
            os.mkdir(root / dir)

        self.store_file = root / dir / f"{name}.pkl"
        self.index_file = root / dir / f"{name}.index"

        if not (self.index_file.exists() and self.store_file.exists()):
            self.store = None
        else:
            with open(str(self.store_file), "rb") as f:
                self.store = pickle.load(f)
            self.store.index = faiss.read_index(str(self.index_file))

    def _persist(self):
        faiss.write_index(self.store.index, str(self.index_file))
        index = self.store.index
        self.store.index = None
        with open(self.store_file, "wb") as f:
            pickle.dump(self.store, f)
        self.store.index = index

    def write(self, schemas: dict):
        """batch writing all table schemas into long-term memory
        schemas should be
          {
            'table description': 'table schema',
            'table description': 'table schema'
          }
        """
        if config.get("llm.openai.azure") is not None:
            embeddings = OpenAIEmbeddings(
                deployment=config.get("llm.openai.azure.embedding_deployment"),
                openai_api_base=config.get(
                    "llm.openai.api_base", default="https://api.openai.com/v1"
                ),
                openai_api_key=config.get("llm.openai.api_key"),
                openai_api_type="azure",
                chunk_size=16,
            )
        else:
            embeddings = OpenAIEmbeddings(
                openai_api_key=config.get("llm.openai.api_key"), chunk_size=16
            )

        self.store = FAISS.from_texts(
            list(schemas.keys()),
            embeddings,
            metadatas=[{k: schemas[k]} for k in schemas],
        )
        self._persist()

    def search(self, query, k=10):
        """search top k tables that most relating to the query"""
        r = self.store.similarity_search(query, k=k)
        return str("\n".join([f"{x.metadata}" for x in r]))

datagpt/memory/faiss.py
digai-co/DataGPT
from datagpt.config import config
from datagpt.log import logger
from datagpt.tool.db.mysql import MySQL
from datagpt.tool.db.postgres import PostgreSQL


def get_db():
    db_type = config.get("database.type")
    db_uri = config.get("database.uri")
    if db_type == "postgresql":
        return PostgreSQL(db_uri)
    elif db_type == "mysql":
        return MySQL(db_uri)
    else:
        logger.exception("No vaild database config")


db = get_db()

__all__ = ["db"]

datagpt/tool/db/__init__.py
digai-co/DataGPT
from abc import ABCMeta, abstractmethod
from typing import Iterable, List, Optional


class TableInfo:
    name: str
    comment: str

    def __init__(self, name: str, comment: Optional[str] = None):
        self.name = name
        self.comment = comment

    def __str__(self):
        return f"TableInfo(name={self.name}, comment={self.comment})"


class FieldInfo:
    col_name: str
    col_type: str
    col_length: Optional[int]
    col_default: str
    is_nullable: bool
    col_comment: str

    def __init__(
        self,
        col_name: str,
        col_type: str,
        col_length: Optional[int],
        col_default: str,
        is_nullable: bool = False,
        col_comment: Optional[str] = None,
    ):
        self.col_name = col_name
        self.col_type = col_type
        self.col_length = col_length
        self.col_default = col_default
        self.is_nullable = is_nullable
        self.col_comment = col_comment

    def __str__(self):
        return f"FieldInfo(col_name={self.col_name}, col_type={self.col_type}, col_length={self.col_length}, col_default={self.col_default}, is_nullable={self.is_nullable}, col_comment={self.col_comment})"


class RDBMS(metaclass=ABCMeta):
    @abstractmethod
    def get_type(self) -> str:
        pass

    @abstractmethod
    def get_tables(self) -> Iterable[TableInfo]:
        pass

    @abstractmethod
    def get_fields(
        self, table_name: str, schema: str = "public"
    ) -> Iterable[FieldInfo]:
        pass

    @abstractmethod
    def run(self, command: str) -> List:
        pass

datagpt/tool/db/base.py
digai-co/DataGPT
from typing import Iterable, List
from datagpt.tool.db.base import RDBMS, FieldInfo, TableInfo

from urllib.parse import urlparse
import mysql.connector

from datagpt.log import logger


class MySQL(RDBMS):
    def __init__(self, uri: str):
        self._uri = uri

        o = urlparse(uri)
        config = {
            "host": o.hostname,
            "port": o.port,
            "user": o.username,
            "password": o.password,
            "database": o.path[1:],
        }

        self._conn = mysql.connector.connect(**config)

    def get_type(self) -> str:
        return "MySQL"

    def get_tables(self) -> Iterable[TableInfo]:
        cur = self._conn.cursor()
        cur.execute(
            """
            SELECT table_name, table_comment
            FROM information_schema.tables
            WHERE table_schema = DATABASE();
        """
        )
        tb_infos = cur.fetchall()
        cur.close()

        return [TableInfo(name, comment) for name, comment in tb_infos]

    def get_fields(
        self, table_name: str, schema: str = "public"
    ) -> Iterable[FieldInfo]:
        cur = self._conn.cursor()
        cur.execute(
            """
            SELECT COLUMN_NAME, DATA_TYPE, CHARACTER_MAXIMUM_LENGTH, COLUMN_DEFAULT, IS_NULLABLE, COLUMN_COMMENT
            FROM information_schema.columns
            WHERE table_schema = DATABASE() AND table_name = %s
            ORDER BY ORDINAL_POSITION;
        """,
            (table_name,),
        )
        fields = cur.fetchall()
        cur.close()

        return [
            FieldInfo(name, type, length, default, nullable, comment)
            for name, type, length, default, nullable, comment in fields
        ]

    def run(self, command: str) -> List:
        try:
            cur = self._conn.cursor()
            cur.execute(command)
            rows = cur.fetchall()
            cols = cur.description
            rows.insert(0, tuple([col[0] for col in cols]))
            return rows
        except Exception as e:
            logger.error("run sql error: " + str(e))
            return "SqlError"
        finally:
            cur.close()

datagpt/tool/db/mysql.py
digai-co/DataGPT
from typing import Iterable, List

from psycopg_pool import ConnectionPool

from datagpt.tool.db.base import RDBMS, FieldInfo, TableInfo
from datagpt.log import logger


class PostgreSQL(RDBMS):
    def __init__(self, uri: str):
        self._pool = ConnectionPool(uri)

    def get_type(self) -> str:
        return "PostgreSQL"

    def get_tables(self) -> Iterable[TableInfo]:
        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                tb_infos = cur.execute(
                    """
                    SELECT tb.table_name, d.description
                    FROM information_schema.tables tb
                        LEFT JOIN pg_namespace b ON b.nspname = tb.table_schema
                        LEFT JOIN pg_class c ON c.relname = tb.table_name AND c.relnamespace = b.oid
                        LEFT JOIN pg_description d ON d.objoid = c.oid AND d.objsubid = '0'
                    WHERE tb.table_schema = 'public' and tb.table_type = 'BASE TABLE';
                """
                ).fetchall()
                return [TableInfo(name, comment) for name, comment in tb_infos]

    def get_fields(
        self, table_name: str, schema: str = "public"
    ) -> Iterable[FieldInfo]:
        with self._pool.connection() as conn:
            with conn.cursor() as cur:
                fields = cur.execute(
                    """
                    SELECT
                        c.column_name,
                        c.udt_name,
                        character_maximum_length,
                        c.column_default,
                        c.is_nullable,
                        pgd.description AS column_comment
                    FROM
                        information_schema.columns c
                    LEFT JOIN
                        pg_catalog.pg_description pgd ON pgd.objoid = (select oid from pg_class where relnamespace=(select oid from pg_namespace where nspname=%(schema)s) and relname=%(table)s) AND pgd.objsubid = c.ordinal_position
                    WHERE
                        c.table_schema = %(schema)s AND c.table_name = %(table)s
                    ORDER BY
                        c.ordinal_position;
                """,
                    {"table": table_name, "schema": schema},
                ).fetchall()

                return [
                    FieldInfo(name, type, length, default, nullable, comment)
                    for name, type, length, default, nullable, comment in fields
                ]

    def run(self, command: str) -> List:
        try:
            with self._pool.connection() as conn:
                with conn.cursor() as cur:
                    rows = cur.execute(command).fetchall()
                    cols = cur.description
                    rows.insert(0, tuple([col.name for col in cols]))
                    return rows
        except Exception as e:
            logger.error("run sql error: " + str(e))
            return "SqlError"

datagpt/tool/db/postgres.py
digai-co/DataGPT
import openai
from datagpt.log import logger
from datagpt.config import config


class OpenAI:
    """LLM api to OpenAI ChatGPT models."""

    def __init__(self) -> None:
        if config.get("llm.openai.azure") is not None:
            openai.api_type = "azure"
            openai.api_version = config.get(
                "llm.openai.azure.api_version", default="2023-05-15"
            )
        openai.api_base = config.get(
            "llm.openai.api_base", default="https://api.openai.com/v1"
        )
        openai.api_key = config.get("llm.openai.api_key")

    def completion(self, messages: list[dict]) -> str:
        if openai.api_type == "azure":
            engine = config.get("llm.openai.azure.engine")
            chat_completion = openai.ChatCompletion.create(
                engine=engine, messages=messages, temperature=0
            )
        else:
            model = config.get("llm.openai.api_model")
            chat_completion = openai.ChatCompletion.create(
                model=model, messages=messages, temperature=0
            )
        return chat_completion.choices[0].message.content

    def ask(self, prompt: str) -> str:
        messages = [{"role": "user", "content": prompt}]
        return self.completion(messages)
datagpt/tool/llm_openai.py
digai-co/DataGPT
from pathlib import Path


def get_project_root():
    current_path = Path.cwd()
    while True:
        if (current_path / ".project_root").exists():
            return current_path
        parent_path = current_path.parent
        if parent_path == current_path:
            raise Exception("Project root not found.")
        current_path = parent_path


root = get_project_root()

datagpt/util.py
digai-co/DataGPT
import sys, os

sys.path.append((os.path.abspath(os.path.join(os.path.dirname(__file__), "../"))))


import json
import gradio as gr

from datagpt.datagpt import DataGPT

# Config Mode: Genrating echart option config directly
# Function Mode:Generating option render function which be called to generate option config
render_mode = "Config"
datagpt = DataGPT()


def search(prompt):
    if prompt == "error":
        raise gr.Error("Error!")
    sql, data, chart = datagpt.run(prompt, render_mode)
    return render_chart(sql, data, chart)


def reexcute_sql(prompt, sql):
    sql, data, chart = datagpt.run_sql(prompt, sql, render_mode)
    return render_chart(sql, data, chart)


def render_chart(sql, data, chart):
    if sql == "BeyondError":
        raise gr.Error("Unable to generate SQL statement, please ask in another way!")
    elif data == "SqlError":
        return [
            sql,  # SQL statement
            gr.update(visible=True),
            "",  # chart iframe
            "",  # chart source code
            "",  # data table
            [
                {
                    "": "",
                    "SQL execution encountered an error, You can modify it below and execute it again.": "red",
                },
            ],
            gr.Highlight.update(visible=True),
        ]
    else:
        chart_iframe, chart_source_code = gen_chart_iframe(chart, json.dumps(data))
        return [
            sql,  # SQL statement
            gr.update(visible=True),  # sql data result
            chart_iframe,  # chart iframe
            chart_source_code,  # chart source code
            gen_table_iframe(data),  # data table
            [],
            gr.Highlight.update(visible=False),
        ]


def gen_chart_iframe(echart_option, result_data):
    html = ""
    # Solution 1: Fill with AI-generated Options
    if render_mode == "Config":
        echart_option = json.loads(echart_option.encode("utf-8"))
        echart_option = echart_option["option"]

        html = """
<head>
    <title>DataGPT</title>
        <script src="https://cdn.jsdelivr.net/npm/echarts@5.4.1/dist/echarts.min.js"></script>
</head>
    <body>
        <div id="main" style="width: 100%;height:350px;"></div>
            <script>
                var chart = echarts.init(document.getElementById("main"));
                var option = {option}   
                chart.setOption(option);
            </script>
    </body>
            """
        html = html.format(option=json.dumps(echart_option))

    # Solution 2: Using AI-generated dynamic functions combined with real data to dynamically generate options
    elif render_mode == "Function":
        html = """
<head>
    <title>DataGPT</title>
        <script src="https://cdn.jsdelivr.net/npm/echarts@5.4.1/dist/echarts.min.js"></script>
</head>
<body>
        <div id="main" style="width: 100%;height:350px;"></div>
            <script>
                var chart = echarts.init(document.getElementById("main"));
                {generate_option}
                var sql_result = {result_data}
                var option = generateOption(sql_result);  
                console.log(JSON.stringify(option));
                chart.setOption(option);
            </script>
</body>
            """
        html = html.format(generate_option=echart_option, result_data=result_data)

    return [
        f"""<iframe style="width: 100%; height: 400px" srcdoc='{html}'></iframe>""",
        f"""```html{html}""",
    ]


def gen_table_iframe(data):
    html = """
<head>
</head>
<body>
  <table style="width:100%; border-radius:10px; border:#eee 1px solid; border-spacing:0px; border-collapse:separate; font-size:13px; font-family:Arial;">
    <tr style="background-color:transparent; border:0; height:40px">
        {tds}
    </tr>
    {trs}   
  </table>
</body>
"""
    tds_html = ""
    trs_html = ""
    columns = data[0]

    # render headers
    for col in columns:
        td = """<td style="padding-left:10px;"><strong>{value}<strong></td>"""
        td = td.format(value=str(col))
        tds_html += td

    # render rows
    if len(data) >= 1:
        data = data[1:]
        index = 0
        for row in data:
            tr = """
                <tr style="background-color:{bg_color};border:0;height:40px;">
                    {tds}
                </tr>"""

            bg_color = ""
            if index % 2 == 1:
                bg_color = "transparent"
            else:
                bg_color = "#f9fafb"

            tds = ""
            for col in row:
                td = """<td style="padding-left:10px;">{value}</td>"""
                td = td.format(value=str(col))
                tds += td
            tr = tr.format(tds=tds, bg_color=bg_color)
            trs_html += tr
            index += 1

    html = html.format(tds=tds_html, trs=trs_html)
    return (
        f"""<iframe style="width: 100%; min-height: 640px" srcdoc='{html}'></iframe>"""
    )


def change_mode(mode):
    global render_mode
    render_mode = mode
    if mode == "Config":
        return gr.update(info="Config Mode: Genrating echart option config directly")
    elif mode == "Function":
        return gr.update(
            info="Function Mode:Generating option render function which be called to generate option config"
        )


def launch_app():
    chart_iframe = gr.HTML("Chart")
    chart_source_code = gr.Markdown("Chart source code")

    with gr.Blocks(title="DataGPT", css="label {font-size:12px}") as ui:
        gr.Markdown(
            """
            # DataGPT
            AI Agent for Natural Language Queries to Generate Charts from a Database.
            """
        )

        with gr.Row():
            with gr.Column(scale=2):
                input_prompt = gr.Textbox(
                    value="What are the specific grades of each student in each course?",
                    label="Question",
                    placeholder="Please input question",
                )
                radio = gr.Radio(
                    choices=["Config", "Function"],
                    type="value",
                    value="Config",
                    label="Render Mode",
                    info="Option Config Mode:Genrating echart option config directly",
                )
                btn = gr.Button(value="Submit", variant="primary")
                error_info = gr.HighlightedText(
                    label="",
                    visible=False,
                    combine_adjacent=True,
                    show_legend=True,
                    color_map={
                        "SQL execution encountered an error, You can modify it below and execute it again.": "red"
                    },
                )
                lbl_sql = gr.Code(
                    value="The sql will be displayed here.",
                    language="markdown",
                    interactive=True,
                    lines=1,
                )
                btn_exe_sql = gr.Button(
                    value="Modify and Re-execute the SQL Statement",
                    size="sm",
                    visible=False,
                )
            with gr.Column(scale=2):
                demo = gr.TabbedInterface(
                    [chart_iframe, chart_source_code], ["Chart", "Source Code"]
                )
        with gr.Row():
            table_iframe = gr.HTML("")

        radio.change(change_mode, inputs=[radio], outputs=[radio])
        btn.click(
            search,
            inputs=input_prompt,
            outputs=[
                lbl_sql,
                btn_exe_sql,
                chart_iframe,
                chart_source_code,
                table_iframe,
                error_info,
                error_info,
            ],
        )
        btn_exe_sql.click(
            reexcute_sql,
            inputs=[input_prompt, lbl_sql],
            outputs=[
                lbl_sql,
                btn_exe_sql,
                chart_iframe,
                chart_source_code,
                table_iframe,
                error_info,
                error_info,
            ],
        )
        ui.launch(share=True)


launch_app()

server/app.py
digai-co/DataGPT