Update pipeline.py

pipeline.py

@@ -13,7 +13,7 @@ from smolagents import CodeAgent, DuckDuckGoSearchTool, ManagedAgent, LiteLLMMod
 from pydantic import BaseModel, ValidationError, validator
 from mistralai import Mistral
 from langchain.prompts import PromptTemplate
-
+from langchain_google_genai import ChatGoogleGenerativeAI
 # Import chains and tools
 from classification_chain import get_classification_chain
 from cleaner_chain import get_cleaner_chain
@@ -25,6 +25,13 @@ from prompts import classification_prompt, refusal_prompt, tailor_prompt
 mistral_api_key = os.environ.get("MISTRAL_API_KEY")
 client = Mistral(api_key=mistral_api_key)
 
+gemini_llm = ChatGoogleGenerativeAI(
+    model="gemini-1.5-pro",
+    temperature=0.5,
+    max_retries=2,
+    google_api_key=os.environ.get("GEMINI_API_KEY"),
+    # Additional parameters or safety_settings can be added here if needed
+)
 # Load spaCy model for NER and download it if not already installed
 def install_spacy_model():
     try:
@@ -131,25 +138,19 @@ def build_or_load_vectorstore(csv_path: str, store_dir: str) -> FAISS:
         return vectorstore
 
 # Function to build RAG chain
-def build_rag_chain(llm_model: LiteLLMModel, vectorstore: FAISS) -> RetrievalQA:
-    class GeminiLangChainLLM(LLM):
-        def _call(self, prompt: str, stop: Optional[list] = None, **kwargs) -> str:
-            messages = [{"role": "user", "content": prompt}]
-            return llm_model(messages, stop_sequences=stop)
-
-        @property
-        def _llm_type(self) -> str:
-            return "custom_gemini"
-    
-    retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})
-    gemini_as_llm = GeminiLangChainLLM()
-    rag_chain = RetrievalQA.from_chain_type(
-        llm=gemini_as_llm,
-        chain_type="stuff",
-        retriever=retriever,
-        return_source_documents=True
-    )
-    return rag_chain
+def build_rag_chain(vectorstore: FAISS) -> RetrievalQA:
+    """Build RAG chain using the Gemini LLM directly without a custom class."""
+    try:
+        retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})
+        chain = RetrievalQA.from_chain_type(
+            llm=gemini_llm,  # Directly use the ChatGoogleGenerativeAI instance
+            chain_type="stuff",
+            retriever=retriever,
+            return_source_documents=True
+        )
+        return chain
+    except Exception as e:
+        raise RuntimeError(f"Error building RAG chain: {str(e)}")
 
 # Function to perform web search using DuckDuckGo
 def do_web_search(query: str) -> str: