diff --git a/src/conductor/engine/pricing.py b/src/conductor/engine/pricing.py index 4e6787ba..ac1dd5fb 100644 --- a/src/conductor/engine/pricing.py +++ b/src/conductor/engine/pricing.py @@ -30,7 +30,7 @@ def _warn_fuzzy_match(requested: str, matched_key: str, strategy: str) -> None: _FUZZY_MATCH_WARNED.add(requested) logger.warning( "Pricing for model %r resolved via %s fallback to %r. " - "Cost and context_window metadata may be inaccurate. " + "Cost calculation may be inaccurate. " "Add %r to DEFAULT_PRICING or pass an override to silence this warning.", requested, strategy, @@ -41,66 +41,59 @@ def _warn_fuzzy_match(requested: str, matched_key: str, strategy: str) -> None: @dataclass(frozen=True) class ModelPricing: - """Pricing and metadata per model. + """Pricing per model. Attributes: input_per_mtok: Cost per million input tokens (USD). output_per_mtok: Cost per million output tokens (USD). cache_read_per_mtok: Cost per million cache read tokens (USD). cache_write_per_mtok: Cost per million cache write tokens (USD). - context_window: Context window size in tokens, or None if unknown. """ input_per_mtok: float output_per_mtok: float cache_read_per_mtok: float = 0.0 cache_write_per_mtok: float = 0.0 - context_window: int | None = None -# Default model table (pricing + context window metadata) -# Sources: OpenAI pricing page, Anthropic pricing page, provider docs +# Default model table (pricing only). +# Context-window metadata is sourced from each provider's SDK at runtime via +# ``AgentProvider.get_max_prompt_tokens()`` — see ``providers/base.py``. +# Sources: OpenAI pricing page, Anthropic pricing page, provider docs. DEFAULT_PRICING: dict[str, ModelPricing] = { # OpenAI / Copilot models - "gpt-4-turbo": ModelPricing( - input_per_mtok=10.00, output_per_mtok=30.00, context_window=128_000 - ), - "gpt-4o": ModelPricing(input_per_mtok=2.50, output_per_mtok=10.00, context_window=128_000), - "gpt-4o-mini": ModelPricing(input_per_mtok=0.15, output_per_mtok=0.60, context_window=128_000), - "gpt-4.1": ModelPricing(input_per_mtok=2.00, output_per_mtok=8.00, context_window=1_047_576), - "gpt-4.1-mini": ModelPricing( - input_per_mtok=0.15, output_per_mtok=0.60, context_window=1_047_576 - ), - "gpt-4": ModelPricing(input_per_mtok=30.00, output_per_mtok=60.00, context_window=8_192), - "gpt-3.5-turbo": ModelPricing(input_per_mtok=0.50, output_per_mtok=1.50, context_window=16_385), - "gpt-5.2": ModelPricing(input_per_mtok=2.00, output_per_mtok=8.00, context_window=400_000), - "gpt-5.1": ModelPricing(input_per_mtok=2.00, output_per_mtok=8.00, context_window=400_000), + "gpt-4-turbo": ModelPricing(input_per_mtok=10.00, output_per_mtok=30.00), + "gpt-4o": ModelPricing(input_per_mtok=2.50, output_per_mtok=10.00), + "gpt-4o-mini": ModelPricing(input_per_mtok=0.15, output_per_mtok=0.60), + "gpt-4.1": ModelPricing(input_per_mtok=2.00, output_per_mtok=8.00), + "gpt-4.1-mini": ModelPricing(input_per_mtok=0.15, output_per_mtok=0.60), + "gpt-4": ModelPricing(input_per_mtok=30.00, output_per_mtok=60.00), + "gpt-3.5-turbo": ModelPricing(input_per_mtok=0.50, output_per_mtok=1.50), + "gpt-5.2": ModelPricing(input_per_mtok=2.00, output_per_mtok=8.00), + "gpt-5.1": ModelPricing(input_per_mtok=2.00, output_per_mtok=8.00), # O-series - "o1": ModelPricing(input_per_mtok=15.00, output_per_mtok=60.00, context_window=200_000), - "o1-mini": ModelPricing(input_per_mtok=3.00, output_per_mtok=12.00, context_window=128_000), - "o1-preview": ModelPricing(input_per_mtok=15.00, output_per_mtok=60.00, context_window=128_000), - "o3-mini": ModelPricing(input_per_mtok=1.10, output_per_mtok=4.40, context_window=200_000), + "o1": ModelPricing(input_per_mtok=15.00, output_per_mtok=60.00), + "o1-mini": ModelPricing(input_per_mtok=3.00, output_per_mtok=12.00), + "o1-preview": ModelPricing(input_per_mtok=15.00, output_per_mtok=60.00), + "o3-mini": ModelPricing(input_per_mtok=1.10, output_per_mtok=4.40), # Claude 4.5 Series (newest) "claude-opus-4-5": ModelPricing( input_per_mtok=5.00, output_per_mtok=25.00, cache_read_per_mtok=0.50, cache_write_per_mtok=6.25, - context_window=200_000, ), "claude-sonnet-4-5": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "claude-haiku-4-5": ModelPricing( input_per_mtok=1.00, output_per_mtok=5.00, cache_read_per_mtok=0.10, cache_write_per_mtok=1.25, - context_window=200_000, ), # Short aliases for Claude 4.5 Series (used in workflow files) "opus-4.5": ModelPricing( @@ -108,21 +101,18 @@ class ModelPricing: output_per_mtok=25.00, cache_read_per_mtok=0.50, cache_write_per_mtok=6.25, - context_window=200_000, ), "sonnet-4.5": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "haiku-4.5": ModelPricing( input_per_mtok=1.00, output_per_mtok=5.00, cache_read_per_mtok=0.10, cache_write_per_mtok=1.25, - context_window=200_000, ), # Claude 4.6 Series "claude-opus-4.6": ModelPricing( @@ -130,21 +120,18 @@ class ModelPricing: output_per_mtok=25.00, cache_read_per_mtok=0.50, cache_write_per_mtok=6.25, - context_window=1_000_000, ), "claude-opus-4.6-1m": ModelPricing( input_per_mtok=5.00, output_per_mtok=25.00, cache_read_per_mtok=0.50, cache_write_per_mtok=6.25, - context_window=1_000_000, ), "claude-sonnet-4.6": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=1_000_000, ), # Claude 4 Series "claude-opus-4": ModelPricing( @@ -152,21 +139,18 @@ class ModelPricing: output_per_mtok=75.00, cache_read_per_mtok=1.50, cache_write_per_mtok=18.75, - context_window=200_000, ), "claude-sonnet-4": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "claude-haiku-4": ModelPricing( input_per_mtok=0.25, output_per_mtok=1.25, cache_read_per_mtok=0.03, cache_write_per_mtok=0.30, - context_window=200_000, ), # Claude 3.x Series "claude-3-7-sonnet": ModelPricing( @@ -174,69 +158,59 @@ class ModelPricing: output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "claude-3.7-sonnet": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "claude-3-5-sonnet": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "claude-3.5-sonnet": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "claude-3-5-haiku": ModelPricing( input_per_mtok=0.80, output_per_mtok=4.00, cache_read_per_mtok=0.08, cache_write_per_mtok=1.00, - context_window=200_000, ), "claude-3.5-haiku": ModelPricing( input_per_mtok=0.80, output_per_mtok=4.00, cache_read_per_mtok=0.08, cache_write_per_mtok=1.00, - context_window=200_000, ), "claude-3-opus": ModelPricing( input_per_mtok=15.00, output_per_mtok=75.00, cache_read_per_mtok=1.50, cache_write_per_mtok=18.75, - context_window=200_000, ), "claude-3-sonnet": ModelPricing( input_per_mtok=3.00, output_per_mtok=15.00, cache_read_per_mtok=0.30, cache_write_per_mtok=3.75, - context_window=200_000, ), "claude-3-haiku": ModelPricing( input_per_mtok=0.25, output_per_mtok=1.25, cache_read_per_mtok=0.03, cache_write_per_mtok=0.30, - context_window=200_000, ), # Gemini "gemini-3.1-pro-preview": ModelPricing( input_per_mtok=1.25, output_per_mtok=5.00, - context_window=1_000_000, ), } diff --git a/src/conductor/engine/workflow.py b/src/conductor/engine/workflow.py index 1037b054..c6076aad 100644 --- a/src/conductor/engine/workflow.py +++ b/src/conductor/engine/workflow.py @@ -766,16 +766,64 @@ async def _execute_subworkflow_with_inputs( usage = child_engine.usage_tracker.get_summary() return output, usage - def _get_context_window_for_agent(self, agent: AgentDef) -> int | None: - """Return the context window size for an agent's model.""" - from conductor.engine.pricing import get_pricing - - model = agent.model - if not model: + async def _get_provider_for_agent(self, agent: AgentDef) -> AgentProvider | None: + """Resolve the provider that will (or did) execute ``agent``. + + Mirrors the executor-resolution logic in ``_get_executor_for_agent`` + so context-window metadata lookups go through the same provider that + handles execution. Returns ``None`` only when no provider can be + determined (e.g. transient registry failures); callers must treat + ``None`` as "metadata unavailable". + """ + if self._registry is not None: + try: + return await self._registry.get_provider(agent) + except Exception as e: + logger.debug("Provider lookup via registry failed for %s: %s", agent.name, e) + return None + return self._single_provider + + async def _get_context_window_for_agent( + self, agent: AgentDef, output: AgentOutput | None = None + ) -> int | None: + """Return the SDK-reported max prompt tokens for an agent. + + Tries each candidate model in priority order — the model the SDK + actually used (``output.model``), the agent's configured model, the + workflow's runtime default — and returns the first non-``None`` + result. This is a real fallback chain: if ``output.model`` is an + SDK-specific variant the provider doesn't know about, the lookup + retries with ``agent.model`` before giving up. + + Returns ``None`` when no candidate resolves, no provider can be + reached, or the provider's metadata call fails — context-window + metadata is best-effort and must never break workflow execution. + """ + provider = await self._get_provider_for_agent(agent) + if provider is None: return None - - pricing = get_pricing(model) - return pricing.context_window if pricing else None + candidates: list[str] = [] + if output is not None and output.model: + candidates.append(output.model) + if agent.model and agent.model not in candidates: + candidates.append(agent.model) + default = self.config.workflow.runtime.default_model + if default and default not in candidates: + candidates.append(default) + for model in candidates: + try: + value = await provider.get_max_prompt_tokens(model) + except Exception as e: + logger.debug( + "get_max_prompt_tokens(%r) raised on provider for agent %s: %s", + model, + agent.name, + e, + ) + continue + if value is not None: + return value + return None async def run(self, inputs: dict[str, Any]) -> dict[str, Any]: """Execute the workflow from entry_point to $end. @@ -1520,7 +1568,9 @@ async def _execute_loop(self, current_agent_name: str) -> dict[str, Any]: "agent_name": agent.name, "iteration": agent_execution_count, "agent_type": agent.type or "agent", - "context_window_max": self._get_context_window_for_agent(agent), + "context_window_max": await self._get_context_window_for_agent( + agent + ), }, ) @@ -1844,7 +1894,9 @@ async def _execute_loop(self, current_agent_name: str) -> dict[str, Any]: "output": output.content, "output_keys": output_keys, "context_window_used": output.input_tokens, - "context_window_max": self._get_context_window_for_agent(agent), + "context_window_max": await self._get_context_window_for_agent( + agent, output + ), }, ) @@ -2472,7 +2524,9 @@ async def execute_single_agent(agent: AgentDef) -> tuple[str, Any]: "tokens": output.tokens_used, "cost_usd": usage.cost_usd, "context_window_used": output.input_tokens, - "context_window_max": self._get_context_window_for_agent(agent), + "context_window_max": await self._get_context_window_for_agent( + agent, output + ), }, ) diff --git a/src/conductor/providers/base.py b/src/conductor/providers/base.py index 20d8c915..2b56c6b5 100644 --- a/src/conductor/providers/base.py +++ b/src/conductor/providers/base.py @@ -7,8 +7,9 @@ from __future__ import annotations import asyncio +import re from abc import ABC, abstractmethod -from collections.abc import Callable +from collections.abc import Callable, Iterable from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -20,6 +21,46 @@ EventCallback = Callable[[str, dict[str, Any]], None] +# Suffixes that providers may strip when matching aliased model names against +# their SDK's canonical IDs (e.g. "claude-3-5-sonnet-latest" -> base name). +_VERSION_SUFFIX_RE = re.compile(r"-(\d{8}|latest|preview)$") + + +def match_model_id(requested: str, known_ids: Iterable[str]) -> str | None: + """Find the canonical SDK ID matching a possibly aliased model name. + + Match strategies, in order: + + 1. Exact match. + 2. Boundary prefix match (longest first), in either direction. Handles + both ``"claude-3-5-sonnet-20241022"`` for requested + ``"claude-3-5-sonnet"`` *and* the reverse, where the SDK lists a + dated ID and the user specified the base name. + 3. Suffix-strip (``-YYYYMMDD``, ``-latest``, ``-preview``) on the + requested name, then re-try strategies 1 and 2. + + Returns the matching SDK ID, or ``None`` if no strategy succeeds. + """ + ids = [str(i) for i in known_ids] + if not ids: + return None + if requested in ids: + return requested + sorted_ids = sorted(ids, key=lambda s: len(s), reverse=True) + for known in sorted_ids: + if requested.startswith(known + "-") or known.startswith(requested + "-"): + return known + simplified = _VERSION_SUFFIX_RE.sub("", requested) + if simplified == requested: + return None + if simplified in ids: + return simplified + for known in sorted_ids: + if simplified.startswith(known + "-") or known.startswith(simplified + "-"): + return known + return None + + @dataclass class AgentOutput: """Normalized output from any SDK provider. @@ -144,3 +185,33 @@ async def close(self) -> None: such as HTTP clients or session state. """ ... + + async def get_max_prompt_tokens(self, model: str) -> int | None: + """Return the SDK-reported maximum input (prompt) tokens for ``model``. + + This is the authoritative cap on prompt size enforced by the underlying + SDK or backend (e.g. the Copilot ``max_prompt_tokens`` field, or the + Anthropic ``max_input_tokens`` field). It is typically lower than the + model's theoretical context window — for example, the Copilot SDK + currently caps most GPT-5 variants at 128K despite a 400K model max. + + Implementations should: + + * Query their SDK's model-listing endpoint (cached after the first call). + * Return ``None`` when the model is unknown to the provider, when the + SDK call fails, or when no metadata is available. + * Never raise — context-window metadata is best-effort and must not + interrupt workflow execution. + + The default implementation returns ``None``, which causes the + dashboard's context-window bar to be hidden and any future enforcement + to be skipped — both safe degradations. + + Args: + model: The model identifier as it would be sent to the SDK + (e.g. ``"gpt-5.2"``, ``"claude-sonnet-4-5-20250929"``). + + Returns: + The maximum prompt (input) tokens the SDK will accept, or ``None``. + """ + return None diff --git a/src/conductor/providers/claude.py b/src/conductor/providers/claude.py index 4f8af2bf..b3feec81 100644 --- a/src/conductor/providers/claude.py +++ b/src/conductor/providers/claude.py @@ -30,7 +30,7 @@ from conductor.exceptions import ProviderError, ValidationError from conductor.executor.output import validate_output -from conductor.providers.base import AgentOutput, AgentProvider, EventCallback +from conductor.providers.base import AgentOutput, AgentProvider, EventCallback, match_model_id if TYPE_CHECKING: from conductor.config.schema import AgentDef, OutputField @@ -39,13 +39,14 @@ # Try to import the Anthropic SDK try: import anthropic - from anthropic import AsyncAnthropic + from anthropic import AnthropicError, AsyncAnthropic ANTHROPIC_SDK_AVAILABLE = True except ImportError: ANTHROPIC_SDK_AVAILABLE = False AsyncAnthropic = None # type: ignore[misc, assignment] anthropic = None # type: ignore[assignment] + AnthropicError = Exception # type: ignore[misc, assignment] logger = logging.getLogger(__name__) @@ -178,6 +179,13 @@ def __init__( self._mcp_servers_config = mcp_servers self._mcp_manager: MCPManager | None = None + # Cache of model_id -> max_input_tokens populated lazily on first + # get_max_prompt_tokens() call. Guarded by an asyncio.Lock to avoid + # racing concurrent first-callers and emitting duplicate models.list() + # requests. + self._max_input_cache: dict[str, int | None] | None = None + self._max_input_cache_lock = asyncio.Lock() + # Initialize the client (sync initialization) self._initialize_client() @@ -288,7 +296,8 @@ async def validate_connection(self) -> bool: async def _log_available_models(self) -> None: """List and log available models, warn if default model is unavailable. - Consolidated from the former _verify_available_models method. + Also seeds ``_max_input_cache`` so the first call to + :meth:`get_max_prompt_tokens` doesn't pay for an extra round-trip. """ if self._client is None: return @@ -297,20 +306,73 @@ async def _log_available_models(self) -> None: # Call client.models.list() to get available models (async) logger.debug("Discovering available Claude models via client.models.list()...") models_page = await self._client.models.list() - available_models = [model.id for model in models_page.data] + except (TimeoutError, AnthropicError, OSError) as e: + logger.warning(f"Could not list available models (discovery failed): {e}") + return - logger.info(f"Available Claude models: {', '.join(available_models)}") + available_models = [model.id for model in models_page.data] + logger.info(f"Available Claude models: {', '.join(available_models)}") - # Warn if default model not in list - if self._default_model not in available_models: - logger.warning( - f"Requested model '{self._default_model}' is not in the list of " - f"available models. API calls may fail. Available: {available_models}" - ) - else: - logger.debug(f"Default model '{self._default_model}' verified in available models") - except Exception as e: - logger.warning(f"Could not list available models (discovery failed): {e}") + # Warn if default model not in list (after stripping aliases like -latest). + if match_model_id(self._default_model, available_models) is None: + logger.warning( + f"Requested model '{self._default_model}' is not in the list of " + f"available models. API calls may fail. Available: {available_models}" + ) + else: + logger.debug(f"Default model '{self._default_model}' verified in available models") + + # Seed the metadata cache so get_max_prompt_tokens() is a pure lookup. + self._install_max_input_cache(models_page.data) + + def _install_max_input_cache(self, models_data: list[Any]) -> None: + """Replace ``_max_input_cache`` with a fresh mapping of id -> max_input.""" + self._max_input_cache = { + info.id: getattr(info, "max_input_tokens", None) for info in models_data + } + + async def get_max_prompt_tokens(self, model: str) -> int | None: + """Return the Anthropic SDK's ``max_input_tokens`` for ``model``. + + On first call, populates a per-instance cache by enumerating + ``client.models.list()``; subsequent calls are dictionary lookups. + ``validate_connection()`` already populates the cache, so callers + that go through normal connection setup never pay for an extra + round-trip. + + Resolves aliases (``-latest``, dated suffixes, base/versioned name + mismatches) via :func:`match_model_id`. Returns ``None`` when the + SDK is unavailable, the model can't be resolved, or the listing + call fails — context-window metadata must never block workflow + execution. + + Note: the value reflects the API's *default* input window. Claude + models with a 1M-context beta require an explicit beta header, + which Conductor does not set today; for those models the API still + reports the default window. + """ + if not ANTHROPIC_SDK_AVAILABLE or self._client is None: + return None + + if self._max_input_cache is None: + # Fetch outside the lock so concurrent callers don't all queue + # behind a slow round-trip; the lock only guards the install. + try: + page = await self._client.models.list() + except (TimeoutError, AnthropicError, OSError) as e: + # Don't cache the failure — let the next call retry. + logger.debug("Failed to list Anthropic models: %s", e) + return None + async with self._max_input_cache_lock: + if self._max_input_cache is None: + self._install_max_input_cache(page.data) + + # The block above either returned early on failure or installed the + # cache, so it's guaranteed non-None here. + cache = self._max_input_cache + assert cache is not None + matched_id = match_model_id(model, cache.keys()) + return cache.get(matched_id) if matched_id is not None else None async def _ensure_mcp_connected(self) -> None: """Connect to MCP servers if configured. @@ -398,11 +460,18 @@ async def close(self) -> None: logger.debug("MCP manager closed") if self._client is not None: - # AsyncAnthropic uses httpx AsyncClient internally which should be closed - await self._client.close() + # Drop the client reference *before* awaiting close() so any + # in-flight get_max_prompt_tokens() observes None on its next + # access and skips the SDK call. Already-issued requests will + # error and be swallowed by the metadata path's narrow except. + client = self._client self._client = None + await client.close() logger.debug("Claude provider closed") + # Drop cached metadata so a re-initialized provider re-fetches. + self._max_input_cache = None + async def execute( self, agent: AgentDef, diff --git a/src/conductor/providers/copilot.py b/src/conductor/providers/copilot.py index 94878040..fbb8fba6 100644 --- a/src/conductor/providers/copilot.py +++ b/src/conductor/providers/copilot.py @@ -18,7 +18,7 @@ from typing import TYPE_CHECKING, Any from conductor.exceptions import ProviderError, ValidationError -from conductor.providers.base import AgentOutput, AgentProvider, EventCallback +from conductor.providers.base import AgentOutput, AgentProvider, EventCallback, match_model_id if TYPE_CHECKING: from conductor.config.schema import AgentDef, OutputField @@ -1775,6 +1775,34 @@ async def close(self) -> None: self._call_history.clear() self._retry_history.clear() + async def get_max_prompt_tokens(self, model: str) -> int | None: + """Return the Copilot SDK's ``max_prompt_tokens`` for ``model``. + + Queries ``client.list_models()`` (cached internally by the SDK), + resolves any aliases (e.g. ``-latest``, dated suffixes, base-name + vs versioned-name) via :func:`match_model_id`, and returns + ``capabilities.limits.max_prompt_tokens`` for the matched entry. + + Returns ``None`` in mock-handler mode, when the SDK is unavailable, + when no match is found, or when the SDK call fails — context-window + metadata must never block workflow execution. + """ + if self._mock_handler is not None or not COPILOT_SDK_AVAILABLE: + return None + try: + await self._ensure_client_started() + models = await self._client.list_models() + except (TimeoutError, ProviderError, OSError, RuntimeError) as e: + logger.debug("Failed to list Copilot models for %r: %s", model, e) + return None + by_id = {info.id: info for info in models} + matched_id = match_model_id(model, by_id.keys()) + if matched_id is None: + return None + info = by_id[matched_id] + limits = getattr(info.capabilities, "limits", None) + return getattr(limits, "max_prompt_tokens", None) + def get_session_ids(self) -> dict[str, str]: """Get tracked session IDs for all executed agents. diff --git a/tests/test_engine/test_context_window_events.py b/tests/test_engine/test_context_window_events.py index c24e3562..ec31bd7e 100644 --- a/tests/test_engine/test_context_window_events.py +++ b/tests/test_engine/test_context_window_events.py @@ -1,4 +1,10 @@ -"""Tests that workflow events include context_window fields.""" +"""Tests that workflow events include context_window fields. + +Context-window metadata is now sourced from each provider's SDK at runtime +(``AgentProvider.get_max_prompt_tokens``). In mock-handler mode the Copilot +provider has no SDK to query and returns ``None`` by default, so these tests +monkeypatch the provider method to inject the values being asserted. +""" from __future__ import annotations @@ -45,6 +51,18 @@ def _make_emitter_and_collector() -> tuple[WorkflowEventEmitter, EventCollector] return emitter, collector +def _provider_with_max_prompt(values: dict[str, int | None]) -> CopilotProvider: + """Build a mock-handler Copilot provider whose ``get_max_prompt_tokens`` + returns values from ``values`` (or ``None`` for unknown models).""" + provider = CopilotProvider(mock_handler=lambda a, p, c: {"answer": "hi", "result": a.name}) + + async def fake_get_max_prompt_tokens(model: str) -> int | None: + return values.get(model) + + provider.get_max_prompt_tokens = fake_get_max_prompt_tokens # type: ignore[method-assign] + return provider + + class TestAgentStartedContextWindow: """agent_started event includes context_window_max.""" @@ -70,7 +88,7 @@ async def test_agent_started_has_context_window_max(self) -> None: ], output={"answer": "{{ a1.output.answer }}"}, ) - provider = CopilotProvider(mock_handler=lambda a, p, c: {"answer": "hi"}) + provider = _provider_with_max_prompt({"gpt-4o": 128000}) engine = WorkflowEngine(config, provider, event_emitter=emitter) await engine.run({}) @@ -104,7 +122,7 @@ async def test_agent_completed_has_context_window_fields(self) -> None: ], output={"answer": "{{ a1.output.answer }}"}, ) - provider = CopilotProvider(mock_handler=lambda a, p, c: {"answer": "hi"}) + provider = _provider_with_max_prompt({"gpt-4o": 128000}) engine = WorkflowEngine(config, provider, event_emitter=emitter) await engine.run({}) @@ -115,7 +133,7 @@ async def test_agent_completed_has_context_window_fields(self) -> None: class TestContextWindowNoneForUnknownModel: - """context_window_max is None when model is unknown.""" + """context_window_max is None when the provider has no metadata for the model.""" @pytest.mark.asyncio async def test_unknown_model_returns_none(self) -> None: @@ -139,7 +157,8 @@ async def test_unknown_model_returns_none(self) -> None: ], output={"answer": "{{ a1.output.answer }}"}, ) - provider = CopilotProvider(mock_handler=lambda a, p, c: {"answer": "hi"}) + # Empty metadata table — every lookup returns None. + provider = _provider_with_max_prompt({}) engine = WorkflowEngine(config, provider, event_emitter=emitter) await engine.run({}) @@ -184,7 +203,7 @@ async def test_parallel_agent_completed_has_context_window_fields(self) -> None: ], output={"result": "done"}, ) - provider = CopilotProvider(mock_handler=lambda a, p, c: {"result": a.name}) + provider = _provider_with_max_prompt({"gpt-4o": 128000}) engine = WorkflowEngine(config, provider, event_emitter=emitter) await engine.run({}) @@ -229,10 +248,146 @@ async def test_parallel_agent_unknown_model_context_window_none(self) -> None: ], output={"result": "done"}, ) - provider = CopilotProvider(mock_handler=lambda a, p, c: {"result": "ok"}) + provider = _provider_with_max_prompt({}) engine = WorkflowEngine(config, provider, event_emitter=emitter) await engine.run({}) events = collector.of_type("parallel_agent_completed") for event in events: assert event.data["context_window_max"] is None + + +class TestContextWindowResolutionOrder: + """The model is resolved from output.model first, then agent.model, then default.""" + + @pytest.mark.asyncio + async def test_default_model_used_when_agent_has_no_model(self) -> None: + emitter, collector = _make_emitter_and_collector() + config = WorkflowConfig( + workflow=WorkflowDef( + name="test-default", + entry_point="a1", + runtime=RuntimeConfig(provider="copilot", default_model="gpt-4o"), + context=ContextConfig(mode="accumulate"), + limits=LimitsConfig(max_iterations=10), + ), + agents=[ + AgentDef( + name="a1", + prompt="Hello", + output={"answer": OutputField(type="string")}, + routes=[RouteDef(to="$end")], + ), + ], + output={"answer": "{{ a1.output.answer }}"}, + ) + provider = _provider_with_max_prompt({"gpt-4o": 128000}) + engine = WorkflowEngine(config, provider, event_emitter=emitter) + await engine.run({}) + + # agent_started has no output yet, but should resolve via default_model. + assert collector.first("agent_started").data["context_window_max"] == 128000 + + @pytest.mark.asyncio + async def test_output_model_preferred_over_configured(self) -> None: + emitter, collector = _make_emitter_and_collector() + config = WorkflowConfig( + workflow=WorkflowDef( + name="test-output-model", + entry_point="a1", + runtime=RuntimeConfig(provider="copilot"), + context=ContextConfig(mode="accumulate"), + limits=LimitsConfig(max_iterations=10), + ), + agents=[ + AgentDef( + name="a1", + model="gpt-4o", + prompt="Hello", + output={"answer": OutputField(type="string")}, + routes=[RouteDef(to="$end")], + ), + ], + output={"answer": "{{ a1.output.answer }}"}, + ) + + # Mock handler reports a different model than the agent requested + # (e.g. the SDK aliased or substituted it). + def handler(agent, prompt, context): # type: ignore[no-untyped-def] + return {"answer": "hi"} + + provider = CopilotProvider(mock_handler=handler) + + async def fake_get_max_prompt_tokens(model: str) -> int | None: + return {"gpt-4o": 128000, "gpt-5.2": 400000}.get(model) + + provider.get_max_prompt_tokens = fake_get_max_prompt_tokens # type: ignore[method-assign] + + # Force the AgentOutput.model field via a wrapper. The simplest hook + # here is to set the model on the mock-execute return — done by + # patching the provider's execute to override output.model. + original_execute = provider.execute + + async def execute_with_model(*args, **kwargs): # type: ignore[no-untyped-def] + output = await original_execute(*args, **kwargs) + output.model = "gpt-5.2" + return output + + provider.execute = execute_with_model # type: ignore[method-assign] + + engine = WorkflowEngine(config, provider, event_emitter=emitter) + await engine.run({}) + + # agent_started runs before execution, no output yet — uses agent.model + assert collector.first("agent_started").data["context_window_max"] == 128000 + # agent_completed has output.model — uses that + assert collector.first("agent_completed").data["context_window_max"] == 400000 + + @pytest.mark.asyncio + async def test_falls_back_to_agent_model_when_output_model_unknown(self) -> None: + """If output.model is an SDK-unknown variant (e.g. a reasoning-effort + tier the provider doesn't list), the chain retries with agent.model + rather than returning None.""" + emitter, collector = _make_emitter_and_collector() + config = WorkflowConfig( + workflow=WorkflowDef( + name="test-fallback", + entry_point="a1", + runtime=RuntimeConfig(provider="copilot"), + context=ContextConfig(mode="accumulate"), + limits=LimitsConfig(max_iterations=10), + ), + agents=[ + AgentDef( + name="a1", + model="claude-opus-4.7", + prompt="Hello", + output={"answer": OutputField(type="string")}, + routes=[RouteDef(to="$end")], + ), + ], + output={"answer": "{{ a1.output.answer }}"}, + ) + + provider = CopilotProvider(mock_handler=lambda a, p, c: {"answer": "hi"}) + + # Provider only knows the base name, not the reasoning-effort variant. + async def fake_get_max_prompt_tokens(model: str) -> int | None: + return {"claude-opus-4.7": 200_000}.get(model) + + provider.get_max_prompt_tokens = fake_get_max_prompt_tokens # type: ignore[method-assign] + + original_execute = provider.execute + + async def execute_with_variant_model(*args, **kwargs): # type: ignore[no-untyped-def] + output = await original_execute(*args, **kwargs) + output.model = "claude-opus-4.7-xhigh" # SDK doesn't know this name + return output + + provider.execute = execute_with_variant_model # type: ignore[method-assign] + + engine = WorkflowEngine(config, provider, event_emitter=emitter) + await engine.run({}) + + # output.model returned None; chain fell back to agent.model. + assert collector.first("agent_completed").data["context_window_max"] == 200_000 diff --git a/tests/test_providers/test_base.py b/tests/test_providers/test_base.py index 86a8cbea..aaf4910f 100644 --- a/tests/test_providers/test_base.py +++ b/tests/test_providers/test_base.py @@ -1,6 +1,6 @@ """Unit tests for the AgentProvider ABC and AgentOutput dataclass.""" -from conductor.providers.base import AgentOutput +from conductor.providers.base import AgentOutput, match_model_id class TestAgentOutput: @@ -67,3 +67,54 @@ def test_agent_output_with_complex_content(self) -> None: assert output.content["analysis"]["score"] == 8.5 assert output.content["analysis"]["issues"] == ["minor", "cosmetic"] assert output.content["analysis"]["approved"] is True + + +class TestMatchModelId: + """Unit tests for the alias-aware model ID matcher.""" + + def test_exact_match(self) -> None: + assert match_model_id("gpt-4o", ["gpt-4o", "gpt-4.1"]) == "gpt-4o" + + def test_returns_none_when_no_known_ids(self) -> None: + assert match_model_id("gpt-4o", []) is None + + def test_returns_none_for_unrelated_name(self) -> None: + assert match_model_id("totally-different", ["gpt-4o"]) is None + + def test_versioned_suffix_matches_base(self) -> None: + # Requested name has a dated suffix; SDK lists the base name. + assert ( + match_model_id("claude-3-5-sonnet-20241022", ["claude-3-5-sonnet"]) + == "claude-3-5-sonnet" + ) + + def test_base_matches_versioned_sdk_id(self) -> None: + # Requested name is the base; SDK lists a dated/aliased variant. + assert ( + match_model_id("claude-3-5-sonnet", ["claude-3-5-sonnet-20241022"]) + == "claude-3-5-sonnet-20241022" + ) + + def test_latest_alias_strips_and_matches(self) -> None: + assert ( + match_model_id("claude-3-5-sonnet-latest", ["claude-3-5-sonnet-20241022"]) + == "claude-3-5-sonnet-20241022" + ) + + def test_preview_alias_strips_and_matches(self) -> None: + assert match_model_id("gemini-3.1-pro-preview", ["gemini-3.1-pro"]) == "gemini-3.1-pro" + + def test_longest_match_wins(self) -> None: + # "o1-mini-20240101" must match "o1-mini" (longer), not "o1". + assert match_model_id("o1-mini-20240101", ["o1", "o1-mini"]) == "o1-mini" + + def test_boundary_check_prevents_cross_family_match(self) -> None: + # "claude-opus-4.7" must NOT match "claude-opus-4" (different family). + # No valid match -> None. + assert match_model_id("claude-opus-4.7", ["claude-opus-4"]) is None + + def test_boundary_check_prevents_cross_family_with_suffix(self) -> None: + assert match_model_id("claude-opus-4.7-high", ["claude-opus-4"]) is None + + def test_unknown_after_suffix_strip_returns_none(self) -> None: + assert match_model_id("totally-different-latest", ["gpt-4o"]) is None diff --git a/tests/test_providers/test_claude.py b/tests/test_providers/test_claude.py index bab3fabe..8365ac3e 100644 --- a/tests/test_providers/test_claude.py +++ b/tests/test_providers/test_claude.py @@ -2302,3 +2302,140 @@ def __init__(self) -> None: # Verify retry-after header was used (delay should be 5.0) assert len(provider._retry_history) == 1 assert provider._retry_history[0]["delay"] == 5.0 + + +@patch("conductor.providers.claude.ANTHROPIC_SDK_AVAILABLE", True) +@patch("conductor.providers.claude.AsyncAnthropic") +class TestClaudeGetMaxPromptTokens: + """Tests for ClaudeProvider.get_max_prompt_tokens.""" + + @pytest.mark.asyncio + async def test_returns_max_input_tokens_for_known_model( + self, mock_anthropic_class: Mock + ) -> None: + mock_client = Mock() + mock_client.models.list = AsyncMock( + return_value=Mock( + data=[ + Mock(id="claude-sonnet-4-5", max_input_tokens=200_000), + Mock(id="claude-opus-4-5", max_input_tokens=200_000), + ] + ) + ) + mock_anthropic_class.return_value = mock_client + + provider = ClaudeProvider() + assert await provider.get_max_prompt_tokens("claude-sonnet-4-5") == 200_000 + + @pytest.mark.asyncio + async def test_returns_none_for_unknown_model(self, mock_anthropic_class: Mock) -> None: + mock_client = Mock() + mock_client.models.list = AsyncMock(return_value=Mock(data=[])) + mock_anthropic_class.return_value = mock_client + + provider = ClaudeProvider() + assert await provider.get_max_prompt_tokens("unknown-x") is None + + @pytest.mark.asyncio + async def test_sdk_failure_returns_none_and_does_not_cache( + self, mock_anthropic_class: Mock + ) -> None: + """An SDK exception is swallowed and not cached, so a later call retries.""" + from anthropic import APIConnectionError + + # APIConnectionError requires a request kwarg; build a minimal one. + err = APIConnectionError(request=Mock()) + + mock_client = Mock() + # First call raises, second call succeeds — proves the failure isn't + # cached as "no metadata" forever. + mock_client.models.list = AsyncMock( + side_effect=[ + err, + Mock(data=[Mock(id="claude-sonnet-4-5", max_input_tokens=200_000)]), + ] + ) + mock_anthropic_class.return_value = mock_client + + provider = ClaudeProvider() + assert await provider.get_max_prompt_tokens("claude-sonnet-4-5") is None + assert await provider.get_max_prompt_tokens("claude-sonnet-4-5") == 200_000 + assert mock_client.models.list.await_count == 2 + + @pytest.mark.asyncio + async def test_unexpected_exception_propagates(self, mock_anthropic_class: Mock) -> None: + """Non-SDK exceptions (programming errors) are not swallowed by the + provider — they bubble up so the engine's outer safety net handles them.""" + mock_client = Mock() + mock_client.models.list = AsyncMock(side_effect=RuntimeError("bug")) + mock_anthropic_class.return_value = mock_client + + provider = ClaudeProvider() + with pytest.raises(RuntimeError): + await provider.get_max_prompt_tokens("claude-sonnet-4-5") + + @pytest.mark.asyncio + async def test_alias_resolves_via_match_model_id(self, mock_anthropic_class: Mock) -> None: + """``-latest`` and dated suffix aliases resolve to the SDK's listed ID.""" + mock_client = Mock() + mock_client.models.list = AsyncMock( + return_value=Mock( + data=[ + Mock(id="claude-3-5-sonnet-20241022", max_input_tokens=200_000), + ] + ) + ) + mock_anthropic_class.return_value = mock_client + + provider = ClaudeProvider() + # `-latest` strips to base, then prefix-matches the dated SDK ID. + assert await provider.get_max_prompt_tokens("claude-3-5-sonnet-latest") == 200_000 + # The base name (no dated suffix) also matches the dated SDK ID. + assert await provider.get_max_prompt_tokens("claude-3-5-sonnet") == 200_000 + + @pytest.mark.asyncio + async def test_caches_after_first_call(self, mock_anthropic_class: Mock) -> None: + """Second call must hit the cache, not the SDK.""" + mock_client = Mock() + mock_client.models.list = AsyncMock( + return_value=Mock(data=[Mock(id="claude-sonnet-4-5", max_input_tokens=200_000)]) + ) + mock_anthropic_class.return_value = mock_client + + provider = ClaudeProvider() + await provider.get_max_prompt_tokens("claude-sonnet-4-5") + await provider.get_max_prompt_tokens("claude-sonnet-4-5") + await provider.get_max_prompt_tokens("anything-else") + + assert mock_client.models.list.await_count == 1 + + @pytest.mark.asyncio + async def test_validate_connection_seeds_cache(self, mock_anthropic_class: Mock) -> None: + """``validate_connection()`` populates the cache so the first + ``get_max_prompt_tokens()`` call is a pure dict lookup.""" + mock_client = Mock() + mock_client.models.list = AsyncMock( + return_value=Mock(data=[Mock(id="claude-sonnet-4-5", max_input_tokens=200_000)]) + ) + mock_anthropic_class.return_value = mock_client + + provider = ClaudeProvider() + assert await provider.validate_connection() is True + # validate_connection itself called list (once for the API check, then + # _log_available_models reuses the response). Reset the counter to + # prove get_max_prompt_tokens doesn't add another call. + before = mock_client.models.list.await_count + assert await provider.get_max_prompt_tokens("claude-sonnet-4-5") == 200_000 + assert mock_client.models.list.await_count == before + + @pytest.mark.asyncio + @patch("conductor.providers.claude.ANTHROPIC_SDK_AVAILABLE", False) + async def test_returns_none_when_sdk_unavailable(self, mock_anthropic_class: Mock) -> None: + # Need a workaround: ANTHROPIC_SDK_AVAILABLE is False so __init__ + # raises. Build an instance bypassing the init guard by patching + # only at call time. + with patch("conductor.providers.claude.ANTHROPIC_SDK_AVAILABLE", True): + provider = ClaudeProvider() + + with patch("conductor.providers.claude.ANTHROPIC_SDK_AVAILABLE", False): + assert await provider.get_max_prompt_tokens("claude-sonnet-4-5") is None diff --git a/tests/test_providers/test_context_window.py b/tests/test_providers/test_context_window.py deleted file mode 100644 index cc2b2d6f..00000000 --- a/tests/test_providers/test_context_window.py +++ /dev/null @@ -1,86 +0,0 @@ -"""Tests for context window lookups via the unified pricing registry.""" - -from __future__ import annotations - -from conductor.engine.pricing import DEFAULT_PRICING, get_pricing - - -def _context_window(model: str) -> int | None: - """Helper: look up context window via the pricing registry.""" - pricing = get_pricing(model) - return pricing.context_window if pricing else None - - -class TestExactMatch: - """Exact model name lookups.""" - - def test_claude_sonnet_4(self) -> None: - assert _context_window("claude-sonnet-4") == 200_000 - - def test_claude_opus_4_6_1m(self) -> None: - assert _context_window("claude-opus-4.6-1m") == 1_000_000 - - def test_gpt_4o(self) -> None: - assert _context_window("gpt-4o") == 128_000 - - def test_gpt_4_legacy(self) -> None: - assert _context_window("gpt-4") == 8_192 - - def test_gpt_4_1(self) -> None: - assert _context_window("gpt-4.1") == 1_047_576 - - def test_short_alias(self) -> None: - assert _context_window("sonnet-4.5") == 200_000 - - def test_gemini(self) -> None: - assert _context_window("gemini-3.1-pro-preview") == 1_000_000 - - -class TestPrefixMatch: - """Prefix-based fuzzy matching.""" - - def test_dated_suffix(self) -> None: - assert _context_window("claude-sonnet-4-20250514") == 200_000 - - def test_latest_suffix(self) -> None: - assert _context_window("claude-3-5-sonnet-latest") == 200_000 - - def test_preview_suffix(self) -> None: - assert _context_window("claude-3-5-sonnet-preview") == 200_000 - - def test_o1_mini_prefers_longer_key(self) -> None: - """o1-mini must match 'o1-mini' (128K), not 'o1' (200K).""" - assert _context_window("o1-mini") == 128_000 - - def test_o1_mini_with_date(self) -> None: - assert _context_window("o1-mini-20240101") == 128_000 - - def test_dot_and_dash_notation(self) -> None: - assert _context_window("claude-3.5-sonnet") == 200_000 - assert _context_window("claude-3-5-sonnet") == 200_000 - - -class TestUnknownModel: - """Unknown models return None.""" - - def test_completely_unknown(self) -> None: - assert _context_window("totally-unknown-model") is None - - def test_empty_string(self) -> None: - assert _context_window("") is None - - def test_partial_gpt(self) -> None: - assert _context_window("gpt") is None - - -class TestTableConsistency: - """Sanity checks on the registry.""" - - def test_all_context_windows_positive(self) -> None: - for model, pricing in DEFAULT_PRICING.items(): - if pricing.context_window is not None: - assert pricing.context_window > 0, f"{model} has non-positive context window" - - def test_all_have_context_window(self) -> None: - for model, pricing in DEFAULT_PRICING.items(): - assert pricing.context_window is not None, f"{model} is missing context_window" diff --git a/tests/test_providers/test_copilot.py b/tests/test_providers/test_copilot.py index e1d0664e..cc90515e 100644 --- a/tests/test_providers/test_copilot.py +++ b/tests/test_providers/test_copilot.py @@ -695,3 +695,95 @@ def test_runs_on_unix(self, monkeypatch: pytest.MonkeyPatch) -> None: # early — it proceeded past the guard and attempted the import. with contextlib.suppress(ModuleNotFoundError): provider._fix_pipe_blocking_mode() + + +class TestGetMaxPromptTokens: + """Tests for CopilotProvider.get_max_prompt_tokens.""" + + @staticmethod + def _make_model(model_id: str, max_prompt_tokens: int) -> Any: + from types import SimpleNamespace + + return SimpleNamespace( + id=model_id, + capabilities=SimpleNamespace( + limits=SimpleNamespace(max_prompt_tokens=max_prompt_tokens) + ), + ) + + @staticmethod + def _provider_with_list_models(list_models_impl: Any) -> CopilotProvider: + """Build a provider with the SDK short-circuit disabled and a fake client. + + Uses ``stub_handler`` for ``mock_handler`` to satisfy the constructor, + then nulls ``_mock_handler`` so ``get_max_prompt_tokens`` falls through + to the SDK path. ``_started=True`` skips ``_ensure_client_started``. + """ + from types import SimpleNamespace + + provider = CopilotProvider(mock_handler=stub_handler) + provider._mock_handler = None + provider._client = SimpleNamespace(list_models=list_models_impl) + provider._started = True + return provider + + @pytest.mark.asyncio + async def test_mock_handler_mode_returns_none(self) -> None: + """Mock-handler mode has no SDK to query — must return None.""" + provider = CopilotProvider(mock_handler=stub_handler) + assert await provider.get_max_prompt_tokens("gpt-4o") is None + + @pytest.mark.asyncio + async def test_returns_max_prompt_tokens_for_known_model(self) -> None: + async def list_models() -> list[Any]: + return [self._make_model("gpt-4o", 128000)] + + provider = self._provider_with_list_models(list_models) + assert await provider.get_max_prompt_tokens("gpt-4o") == 128000 + + @pytest.mark.asyncio + async def test_returns_none_for_unknown_model(self) -> None: + async def list_models() -> list[Any]: + return [] + + provider = self._provider_with_list_models(list_models) + assert await provider.get_max_prompt_tokens("anything") is None + + @pytest.mark.asyncio + async def test_oserror_returns_none_and_does_not_cache(self) -> None: + """A transport-level error is swallowed; the next call retries.""" + calls = 0 + + async def list_models() -> list[Any]: + nonlocal calls + calls += 1 + if calls == 1: + raise OSError("network down") + return [self._make_model("gpt-4o", 128000)] + + provider = self._provider_with_list_models(list_models) + assert await provider.get_max_prompt_tokens("gpt-4o") is None + assert await provider.get_max_prompt_tokens("gpt-4o") == 128000 + + @pytest.mark.asyncio + async def test_unexpected_exception_propagates(self) -> None: + """Non-SDK exceptions (programming errors) are not swallowed by the + provider — they bubble up so the engine's outer safety net handles them.""" + + async def list_models() -> list[Any]: + raise ValueError("bug") + + provider = self._provider_with_list_models(list_models) + with pytest.raises(ValueError): + await provider.get_max_prompt_tokens("gpt-4o") + + @pytest.mark.asyncio + async def test_alias_resolves_via_match_model_id(self) -> None: + """Versioned-suffix aliases resolve to the SDK's listed ID.""" + + async def list_models() -> list[Any]: + return [self._make_model("claude-3-5-sonnet", 200_000)] + + provider = self._provider_with_list_models(list_models) + assert await provider.get_max_prompt_tokens("claude-3-5-sonnet-latest") == 200_000 + assert await provider.get_max_prompt_tokens("claude-3-5-sonnet-20241022") == 200_000