Refactor and cleanup of project for fork.

README.md

@@ -118,6 +118,8 @@ pip install -e .
 
 # Optional: Install with RouterR1 support (requires GPU)
 # RouterR1 is tested with vllm==0.6.3 (torch==2.4.0); the extra pins these versions.
+# NOTE: these pins have no wheels for Python >= 3.14 — use Python 3.10–3.13 for the
+# [router-r1]/[all] extras. The base package supports Python 3.10–3.14.
 pip install -e ".[router-r1]"
 
 # Optional: Install all optional dependencies
@@ -130,6 +132,13 @@ pip install -e ".[all]"
 pip install llmrouter-lib
 ```
 
+> **🔒 Security note — choose Python 3.10–3.13 for production.** `litellm` (the LLM API-calling
+> layer) dropped Python 3.14 support at 1.83.8, so on 3.14 it is capped at 1.83.7, which carries known
+> CVEs and pins vulnerable transitive deps. On **Python 3.10–3.13** the dependency floor resolves
+> `litellm` to its fully-patched line and `pip-audit` reports no known vulnerabilities. Python 3.14 is
+> supported for development; for security-sensitive deployments use 3.10–3.13 until `litellm` restores
+> 3.14 support. See [.agent-reviews/dependency-audit.md](.agent-reviews/dependency-audit.md).
+
 ### 🔑 Setting Up API Keys
 
 LLMRouter requires API keys to make LLM API calls for inference, chat, and data generation. Set the `API_KEYS` environment variable using one of the following formats:

custom_routers/randomrouter/router.py

@@ -10,7 +10,7 @@
 """
 
 import random
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List
 import torch.nn as nn
 
 from llmrouter.models.meta_router import MetaRouter

custom_routers/thresholdrouter/router.py

@@ -5,7 +5,6 @@
 from typing import Any, Dict, List
 import torch
 import torch.nn as nn
-import numpy as np
 
 from llmrouter.models.meta_router import MetaRouter
 
@@ -129,9 +128,9 @@ def route_single(self, query_input: Dict[str, Any]) -> Dict[str, Any]:
             if not isinstance(embedding, torch.Tensor):
                 embedding = torch.tensor(embedding, dtype=torch.float32)
         elif hasattr(self, 'query_embeddings') and 'query' in query_input:
-            # Try to get from loaded embeddings (if available)
-            query = query_input['query']
+            # Try to get from loaded embeddings (if available).
             # This is a simplified version - real implementation would hash or lookup
+            # the query embedding here; for now we require an explicit 'embedding'.
             raise ValueError(
                 "Query embedding not provided. "
                 "Pass 'embedding' in query_input or implement embedding generation."

llmrouter/cli/router_chat.py

@@ -9,7 +9,7 @@
 import argparse
 import os
 import yaml
-from typing import Dict, Any, Optional
+from typing import Any, Optional
 
 import gradio as gr
 

llmrouter/data/api_calling_evaluation.py

@@ -21,19 +21,16 @@
 import sys
 import time
 import json
-import ast
 import re
 import argparse
 import yaml
-from typing import Dict, List, Tuple, Optional, Union
+from typing import Dict, List
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor, as_completed
-import threading
 from pathlib import Path
 
 import pandas as pd
 import numpy as np
-import torch
 from tqdm import tqdm
 
 # Allow importing local helper packages under repo `data/` (e.g., `human_eval`, `mbpp`)
@@ -45,16 +42,13 @@
 # Import utils
 from llmrouter.utils import (
     setup_environment,
-    format_mc_prompt, format_gsm8k_prompt, format_math_prompt,
-    format_commonsense_qa_prompt, format_mbpp_prompt, format_humaneval_prompt,
-    generate_task_query, ProgressTracker, to_tensor, clean_df,
-    process_final_data, call_api
+    generate_task_query, ProgressTracker, call_api
 )
 from llmrouter.utils.data_processing import process_unified_embeddings_and_routing
 from llmrouter.data.data_loader import DataLoader
 
 # Import evaluation functions
-from llmrouter.utils import f1_score, exact_match_score, get_bert_score, evaluate_code, cem_score
+from llmrouter.utils import f1_score, exact_match_score, get_bert_score, cem_score
 from llmrouter.utils.evaluation import last_boxed_only_string, remove_boxed, is_equiv
 try:
     from human_eval.evaluate_functional_correctness import entry_point_item
@@ -321,7 +315,7 @@ def eval_perf(metric, prediction, ground_truth, task_name, task_id=None):
                 answer = remove_boxed(string_in_last_boxed)
                 if is_equiv(answer, ground_truth_processed):
                     return 1
-        except Exception as e:
+        except Exception:
             return 0
         return 0
 

llmrouter/data/data.py

@@ -53,11 +53,11 @@
     ```
 """
 
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional
 from abc import ABC, abstractmethod
 from enum import Enum
 
-from pydantic import BaseModel, ConfigDict, Field, field_validator
+from pydantic import BaseModel, Field, field_validator
 
 
 class DataFormatType(Enum):

llmrouter/data/data_generation.py

@@ -31,18 +31,13 @@
 import json
 import argparse
 import yaml
-from pathlib import Path
-from typing import Dict, List, Tuple, Optional, Union
-from collections import defaultdict
 
-import pandas as pd
 import numpy as np
-from tqdm import tqdm
 from datasets import load_dataset
 
 # Import utils
 from llmrouter.utils import (
-    setup_environment, TASK_DESCRIPTIONS, CASE_NUM
+    setup_environment, CASE_NUM
 )
 from llmrouter.data.data_loader import DataLoader
 from llmrouter.data import batch_vlm_describe_images

llmrouter/data/data_loader.py

@@ -1,6 +1,6 @@
 import os
 import json
-from llmrouter.utils import load_csv, load_jsonl, jsonl_to_csv, load_pt
+from llmrouter.utils import load_jsonl, jsonl_to_csv, load_pt
 
 
 def load_json_file(path: str):

llmrouter/data/generate_llm_embeddings.py

@@ -18,8 +18,7 @@
 import json
 import argparse
 import yaml
-from pathlib import Path
-from typing import Dict, List
+from typing import Dict
 
 from llmrouter.utils import setup_environment, get_longformer_embedding
 from llmrouter.data.data_loader import DataLoader

llmrouter/models/__init__.py

@@ -1,107 +1,70 @@
-from .meta_router import MetaRouter
-from .base_trainer import BaseTrainer
-
-from .smallest_llm import SmallestLLM
-from .largest_llm import LargestLLM
-
-from .knnrouter import KNNRouter
-from .knnrouter import KNNRouterTrainer
-
-from .svmrouter import SVMRouter
-from .svmrouter import SVMRouterTrainer
-
-from .mlprouter import MLPRouter
-from .mlprouter import MLPTrainer
-
-from .mfrouter import MFRouter
-from .mfrouter import MFRouterTrainer
-
-from .elorouter import EloRouter
-from .elorouter import EloRouterTrainer
-
-from .automix import AutomixRouter
-from .automix import AutomixRouterTrainer
-
-from .routerdc import DCRouter
-from .routerdc import DCTrainer
-
-from .hybrid_llm import HybridLLMRouter
-from .hybrid_llm import HybridLLMTrainer
-
-try:
-    from .graphrouter import GraphRouter
-    from .graphrouter import GraphTrainer
-except Exception:
-    GraphRouter = None
-    GraphTrainer = None
-
-try:
-    from .causallm_router import CausalLMRouter
-    from .causallm_router import CausalLMTrainer
-except Exception:
-    CausalLMRouter = None
-    CausalLMTrainer = None
-
-try:
-    from .router_r1 import RouterR1
-except Exception:
-    RouterR1 = None
-
-try:
-    from .gmtrouter import GMTRouter
-    from .gmtrouter import GMTRouterTrainer
-except Exception:
-    GMTRouter = None
-    GMTRouterTrainer = None
-
-try:
-    from .personalizedrouter import PersonalizedRouter
-    from .personalizedrouter import PersonalizedRouterTrainer
-except Exception:
-    PersonalizedRouter = None
-    PersonalizedRouterTrainer = None
-
-__all__ = [
-    "MetaRouter",
-    "BaseTrainer",
-    "SmallestLLM",
-    "LargestLLM",
-
-    "KNNRouter",
-    "KNNRouterTrainer",
-
-    "SVMRouter",
-    "SVMRouterTrainer",
-
-    "MLPRouter",
-    "MLPTrainer",
-
-    "MFRouter",
-    "MFRouterTrainer",
-
-    "EloRouter",
-    "EloRouterTrainer",
-
-    "DCRouter",
-    "DCTrainer",
-
-    "AutomixRouter",
-    "AutomixRouterTrainer",
-
-    "HybridLLMRouter",
-    "HybridLLMTrainer",
-
-    "GraphRouter",
-    "GraphTrainer",
-
-    "CausalLMRouter",
-    "CausalLMTrainer",
-
-    "RouterR1",
-
-    "GMTRouter",
-    "GMTRouterTrainer",
-
-    "PersonalizedRouter",
-    "PersonalizedRouterTrainer",
-]
+"""Router and trainer registry.
+
+Names are resolved **lazily** (PEP 562 ``__getattr__``) so that
+``import llmrouter.models`` does not eagerly import torch / torch-geometric /
+transformers / peft for every router. Importing a specific name — e.g.
+``from llmrouter.models import KNNRouter`` — loads only that router's submodule
+(and hence only its dependencies). Optional routers whose heavy dependencies are
+unavailable resolve to ``None``, preserving the previous behavior.
+"""
+import importlib
+
+# public attribute name -> (submodule, is_optional)
+# optional entries resolve to None when their (heavy/GPU) dependencies are missing.
+_REGISTRY = {
+    "MetaRouter": ("meta_router", False),
+    "BaseTrainer": ("base_trainer", False),
+    "SmallestLLM": ("smallest_llm", False),
+    "LargestLLM": ("largest_llm", False),
+    "KNNRouter": ("knnrouter", False),
+    "KNNRouterTrainer": ("knnrouter", False),
+    "SVMRouter": ("svmrouter", False),
+    "SVMRouterTrainer": ("svmrouter", False),
+    "MLPRouter": ("mlprouter", False),
+    "MLPTrainer": ("mlprouter", False),
+    "MFRouter": ("mfrouter", False),
+    "MFRouterTrainer": ("mfrouter", False),
+    "EloRouter": ("elorouter", False),
+    "EloRouterTrainer": ("elorouter", False),
+    "AutomixRouter": ("automix", False),
+    "AutomixRouterTrainer": ("automix", False),
+    "DCRouter": ("routerdc", False),
+    "DCTrainer": ("routerdc", False),
+    "HybridLLMRouter": ("hybrid_llm", False),
+    "HybridLLMTrainer": ("hybrid_llm", False),
+    "GraphRouter": ("graphrouter", True),
+    "GraphTrainer": ("graphrouter", True),
+    "CausalLMRouter": ("causallm_router", True),
+    "CausalLMTrainer": ("causallm_router", True),
+    "RouterR1": ("router_r1", True),
+    "GMTRouter": ("gmtrouter", True),
+    "GMTRouterTrainer": ("gmtrouter", True),
+    "PersonalizedRouter": ("personalizedrouter", True),
+    "PersonalizedRouterTrainer": ("personalizedrouter", True),
+}
+
+__all__ = list(_REGISTRY)
+
+
+def __getattr__(name):
+    """Lazily import and cache a registered router/trainer on first access."""
+    try:
+        submodule, optional = _REGISTRY[name]
+    except KeyError:
+        raise AttributeError(
+            f"module {__name__!r} has no attribute {name!r}"
+        ) from None
+    try:
+        module = importlib.import_module(f".{submodule}", __name__)
+        value = getattr(module, name)
+    except Exception:
+        if optional:
+            value = None  # optional router with missing deps -> None (legacy behavior)
+        else:
+            raise
+    globals()[name] = value  # cache so later lookups bypass __getattr__
+    return value
+
+
+def __dir__():
+    return sorted(__all__)

llmrouter/models/automix/model.py

@@ -10,7 +10,7 @@
 Adapted for LLMRouter framework with PyTorch nn.Module interface.
 """
 
-from typing import List, Union
+from typing import List
 
 import numpy as np
 import pandas as pd
@@ -230,11 +230,9 @@ def safe_max_key(d):
                 """Find key with maximum value, handling inf/-inf."""
                 max_val = float('-inf')
                 max_key = None
-                has_inf = False
                 for k, v in d.items():
                     if isinstance(v, float):
                         if v == float('inf'):
-                            has_inf = True
                             max_key = k  # Prefer inf over finite values
                             break
                     if v > max_val: