Skip to content

pyagent-router API Reference

pyagent_router.scorer.DifficultyScorer

Heuristic-based task difficulty scorer.

Scores tasks on a 1-10 scale using multiple signals: - Token length - Keyword complexity - Question structure - Required reasoning depth

Parameters:

Name Type Description Default
custom_signals dict[str, Any] | None

Optional dict of custom signal functions. Each function takes a task string and returns a float 0-1.

 None
Source code in packages/pyagent-router/src/pyagent_router/scorer.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
class DifficultyScorer:
    """Heuristic-based task difficulty scorer.

    Scores tasks on a 1-10 scale using multiple signals:
    - Token length
    - Keyword complexity
    - Question structure
    - Required reasoning depth

    Args:
        custom_signals: Optional dict of custom signal functions.
            Each function takes a task string and returns a float 0-1.
    """

    def __init__(
        self,
        custom_signals: dict[str, Any] | None = None,
    ) -> None:
        self._custom_signals = custom_signals or {}

    def score(self, task: str) -> DifficultyScore:
        """Score the difficulty of a task string."""
        signals: dict[str, float] = {}

        # Signal 1: Length complexity (longer = harder, up to a point)
        word_count = len(task.split())
        signals["length"] = min(word_count / 200, 1.0)

        # Signal 2: Keyword complexity
        task_lower = task.lower()
        complex_hits = sum(1 for kw in _COMPLEX_KEYWORDS if kw in task_lower)
        simple_hits = sum(1 for kw in _SIMPLE_KEYWORDS if kw in task_lower)
        signals["keywords"] = min(complex_hits / 3, 1.0) - min(simple_hits / 3, 0.5)
        signals["keywords"] = max(0.0, signals["keywords"])

        # Signal 3: Multi-part questions (numbered steps, multiple questions)
        multi_part = len(re.findall(r"\d+\.", task)) + task.count("?") - 1
        signals["multi_part"] = min(max(multi_part, 0) / 5, 1.0)

        # Signal 4: Code/math indicators
        code_indicators = sum(
            1 for marker in ["```", "def ", "class ", "function", "import"] if marker in task
        )
        math_indicators = sum(
            1 for marker in ["∑", "∫", "equation", "formula", "proof"] if marker in task_lower
        )
        signals["technical"] = min((code_indicators + math_indicators) / 3, 1.0)

        # Custom signals
        for name, fn in self._custom_signals.items():
            signals[name] = float(fn(task))

        # Weighted average → 1-10 scale
        weights = {"length": 0.15, "keywords": 0.35, "multi_part": 0.25, "technical": 0.25}
        weighted_sum = sum(signals.get(k, 0) * w for k, w in weights.items())

        # Add custom signals with equal weight
        if self._custom_signals:
            custom_avg = sum(signals.get(k, 0) for k in self._custom_signals) / len(
                self._custom_signals
            )
            weighted_sum = weighted_sum * 0.7 + custom_avg * 0.3

        raw_score = max(1, min(10, int(weighted_sum * 10) + 1))

        category = "easy" if raw_score <= 3 else "medium" if raw_score <= 6 else "hard"

        return DifficultyScore(score=raw_score, signals=signals, category=category)

score(task)

Score the difficulty of a task string.

Source code in packages/pyagent-router/src/pyagent_router/scorer.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def score(self, task: str) -> DifficultyScore:
    """Score the difficulty of a task string."""
    signals: dict[str, float] = {}

    # Signal 1: Length complexity (longer = harder, up to a point)
    word_count = len(task.split())
    signals["length"] = min(word_count / 200, 1.0)

    # Signal 2: Keyword complexity
    task_lower = task.lower()
    complex_hits = sum(1 for kw in _COMPLEX_KEYWORDS if kw in task_lower)
    simple_hits = sum(1 for kw in _SIMPLE_KEYWORDS if kw in task_lower)
    signals["keywords"] = min(complex_hits / 3, 1.0) - min(simple_hits / 3, 0.5)
    signals["keywords"] = max(0.0, signals["keywords"])

    # Signal 3: Multi-part questions (numbered steps, multiple questions)
    multi_part = len(re.findall(r"\d+\.", task)) + task.count("?") - 1
    signals["multi_part"] = min(max(multi_part, 0) / 5, 1.0)

    # Signal 4: Code/math indicators
    code_indicators = sum(
        1 for marker in ["```", "def ", "class ", "function", "import"] if marker in task
    )
    math_indicators = sum(
        1 for marker in ["∑", "∫", "equation", "formula", "proof"] if marker in task_lower
    )
    signals["technical"] = min((code_indicators + math_indicators) / 3, 1.0)

    # Custom signals
    for name, fn in self._custom_signals.items():
        signals[name] = float(fn(task))

    # Weighted average → 1-10 scale
    weights = {"length": 0.15, "keywords": 0.35, "multi_part": 0.25, "technical": 0.25}
    weighted_sum = sum(signals.get(k, 0) * w for k, w in weights.items())

    # Add custom signals with equal weight
    if self._custom_signals:
        custom_avg = sum(signals.get(k, 0) for k in self._custom_signals) / len(
            self._custom_signals
        )
        weighted_sum = weighted_sum * 0.7 + custom_avg * 0.3

    raw_score = max(1, min(10, int(weighted_sum * 10) + 1))

    category = "easy" if raw_score <= 3 else "medium" if raw_score <= 6 else "hard"

    return DifficultyScore(score=raw_score, signals=signals, category=category)

pyagent_router.scorer.DifficultyScore dataclass

Result of difficulty scoring.

Attributes:

Name Type Description
score int

Difficulty score from 1 (trivial) to 10 (extremely hard).
signals dict[str, float]

Dictionary of individual signal scores that contributed.
category str

Human-readable difficulty category.
Source code in packages/pyagent-router/src/pyagent_router/scorer.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
@dataclass(frozen=True)
class DifficultyScore:
    """Result of difficulty scoring.

    Attributes:
        score: Difficulty score from 1 (trivial) to 10 (extremely hard).
        signals: Dictionary of individual signal scores that contributed.
        category: Human-readable difficulty category.
    """

    score: int
    signals: dict[str, float] = field(default_factory=dict)
    category: str = ""

    @property
    def is_easy(self) -> bool:
        return self.score <= 3

    @property
    def is_medium(self) -> bool:
        return 4 <= self.score <= 6

    @property
    def is_hard(self) -> bool:
        return self.score >= 7

pyagent_router.estimator.CostEstimator

Estimate LLM call costs based on model pricing registry.

Parameters:

Name Type Description Default
pricing dict[str, ModelPricing] | None

Optional custom pricing dict. Defaults to built-in pricing table.

 None
default_output_ratio float

Estimated output/input token ratio when output length unknown.

 0.5
Source code in packages/pyagent-router/src/pyagent_router/estimator.py 
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
class CostEstimator:
    """Estimate LLM call costs based on model pricing registry.

    Args:
        pricing: Optional custom pricing dict. Defaults to built-in pricing table.
        default_output_ratio: Estimated output/input token ratio when output length unknown.
    """

    def __init__(
        self,
        pricing: dict[str, ModelPricing] | None = None,
        default_output_ratio: float = 0.5,
    ) -> None:
        self._pricing = pricing or dict(DEFAULT_PRICING)
        self._output_ratio = default_output_ratio

    def estimate(
        self,
        model: str,
        input_tokens: int,
        output_tokens: int | None = None,
    ) -> CostEstimate:
        """Estimate cost for a single LLM call.

        Args:
            model: Model name (must be in pricing registry).
            input_tokens: Number of input tokens.
            output_tokens: Number of output tokens. If None, estimated from input.

        Returns:
            CostEstimate with breakdown.

        Raises:
            KeyError: If model not found in pricing registry.
        """
        if model not in self._pricing:
            raise KeyError(
                f"Model '{model}' not in pricing registry. "
                f"Available: {', '.join(sorted(self._pricing.keys()))}"
            )

        pricing = self._pricing[model]
        est_output = (
            output_tokens if output_tokens is not None else int(input_tokens * self._output_ratio)
        )

        return CostEstimate(
            model=model,
            input_tokens=input_tokens,
            output_tokens=est_output,
            input_cost=input_tokens * pricing.input_per_million / 1_000_000,
            output_cost=est_output * pricing.output_per_million / 1_000_000,
        )

    def estimate_from_text(self, model: str, text: str) -> CostEstimate:
        """Estimate cost from raw text (approximates 4 chars per token)."""
        input_tokens = len(text) // 4
        return self.estimate(model, input_tokens)

    def compare(self, text: str, models: list[str] | None = None) -> list[CostEstimate]:
        """Compare costs across multiple models for the same input.

        Args:
            text: The input text.
            models: Models to compare. Defaults to all registered models.

        Returns:
            List of CostEstimates sorted by total_cost ascending.
        """
        target_models = models or list(self._pricing.keys())
        estimates = [self.estimate_from_text(m, text) for m in target_models if m in self._pricing]
        return sorted(estimates, key=lambda e: e.total_cost)

    @property
    def available_models(self) -> list[str]:
        return sorted(self._pricing.keys())

compare(text, models=None)

Compare costs across multiple models for the same input.

Parameters:

Name Type Description Default
text str

The input text.

 required
models list[str] | None

Models to compare. Defaults to all registered models.

 None

Returns:

Type Description
list[CostEstimate]

List of CostEstimates sorted by total_cost ascending.
Source code in packages/pyagent-router/src/pyagent_router/estimator.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
def compare(self, text: str, models: list[str] | None = None) -> list[CostEstimate]:
    """Compare costs across multiple models for the same input.

    Args:
        text: The input text.
        models: Models to compare. Defaults to all registered models.

    Returns:
        List of CostEstimates sorted by total_cost ascending.
    """
    target_models = models or list(self._pricing.keys())
    estimates = [self.estimate_from_text(m, text) for m in target_models if m in self._pricing]
    return sorted(estimates, key=lambda e: e.total_cost)

estimate(model, input_tokens, output_tokens=None)

Estimate cost for a single LLM call.

Parameters:

Name Type Description Default
model str

Model name (must be in pricing registry).

 required
input_tokens int

Number of input tokens.

 required
output_tokens int | None

Number of output tokens. If None, estimated from input.

 None

Returns:

Type Description
CostEstimate

CostEstimate with breakdown.

Raises:

Type Description
KeyError

If model not found in pricing registry.
Source code in packages/pyagent-router/src/pyagent_router/estimator.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def estimate(
    self,
    model: str,
    input_tokens: int,
    output_tokens: int | None = None,
) -> CostEstimate:
    """Estimate cost for a single LLM call.

    Args:
        model: Model name (must be in pricing registry).
        input_tokens: Number of input tokens.
        output_tokens: Number of output tokens. If None, estimated from input.

    Returns:
        CostEstimate with breakdown.

    Raises:
        KeyError: If model not found in pricing registry.
    """
    if model not in self._pricing:
        raise KeyError(
            f"Model '{model}' not in pricing registry. "
            f"Available: {', '.join(sorted(self._pricing.keys()))}"
        )

    pricing = self._pricing[model]
    est_output = (
        output_tokens if output_tokens is not None else int(input_tokens * self._output_ratio)
    )

    return CostEstimate(
        model=model,
        input_tokens=input_tokens,
        output_tokens=est_output,
        input_cost=input_tokens * pricing.input_per_million / 1_000_000,
        output_cost=est_output * pricing.output_per_million / 1_000_000,
    )

estimate_from_text(model, text)

Estimate cost from raw text (approximates 4 chars per token).

Source code in packages/pyagent-router/src/pyagent_router/estimator.py 
115
116
117
118
def estimate_from_text(self, model: str, text: str) -> CostEstimate:
    """Estimate cost from raw text (approximates 4 chars per token)."""
    input_tokens = len(text) // 4
    return self.estimate(model, input_tokens)

pyagent_router.estimator.CostEstimate dataclass

Estimated cost for a single LLM call.

Attributes:

Name Type Description
model str

Model name.
input_tokens int

Estimated input tokens.
output_tokens int

Estimated output tokens.
input_cost float

Cost for input tokens in USD.
output_cost float

Cost for output tokens in USD.
total_cost float

Total estimated cost in USD.
Source code in packages/pyagent-router/src/pyagent_router/estimator.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
@dataclass(frozen=True)
class CostEstimate:
    """Estimated cost for a single LLM call.

    Attributes:
        model: Model name.
        input_tokens: Estimated input tokens.
        output_tokens: Estimated output tokens.
        input_cost: Cost for input tokens in USD.
        output_cost: Cost for output tokens in USD.
        total_cost: Total estimated cost in USD.
    """

    model: str
    input_tokens: int
    output_tokens: int
    input_cost: float
    output_cost: float

    @property
    def total_cost(self) -> float:
        return self.input_cost + self.output_cost

pyagent_router.estimator.ModelPricing dataclass

Pricing for a model per 1M tokens.

Source code in packages/pyagent-router/src/pyagent_router/estimator.py 
12
13
14
15
16
17
@dataclass(frozen=True)
class ModelPricing:
    """Pricing for a model per 1M tokens."""

    input_per_million: float
    output_per_million: float

pyagent_router.selector.ModelSelector

Select the optimal model based on task difficulty and cost.

Strategy: find the cheapest model whose difficulty range covers the task and whose capabilities match the required capability (if specified).

Parameters:

Name Type Description Default
specs list[ModelSpec] | None

List of ModelSpec definitions. Defaults to built-in specs.

 None
cost_estimator CostEstimator | None

CostEstimator instance. Created automatically if None.

 None
scorer DifficultyScorer | None

DifficultyScorer instance. Created automatically if None.

 None
Source code in packages/pyagent-router/src/pyagent_router/selector.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
class ModelSelector:
    """Select the optimal model based on task difficulty and cost.

    Strategy: find the cheapest model whose difficulty range covers the task
    and whose capabilities match the required capability (if specified).

    Args:
        specs: List of ModelSpec definitions. Defaults to built-in specs.
        cost_estimator: CostEstimator instance. Created automatically if None.
        scorer: DifficultyScorer instance. Created automatically if None.
    """

    def __init__(
        self,
        specs: list[ModelSpec] | None = None,
        cost_estimator: CostEstimator | None = None,
        scorer: DifficultyScorer | None = None,
    ) -> None:
        self._specs = specs or list(DEFAULT_MODEL_SPECS)
        self._estimator = cost_estimator or CostEstimator()
        self._scorer = scorer or DifficultyScorer()

    def select(
        self,
        task: str,
        required_capability: Capability | None = None,
    ) -> SelectionResult:
        """Select the best model for a given task.

        Args:
            task: The task text to analyze.
            required_capability: Optional capability filter.

        Returns:
            SelectionResult with chosen model and reasoning.
        """
        difficulty = self._scorer.score(task)

        # Filter specs by difficulty range and capability
        candidates: list[ModelSpec] = []
        for spec in self._specs:
            if spec.min_difficulty <= difficulty.score <= spec.max_difficulty and (
                required_capability is None or required_capability in spec.capabilities
            ):
                candidates.append(spec)

        if not candidates:
            # Fallback to the most capable model
            candidates = [self._specs[-1]]

        # Estimate cost for each candidate and pick cheapest
        scored: list[tuple[ModelSpec, CostEstimate]] = []
        for spec in candidates:
            try:
                estimate = self._estimator.estimate_from_text(spec.name, task)
                scored.append((spec, estimate))
            except KeyError:
                continue

        if not scored:
            # Last resort fallback
            spec = self._specs[0]
            estimate = CostEstimate(spec.name, len(task) // 4, len(task) // 8, 0.0, 0.0)
            scored = [(spec, estimate)]

        scored.sort(key=lambda x: x[1].total_cost)
        chosen_spec, chosen_cost = scored[0]
        alternatives = [s.name for s, _ in scored[1:]]

        reason = (
            f"Difficulty {difficulty.score}/10 ({difficulty.category}) → "
            f"{chosen_spec.name} (cheapest candidate at ${chosen_cost.total_cost:.6f})"
        )

        return SelectionResult(
            model=chosen_spec.name,
            difficulty=difficulty,
            cost_estimate=chosen_cost,
            reason=reason,
            alternatives=alternatives,
        )

select(task, required_capability=None)

Select the best model for a given task.

Parameters:

Name Type Description Default
task str

The task text to analyze.

 required
required_capability Capability | None

Optional capability filter.

 None

Returns:

Type Description
SelectionResult

SelectionResult with chosen model and reasoning.
Source code in packages/pyagent-router/src/pyagent_router/selector.py 
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def select(
    self,
    task: str,
    required_capability: Capability | None = None,
) -> SelectionResult:
    """Select the best model for a given task.

    Args:
        task: The task text to analyze.
        required_capability: Optional capability filter.

    Returns:
        SelectionResult with chosen model and reasoning.
    """
    difficulty = self._scorer.score(task)

    # Filter specs by difficulty range and capability
    candidates: list[ModelSpec] = []
    for spec in self._specs:
        if spec.min_difficulty <= difficulty.score <= spec.max_difficulty and (
            required_capability is None or required_capability in spec.capabilities
        ):
            candidates.append(spec)

    if not candidates:
        # Fallback to the most capable model
        candidates = [self._specs[-1]]

    # Estimate cost for each candidate and pick cheapest
    scored: list[tuple[ModelSpec, CostEstimate]] = []
    for spec in candidates:
        try:
            estimate = self._estimator.estimate_from_text(spec.name, task)
            scored.append((spec, estimate))
        except KeyError:
            continue

    if not scored:
        # Last resort fallback
        spec = self._specs[0]
        estimate = CostEstimate(spec.name, len(task) // 4, len(task) // 8, 0.0, 0.0)
        scored = [(spec, estimate)]

    scored.sort(key=lambda x: x[1].total_cost)
    chosen_spec, chosen_cost = scored[0]
    alternatives = [s.name for s, _ in scored[1:]]

    reason = (
        f"Difficulty {difficulty.score}/10 ({difficulty.category}) → "
        f"{chosen_spec.name} (cheapest candidate at ${chosen_cost.total_cost:.6f})"
    )

    return SelectionResult(
        model=chosen_spec.name,
        difficulty=difficulty,
        cost_estimate=chosen_cost,
        reason=reason,
        alternatives=alternatives,
    )

pyagent_router.selector.ModelSpec dataclass

Specification of a model's capabilities and constraints.

Attributes:

Name Type Description
name str

Model identifier (must match pricing registry).
min_difficulty int

Minimum difficulty score this model should handle.
max_difficulty int

Maximum difficulty score this model should handle.
capabilities set[Capability]

Set of capabilities this model excels at.
max_context int

Maximum context window in tokens.
Source code in packages/pyagent-router/src/pyagent_router/selector.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
@dataclass
class ModelSpec:
    """Specification of a model's capabilities and constraints.

    Attributes:
        name: Model identifier (must match pricing registry).
        min_difficulty: Minimum difficulty score this model should handle.
        max_difficulty: Maximum difficulty score this model should handle.
        capabilities: Set of capabilities this model excels at.
        max_context: Maximum context window in tokens.
    """

    name: str
    min_difficulty: int = 1
    max_difficulty: int = 10
    capabilities: set[Capability] = field(default_factory=lambda: {Capability.GENERAL})
    max_context: int = 128_000

pyagent_router.selector.SelectionResult dataclass

Result of model selection.

Attributes:

Name Type Description
model str

Selected model name.
difficulty DifficultyScore

The difficulty assessment.
cost_estimate CostEstimate

Estimated cost for this model.
reason str

Human-readable explanation of why this model was chosen.
alternatives list[str]

Other models that were considered.
Source code in packages/pyagent-router/src/pyagent_router/selector.py 
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
@dataclass(frozen=True)
class SelectionResult:
    """Result of model selection.

    Attributes:
        model: Selected model name.
        difficulty: The difficulty assessment.
        cost_estimate: Estimated cost for this model.
        reason: Human-readable explanation of why this model was chosen.
        alternatives: Other models that were considered.
    """

    model: str
    difficulty: DifficultyScore
    cost_estimate: CostEstimate
    reason: str
    alternatives: list[str] = field(default_factory=list)

pyagent_router.selector.Capability

Bases: StrEnum

Model capabilities for filtering.

Source code in packages/pyagent-router/src/pyagent_router/selector.py 
16
17
18
19
20
21
22
23
24
class Capability(StrEnum):
    """Model capabilities for filtering."""

    CODE = "code"
    MATH = "math"
    REASONING = "reasoning"
    CREATIVE = "creative"
    GENERAL = "general"
    VISION = "vision"

pyagent_router.middleware.RouterMiddleware

Middleware that wraps agents with automatic model routing.

Usage

middleware = RouterMiddleware(model_registry={"gpt-4o": my_gpt4o, "gpt-4o-mini": my_mini}) routed_agent = middleware.wrap(my_agent)

routed_agent now auto-selects model per call

Parameters:

Name Type Description Default
model_registry dict[str, LLMCallable]

Mapping of model names to LLM callables.

 required
selector ModelSelector | None

Optional ModelSelector. Created with defaults if None.

 None
required_capability Capability | None

Optional default capability filter.

 None
Source code in packages/pyagent-router/src/pyagent_router/middleware.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
class RouterMiddleware:
    """Middleware that wraps agents with automatic model routing.

    Usage:
        middleware = RouterMiddleware(model_registry={"gpt-4o": my_gpt4o, "gpt-4o-mini": my_mini})
        routed_agent = middleware.wrap(my_agent)
        # routed_agent now auto-selects model per call

    Args:
        model_registry: Mapping of model names to LLM callables.
        selector: Optional ModelSelector. Created with defaults if None.
        required_capability: Optional default capability filter.
    """

    def __init__(
        self,
        model_registry: dict[str, LLMCallable],
        selector: ModelSelector | None = None,
        required_capability: Capability | None = None,
    ) -> None:
        self._registry = model_registry
        self._selector = selector or ModelSelector()
        self._capability = required_capability

    def wrap(self, agent: Agent) -> RoutedAgent:
        """Wrap an agent with routing capabilities."""
        return RoutedAgent(
            agent=agent,
            selector=self._selector,
            model_registry=self._registry,
            required_capability=self._capability,
        )

    def wrap_all(self, agents: list[Agent]) -> list[RoutedAgent]:
        """Wrap multiple agents."""
        return [self.wrap(a) for a in agents]

wrap(agent)

Wrap an agent with routing capabilities.

Source code in packages/pyagent-router/src/pyagent_router/middleware.py 
105
106
107
108
109
110
111
112
def wrap(self, agent: Agent) -> RoutedAgent:
    """Wrap an agent with routing capabilities."""
    return RoutedAgent(
        agent=agent,
        selector=self._selector,
        model_registry=self._registry,
        required_capability=self._capability,
    )

wrap_all(agents)

Wrap multiple agents.

Source code in packages/pyagent-router/src/pyagent_router/middleware.py 
114
115
116
def wrap_all(self, agents: list[Agent]) -> list[RoutedAgent]:
    """Wrap multiple agents."""
    return [self.wrap(a) for a in agents]

pyagent_router.middleware.RoutedAgent

Bases: Agent

An agent wrapper that routes each call through ModelSelector.

The routing decision is recorded in metadata for tracing.

Parameters:

Name Type Description Default
agent Agent

The original agent.

 required
selector ModelSelector

ModelSelector to use for routing decisions.

 required
model_registry dict[str, LLMCallable]

Mapping of model names to LLM callables.

 required
required_capability Capability | None

Optional capability filter for model selection.

 None
Source code in packages/pyagent-router/src/pyagent_router/middleware.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
class RoutedAgent(Agent):
    """An agent wrapper that routes each call through ModelSelector.

    The routing decision is recorded in metadata for tracing.

    Args:
        agent: The original agent.
        selector: ModelSelector to use for routing decisions.
        model_registry: Mapping of model names to LLM callables.
        required_capability: Optional capability filter for model selection.
    """

    def __init__(
        self,
        agent: Agent,
        selector: ModelSelector,
        model_registry: dict[str, LLMCallable],
        required_capability: Capability | None = None,
    ) -> None:
        super().__init__(
            name=agent.name,
            llm=agent.llm,
            system_prompt=agent.system_prompt,
            description=agent.description,
        )
        self._original_agent = agent
        self._selector = selector
        self._model_registry = model_registry
        self._required_capability = required_capability
        self.routing_log: list[SelectionResult] = []

    async def run(self, messages: list[Message]) -> Message:
        """Route the call to the optimal model, then execute."""
        # Extract task text from messages for difficulty scoring
        task_text = " ".join(m.content for m in messages if m.content)

        selection = self._selector.select(task_text, self._required_capability)
        self.routing_log.append(selection)

        # Swap LLM to the selected model if available in registry
        llm = self._model_registry.get(selection.model, self._original_agent.llm)

        # Create a temporary agent with the routed LLM
        routed = Agent(
            name=self._original_agent.name,
            llm=llm,
            system_prompt=self._original_agent.system_prompt,
            description=self._original_agent.description,
        )
        result = await routed.run(messages)

        # Attach routing metadata to the message
        result = Message(
            role=result.role,
            content=result.content,
            name=result.name,
            metadata={
                **result.metadata,
                "routed_model": selection.model,
                "difficulty": selection.difficulty.score,
                "estimated_cost": selection.cost_estimate.total_cost,
                "reason": selection.reason,
            },
        )
        return result

run(messages) async

Route the call to the optimal model, then execute.

Source code in packages/pyagent-router/src/pyagent_router/middleware.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
async def run(self, messages: list[Message]) -> Message:
    """Route the call to the optimal model, then execute."""
    # Extract task text from messages for difficulty scoring
    task_text = " ".join(m.content for m in messages if m.content)

    selection = self._selector.select(task_text, self._required_capability)
    self.routing_log.append(selection)

    # Swap LLM to the selected model if available in registry
    llm = self._model_registry.get(selection.model, self._original_agent.llm)

    # Create a temporary agent with the routed LLM
    routed = Agent(
        name=self._original_agent.name,
        llm=llm,
        system_prompt=self._original_agent.system_prompt,
        description=self._original_agent.description,
    )
    result = await routed.run(messages)

    # Attach routing metadata to the message
    result = Message(
        role=result.role,
        content=result.content,
        name=result.name,
        metadata={
            **result.metadata,
            "routed_model": selection.model,
            "difficulty": selection.difficulty.score,
            "estimated_cost": selection.cost_estimate.total_cost,
            "reason": selection.reason,
        },
    )
    return result