๐Ÿงญ open source ยท MIT licensed

llm-router

Given a prompt and a list of available models, which model is the best fit? A lightweight classifier that routes prompts to the right LLM by task category โ€” built for agentic AI systems.

How it works

Prompt in, task category + confidence out โ€” then the best available model for that category.

Prompt
โ†’
Sentence Transformer
โ†’
Embedding
โ†’
Task Classifier
โ†’
Category + Confidence
โ†’
Model Registry Filter
โ†’
Ranked Candidates

Task categories

The classifier is trained on task categories only โ€” never specific model names โ€” so the model registry can change without retraining it.

simple
Quick facts, arithmetic, definitions
coding
Write, debug, explain, refactor code
reasoning
Trade-offs, multi-step logic, analysis
security
Vulnerability review, threat modeling
summarization
Condense documents, extract key points

Evaluation

Trained on ~1,700 examples (340/category) โ€” see the dataset card for how it was built.

99.6%
Validation accuracy
99.6%
Validation macro F1
98.8%
Test accuracy
98.8%
Test macro F1

Usage

Load the artifact directly from the Hub, or use the llm_router Python package.

from huggingface_hub import hf_hub_download
from sentence_transformers import SentenceTransformer
import joblib

path = hf_hub_download(repo_id="ai-mitra/llm-router", filename="classifier.joblib")
payload = joblib.load(path)
classifier, classes = payload["classifier"], payload["classes"]

embedder = SentenceTransformer(payload["embedding_model"])
embedding = embedder.encode(["Explain what a Python decorator does"])
probabilities = classifier.predict_proba(embedding)[0]

print(dict(zip(classes, probabilities)))
# {'coding': 0.87, 'reasoning': 0.06, 'security': 0.02, 'simple': 0.03, 'summarization': 0.02}

Resources