Update gradio1final.py

gradio1final.py

@@ -1,685 +1,685 @@
-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
-
-
-
-
-os.environ['GOOGLE_API_KEY'] = 'AIzaSyDLwJAr5iXL2Weaw1XphFNSeijytqOSbDg'
-os.environ['OPENAI_API_KEY'] = 'sk-proj-Kh4UIWkfDxSGppQpooxXT3BlbkFJATohXqhpkJE6MqliIkmU'
-# Set your API keys
-
-OPENAI_API_KEY = 'sk-proj-Kh4UIWkfDxSGppQpooxXT3BlbkFJATohXqhpkJE6MqliIkmU'
-GOOGLE_API_KEY = 'AIzaSyDLwJAr5iXL2Weaw1XphFNSeijytqOSbDg'
-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='postgres',
-            user='shivanshu',
-            password='root',
-            host='34.170.181.105',
-            port='5432'
-        )
-        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(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}")
-
-
-
-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:
-        # 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), "", []
-   
-
-
-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.Image(value=image_path, show_label=False, height=50, width=50)
-            # 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
+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
+
+
+
+
+os.environ['GOOGLE_API_KEY'] = 'AIzaSyDLwJAr5iXL2Weaw1XphFNSeijytqOSbDg'
+os.environ['OPENAI_API_KEY'] = 'sk-proj-Kh4UIWkfDxSGppQpooxXT3BlbkFJATohXqhpkJE6MqliIkmU'
+# Set your API keys
+
+OPENAI_API_KEY = 'sk-proj-Kh4UIWkfDxSGppQpooxXT3BlbkFJATohXqhpkJE6MqliIkmU'
+GOOGLE_API_KEY = 'AIzaSyDLwJAr5iXL2Weaw1XphFNSeijytqOSbDg'
+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='postgres',
+            user='shivanshu',
+            password='root',
+            host='34.170.181.105',
+            port='5432'
+        )
+        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(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}")
+
+
+
+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:
+        # 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"C:\Users\shivanshu.t\Downloads\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.Image(value=image_path, show_label=False, height=50, width=50)
+            # 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)