app.py

import streamlit as st
import os
from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
from pinecone import Pinecone, ServerlessSpec
from langchain.vectorstores import Pinecone as PineconeVectorStore
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from dotenv import load_dotenv
load_dotenv()

# Constants
PAGE_TITLE = "PTE Assistant"
PAGE_ICON = "🎓"
PDF_PATH = "maindataset.pdf"
EMBEDDING_MODEL = "models/embedding-001"
LLM_MODEL = "gemini-1.5-flash"

# Initialize session state
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

def set_page_config():
    st.set_page_config(page_title=PAGE_TITLE, page_icon=PAGE_ICON, layout="wide")

def setup_sidebar():
    st.sidebar.title(PAGE_TITLE)
    st.sidebar.info("This is designed to help you with PTE exam preparation. Ask any question related to PTE modules, strategies or practice tips")

@st.cache_resource
def initialize_rag_components():
    embeddings = GoogleGenerativeAIEmbeddings(model=EMBEDDING_MODEL)
    Pinecone(api_key= "7f389225-3c39-4cd0-afc2-d3fafc869cd2")
    docsearch = PineconeVectorStore.from_existing_index(
        index_name="quickstart",
        embedding=embeddings
    )

    retriever = docsearch.as_retriever(search_type="similarity", search_kwargs={"k": 3})
    llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0.01, max_output_tokens=500)

    system_prompt = """
    You are an advanced AI assistant specialized in PTE (Pearson Test of English) exam preparation. Your role is to provide expert guidance, explanations, and strategies to help students excel in all aspects of the PTE exam.
Core Responsibilities:

Provide accurate, detailed information about PTE exam structure, scoring, and recent updates.
Offer tailored advice and strategies for each PTE section: Speaking, Writing, Reading, and Listening.
Suggest effective study plans and time management techniques.
Provide constructive feedback on practice responses (when given).

Guidelines for Responses:

Use the following retrieved context to inform your answers: {context}
If the context doesn't provide sufficient information or
If you don't know the answer or are unsure, clearly state this and suggest reliable resources for further information.
Tailor your language complexity to the user's apparent level of understanding.
Be concise yet thorough. Aim for clear, actionable advice.
Use bullet points or numbered lists for step-by-step instructions or multiple tips.

Ethical Considerations:

Never provide or encourage cheating methods.
Emphasize the importance of genuine language skill development over exam tricks.
Respect copyright; don't reproduce exact questions from official PTE materials.
"""

    prompt = ChatPromptTemplate.from_messages([
        ("system", system_prompt),
        ("human", "{input}")
    ])
    question_answer_chain = create_stuff_documents_chain(llm, prompt)
    return create_retrieval_chain(retriever, question_answer_chain)

def display_chat_history():
    for role, message in st.session_state.chat_history:
        st.write(f"{'👤 **You:**' if role == 'You' else '🤖 **Assistant:**'} {message}")

def main():
    set_page_config()
    setup_sidebar()
    st.title(f"🎓 {PAGE_TITLE} using AI ")

    rag_chain = initialize_rag_components()

    st.write("Chat with the PTE Assistant:")
    user_input = st.chat_input("Your question:")

    if user_input:
        with st.spinner("Thinking..."):
            response = rag_chain.invoke({"input": user_input})
        
        if response:
            st.session_state.chat_history.extend([
                ("You", user_input),
                ("Assistant", response["answer"])
            ])

    display_chat_history()

if __name__ == "__main__":
    main()