qa_chain.py

import os
from langchain_groq import ChatGroq
from langchain_core.prompts import ChatPromptTemplate 
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnableMap
from retrieval import retriever
from re_ranker import rerank_documents
from query_expansion import query_expansion_chain
from dotenv import load_dotenv
load_dotenv()


GROQ_API = os.getenv("GROQ_API_KEY")


llm = ChatGroq(
    model="openai/gpt-oss-120b",
    api_key= GROQ_API 

)

prompt = ChatPromptTemplate.from_template(
    """Answer the question based on the following context below.
    
    Context:
    {context}
    
    Question:
    {question}
    
    Answer:"""
)

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

def get_context(x):
    expanded_query = query_expansion_chain.invoke({"query": x["question"]})
    retrieved_docs = retriever.invoke(expanded_query)
    reranked_docs = rerank_documents(x["question"], retrieved_docs, top_n=3)
    return format_docs(reranked_docs)


rag_pipeline = (
    RunnableMap({
        "context": get_context,
        "question": lambda x: x["question"],
    })
    | prompt
    | llm
    | StrOutputParser()
    
)

if __name__=="__main__":
    question = "covering index"
    answer = rag_pipeline.invoke({"question": question})
    print(f"Question: {question}\n")
    print(f"Answer: {answer}")