Code optimization

app/frontend/gradio_app.py

@@ -3,8 +3,14 @@ app/frontend/gradio_app.py — Full page warm gray UI
 """
 import os
 import gradio as gr
+from datetime import datetime, timezone, timedelta
 from langchain_core.messages import HumanMessage
 
+IST = timezone(timedelta(hours=5, minutes=30))
+
+def _now_ist() -> str:
+    return datetime.now(IST).strftime("%d %b %Y  %I:%M:%S %p IST")
+
 os.environ["GRADIO_MODE"] = "true"
 os.environ["HITL_ENABLED"] = os.getenv("HITL_ENABLED", "true")
 
@@ -21,9 +27,18 @@ _pending_hitl_state: AgentState | None = None
 
 
 def run_graph(query: str) -> AgentState:
-    # Just pass the query — graph manages its own message history via state
+    # Restore persisted messages from MemorySaver checkpoint
+    # This gives the LLM full clean conversation history across turns
+    prior_messages = []
+    try:
+        checkpoint = _graph.get_state(_thread_config)
+        if checkpoint and checkpoint.values:
+            prior_messages = checkpoint.values.get("messages", [])
+    except Exception:
+        pass  # First turn — no checkpoint yet
+
     initial_state: AgentState = {
-        "messages": [],  # MemorySaver restores history; safety_node adds HumanMessage
+        "messages": prior_messages,  # Prior safe history; safety_node adds current HumanMessage
         "query": query,
         "route": "", "rag_context": "", "tool_calls": [], "tool_results": [],
         "response": "", "retry_count": 0, "hitl_approved": False,
@@ -75,6 +90,10 @@ def handle_submit(user_message, chat_history):
         return chat_history, "", "*Waiting for a query...*", "", gr.update(visible=False), gr.update(value="")
 
     chat_history = chat_history + [user_msg(user_message)]
+    print(f"\n{'─'*60}")
+    print(f"[QUERY] {user_message}")
+    print(f"[TIME]  {_now_ist()}")
+    print(f"{'─'*60}")
     try:
         fs = run_graph(user_message)
         route = fs.get("route", "")

app/frontend/gradio_app_hf.py

@@ -25,6 +25,7 @@ if "dotenv" not in sys.modules:
     sys.modules["dotenv"] = MagicMock()
 
 import gradio as gr
+from datetime import datetime, timezone, timedelta
 from langchain_core.messages import HumanMessage
 
 from app.graph.builder import build_graph
@@ -32,6 +33,10 @@ from app.state import AgentState
 from app.nodes.hitl import HITLPauseException
 from app.frontend.css import CSS
 
+IST = timezone(timedelta(hours=5, minutes=30))
+
+def _now_ist() -> str:
+    return datetime.now(IST).strftime("%d %b %Y  %I:%M:%S %p IST")
 
 # ── Graph singleton ────────────────────────────────────────────────────────
 _graph         = build_graph()
@@ -43,9 +48,17 @@ _pending_hitl_state: AgentState | None = None
 # ── Core runner ────────────────────────────────────────────────────────────
 
 def run_graph(query: str) -> AgentState:
-    # messages=[] — MemorySaver restores prior history; safety_node adds HumanMessage
+    # Restore persisted messages from MemorySaver checkpoint
+    prior_messages = []
+    try:
+        checkpoint = _graph.get_state(_thread_config)
+        if checkpoint and checkpoint.values:
+            prior_messages = checkpoint.values.get("messages", [])
+    except Exception:
+        pass  # First turn — no checkpoint yet
+
     initial_state: AgentState = {
-        "messages":         [],
+        "messages":         prior_messages,  # Prior safe history; safety_node adds current HumanMessage
         "query":            query,
         "route":            "",
         "rag_context":      "",
@@ -109,6 +122,10 @@ def handle_submit(user_message, chat_history):
         return chat_history, "", "*Waiting for a query...*", "", gr.update(visible=False), gr.update(value="")
 
     chat_history = chat_history + [user_msg(user_message)]
+    print(f"\n{'─'*60}")
+    print(f"[QUERY] {user_message}")
+    print(f"[TIME]  {_now_ist()}")
+    print(f"{'─'*60}")
     try:
         fs    = run_graph(user_message)
         route = fs.get("route", "")