llm-router

llm-router

A lightweight task classifier that predicts which category a prompt belongs to (simple, coding, reasoning, security, summarization), for routing prompts to the most suitable LLM in an agentic AI system. It predicts a task category only -- it has no knowledge of any specific LLM provider or model name, so the model registry it routes to can change without ever retraining this classifier.

Built with the llm-router Python package: sentence-transformers/all-MiniLM-L6-v2 embeddings feeding a scikit-learn LogisticRegression classifier.

Labels

coding, reasoning, security, simple, summarization

Evaluation

Full per-class precision/recall/F1 and confusion matrices are in metadata.json -> evaluation. Trained on ai-mitra/llm-router-dataset (~1700 examples, 340 per category).

Usage

pip install llm-router  # or: sentence-transformers scikit-learn joblib huggingface_hub

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)))

Or, using the llm_router package directly:

from llm_router.embeddings import SentenceTransformerEmbedder
from llm_router.classifier import TaskClassifier

embedder = SentenceTransformerEmbedder("sentence-transformers/all-MiniLM-L6-v2")
classifier = TaskClassifier.from_artifact(
    hf_hub_download(repo_id="ai-mitra/llm-router", filename="classifier.joblib"),
    embedder=embedder,
    confidence_threshold=0.45,
)
result = classifier.predict("Explain what a Python decorator does")
print(result.category, result.confidence)

Files

Confidence threshold

Softmax confidence over 5 balanced classes at this dataset size commonly lands in the 0.45-0.95 range for correct predictions, not uniformly near 1.0. 0.45 is a reasonable default cutoff below which a prediction should be treated as "not confident enough" rather than trusted outright -- recalibrate against your own held-out prompts if you fine-tune or replace this artifact.

Scope

This model answers "what kind of task is this prompt," nothing else. It does not call any LLM, does not know about API keys or pricing, and is not a health/availability signal for any model endpoint. See the llm-router package for how this classifier fits into full prompt -> model selection.

License

MIT