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)
google_api_key = os.environ.get('GOOGLE_API_KEY')
openai_api_key = os.environ.get('OPENAI_API_KEY')
# Ensure that the API keys are set
if not google_api_key or not openai_api_key:
raise ValueError("API keys for Google and OpenAI must be set as environment variables.")
# Set your API keys
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'
')
gr.Image(value=image_path, show_label=False, height=50, width=50)
# gr.HTML(f'
')
gr.Markdown(
"""
# Natural Language Query for Network Data
This app generates SQL queries from user queries using Google Gemini or OpenAI models.
Click here to view the data:
View Spreadsheet
"""
)
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)