app.py

import streamlit as st
from langchain_pinecone import PineconeVectorStore

# from langchain_openai import OpenAI, ChatOpenAI
# from langchain_google_genai import ChatGoogleGenerativeAI

import openai

# from langchain.prompts import PromptTemplate
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.output_parsers import PydanticOutputParser
from dotenv import load_dotenv
import os

from pydantic import BaseModel, Field
from typing import List, Union, Tuple, Any
from typing_extensions import Literal


# Function to download Hugging Face embeddings
def download_hugging_face_embeddings():
    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
    return embeddings

# Load environment variables
def load_env_variables():
    load_dotenv()
    PINECONE_API_KEY = os.environ.get('PINECONE_API_KEY')
    OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
    GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")

    os.environ["PINECONE_API_KEY"] = PINECONE_API_KEY
    os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
    os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY

# utils for llm 
def construct_prompt(prompt: str, role: str):
   return{
          "role": role,
          "content": prompt
        }


# Function to initialize Pinecone vector store and retriever
def initialize_vector_store(embeddings):
    index_name = "yolotest"
    vector_store = PineconeVectorStore.from_existing_index(index_name=index_name, embedding=embeddings)
    # retriever = vector_store.as_retriever(search_type="similarity", search_kwargs={"k": 5})
    return vector_store

# Pydantic Schemas
class MetadataSchema(BaseModel):
    page: float = Field(..., description="Page number of the document")
    page_label: str = Field(..., description="Page label of the document")
    total_pages: float = Field(..., description="Total pages in the document")
    source: str = Field(..., description="Source file path of the document")
    score: float = Field(..., description="Similarity score of the retrieved document")

class DocumentSchema(BaseModel):
    metadata: MetadataSchema = Field(..., description="Filtered metadata of the document")
    page_content: str = Field(..., description="Content of the document page")

class RetrievedDocsSchema(BaseModel):
    documents: List[DocumentSchema]


# Define the individual Question model
class Question(BaseModel):
    question: str = Field(..., description="The question prompt that the user needs to answer.")
    type: Literal['fill_missing', 'MCQ', 'short_answer'] = Field(..., description="The type of question: fill_missing, MCQ, or short_answer.")
    options: Union[List[str], None] = Field(None, description="The options for the question, used only for MCQ type.")

# Define the Questions model that contains a list of Question objects and the total number of questions
class Questions(BaseModel):
    no_of_questions: int = Field(..., description="The total number of questions generated.")
    questions: List[Question] = Field(..., description="A list of Question objects.")

def initialize_prompts(parser):
    system_prompt = """
    You are an AI assistant specializing in question generation. Your role is to analyze the provided query and context to create structured questions that match the given instructions.

    Follow the specified question type format and ensure clarity, correctness, and alignment with the context.
    """

    document_prompt = """
    **Query**: {query}

    **Context**: {context}

    **Question Type Instructions:**
    - If a **specific question type** is provided, generate questions **only in that type**: {question_type}.
    - Do **not** generate other question types if a type is specified.
    - If 'MCQ', provide at least **3 options** per question.
    - If 'fill_missing', leave a **blank space** for the missing word.
    - If 'short_answer', ensure the **answer is clear** from the context.
    - If no type is specified (or 'general' is selected), generate a **variety** of question types.

    Ensure that the generated questions align with the query and retrieved context.
    """

    footer_prompt = f"""
    **Output Format:**
    {parser.get_format_instructions()}

    **Important:**
    - Ensure all questions match the requested type.
    - Do not generate options unless the type is 'MCQ'.
    - Maintain high accuracy and relevance to the query and context.
    """

    return system_prompt, document_prompt, footer_prompt


# Initialize components before user query
embeddings = download_hugging_face_embeddings()
load_env_variables()
vector_store = initialize_vector_store(embeddings)
parser = PydanticOutputParser(pydantic_object=Questions)
prompt_template = initialize_prompts(parser)


# Function for retrieving with score
def retrieve_and_format_results(vector_store: Any, query: str, k: int = 5, filter: dict = {}) -> RetrievedDocsSchema:
    """
    Retrieves documents using similarity search and formats them into the Pydantic schema.

    Args:
        vector_store (Any): The vector store used for retrieval.
        query (str): The search query.
        k (int): Number of documents to retrieve.
        filter (dict): Optional filter for the search.

    Returns:
        RetrievedDocsSchema: A structured schema containing documents and metadata.
    """
    # Ensure query is valid
    if not query or not isinstance(query, str):
        raise ValueError("Query must be a non-empty string.")

    # Retrieve documents with similarity scores
    try:
        retrieved_docs = vector_store.similarity_search_with_score(query, k=k, filter=filter)
    except Exception as e:
        raise RuntimeError(f"Error retrieving documents: {e}")

    # Convert retrieved documents into the Pydantic schema
    documents_list = [
        DocumentSchema(
            metadata=MetadataSchema(
                page=doc.metadata.get("page", 0),
                page_label=doc.metadata.get("page_label", ""),
                total_pages=doc.metadata.get("total_pages", 0),
                source=doc.metadata.get("source", ""),
                score=float(score) if score is not None else 0.0  # Ensure score is a float
            ),
            page_content=doc.page_content.strip()  # Trim extra spaces
        )
        for doc, score in retrieved_docs if doc.page_content.strip()  # Filter out empty docs
    ]

    return RetrievedDocsSchema(documents=documents_list)

# LLM to be used
def generate_text_with_schema(prompt: str, response_format, chat_history: list = None, max_output_tokens: int = None, temperature: float = 0):

    client = openai.OpenAI()

    # Ensure chat_history is initialized and includes the prompt
    if chat_history is None:
        chat_history = []

    chat_history.append({"role": "user", "content": prompt})  # Add user message

    response = client.beta.chat.completions.parse(
        model="gpt-4o-2024-08-06",
        messages=chat_history,  # Pass the updated history
        max_tokens=max_output_tokens,
        temperature=temperature,
        response_format=response_format
    )

    return response.choices[0].message.parsed


# Process and Generate Questions
def process_and_generate_questions2(vector_store, query: str, parser, response_format, question_type: str = "general",
                                   k: int = 5, filter: dict = {}, max_output_tokens: int = None, temperature: float = 0):
    """
    Retrieves relevant documents, formats prompts, and generates questions in a structured format.

    Args:
        vector_store (Any): The vector store used for retrieval.
        query (str): The search query.
        parser: Parser object to INST the retriever formatting.
        response_format: Expected response format for structured output.
        question_type (str): Type of questions to generate ('MCQ', 'fill_missing', 'short_answer', 'general').
        k (int): Number of documents to retrieve.
        filter (dict): Optional filter for retrieval.
        max_output_tokens (int): Maximum number of tokens in the response.
        temperature (float): Temperature setting for response generation.

    Returns:
        Parsed structured response containing generated questions.
    """
    retrieved_docs_schema = retrieve_and_format_results(vector_store, query, k, filter)
    retrieved_docs = retrieved_docs_schema.documents
    if not retrieved_docs:
        return {"error": "No relevant documents found for the query."}

    system_prompt, document_prompt_template, footer_prompt = initialize_prompts(parser)
    # print(f"\n\n sys :{system_prompt} ")

    chat_history = [construct_prompt(prompt=system_prompt, role="system")]
    # print(f"\n\n hist before :{chat_history} ")

    context = "\n\n".join([doc.page_content for doc in retrieved_docs])
    # print(f"\n\n context :{context} ")

    try:
        document_prompt = document_prompt_template.format(query=query, context=context, question_type=question_type)
    except KeyError as e:
        raise ValueError(f"Missing format key in document prompt template: {e}")
    # print(f"\n\n document_prompt :{document_prompt} ")
    # print(question_type)


    full_prompt = f"\n\n{document_prompt}\n\n{footer_prompt}"
    # print(f"\n\n full_prompt :{full_prompt} ")

    response = generate_text_with_schema(
        prompt=full_prompt,
        chat_history = chat_history,
        response_format=response_format,
        max_output_tokens=max_output_tokens,
        temperature=temperature
    )
    chat_history.append(construct_prompt(prompt=response, role="assistant"))
    print(f"\n\n hist after:{chat_history} ")


    return response, full_prompt, retrieved_docs_schema, chat_history


def main():
    st.set_page_config(layout="wide")
    st.title("A Simple RAG App to Generate Questions in Specific Formats")

    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []

    col1, col2, col3, col4 = st.columns([1, 2, 2, 2])
    
    with col1:
        st.subheader("Chat History")
        with st.expander("Show/Hide Chat History", expanded=True):
            for message in st.session_state.chat_history:
                st.markdown(f"**{message['role'].capitalize()}**: {message['content']}")
    
    with col2:
        st.subheader("Retrieved Documents")
        retrieved_docs_schema = RetrievedDocsSchema(documents=[])
        for doc in retrieved_docs_schema.documents:
            with st.expander(f"Source: {doc.metadata.source}, Page {doc.metadata.page}/{doc.metadata.total_pages} (Score: {doc.metadata.score:.4f})"):
                st.text_area("Content:", doc.page_content, height=150)
    
    with col3:
        st.subheader("Generated Prompt")
        st.write("Final prompt will appear here.")
    
    with col4:
        st.subheader("Generated Response")
        st.write("Response will be displayed here.")

    question_type = st.selectbox("Select Question Type", ["general", "MCQ", "fill_missing", "short_answer"])
    query = st.chat_input("Enter your query: ")
    
    if query and question_type:
        with st.spinner('Generating questions...'):
            st.write(f"**Your query:** {query}")
            response, full_prompt, retrieved_docs_schema, chat_history = process_and_generate_questions2(
                vector_store, query, parser, response_format=Questions, question_type=question_type
            )
            
            with col3:
                st.write(f"**Final prompt is:** {full_prompt}")
            
            with col2:
                for doc in retrieved_docs_schema.documents:
                    with st.expander(f"Source: {doc.metadata.source}, Page {doc.metadata.page}/{doc.metadata.total_pages} (Score: {doc.metadata.score:.4f})"):
                        st.text_area("Content:", doc.page_content, height=150)
            
            with col1:
                st.session_state.chat_history.append({"role": "user", "content": query})
                st.session_state.chat_history.append({"role": "assistant", "content": str(response)})
                st.markdown(f"**Assistant**: {response}")
            
            with col4:
                st.write(response)

if __name__ == "__main__":
    main()