Update src/rag_pipeline.py

src/rag_pipeline.py

@@ -3,7 +3,6 @@ from datasets import load_dataset
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
 from langchain_text_splitters import CharacterTextSplitter
-from langchain.chat_models import ChatGoogleGenerativeAI
 from langchain_core.documents import Document
 from langgraph.graph import START, StateGraph
 from langgraph.checkpoint.memory import MemorySaver
@@ -12,6 +11,7 @@ from langchain_core.prompts import ChatPromptTemplate
 from huggingface_hub import login
 from dotenv import load_dotenv
 from typing import TypedDict, List
+import google.generativeai as genai  # Official Google Gemini SDK
 
 # Load environment variables
 load_dotenv()
@@ -19,6 +19,9 @@ load_dotenv()
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
 HF_TOKEN = os.getenv("HF_TOKEN")
 
+# Configure Google Gemini SDK
+genai.configure(api_key=GOOGLE_API_KEY)
+
 # Authenticate Hugging Face
 if HF_TOKEN:
     try:
@@ -29,7 +32,6 @@ if HF_TOKEN:
 else:
     print("⚠️ No HF_TOKEN found in .env file. Using public mode.")
 
-
 # --- STATE DEFINITION ---
 class RAGState(TypedDict):
     question: str
@@ -38,6 +40,22 @@ class RAGState(TypedDict):
     chat_history: List[str]
     source_documents: List[Document]
 
+# --- LLM Wrapper ---
+class GeminiLLMWrapper:
+    """
+    A simple wrapper around google-generativeai chat API to mimic
+    the ChatGoogleGenerativeAI interface for compatibility with app.py.
+    """
+    def invoke(self, prompt: str):
+        response = genai.chat.create(
+            model="models/gemini-2.5-flash",
+            messages=[{"role": "user", "content": prompt}]
+        )
+        # Wrap the response to have a .content attribute
+        class Result:
+            content = response.last
+        return Result()
+
 
 def build_rag_pipeline():
     """Builds a LangGraph-based RAG pipeline compatible with LangChain 1.x."""
@@ -65,7 +83,7 @@ def build_rag_pipeline():
     retriever = vector_db.as_retriever(search_kwargs={"k": 3})
 
     # --- LLM ---
-    llm = ChatGoogleGenerativeAI(model="models/gemini-2.5-flash", google_api_key=GOOGLE_API_KEY)
+    llm = GeminiLLMWrapper()  # Use wrapper instead of ChatGoogleGenerativeAI
 
     # --- PROMPT TEMPLATE ---
     prompt = ChatPromptTemplate.from_template(