python-rerouting-library 0.2.1__py3-none-any.whl

python_rerouting_library/router.py

@@ -0,0 +1,227 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from time import perf_counter
+from typing import Any
+
+import joblib
+
+
+DEFAULT_EMBEDDING_MODEL = (
+    "sentence-transformers/all-MiniLM-L6-v2"
+)
+
+
+@dataclass(frozen=True)
+class RouteDecision:
+    label: str
+    confidence: float | None
+    complex_probability: float | None
+    latency_ms: float
+
+
+class Router:
+    def __init__(
+        self,
+        classifier_path: str | Path,
+        *,
+        embedding_model_name: str | None = None,
+        simple_threshold: float = 0.40,
+        complex_threshold: float = 0.60,
+        device: str = "cpu",
+    ) -> None:
+        self.classifier_path = Path(
+            classifier_path
+        )
+
+        if not self.classifier_path.is_file():
+            raise FileNotFoundError(
+                "Router classifier not found: "
+                f"{self.classifier_path}. "
+                "python-rerouting-library does not ship "
+                "with a default complexity classifier. "
+                "Train one with "
+                "'python -m python_rerouting_library.training' "
+                "or provide a compatible classifier_path."
+            )
+
+        if not (
+            0.0
+            < simple_threshold
+            < complex_threshold
+            < 1.0
+        ):
+            raise ValueError(
+                "Thresholds must satisfy: "
+                "0 < simple_threshold "
+                "< complex_threshold < 1"
+            )
+
+        self.simple_threshold = float(
+            simple_threshold
+        )
+
+        self.complex_threshold = float(
+            complex_threshold
+        )
+
+        self.device = device
+
+        artifact = joblib.load(
+            self.classifier_path
+        )
+
+        if isinstance(artifact, dict):
+            if "classifier" not in artifact:
+                raise ValueError(
+                    "Router artifact does not "
+                    "contain a 'classifier'."
+                )
+
+            self.classifier = artifact[
+                "classifier"
+            ]
+
+            artifact_model_name = artifact.get(
+                "embedding_model_name",
+                DEFAULT_EMBEDDING_MODEL,
+            )
+
+        else:
+            self.classifier = artifact
+
+            artifact_model_name = (
+                DEFAULT_EMBEDDING_MODEL
+            )
+
+        self.embedding_model_name = (
+            embedding_model_name
+            or artifact_model_name
+            or DEFAULT_EMBEDDING_MODEL
+        )
+
+        self._embedding_model: Any | None = None
+
+    def _get_embedding_model(
+        self,
+    ) -> Any:
+        if self._embedding_model is None:
+            try:
+                from sentence_transformers import (
+                    SentenceTransformer,
+                )
+            except ImportError as exc:
+                raise RuntimeError(
+                    "Sentence Transformers could not be "
+                    "loaded. Semantic routing requires the "
+                    "'sentence-transformers' dependency and "
+                    "a compatible Python environment."
+                ) from exc
+
+            self._embedding_model = (
+                SentenceTransformer(
+                    self.embedding_model_name,
+                    device=self.device,
+                )
+            )
+
+        return self._embedding_model
+
+    def warmup(self) -> None:
+        model = self._get_embedding_model()
+
+        model.encode(
+            ["router warmup"],
+            normalize_embeddings=True,
+            show_progress_bar=False,
+        )
+
+    def _get_complex_probability(
+        self,
+        embedding,
+    ) -> float:
+        probabilities = (
+            self.classifier.predict_proba(
+                embedding
+            )[0]
+        )
+
+        classes = list(
+            self.classifier.classes_
+        )
+
+        if "complex" not in classes:
+            raise ValueError(
+                "Classifier does not contain "
+                "the 'complex' class."
+            )
+
+        complex_index = classes.index(
+            "complex"
+        )
+
+        return float(
+            probabilities[complex_index]
+        )
+
+    def route(
+        self,
+        query: str,
+    ) -> RouteDecision:
+        if not isinstance(query, str):
+            raise TypeError(
+                "Query must be a string."
+            )
+
+        query = query.strip()
+
+        if not query:
+            raise ValueError(
+                "Query must not be empty."
+            )
+
+        start = perf_counter()
+
+        model = self._get_embedding_model()
+
+        embedding = model.encode(
+            [query],
+            normalize_embeddings=True,
+            show_progress_bar=False,
+        )
+
+        probability = (
+            self._get_complex_probability(
+                embedding
+            )
+        )
+
+        if probability < self.simple_threshold:
+            label = "simple"
+            confidence = (
+                1.0 - probability
+            )
+
+        elif probability > self.complex_threshold:
+            label = "complex"
+            confidence = probability
+
+        else:
+            label = "uncertain"
+
+            confidence = (
+                abs(probability - 0.5)
+                * 2.0
+            )
+
+        latency_ms = (
+            perf_counter() - start
+        ) * 1000.0
+
+        return RouteDecision(
+            label=label,
+            confidence=confidence,
+            complex_probability=probability,
+            latency_ms=latency_ms,
+        )

python_rerouting_library/training.py

@@ -0,0 +1,177 @@
+ from __future__ import annotations
+
+import argparse
+import csv
+from pathlib import Path
+
+import joblib
+from sentence_transformers import SentenceTransformer
+from sklearn.linear_model import LogisticRegression
+
+from .router import DEFAULT_EMBEDDING_MODEL
+
+
+def load_training_data(
+    csv_path: Path,
+) -> tuple[list[str], list[str]]:
+    queries: list[str] = []
+    labels: list[str] = []
+
+    with csv_path.open(
+        "r",
+        encoding="utf-8",
+        newline="",
+    ) as f:
+        reader = csv.DictReader(f)
+
+        required = {"label", "query"}
+
+        if not required.issubset(
+            set(reader.fieldnames or [])
+        ):
+            raise ValueError(
+                "CSV must contain 'label' and 'query' columns."
+            )
+
+        for row in reader:
+            label = row["label"].strip().lower()
+            query = row["query"].strip()
+
+            if label not in {"simple", "complex"}:
+                raise ValueError(
+                    f"Invalid label: {label!r}"
+                )
+
+            if not query:
+                raise ValueError(
+                    "Encountered an empty query."
+                )
+
+            labels.append(label)
+            queries.append(query)
+
+    if not queries:
+        raise ValueError(
+            "Training CSV is empty."
+        )
+
+    return queries, labels
+
+
+def train_router(
+    csv_path: str | Path,
+    output_path: str | Path,
+    *,
+    embedding_model_name: str = DEFAULT_EMBEDDING_MODEL,
+) -> Path:
+    csv_path = Path(csv_path)
+    output_path = Path(output_path)
+
+    queries, labels = load_training_data(
+        csv_path
+    )
+
+    print(
+        f"Loaded {len(queries)} training examples: "
+        f"{labels.count('simple')} simple / "
+        f"{labels.count('complex')} complex"
+    )
+
+    print(
+        f"Loading embedding model: "
+        f"{embedding_model_name}"
+    )
+
+    encoder = SentenceTransformer(
+        embedding_model_name,
+        device="cpu",
+    )
+
+    print("Generating embeddings...")
+
+    embeddings = encoder.encode(
+        queries,
+        convert_to_numpy=True,
+        normalize_embeddings=True,
+        show_progress_bar=True,
+    )
+
+    print(
+        "Training Logistic Regression classifier..."
+    )
+
+    classifier = LogisticRegression(
+        max_iter=1000,
+        class_weight="balanced",
+        solver="liblinear",
+        random_state=42,
+    )
+
+    classifier.fit(
+        embeddings,
+        labels,
+    )
+
+    artifact = {
+        "classifier": classifier,
+        "embedding_model_name": embedding_model_name,
+        "training_examples": len(queries),
+        "classes": list(classifier.classes_),
+    }
+
+    output_path.parent.mkdir(
+        parents=True,
+        exist_ok=True,
+    )
+
+    joblib.dump(
+        artifact,
+        output_path,
+    )
+
+    print(
+        f"Saved router classifier to: "
+        f"{output_path.resolve()}"
+    )
+
+    return output_path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Train the MiniLM + Logistic Regression "
+            "simple/complex router."
+        )
+    )
+
+    parser.add_argument(
+        "--csv",
+        type=Path,
+        required=True,
+    )
+
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default=Path(
+            "artifacts/router_classifier.joblib"
+        ),
+    )
+
+    parser.add_argument(
+        "--model",
+        default=DEFAULT_EMBEDDING_MODEL,
+    )
+
+    args = parser.parse_args()
+
+    train_router(
+        csv_path=args.csv,
+        output_path=args.output,
+        embedding_model_name=args.model,
+    )
+
+
+if __name__ == "__main__":
+    main()