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()