app.py

import os
import re
import json
import openai
import psycopg2
import google.generativeai as genai
import gradio as gr
from langchain_community.document_loaders.csv_loader import CSVLoader
from langchain.vectorstores import FAISS
from langchain.embeddings.base import Embeddings
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
from langchain_core.prompts import ChatPromptTemplate
from langchain.output_parsers.json import SimpleJsonOutputParser
from langchain.memory import ConversationBufferMemory
from langchain.chains import LLMChain
from langchain_core.prompts import (
    ChatPromptTemplate,
    MessagesPlaceholder,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain import OpenAI



OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')
GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')
openai.api_key = OPENAI_API_KEY
genai.configure(api_key=GOOGLE_API_KEY)




class OpenAIEmbeddings(Embeddings):
    def embed_documents(self, texts):
        response = openai.Embedding.create(
            model="text-embedding-ada-002",
            api_key=OPENAI_API_KEY,
            input=texts
        )
        embeddings = [e["embedding"] for e in response["data"]]
        return embeddings




    def embed_query(self, text):
        response = openai.Embedding.create(
            model="text-embedding-ada-002",
            api_key=OPENAI_API_KEY,
            input=[text]
        )
        embedding = response["data"][0]["embedding"]
        return embedding




class GeminiEmbeddings(GoogleGenerativeAIEmbeddings):
    def __init__(self, api_key):
        super().__init__(model="models/text-embedding-004", google_api_key=GOOGLE_API_KEY)




def extract_entities_openai(query):
    prompt = f"""
Find the entities from the query given below enclosed in triple quotes. Make sure to ONLY return response in JSON format, with the key as "KPI" and extracted entities as list. Example Output JSON:
{{
  "KPI": ["Extracted Entity 1", "Extracted Entity 2", ....]
}}


Query begins here:
\"\"\"{query}\"\"\"
"""
    response = openai.ChatCompletion.create(
        model="gpt-4",
        messages=[
            {"role": "user", "content": prompt}
        ],
        temperature=0
    )
    return response['choices'][0]['message']['content'].strip()



def extract_entities_gemini(query):
    try:
        llm_content_summary = ChatGoogleGenerativeAI(
            model="gemini-1.5-pro",
            temperature=0,
            max_tokens=250,
            timeout=None,
            max_retries=2,
            google_api_key=GOOGLE_API_KEY,
        )




        prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "system", """Use the instructions below to generate responses based on user inputs. Return the answer as a JSON object.""",
                ),
                ("user", f"""Find the entities from the query given below enclosed in triple quotes. Make sure to ONLY return response in JSON format, with the key as "KPI" and extracted entities as list. Example Output JSON:
"KPI": ["Extracted Entity 1", "Extracted Entity 2", ....]
Query begins here:
\"\"\"{query}\"\"\"
"""),
            ]
        )




        json_parser = SimpleJsonOutputParser()
        chain = prompt | llm_content_summary | json_parser
        response = chain.invoke({"input": query})
       
        if isinstance(response, dict):
            response = json.dumps(response)
       
        return response
    except Exception as e:
        print(f"Error generating content summary: {e}")
        return None




def fetch_table_schema():
    try:
        conn = psycopg2.connect(
            dbname= os.getenv('dbname'),
            user=os.getenv('user'),
            password=os.getenv('password'),
            host=os.getenv('host'),
            port=os.getenv('port')
        )
        cursor = conn.cursor()
        table_name = 'network'
        query = f"""
        SELECT
            column_name,
            data_type,
            character_maximum_length,
            is_nullable
        FROM
            information_schema.columns
        WHERE
            table_name = '{table_name}';
        """
        cursor.execute(query)
        rows = cursor.fetchall()
        cursor.close()
        schema_dict = {
            row[0]: {
                'data_type': row[1],
                'character_maximum_length': row[2],
                'is_nullable': row[3]
            }
            for row in rows
        }
        return schema_dict
    except Exception as e:
        print(f"Error fetching table schema: {e}")
        return {}




column_names=[]

def extract_column_names(sql_query):
    # Use a regular expression to extract the part of the query between SELECT and FROM
    pattern = r'SELECT\s+(.*?)\s+FROM'
    match = re.search(pattern, sql_query, re.IGNORECASE | re.DOTALL)
    if not match:
        return []


    columns_part = match.group(1).strip()

    # Split columns based on commas that are not within parentheses
    column_names = re.split(r',\s*(?![^()]*\))', columns_part)

    # Process each column to handle aliases and functions
    clean_column_names = []
    for col in column_names:
        # Remove any function wrappers (e.g., TRIM, COUNT, etc.)
        col = re.sub(r'\b\w+\((.*?)\)', r'\1', col)

        # Remove any aliases (i.e., words following 'AS')
        col = re.split(r'\s+AS\s+', col, flags=re.IGNORECASE)[-1]

        # Strip any remaining whitespace or backticks/quotes
        col = col.strip(' `"[]')
        clean_column_names.append(col)
    return clean_column_names


def process_sublist(sublist,similarity_threshold):
            processed_list = sublist[0:3]
            for index in range(3, len(sublist)):
                if sublist[index]['similarity'] >= similarity_threshold:
                    processed_list.append(sublist[index])
                else:
                    break
            return processed_list


from langchain_community.chat_models import ChatOpenAI
# Initialize the OpenAI model and memory
openai_model = ChatOpenAI(model='gpt-4', temperature=0, api_key=OPENAI_API_KEY)
openai_memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

def generate_sql_query_openai(description, table_schema, vector_store, og_table_schema):
    global openai_model
    global openai_memory
    global column_names
    try:
        user_input = description
        print(user_input)
        entities = extract_entities_openai(user_input)
        entities_obj = json.loads(entities)
        kpis = entities_obj['KPI']

        # Fetch similar documents
        final_results = []
        similar_kpis = []
        for kpi in kpis:
            docs = vector_store.similarity_search_with_score(kpi, k=6)
            results = []
            count = 1
            for doc, distance in docs:
                name_desc = doc.page_content.split("\nDescription: ")
                name = name_desc[0].replace("Name: ", "")
                description = name_desc[1] if len(name_desc) > 1 else "No description available."
                similarity_score = round((1 - distance) * 100, 2)  # Convert distance to similarity score in percent
                results.append({"name": name, "description": description, "similarity": similarity_score, "Index": count})
                similar_kpis.append({"name": name, "similarity": similarity_score})
                count += 1  # Increment the count for each result
            final_results.append(results)

        # Process results
        similarity_threshold = 75.0

        processed_sublists = [process_sublist(sublist,similarity_threshold) for sublist in final_results]
        flattened_results = [item for sublist in processed_sublists for item in sublist]
       
        user_input_entities = [item['name'] for item in flattened_results]
        print("BYE1",user_input_entities)
        
        try:
            # Strip whitespace and ensure case matches for comparison
            user_input_entities = [key.strip() for key in user_input_entities]
            og_table_schema_keys = [key.strip() for key in og_table_schema.keys()]

            # Check and create user_input_table_schema
            table_schema = {key: og_table_schema[key] for key in user_input_entities if key in og_table_schema_keys}
            print("BYE3",table_schema)
        except Exception as e:
            print(f"An error occurred: {e}")


        table_schema = json.dumps(table_schema)
        table_schema = table_schema.replace('{', '[')
        table_schema = table_schema.replace('}', ']')
        print("GG123")
        system_message_template1 = f"""Generate the PostgreSQL query for the following task: {description}.
        The connection with the database is already setup and the table is called network.
        Enclose column names in double quotes ("), but do not use escape characters (e.g., "\").
        Do not assign aliases to the columns.
        Do not calculate new columns, unless specifically called to.
        Return only the PostgreSQL query, nothing else.
        The list of all the columns is as follows: {table_schema}
        Make sure the response should strictly follow JSON format. The key should be "Query" and the value should be the Postgresql query.
        Example Output JSON:
        ["Query": PostgreSQL Executable query]"""
               
        system_message_template = system_message_template1
        print(system_message_template)
        # Create the ChatPromptTemplate
        print("GG1234")
        prompt = ChatPromptTemplate(
            messages=[
                SystemMessagePromptTemplate.from_template(system_message_template),
                # Placeholder for chat history
                MessagesPlaceholder(variable_name="chat_history"),
                # User's question will be dynamically inserted
                HumanMessagePromptTemplate.from_template("""
                    {question}
                """)
            ]
        )



        print("GG12345")
        conversation = LLMChain(
            llm=openai_model,
            prompt=prompt,
            verbose=True,
            memory=openai_memory
        )




        print(prompt)
        response = conversation.invoke({'question': user_input})
        response = response['text']
        response = response.replace('\\', '')
        print(response)
        print("************************************************")
        print(type(response))
       
        # Regular expression pattern to extract the query string
        pattern = r'{"Query":\s*"(.*?)"\s*}'
        # Extract the content

        sql_query = None
        match = re.search(pattern, response)
        if match:
            sql_query = match.group(1)
            print(sql_query)
        
        print("jiji1")
        if sql_query:
            # Fetch data from database
            results = fetch_data_from_db(sql_query)
            print(results)
            column_names1 = extract_column_names(sql_query)
            column_names=column_names1
            print("jiji",column_names)

        else:
            column_names=[]
            results=[]

        pattern = r'```.*?```(.*)'
        match = re.search(pattern, response, re.DOTALL)
        print("GG",match)
        if match:
            response = match.group(1).strip()
            print("GG1",response)

        print(type(user_input_entities))
        print(type(response))
        print(sql_query)
        print(type(results))
        print("Process completed.")
        return user_input_entities, response, sql_query, results
   
    except Exception as e:
        print(f"Error generating SQL query: {e}")



gemini_model = ChatGoogleGenerativeAI(model='gemini-1.5-pro-001',temperature=0,google_api_key = GOOGLE_API_KEY)
gemini_memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

def generate_sql_query_gemini(description, table_schema,vector_store,og_table_schema):
    global gemini_model
    global gemini_memory
    global column_names
    try:
        user_input = description
        print(user_input)
        entities = extract_entities_gemini(user_input)
        entities_obj = json.loads(entities)
        kpis = entities_obj['KPI']

# Fetch similar documents
        final_results = []
        similar_kpis = []
        for kpi in kpis:
            docs = vector_store.similarity_search_with_score(kpi, k=6)
            results = []
            count = 1
            for doc, distance in docs:
                name_desc = doc.page_content.split("\nDescription: ")
                name = name_desc[0].replace("Name: ", "")
                description = name_desc[1] if len(name_desc) > 1 else "No description available."
                similarity_score = round((1 - distance) * 100, 2)  # Convert distance to similarity score in percent
                results.append({"name": name, "description": description, "similarity": similarity_score, "Index": count})
                similar_kpis.append({"name": name, "similarity": similarity_score})
                count += 1  # Increment the count for each result
            final_results.append(results)

        # Process results
        similarity_threshold = 75.0

        processed_sublists = [process_sublist(sublist,similarity_threshold) for sublist in final_results]
        flattened_results = [item for sublist in processed_sublists for item in sublist]
       
        user_input_entities = [item['name'] for item in flattened_results]
        print("BYE1",user_input_entities)
        
        try:
            # Strip whitespace and ensure case matches for comparison
            user_input_entities = [key.strip() for key in user_input_entities]
            og_table_schema_keys = [key.strip() for key in og_table_schema.keys()]

            # Check and create user_input_table_schema
            table_schema = {key: og_table_schema[key] for key in user_input_entities if key in og_table_schema_keys}
            print("BYE3",table_schema)
        except Exception as e:
            print(f"An error occurred: {e}")


        table_schema = json.dumps(table_schema)
        table_schema = table_schema.replace('{', '[')
        table_schema = table_schema.replace('}', ']')
        system_message_template1 = f"""Generate the PostgreSQL query for the following task: {description}.
        The connection with the database is already setup and the table is called network.
        Enclose column names in double quotes ("), but do not use escape characters (e.g., "\").
        Do not assign aliases to the columns.
        Do not calculate new columns, unless specifically called to.
        Return only the PostgreSQL query, nothing else.
        The list of all the columns is as follows: {table_schema}
        Make sure the response should strictly follow JSON format. The key should be "Query" and the value should be the Postgresql query.
        Example Output JSON:
        ["Query": PostgreSQL Executable query]"""
               
        system_message_template = system_message_template1
        print(system_message_template)
        # Create the ChatPromptTemplate
        prompt = ChatPromptTemplate(
            messages=[
                SystemMessagePromptTemplate.from_template(system_message_template),
                # Placeholder for chat history
                MessagesPlaceholder(variable_name="chat_history"),
                # User's question will be dynamically inserted
                HumanMessagePromptTemplate.from_template("""
                    {question}
                """)
            ]
        )

        conversation = LLMChain(
            llm=gemini_model,
            prompt=prompt,
            verbose=True,
            memory=gemini_memory
        )


        print(prompt)
        response = conversation.invoke({'question': user_input})
        response = response['text']
        response = response.replace('\\', '')
        print(response)
       
        # Pattern to extract SQL query from the response
        patterns = [
            r"""```json\n{\s*"Query":\s*"(.*?)"}\n```""",
            r"""```json\n{\s*"Query":\s*"(.*?)"}\s*```""",
            r"""```json\s*{\s*"Query":\s*"(.*?)"\s*}\s*```""",
            r"""```json\s*{\s*"Query":\s*"(.*?)"\s*}```""",
            r"""```json\n{\n\s*"Query":\s*"(.*?)"\n}\n```""",
            r"""```json\s*\{\s*['"]Query['"]:\s*['"](.*?)['"]\s*\}\s*```""",
            r"""```json\s*\{\s*['"]Query['"]\s*:\s*['"](.*?)['"]\s*\}\s*```""",
            r"""```json\s*{\s*"Query"\s*:\s*"(.*?)"\s*}\s*```""",
            r"""```json\s*{\s*"Query"\s*:\s*\"(.*?)\"\s*}\s*```""",
            r"""\"Query\"\s*:\s*\"(.*?)\"""",
            r"""```json\s*\{\s*\"Query\":\s*\"(.*?)\"\s*\}\s*```""",
            r"""['"]Query['"]\s*:\s*['"](.*?)['"]""",
            r"""```json\s*{\s*"Query":\s*"(.*?)"}\s*```""",
]
        sql_query = None
        for pattern in patterns:
            matches = re.findall(pattern, response, re.DOTALL)
            if matches:
                sql_query = matches[0]
        
        print("jiji1")
        if sql_query:
            # Fetch data from database
            results = fetch_data_from_db(sql_query)
            print(results)
            column_names1 = extract_column_names(sql_query)
            column_names=column_names1
            print("jiji",column_names)

        else:
            column_names=[]
            results=[]

        pattern = r'```.*?```(.*)'
        match = re.search(pattern, response, re.DOTALL)
        print("GG",match)
        if match:
            response = match.group(1).strip()
            print("GG1",response)

        print(type(user_input_entities))
        print(type(response))
        print(type(sql_query))
        print(type(results))
        print("Process completed.")
        return user_input_entities, response, sql_query, results
   
    except Exception as e:
        print(f"Error generating SQL query: {e}")




def fetch_data_from_db(query):
    try:
        conn = psycopg2.connect(
            dbname='postgres',
            user='shivanshu',
            password='root',
            host='34.170.181.105',
            port='5432'
        )
        cursor = conn.cursor()
        cursor.execute(query)
        results = cursor.fetchall()
        cursor.close()
        conn.close()
        print(results)
        return results
    except Exception as e:
        print(f"Error fetching data from database: {e}")
        return []




def process_gradio(query, model_type):
    try:
        # URL of the CSV file hosted on Hugging Face
        # csv_url = 'https://huggingface.co/spaces/Shivanshutripathi/AutoBI/blob/main/des.csv'
        
        # # Fetch the CSV file content
        # response = requests.get(csv_url)
        # response.raise_for_status()  # Check if the request was successful
        
        # # Load the CSV file content into a pandas DataFrame
        # csv_data = StringIO(response.text)
        # documents = pd.read_csv(csv_data)
        
    # try:
    #     # Load the CSV file
        csv_loader = CSVLoader(file_path='des.csv')
        documents = csv_loader.load()

        # Define the vector DB paths
        vector_db_path_gemini = "faiss_index_gemini"
        vector_db_path_openai = "faiss_index_openai"

        # Check if the directory paths exist, if not, create them
        os.makedirs(vector_db_path_gemini, exist_ok=True)
        os.makedirs(vector_db_path_openai, exist_ok=True)

        # Determine the model to use
        if model_type == 'gemini':
            vector_db_path = vector_db_path_gemini
            embeddings = GoogleGenerativeAIEmbeddings(model="models/text-embedding-004", api_key=GOOGLE_API_KEY)
        else:
            vector_db_path = vector_db_path_openai
            embeddings = OpenAIEmbeddings()

        # Check if the FAISS index already exists
        index_file_path = os.path.join(vector_db_path, "index")
        if os.path.exists(index_file_path):
            vector_store = FAISS.load_local(vector_db_path, embeddings, allow_dangerous_deserialization=True)
        else:
            texts = [doc.page_content for doc in documents]
            vector_store = FAISS.from_texts(texts, embeddings)
            vector_store.save_local(vector_db_path)


        og_table_schema = fetch_table_schema()
        new_table_schema = {}


        # Generate SQL query
        if model_type == 'gemini':
            user_input_entities, response, sql_query, results = generate_sql_query_gemini(query, new_table_schema, vector_store, og_table_schema)
        else:
            user_input_entities, response, sql_query, results = generate_sql_query_openai(query, new_table_schema, vector_store, og_table_schema)


        return user_input_entities or {}, response or "", sql_query or "", results or {}




    except Exception as e:
        # Ensure the function still returns four values, even in case of an error
        return {}, str(e), "", []

# image_path = r"incedo-logo.png"

with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=1):
            # Add the image in the left corner
            # gr.Image(value=image_path, show_label=False, height=45, width=45)
            # gr.HTML(f'<img src="{image_link}" alt="Logo" height="50" width="50">')
            gr.HTML(
                """
                <div style="text-align: left; padding: 10px;">
                    <img src="https://mma.prnewswire.com/media/1807312/incedo_Logo.jpg" alt="Logo" height="50" width="80">
                </div>
                """
            )
            # gr.HTML(f'<img src="{image_path}" alt="Logo" height="50" width="100" style="display: block; margin-left: auto; margin-right: auto;">')
    
    gr.Markdown(
        """
        # Natural Language Query for Network Data
        <p style="font-size: 16px;">
        This app generates SQL queries from user queries using Google Gemini or OpenAI models.
        </p>
        <p style="font-size: 16px;">
            Click here to view the data: 
            <a href="https://docs.google.com/spreadsheets/d/1uYeHbqzz1NKL8e4tlzbIk8K5qgLfY_To-pjRjOGjWQg/edit?usp=sharing" target="_blank" style="color: #0066cc; text-decoration: none;">
                View Spreadsheet
            </a>
        </p>
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            query_input = gr.Textbox(label="Enter your query")
            model_input = gr.Radio(choices=["gemini", "openai"], label="Model Type", value="gemini")
            temperature_slider = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.7, label="Temperature", interactive=True)
            submit_button = gr.Button("Submit")

            gr.Examples(
                examples=[
                    ["What is the average latency for each network engineer?", "openai"],
                    ["Can you tell me what is the average RRC setup attempts, how is it distributed across time of the day? Can you also show me this distribution both for weekdays as well as weekends. Show me the data sorted by time of the day.", "openai"],
                    ["What is the average PRB utilization for each Network engineer, and give me this average both for the weekends as well for the weekdays. Show this for each network engineer averaged across four time periods, 12 midnight  - 6am , 6am -12 noon , 12 noon - 6pm , 6pm - 12 midnight.  Finally show me this data sorted by network engineer as well time periods.", "openai"]
                ],
                inputs=[query_input, model_input]
            )
       
        with gr.Column(scale=2):
            output_results = gr.DataFrame(label="Query Results")
            output_sql_query = gr.Textbox(label="Generated SQL Query")
            output_response = gr.Textbox(label="Similar Entities")
            # output_user_input_schema = gr.DataFrame(label="User Input Schema")
            output_user_input_schema = gr.JSON(label="Retrived KPIs")


           
    # Define the button click action
    def update_dataframe(query_input, model_input):
        global column_names

        # Process the query and model input
        user_input_entities, response, sql_query, results = process_gradio(query_input, model_input)
        
        # Check if column_names is not empty
        if column_names:
            output_results.headers = column_names  # Set headers with dynamic column names
        else:
            output_results.headers = []  # Set headers to an empty list for an empty DataFrame

        # Return the processed results
        return user_input_entities, response, sql_query, results




    def update_dataframe(query_input, model_input):
        global column_names

        # Process the query and model input
        user_input_entities, response, sql_query, results = process_gradio(query_input, model_input)
        
        # Check if column_names is not empty
        if column_names:
            output_results.headers = column_names  # Set headers with dynamic column names
        else:
            output_results.headers = []  # Set headers to an empty list for an empty DataFrame

        # Return the processed results
        return user_input_entities, response, sql_query, results


    submit_button.click(
        fn=update_dataframe,
        inputs=[query_input, model_input],
        outputs=[output_user_input_schema, output_response, output_sql_query, output_results]
    )

# Launch the app
demo.launch(debug=True)