Update pipeline.py

pipeline.py

@@ -32,10 +32,10 @@ gemini_llm = ChatGoogleGenerativeAI(
     model="gemini-1.5-pro",
     temperature=0,
     max_retries=2,
-    # You can add additional parameters or safety_settings here if needed
+    # Additional parameters or safety_settings can be added here if needed
 )
 
-# Initialize LiteLLM model for web search (if needed)
+# Initialize ManagedAgent for web search using Gemini
 pydantic_agent = ManagedAgent( 
     llm=ChatGoogleGenerativeAI(
         model="gemini-1.5-pro",
@@ -191,29 +191,12 @@ def build_or_load_vectorstore(csv_path: str, store_dir: str) -> FAISS:
     except Exception as e:
         raise RuntimeError(f"Error building/loading vector store: {str(e)}")
 
-class GeminiLangChainLLM(LLM):
-    def _call(self, prompt: str, stop: Optional[list] = None, **kwargs) -> str:
-        """Call the Gemini model using ChatGoogleGenerativeAI and ensure string output."""
-        try:
-            # Construct message list for the Gemini model
-            messages = [("human", prompt)]
-            ai_msg = gemini_llm.invoke(messages)
-            return ai_msg.content.strip() if ai_msg and ai_msg.content else str(prompt)
-        except Exception as e:
-            print(f"Error in GeminiLangChainLLM._call: {e}")
-            return str(prompt)  # Fallback to returning the prompt
-
-    @property
-    def _llm_type(self) -> str:
-        return "custom_gemini"
-
 def build_rag_chain(vectorstore: FAISS) -> RetrievalQA:
-    """Build RAG chain with enhanced error handling."""
+    """Build RAG chain using the Gemini LLM directly without a custom class."""
     try:
         retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})
-        gemini_llm_instance = GeminiLangChainLLM()
         chain = RetrievalQA.from_chain_type(
-            llm=gemini_llm_instance,
+            llm=gemini_llm,  # Directly use the ChatGoogleGenerativeAI instance
             chain_type="stuff",
             retriever=retriever,
             return_source_documents=True