import logging

import streamlit as st

from app.chat import callbacks, config, utils
from app.chat.data_models import Message, MessagePayload, MessageType
from app.engine.engine import generate_query
from app.engine.sql_engine import Database

logger = logging.getLogger(__name__)


utils.setup_look("Mowa SQL Copilot")
db = Database()


def chat_interface():
    with st.sidebar:
        st.selectbox(
            "Language",
            ("English", "Polish"),
            key="lang",
            label_visibility="hidden",
            on_change=callbacks.on_lang_change,
        )

        st.markdown("***")
        table_name = st.selectbox(config.LANG_DICT[st.session_state["lang"]]["check_schema"], db.tables)
        _, center, _ = st.columns([0.1, 1, 0.1])
        with center:
            st.dataframe(
                db.schema_dict[table_name],
                hide_index=False,
            )

        with st.expander(config.LANG_DICT[st.session_state["lang"]]["choose_tables"]):
            st.multiselect(
                "Tables",
                db.tables,
                default=db.tables,
                key="chosen_tables",
                label_visibility="hidden",
                on_change=callbacks.on_table_change,
            )

    chat_placeholder = st.empty()
    with chat_placeholder.container():
        with st.chat_message("bot", avatar=utils.URL_TO_ICON):
            st.write(config.LANG_DICT[st.session_state["lang"]]["welcome_msg"])
        utils.show_messages()
        if len(st.session_state[config.HISTORY]) > 0:
            is_active_editor = st.session_state[config.BOT_MESSAGE][-1].payload.active or False
        else:
            is_active_editor = False
        user_input = st.chat_input(
            placeholder=config.LANG_DICT[st.session_state["lang"]]["placeholder"], disabled=is_active_editor
        )
        if user_input:
            logger.debug("on user input")
            st.chat_input(
                placeholder=config.LANG_DICT[st.session_state["lang"]]["placeholder"],
                disabled=True,
                key="disabled_chat_input",
            )
            user_message = Message(
                message_type=MessageType.USER,
                payload=MessagePayload(text=user_input),
            )
            st.session_state[config.USER_MESSAGE].append(user_message)

            with st.chat_message("user"):
                st.write(user_message.payload.text)

            st.info(config.LANG_DICT[st.session_state["lang"]]["sending_msg"])
            query = (
                generate_query(
                    text=user_message.payload.text, history=st.session_state.history, lang=st.session_state["lang"]
                )
                .strip()
                .split(";")[0]
            )

            ai_message = Message(
                message_type=MessageType.AI,
                payload=MessagePayload(code=query, active=True),
            )
            st.session_state.history.append(
                {"type": "Human", "text": user_message.payload.text},
            )
            callbacks.update_session_state_on_response(ai_message)

            st.rerun()

        col1, col2 = st.columns([0.25, 1])
        with col1:
            st.button(config.LANG_DICT[st.session_state["lang"]]["clear_button"], on_click=callbacks.on_clear_click)
        with col2:
            if is_active_editor:
                st.button(
                    config.LANG_DICT[st.session_state["lang"]]["change_button"], on_click=callbacks.on_regen_click
                )


if __name__ == "__main__":
    st.title("Mowa SQL Copilot")
    config.initialize_streamlit()
    chat_interface()

app/Chat.py
mowa-ai/sql-copilot
import logging
import uuid

import streamlit as st

from app.chat import config
from app.chat.data_models import Message
from app.engine.sql_engine import Database

logger = logging.getLogger(__name__)


def on_clear_click() -> None:
    logger.debug("on_clear_click")
    del st.session_state[config.USER_MESSAGE][:]
    del st.session_state[config.BOT_MESSAGE][:]
    del st.session_state[config.HISTORY][:]
    st.session_state.chat_id = uuid.uuid4().hex


def on_table_change() -> None:
    logger.debug("on_save_click")
    db = Database()
    db.update_schema(st.session_state["chosen_tables"])
    logger.debug(f"Tables updated: {db.schema}")


def on_regen_click() -> None:
    logger.debug("on_regen_click")
    st.session_state[config.BOT_MESSAGE].pop()
    st.session_state.history.pop()
    st.session_state[config.USER_MESSAGE].pop()
    st.session_state.history.pop()


def on_lang_change() -> None:
    logger.debug("on_lang_change")
    on_clear_click()


def update_session_state_on_response(response: Message, update=False) -> None:
    logger.info("Updating session state on response")
    logger.debug(f"Response payload: {response}")
    if update:
        st.session_state[config.BOT_MESSAGE].pop()
        st.session_state.history.pop()

    st.session_state[config.BOT_MESSAGE].append(response)
    record = {"type": "AI"}
    if response.payload.text:
        record["text"] = response.payload.text
    if response.payload.code:
        record["code"] = response.payload.code
    if response.payload.df is not None:
        record["df"] = response.payload.df
    if response.payload.active is None:
        record["active"] = True

    st.session_state.history.append(record)
    logger.debug("Updating session state on response - done")

app/chat/callbacks.py
mowa-ai/sql-copilot
import logging
import uuid

import streamlit as st

logger = logging.getLogger(__name__)

USER_MESSAGE = "user_message"
BOT_MESSAGE = "bot_message"
HISTORY = "history"
LANG_DICT = {
    "English": {
        "check_schema": "Check schema",
        "choose_tables": "Tables for assistant",
        "placeholder": "Your message",
        "clear_button": "Clear messages",
        "change_button": "Change request",
        "sending_msg": "Sending message to AI...",
        "welcome_msg": """
Hello! I'm MowaSQL, a computer-based SQL Copilot created by [mowaAI](https://www.mowa.ai/). Visit my creators' webpage and don't hesitate to reach them.

I'm designed to understand text in natural language and translate it to SQL, which means I can assist you in analyzing various datasets. Ask a question, check SQL code I've generated, edit if necessary and run to get a result.

What would you like to ask?
""",
    },
    "Polish": {
        "check_schema": "Sprawdź schemat",
        "choose_tables": "Tabele dla Asystenta",
        "placeholder": "Twoja wiadomość",
        "clear_button": "Wyczyść",
        "change_button": "Zmień zapytanie",
        "sending_msg": "Wysyłanie wiadomości do AI...",
        "welcome_msg": """
Cześć! Jestem MowaSQL, twój personalny SQL Copilot, stworzony przez [mowaAI](https://www.mowa.ai/). Odwiedź stronę internetową moich twórców i nie wahaj się z nimi skontaktować.

Jestem stworzony do rozumienia tekstu w języku naturalnym i tłumaczenia go na SQL, co oznacza, że mogę pomóc Ci w analizie różnych zbiorów danych. Zadaj pytanie, sprawdź wygenerowany kod SQL, w razie potrzeby edytuj i uruchom, aby uzyskać wynik.

O co chciałbyś zapytać?
""",
    },
}


def initialize_streamlit() -> None:
    st.session_state.setdefault(USER_MESSAGE, [])
    st.session_state.setdefault("lang", "English")
    st.session_state.setdefault(BOT_MESSAGE, [])
    st.session_state.setdefault(HISTORY, [])
    st.session_state.setdefault("chat_id", uuid.uuid4().hex)
    st.session_state.setdefault("tables", uuid.uuid4().hex)

app/chat/config.py
mowa-ai/sql-copilot
import datetime
from dataclasses import dataclass, field
from enum import Enum

import pandas as pd


class MessageType(Enum):
    AI = "AI"
    USER = "User"


@dataclass
class MessagePayload:
    text: str | None = None
    code: str | None = None
    df: pd.DataFrame | None = None
    active: bool | None = None


@dataclass
class Message:
    message_type: MessageType
    payload: MessagePayload
    timestamp: datetime.datetime = field(default_factory=datetime.datetime.now)

app/chat/data_models.py
mowa-ai/sql-copilot
import base64
import logging
from pathlib import Path

import pandas as pd
import sqlparse
import streamlit as st
from code_editor import code_editor
from PIL import Image

from app.chat import callbacks, config
from app.chat.data_models import Message, MessagePayload, MessageType
from app.engine.engine import execute_query

logger = logging.getLogger(__name__)
URL_TO_ICON = "https://storage.googleapis.com/assets_mvp/Icon.png"
pd.options.display.float_format = "{:.2f}".format


code_buttons = [
    {
        "name": "Run",
        "feather": "Play",
        "alwaysOn": True,
        "primary": True,
        "hasText": True,
        "showWithIcon": True,
        "commands": ["submit"],
        "style": {"bottom": "0.44rem", "right": "0.4rem"},
    },
]


def setup_look(
    page_title,
    icon_path="imgs/Icon.png",
):
    im = Image.open(icon_path)
    st.set_page_config(page_title=page_title, page_icon=im, initial_sidebar_state="expanded")
    hide_footer_style = """
       <style>
       #MainMenu {visibility: hidden; }
       footer {visibility: hidden;}
       </style>
       """
    st.markdown(hide_footer_style, unsafe_allow_html=True)


def run_code_editor(sql_code):
    response_dict = code_editor(
        sql_code,
        lang="sql",
        buttons=code_buttons,
        focus=True,
    )

    if response_dict["type"] == "submit" and len(response_dict["text"]) != 0:
        try:
            df = execute_query(response_dict["text"])
            df = df.loc[:, ~df.columns.duplicated()]
            st.dataframe(df, hide_index=True)
        except Exception as exeption:
            st.write(str(exeption))
            return

        ai_message = Message(
            message_type=MessageType.AI,
            payload=MessagePayload(code=response_dict["text"], df=df, active=False),
        )
        callbacks.update_session_state_on_response(ai_message, update=True)
        response_dict = None
        st.rerun()


def show_bot_message(msg: Message, active: bool) -> None:
    if msg.payload.code or msg.payload.text:
        with st.chat_message("bot", avatar=URL_TO_ICON):
            if msg.payload.text:
                st.write(msg.payload.text)
            if msg.payload.code:
                sql_code = sqlparse.format(msg.payload.code, reindent=True, keyword_case="upper")
                if active:
                    run_code_editor(sql_code)
                else:
                    st.code(sql_code, language="sql")
                if msg.payload.df is not None:
                    st.dataframe(
                        msg.payload.df,
                        hide_index=True,
                    )


def show_messages() -> None:
    messages = st.session_state[config.BOT_MESSAGE] + st.session_state[config.USER_MESSAGE]
    messages.sort(key=lambda x: x.timestamp)
    for msg in messages:
        if msg.message_type == MessageType.AI:
            show_bot_message(msg, msg.payload.active)
        else:
            with st.chat_message("user"):
                st.write(msg.payload.text)

app/chat/utils.py
mowa-ai/sql-copilot
from typing import List

import pandas as pd

from app.engine.llm_chain import QueryChain
from app.engine.sql_engine import Database


def generate_query(text: str, history: List = None, lang: str = "English") -> str:
    db = Database()
    chain = QueryChain(db.schema, lang=lang)
    output = chain.generate_query(text, history=history)
    return output


def execute_query(query: str) -> pd.DataFrame:
    db = Database()
    output = db.execute_query(query=query)
    return output

app/engine/engine.py
mowa-ai/sql-copilot
from typing import List

from langchain.chains import LLMChain
from langchain.llms import OpenAI
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
from langchain.schema.messages import AIMessage, BaseMessage, HumanMessage
from loguru import logger

from app.engine.prompts import LLM_TEMPLATE_EN, LLM_TEMPLATE_PL

USE_LAST_N_MESSAGES_AS_MEMORY = 10


class QueryChain:
    def __init__(self, schema, lang):
        self.llm = OpenAI(temperature=0)

        if lang == "English":
            self.template = LLM_TEMPLATE_EN.format(schema=schema)
        elif lang == "Polish":
            self.template = LLM_TEMPLATE_PL.format(schema=schema)
        else:
            raise ValueError("Wrong lang value, it should be English or Polish.")
        self.prompt = PromptTemplate.from_template(self.template)

    def process_messages(self, history: List[dict]) -> List[BaseMessage]:
        memory = []
        for message in history[-USE_LAST_N_MESSAGES_AS_MEMORY:]:
            if message["type"] == "Human":
                new_record = HumanMessage(content=message["text"])
            elif message["type"] == "AI":
                new_record = AIMessage(content=message["code"])
            else:
                raise ValueError("Wrong message type in the history")
            memory.append(new_record)
        return memory

    def generate_query(self, question: str, history: List) -> str:
        logger.info("Query generation: START.")
        memory = ConversationBufferMemory(memory_key="chat_history")
        memory.chat_memory.messages = self.process_messages(history=history)
        chain = LLMChain(llm=self.llm, prompt=self.prompt, verbose=True, memory=memory)
        answer = chain({"question": question})
        answer = answer["text"]
        logger.info("Query generation: DONE")
        return answer

app/engine/llm_chain.py
mowa-ai/sql-copilot
LLM_TEMPLATE_EN = """
Here is a database schema:
department:
CREATE TABLE "department" (
"Department_ID" INTEGER PRIMARY KEY,
  "Name" TEXT,
  "Creation" TEXT,
  "Ranking" INTEGER,
  "Budget_in_Billions" REAL,
  "Num_Employees" REAL
)

head:
CREATE TABLE "head" (
"head_ID" INTEGER PRIMARY KEY,
  "name" TEXT,
  "born_state" TEXT,
  "age" REAL
)

management:
CREATE TABLE "management" (
"department_ID" INTEGER PRIMARY KEY,
  "head_ID" INTEGER,
  "temporary_acting" TEXT
)

Previous conversation:

Please write me a SQL statement that answers the following question: How many heads of the departments are older than 56 ? [/INST] select count ( * ) from head where age > 56; [INST] Here is a database schema:
book_club:
CREATE TABLE "book_club" (
"book_club_id" INTEGER PRIMARY KEY,
  "Year" INTEGER,
  "Author_or_Editor" TEXT,
  "Book_Title" TEXT,
  "Publisher" TEXT,
  "Category" TEXT,
  "Result" TEXT
)

culture_company:
CREATE TABLE "culture_company" (
"Company_name" TEXT PRIMARY KEY,
  "Type" TEXT,
  "Incorporated_in" TEXT,
  "Group_Equity_Shareholding" REAL,
  "book_club_id" TEXT,
  "movie_id" TEXT
)

movie:
CREATE TABLE "movie" (
"movie_id" INTEGER PRIMARY KEY,
  "Title" TEXT,
  "Year" INTEGER,
  "Director" TEXT,
  "Budget_million" REAL,
  "Gross_worldwide" INTEGER
)

Previous conversation:
Human: Show me movies for 2022
AI: SELECT Title,
       Year
FROM movie
WHERE Year = "2022";
Human: What is the number of movies per year?
AI: SELECT count(*)
FROM cust_order
GROUP BY Year;
Human: Show me company movie pairs
AI: SELECT Company_name,
       Title
FROM culture_company as c
LEFT JOIN movie as m
ON c.movie_id = m.movie_id;

Please write me a SQL statement that answers the following question: What are all company names that have a corresponding movie directed in the year 1999? [/INST] select culture_company.company_name from movie inner join culture_company on movie.movie_id = culture_company.movie_id where movie.year = 1999; [INST] Here is a database schema:
station:
CREATE TABLE "station" (
"id" INTEGER PRIMARY KEY,
  "name" TEXT,
  "lat" REAL,
  "long" REAL,
  "dock_count" INTEGER,
  "city" TEXT,
  "installation_date" TEXT
)

status:
CREATE TABLE "status" (
"station_id" INTEGER,
  "bikes_available" INTEGER,
  "time" TEXT,
)

trip:
CREATE TABLE "trip" (
"id" INTEGER PRIMARY KEY,
  "duration" INTEGER,
  "start_date" TEXT,
  "start_station_name" TEXT,
  "start_station_id" INTEGER,
  "end_date" TEXT,
  "end_station_name" TEXT,
  "end_station_id" INTEGER,
  "bike_id" INTEGER,
  "subscription_type" TEXT,
  "zip_code" INTEGER,
)

weather:
CREATE TABLE "weather" (
"date" TEXT,
  "max_temperature_f" INTEGER,
  "mean_temperature_f" INTEGER,
  "min_temperature_f" INTEGER,
  "mean_dew_point_f" INTEGER,
  "max_dew_point_f" INTEGER,
  "min_dew_point_f" INTEGER,
  "max_humidity" INTEGER,
  "mean_humidity" INTEGER,
  "max_sea_level_pressure_inches" REAL,
  "mean_sea_level_pressure_inches" REAL,
  "min_sea_level_pressure_inches" REAL,
  "max_visibility_miles" INTEGER,
  "mean_visibility_miles" INTEGER,
  "min_visibility_miles" INTEGER,
  "max_wind_Speed_mph" INTEGER,
  "mean_wind_speed_mph" INTEGER,
  "max_gust_speed_mph" INTEGER,
  "precipitation_inches" INTEGER,
  "cloud_cover" TEXT,
  "events" REAL,
  "wind_dir_degrees" INTEGER,
  "zip_code" INTEGER,
)


Please write me a SQL statement that answers the following question: What are the ids of stations that have latitude above 37.4 and never had bike availability below 7? [/INST] select id from station where lat > 37.4 except select station_id from status group by station_id having min ( bikes_available ) < 7; [INST] Here is a database schema:
{schema}

Previous conversation:
{{chat_history}}

Please write me a SQLite statement that answers the following question: {{question}}
Do not provide any commentary or explanation of what the code does, just the SQL statement ending in a semicolon.
Answer:
"""

LLM_TEMPLATE_PL = """
Oto schemat bazy danych:
department:
CREATE TABLE "department" (
"Department_ID" INTEGER PRIMARY KEY,
  "Name" TEXT,
  "Creation" TEXT,
  "Ranking" INTEGER,
  "Budget_in_Billions" REAL,
  "Num_Employees" REAL
)

head:
CREATE TABLE "head" (
"head_ID" INTEGER PRIMARY KEY,
  "name" TEXT,
  "born_state" TEXT,
  "age" REAL
)

management:
CREATE TABLE "management" (
"department_ID" INTEGER PRIMARY KEY,
  "head_ID" INTEGER,
  "temporary_acting" TEXT
)

Poprzednia konwersacja:

Proszę o napisanie instrukcji SQL odpowiadającej na pytanie: Ilu kierowników działów ma więcej niż 56 lat? [/INST] select count ( * ) from head where age > 56; [INST] Oto schemat bazy danych:
book_club:
CREATE TABLE "book_club" (
"book_club_id" INTEGER PRIMARY KEY,
  "Year" INTEGER,
  "Author_or_Editor" TEXT,
  "Book_Title" TEXT,
  "Publisher" TEXT,
  "Category" TEXT,
  "Result" TEXT
)

culture_company:
CREATE TABLE "culture_company" (
"Company_name" TEXT PRIMARY KEY,
  "Type" TEXT,
  "Incorporated_in" TEXT,
  "Group_Equity_Shareholding" REAL,
  "book_club_id" TEXT,
  "movie_id" TEXT
)

movie:
CREATE TABLE "movie" (
"movie_id" INTEGER PRIMARY KEY,
  "Title" TEXT,
  "Year" INTEGER,
  "Director" TEXT,
  "Budget_million" REAL,
  "Gross_worldwide" INTEGER
)

Poprzednia konwersacja:
Human: Pokaż mi filmy z 2022
AI: SELECT Title,
       Year
FROM movie
WHERE Year = "2022";
Human: Jaka jest liczba filmów rocznie?
AI: SELECT count(*)
FROM cust_order
GROUP BY Year;
Human: Pokaż mi pary filmów firma
AI: SELECT Company_name,
       Title
FROM culture_company as c
LEFT JOIN movie as m
ON c.movie_id = m.movie_id;

Proszę o napisanie instrukcji SQL odpowiadającej na pytanie: Jakie są nazwy wszystkich firm, którym odpowiada film wyreżyserowany w 1999 roku? [/INST] select culture_company.company_name from movie inner join culture_company on movie.movie_id = culture_company.movie_id where movie.year = 1999; [INST] Oto schemat bazy danych:
station:
CREATE TABLE "station" (
"id" INTEGER PRIMARY KEY,
  "name" TEXT,
  "lat" REAL,
  "long" REAL,
  "dock_count" INTEGER,
  "city" TEXT,
  "installation_date" TEXT
)

status:
CREATE TABLE "status" (
"station_id" INTEGER,
  "bikes_available" INTEGER,
  "time" TEXT,
)

trip:
CREATE TABLE "trip" (
"id" INTEGER PRIMARY KEY,
  "duration" INTEGER,
  "start_date" TEXT,
  "start_station_name" TEXT,
  "start_station_id" INTEGER,
  "end_date" TEXT,
  "end_station_name" TEXT,
  "end_station_id" INTEGER,
  "bike_id" INTEGER,
  "subscription_type" TEXT,
  "zip_code" INTEGER,
)

weather:
CREATE TABLE "weather" (
"date" TEXT,
  "max_temperature_f" INTEGER,
  "mean_temperature_f" INTEGER,
  "min_temperature_f" INTEGER,
  "mean_dew_point_f" INTEGER,
  "max_dew_point_f" INTEGER,
  "min_dew_point_f" INTEGER,
  "max_humidity" INTEGER,
  "mean_humidity" INTEGER,
  "max_sea_level_pressure_inches" REAL,
  "mean_sea_level_pressure_inches" REAL,
  "min_sea_level_pressure_inches" REAL,
  "max_visibility_miles" INTEGER,
  "mean_visibility_miles" INTEGER,
  "min_visibility_miles" INTEGER,
  "max_wind_Speed_mph" INTEGER,
  "mean_wind_speed_mph" INTEGER,
  "max_gust_speed_mph" INTEGER,
  "precipitation_inches" INTEGER,
  "cloud_cover" TEXT,
  "events" REAL,
  "wind_dir_degrees" INTEGER,
  "zip_code" INTEGER,
)

Proszę o napisanie instrukcji SQL odpowiadającej na pytanie: Jakie są identyfikatory stacji, które mają szerokość geograficzną powyżej 37,4 i nigdy nie miały dostępności rowerów poniżej 7? [/INST] select id from station where lat > 37.4 except select station_id from status group by station_id having min ( bikes_available ) < 7; [INST] Oto schemat bazy danych:
{schema}

Poprzednia konwersacja:
{{chat_history}}

Proszę o napisanie instrukcji SQL odpowiadającej na pytanie: {{question}}
Nie podawaj żadnego komentarza ani wyjaśnienia działania kodu, wystarczy instrukcja SQL zakończona średnikiem.
Odpowiedź:
"""

app/engine/prompts.py
mowa-ai/sql-copilot
import os
import sqlite3

import pandas as pd

PATH_TO_DATA = os.environ["PATH_TO_DATA"]


class Database(object):
    def __init__(self):
        self.con = sqlite3.connect(PATH_TO_DATA)
        self.tables = None
        self.schema, self.schema_dict = None, None
        self.update_schema()

    def execute_query(self, query: str) -> pd.DataFrame:
        df = pd.read_sql_query(query, self.con)
        return df

    def update_schema(self, chosen_tables=None) -> None:
        sql = "SELECT name FROM sqlite_master WHERE type='table';"
        self.tables = chosen_tables or [x[0] for x in self.execute_query(sql).values]
        self.schema = ""
        self.schema_dict = {}
        for t in self.tables:
            cursor = self.con.execute("SELECT sql FROM sqlite_master WHERE name=?;", [t])
            sql = cursor.fetchone()[0]
            cursor.close()
            self.schema += f"{t}:\n" + sql + "\n\n"
            self.schema_dict[t] = pd.read_sql_query(f"PRAGMA table_info({t})", con=self.con)[["name", "type"]]

app/engine/sql_engine.py
mowa-ai/sql-copilot