Update app.py

app.py

@@ -6,58 +6,86 @@ import platform
 import psutil
 import multiprocessing
 import time
+import uuid # For generating unique session IDs
 
 app = FastAPI()
 
 # === Model Config ===
-REPO_ID = "TheBloke/phi-2-GGUF"
-FILENAME = "phi-2.Q3_K_M.gguf"
+# Switched to TinyLlama-1.1B-Chat-v1.0 for better CPU performance
+REPO_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0-GGUF"
+FILENAME = "tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf" # Q4_K_M is a good balance of size and quality
 MODEL_DIR = "models"
 MODEL_PATH = os.path.join(MODEL_DIR, FILENAME)
 
 # === Download if model not available ===
 if not os.path.exists(MODEL_PATH):
     print(f"⬇️ Downloading {FILENAME} from Hugging Face...")
-    model_path = hf_hub_download(
-        repo_id=REPO_ID,
-        filename=FILENAME,
-        cache_dir=MODEL_DIR,
-        local_dir=MODEL_DIR,
-        local_dir_use_symlinks=False
-    )
+    try:
+        model_path = hf_hub_download(
+            repo_id=REPO_ID,
+            filename=FILENAME,
+            cache_dir=MODEL_DIR,
+            local_dir=MODEL_DIR,
+            local_dir_use_symlinks=False
+        )
+        print(f"✅ Model downloaded to: {model_path}")
+    except Exception as e:
+        print(f"❌ Error downloading model: {e}")
+        # Exit or handle error appropriately if model download fails
+        exit(1)
 else:
     print(f"✅ Model already available at: {MODEL_PATH}")
     model_path = MODEL_PATH
 
 # === Optimal thread usage ===
+# psutil.cpu_count(logical=True) gives the number of logical cores (threads)
+# psutil.cpu_count(logical=False) gives the number of physical cores
+# For llama.cpp, n_threads often performs best when set to the number of physical cores,
+# or slightly more, but not exceeding logical cores. Experimentation is key.
 logical_cores = psutil.cpu_count(logical=True)
 physical_cores = psutil.cpu_count(logical=False)
-recommended_threads = 4
+# A common recommendation is to use physical cores or physical_cores * 2
+# Let's try physical_cores for a start, or a fixed value if physical_cores is too low.
+recommended_threads = max(1, physical_cores) # Ensure at least 1 thread
+
+print(f"Detected physical cores: {physical_cores}, logical cores: {logical_cores}")
+print(f"Using n_threads: {recommended_threads}")
 
 # === Load the model ===
-llm = Llama(
-    model_path=model_path,
-    n_ctx=2048,  # Can increase depending on memory
-    n_threads=recommended_threads,
-    use_mlock=True,  # lock model in RAM for faster access
-    n_gpu_layers=0,  # CPU only, use >0 if GPU is present
-    chat_format="chatml",  # for Hermes 2
-    verbose=False
-)
+try:
+    llm = Llama(
+        model_path=model_path,
+        n_ctx=1024,  # Reduced context for TinyLlama, can increase if memory allows and context is critical
+        n_threads=recommended_threads,
+        use_mlock=True,  # Lock model in RAM for faster access (good for stability on CPU)
+        n_gpu_layers=0,  # CPU only, keep at 0 for Hugging Face free tier
+        chat_format="chatml",  # TinyLlama Chat uses ChatML format
+        verbose=False
+    )
+    print("🚀 Llama model loaded successfully!")
+except Exception as e:
+    print(f"❌ Error loading Llama model: {e}")
+    exit(1)
+
+# === Global dictionary to store chat histories per session ===
+# In a production environment, this should be replaced with a persistent storage
+# like Redis, a database, or a dedicated session management system.
+chat_histories = {}
 
 @app.get("/")
 def root():
-    return {"message": "✅ Hermes 2 Mistral API is live and optimized!"}
+    return {"message": "✅ Data Analysis AI API is live and optimized!"}
 
 @app.get("/get_sys")
 def get_sys_specs():
+    """Returns system specifications including CPU, RAM, and OS details."""
     memory = psutil.virtual_memory()
     return {
         "CPU": {
             "physical_cores": physical_cores,
             "logical_cores": logical_cores,
-            "max_freq_mhz": psutil.cpu_freq().max,
-            "cpu_usage": psutil.cpu_percent(interval=1)
+            "max_freq_mhz": psutil.cpu_freq().max if psutil.cpu_freq() else "N/A",
+            "cpu_usage_percent": psutil.cpu_percent(interval=1) # CPU usage over 1 second
         },
         "RAM": {
             "total_GB": round(memory.total / (1024 ** 3), 2),
@@ -67,45 +95,94 @@ def get_sys_specs():
         "System": {
             "platform": platform.platform(),
             "architecture": platform.machine(),
-            "python": platform.python_version()
+            "python_version": platform.python_version()
+        },
+        "Model_Config": {
+            "model_name": FILENAME,
+            "n_ctx": llm.n_ctx(),
+            "n_threads": llm.n_threads(),
+            "use_mlock": llm.use_mlock()
         }
     }
 
 @app.get("/process_list")
 def process_list():
-    time.sleep(1)  # Let CPU settle
+    """Returns a list of processes consuming significant CPU."""
+    time.sleep(1)  # Let CPU settle for accurate measurement
     processes = []
-    for proc in psutil.process_iter(['pid', 'name']):
+    for proc in psutil.process_iter(['pid', 'name', 'cpu_percent', 'memory_percent']):
         try:
             cpu = proc.cpu_percent()
-            if cpu > 10:
+            mem = proc.memory_percent()
+            # Filter processes using more than 5% CPU or 2% memory
+            if cpu > 5 or mem > 2:
                 processes.append({
                     "pid": proc.pid,
                     "name": proc.name(),
-                    "cpu_percent": cpu
+                    "cpu_percent": round(cpu, 2),
+                    "memory_percent": round(mem, 2)
                 })
-        except (psutil.NoSuchProcess, psutil.AccessDenied):
+        except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
             pass
+    # Sort by CPU usage descending
+    processes.sort(key=lambda x: x['cpu_percent'], reverse=True)
     return {"heavy_processes": processes}
 
 @app.post("/generate")
 async def generate(request: Request):
+    """
+    Generates a response from the LLM, maintaining chat context.
+    Expects a JSON body with 'prompt' and optionally 'session_id'.
+    If 'session_id' is not provided, a new one will be generated.
+    """
     data = await request.json()
     prompt = data.get("prompt", "").strip()
-    print("🧾 Prompt received:", prompt)
-
-    response = llm.create_chat_completion(
-        messages=[
-            {"role": "system", "content": "You are a helpful AI assistant."},
-            {"role": "user", "content": prompt}
-        ],
-        max_tokens=2048,
-        temperature=0.7,
-        stop=["</s>"]
-    )
+    session_id = data.get("session_id")
 
-    llm.reset()  # Free memory after response
+    if not prompt:
+        return {"error": "Prompt cannot be empty"}, 400
 
-    return {
-        "response": response["choices"][0]["message"]["content"].strip()
-    }
+    # Generate a new session ID if not provided (for new conversations)
+    if not session_id:
+        session_id = str(uuid.uuid4())
+        # Initialize chat history for a new session with a system message
+        chat_histories[session_id] = [
+            {"role": "system", "content": "You are a helpful AI assistant for data analysis. Provide concise and actionable suggestions based on the data provided or questions asked. Keep your responses focused on data insights and actionable steps for report generation."}
+        ]
+        print(f"🆕 New session created: {session_id}")
+    elif session_id not in chat_histories:
+        # If a session ID is provided but not found, re-initialize it
+        chat_histories[session_id] = [
+            {"role": "system", "content": "You are a helpful AI assistant for data analysis. Provide concise and actionable suggestions based on the data provided or questions asked. Keep your responses focused on data insights and actionable steps for report generation."}
+        ]
+        print(f"⚠️ Session ID {session_id} not found, re-initializing history.")
+
+    print(f"🧾 Prompt received for session {session_id}: {prompt}")
+
+    # Add the user's new message to the history for this session
+    chat_histories[session_id].append({"role": "user", "content": prompt})
+
+    try:
+        # Pass the entire chat history for context
+        response = llm.create_chat_completion(
+            messages=chat_histories[session_id],
+            max_tokens=512,  # Limit response length for faster generation
+            temperature=0.7, # Adjust temperature for creativity vs. coherence (0.0-1.0)
+            stop=["</s>"] # Stop sequence for TinyLlama Chat
+        )
+
+        ai_response_content = response["choices"][0]["message"]["content"].strip()
+
+        # Add the AI's response to the history for future turns
+        chat_histories[session_id].append({"role": "assistant", "content": ai_response_content})
+
+        return {
+            "response": ai_response_content,
+            "session_id": session_id # Return the session_id so the client can use it for subsequent requests
+        }
+    except Exception as e:
+        print(f"❌ Error during generation for session {session_id}: {e}")
+        # Remove the last user message from history if generation failed to prevent bad state
+        if chat_histories[session_id] and chat_histories[session_id][-1]["role"] == "user":
+            chat_histories[session_id].pop()
+        return {"error": f"Failed to generate response: {e}. Please try again.", "session_id": session_id}, 500