diff --git a/custom_routers/taskawarerouter/README.md b/custom_routers/taskawarerouter/README.md new file mode 100644 index 00000000..33fb79db --- /dev/null +++ b/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 \ No newline at end of file diff --git a/custom_routers/taskawarerouter/config.yaml b/custom_routers/taskawarerouter/config.yaml new file mode 100644 index 00000000..d1b03790 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/custom_routers/taskawarerouter/router.py b/custom_routers/taskawarerouter/router.py new file mode 100644 index 00000000..41135c3e --- /dev/null +++ b/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] \ No newline at end of file diff --git a/custom_routers/taskawarerouter/test_router.py b/custom_routers/taskawarerouter/test_router.py new file mode 100644 index 00000000..d651a492 --- /dev/null +++ b/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 \ No newline at end of file