Update src/chains/qa_chain_gemini.py

src/chains/qa_chain_gemini.py

@@ -1 +1,34 @@
-# Placeholder for Gemini RAG chain
+from langchain_chroma import Chroma
+from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
+from langchain.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.chains import create_history_aware_retriever, create_retrieval_chain
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+from src.agents.prompts import RAG_PROMPT
+
+
+def build_gemini_rag_chain(pdf_path: str):
+    # Load and split documents
+    loader = PyPDFLoader(pdf_path)
+    documents = loader.load()
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    texts = splitter.split_documents(documents)
+
+    # Create vectorstore
+    vectorstore = Chroma.from_documents(texts, embedding=GoogleGenerativeAIEmbeddings(model="models/embedding-001"))
+    retriever = vectorstore.as_retriever()
+
+    # Make retriever history-aware
+    contextualize_q_prompt = ChatPromptTemplate.from_messages([
+        ("system", "Given a chat history and the latest user question, rewrite it as a standalone question."),
+        MessagesPlaceholder("chat_history"),
+        ("human", "{input}"),
+    ])
+
+    model = ChatGoogleGenerativeAI(model="gemini-2.0-flash", convert_system_message_to_human=True)
+    history_aware_retriever = create_history_aware_retriever(model, retriever, contextualize_q_prompt)
+
+    # Build RAG chain
+    qa_chain = create_stuff_documents_chain(model, RAG_PROMPT)
+    return create_retrieval_chain(history_aware_retriever, qa_chain)