modify main (#2)

- modify main (85b34987f65ac2d0ae5e793d781a3dc4beb529ac)


Co-authored-by: Traikia Ilyes <fansa34@users.noreply.huggingface.co>

app.py

@@ -4,7 +4,8 @@ from rag_pipeline import full_rag_pipeline
 from langchain_google_genai import GoogleGenerativeAI
 import os
 from dotenv import load_dotenv
-
+import psutil
+import os
 # Load environment variables from .env file
 load_dotenv()
 
@@ -17,8 +18,28 @@ expanding_retriever = prepare_environment_and_retriever()
 
 @app.get("/ask")
 def ask_question(question: str, with_citations: bool = Query(False, description="Include citations in the response")):
-    response = full_rag_pipeline(question, expanding_retriever,clean_all_citations=with_citations)
-    return {"question": question, "answer": response}
+    process = psutil.Process(os.getpid())
+    start_time = time.time()
+    start_mem = process.memory_info().rss  # RAM in bytes
+
+    # Run RAG
+    response = full_rag_pipeline(question, expanding_retriever, clean_all_citations=with_citations)
+
+    end_time = time.time()
+    end_mem = process.memory_info().rss
+
+    # Metrics
+    latency = end_time - start_time
+    ram_used_mb = (end_mem - start_mem) / (1024 ** 2)  # Convert to MB
+
+    return {
+        "question": question,
+        "answer": response,
+        "metrics": {
+            "latency_seconds": round(latency, 3),
+            "ram_usage_delta_mb": round(ram_used_mb, 2)
+        }
+    }
 @app.get("/generate_title")
 def generate_title(first_question: str = Query(..., description="The first question to generate a title from")):
     # Initialize the LLM - using the same model as in prepare_env.py