Update app.py

app.py

@@ -2,24 +2,24 @@ import os
 import sys
 import logging
 import gradio as gr
-from huggingface_hub import InferenceClient
 import re
 import numpy as np
 from sklearn.feature_extraction.text import CountVectorizer
 from sklearn.naive_bayes import MultinomialNB
 import asyncio
 from crewai import Agent, Task, Crew
+from langchain.llms import HuggingFaceHub
 import random
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Initialize the client with the Mistral-7B-Instruct-v0.2 model
+# Set up the Mistral model using HuggingFaceHub
 try:
-    client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.2")
+    mistral_llm = HuggingFaceHub(repo_id="mistralai/Mistral-7B-Instruct-v0.2", model_kwargs={"temperature": 0.7, "max_new_tokens": 500})
 except Exception as e:
-    logger.error(f"Failed to initialize InferenceClient: {e}")
+    logger.error(f"Failed to initialize Mistral model: {e}")
     sys.exit(1)
 
 # Shared context for both agents
@@ -69,14 +69,6 @@ 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)
 
-async def generate_response(prompt):
-    try:
-        response = await client.text_generation(prompt, max_new_tokens=500, temperature=0.7)
-        return response
-    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)
     
@@ -94,7 +86,8 @@ communication_expert_crew = Agent(
     goal='Interpret and rephrase user queries with empathy and respect',
     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,
+    llm=mistral_llm
 )
 
 response_expert_crew = Agent(
@@ -102,32 +95,8 @@ response_expert_crew = Agent(
     goal='Provide accurate, helpful, and emotionally intelligent responses to user queries',
     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
-)
-
-# Task definitions
-rephrase_task = Task(
-    description="Rephrase the user query with empathy and respect",
-    agent=communication_expert_crew,
-    async_execution=True,
-    context=[
-        {
-            "role": "system",
-            "content": SHARED_CONTEXT
-        }
-    ]
-)
-
-response_task = Task(
-    description="Provide an accurate and helpful response to the user query",
-    agent=response_expert_crew,
-    async_execution=True,
-    context=[
-        {
-            "role": "system",
-            "content": SHARED_CONTEXT
-        }
-    ]
+    allow_delegation=False,
+    llm=mistral_llm
 )
 
 # Main function
@@ -142,12 +111,17 @@ async def zerodha_support(message, history, username):
 
         # Use crewAI for initial query rephrasing
         try:
-            rephrase_task.update_task(f"Rephrase the following user query with empathy and respect: '{sanitized_message}'")
+            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()
         except Exception as e:
             logger.error(f"Error in CrewAI rephrasing: {e}")
@@ -155,12 +129,17 @@ async def zerodha_support(message, history, username):
 
         # Generate response using Response Expert
         try:
-            response_task.update_task(f"Provide an accurate and helpful response to the user query: '{rephrased_query}'")
+            response_task = Task(
+                description=f"Provide an accurate and helpful response to the user query: '{rephrased_query}'",
+                agent=response_expert_crew
+            )
+
             crew = Crew(
                 agents=[response_expert_crew],
                 tasks=[response_task],
                 verbose=2
             )
+
             response = await crew.kickoff()
         except Exception as e:
             logger.error(f"Error in CrewAI response generation: {e}")
@@ -238,11 +217,10 @@ with gr.Blocks(css=custom_css) as demo:
                 submit_button = gr.Button("Send")
             
             with gr.Column(scale=1):
-                with gr.Box():
-                    gr.HTML("<div class='user-info'>")
-                    user_display = gr.Textbox(label="Logged in as", interactive=False)
-                    ticket_number = gr.Textbox(label="Support Ticket", value=generate_ticket_number(), interactive=False)
-                    gr.HTML("</div>")
+                gr.HTML("<div class='user-info'>")
+                user_display = gr.Textbox(label="Logged in as", interactive=False)
+                ticket_number = gr.Textbox(label="Support Ticket", value=generate_ticket_number(), interactive=False)
+                gr.HTML("</div>")
                 
                 gr.HTML("<h4>Quick Links</h4>")
                 gr.HTML("<ul><li><a href='https://zerodha.com/varsity/' target='_blank'>Zerodha Varsity</a></li><li><a href='https://console.zerodha.com/' target='_blank'>Zerodha Console</a></li></ul>")