cost_benefit modified

src/cost_benefit.py

@@ -3,48 +3,58 @@ import subprocess
 import time
 import requests
 
-
-def get_best_model(weights: dict, runtime_env: str) -> dict:
-    #placeholders
-    models = {
-        "llama3.2": {"size": 2.5, "token_cost": 0.0001, "speed": 30},
-        "mistral": {"size": 4.2, "token_cost": 0.0002, "speed": 50},
-        "gemini-2.0-flash": {"size": 6.1, "token_cost": 0.0005, "speed": 60},
-        "gemini-2.5-pro-preview-03-25": {"size": 8.2, "token_cost": 0.002, "speed": 45}
-    }
-
-    penalty = {
-        "gpu": 1.0,
-        "cpu-local": 2.0,
-        "cloud-only": 1.5
+def detect_available_budget(runtime_env: str) -> int:
+    import torch
+    if "local" in runtime_env and torch.cuda.is_available():
+        total_vram_mb = torch.cuda.get_device_properties(0).total_memory // (1024 ** 2)
+        return min(total_vram_mb, 100)
+    else:
+        return 100
+
+
+def get_best_model(runtime_env: str, use_local_only=False, use_api_only=False) -> dict:
+    # Model info (cost, tokens/sec, type)
+    static_costs = {
+        "llama3.2": {"size": 20, "token_cost": 0.0001, "tokens_sec": 30, "type": "local"},
+        "mistral": {"size": 40, "token_cost": 0.0002, "tokens_sec": 50, "type": "local"},
+        "gemini-2.0-flash": {"size": 60, "token_cost": 0.0005, "tokens_sec": 60, "type": "api"},
+        "gemini-2.5-pro-preview-03-25": {"size": 80, "token_cost": 0.002, "tokens_sec": 45, "type": "api"}
     }
 
-    best_model = None
-    best_score = float("-inf")  # Track max score
+    def detect_available_budget(runtime_env: str) -> int:
+        import torch
+        if "local" in runtime_env and torch.cuda.is_available():
+            total_vram_mb = torch.cuda.get_device_properties(0).total_memory // (1024 ** 2)
+            return min(total_vram_mb, 100)
+        else:
+            return 100
 
-    for model, metrics in models.items():
-        p = penalty.get(runtime_env, 2.0)
+    budget = detect_available_budget(runtime_env)
 
-        cost_score = (
-            weights["w_size"] * metrics["size"] * p +
-            weights["w_token_cost"] * metrics["token_cost"] * p +
-            weights["w_speed"] * (100 - metrics["speed"])
-        )
-        benefit_score = weights["w_speed"] * metrics["speed"]
-
-        decision_score = benefit_score / cost_score if cost_score != 0 else 0
-
-        if decision_score > best_score:
-            best_score = decision_score
+    best_model = None
+    best_speed = -1
+
+    for model, info in static_costs.items():
+        if info["size"] > budget:
+            continue
+        if use_local_only and info["type"] != "local":
+            continue
+        if use_api_only and info["type"] != "api":
+            continue
+        if info["tokens_sec"] > best_speed:
             best_model = model
+            best_speed = info["tokens_sec"]
 
     if not best_model:
-        return "No suitable model found"
+        return {
+            "model": "llama3.2",
+            "token_cost": static_costs["llama3.2"]["token_cost"],
+            "tokens_sec": static_costs["llama3.2"]["tokens_sec"],
+            "note": "Defaulted due to no models fitting filters"
+        }
 
     return {
         "model": best_model,
-        "score": best_score,
-        "token_cost": models[best_model]["token_cost"],
-        "tokens_sec": models[best_model]["speed"],
-        "output": f"Sample output from {best_model}"
+        "token_cost": static_costs[best_model]["token_cost"],
+        "tokens_sec": static_costs[best_model]["tokens_sec"]
     }

src/manager/config/model_selector.py

@@ -7,24 +7,13 @@ load_dotenv()
 def choose_best_model(return_full=False):
     env = detect_runtime_environment()
     print(f"[INFO] Runtime Environment: {env}")
-    
-    weights = {
-        "w_size": 0.1,
-        "w_token_cost": 100,
-        "w_speed": 0.5
-    }
 
-    result = get_best_model(weights, env)
+    result = get_best_model(env)
 
-    if isinstance(result, str) or not result.get("model"):
-        if env == "cpu-local":
-            if os.getenv("GEMINI_KEY"):
-                print("[INFO] Falling back to Gemini for cpu-local.")
-                return {"model": "gemini-2.0-flash"} if return_full else "gemini-2.0-flash"
-            else:
-                print("[WARN] GOOGLE_API_KEY missing. Falling back to llama3.2.")
-                return {"model": "llama3.2"} if return_full else "llama3.2"
-        return {"model": "llama3.2"} if return_full else "llama3.2"
+    if not result.get("model"):
+        print("[WARN] No model found under budget — using fallback.")
+        fallback_model = "gemini-2.0-flash" if os.getenv("GEMINI_KEY") else "llama3.2"
+        return {"model": fallback_model} if return_full else fallback_model
 
-    print(f"[INFO] Auto-selected model: {result['model']}")
+    print(f"[INFO] Auto-selected model: {result['model']} (token cost: {result['token_cost']}, tokens/sec: {result['tokens_sec']})")
     return result if return_full else result["model"]

src/tools/default_tools/test_cost/agent_creator_tool.py

@@ -109,34 +109,39 @@ class AgentCreator():
 
     def run(self, **kwargs):
         print("Running Agent Creator")
+
         agent_name = kwargs.get("agent_name")
+        base_model = kwargs.get("base_model")
+
+        # NEW: read flags from kwargs
+        use_local_only = kwargs.get("use_local_only", False)
+        use_api_only = kwargs.get("use_api_only", False)
 
-        # Get full model info (not just name)
-        model_info = choose_best_model(return_full=True)
-        base_model = kwargs.get("base_model") or choose_best_model()
-        base_model = model_info["model"]
-        token_cost = model_info.get("token_cost", 0.0001)
-        speed = model_info.get("tokens_sec", 30)
-        score = model_info.get("score", 1)
+        if not base_model:
+            env = detect_runtime_environment()
+            print(f"\n[DEBUG] Detected Runtime Environment: {env}")
+            
+            from src.cost_benefit import get_best_model
+            model_meta = get_best_model(
+                runtime_env=env,
+                use_local_only=use_local_only,
+                use_api_only=use_api_only
+            )
+            base_model = model_meta["model"]
+        else:
+            model_meta = {"model": base_model}
 
-        env = detect_runtime_environment()
-        print(f"\n[DEBUG] Detected Runtime Environment: {env}")
         print(f"[DEBUG] Selected Model: {base_model}")
-        print(f"[DEBUG] Token Cost: {token_cost}, Speed: {speed}, Score: {score}")
+
+        if base_model not in self.inputSchema["creates"]["types"]:
+            print(f"[WARN] Auto-selected model '{base_model}' not in schema. Falling back to gemini-2.0-flash")
+            base_model = "gemini-2.0-flash"
 
         system_prompt = kwargs.get("system_prompt")
         description = kwargs.get("description")
-        #create_cost = self.inputSchema["creates"]["types"][base_model]["create_cost"]
-        #if base_model not in self.inputSchema["creates"]["types"]:
-        #    print(f"[WARN] Auto-selected model '{base_model}' not in schema. Falling back to gemini-2.0-flash")
-        #    base_model = "gemini-2.0-flash"
-        #invoke_cost = self.inputSchema["creates"]["types"][base_model]["invoke_cost"]
-
-        # Dynamically calculated costs
-        create_cost = round(10 + (token_cost * 10000) + (50 / (speed + 1)), 2)
-        invoke_cost = round(create_cost * 2, 2)
 
-        print(f"[INFO] Assigned Create Cost: {create_cost}, Invoke Cost: {invoke_cost}")
+        create_cost = self.inputSchema["creates"]["types"][base_model]["create_cost"]
+        invoke_cost = self.inputSchema["creates"]["types"][base_model]["invoke_cost"]
 
         agent_manager = AgentManager()
         try:
@@ -157,8 +162,7 @@ class AgentCreator():
 
         return {
             "status": "success",
-            "message": "Agent successfully created",
+            "message": f"Agent '{agent_name}' created using model '{base_model}'",
+            "model_info": model_meta,
             "remaining_budget": remaining_budget,
         }
-
-