query-oracle 0.3.0__py3-none-any.whl

llm_router/__init__.py

@@ -0,0 +1,14 @@
+from .models import (
+    ClassificationResult, EffortLevel, ProviderType,
+    QueryTier, RouterConfig, RouterResponse,
+    MODEL_PRICING, estimate_cost,
+)
+from .router import QueryRouter
+from .classifier import PromptClassifier, DistilBERTClassifier
+
+__all__ = [
+    "QueryRouter", "RouterConfig", "RouterResponse",
+    "ClassificationResult", "QueryTier", "EffortLevel", "ProviderType",
+    "PromptClassifier", "DistilBERTClassifier",
+    "MODEL_PRICING", "estimate_cost",
+]

llm_router/classifier.py

@@ -0,0 +1,301 @@
+"""
+Query classifier.
+
+Two backends, same interface:
+  PromptClassifier   — calls a small LLM (default, works immediately)
+  DistilBERTClassifier — loads a fine-tuned local model (zero API cost,
+                         ~10ms inference; requires a trained model directory)
+
+Swap backends by passing a different classifier instance to QueryRouter.
+"""
+from __future__ import annotations
+import json
+import time
+from pathlib import Path
+from abc import ABC, abstractmethod
+
+from .models import ClassificationResult, EffortLevel, QueryTier, RouterConfig, estimate_cost
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Abstract base
+# ─────────────────────────────────────────────────────────────────────────────
+
+class BaseClassifier(ABC):
+    @abstractmethod
+    def classify(self, query: str) -> ClassificationResult: ...
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Option A: Prompt-based (works out of the box, any provider)
+# ─────────────────────────────────────────────────────────────────────────────
+
+CLASSIFICATION_PROMPT = """\
+You are a query classification engine. Analyse the user query and return JSON:
+
+{{
+  "tier":           "fast" | "balanced" | "deep",
+  "effort":         "low"  | "medium"   | "high",
+  "facts_ratio":    <float 0.0-1.0>,
+  "judgment_ratio": <float 0.0-1.0>,
+  "confidence":     <float 0.0-1.0>,
+  "reasoning":      "<one sentence>"
+}}
+
+Rules:
+- facts_ratio + judgment_ratio = 1.0
+- fast     -> factual lookups, news, nutrition, simple maths
+- balanced -> moderate analysis, explanations with nuance
+- deep     -> ideation, strategy, ethics, complex design
+- effort follows tier by default but may be raised one level for unusually nuanced queries.
+
+Return ONLY the JSON. No text outside it.
+
+Query: {query}"""
+
+
+class PromptClassifier(BaseClassifier):
+    """
+    Classifies queries via a prompt call to a small LLM.
+    Works immediately with any API key. Logs every result for future fine-tuning.
+    """
+
+    def __init__(self, config: RouterConfig, client=None):
+        self.config = config
+        self._client = client
+        if config.log_classifications:
+            Path(config.log_path).parent.mkdir(parents=True, exist_ok=True)
+
+    def classify(self, query: str) -> ClassificationResult:
+        raw, in_tok, out_tok = self._call_model(query)
+        result = self._parse(raw)
+        result.classifier_input_tokens  = in_tok
+        result.classifier_output_tokens = out_tok
+        result.classifier_cost_usd      = estimate_cost(
+            self.config.classifier_model, in_tok, out_tok
+        )
+        if self.config.log_classifications:
+            self._log(query, result)
+        return result
+
+    def _call_model(self, query: str) -> tuple[str, int, int]:
+        """Returns (text, input_tokens, output_tokens)."""
+        from .models import ProviderType
+        provider = self.config.classifier_provider
+        prompt   = CLASSIFICATION_PROMPT.format(query=query)
+
+        if provider == ProviderType.ANTHROPIC:
+            import anthropic
+            client = self._client or anthropic.Anthropic()
+            r = client.messages.create(
+                model=self.config.classifier_model,
+                max_tokens=512,
+                messages=[{"role": "user", "content": prompt}],
+            )
+            return (
+                r.content[0].text.strip(),
+                r.usage.input_tokens,
+                r.usage.output_tokens,
+            )
+
+        elif provider == ProviderType.OPENAI:
+            from openai import OpenAI
+            client = self._client or OpenAI()
+            r = client.chat.completions.create(
+                model=self.config.classifier_model,
+                max_tokens=512,
+                messages=[{"role": "user", "content": prompt}],
+            )
+            return (
+                r.choices[0].message.content.strip(),
+                r.usage.prompt_tokens,
+                r.usage.completion_tokens,
+            )
+
+        elif provider == ProviderType.GEMINI:
+            # BUG 9 fix: honour self._client when injected; otherwise configure from env.
+            if self._client:
+                model_instance = self._client
+            else:
+                import os
+                import google.generativeai as genai
+                api_key = os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
+                if api_key:
+                    genai.configure(api_key=api_key)
+                model_instance = genai.GenerativeModel(self.config.classifier_model)
+            r = model_instance.generate_content(prompt)
+            in_tok  = (
+                getattr(r.usage_metadata, "prompt_token_count", 0)
+                if hasattr(r, "usage_metadata") else 0
+            )
+            out_tok = (
+                getattr(r.usage_metadata, "candidates_token_count", 0)
+                if hasattr(r, "usage_metadata") else 0
+            )
+            return r.text.strip(), in_tok, out_tok
+
+        raise ValueError(f"Unknown classifier provider: {provider}")
+
+    # ── BUG 7 fix: coerce helpers tolerate uppercase / unexpected values ──────
+
+    @staticmethod
+    def _coerce_tier(raw: str) -> QueryTier:
+        try:
+            return QueryTier(raw.strip().lower())
+        except ValueError:
+            return QueryTier.BALANCED
+
+    @staticmethod
+    def _coerce_effort(raw: str) -> EffortLevel:
+        try:
+            return EffortLevel(raw.strip().lower())
+        except ValueError:
+            return EffortLevel.MEDIUM
+
+    def _parse(self, raw: str) -> ClassificationResult:
+        if raw.startswith("```"):
+            lines = raw.splitlines()
+            raw = "\n".join(lines[1:-1] if lines[-1] == "```" else lines[1:])
+        try:
+            d = json.loads(raw)
+        except json.JSONDecodeError:
+            return ClassificationResult(
+                tier=QueryTier.BALANCED, effort=EffortLevel.MEDIUM,
+                facts_ratio=0.5, judgment_ratio=0.5,
+                reasoning=f"Parse error — defaulted. Raw: {raw[:80]}",
+                confidence=0.0,
+            )
+        return ClassificationResult(
+            tier=self._coerce_tier(d.get("tier", "balanced")),
+            effort=self._coerce_effort(d.get("effort", "medium")),
+            facts_ratio=float(d.get("facts_ratio", 0.5)),
+            judgment_ratio=float(d.get("judgment_ratio", 0.5)),
+            reasoning=d.get("reasoning", ""),
+            confidence=float(d.get("confidence", 0.8)),
+        )
+
+    def _log(self, query: str, result: ClassificationResult) -> None:
+        entry = {
+            "timestamp":      time.time(),
+            "query":          query,
+            "tier":           result.tier.value,
+            "effort":         result.effort.value,
+            "facts_ratio":    result.facts_ratio,
+            "judgment_ratio": result.judgment_ratio,
+            "confidence":     result.confidence,
+            "reasoning":      result.reasoning,
+        }
+        with open(self.config.log_path, "a", encoding="utf-8") as f:
+            f.write(json.dumps(entry) + "\n")
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Option B: Fine-tuned DistilBERT (zero API cost, ~10ms, runs locally)
+# ─────────────────────────────────────────────────────────────────────────────
+
+LABEL2TIER = {0: QueryTier.FAST, 1: QueryTier.BALANCED, 2: QueryTier.DEEP}
+TIER2EFFORT = {
+    QueryTier.FAST:     EffortLevel.LOW,
+    QueryTier.BALANCED: EffortLevel.MEDIUM,
+    QueryTier.DEEP:     EffortLevel.HIGH,
+}
+
+
+class DistilBERTClassifier(BaseClassifier):
+    """
+    Zero-cost local classifier using a fine-tuned DistilBERT model.
+
+    Usage:
+        clf = DistilBERTClassifier("./training/query-classifier-final")
+        router = QueryRouter(classifier=clf)
+
+    Train your model first using the scripts in /training/.
+    Requires: pip install transformers torch  (or onnxruntime for ONNX)
+    """
+
+    def __init__(self, model_dir: str, use_onnx: bool = False):
+        self.model_dir = model_dir
+        self.use_onnx  = use_onnx
+        self._model    = None
+        self._tokenizer = None
+        self._ort_session = None
+        self._load()
+
+    def _load(self):
+        if self.use_onnx:
+            self._load_onnx()
+        else:
+            self._load_pytorch()
+
+    def _load_pytorch(self):
+        try:
+            from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification
+            import torch
+            self._tokenizer = DistilBertTokenizerFast.from_pretrained(self.model_dir)
+            self._model = DistilBertForSequenceClassification.from_pretrained(self.model_dir)
+            self._model.eval()
+            self._torch = torch
+        except ImportError:
+            raise ImportError(
+                "Install pytorch + transformers: pip install transformers torch"
+            )
+
+    def _load_onnx(self):
+        try:
+            import onnxruntime as ort
+            from transformers import DistilBertTokenizerFast
+            self._tokenizer = DistilBertTokenizerFast.from_pretrained(self.model_dir)
+            onnx_path = f"{self.model_dir}/model.onnx"
+            self._ort_session = ort.InferenceSession(onnx_path)
+        except ImportError:
+            raise ImportError(
+                "Install onnxruntime: pip install onnxruntime"
+            )
+
+    def classify(self, query: str) -> ClassificationResult:
+        if self.use_onnx:
+            return self._classify_onnx(query)
+        return self._classify_pytorch(query)
+
+    def _classify_pytorch(self, query: str) -> ClassificationResult:
+        import torch
+        inputs = self._tokenizer(
+            query, return_tensors="pt", truncation=True,
+            padding="max_length", max_length=128,
+        )
+        with torch.no_grad():
+            logits = self._model(**inputs).logits
+        probs     = torch.softmax(logits, dim=-1)[0]
+        label_id  = int(probs.argmax())
+        confidence = float(probs[label_id])
+        tier = LABEL2TIER[label_id]
+        return ClassificationResult(
+            tier=tier, effort=TIER2EFFORT[tier],
+            facts_ratio=float(probs[0]),
+            judgment_ratio=float(probs[2]),
+            reasoning="DistilBERT local classifier",
+            confidence=confidence,
+        )
+
+    def _classify_onnx(self, query: str) -> ClassificationResult:
+        import numpy as np
+        inputs = self._tokenizer(
+            query, return_tensors="np", truncation=True,
+            padding="max_length", max_length=128,
+        )
+        logits = self._ort_session.run(
+            None,
+            {"input_ids": inputs["input_ids"], "attention_mask": inputs["attention_mask"]},
+        )[0]
+        # BUG 8 fix: numerically stable softmax (subtract max before exp, keep axes)
+        logits_shifted = logits - logits.max(axis=-1, keepdims=True)
+        probs = np.exp(logits_shifted) / np.exp(logits_shifted).sum(axis=-1, keepdims=True)
+        label_id   = int(probs.argmax())
+        confidence = float(probs[0][label_id])
+        tier = LABEL2TIER[label_id]
+        return ClassificationResult(
+            tier=tier, effort=TIER2EFFORT[tier],
+            facts_ratio=float(probs[0][0]),
+            judgment_ratio=float(probs[0][2]),
+            reasoning="DistilBERT ONNX local classifier",
+            confidence=confidence,
+        )

llm_router/mcp_server.py

@@ -0,0 +1,103 @@
+"""
+MCP server — package-level module.
+
+Installed entry point: query-oracle-mcp
+  (runs via: python -m llm_router.mcp_server  OR  query-oracle-mcp)
+
+Claude Code / Claude Desktop config (~/.claude/claude_desktop_config.json):
+
+  {
+    "mcpServers": {
+      "query-oracle": {
+        "command": "query-oracle-mcp",
+        "env": {
+          "ANTHROPIC_API_KEY": "sk-ant-...",
+          "OPENAI_API_KEY":    "sk-...",
+          "GEMINI_API_KEY":    "AIza..."
+        }
+      }
+    }
+  }
+
+Or via uvx (no install needed):
+  "command": "uvx",
+  "args": ["--from", "query-oracle[mcp]", "query-oracle-mcp"]
+"""
+from __future__ import annotations
+
+try:
+    from mcp.server.fastmcp import FastMCP
+except ImportError:
+    try:
+        from mcp.server import FastMCP  # type: ignore[no-redef]
+    except ImportError:
+        raise ImportError("Install the MCP SDK: pip install 'query-oracle[mcp]'")
+
+from .models import RouterConfig
+from .router import QueryRouter
+from .providers import AnthropicProvider, OpenAIProvider, GeminiProvider
+
+mcp = FastMCP("query-oracle")
+
+
+def _provider(name: str, config: RouterConfig):
+    n = name.lower()
+    if n == "openai":
+        return OpenAIProvider(config)
+    if n == "gemini":
+        return GeminiProvider(config)
+    return AnthropicProvider(config)
+
+
+@mcp.tool()
+def route(query: str, provider: str = "anthropic") -> str:
+    """
+    Route a query to the optimal LLM model automatically.
+
+    Classifies the query on a facts-to-judgment spectrum and dispatches it to:
+    - FAST  (Haiku / GPT-4o-mini / Gemini Flash)   for factual lookups
+    - BALANCED (Sonnet / GPT-4o / Gemini Pro)       for analysis
+    - DEEP  (Opus+thinking / o1-high / Gemini Thinking) for strategy & design
+
+    Args:
+        query:    The user question or task.
+        provider: LLM provider — anthropic | openai | gemini  (default: anthropic)
+    """
+    config = RouterConfig()
+    router = QueryRouter(config=config, provider=_provider(provider, config))
+    r = router.route(query)
+    meta = (
+        f"[query-oracle] tier={r.tier.value} | model={r.model_used} | "
+        f"thinking={r.extended_thinking_used} | {r.latency_ms:.0f}ms | "
+        f"cost=${r.cost_usd:.6f}\n\n"
+    )
+    return meta + r.content
+
+
+@mcp.tool()
+def classify(query: str) -> str:
+    """
+    Classify a query without executing it — returns the routing decision only.
+
+    Args:
+        query: The query to classify.
+    """
+    config = RouterConfig()
+    router = QueryRouter(config=config)
+    c = router.classifier.classify(query)
+    return (
+        f"Tier:           {c.tier.value}\n"
+        f"Effort:         {c.effort.value}\n"
+        f"Facts ratio:    {c.facts_ratio:.0%}\n"
+        f"Judgment ratio: {c.judgment_ratio:.0%}\n"
+        f"Confidence:     {c.confidence:.0%}\n"
+        f"Reasoning:      {c.reasoning}"
+    )
+
+
+def main() -> None:
+    mcp.run()
+
+
+if __name__ == "__main__":
+    main()

llm_router/models.py

@@ -0,0 +1,159 @@
+"""
+Core data models for query-oracle.
+"""
+from __future__ import annotations
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Optional
+
+
+class QueryTier(str, Enum):
+    FAST     = "fast"      # Factual, simple, low latency
+    BALANCED = "balanced"  # Moderate reasoning
+    DEEP     = "deep"      # Heavy ideation, judgment, creativity
+
+
+class EffortLevel(str, Enum):
+    LOW    = "low"
+    MEDIUM = "medium"
+    HIGH   = "high"
+
+
+class ProviderType(str, Enum):
+    ANTHROPIC = "anthropic"
+    OPENAI    = "openai"
+    GEMINI    = "gemini"
+
+
+# Pricing per 1 M tokens (input, output) in USD — update as providers change rates.
+MODEL_PRICING: dict[str, tuple[float, float]] = {
+    # Anthropic
+    "claude-haiku-4-5":   (0.80,   4.00),
+    "claude-sonnet-4-5":  (3.00,  15.00),
+    "claude-opus-4-5":    (15.00, 75.00),
+    # OpenAI
+    "gpt-4o-mini":        (0.15,   0.60),
+    "gpt-4o":             (2.50,  10.00),
+    "o1":                 (15.00, 60.00),
+    # Gemini
+    "gemini-1.5-flash":              (0.075,  0.30),
+    "gemini-1.5-pro":                (1.25,   5.00),
+    "gemini-2.0-flash-thinking-exp": (0.00,   0.00),
+}
+
+
+def estimate_cost(model: str, input_tokens: int, output_tokens: int) -> float:
+    """Return estimated cost in USD for a single call."""
+    if model not in MODEL_PRICING:
+        return 0.0
+    in_rate, out_rate = MODEL_PRICING[model]
+    return (input_tokens * in_rate + output_tokens * out_rate) / 1_000_000
+
+
+@dataclass
+class ClassificationResult:
+    tier:           QueryTier
+    effort:         EffortLevel
+    facts_ratio:    float   # 0.0 = pure judgment → 1.0 = pure facts
+    judgment_ratio: float
+    reasoning:      str
+    confidence:     float
+    # Classifier's own token usage (0 for local DistilBERT)
+    classifier_input_tokens:  int   = 0
+    classifier_output_tokens: int   = 0
+    classifier_cost_usd:      float = 0.0
+
+
+@dataclass
+class RouterConfig:
+    """
+    All routing behaviour is driven from here.
+    Swap model IDs, enable thinking, change token budgets —
+    without touching provider or classifier code.
+    """
+
+    # ── Anthropic model IDs ──────────────────────────────────────────────
+    anthropic_models: dict[QueryTier, str] = field(default_factory=lambda: {
+        QueryTier.FAST:     "claude-haiku-4-5",
+        QueryTier.BALANCED: "claude-sonnet-4-5",
+        QueryTier.DEEP:     "claude-opus-4-5",
+    })
+
+    # ── OpenAI model IDs ─────────────────────────────────────────────────
+    openai_models: dict[QueryTier, str] = field(default_factory=lambda: {
+        QueryTier.FAST:     "gpt-4o-mini",
+        QueryTier.BALANCED: "gpt-4o",
+        QueryTier.DEEP:     "o1",
+    })
+
+    # ── Gemini model IDs ─────────────────────────────────────────────────
+    gemini_models: dict[QueryTier, str] = field(default_factory=lambda: {
+        QueryTier.FAST:     "gemini-1.5-flash",
+        QueryTier.BALANCED: "gemini-1.5-pro",
+        QueryTier.DEEP:     "gemini-2.0-flash-thinking-exp",
+    })
+
+    # ── Extended thinking (Anthropic only) ───────────────────────────────
+    extended_thinking_map: dict[QueryTier, bool] = field(default_factory=lambda: {
+        QueryTier.FAST:     False,
+        QueryTier.BALANCED: False,
+        QueryTier.DEEP:     True,
+    })
+    thinking_budget_map: dict[QueryTier, int] = field(default_factory=lambda: {
+        QueryTier.FAST:     0,
+        QueryTier.BALANCED: 0,
+        QueryTier.DEEP:     8000,
+    })
+
+    # ── Reasoning effort (OpenAI o-series) ───────────────────────────────
+    openai_reasoning_effort: dict[QueryTier, Optional[str]] = field(default_factory=lambda: {
+        QueryTier.FAST:     None,
+        QueryTier.BALANCED: None,
+        QueryTier.DEEP:     "high",
+    })
+
+    # ── Max output tokens ────────────────────────────────────────────────
+    # DEEP must be > thinking_budget_map[DEEP] + answer headroom.
+    # Default: 8000 thinking + 8000 answer = 16000.
+    max_tokens_map: dict[QueryTier, int] = field(default_factory=lambda: {
+        QueryTier.FAST:     1024,
+        QueryTier.BALANCED: 4096,
+        QueryTier.DEEP:     16000,
+    })
+
+    # ── Classifier settings ──────────────────────────────────────────────
+    classifier_provider: ProviderType = ProviderType.ANTHROPIC
+    classifier_model:    str          = "claude-haiku-4-5"
+
+    # ── Logging (builds your fine-tuning dataset passively) ─────────────
+    log_classifications: bool = True
+    log_path:            str  = "logs/classifications.jsonl"
+
+    def __post_init__(self) -> None:
+        for tier in QueryTier:
+            budget = self.thinking_budget_map.get(tier, 0)
+            max_t  = self.max_tokens_map.get(tier, 0)
+            if budget > 0 and max_t <= budget:
+                raise ValueError(
+                    f"RouterConfig: max_tokens_map[{tier.value}]={max_t} must be "
+                    f"greater than thinking_budget_map[{tier.value}]={budget}. "
+                    f"Thinking tokens count against max_tokens, leaving only "
+                    f"{max_t - budget} tokens for the actual answer. "
+                    f"Set max_tokens_map[{tier.value}] to at least {budget + 1024}."
+                )
+
+
+@dataclass
+class RouterResponse:
+    content:                str
+    tier:                   QueryTier
+    effort:                 EffortLevel
+    provider:               ProviderType
+    model_used:             str
+    extended_thinking_used: bool
+    classification:         ClassificationResult
+    input_tokens:           int   = 0
+    output_tokens:          int   = 0
+    latency_ms:             float = 0.0
+    cost_usd:               float = 0.0   # completion-only cost
+    total_cost_usd:         float = 0.0   # completion + classifier cost

llm_router/providers/__init__.py

@@ -0,0 +1,9 @@
+from .base import BaseProvider, CompletionResult
+from .anthropic_provider import AnthropicProvider
+from .openai_provider import OpenAIProvider
+from .gemini_provider import GeminiProvider
+
+__all__ = [
+    "BaseProvider", "CompletionResult",
+    "AnthropicProvider", "OpenAIProvider", "GeminiProvider",
+]

llm_router/providers/anthropic_provider.py

@@ -0,0 +1,50 @@
+"""Anthropic provider — Haiku / Sonnet / Opus with optional extended thinking."""
+from __future__ import annotations
+import anthropic
+from ..models import EffortLevel, ProviderType, QueryTier, RouterConfig
+from .base import BaseProvider, CompletionResult
+
+
+class AnthropicProvider(BaseProvider):
+    def __init__(self, config: RouterConfig, api_key: str | None = None):
+        self.config = config
+        self.client = anthropic.Anthropic(api_key=api_key)  # falls back to env var
+
+    @property
+    def provider_type(self) -> ProviderType:
+        return ProviderType.ANTHROPIC
+
+    def complete(self, query: str, tier: QueryTier, effort: EffortLevel) -> CompletionResult:
+        model      = self.config.anthropic_models[tier]
+        thinking   = self.config.extended_thinking_map[tier]
+        budget     = self.config.thinking_budget_map[tier]
+        max_tokens = self.config.max_tokens_map[tier]
+
+        if thinking and budget > 0:
+            return self._with_thinking(query, model, budget, max_tokens)
+        return self._standard(query, model, max_tokens)
+
+    def _standard(self, query: str, model: str, max_tokens: int) -> CompletionResult:
+        r = self.client.messages.create(
+            model=model, max_tokens=max_tokens,
+            messages=[{"role": "user", "content": query}],
+        )
+        text = "".join(b.text for b in r.content if hasattr(b, "text"))
+        return CompletionResult(
+            content=text, model_used=model, provider=ProviderType.ANTHROPIC,
+            extended_thinking_used=False,
+            input_tokens=r.usage.input_tokens, output_tokens=r.usage.output_tokens,
+        )
+
+    def _with_thinking(self, query: str, model: str, budget: int, max_tokens: int) -> CompletionResult:
+        r = self.client.messages.create(
+            model=model, max_tokens=max_tokens,
+            thinking={"type": "enabled", "budget_tokens": budget},
+            messages=[{"role": "user", "content": query}],
+        )
+        text = "".join(b.text for b in r.content if b.type == "text")
+        return CompletionResult(
+            content=text, model_used=model, provider=ProviderType.ANTHROPIC,
+            extended_thinking_used=True,
+            input_tokens=r.usage.input_tokens, output_tokens=r.usage.output_tokens,
+        )

llm_router/providers/base.py

@@ -0,0 +1,24 @@
+"""Abstract provider interface — implement this to add any LLM backend."""
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from ..models import EffortLevel, ProviderType, QueryTier
+
+
+@dataclass
+class CompletionResult:
+    content:                str
+    model_used:             str
+    provider:               ProviderType
+    extended_thinking_used: bool
+    input_tokens:           int
+    output_tokens:          int
+
+
+class BaseProvider(ABC):
+    @abstractmethod
+    def complete(self, query: str, tier: QueryTier, effort: EffortLevel) -> CompletionResult: ...
+
+    @property
+    @abstractmethod
+    def provider_type(self) -> ProviderType: ...

llm_router/providers/gemini_provider.py

@@ -0,0 +1,53 @@
+"""
+Google Gemini provider — gemini-flash / gemini-pro / gemini-thinking.
+
+Requires: pip install google-generativeai
+"""
+from __future__ import annotations
+from ..models import EffortLevel, ProviderType, QueryTier, RouterConfig
+from .base import BaseProvider, CompletionResult
+
+try:
+    import google.generativeai as genai
+except ImportError:
+    genai = None  # type: ignore
+
+
+class GeminiProvider(BaseProvider):
+    def __init__(self, config: RouterConfig, api_key: str | None = None):
+        if genai is None:
+            raise ImportError(
+                "Install the Gemini SDK: pip install google-generativeai"
+            )
+        import os
+        key = api_key or os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY")
+        if not key:
+            raise ValueError(
+                "Set GEMINI_API_KEY (or GOOGLE_API_KEY) environment variable."
+            )
+        genai.configure(api_key=key)
+        self.config = config
+
+    @property
+    def provider_type(self) -> ProviderType:
+        return ProviderType.GEMINI
+
+    def complete(self, query: str, tier: QueryTier, effort: EffortLevel) -> CompletionResult:
+        model_id   = self.config.gemini_models[tier]
+        max_tokens = self.config.max_tokens_map[tier]
+
+        model = genai.GenerativeModel(model_id)
+        generation_config = genai.types.GenerationConfig(max_output_tokens=max_tokens)
+
+        response = model.generate_content(query, generation_config=generation_config)
+        text = response.text or ""
+
+        usage = getattr(response, "usage_metadata", None)
+        return CompletionResult(
+            content=text,
+            model_used=model_id,
+            provider=ProviderType.GEMINI,
+            extended_thinking_used=("thinking" in model_id),
+            input_tokens=getattr(usage, "prompt_token_count", 0),
+            output_tokens=getattr(usage, "candidates_token_count", 0),
+        )

llm_router/providers/openai_provider.py

@@ -0,0 +1,55 @@
+"""
+OpenAI provider — gpt-4o-mini / gpt-4o / o1.
+
+Requires: pip install openai
+"""
+from __future__ import annotations
+from ..models import EffortLevel, ProviderType, QueryTier, RouterConfig
+from .base import BaseProvider, CompletionResult
+
+try:
+    from openai import OpenAI
+except ImportError:
+    OpenAI = None  # type: ignore
+
+
+class OpenAIProvider(BaseProvider):
+    def __init__(self, config: RouterConfig, api_key: str | None = None):
+        if OpenAI is None:
+            raise ImportError("Install the openai package: pip install openai")
+        self.config = config
+        self.client = OpenAI(api_key=api_key)  # falls back to OPENAI_API_KEY
+
+    @property
+    def provider_type(self) -> ProviderType:
+        return ProviderType.OPENAI
+
+    def complete(self, query: str, tier: QueryTier, effort: EffortLevel) -> CompletionResult:
+        model      = self.config.openai_models[tier]
+        max_tokens = self.config.max_tokens_map[tier]
+        reasoning  = self.config.openai_reasoning_effort.get(tier)
+
+        kwargs: dict = {
+            "model": model,
+            "messages": [{"role": "user", "content": query}],
+        }
+
+        # o-series models use max_completion_tokens and support reasoning_effort
+        if model.startswith("o"):
+            kwargs["max_completion_tokens"] = max_tokens
+            if reasoning:
+                kwargs["reasoning_effort"] = reasoning
+        else:
+            kwargs["max_tokens"] = max_tokens
+
+        r = self.client.chat.completions.create(**kwargs)
+        text = r.choices[0].message.content or ""
+
+        return CompletionResult(
+            content=text,
+            model_used=model,
+            provider=ProviderType.OPENAI,
+            extended_thinking_used=(model.startswith("o") and reasoning == "high"),
+            input_tokens=r.usage.prompt_tokens if r.usage else 0,
+            output_tokens=r.usage.completion_tokens if r.usage else 0,
+        )

llm_router/router.py

@@ -0,0 +1,77 @@
+"""
+QueryRouter — the single entry point.
+
+from llm_router import QueryRouter
+from llm_router.providers import OpenAIProvider
+
+router = QueryRouter(provider=OpenAIProvider(config))
+response = router.route("Design an intent-driven AI OS for non-technical users.")
+print(response.content, response.tier, response.model_used)
+"""
+from __future__ import annotations
+import asyncio
+import time
+
+from .classifier import BaseClassifier, PromptClassifier
+from .models import RouterConfig, RouterResponse, estimate_cost
+from .providers.base import BaseProvider
+
+
+class QueryRouter:
+    """
+    Classify → Route → Execute.
+
+    Parameters
+    ----------
+    config     : RouterConfig, optional  — override model IDs, budgets, etc.
+    provider   : BaseProvider            — which LLM backend to execute on
+                 (AnthropicProvider / OpenAIProvider / GeminiProvider)
+    classifier : BaseClassifier, optional — swap in DistilBERTClassifier once trained
+    """
+
+    def __init__(
+        self,
+        config:     RouterConfig  | None = None,
+        provider:   BaseProvider  | None = None,
+        classifier: BaseClassifier| None = None,
+    ):
+        self.config = config or RouterConfig()
+
+        # Default provider: Anthropic (reads ANTHROPIC_API_KEY from env)
+        if provider is None:
+            from .providers.anthropic_provider import AnthropicProvider
+            provider = AnthropicProvider(self.config)
+        self.provider = provider
+
+        # Default classifier: prompt-based (uses same provider as classifier_provider in config)
+        self.classifier = classifier or PromptClassifier(self.config)
+
+    def route(self, query: str) -> RouterResponse:
+        t0 = time.monotonic()
+
+        classification = self.classifier.classify(query)
+        completion     = self.provider.complete(
+            query, tier=classification.tier, effort=classification.effort
+        )
+
+        cost_usd = estimate_cost(
+            completion.model_used, completion.input_tokens, completion.output_tokens
+        )
+        return RouterResponse(
+            content=completion.content,
+            tier=classification.tier,
+            effort=classification.effort,
+            provider=completion.provider,
+            model_used=completion.model_used,
+            extended_thinking_used=completion.extended_thinking_used,
+            classification=classification,
+            input_tokens=completion.input_tokens,
+            output_tokens=completion.output_tokens,
+            latency_ms=(time.monotonic() - t0) * 1000,
+            cost_usd=cost_usd,
+            total_cost_usd=cost_usd + classification.classifier_cost_usd,
+        )
+
+    async def async_route(self, query: str) -> RouterResponse:
+        """Non-blocking version of route() — runs the synchronous call in a thread pool."""
+        return await asyncio.to_thread(self.route, query)

query_oracle-0.3.0.dist-info/METADATA

@@ -0,0 +1,299 @@
+Metadata-Version: 2.4
+Name: query-oracle
+Version: 0.3.0
+Summary: Automatic LLM model and effort-level routing based on query classification.
+Author-email: Heman Phadvas <tailung.hp@gmail.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/hemanpadvas2002/query-oracle
+Project-URL: Repository, https://github.com/hemanpadvas2002/query-oracle
+Project-URL: Issues, https://github.com/hemanpadvas2002/query-oracle/issues
+Project-URL: Live API, https://query-oracle-production.up.railway.app
+Keywords: llm,routing,anthropic,openai,gemini,claude,mcp,model-selection,cost-optimization
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: anthropic>=0.40.0
+Requires-Dist: python-dotenv>=1.0.0
+Provides-Extra: server
+Requires-Dist: fastapi>=0.110.0; extra == "server"
+Requires-Dist: uvicorn[standard]>=0.27.0; extra == "server"
+Provides-Extra: mcp
+Requires-Dist: mcp<2.0.0,>=1.0.0; extra == "mcp"
+Provides-Extra: openai
+Requires-Dist: openai>=1.30.0; extra == "openai"
+Provides-Extra: gemini
+Requires-Dist: google-generativeai>=0.7.0; extra == "gemini"
+Provides-Extra: train
+Requires-Dist: transformers>=4.40.0; extra == "train"
+Requires-Dist: datasets>=2.19.0; extra == "train"
+Requires-Dist: torch>=2.2.0; extra == "train"
+Requires-Dist: scikit-learn>=1.4.0; extra == "train"
+Requires-Dist: optimum>=1.19.0; extra == "train"
+Requires-Dist: onnxruntime>=1.17.0; extra == "train"
+Provides-Extra: all
+Requires-Dist: query-oracle[gemini,mcp,openai,server]; extra == "all"
+Dynamic: license-file
+
+# query-oracle
+
+> Automatic LLM routing — the right model, the right effort, zero manual selection.
+
+[![CI](https://github.com/hemanpadvas2002/query-oracle/actions/workflows/ci.yml/badge.svg)](https://github.com/hemanpadvas2002/query-oracle/actions/workflows/ci.yml)
+[![Live](https://img.shields.io/badge/API-live%20on%20Railway-brightgreen)](https://query-oracle-production.up.railway.app/health)
+[![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue)](https://www.python.org/)
+
+---
+
+## Live demo
+
+> **Illustrative** — actual output will vary by provider, model, and query.
+
+```python
+from llm_router import QueryRouter
+
+r = QueryRouter().route("Design a fault-tolerant event streaming architecture.")
+print(f"tier={r.tier.value}  model={r.model_used}  thinking={r.extended_thinking_used}")
+print(f"cost=${r.cost_usd:.5f}  total=${r.total_cost_usd:.5f}  latency={r.latency_ms:.0f}ms")
+print(r.classification.reasoning)
+```
+
+---
+
+## What it does
+
+query-oracle sits in front of your LLM calls and automatically decides which model deserves the query. Factual lookups go to Haiku or GPT-4o-mini in under a second; open-ended design problems get routed to Opus with extended thinking or o1 with high reasoning effort. The classification itself costs a fraction of a cent and the routing decision is logged so you can fine-tune a local DistilBERT classifier later — eventually dropping the classification API cost to zero.
+
+It ships as a Python library, a live REST API, a Claude Code MCP plugin, a VS Code / Cursor extension, an OpenAI Custom GPT Action, and a reusable GitHub Actions workflow. Pick whichever integration fits your stack.
+
+---
+
+## Quick install — pick your platform
+
+### Python library
+
+```bash
+pip install "query-oracle"
+```
+
+With optional extras:
+
+```bash
+pip install "query-oracle[server]"     # FastAPI REST server
+pip install "query-oracle[mcp]"        # Claude Code / Desktop MCP plugin
+pip install "query-oracle[openai]"     # OpenAI provider
+pip install "query-oracle[gemini]"     # Gemini provider
+pip install "query-oracle[all]"        # everything
+```
+
+### Claude Code / Claude Desktop (MCP)
+
+```bash
+pip install "query-oracle[mcp]"
+claude mcp add query-oracle -- query-oracle-mcp
+```
+
+Then in any Claude conversation:
+
+```
+route "Design a real-time fraud detection pipeline."
+classify "What is the capital of France?"
+```
+
+Or add to `~/.claude/claude_desktop_config.json` manually:
+
+```json
+{
+  "mcpServers": {
+    "query-oracle": {
+      "command": "query-oracle-mcp",
+      "env": { "ANTHROPIC_API_KEY": "sk-ant-..." }
+    }
+  }
+}
+```
+
+---
+
+### ChatGPT — Custom GPT Action
+
+Paste this URL into **GPT builder → Configure → Actions → Import from URL**:
+
+```
+https://query-oracle-production.up.railway.app/openapi.json
+```
+
+The server is live — no setup required. See [`openai-plugin/README.md`](openai-plugin/README.md) for how to self-host and add authentication.
+
+---
+
+### Cursor / VS Code Extension
+
+The extension starts the REST server automatically — no manual uvicorn command.
+
+```bash
+cd vscode-extension
+npm install
+npm run package                               # → query-oracle-1.0.0.vsix
+code --install-extension query-oracle-1.0.0.vsix
+```
+
+Press `Cmd/Ctrl+Shift+L` to open the query input. Responses appear in the **LLM Query Router** output panel with tier, model, latency, and cost.
+
+---
+
+### GitHub Copilot / Codex
+
+Copy `.github/copilot-instructions.md` into your own repo's `.github/` folder. Copilot will read it automatically in VS Code and JetBrains and stop suggesting hard-coded model names.
+
+For GitHub Actions / Copilot Workspace tasks:
+
+```yaml
+jobs:
+  design:
+    uses: hemanpadvas2002/query-oracle/.github/workflows/copilot-router.yml@main
+    with:
+      query: "Design a zero-downtime database migration strategy."
+    secrets:
+      ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+    # outputs: response, tier, model_used, cost_usd
+```
+
+---
+
+## How routing works
+
+| Tier | When | Anthropic | OpenAI | Gemini |
+|---|---|---|---|---|
+| `fast` | Factual lookups, maths, news, nutrition | claude-haiku-4-5 | gpt-4o-mini | gemini-1.5-flash |
+| `balanced` | Analysis, explanations, moderate reasoning | claude-sonnet-4-5 | gpt-4o | gemini-1.5-pro |
+| `deep` | Strategy, design, ethics, complex ideation | claude-opus-4-5 + thinking | o1 (high) | gemini-2.0-flash-thinking |
+
+The classifier sends the query to a small model (Haiku by default) with a structured prompt that returns `tier`, `effort`, `facts_ratio`, `judgment_ratio`, `confidence`, and a one-sentence `reasoning`. Every result is logged to `logs/classifications.jsonl` — this passively builds the labelled dataset for local DistilBERT fine-tuning.
+
+---
+
+## Response fields
+
+```python
+RouterResponse(
+    content      = "...",
+    tier         = QueryTier.DEEP,
+    effort       = EffortLevel.HIGH,
+    provider     = ProviderType.ANTHROPIC,
+    model_used   = "claude-opus-4-5",
+    extended_thinking_used = True,
+    input_tokens  = 312,
+    output_tokens = 891,
+    latency_ms   = 3241.4,
+    cost_usd     = 0.02184,   # completion cost only
+    total_cost_usd = 0.02188, # completion + classifier cost
+    classification = ClassificationResult(
+        tier           = QueryTier.DEEP,
+        effort         = EffortLevel.HIGH,
+        facts_ratio    = 0.12,
+        judgment_ratio = 0.88,
+        confidence     = 0.94,
+        reasoning      = "Complex distributed systems design — strategy tier warranted",
+        classifier_input_tokens  = 85,
+        classifier_output_tokens = 47,
+        classifier_cost_usd      = 0.0000456,
+    ),
+)
+```
+
+---
+
+## Live REST API
+
+**Base URL:** `https://query-oracle-production.up.railway.app`
+
+```bash
+# Health check (no auth required)
+curl https://query-oracle-production.up.railway.app/health
+# {"status":"ok"}
+
+# Classify only (auth required when QUERY_ORACLE_API_KEY is set)
+curl -s -X POST https://query-oracle-production.up.railway.app/classify \
+  -H "Authorization: Bearer $QUERY_ORACLE_API_KEY" \
+  -H "Content-Type: application/json" \
+  -d '{"query": "How many calories in a banana?"}' | jq .
+
+# Route and get a full response
+curl -s -X POST https://query-oracle-production.up.railway.app/route \
+  -H "Authorization: Bearer $QUERY_ORACLE_API_KEY" \
+  -H "Content-Type: application/json" \
+  -d '{"query": "Explain backpressure in reactive systems.", "provider": "anthropic"}' \
+  | jq '{tier, model_used, cost_usd, total_cost_usd, latency_ms}'
+```
+
+Interactive docs: [`/docs`](https://query-oracle-production.up.railway.app/docs)
+
+---
+
+## Securing your deployment
+
+By default the server runs in open mode (dev-only). Before exposing it publicly:
+
+**1. Set an API key**
+
+```bash
+export QUERY_ORACLE_API_KEY="your-secret-key"
+```
+
+All requests to `/route` and `/classify` then require:
+
+```
+Authorization: Bearer your-secret-key
+```
+
+The `/health` endpoint stays open. The server logs a warning at startup if no key is set.
+
+**2. Restrict CORS origins**
+
+```bash
+export CORS_ORIGINS="https://yourdomain.com,https://app.yourdomain.com"
+```
+
+When unset, CORS defaults to `*` (any origin). Set it to your specific domains in production.
+
+**3. Rate limiting**
+
+The `/route` endpoint enforces 30 requests per minute per client IP in-process. For heavier traffic, put an API gateway (e.g. Nginx, Cloudflare, Railway gateway) in front.
+
+**Railway environment variables:**
+
+```
+QUERY_ORACLE_API_KEY   → your secret key
+ANTHROPIC_API_KEY      → sk-ant-...
+OPENAI_API_KEY         → sk-...   (optional)
+GEMINI_API_KEY         → AIza...  (optional)
+CORS_ORIGINS           → https://yourdomain.com
+```
+
+---
+
+## Train your own classifier (coming soon)
+
+Once `logs/classifications.jsonl` accumulates ~500 entries, run `python training/train.py --data training/data/queries.csv` to fine-tune a local DistilBERT model. Swap it in with `QueryRouter(classifier=DistilBERTClassifier("training/query-classifier-final"))` and classification drops to ~10 ms with zero API cost.
+
+---
+
+## Contributing
+
+Open an issue or PR — the codebase is intentionally small. Adding a new provider means subclassing `BaseProvider` and implementing one method; the routing logic, classifier, and all integrations stay unchanged.
+
+---
+
+## Licence
+
+MIT — see [LICENSE](LICENSE).

query_oracle-0.3.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+llm_router/__init__.py,sha256=VSDa8-gl7rDuFSmdhMiU4jJtDA6aNz_bhk-2vMUAjZw,477
+llm_router/classifier.py,sha256=nJMaAdB-xRaeScbgmEmiFl2FgTOmuMhCTWyIeNguevs,12320
+llm_router/mcp_server.py,sha256=0XpkHiNEGh6YyFpLEnNwr_kMeICjmI7nxh0t8py4o8Y,2912
+llm_router/models.py,sha256=xb-GjPank1rWXNeiFnhVoH3L7-BRDTEemq_PuAt1Rvw,6393
+llm_router/router.py,sha256=n4F2aY0yBevvdBZnOebxnEBtZK6THGZQFOIJnS7kPSE,2880
+llm_router/providers/__init__.py,sha256=iZ3zMXUMYRaMpNPkSkgbqHQWGvGosFMW5pSCQhjPNLw,303
+llm_router/providers/anthropic_provider.py,sha256=l2sV3N2dbcsTeslqc7pXgXxDBpZlz9fMjUlIHDuwjXI,2282
+llm_router/providers/base.py,sha256=AEcukPfowcc8Fd45WCLfcj3EQreIQ2GpzrdotI0VFbk,708
+llm_router/providers/gemini_provider.py,sha256=4FFZHE74vd18wuJjze1sphUSlwz4Yvi1hHUe0EU5peM,1912
+llm_router/providers/openai_provider.py,sha256=LAQc9bG4e6Ve2H-jJbvnuBuy_B9crak4kavCrdRlSjg,1933
+query_oracle-0.3.0.dist-info/licenses/LICENSE,sha256=Odee6ZHHkgDBoj11Z4LLnTeO84n-HqR-OMNSpn5SW-o,1070
+query_oracle-0.3.0.dist-info/METADATA,sha256=R7CrmWfbqm44XcTpdKUrjyHnRRtCtIkIp8nCiLxpyo4,10642
+query_oracle-0.3.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+query_oracle-0.3.0.dist-info/entry_points.txt,sha256=-C5SFkPBsUIPoRTjRtLS5FW1EWYOd_NPMIVS5kC5WsY,64
+query_oracle-0.3.0.dist-info/top_level.txt,sha256=BJJqmvn9i4ijWlUe0K4GgX1wTgjOnnwRfXfbEKwn5Ao,11
+query_oracle-0.3.0.dist-info/RECORD,,

query_oracle-0.3.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

query_oracle-0.3.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+query-oracle-mcp = llm_router.mcp_server:main

query_oracle-0.3.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Heman Phadvas
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

query_oracle-0.3.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+llm_router