to handel pydantic error

pipeline.py

@@ -9,12 +9,13 @@ from langchain.docstore.document import Document
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
 from langchain.chains import RetrievalQA
-from smolagents import CodeAgent, DuckDuckGoSearchTool, ManagedAgent, LiteLLMModel
+from smolagents import DuckDuckGoSearchTool, ManagedAgent
 from pydantic import BaseModel, Field, ValidationError, validator
 from mistralai import Mistral
-from langchain.prompts import PromptTemplate
 
-# Import chains and tools
+# Import Google Gemini model
+from langchain_google_genai import ChatGoogleGenerativeAI
+
 from classification_chain import get_classification_chain
 from cleaner_chain import get_cleaner_chain
 from refusal_chain import get_refusal_chain
@@ -25,10 +26,25 @@ 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)
 
-# Initialize LiteLLM model for web search
-pydantic_agent = LiteLLMModel(model_id="gemini/gemini-pro", api_key=os.environ.get("GEMINI_API_KEY"))
+# Setup ChatGoogleGenerativeAI for Gemini
+# Ensure GOOGLE_API_KEY is set in your environment variables.
+gemini_llm = ChatGoogleGenerativeAI(
+    model="gemini-1.5-pro",
+    temperature=0,
+    max_retries=2,
+    # You can add additional parameters or safety_settings here if needed
+)
+
+# Initialize LiteLLM model for web search (if needed)
+pydantic_agent = ManagedAgent( 
+    llm=ChatGoogleGenerativeAI(
+        model="gemini-1.5-pro",
+        temperature=0,
+        max_retries=2,
+    ),
+    tools=[DuckDuckGoSearchTool()]
+)
 
-# Pydantic models for validation and type safety
 class QueryInput(BaseModel):
     query: str = Field(..., min_length=1, description="The input query string")
     
@@ -45,7 +61,6 @@ class ModerationResult(BaseModel):
     categories: Dict[str, bool] = Field(default_factory=dict, description="Detected content categories")
     original_text: str = Field(..., description="The original input text")
 
-# Load spaCy model for NER
 def install_spacy_model():
     try:
         spacy.load("en_core_web_sm")
@@ -58,6 +73,22 @@ def install_spacy_model():
 install_spacy_model()
 nlp = spacy.load("en_core_web_sm")
 
+def sanitize_message(message: Any) -> str:
+    """Sanitize message input to ensure it's a valid string."""
+    try:
+        if hasattr(message, 'content'):
+            return str(message.content).strip()
+        if isinstance(message, dict) and 'content' in message:
+            return str(message['content']).strip()
+        if isinstance(message, list) and len(message) > 0:
+            if isinstance(message[0], dict) and 'content' in message[0]:
+                return str(message[0]['content']).strip()
+            if hasattr(message[0], 'content'):
+                return str(message[0].content).strip()
+        return str(message).strip()
+    except Exception as e:
+        raise RuntimeError(f"Error in sanitize function: {str(e)}")
+
 def extract_main_topic(query: str) -> str:
     try:
         query_input = QueryInput(query=query)
@@ -160,55 +191,76 @@ 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)}")
 
-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)
+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"
-    
+    @property
+    def _llm_type(self) -> str:
+        return "custom_gemini"
+
+def build_rag_chain(vectorstore: FAISS) -> RetrievalQA:
+    """Build RAG chain with enhanced error handling."""
     try:
         retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 3})
-        gemini_as_llm = GeminiLangChainLLM()
-        return RetrievalQA.from_chain_type(
-            llm=gemini_as_llm,
+        gemini_llm_instance = GeminiLangChainLLM()
+        chain = RetrievalQA.from_chain_type(
+            llm=gemini_llm_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)}")
-
-
-def sanitize_message(message: Any) -> str:
-    """Sanitize message input to ensure it's a valid string."""
+    
+def do_web_search(query: str) -> str:
     try:
-        if hasattr(message, 'content'):
-            return str(message.content)
-        if isinstance(message, (list, dict)):
-            return str(message)
-        return str(message)
+        search_tool = DuckDuckGoSearchTool()
+        search_agent = ManagedAgent(llm=gemini_llm, tools=[search_tool])
+        search_result = search_agent.run(f"Search for information about: {query}")
+        return str(search_result).strip()
     except Exception as e:
-        raise RuntimeError(f"Error in sanitize function: {str(e)}")
+        print(f"Web search failed: {e}")
+        return ""
+
+def merge_responses(csv_answer: str, web_answer: str) -> str:
+    try:
+        if not csv_answer and not web_answer:
+            return "I apologize, but I couldn't find any relevant information."
         
-    
+        if not web_answer:
+            return csv_answer
+            
+        if not csv_answer:
+            return web_answer
+            
+        return f"{csv_answer}\n\nAdditional information from web search:\n{web_answer}"
+    except Exception as e:
+        print(f"Error merging responses: {e}")
+        return csv_answer or web_answer or "I apologize, but I couldn't process the information properly."
+
 def run_pipeline(query: str) -> str:
     try:
-        query = sanitize_message(query)
-        topic=extract_main_topic(query)
-        moderation_result = moderate_text(query)
-        try:
-            if not moderation_result.is_safe:
-                return "Sorry, this query contains harmful or inappropriate content."
-        except Exception as e:
-            raise RuntimeError(f"Error in run_runpipeline check moderation: {str(e)}")
-        try:
-            classification = classify_query(moderation_result.original_text)
-        except Exception as e:
-            raise RuntimeError(f"Error in run_runpipeline check classify_query: {str(e)}")
+        print(query)
+        sanitized_query = sanitize_message(query)
+        query_input = QueryInput(query=sanitized_query)
+        
+        topic = extract_main_topic(query_input.query)
+        moderation_result = moderate_text(query_input.query)
+        
+        if not moderation_result.is_safe:
+            return "Sorry, this query contains harmful or inappropriate content."
+            
+        classification = classify_query(moderation_result.original_text)
         
         if classification == "OutOfScope":
             refusal_text = refusal_chain.run({"topic": topic})
@@ -216,22 +268,37 @@ def run_pipeline(query: str) -> str:
 
         if classification == "Wellness":
             rag_result = wellness_rag_chain({"query": moderation_result.original_text})
-            csv_answer = rag_result["result"].strip()
+            if isinstance(rag_result, dict) and "result" in rag_result:
+                csv_answer = str(rag_result["result"]).strip()
+            else:
+                csv_answer = str(rag_result).strip()
             web_answer = "" if csv_answer else do_web_search(moderation_result.original_text)
             final_merged = merge_responses(csv_answer, web_answer)
             return tailor_chain.run({"response": final_merged}).strip()
 
         if classification == "Brand":
             rag_result = brand_rag_chain({"query": moderation_result.original_text})
-            csv_answer = rag_result["result"].strip()
+            if isinstance(rag_result, dict) and "result" in rag_result:
+                csv_answer = str(rag_result["result"]).strip()
+            else:
+                csv_answer = str(rag_result).strip()
             final_merged = merge_responses(csv_answer, "")
             return tailor_chain.run({"response": final_merged}).strip()
         
         refusal_text = refusal_chain.run({"topic": topic})
         return tailor_chain.run({"response": refusal_text}).strip()
-    except Exception as e:
-        raise RuntimeError(f"Error in run_runpipeline: {str(e)}")
         
+    except ValidationError as e:
+        raise ValueError(f"Input validation failed: {str(e)}")
+    except Exception as e:
+        raise RuntimeError(f"Error in run_pipeline: {str(e)}")
+
+def run_with_chain(query: str) -> str:
+    try:
+        return run_pipeline(query)
+    except Exception as e:
+        print(f"Error in run_with_chain: {str(e)}")
+        return "I apologize, but I encountered an error processing your request. Please try again."
 
 # Initialize chains and vectorstores
 classification_chain = get_classification_chain()
@@ -247,12 +314,7 @@ brand_store_dir = "faiss_brand_store"
 wellness_vectorstore = build_or_load_vectorstore(wellness_csv, wellness_store_dir)
 brand_vectorstore = build_or_load_vectorstore(brand_csv, brand_store_dir)
 
-gemini_llm = LiteLLMModel(model_id="gemini/gemini-pro", api_key=os.environ.get("GEMINI_API_KEY"))
-wellness_rag_chain = build_rag_chain(gemini_llm, wellness_vectorstore)
-brand_rag_chain = build_rag_chain(gemini_llm, brand_vectorstore)
+wellness_rag_chain = build_rag_chain(wellness_vectorstore)
+brand_rag_chain = build_rag_chain(brand_vectorstore)
 
 print("Pipeline initialized successfully!")
-
-
-def run_with_chain(query: str) -> str:
-    return run_pipeline(query)