More improvements.

modules/utils.py

@@ -4,6 +4,7 @@ import datetime
 import time
 import config
 import threading
+import re
 
 SAFETY_SETTINGS = [
     {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
@@ -15,6 +16,9 @@ SAFETY_SETTINGS = [
 # Thread-local storage for logging context
 _log_context = threading.local()
 
+# Cache for dynamic pricing from AI model selection
+PRICING_CACHE = {}
+
 def set_log_file(filepath):
     _log_context.log_file = filepath
 
@@ -136,32 +140,85 @@ def get_latest_run_folder(base_name):
     runs.sort(key=lambda x: int(x.split('_')[1]) if x.split('_')[1].isdigit() else 0)
     return os.path.join(base_name, runs[-1])
 
+def update_pricing(model_name, cost_str):
+    """Parses cost string from AI selection and updates cache."""
+    if not model_name or not cost_str or cost_str == 'N/A': return
+    
+    try:
+        # Look for patterns like "$0.075 Input" or "$3.50/1M"
+        # Default to 0.0
+        in_cost = 0.0
+        out_cost = 0.0
+        
+        # Extract all float-like numbers following a $ sign
+        prices = re.findall(r'(?:\$|USD)\s*([0-9]+\.?[0-9]*)', cost_str, re.IGNORECASE)
+        
+        if len(prices) >= 2:
+            in_cost = float(prices[0])
+            out_cost = float(prices[1])
+        elif len(prices) == 1:
+            in_cost = float(prices[0])
+            out_cost = in_cost * 3 # Rough heuristic if only one price provided
+            
+        if in_cost > 0:
+            PRICING_CACHE[model_name] = {"input": in_cost, "output": out_cost}
+            # log("SYSTEM", f"Updated pricing for {model_name}: In=${in_cost} | Out=${out_cost}")
+    except:
+        pass
+
+def calculate_cost(model_label, input_tokens, output_tokens, image_count=0):
+    cost = 0.0
+    m = model_label.lower()
+    
+    # Check dynamic cache first
+    if model_label in PRICING_CACHE:
+        rates = PRICING_CACHE[model_label]
+        cost = (input_tokens / 1_000_000 * rates['input']) + (output_tokens / 1_000_000 * rates['output'])
+    elif 'imagen' in m or image_count > 0:
+        cost = (image_count * 0.04)
+    else:
+        # Fallbacks
+        if 'flash' in m:
+            cost = (input_tokens / 1_000_000 * 0.075) + (output_tokens / 1_000_000 * 0.30)
+        elif 'pro' in m or 'logic' in m:
+            cost = (input_tokens / 1_000_000 * 3.50) + (output_tokens / 1_000_000 * 10.50)
+            
+    return round(cost, 6)
+
 def log_usage(folder, model_label, usage_metadata=None, image_count=0):
     if not folder or not os.path.exists(folder): return
     
     log_path = os.path.join(folder, "usage_log.json")
     
-    entry = {
-        "timestamp": int(time.time()),
-        "model": model_label,
-        "input_tokens": 0,
-        "output_tokens": 0,
-        "images": image_count
-    }
+    input_tokens = 0
+    output_tokens = 0
     
     if usage_metadata:
         try:
-            entry["input_tokens"] = usage_metadata.prompt_token_count
-            entry["output_tokens"] = usage_metadata.candidates_token_count
+            input_tokens = usage_metadata.prompt_token_count
+            output_tokens = usage_metadata.candidates_token_count
         except: pass
 
+    # Calculate Cost
+    cost = calculate_cost(model_label, input_tokens, output_tokens, image_count)
+
+    entry = {
+        "timestamp": int(time.time()),
+        "date": datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
+        "model": model_label,
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "images": image_count,
+        "cost": round(cost, 6)
+    }
+
     data = {"log": [], "totals": {"input_tokens": 0, "output_tokens": 0, "images": 0, "est_cost_usd": 0.0}}
     
     if os.path.exists(log_path):
         try:
             loaded = json.load(open(log_path, 'r'))
             if isinstance(loaded, list): data["log"] = loaded
-            else: data = loaded
+            elif isinstance(loaded, dict): data = loaded
         except: pass
         
     data["log"].append(entry)
@@ -171,25 +228,28 @@ def log_usage(folder, model_label, usage_metadata=None, image_count=0):
     t_out = sum(x.get('output_tokens', 0) for x in data["log"])
     t_img = sum(x.get('images', 0) for x in data["log"])
     
-    cost = 0.0
+    total_cost = 0.0
     for x in data["log"]:
-        m = x.get('model', '').lower()
-        i = x.get('input_tokens', 0)
-        o = x.get('output_tokens', 0)
-        imgs = x.get('images', 0)
-        
-        if 'flash' in m:
-            cost += (i / 1_000_000 * 0.075) + (o / 1_000_000 * 0.30)
-        elif 'pro' in m or 'logic' in m:
-            cost += (i / 1_000_000 * 3.50) + (o / 1_000_000 * 10.50)
-        elif 'imagen' in m or imgs > 0:
-            cost += (imgs * 0.04)
+        if 'cost' in x:
+            total_cost += x['cost']
+        else:
+            # Fallback calculation for old logs without explicit cost field
+            c = 0.0
+            mx = x.get('model', '').lower()
+            ix = x.get('input_tokens', 0)
+            ox = x.get('output_tokens', 0)
+            imgx = x.get('images', 0)
+            
+            if 'flash' in mx: c = (ix / 1_000_000 * 0.075) + (ox / 1_000_000 * 0.30)
+            elif 'pro' in mx or 'logic' in mx: c = (ix / 1_000_000 * 3.50) + (ox / 1_000_000 * 10.50)
+            elif 'imagen' in mx or imgx > 0: c = (imgx * 0.04)
+            total_cost += c
             
     data["totals"] = {
         "input_tokens": t_in,
         "output_tokens": t_out,
         "images": t_img,
-        "est_cost_usd": round(cost, 4)
+        "est_cost_usd": round(total_cost, 4)
     }
     
     with open(log_path, 'w') as f: json.dump(data, f, indent=2)