Add ram usage endpoint

app.py

@@ -12,6 +12,7 @@ import logging
 import os
 import faiss
 import asyncio
+import psutil  # Added for RAM tracking
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -105,7 +106,6 @@ def retrieve_context(query, top_k=2):
 try:
     with open("cv_text.txt", "r", encoding="utf-8") as f:
         full_cv_text = f.read()
-    # Ensure full_cv_text is a string
     if not isinstance(full_cv_text, str):
         full_cv_text = str(full_cv_text)
     logger.info("CV text loaded successfully")
@@ -129,15 +129,12 @@ async def stream_response(query, history):
         f"CV: {full_cv_text}"
     )
 
-    # Ensure system_prompt is a string and debug its state
     if not isinstance(system_prompt, str):
         system_prompt = str(system_prompt)
     logger.info(f"System prompt type: {type(system_prompt)}, length: {len(system_prompt)}")
 
-    # Combine system prompt, history, and current query
     messages = [{"role": "system", "content": system_prompt}] + history + [{"role": "user", "content": query}]
 
-    # Estimate token counts and truncate history if necessary
     try:
         system_tokens = len(generator.tokenize(system_prompt.encode('utf-8'), add_bos=True, special=True))
         query_tokens = len(generator.tokenize(query.encode('utf-8'), add_bos=False, special=True))
@@ -148,19 +145,16 @@ async def stream_response(query, history):
         yield "data: [DONE]\n\n"
         return
 
-    total_tokens = system_tokens + query_tokens + sum(history_tokens) + len(history) * 10 + 10  # Rough estimate for formatting
-
-    max_allowed_tokens = generator.n_ctx() - 512 - 100  # max_tokens=512, safety_margin=100
+    total_tokens = system_tokens + query_tokens + sum(history_tokens) + len(history) * 10 + 10
+    max_allowed_tokens = generator.n_ctx() - 512 - 100
 
     while total_tokens > max_allowed_tokens and history:
         removed_msg = history.pop(0)
         removed_tokens = len(generator.tokenize(removed_msg["content"].encode('utf-8'), add_bos=False, special=True))
         total_tokens -= (removed_tokens + 10)
 
-    # Reconstruct messages after possible truncation
     messages = [{"role": "system", "content": system_prompt}] + history + [{"role": "user", "content": query}]
 
-    # Generate response with lock
     async with model_lock:
         try:
             for chunk in generator.create_chat_completion(
@@ -187,6 +181,20 @@ class QueryRequest(BaseModel):
     query: str
     history: list[dict]
 
+# RAM Usage Tracking Function
+def get_ram_usage():
+    memory = psutil.virtual_memory()
+    total_ram = memory.total / (1024 ** 3)  # Convert to GB
+    used_ram = memory.used / (1024 ** 3)   # Convert to GB
+    free_ram = memory.available / (1024 ** 3)  # Convert to GB
+    percent_used = memory.percent
+    return {
+        "total_ram_gb": round(total_ram, 2),
+        "used_ram_gb": round(used_ram, 2),
+        "free_ram_gb": round(free_ram, 2),
+        "percent_used": percent_used
+    }
+
 @app.post("/api/predict")
 async def predict(request: QueryRequest):
     query = request.query
@@ -208,6 +216,16 @@ async def model_info():
         "faiss_index_dim": cv_embeddings.shape[1],
     }
 
+@app.get("/ram_usage")
+async def ram_usage():
+    """Endpoint to get current RAM usage."""
+    try:
+        ram_stats = get_ram_usage()
+        return ram_stats
+    except Exception as e:
+        logger.error(f"Error retrieving RAM usage: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error retrieving RAM usage: {str(e)}")
+
 @app.on_event("startup")
 async def warm_up_model():
     logger.info("Warming up the model...")
@@ -215,4 +233,7 @@ async def warm_up_model():
     dummy_history = []
     async for _ in stream_response(dummy_query, dummy_history):
         pass
-    logger.info("Model warm-up completed.")
+    logger.info("Model warm-up completed.")
+    # Log initial RAM usage
+    ram_stats = get_ram_usage()
+    logger.info(f"Initial RAM usage after startup: {ram_stats}")

requirements.txt

@@ -6,4 +6,5 @@ numpy==1.26.4
 llama-cpp-python==0.3.1
 huggingface_hub==0.30.1
 faiss-cpu==1.8.0
-asyncio
+asyncio
+psutil