Update app.py

app.py

@@ -18,6 +18,8 @@ import gspread
 from oauth2client.service_account import ServiceAccountCredentials
 from typing import Optional, List, Dict, Any
 from pydantic import BaseModel, Field
+import time
+import tracemalloc
 
 # Custom Configuration for Pydantic
 class Config:
@@ -164,16 +166,52 @@ graphrag_chain = GraphQAChain.from_llm(
     verbose=True
 )
 
+# Define functions to measure memory and time
+def measure_memory_and_time(func):
+    def wrapper(*args, **kwargs):
+        tracemalloc.start()  # Start memory tracking
+        start_time = time.time()  # Start time tracking
+        
+        result = func(*args, **kwargs)
+        
+        current, peak = tracemalloc.get_traced_memory()  # Get memory usage
+        tracemalloc.stop()
+        
+        end_time = time.time()  # End time tracking
+        elapsed_time = end_time - start_time  # Calculate elapsed time
+        
+        return result, elapsed_time, peak / 1024  # Convert memory to KB
+    return wrapper
+
+# Modified functions to use memory and time measurement
+@measure_memory_and_time
+def get_rag_response(question: str) -> str:
+    context = get_context(question)
+    response = rag_llm.predict(prompt.format(context=context, question=question))
+    return response
+
+@measure_memory_and_time
 def get_graphrag_response(question: str) -> str:
     system_prompt = "You are a Thai rice assistant that gives concise and direct answers. Do not explain the process, just provide the answer, provide the answer only in Thai."
     formatted_question = f"System Prompt: {system_prompt}\n\nQuestion: {question}"
     response = graphrag_chain.run(formatted_question)
     return response
 
-def compare_models(question: str) -> tuple[str, str]:
-    rag_response = get_rag_response(question)
-    graphrag_response = get_graphrag_response(question)
-    return rag_response, graphrag_response
+# Modify compare_models to collect and display metrics
+def compare_models(question: str) -> dict:
+    rag_response, rag_time, rag_memory = get_rag_response(question)
+    graphrag_response, graphrag_time, graphrag_memory = get_graphrag_response(question)
+    
+    # Combine responses with metrics
+    results = {
+        "RAG Response": rag_response,
+        "RAG Time (s)": round(rag_time, 2),
+        "RAG Memory (KB)": round(rag_memory, 2),
+        "GraphRAG Response": graphrag_response,
+        "GraphRAG Time (s)": round(graphrag_time, 2),
+        "GraphRAG Memory (KB)": round(graphrag_memory, 2),
+    }
+    return results
 
 # Gradio Interface
 with gr.Blocks() as demo:
@@ -181,16 +219,32 @@ with gr.Blocks() as demo:
         with gr.Column():
             question_input = gr.Textbox(label="ถามคำถามเกี่ยวกับข้าว:", placeholder="Enter your question about Thai rice")
             submit_btn = gr.Button(value="ถาม")
-            rag_output = gr.Textbox(label="Model A", interactive=False)
-            graphrag_output = gr.Textbox(label="Model B", interactive=False)
+            rag_output = gr.Textbox(label="RAG Response", interactive=False)
+            rag_time_output = gr.Textbox(label="RAG Time (s)", interactive=False)
+            rag_memory_output = gr.Textbox(label="RAG Memory (KB)", interactive=False)
+            graphrag_output = gr.Textbox(label="GraphRAG Response", interactive=False)
+            graphrag_time_output = gr.Textbox(label="GraphRAG Time (s)", interactive=False)
+            graphrag_memory_output = gr.Textbox(label="GraphRAG Memory (KB)", interactive=False)
             feedback = gr.Radio(label="Which response is better?", choices=["A ดีกว่า", "B ดีกว่า", "เท่ากัน", "แย่ทั้งคู่"])
             submit_feedback = gr.Button(value="Submit Feedback")
 
+    # Update Gradio app with time and memory results
+    def display_results(question):
+        results = compare_models(question)
+        return (
+            results["RAG Response"], results["RAG Time (s)"], results["RAG Memory (KB)"],
+            results["GraphRAG Response"], results["GraphRAG Time (s)"], results["GraphRAG Memory (KB)"]
+        )
+
     # Event handlers
-    submit_btn.click(fn=compare_models, inputs=[question_input], outputs=[rag_output, graphrag_output])
+    submit_btn.click(
+        fn=display_results,
+        inputs=[question_input],
+        outputs=[rag_output, rag_time_output, rag_memory_output, graphrag_output, graphrag_time_output, graphrag_memory_output]
+    )
     submit_feedback.click(fn=add_review, inputs=[question_input, rag_output, graphrag_output, feedback])
     demo.load(fn=load_data, inputs=None)
 
 # Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()