app.py

import streamlit as st
import streamlit_chat
import json
import os
import datetime
from pymongo import MongoClient
from bson import ObjectId
from dotenv import load_dotenv
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import ChatPromptTemplate

st.set_page_config(layout="wide", page_title="IOCL Chatbot", page_icon="📄")
load_dotenv()

# MongoDB connection setup
MONGO_URI = os.getenv("MONGO_URI")
client = MongoClient(MONGO_URI)
db = client["chatbot_db"]
chat_sessions = db["chat_sessions"]

# Set LLM models
FLASH_API = os.getenv("FLASH_API")
embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001", google_api_key=FLASH_API)
llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0, max_tokens=None, google_api_key=FLASH_API)

# Load the extracted JSON data
with open('iocl_extracted_2.json', 'r') as file:
    extracted_data = json.load(file)
reference_data = json.dumps(extracted_data)

# Initialize session state for current chat session
if 'current_chat_id' not in st.session_state:
    st.session_state['current_chat_id'] = None
if 'chat_history' not in st.session_state:
    st.session_state['chat_history'] = []
if 'regenerate' not in st.session_state:
    st.session_state['regenerate'] = False  # Track regenerate button state


# Function to create a new chat session in MongoDB
def create_new_chat_session():
    new_session = {
        "created_at": datetime.datetime.utcnow(),
        "messages": []  # Empty at first
    }
    session_id = chat_sessions.insert_one(new_session).inserted_id
    return str(session_id)


# Function to load a chat session by MongoDB ID
# Function to load the chat session by MongoDB ID (load full history for display)
def load_chat_session(session_id):
    session = chat_sessions.find_one({"_id": ObjectId(session_id)})
    if session:
        # Load the full chat history (no slicing here)
        st.session_state['chat_history'] = session['messages']



# Function to update chat session in MongoDB (store last 15 question-answer pairs)
# Function to update chat session in MongoDB (store entire chat history)
def update_chat_session(session_id, question, answer):
    # Append the new question-answer pair to the full messages array
    chat_sessions.update_one(
        {"_id": ObjectId(session_id)},
        {"$push": {"messages": {"$each": [{"question": question, "answer": answer}]}}}
    )


# Function to replace the last response in MongoDB
def replace_last_response_in_mongo(session_id, new_answer):
    last_message_index = len(st.session_state['chat_history']) - 1
    if last_message_index >= 0:
        # Replace the last response in MongoDB
        chat_sessions.update_one(
            {"_id": ObjectId(session_id)},
            {"$set": {f"messages.{last_message_index}.answer": new_answer}}
        )

# Function to regenerate the response
def regenerate_response():
    if st.session_state['chat_history']:
        last_question = st.session_state['chat_history'][-1]["question"]  # Get the last question
        # Exclude the last response from the history when sending the question to LLM
        previous_history = st.session_state['chat_history'][:-1]  # Exclude the last Q&A pair

        with st.spinner("Please wait, regenerating the response!"):
            # Generate a new response for the last question using only the previous history
            new_reply = generate_summary(str(reference_data), last_question, previous_history)

        # Replace the last response in the session state with the new reply
        st.session_state['chat_history'][-1]["answer"] = new_reply

        # Update MongoDB with the new response
        if st.session_state['current_chat_id']:
            replace_last_response_in_mongo(st.session_state['current_chat_id'], new_reply)

        st.session_state['regenerate'] = False  # Reset regenerate flag
        st.rerun()


# Function to generate a detailed response based on the query and JSON data


# When generating a response, pass only the latest 15 messages to the LLM
def generate_summary(reference_data, query, chat_history):
    try:
        # Escape curly braces in the JSON data to avoid conflicts with prompt placeholders
        escaped_reference_data = reference_data.replace("{", "{{").replace("}", "}}")

        # Limit the history sent to the LLM to the latest 15 question-answer pairs
        limited_history = chat_history[-15:] if len(chat_history) > 15 else chat_history

        # Create conversation history for the LLM, only using the last 15 entries
        history_text = "\n".join([f"User: {q['question']}\nLLM: {q['answer']}" for q in limited_history])

        # Define the system and user prompts including the limited history
        prompt = ChatPromptTemplate.from_messages([
            ("system", """You are a chatbot who specialises in answering questions related to Indan Oil Corporation Limitied(IOCL).
                            This is the extracted data from the Indian Oil Corporation Limited (IOCL) website. The extracted data contains detailed information about the company and its various operations. You will be provided with a query, and you must use this data to answer it comprehensively.
                            Additionally, the conversation history may contain relevant context or prior queries and responses. Use this history to ensure your answer is accurate and coherent, building on previous information if necessary.
                          
                            Key Guidelines:
                            1.Accuracy is paramount: If the extracted data or conversation history does not contain the information required to answer the query, clearly state, "The answer is not available in the context." Do not attempt to provide a speculative or incorrect response.
                            2.Be detailed: Provide clear, concise, and thorough answers without omitting any relevant information from the extracted data.
                            3.Avoid quoting field names: When responding, avoid directly quoting or referencing field names or formats from the extracted data. Instead, present the information naturally, as if summarizing or interpreting the data. Try to give the answer in points.
                            4.Use the conversation history: When applicable, refer to earlier parts of the conversation to ensure consistency and accuracy in your response.
                            5.Sometime a query might be a followup question,without proper context in that case try using previous conversation history and try to utilise latest messages to answer it if possible.
                            6.Answer the queries in conversational style.
                            7.If any links or source PDFs links are available within the extracted data part which you are referring to answer the query, you must include these in your responses, suggesting that users can refer to them for more detailed information.
                                          """),
            ("human", f'''
                            Previous Conversation History: \n{history_text}\n
                            "Extracted Data": \n{escaped_reference_data}\n
                            "Query":\n {query}\n
                            '''
             )
        ])

        # Chain the prompt with LLM for response generation
        chain = prompt | llm
        result = chain.invoke({"Extracted Data": escaped_reference_data, "Query": query})

        # Return the generated response
        return result.content

    except Exception as e:
        st.error(f"Error answering your question: {e}")
        return "Error answering question."



# Sidebar for showing chat sessions and creating new sessions
st.sidebar.header("Chat Sessions")

# Button for creating a new chat
if st.sidebar.button("New Chat"):
    new_chat_id = create_new_chat_session()
    st.session_state['current_chat_id'] = new_chat_id
    st.session_state['chat_history'] = []

# List existing chat sessions with delete button (dustbin icon)
existing_sessions = chat_sessions.find().sort("created_at", -1)
for session in existing_sessions:
    session_id = str(session['_id'])
    session_date = session['created_at'].strftime("%Y-%m-%d %H:%M:%S")
    col1, col2 = st.sidebar.columns([8, 1])

    # Display session name
    with col1:
        if st.button(f"Session {session_date}", key=session_id):
            st.session_state['current_chat_id'] = session_id
            load_chat_session(session_id)

    # Display delete icon (dustbin)
    with col2:
        if st.button("🗑️", key=f"delete_{session_id}"):
            chat_sessions.delete_one({"_id": ObjectId(session_id)})
            st.rerun()  # Refresh the app to remove the deleted session from the sidebar

# Main chat interface
st.markdown('<div class="fixed-header"><h1>Welcome To IOCL Chatbot</h1></div>', unsafe_allow_html=True)
st.markdown("<hr>", unsafe_allow_html=True)

# Input box for the question
user_question = st.chat_input(f"Ask a Question related to IOCL Website")

if user_question:
    # Automatically create a new session if none exists
    if not st.session_state['current_chat_id']:
        new_chat_id = create_new_chat_session()
        st.session_state['current_chat_id'] = new_chat_id

    with st.spinner("Please wait, I am thinking!!"):
        # Store the user's question and get the assistant's response
        reply = generate_summary(str(reference_data), user_question, st.session_state['chat_history'])


    # Append the new question-answer pair to chat history
        st.session_state['chat_history'].append({"question": user_question, "answer": reply})

        # Update the current chat session in MongoDB
        if st.session_state['current_chat_id']:
            update_chat_session(st.session_state['current_chat_id'], user_question, reply)

# Display the updated chat history (show last 15 question-answer pairs)
for i, pair in enumerate(st.session_state['chat_history']):
    question = pair["question"]
    answer = pair["answer"]
    streamlit_chat.message(question, is_user=True, key=f"chat_message_user_{i}")
    streamlit_chat.message(answer, is_user=False, key=f"chat_message_assistant_{i}")

# Display regenerate button under the last response
if st.session_state['chat_history'] and not st.session_state['regenerate']:
    if st.button("🔄 Regenerate", key="regenerate_button"):
        st.session_state['regenerate'] = True
        regenerate_response()