Update app.py

app.py

@@ -142,28 +142,22 @@ def init_db():
         conn.commit()
 
 def initialize_llm():
-    """Initialize the model in the simplest way possible"""
+    """Initialize the LLM using InferenceClient."""
     try:
-        # Get API token
+        # Get API token from environment variable
         api_token = os.environ.get('HF_TOKEN')
         if not api_token:
             raise ValueError("No API token found")
 
-        # Initialize with minimal settings
-        llm = HuggingFaceHub(
-            repo_id="mistralai/Mistral-7B-Instruct-v0.3",  # Using a simpler model for testing
-            huggingfacehub_api_token=api_token,
-            task="text-generation"        )
+        # Initialize the InferenceClient
+        client = InferenceClient(api_key=api_token)
 
-        # Quick test
-        test = llm("Say hello")
-        print(f"Test response: {test}")
-
-        return llm
+        print("LLM initialized successfully!")
+        return client
 
     except Exception as e:
         print(f"LLM initialization error: {str(e)}")
-        raise
+        return None
 
         
 class ChatSession:
@@ -278,14 +272,12 @@ try:
     print("Database initialized successfully")
 
     print("Initializing LLM...")
-    llm = initialize_llm()
-    if llm is None:
+    llm_client = initialize_llm()
+    if llm_client is None:
         raise ValueError("LLM initialization failed")
     print("LLM initialized successfully")
 
-    print("Creating LLM chain...")
-    llm_chain = LLMChain(llm=llm, prompt=prompt)
-    print("LLM chain created successfully")
+
 
 except Exception as e:
     print(f"Fatal initialization error: {e}")
@@ -472,24 +464,41 @@ def get_chat_list():
 @app.route("/api/chat", methods=["POST"])
 def chat():
     try:
-        # Get the user's message
+        # Get the user's message from the request
         data = request.json
         user_input = data.get("message", "")
         print(f"Received message: {user_input}")
 
+        if not user_input:
+            return jsonify({
+                "success": False,
+                "response": "No message provided."
+            })
+
         # Very simple prompt for testing
-        prompt = f"You are a helpful assistant. Keep your response short and simple. User says: {user_input}"
+        messages = [
+            {"role": "user", "content": user_input}
+        ]
 
         try:
-            # Get response directly from model
-            response = llm(prompt)
-            print(f"Raw response received: {response}")
+            # Get response from the model using InferenceClient
+            if llm_client:
+                completion = llm_client.chat.completions.create(
+                    model="mistralai/Mistral-7B-Instruct-v0.3",  # Replace with the appropriate model
+                    messages=messages,
+                    max_tokens=150  # Adjust the max tokens as needed
+                )
+                response = completion['choices'][0]['message']['content']
+                print(f"Raw response received: {response}")
+
+                # Return the model's response
+                return jsonify({
+                    "success": True,
+                    "response": response
+                })
 
-            # Return whatever response we get
-            return jsonify({
-                "success": True,
-                "response": response
-            })
+            else:
+                raise ValueError("LLM client not initialized properly.")
 
         except Exception as e:
             print(f"Model error: {str(e)}")