workers.py

from PySide6.QtCore import QThread, Signal
from unstructured.partition.auto import partition
import chromadb
from chromadb.config import DEFAULT_TENANT, DEFAULT_DATABASE, Settings
import json
import pandas as pd

import models


class DatabaseConnectionWorker(QThread):
    connection_result = Signal(
        bool, object, object, object, object, object, object, object
    )

    def __init__(self, host, dbname, user, password, port, api_key, db_type):
        super().__init__()
        self.host = host
        self.dbname = dbname
        self.user = user
        self.password = password
        self.port = port
        self.api_key = api_key
        self.db_type = db_type

    def run(self):
        try:
            db_connection = models.init_database(
                self.user,
                self.password,
                self.host,
                self.port,
                self.dbname,
                self.db_type,
            )
            sql_model = models.get_sql_chain(self.api_key)
            explain_model = models.get_response(self.api_key)
            visual_model = models.generate_plotly_code(self.api_key)
            ehancer_model = models.better_question(self.api_key)
            table_names = db_connection.get_usable_table_names()
            table_names_description = {
                name: db_connection.get_table_info([name]) for name in table_names
            }
            print("The Table Names:",table_names)
            self.connection_result.emit(
                True,
                db_connection,
                sql_model,
                explain_model,
                table_names_description,
                visual_model,
                ehancer_model,
                table_names,
            )
        except Exception as e:
            print(f"Error: {e}")
            self.connection_result.emit(
                False, None, None, None, None, None, None, None
            ) 


class AIProcessingWorkerSQL(QThread):
    response_generated = Signal(
        str, object, object
    ) 
    def __init__(self, sql_model, user_text, chat_history, db, schema, db_type):
        super().__init__()
        self.sql_model = sql_model
        self.chat_history = chat_history
        self.user_text = user_text
        self.db = db
        self.data = None
        self.schema = schema
        self.sql_command = None
        self.db_type = db_type

    def db_run(self):
        try:
            self.data = self.db.run(self.sql_command, fetch="cursor").fetchall()
            if len(self.data) == 0:
                return "the output was empty"
            return "work done"
        except Exception as e:
            print(f"Error: {e}")
            self.data = "the SQL query encountered an error"
            return e

    def run(self):
        query_result = ""
        for attempt in range(5):
            print("sql worker++++", attempt)
            ai_response = self.sql_model.invoke(
                {
                    "schema": self.schema,
                    "chat_history": self.chat_history,
                    "question": self.user_text,
                    "db_type": self.db_type,
                }
            )["text"]
            ai_response = ai_response.replace("```sql", "").replace("```", "")
            self.sql_command = ai_response
            query_result = self.db_run()
            print(self.sql_command)

            if query_result == "work done":
                break

            ai_response = self.sql_model.invoke(
                {
                    "schema": self.schema,
                    "chat_history": self.chat_history,
                    "question": "For this question: "
                    + str(self.user_text)
                    + f" the query '{self.sql_command}' encountered an error: {query_result}",
                    "db_type": self.db_type,
                }
            )["text"]
            query_result = self.db_run()
            print(self.sql_command)

            if query_result == "work done":
                break
        print("sql worker++++", query_result)
        self.response_generated.emit(ai_response, self.data, self.user_text)


class AIProcessingWorkerExplain(QThread):
    responce_generated = Signal(str)

    def __init__(
        self, explain_model, user_text, sql_command, data, chat_history, schema
    ):
        super().__init__()
        self.chat_history = chat_history
        self.explain_model = explain_model
        self.user_text = user_text
        self.sql_command = sql_command
        self.data = data
        self.schema = schema

    def run(self):
        formatted_string = self.explain_model.invoke(
            {
                "schema": self.schema,
                "chat_history": self.chat_history,
                "query": self.sql_command,
                "question": self.user_text["enhanced_question"],
                "response": self.data,
            }
        )["text"]
        self.responce_generated.emit(formatted_string)


class AIProcessingWorkerUserExpresstion(QThread):
    responce_ehancer_ = Signal(object)

    def __init__(self, user_input, model, chat_history, schema, chroma_db):
        super().__init__()
        self.user_input = user_input
        self.model = model
        self.chat_history = chat_history
        self.table_names = schema
        if chroma_db is not None:
            self.chroma_db = chroma_db
        else:
            self.chroma_db = None

    def clean_json(self, text):
        return text.replace("```json", "").replace("```", "")

    def check_if_collection_exist(self):
        try:
            self.chroma_db.get_collection("sql_data")
            return True
        except Exception as e:
            return False

    def run(self):
        try:
            document_info = ""
            if self.chroma_db is not None and self.check_if_collection_exist():
                document_info = self.chroma_db.get_collection("sql_data").query(
                    query_texts=[self.user_input],
                    n_results=5,
                )
                print(document_info)

            for attempt in range(5):
                response = self.model.invoke(
                    {
                        "user_input": self.user_input,
                        "chat_history": self.chat_history,
                        "tables": self.table_names,
                        "documnet": document_info
                        if document_info == ""
                        else document_info["documents"],
                    }
                )["text"]
                cleaned_response = self.clean_json(response)
                print(cleaned_response)
                try:
                    result = json.loads(cleaned_response)
                    self.responce_ehancer_.emit(result)
                    break
                except json.JSONDecodeError as decode_error:
                    # Retry with updated input if there's an error
                    retry_input = f"Question: {self.user_input}, Response: {cleaned_response}, Error: {decode_error}"
                    retry_response = self.model.invoke(
                        {
                            "user_input": retry_input,
                            "chat_history": self.chat_history,
                            "tables": self.table_names,
                            "documnet": document_info
                            if document_info == ""
                            else document_info["documents"],
                        }
                    )["text"]
                    cleaned_retry_response = self.clean_json(retry_response)
                    result = json.loads(cleaned_retry_response)
                    self.responce_ehancer_.emit(result)
                    break

        except Exception as e:
            print(f"Error: {e}")
            self.responce_ehancer_.emit("User API key is invalid")


class AIProcessingWorkerGraph(QThread):
    responce_generated = Signal(str)

    def __init__(self, visual_model, user_text, data: pd.DataFrame, sql, complate=""):
        super().__init__()
        self.visual_model = visual_model
        self.user_text = user_text
        self.data = data
        self.complate = complate
        self.sql_command = sql

    def run(self):
        formatted_string = self.visual_model.invoke(
            {
                "question": str(self.user_text)
                if self.complate == ""
                else str(self.user_text) + " do not for get that " + self.complate,
                "sql": self.sql_command,
                "df_metadata": self.data.dtypes,
                "sample_data": self.data.head(),
            }
        )["text"]
        self.responce_generated.emit(
            formatted_string.replace("```python", "").replace("```", "")
        )


class ChromaDBClient(QThread):
    responce_generated = Signal(object)

    def __init__(self, document):
        super().__init__()
        self.document = document
        self.client = chromadb.PersistentClient(
            settings=Settings(),
            tenant=DEFAULT_TENANT,
            database=DEFAULT_DATABASE,
        )

    def run(self):
        try:
            print("process started")
            element_dict = partition(self.document)
            print("document partitioned")
            client = chromadb.PersistentClient(
                path="chroma",  # Explicitly set the path
                settings=Settings(),
                tenant=DEFAULT_TENANT,
                database=DEFAULT_DATABASE,
            )
            collection = client.get_or_create_collection(
                name="sql_data", metadata={"hnsw:space": "cosine"}
            )
            k = 1
            for element in element_dict:
                collection.upsert(
                    documents=[element.text],
                    ids=[str(k)],
                )
                k += 1
            print("process complete")
            self.responce_generated.emit(client)
        except Exception as e:
            print(f"An error occurred: {e}")
            self.responce_generated.emit(None)