AI Prompt Classification and Routing Project

"""Route prompts to specialised handlers with a review threshold."""

from __future__ import annotations

from dataclasses import dataclass
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline


PROMPTS = {
    "summarize": ["summarize this article", "make a short summary", "give me the key points", "condense this report", "shorten the following text"],
    "translate": ["translate this to Hindi", "convert into French", "provide a Spanish translation", "translate the paragraph", "change this into English"],
    "code_help": ["debug this Python function", "explain this SQL error", "write a unit test", "fix my code", "review this program"],
    "data_analysis": ["calculate monthly revenue", "analyse customer churn", "show the sales trend", "compare conversion rates", "build a KPI table"],
}


@dataclass(frozen=True)
class Decision:
    route: str | None
    confidence: float
    needs_review: bool


def train_prompt_router() -> Pipeline:
    texts, labels = [], []
    for label, examples in PROMPTS.items():
        texts.extend(examples)
        labels.extend([label] * len(examples))
    return Pipeline([("tfidf", TfidfVectorizer(ngram_range=(1, 2))),
                     ("classifier", LogisticRegression(max_iter=1000, random_state=42))]).fit(texts, labels)


def route_prompt(model: Pipeline, prompt: str, threshold: float = 0.42) -> Decision:
    probabilities = model.predict_proba([prompt])[0]
    index = int(probabilities.argmax())
    confidence = float(probabilities[index])
    return Decision(str(model.classes_[index]) if confidence >= threshold else None, confidence, confidence < threshold)


def main() -> None:
    model = train_prompt_router()
    for prompt in ("Please summarize this report", "Why is my Python test failing?", "Book a table for dinner"):
        print(prompt, "->", route_prompt(model, prompt))


if __name__ == "__main__":
    main()