Update app.py

app.py

@@ -9,21 +9,11 @@ from sklearn.feature_extraction.text import CountVectorizer
 from sklearn.naive_bayes import MultinomialNB
 import asyncio
 from crewai import Agent as CrewAgent, Task, Crew
-import openai
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Check for OpenAI API key
-if 'OPENAI_API_KEY' not in os.environ:
-    logger.error("OPENAI_API_KEY environment variable is not set.")
-    logger.info("Please set the OPENAI_API_KEY environment variable before running this script.")
-    sys.exit(1)
-
-# Set OpenAI API key
-openai.api_key = os.environ['OPENAI_API_KEY']
-
 # Initialize the client with the Mistral-7B-Instruct-v0.2 model
 try:
     client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")
@@ -31,7 +21,7 @@ except Exception as e:
     logger.error(f"Failed to initialize InferenceClient: {e}")
     sys.exit(1)
 
-# Shared context for agents
+# Shared context for both agents
 SHARED_CONTEXT = """You are part of a multi-agent system designed to provide respectful, empathetic, and accurate support for Zerodha, a leading Indian financial services company. Your role is crucial in ensuring all interactions uphold the highest standards of customer service while maintaining Zerodha's excellent reputation.
 
 Key points about Zerodha:
@@ -78,24 +68,44 @@ def check_confidence(response):
     uncertain_phrases = ["I'm not sure", "It's possible", "I don't have enough information"]
     return not any(phrase.lower() in response.lower() for phrase in uncertain_phrases)
 
-# Agent setup using CrewAI only
+async def generate_response(prompt):
+    try:
+        return await client.text_generation(prompt, max_new_tokens=500, temperature=0.7)
+    except Exception as e:
+        logger.error(f"Error generating response: {e}")
+        return "I apologize, but I'm having trouble generating a response at the moment. Please try again later."
+
+def post_process_response(response):
+    response = re.sub(r'\b(stupid|dumb|idiotic|foolish)\b', 'mistaken', response, flags=re.IGNORECASE)
+    
+    if not re.search(r'(Thank you|Is there anything else|Hope this helps|Let me know if you need more information)\s*$', response, re.IGNORECASE):
+        response += "\n\nIs there anything else I can help you with regarding Zerodha's services?"
+    
+    if re.search(r'\b(invest|trade|buy|sell|market)\b', response, re.IGNORECASE):
+        response += "\n\nPlease note that this information is for educational purposes only and should not be considered as financial advice. Always do your own research and consider consulting with a qualified financial advisor before making investment decisions."
+    
+    return response
+
+# CrewAI setup
 communication_expert_crew = CrewAgent(
     role='Communication Expert',
     goal='Interpret and rephrase user queries with empathy and respect',
-    backstory=SHARED_CONTEXT,
+    backstory="""You are an expert in communication, specializing in understanding and rephrasing queries to ensure they are interpreted in the most positive and constructive light. Your role is crucial in setting the tone for respectful and empathetic interactions.""",
     verbose=True,
-    allow_delegation=False
+    allow_delegation=False,
+    tools=[generate_response]
 )
 
 response_expert_crew = CrewAgent(
     role='Response Expert',
     goal='Provide accurate, helpful, and emotionally intelligent responses to user queries',
-    backstory=SHARED_CONTEXT,
+    backstory="""You are an expert in Zerodha's services and policies, with a keen ability to provide comprehensive and empathetic responses. Your role is to ensure that all user queries are addressed accurately while maintaining a respectful and supportive tone.""",
     verbose=True,
-    allow_delegation=False
+    allow_delegation=False,
+    tools=[generate_response]
 )
 
-# Main function for handling user queries
+# Main function
 async def zerodha_support(message, history):
     try:
         sanitized_message = sanitize_input(message)
@@ -105,33 +115,30 @@ async def zerodha_support(message, history):
         
         sanitized_message = redact_sensitive_info(sanitized_message)
 
-        # CrewAI for rephrasing query
-        rephrase_task = Task(
-            description=f"Rephrase the following user query with empathy and respect: '{sanitized_message}'",
-            agent=communication_expert_crew
-        )
-        
-        crew = Crew(
-            agents=[communication_expert_crew],
-            tasks=[rephrase_task],
-            verbose=2
-        )
-
-        rephrased_query = await crew.kickoff()
-
-        # CrewAI for generating the response
-        response_task = Task(
-            description=f"Generate a comprehensive and empathetic response to: '{rephrased_query}'",
-            agent=response_expert_crew
-        )
-        
-        response_crew = Crew(
-            agents=[response_expert_crew],
-            tasks=[response_task],
-            verbose=2
-        )
-
-        response = await response_crew.kickoff()
+        # Use crewAI for initial query rephrasing
+        try:
+            rephrase_task = Task(
+                description=f"Rephrase the following user query with empathy and respect: '{sanitized_message}'",
+                agent=communication_expert_crew
+            )
+
+            crew = Crew(
+                agents=[communication_expert_crew],
+                tasks=[rephrase_task],
+                verbose=2
+            )
+
+            rephrased_query = crew.kickoff()
+        except Exception as e:
+            logger.error(f"Error in CrewAI rephrasing: {e}")
+            rephrased_query = sanitized_message  # Fallback to original message if rephrasing fails
+
+        # Generate response using Mistral model
+        try:
+            response = await generate_response(f"{SHARED_CONTEXT}\n\nUser: {rephrased_query}\nAssistant:")
+        except Exception as e:
+            logger.error(f"Error in response generation: {e}")
+            response = "I apologize, but I'm having trouble generating a response at the moment. Please try again later."
 
         if not check_response_content(response):
             response += "\n\nPlease note that I cannot provide specific investment advice or guarantee returns. For personalized guidance, please consult with a qualified financial advisor."
@@ -166,4 +173,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()