feat: add TaskAwareRouter - routes queries by task type and complexit…

custom_routers/taskawarerouter/README.md

@@ -0,0 +1,102 @@
+# TaskAwareRouter
+
+TaskAwareRouter is a custom inference router for LLMRouter that selects an appropriate LLM based on task domain and estimated complexity that selects the optimal LLM based on
+two signals extracted from the user query:
+
+1. **Task type** — coding / design / planning / research
+2. **Complexity** — simple / complex
+
+## How it works
+
+A small cheap LLM (qwen2.5-7b) judges the incoming query and returns
+task type and complexity as JSON. The router then looks up the right
+model from a task map built from available llm_data.
+
+If the judge LLM fails (timeout, missing API key, network error),
+the router falls back to keyword matching and safe defaults.
+If task classification fails, the router falls back to rule-based keyword routing to ensure deterministic behavior.
+
+## PIPELINE
+User Query
+      │
+      ▼
+Judge LLM (Qwen2.5-7B)
+      │
+      ▼
+Task + Complexity
+      │
+      ▼
+Routing Table
+      │
+      ▼
+Selected LLM
+
+
+## Routing table
+
+| Task     | Simple                      | Complex                          |
+|----------|-----------------------------|----------------------------------|
+| coding   | mistral-7b-instruct-v0.3    | llama3-70b-instruct              |
+| design   | llama-3.1-8b-instruct       | mixtral-8x22b-instruct-v0.1      |
+| planning | qwen2.5-7b-instruct         | llama-3.3-nemotron-super-49b-v1  |
+| research | qwen2.5-7b-instruct         | mixtral-8x22b-instruct-v0.1      |
+| language | llama-3.1-8b-instruct       | mixtral-8x22b-instruct-v0.1      |
+
+## Free tier limitation
+
+During development and validation, only `meta/llama-3.1-8b-instruct`
+was available under the free NVIDIA NIM tier. All task_map entries
+currently point to this model. The routing logic correctly selects
+model tiers — swap in a full-access API key to activate
+differentiated routing across all models.
+
+
+
+## Usage
+
+```bash
+llmrouter infer \
+  --router taskawarerouter \
+  --config custom_routers/taskawarerouter/config.yaml \
+  --query "Build a production grade authentication system" \
+  --route-only
+```
+
+## Motivation
+
+No existing router in LLMRouter routes by task domain and complexity
+combined. This router fills that gap using a lightweight LLM judge
+instead of hardcoded keywords or training data.
+
+
+
+## Running tests
+
+```bash
+python -m pytest custom_routers/taskawarerouter/test_router.py -v
+```
+
+## Limitations
+
+- Task map model names must exist in your llm_data config
+- Judge LLM requires a valid API_KEYS environment variable
+- Task classification accuracy depends on judge model quality
+- Free NVIDIA tier limits available models for validation
+
+## Why this router
+
+No existing router in LLMRouter routes by task domain and complexity
+combined. This router fills that gap using a lightweight LLM judge
+instead of hardcoded keywords or training data. It handles any human
+language naturally without requiring predefined keyword lists.
+
+## Future Work
+
+- Confidence-based routing
+- Cost-aware model selection
+- Latency-aware routing
+- User feedback learning
+
+## Author
+Vidhursh Kumar V
+GitHub: @Vidhursh-16

custom_routers/taskawarerouter/config.yaml

@@ -0,0 +1,20 @@
+# Config for TaskAwareRouter
+data_path:
+  query_data_train: 'data/example_data/query_data/default_query_train.jsonl'
+  query_data_test: 'data/example_data/query_data/default_query_test.jsonl'
+  query_embedding_data: 'data/example_data/routing_data/query_embeddings_longformer.pt'
+  routing_data_train: 'data/example_data/routing_data/default_routing_train_data.jsonl'
+  routing_data_test: 'data/example_data/routing_data/default_routing_test_data.jsonl'
+  llm_data: 'data/example_data/llm_candidates/default_llm.json'
+  llm_embedding_data: 'data/example_data/llm_candidates/default_llm_embeddings.json'
+
+router:
+  judge_model: 'llama-3.1-8b-instruct'
+  judge_api_model: 'meta/llama-3.1-8b-instruct'
+  judge_timeout: 30
+
+metric:
+  weights:
+    performance: 1
+    cost: 0
+    llm_judge: 0

custom_routers/taskawarerouter/router.py

@@ -0,0 +1,206 @@
+# ── Standard library imports ──────────────────────────────────────────
+import os
+import json
+
+
+import urllib.request
+import urllib.error
+
+import yaml
+
+# ── LLMRouter imports ─────────────────────────────────────────────────
+from llmrouter.models.meta_router import MetaRouter
+import torch.nn as nn
+
+# ── HTTP library to call judge LLM ────────────────────────────────────
+import urllib.request
+
+
+# ═════════════════════════════════════════════════════════════════════
+class TaskAwareRouter(MetaRouter):
+    """
+    Routes queries to the right LLM based on:
+      1. Task type  (coding / design / planning / research)
+      2. Complexity (simple / complex)
+
+    A small cheap LLM judges both signals from the query.
+    If judge fails, falls back to safe defaults.
+    """
+
+    # ── Startup ───────────────────────────────────────────────────────
+    def __init__(self, yaml_path: str):
+        super().__init__(model=nn.Identity(), yaml_path=yaml_path)
+
+        # Real model names loaded from config → no hardcoding
+        self.llm_names = list(self.llm_data.keys())
+
+        # Judge model = smallest cheapest available
+        import yaml
+        with open(yaml_path, "r") as f:
+             _cfg = yaml.safe_load(f)
+             router_cfg = _cfg.get("router", {})
+             self.judge_model     = router_cfg.get("judge_model", self.llm_names[0])
+             self.judge_api_model = router_cfg.get("judge_api_model", "")
+             self.judge_timeout   = router_cfg.get("judge_timeout", 30)
+
+        # Task → [simple model, complex model]
+        # Built from actual available models by price
+        self.task_map = {
+    "coding": {
+        "simple":  "mistral-7b-instruct-v0.3",
+        "complex": "llama3-70b-instruct",
+    },
+    "design": {
+        "simple":  "llama-3.1-8b-instruct",
+        "complex": "mixtral-8x22b-instruct-v0.1",
+    },
+    "planning": {
+        "simple":  "qwen2.5-7b-instruct",
+        "complex": "llama-3.3-nemotron-super-49b-v1",
+    },
+    "research": {
+        "simple":  "qwen2.5-7b-instruct",
+        "complex": "mixtral-8x22b-instruct-v0.1",
+    },
+    "language": {
+        "simple":  "llama-3.1-8b-instruct",
+        "complex": "mixtral-8x22b-instruct-v0.1",
+    },
+}
+
+        # Keyword fallback if judge LLM fails
+        self.complex_keywords = [
+            "advanced", "complex", "production", "scale",
+            "optimize", "architect", "integrate", "enterprise",
+            "microservices", "distributed", "secure", "pipeline"
+        ]
+
+    # ── Judge LLM: classify task + complexity in one call ─────────────
+    def _llm_judge(self, query: str) -> dict:
+        """
+        Sends query to smallest model.
+        Asks it to return task type and complexity as JSON.
+        """
+        prompt = f"""You are a query classifier. Given a user query, return ONLY a JSON object with two fields:
+- "task": one of ["coding", "design", "planning", "research","language"]
+- "complexity": one of ["simple", "complex"]
+
+Rules:
+- coding    = writing code, debugging, APIs, backend, frontend implementation
+- design    = UI layout, visual design, wireframes, user experience
+- planning  = project plans, roadmaps, strategy, timelines
+- research  = finding information, comparing options, summarizing topics,"what is", "which is better"
+- language  = translation, grammar correction, rewriting, editing, summarizing text
+- simple    = straightforward, single step, clear requirement
+- complex   = multi-step, production grade, architecture level, ambiguous
+Important:
+- "compare", "best", "options", "which", "recommend" → always research
+- "build", "create", "implement", "fix", "write" → always coding
+- "design", "layout", "wireframe", "UI" → always design
+- "plan", "roadmap", "strategy", "timeline" → always planning
+- language = translation, grammar, writing, summarizing text, editing
+
+Query: {query}
+
+Return ONLY the JSON. No explanation. No markdown. Example: {{"task": "coding", "complexity": "simple"}}"""
+
+        # Build API request
+        judge_model_info = self.llm_data.get(self.judge_model, {})
+        api_endpoint      = judge_model_info.get("api_endpoint", "")
+        model_name        = self.judge_api_model    
+        api_key           = os.environ.get("API_KEYS", "")
+
+        payload = json.dumps({
+            "model": model_name,
+            "messages": [{"role": "user", "content": prompt}],
+            "max_tokens": 50,
+            "temperature": 0,
+        }).encode("utf-8")
+
+        req = urllib.request.Request(
+            f"{api_endpoint}/chat/completions",
+            data=payload,
+            headers={
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {api_key}",
+            },
+            method="POST",
+        )
+
+        with urllib.request.urlopen(req, timeout=self.judge_timeout) as resp:
+            result  = json.loads(resp.read().decode())
+            content = result["choices"][0]["message"]["content"].strip()
+            return json.loads(content)
+
+    # ── Keyword fallback if judge fails ───────────────────────────────
+    def _keyword_fallback(self, query: str) -> dict:
+        """
+        Simple rule based fallback.
+        Used only when judge LLM is unavailable.
+        """
+        complexity = "simple"
+        for keyword in self.complex_keywords:
+            if keyword in query.lower():
+                complexity = "complex"
+                break
+        return {"task": "coding", "complexity": complexity}
+
+    # ── Pick model from task map ──────────────────────────────────────
+    def _pick_model(self, task: str, complexity: str) -> str:
+        """
+        Looks up task_map → returns model name.
+        Falls back to first available model if not found.
+        """
+        task_entry = self.task_map.get(task, {})
+        preferred  = task_entry.get(complexity, "")
+
+        if preferred and preferred in self.llm_data:
+            return preferred
+
+        # Safety net → never crash
+        return self.llm_names[0]
+
+    # ── Main function LLMRouter calls ─────────────────────────────────
+    def route_single(self, query_input: dict) -> dict:
+        """
+        Entry point. Called by LLMRouter for every query.
+        """
+        query = query_input.get("query", "")
+
+        # Step 1 → judge task + complexity
+        try:
+            judgment   = self._llm_judge(query)
+            task       = judgment.get("task", "coding")
+            complexity = judgment.get("complexity", "simple")
+        except urllib.error.HTTPError as e:
+            print(f"[TaskAwareRouter] Judge HTTP error: {e.code} {e.reason}")
+            judgment   = self._keyword_fallback(query)
+            task       = judgment["task"]
+            complexity = judgment["complexity"]
+        except urllib.error.URLError as e:
+             print(f"[TaskAwareRouter] Judge connection error: {e.reason}") 
+             judgment   = self._keyword_fallback(query)
+             task       = judgment["task"]
+             complexity = judgment["complexity"]
+        except json.JSONDecodeError as e:
+            print(f"[TaskAwareRouter] Judge response parse error: {e}")
+            judgment   = self._keyword_fallback(query)
+            task       = judgment["task"]
+            complexity = judgment["complexity"]
+
+        # Step 2 → pick right model
+        model = self._pick_model(task, complexity)
+
+        # Step 3 → return result
+        return {
+            "query":         query,
+            "task":          task,
+            "complexity":    complexity,
+            "model_name":    model,
+            "predicted_llm": model,
+        }
+
+    # ── Batch routing ─────────────────────────────────────────────────
+    def route_batch(self, batch: list) -> list:
+        """Run route_single for every query in a list."""
+        return [self.route_single(q) for q in batch]

custom_routers/taskawarerouter/test_router.py

@@ -0,0 +1,34 @@
+import pytest
+import sys
+import os
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../")))
+
+from custom_routers.taskawarerouter.router import TaskAwareRouter
+
+CONFIG = "custom_routers/taskawarerouter/config.yaml"
+
+
+# ── Test 1: keyword fallback works when judge fails ───────────────────
+def test_keyword_fallback():
+    router = TaskAwareRouter(yaml_path=CONFIG)
+    result = router._keyword_fallback("build an enterprise scale pipeline")
+    assert result["task"] == "coding"
+    assert result["complexity"] == "complex"
+
+
+# ── Test 2: unknown task → default model, never crashes ───────────────
+def test_unknown_task_default_model():
+    router = TaskAwareRouter(yaml_path=CONFIG)
+    model = router._pick_model("unknown_task", "simple")
+    assert model in router.llm_names
+
+
+# ── Test 3: successful routing returns required keys ──────────────────
+def test_route_single_returns_required_keys():
+    router = TaskAwareRouter(yaml_path=CONFIG)
+    result = router.route_single({"query": "write a python function"})
+    assert "model_name" in result
+    assert "predicted_llm" in result
+    assert "task" in result
+    assert "complexity" in result