Update app_gradio.py

app_gradio.py

@@ -1,4 +1,8 @@
-# app_gradio.py - Enhanced with streamlined UI
+#### `app_gradio.py`
+This file is significantly updated to manage real-time cost, prevent runaway loops, and handle state more robustly. A new helper class `TokenCostCalculator` is introduced for this purpose.
+
+```python
+# app_gradio.py - Enhanced with real-time cost tracking and safety features
 
 import patch_gradio_jsonschema  # MUST be first
 
@@ -10,6 +14,9 @@ import shutil
 import uuid
 from datetime import datetime
 from dotenv import load_dotenv
+from typing import Any, Dict, List
+from langchain_core.callbacks import BaseCallbackHandler
+from langchain_core.outputs import LLMResult
 
 from memory_manager import MemoryManager, memory_manager
 from graph import triage_app, planner_app, main_app
@@ -20,6 +27,7 @@ load_dotenv()
 log = get_logger(__name__)
 
 OUT_DIR = os.environ.get("OUT_DIR", "/tmp")
+MAX_EXECUTION_PATH_LENGTH = 50 # Safety break for runaway loops
 
 # Create necessary directories
 os.makedirs(OUT_DIR, exist_ok=True)
@@ -27,6 +35,38 @@ os.makedirs("outputs", exist_ok=True)
 os.makedirs("uploads", exist_ok=True)
 os.makedirs("conversations", exist_ok=True)
 
+# --- Real-time Cost Tracking ---
+class TokenCostCalculator(BaseCallbackHandler):
+    """Callback handler to calculate the cost of LLM calls."""
+    def __init__(self):
+        super().__init__()
+        self.total_cost = 0.0
+        self.model_costs = {
+            "gpt-4o": {"input": 0.005 / 1000, "output": 0.015 / 1000},
+            "default": {"input": 0.001 / 1000, "output": 0.002 / 1000} # Fallback
+        }
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        try:
+            model_name = response.llm_output.get("model_name", "default")
+            token_usage = response.llm_output.get("token_usage", {})
+            prompt_tokens = token_usage.get("prompt_tokens", 0)
+            completion_tokens = token_usage.get("completion_tokens", 0)
+
+            costs = self.model_costs.get(model_name, self.model_costs["default"])
+            
+            cost = (prompt_tokens * costs["input"]) + (completion_tokens * costs["output"])
+            self.total_cost += cost
+            log.info(f"LLM call cost: ${cost:.6f} | Total cost: ${self.total_cost:.4f}")
+        except Exception as e:
+            log.warning(f"Could not calculate token cost: {e}")
+
+    def get_total_cost(self) -> float:
+        return self.total_cost
+
+    def reset(self) -> None:
+        self.total_cost = 0.0
+
 # --- State Management ---
 def get_default_state():
     """Initialize clean session state."""
@@ -40,7 +80,9 @@ def get_default_state():
         "experimentResults": {},
         "approved": False,
         "execution_path": [],
-        "chatHistory": []
+        "chatHistory": [],
+        "current_cost": 0.0,
+        "budget_exceeded": False
     }
 
 def reset_all_artifacts_and_context():
@@ -82,10 +124,16 @@ def save_conversation(history, state):
     filename = f"conv_{timestamp}.json"
     filepath = os.path.join("conversations", filename)
     
+    # Ensure state is a dict before saving
+    conversation_data = {
+        "history": history,
+        "state": state if isinstance(state, dict) else {}
+    }
+    
     with open(filepath, 'w', encoding='utf-8') as f:
-        json.dump(history, f, indent=2)
+        json.dump(conversation_data, f, indent=2)
     
-    state_val = state or {}
+    state_val = state if isinstance(state, dict) else {}
     state_val["current_conversation_file"] = filepath
     log.info(f"Conversation saved: {filepath}")
     
@@ -96,13 +144,15 @@ def load_conversation(filepath, state):
         return [], state
     
     with open(filepath, 'r', encoding='utf-8') as f:
-        history = json.load(f)
+        conversation_data = json.load(f)
+    
+    history = conversation_data.get("history", [])
+    loaded_state = conversation_data.get("state", get_default_state())
     
-    state_val = state or {}
-    state_val["current_conversation_file"] = filepath
     log.info(f"Conversation loaded: {filepath}")
     
-    return history, state_val
+    return history, loaded_state
+
 
 def get_saved_conversations():
     if not os.path.exists("conversations"):
@@ -202,8 +252,10 @@ def collect_artifact_paths_from_result(exp_results):
             log.warning(f"Path normalization error: {e}")
             continue
     
+    # Return a unique list of paths
     return list(dict.fromkeys(paths))
 
+
 # --- Core Logic Functions ---
 def start_estimation(message, history, state):
     log.info(f"Starting estimation: '{message}'")
@@ -313,38 +365,56 @@ def execute_main_task(history, state, budget):
         "chatHistory": state.get("chatHistory", []),
         "max_loops": max_loops_calibrated,
         "status_update": "Task approved. Starting...",
-        "rework_cycles": state.get("rework_cycles", 0) or 0,
+        "rework_cycles": 0,
         "pmPlan": state.get("pmPlan", {}) or {},
         "experimentResults": state.get("experimentResults", {}) or {},
         "approved": False,
-        "execution_path": state.get("execution_path", []) or [],
+        "execution_path": [],
+        "current_cost": 0.0,
+        "budget_exceeded": False
     }
     
     aggregated_state = initial_state.copy()
+    cost_tracker = TokenCostCalculator()
     
     try:
-        for step in main_app.stream(initial_state):
+        config = {"callbacks": [cost_tracker]}
+        
+        for step in main_app.stream(initial_state, config=config):
             if not isinstance(step, dict) or not step:
                 continue
             
             node_name = list(step.keys())[0]
             node_output = step.get(node_name)
             
+            # --- Robust State Aggregation ---
             if not isinstance(node_output, dict):
                 if node_output is not None:
                     log.warning(f"Node '{node_name}' returned non-dict: {node_output}")
                     aggregated_state["status_update"] = str(node_output)
                 continue
             
-            # Merge output
             for key, value in node_output.items():
                 if key == 'execution_path':
-                    if isinstance(value, list):
-                        aggregated_state[key] = aggregated_state.get(key, []) + value
+                    current_path = aggregated_state.get(key, [])
+                    if isinstance(current_path, list):
+                         # Ensure we are always extending a list
+                        new_items = value if isinstance(value, list) else [value]
+                        aggregated_state[key] = current_path + new_items
                     else:
-                        aggregated_state[key] = aggregated_state.get(key, []) + [value]
+                        # Recover from corrupted state
+                        aggregated_state[key] = list(current_path) + (value if isinstance(value, list) else [value])
                 else:
                     aggregated_state[key] = value
+
+            # --- Real-time Budget Check ---
+            current_cost = cost_tracker.get_total_cost()
+            aggregated_state["current_cost"] = current_cost
+            if current_cost > float(budget):
+                log.warning(f"Budget exceeded! Current cost ${current_cost:.4f} > Budget ${float(budget):.4f}")
+                aggregated_state["budget_exceeded"] = True
+                aggregated_state["status_update"] = "Budget exceeded, finishing up..."
+                # No need to break here, the graph's `should_continue` will handle it
             
             status = aggregated_state.get("status_update", f"Running: {node_name}...")
             
@@ -356,12 +426,18 @@ def execute_main_task(history, state, budget):
                 aggregated_state, 
                 gr.update(visible=False), 
                 gr.update(interactive=False), 
-                f"🔄 {status}", 
+                f"🔄 {status} | Cost: ${current_cost:.3f}", 
                 [], 
                 update_artifact_list()
             )
-        
-        # Final processing
+            
+            # --- Safety Break for Runaway Loop ---
+            if len(aggregated_state.get("execution_path", [])) > MAX_EXECUTION_PATH_LENGTH:
+                log.error("Execution path limit reached. Terminating due to suspected runaway loop.")
+                aggregated_state["status_update"] = "ERROR: Runaway loop detected."
+                break
+
+        # Final processing after stream finishes or breaks
         final_response = aggregated_state.get("draftResponse", "No response generated.")
         exp_results = aggregated_state.get("experimentResults", {})
         
@@ -394,21 +470,22 @@ def execute_main_task(history, state, budget):
             artifact_md = "### Generated Artifacts:\n" + "\n".join(artifact_lines)
         
         history[-1] = ["", final_response]
-        final_state_json = json.dumps(aggregated_state, indent=2)
+        final_state_json = json.dumps(aggregated_state, indent=2, default=str) # Use default=str for safety
+        final_cost = cost_tracker.get_total_cost()
         
         yield (
             history, 
             final_state_json, 
             gr.update(visible=False), 
             gr.update(interactive=True), 
-            "✅ Task complete", 
+            f"✅ Task complete | Final Cost: ${final_cost:.3f}", 
             exported_paths, 
             artifact_md
         )
     
     except Exception as e:
         log.exception("Execution error")
-        final_state_json = json.dumps(aggregated_state, indent=2)
+        final_state_json = json.dumps(aggregated_state, indent=2, default=str)
         history[-1] = ["", f"❌ Execution failed: {e}"]
         
         yield (
@@ -421,6 +498,7 @@ def execute_main_task(history, state, budget):
             "No artifacts generated."
         )
 
+
 def cancel_task(history, state):
     history.append(["", "❌ Task cancelled."])
     return (