Skip to content

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

Open
Vidhursh-16 wants to merge 1 commit into
ulab-uiuc:mainfrom
Vidhursh-16:feature/task-aware-router
Open

feat: add TaskAwareRouter - routes queries by task type and complexit…#182
Vidhursh-16 wants to merge 1 commit into
ulab-uiuc:mainfrom
Vidhursh-16:feature/task-aware-router

Conversation

@Vidhursh-16

@Vidhursh-16 Vidhursh-16 commented Jun 26, 2026

Copy link
Copy Markdown

…y using LLM judge

Description

Introduced TaskAwareRouter, a custom router that selects the most suitable LLM by analyzing the user's query. It considers both the type of task (coding, design, planning, research, or language) and its complexity (simple or complex) before choosing a model.
based on two signals

  • Task type: coding / design / planning / research / language
  • Complexity: simple / complex

Related Issues

Type of Change

  • New feature (non-breaking change that adds functionality)
  • New router implementation
  • Test addition or update
  • Documentation update
  • Bug fix (non-breaking change that fixes an issue)
  • New feature (non-breaking change that adds functionality)
  • New router implementation
  • Breaking change (fix or feature that would cause existing functionality to change)
  • Documentation update
  • Performance improvement
  • Code refactoring
  • Test addition or update

Changes Made

  • Added TaskAwareRouter in custom_routers/taskawarerouter/
  • Added router.py with LLM judge for task type and complexity detection
  • Added config.yaml with judge_model, judge_api_model, judge_timeout fields
  • Added test_router.py with 3 passing tests
  • Added README.md documenting routing table and free tier limitation

Router Affected (if applicable)

  • knnrouter
  • svmrouter
  • mlprouter
  • mfrouter
  • elorouter
  • routerdc
  • automix
  • hybrid_llm
  • graphrouter
  • causallm_router
  • router_r1
  • gmtrouter
  • knnmultiroundrouter
  • llmmultiroundrouter
  • Other / Core framework

Checklist

  • I have tested my changes locally
  • I have reviewed my own code
  • I have added/updated documentation where necessary
  • I have added tests for my changes (if applicable)
  • My changes do not introduce new warnings or errors
  • Any dependent changes have been merged and published

Test Plan

llmrouter list-routers
llmrouter infer --router taskawarerouter --config custom_routers/taskawarerouter/config.yaml --query "Build a production grade authentication system" --route-only
llmrouter infer --router taskawarerouter --config custom_routers/taskawarerouter/config.yaml --query "Translate this paragraph from English to Japanese" --route-only
python -m pytest custom_routers/taskawarerouter/test_router.py -v

Screenshots (if applicable)

Additional Notes

All task_map entries currently point to llama-3.1-8b-instruct
because only this model was available under free NVIDIA NIM tier
during validation. The routing logic correctly selects model tiers
and supports differentiated routing once full API access is configured.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants