Update app.py

app.py

@@ -8,18 +8,36 @@ 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
+from huggingface_hub import InferenceClient
 import random
+import json
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-# Set up the Mistral model using HuggingFaceHub
+# Function to get Hugging Face API token
+def get_huggingface_api_token():
+    token = os.getenv('HUGGINGFACEHUB_API_TOKEN')
+    if not token:
+        try:
+            with open('config.json', 'r') as config_file:
+                config = json.load(config_file)
+                token = config.get('HUGGINGFACEHUB_API_TOKEN')
+        except (FileNotFoundError, json.JSONDecodeError):
+            logger.warning("Config file not found or invalid. Please provide the Hugging Face API token.")
+    
+    if not token:
+        token = input("Please enter your Hugging Face API token: ")
+    
+    return token
+
+# Initialize the Hugging Face Inference Client
 try:
-    mistral_llm = HuggingFaceHub(repo_id="mistralai/Mistral-7B-Instruct-v0.2", model_kwargs={"temperature": 0.7, "max_new_tokens": 500})
+    hf_token = get_huggingface_api_token()
+    client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.2", token=hf_token)
 except Exception as e:
-    logger.error(f"Failed to initialize Mistral model: {e}")
+    logger.error(f"Failed to initialize Hugging Face client: {e}")
     sys.exit(1)
 
 # Shared context for both agents
@@ -69,6 +87,14 @@ 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)
     
@@ -87,7 +113,7 @@ communication_expert_crew = Agent(
     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,
-    llm=mistral_llm
+    tools=[generate_response]
 )
 
 response_expert_crew = Agent(
@@ -96,7 +122,7 @@ response_expert_crew = Agent(
     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,
-    llm=mistral_llm
+    tools=[generate_response]
 )
 
 # Main function