diff --git a/docs/ext-proc-gateway.md b/docs/ext-proc-gateway.md index 89f0ce43..879d5baf 100644 --- a/docs/ext-proc-gateway.md +++ b/docs/ext-proc-gateway.md @@ -168,6 +168,7 @@ and the pass-through fallback — is identical. - **Per-request configuration.** In Approach 1, all ITS parameters travel in `X-ITS-*` headers; no prior configuration is required. In Approach 2, the upstream LLM is configured once via `POST /configure` and each request supplies only its `budget` in the request body; `X-ITS-*` headers are also accepted - as per-request overrides (header > body > `/configure` default). + as per-request overrides. A field is applied if it is not `None`; otherwise the next tier down + (header > body > `/configure` default) supplies the value. - **Ports.** The values shown are defaults and may be changed: `:8108` (Envoy), `:50051` (ext_proc gRPC), `:8109` (IaaS service), and `:8100` (the upstream LLM / `llm_upstream`). diff --git a/docs/iaas-service.md b/docs/iaas-service.md index 9fed5582..8d1fbf3c 100644 --- a/docs/iaas-service.md +++ b/docs/iaas-service.md @@ -34,8 +34,9 @@ used. | Request body | `budget` | Compute budget | 2 | | `/configure` | `budget`, `endpoint`, `api_key` | Service defaults | 3 (lowest) | -Priority chain: **header > body > service default**. Headers are intended for Envoy -ext_proc routing but are also accepted on the standalone IaaS endpoint. +Priority chain: **header > body > service default**. A field is applied if it is +not `None`; otherwise the next tier down supplies the value. Headers are intended for +Envoy ext_proc routing but are also accepted on the standalone IaaS endpoint. ### Algorithm Selection @@ -57,6 +58,12 @@ tool calls using the `tool_hierarchical` strategy, and text responses fall back exact-content matching. Supply `regex_patterns` to vote on extracted text answers, or set `tool_vote` to pick a different tool-voting strategy. +> **Note on clearing optional fields.** Because `None` means "use the tier below," a +> `/configure` call cannot explicitly reset `tool_vote` back to its default once set. +> To switch voting strategies, supply the new value; to revert to the built-in default +> (`tool_hierarchical`), pass `"tool_hierarchical"` explicitly. `regex_patterns` and +> `exclude_tool_args` can be cleared by passing an empty list (`[]`). + ## API Key Handling API keys can enter the system through three paths: @@ -69,11 +76,12 @@ API keys can enter the system through three paths: **Security properties:** - Keys are **never logged** — the gateway logs endpoint and model but not credentials -- Keys are **hashed** (SHA-256, truncated to 16 hex chars) in LM cache keys to prevent - credential cross-contamination between requests using different API keys -- Keys are **not persisted** to disk — they exist only in memory for the lifetime of the - service process -- On shutdown, all cached LM clients (and their associated keys) are cleared +- Keys are **not persisted** to disk. A key supplied via the `X-ITS-API-Key` header is + request-scoped: it lives only for that request's LM client, which is closed when the + request completes. A key set via `/configure` remains in memory as a service default + (`ITSGateway._default_config`) until replaced or the process restarts. +- Keys are **never shared between requests** — each request builds its own LM client, so + credentials supplied via header or `/configure` cannot cross-contaminate ## Prerequisites @@ -181,7 +189,7 @@ All configurations support: - `endpoint`: OpenAI-compatible API endpoint URL - `api_key`: API key for the provider - `model`: Model identifier -- `alg`: Algorithm name - `"self-consistency"` or `"best-of-n"` +- `alg`: Algorithm name - `"self-consistency"`, `"adaptive-self-consistency"`, or `"beta-self-consistency"` ## Usage Examples @@ -429,15 +437,15 @@ arrives as a burst rather than incrementally. ## Restart and Scaling -- **State**: The service holds an in-memory LM client cache (LRU, default 64 entries), - a gateway instance, and the service config set via `/configure`. No state is persisted +- **State**: The service holds a gateway instance and the service config set via + `/configure`. LM clients are created and discarded per request. No state is persisted to disk. -- **Restart**: Restarting clears all cached LM clients and the service config. - `/configure` must be called again after restart. +- **Restart**: Restarting clears the service config. `/configure` must be called again + after restart. - **Horizontal scaling**: Multiple IaaS instances can run independently. Each maintains - its own LM client cache, config, and gateway. There is no shared state between - instances. A load balancer must route `/configure` to all instances or each instance - must be configured independently. + its own config and gateway. There is no shared state between instances. A load + balancer must route `/configure` to all instances or each instance must be configured + independently. ## API Endpoints @@ -450,7 +458,6 @@ arrives as a burst rather than incrementally. ### Health Check - `GET /docs` - API documentation -- `GET /health` - Service health (if available) ## Troubleshooting @@ -489,7 +496,7 @@ curl -X GET http://localhost:8109/docs **5. Slow Responses** - This is expected behavior for inference-time scaling - Reduce `budget` parameter for faster responses -- Best-of-N with budget=4 typically takes 30-60 seconds +- Self-consistency with budget=4 typically takes 30-60 seconds ### Log Files diff --git a/its_hub/api/__init__.py b/its_hub/api/__init__.py index 3bc1725b..b7ee20c6 100644 --- a/its_hub/api/__init__.py +++ b/its_hub/api/__init__.py @@ -24,7 +24,15 @@ from .orchestrator import AbstractOrchestrator from .reward_models.orm import AbstractOutcomeRewardModel from .reward_models.prm import AbstractProcessRewardModel -from .types import ChatMessage, ChatMessages, GenerationUsage, ITSRequestConfig +from .types import ( + SUPPORTED_ALGORITHMS, + VALID_TOOL_VOTE_OPTIONS, + ChatMessage, + ChatMessages, + GenerationUsage, + ITSRequestConfig, + ITSRequestConfigUpdate, +) __all__ = [ # noqa: RUF022 # Algorithm abstractions @@ -44,6 +52,9 @@ "ChatMessages", "GenerationUsage", "ITSRequestConfig", + "ITSRequestConfigUpdate", + "SUPPORTED_ALGORITHMS", + "VALID_TOOL_VOTE_OPTIONS", # Error types "APIError", "RateLimitError", diff --git a/its_hub/api/gateway.py b/its_hub/api/gateway.py index 0582bc8d..2ed7db66 100644 --- a/its_hub/api/gateway.py +++ b/its_hub/api/gateway.py @@ -6,7 +6,7 @@ from abc import ABC, abstractmethod from typing import Any -from its_hub.api.types import ITSRequestConfig +from its_hub.api.types import ITSRequestConfigUpdate class AbstractGateway(ABC): @@ -21,7 +21,7 @@ class AbstractGateway(ABC): @abstractmethod async def arun_chat_completion( self, - config: ITSRequestConfig, + config: ITSRequestConfigUpdate, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None, tool_choice: str | dict[str, Any] | None = None, @@ -31,7 +31,8 @@ async def arun_chat_completion( """Run chat completion with ITS algorithm asynchronously. Args: - config: Per-request ITS configuration + config: Per-request ITS configuration overlay (merged over the + gateway's service default at run time) messages: OpenAI-format conversation messages tools: Optional tool definitions for function calling tool_choice: Optional tool choice strategy @@ -46,14 +47,9 @@ async def arun_chat_completion( def run_chat_completion( self, - config: ITSRequestConfig, + config: ITSRequestConfigUpdate, messages: list[dict[str, Any]], **kwargs, ) -> dict[str, Any]: """Synchronous wrapper for arun_chat_completion.""" return asyncio.run(self.arun_chat_completion(config, messages, **kwargs)) - - @abstractmethod - async def ashutdown(self) -> None: - """Cleanup resources on service shutdown.""" - pass diff --git a/its_hub/api/types.py b/its_hub/api/types.py index c0203208..ef6b0980 100644 --- a/its_hub/api/types.py +++ b/its_hub/api/types.py @@ -3,7 +3,8 @@ from __future__ import annotations import logging -from dataclasses import dataclass, fields +import re +from dataclasses import dataclass, fields, replace from typing import Literal @@ -18,7 +19,7 @@ class ChatMessage: """ role: Literal["system", "user", "assistant", "tool"] - content: str | list[dict] | None + content: str | list[dict] | None = None tool_calls: list[dict] | None = None # Store as plain dicts tool_call_id: str | None = None @@ -164,26 +165,199 @@ def total_tokens(self) -> int: return self.prompt_tokens + self.completion_tokens -@dataclass +SELF_CONSISTENCY_ALGORITHMS = frozenset( + {"self-consistency", "adaptive-self-consistency", "beta-self-consistency"} +) + +# All currently supported algorithms are self-consistency variants. +SUPPORTED_ALGORITHMS = SELF_CONSISTENCY_ALGORITHMS + +# Valid values for ITSRequestConfig.tool_vote / SelfConsistency.tool_vote. +VALID_TOOL_VOTE_OPTIONS = frozenset( + {"tool_name", "tool_args", "tool_hierarchical", "tool_flat_all"} +) + + +def _validate_optional_fields( + budget: int | None, + alg: str | None, + regex_patterns: list[str] | None, + tool_vote: str | None, + threshold: float | None, + confidence_threshold: float | None, + temperature: float | None, +) -> None: + """Validate the optional scaling fields. Shared by ITSRequestConfig and + ITSRequestConfigUpdate so the two cannot drift. Each check is guarded so + that ``None`` ("not supplied") is always accepted.""" + if budget is not None and not (1 <= budget <= 1000): + raise ValueError("budget must be between 1 and 1000") + if alg is not None and alg not in SUPPORTED_ALGORITHMS: + raise ValueError( + f"Algorithm {alg!r} not supported. Choose from: {SUPPORTED_ALGORITHMS}" + ) + if regex_patterns is not None: + for p in regex_patterns: + try: + re.compile(p) + except re.error as e: + raise ValueError(f"Invalid regex pattern {p!r}: {e}") from e + if threshold is not None and not (0.5 < threshold <= 1.0): + raise ValueError(f"threshold must be in (0.5, 1.0], got: {threshold}") + if confidence_threshold is not None and not (0.5 < confidence_threshold <= 1.0): + raise ValueError( + f"confidence_threshold must be in (0.5, 1.0], got: {confidence_threshold}" + ) + if temperature is not None and not (0.0 <= temperature <= 2.0): + raise ValueError(f"temperature must be in [0.0, 2.0], got: {temperature}") + if tool_vote is not None and tool_vote not in VALID_TOOL_VOTE_OPTIONS: + raise ValueError( + f"tool_vote must be one of {VALID_TOOL_VOTE_OPTIONS}, got: {tool_vote}" + ) + + +def _freeze_list(value: list[str] | None) -> tuple[str, ...] | None: + """Copy a list into an immutable tuple; leave ``None`` unchanged.""" + if value is None: + return None + return tuple(value) + + +def _config_repr(obj: ITSRequestConfig | ITSRequestConfigUpdate) -> str: + return ( + type(obj).__name__ + + "(" + + ", ".join( + f"{f.name}={'***' if f.name == 'api_key' else getattr(obj, f.name)!r}" + for f in fields(obj) + ) + + ")" + ) + + +@dataclass(frozen=True) class ITSRequestConfig: - """Per-request configuration for ITS execution. + """A fully-resolved configuration snapshot, ready to drive one request. - Can be constructed incrementally — headers provide budget/endpoint/api_key, - then model is set from the request body. + ``api_endpoint`` and ``model`` are mandatory — constructing without them + raises. The remaining fields carry system defaults (``budget``, ``alg``, + ``api_key``) or are genuinely optional (``None`` is a meaningful "not set", + e.g. ``temperature=None`` means "don't send it upstream"). """ - budget: int + # LM target api_endpoint: str + model: str + api_key: str | None = None + temperature: float | None = None + # Scaling parameters + budget: int = 4 + alg: str = "self-consistency" + regex_patterns: tuple[str, ...] | None = None + tool_vote: str | None = None + exclude_tool_args: tuple[str, ...] | None = None + threshold: float | None = None + confidence_threshold: float | None = None + + def __post_init__(self): + if not self.api_endpoint: + raise ValueError("api_endpoint must be specified") + if not self.model: + raise ValueError("model must be specified") + _validate_optional_fields( + self.budget, + self.alg, + self.regex_patterns, + self.tool_vote, + self.threshold, + self.confidence_threshold, + self.temperature, + ) + # frozen=True blocks attribute rebinding, not mutation of a referenced + # list. Copy these into tuples so neither the source list nor the + # snapshot can be mutated after validation to affect an in-flight request. + object.__setattr__(self, "regex_patterns", _freeze_list(self.regex_patterns)) + object.__setattr__( + self, "exclude_tool_args", _freeze_list(self.exclude_tool_args) + ) + + def __repr__(self) -> str: + return _config_repr(self) + + +@dataclass +class ITSRequestConfigUpdate: + """A partial configuration overlay. + + Every field is ``Optional``: ``None`` means "not supplied — keep the value + below in the merge hierarchy". Because of this, ``tool_vote`` cannot be + cleared back to its default once set (pass the new value, or + ``"tool_hierarchical"`` explicitly); ``regex_patterns`` and + ``exclude_tool_args`` can be cleared by passing ``[]``. + + Use :meth:`merge` to layer one update over another, and :meth:`resolve` to + materialize a complete :class:`ITSRequestConfig`; mandatory-field presence + is enforced by :class:`ITSRequestConfig` at construction time. + """ + + # LM target + api_endpoint: str | None = None model: str | None = None api_key: str | None = None + temperature: float | None = None + # Scaling parameters + budget: int | None = None + alg: str | None = None + regex_patterns: list[str] | None = None + tool_vote: str | None = None + exclude_tool_args: list[str] | None = None + threshold: float | None = None + confidence_threshold: float | None = None def __post_init__(self): - if self.budget < 1 or self.budget > 1000: - raise ValueError("budget must be between 1 and 1000") + _validate_optional_fields( + self.budget, + self.alg, + self.regex_patterns, + self.tool_vote, + self.threshold, + self.confidence_threshold, + self.temperature, + ) def __repr__(self) -> str: - return ( - f"ITSRequestConfig(budget={self.budget}, " - f"api_endpoint='{self.api_endpoint}', " - f"model={self.model!r}, api_key={'***' if self.api_key else None})" + return _config_repr(self) + + def merge(self, overlay: ITSRequestConfigUpdate) -> ITSRequestConfigUpdate: + """Return a copy of ``self`` with non-``None`` fields of ``overlay`` + applied. ``replace()`` re-fires ``__post_init__``, so the result is + format-validated.""" + return replace( + self, + **{ + f.name: getattr(overlay, f.name) + for f in fields(overlay) + if getattr(overlay, f.name) is not None + }, + ) + + def resolve(self) -> ITSRequestConfig: + """Materialize a complete :class:`ITSRequestConfig`. + + Optional fields are forwarded only when non-``None``; otherwise + the resolved config's defaults (``budget=4``, + ``alg="self-consistency"``, ``api_key=None``) apply. Mandatory-field + presence (``api_endpoint``, ``model``) is enforced by + :class:`ITSRequestConfig` itself at construction time. + """ + optionals = { + f.name: getattr(self, f.name) + for f in fields(self) + if f.name not in {"api_endpoint", "model"} + and getattr(self, f.name) is not None + } + return ITSRequestConfig( + api_endpoint=self.api_endpoint, + model=self.model, + **optionals, ) diff --git a/its_hub/core/algorithms/self_consistency.py b/its_hub/core/algorithms/self_consistency.py index 638884f5..a4fde46b 100644 --- a/its_hub/core/algorithms/self_consistency.py +++ b/its_hub/core/algorithms/self_consistency.py @@ -8,6 +8,7 @@ from dataclasses import dataclass from its_hub.api import ( + VALID_TOOL_VOTE_OPTIONS, AbstractLanguageModel, AbstractOrchestrator, AbstractScalingAlgorithm, @@ -170,16 +171,10 @@ def __init__( ValueError: If tool_vote is not one of the supported options. """ # Validate tool_vote parameter - only validation needed since typing handles the rest - valid_tool_vote_options = { - None, - "tool_name", - "tool_args", - "tool_hierarchical", - "tool_flat_all", - } - if tool_vote not in valid_tool_vote_options: + valid_options = VALID_TOOL_VOTE_OPTIONS | {None} + if tool_vote not in valid_options: raise ValueError( - f"tool_vote must be one of {valid_tool_vote_options}, got: {tool_vote}" + f"tool_vote must be one of {valid_options}, got: {tool_vote}" ) # Set default projection function if provided None self.consistency_space_projection_func = ( diff --git a/its_hub/core/gateway.py b/its_hub/core/gateway.py index f5dec4ce..bcafb7b0 100644 --- a/its_hub/core/gateway.py +++ b/its_hub/core/gateway.py @@ -4,11 +4,14 @@ similar to how HTTP services reuse client instances across requests. """ -import hashlib import logging +import ssl import time -from collections import OrderedDict from typing import Any +from urllib.parse import urlparse + +import aiohttp +import certifi from its_hub.api import ( AbstractGateway, @@ -18,32 +21,19 @@ ChatMessages, GenerationUsage, ITSRequestConfig, + ITSRequestConfigUpdate, ) from its_hub.core.algorithms.adaptive_self_consistency import AdaptiveSelfConsistency from its_hub.core.algorithms.beta_self_consistency import BetaSelfConsistency from its_hub.core.algorithms.self_consistency import ( SelfConsistency, create_regex_projection_function, - validate_regex_patterns, ) from its_hub.core.lms.openai_lm import OpenAICompatibleLanguageModel from its_hub.core.orchestrator import LMOrchestrator logger = logging.getLogger(__name__) -# The self-consistency family shares the same voting surface (regex/tool-vote -# projection) and only differs in how it decides when to stop sampling. -SELF_CONSISTENCY_ALGORITHMS = frozenset( - { - "self-consistency", - "adaptive-self-consistency", - "beta-self-consistency", - } -) - -# All currently supported algorithms are self-consistency variants. -SUPPORTED_ALGORITHMS = SELF_CONSISTENCY_ALGORITHMS - class ITSGateway(AbstractGateway): """Long-lived gateway for running ITS algorithms. @@ -51,16 +41,12 @@ class ITSGateway(AbstractGateway): Initialized once at service startup and reused across all requests. Per-request configuration is passed as arguments to arun_chat_completion(). - Maintains a cache of LM clients keyed by (endpoint, model, hashed_api_key) - to reuse HTTP connections across requests while preventing credential - cross-contamination. - Example: # At service startup gateway = ITSGateway() # Per request - config = ITSRequestConfig( + config = ITSRequestConfigUpdate( budget=10, api_endpoint="http://envoy-cluster/v1", model="gpt-4", @@ -69,145 +55,104 @@ class ITSGateway(AbstractGateway): result = await gateway.arun_chat_completion(config, messages) """ - _DEFAULT_MAX_LM_CACHE_SIZE = 64 - def __init__( self, - algorithm: AbstractScalingAlgorithm | None = None, orchestrator: AbstractOrchestrator | None = None, - max_lm_cache_size: int = _DEFAULT_MAX_LM_CACHE_SIZE, + default_config: ITSRequestConfigUpdate | None = None, ): if orchestrator is None: orchestrator = LMOrchestrator() self._orchestrator = orchestrator - - if algorithm is None: - algorithm = SelfConsistency(orchestrator=orchestrator) - self._algorithm = algorithm - - self._algorithm_name = type(algorithm).__name__ - self._max_lm_cache_size = max_lm_cache_size - self._lm_cache: OrderedDict[ - tuple[str, str, str], OpenAICompatibleLanguageModel - ] = OrderedDict() - logger.info("ITSGateway initialized with algorithm=%s", self._algorithm_name) - - def configure( - self, - alg: str, - regex_patterns: list[str] | None = None, - tool_vote: str | None = None, - exclude_tool_args: list[str] | None = None, - threshold: float | None = None, - confidence_threshold: float | None = None, - ) -> None: - """Configure the gateway's scaling algorithm at runtime. - - ``threshold`` applies only to ``adaptive-self-consistency`` and - ``confidence_threshold`` only to ``beta-self-consistency``; both fall - back to the algorithm's own default when left as ``None``. - - Raises ValueError for unsupported algorithms or invalid options. + # Stored as a partial update (format-validated, never completeness- + # validated): endpoint/model may be absent until a request supplies + # them. System defaults (budget=4, alg="self-consistency") are NOT + # stored here — they live on ITSRequestConfig and are injected at + # resolve() time. + self._default_config = default_config or ITSRequestConfigUpdate() + self._ssl_context = ssl.create_default_context(cafile=certifi.where()) + # One connection pool per upstream endpoint + self._connector_pool: dict[str, aiohttp.TCPConnector] = {} + logger.info("ITSGateway initialized") + + def configure(self, update: ITSRequestConfigUpdate) -> None: + """Merge ``update`` into the service default. + + Format-validated (via ``ITSRequestConfigUpdate.__post_init__`` re-fired + by ``merge``); not completeness-validated — ``api_endpoint``/``model`` + may remain absent until a request supplies them. """ - if alg not in SUPPORTED_ALGORITHMS: - raise ValueError( - f"Algorithm {alg!r} not supported. Choose from: {SUPPORTED_ALGORITHMS}" - ) + self._default_config = self._default_config.merge(update) + logger.info( + "ITSGateway reconfigured with default alg=%s", self._default_config.alg + ) + + def _build_algorithm(self, config: ITSRequestConfig) -> AbstractScalingAlgorithm: + """Construct a fresh scaling algorithm from a config snapshot.""" + alg = config.alg - # The self-consistency family shares projection/tool-vote plumbing. projection_func = None - if regex_patterns: - validate_regex_patterns(regex_patterns) - projection_func = create_regex_projection_function(regex_patterns) + if config.regex_patterns: + projection_func = create_regex_projection_function(config.regex_patterns) common_kwargs = { "consistency_space_projection_func": projection_func, - "exclude_args": exclude_tool_args, + "exclude_args": config.exclude_tool_args, "orchestrator": self._orchestrator, } # Only override the algorithm's default tool_vote when one is explicitly # provided; otherwise let it fall back to DEFAULT_TOOL_VOTE so that # tool-calling responses still vote sensibly with no extra config. - if tool_vote is not None: - common_kwargs["tool_vote"] = tool_vote + if config.tool_vote is not None: + common_kwargs["tool_vote"] = config.tool_vote if alg == "adaptive-self-consistency": - extra = {} if threshold is None else {"threshold": threshold} - algorithm = AdaptiveSelfConsistency(**common_kwargs, **extra) + extra = {} if config.threshold is None else {"threshold": config.threshold} + return AdaptiveSelfConsistency(**common_kwargs, **extra) elif alg == "beta-self-consistency": extra = ( {} - if confidence_threshold is None - else {"confidence_threshold": confidence_threshold} + if config.confidence_threshold is None + else {"confidence_threshold": config.confidence_threshold} ) - algorithm = BetaSelfConsistency(**common_kwargs, **extra) + return BetaSelfConsistency(**common_kwargs, **extra) else: # self-consistency - algorithm = SelfConsistency(**common_kwargs) - - self._algorithm = algorithm - self._algorithm_name = type(algorithm).__name__ - logger.info("ITSGateway reconfigured with algorithm=%s", self._algorithm_name) - - @staticmethod - def _hash_api_key(api_key: str | None) -> str: - return hashlib.sha256((api_key or "").encode()).hexdigest()[:16] - - async def _get_or_create_lm( + return SelfConsistency(**common_kwargs) + + def _get_connector(self, endpoint: str) -> aiohttp.TCPConnector: + """Get or create the pooled connector.""" + parsed = urlparse(endpoint) + origin = f"{parsed.scheme}://{parsed.netloc}" + conn = self._connector_pool.get(origin) + if conn is None: + conn = aiohttp.TCPConnector(ssl=self._ssl_context) + self._connector_pool[origin] = conn + return conn + + def _build_lm( self, - endpoint: str, - model: str, - api_key: str | None = None, + config: ITSRequestConfig, request_id: str | None = None, ) -> OpenAICompatibleLanguageModel: - """Get cached LM client or create new one. - - Clients are cached by (endpoint, model, hashed_api_key) to reuse - HTTP connections while isolating different credentials. - """ - cache_key = (endpoint, model, self._hash_api_key(api_key)) + """Construct a one-shot LM client from a config snapshot.""" log_prefix = f"[{request_id}] " if request_id else "" - - if cache_key in self._lm_cache: - self._lm_cache.move_to_end(cache_key) - logger.debug( - "%sReusing cached LM client: endpoint=%s, model=%s", - log_prefix, - endpoint, - model, - ) - return self._lm_cache[cache_key] - - evicted_lm = None - if len(self._lm_cache) >= self._max_lm_cache_size: - evicted_key, evicted_lm = self._lm_cache.popitem(last=False) - logger.info( - "%sEvicting LM client from cache: endpoint=%s, model=%s", - log_prefix, - evicted_key[0], - evicted_key[1], - ) - logger.info( - "%sCreating new LM client: endpoint=%s, model=%s", + "%sCreating LM client: endpoint=%s, model=%s", log_prefix, - endpoint, - model, + config.api_endpoint, + config.model, ) - lm = OpenAICompatibleLanguageModel( - endpoint=endpoint, - api_key=api_key or "", - model_name=model, + return OpenAICompatibleLanguageModel( + endpoint=config.api_endpoint, + api_key=config.api_key, + model_name=config.model, + temperature=config.temperature, + ssl_context=self._ssl_context, + connector=self._get_connector(config.api_endpoint), ) - self._lm_cache[cache_key] = lm - - if evicted_lm is not None: - await evicted_lm.close() - - return lm async def arun_chat_completion( self, - config: ITSRequestConfig, + config: ITSRequestConfigUpdate, messages: list[dict[str, Any]], tools: list[dict[str, Any]] | None = None, tool_choice: str | dict[str, Any] | None = None, @@ -217,42 +162,41 @@ async def arun_chat_completion( request_id = kwargs.get("request_id") log_prefix = f"[{request_id}] " if request_id else "" - if not config.api_endpoint: - raise ValueError("api_endpoint must be specified in ITSRequestConfig") - if not config.model: - raise ValueError( - "Model must be specified in ITSRequestConfig before running" - ) + # Merge the per-request overlay over the service default, then resolve + # to a complete snapshot. resolve() is the single completeness check + # (raises ValueError if api_endpoint/model are still absent). Nothing + # on `self` is mutated, so a concurrent /configure cannot swap the + # algorithm or close this request's HTTP session mid-flight. + resolved = self._default_config.merge(config).resolve() - lm = await self._get_or_create_lm( - endpoint=config.api_endpoint, - model=config.model, - api_key=config.api_key, - request_id=request_id, - ) + algorithm = self._build_algorithm(resolved) + lm = self._build_lm(resolved, request_id) chat_messages = ChatMessages([ChatMessage.from_dict(msg) for msg in messages]) logger.info( "%sRunning ITS: algorithm=%s, budget=%s, endpoint=%s, model=%s, messages=%s, tools=%s", log_prefix, - self._algorithm_name, - config.budget, - config.api_endpoint, - config.model, + type(algorithm).__name__, + resolved.budget, + resolved.api_endpoint, + resolved.model, len(messages), "yes" if tools else "no", ) t0 = time.monotonic() - result = await self._algorithm.ainfer( - lm=lm, - prompt_or_messages=chat_messages, - budget=config.budget, - return_response_only=False, - tools=tools, - tool_choice=tool_choice, - ) + try: + result = await algorithm.ainfer( + lm=lm, + prompt_or_messages=chat_messages, + budget=resolved.budget, + return_response_only=False, + tools=tools, + tool_choice=tool_choice, + ) + finally: + await lm.close() usage_dict = {} if isinstance(result.usage, GenerationUsage): @@ -272,7 +216,7 @@ async def arun_chat_completion( duration_s, usage_dict, ) - return {"message": result.the_one, "usage": usage_dict} + return {"message": result.the_one, "usage": usage_dict, "alg": resolved.alg} logger.info( "%sITS completed in %.2fs. Usage: %s (selected_index=%s)", @@ -287,20 +231,22 @@ async def arun_chat_completion( "selected_index": result.selected_index, "the_one": result.the_one, "usage": usage_dict, + "alg": resolved.alg, } - async def ashutdown(self) -> None: - """Cleanup resources on service shutdown. + async def aclose(self) -> None: + """Close all pooled connectors. - Closes all cached LM clients (releasing aiohttp sessions) then - clears the cache. + Call on service shutdown to release keep-alive connections. + Idempotent: closing an already-closed connector is a no-op. """ logger.info( - "ITSGateway shutting down, closing %d LM clients", - len(self._lm_cache), + "ITSGateway shutting down, closing %d pooled connectors", + len(self._connector_pool), ) - for lm in self._lm_cache.values(): - await lm.close() - self._lm_cache.clear() + for conn in self._connector_pool.values(): + if not conn.closed: + await conn.close() + self._connector_pool.clear() if hasattr(self._orchestrator, "shutdown"): self._orchestrator.shutdown() diff --git a/its_hub/core/lms/openai_lm.py b/its_hub/core/lms/openai_lm.py index 70c83d64..5125b9b0 100644 --- a/its_hub/core/lms/openai_lm.py +++ b/its_hub/core/lms/openai_lm.py @@ -27,7 +27,7 @@ class OpenAICompatibleLanguageModel(AbstractLanguageModel): def __init__( self, endpoint: str, - api_key: str, + api_key: str | None, model_name: str, system_prompt: str | None = None, is_async: bool = False, # Deprecated: parameter is ignored (always async internally) @@ -44,6 +44,8 @@ def __init__( ssl_context: ssl.SSLContext | None = None, # Raw response preservation include_raw_choices: bool = False, + # Borrowed connection pool (e.g. from ITSGateway); not closed by this LM + connector: aiohttp.BaseConnector | None = None, ): assert max_concurrency == -1 or max_concurrency > 0, ( "max_concurrency must be -1 (unlimited concurrency) or a positive integer" @@ -86,15 +88,13 @@ def __init__( self.ssl_context.check_hostname = False self.ssl_context.verify_mode = ssl.CERT_NONE else: - # For async requests, create SSL context using the same CA bundle as requests - # This ensures aiohttp uses the same certificates as requests library self.ssl_context = ssl.create_default_context(cafile=certifi.where()) - # set up headers for API requests - self.headers = { - "Content-Type": "application/json", - "Authorization": f"Bearer {self.api_key}", - } + # set up headers for API requests; omit Authorization when no API key is + # configured (e.g. a local no-auth vLLM) rather than sending an empty token. + self.headers = {"Content-Type": "application/json"} + if self.api_key: + self.headers["Authorization"] = f"Bearer {self.api_key}" # raw response preservation self.include_raw_choices = include_raw_choices @@ -104,6 +104,7 @@ def __init__( # Session cache: one session per event loop, entry is auto-cleaned via weak references. # Session(s) need to be closed via close or close_session function calls. + self._borrowed_connector = connector self._sessions: weakref.WeakKeyDictionary[ asyncio.AbstractEventLoop, aiohttp.ClientSession ] = weakref.WeakKeyDictionary() @@ -124,10 +125,16 @@ def _get_session(self, loop: asyncio.AbstractEventLoop) -> aiohttp.ClientSession if session is not None and not session.closed: return session - # Create new session for the event loop. - # trust_env=True so aiohttp honors HTTP(S)_PROXY/NO_PROXY/.netrc. - connector = aiohttp.TCPConnector(ssl=self.ssl_context) - session = aiohttp.ClientSession(connector=connector, trust_env=True) + if self._borrowed_connector is not None: + # Borrowed pool: the session uses it but must not close it. + session = aiohttp.ClientSession( + connector=self._borrowed_connector, + connector_owner=False, + trust_env=True, + ) + else: + connector = aiohttp.TCPConnector(ssl=self.ssl_context) + session = aiohttp.ClientSession(connector=connector, trust_env=True) self._sessions[loop] = session return session diff --git a/its_hub/integration/ext_proc/processor.py b/its_hub/integration/ext_proc/processor.py index 8233f71a..5ebdf587 100644 --- a/its_hub/integration/ext_proc/processor.py +++ b/its_hub/integration/ext_proc/processor.py @@ -18,7 +18,7 @@ from envoy.service.ext_proc.v3 import external_processor_pb2_grpc as ext_proc_grpc from envoy.type.v3 import http_status_pb2 -from its_hub.api import ChatMessage, ITSRequestConfig +from its_hub.api import ChatMessage, ITSRequestConfigUpdate from its_hub.core.gateway import ITSGateway _ITS_HEADER_PREFIX = "x-its-" @@ -185,7 +185,7 @@ async def Process( # noqa: N802 — gRPC servicer interface requires this name async def _apply_its( self, - its_config: ITSRequestConfig, + its_config: ITSRequestConfigUpdate, body: bytes, request_id: str, ) -> ext_proc_pb2.ProcessingResponse: @@ -324,7 +324,9 @@ def _headers_continue( ) ) - def _parse_its_headers(self, headers: dict[str, str]) -> ITSRequestConfig | None: + def _parse_its_headers( + self, headers: dict[str, str] + ) -> ITSRequestConfigUpdate | None: """Parse ITS configuration from request headers. Model is NOT extracted from headers - it will be set from request body later. @@ -339,7 +341,7 @@ def _parse_its_headers(self, headers: dict[str, str]) -> ITSRequestConfig | None budget = int(budget_str) api_key = headers.get("x-its-api-key") - return ITSRequestConfig( + return ITSRequestConfigUpdate( budget=budget, api_endpoint=endpoint, api_key=api_key, @@ -349,8 +351,3 @@ def _parse_its_headers(self, headers: dict[str, str]) -> ITSRequestConfig | None except (ValueError, TypeError) as e: logger.error("Failed to parse ITS headers: %s", e) return None - - async def shutdown(self): - """Cleanup resources on service shutdown.""" - logger.info("External Processor shutting down") - await self.gateway.ashutdown() diff --git a/its_hub/integration/ext_proc/server.py b/its_hub/integration/ext_proc/server.py index 2f2eca88..82bfc12f 100644 --- a/its_hub/integration/ext_proc/server.py +++ b/its_hub/integration/ext_proc/server.py @@ -82,8 +82,9 @@ async def serve(port: int = 50051): await server.wait_for_termination() except (KeyboardInterrupt, asyncio.CancelledError): logger.info("Received shutdown signal") - await processor.shutdown() await server.stop(grace=5) + finally: + await processor.gateway.aclose() def _print_config() -> None: diff --git a/its_hub/integration/iaas/app.py b/its_hub/integration/iaas/app.py index 577ec485..b007efaa 100644 --- a/its_hub/integration/iaas/app.py +++ b/its_hub/integration/iaas/app.py @@ -10,13 +10,11 @@ import time import uuid from contextlib import asynccontextmanager -from dataclasses import dataclass from fastapi import FastAPI, Header, HTTPException, status from fastapi.responses import StreamingResponse -from its_hub.api.types import ITSRequestConfig -from its_hub.core.algorithms.self_consistency import SelfConsistency +from its_hub.api.types import ITSRequestConfigUpdate from its_hub.core.gateway import ITSGateway from its_hub.integration.iaas.models import ( ChatCompletionChoice, @@ -29,28 +27,14 @@ logger = logging.getLogger(__name__) -@dataclass -class _ServiceConfig: - """Mutable service-level defaults set via /configure.""" - - endpoint: str = "" - model: str = "" - api_key: str | None = None - budget: int = 4 - temperature: float | None = None - alg: str = "self-consistency" - - class _ServiceState: """Encapsulates all mutable service state (replaces module-level globals).""" def __init__(self): - self.gateway: ITSGateway = ITSGateway(algorithm=SelfConsistency()) - self.config = _ServiceConfig() + self.gateway: ITSGateway = ITSGateway() def reset(self): - self.gateway = ITSGateway(algorithm=SelfConsistency()) - self.config = _ServiceConfig() + self.gateway = ITSGateway() _state = _ServiceState() @@ -59,7 +43,7 @@ def reset(self): @asynccontextmanager async def _lifespan(application: FastAPI): yield - await _state.gateway.ashutdown() + await _state.gateway.aclose() app = FastAPI( @@ -75,26 +59,19 @@ def _build_its_config( its_budget: int | None = None, its_endpoint: str | None = None, its_api_key: str | None = None, -) -> ITSRequestConfig: - """Build ITSRequestConfig merging headers, body, and service defaults. +) -> ITSRequestConfigUpdate: + """Build a per-request ITSRequestConfigUpdate overlay. - Priority: header > body > service default. + Only per-request fields (header > body) are populated; service-default + fields are left ``None`` for the gateway to merge. """ - if its_budget is not None: - budget = its_budget - elif request.budget is not None: - budget = request.budget - else: - budget = _state.config.budget - - api_endpoint = its_endpoint or _state.config.endpoint - api_key = its_api_key if its_api_key is not None else _state.config.api_key - - return ITSRequestConfig( + budget = its_budget if its_budget is not None else request.budget + return ITSRequestConfigUpdate( budget=budget, - api_endpoint=api_endpoint, + api_endpoint=its_endpoint, + api_key=its_api_key, model=request.model, - api_key=api_key, + temperature=request.temperature, ) @@ -102,7 +79,12 @@ def _build_its_config( async def config_service(request: ConfigRequest) -> dict[str, str]: """Configure the IaaS service with language model and scaling algorithm.""" try: - _state.gateway.configure( + config = ITSRequestConfigUpdate( + budget=request.budget, + api_endpoint=request.endpoint, + api_key=request.api_key, + model=request.model, + temperature=request.temperature, alg=request.alg, regex_patterns=request.regex_patterns, tool_vote=request.tool_vote, @@ -110,20 +92,14 @@ async def config_service(request: ConfigRequest) -> dict[str, str]: threshold=request.threshold, confidence_threshold=request.confidence_threshold, ) - _state.config.endpoint = request.endpoint - _state.config.model = request.model - _state.config.api_key = request.api_key - _state.config.alg = request.alg - if request.budget is not None: - _state.config.budget = request.budget - if request.temperature is not None: - _state.config.temperature = request.temperature + _state.gateway.configure(config) + resolved = _state.gateway._default_config logger.info( "Configured IaaS: model=%s, alg=%s, budget=%s", - request.model, - request.alg, - _state.config.budget, + resolved.model, + resolved.alg, + resolved.budget, ) return { "status": "success", @@ -145,11 +121,12 @@ async def config_service(request: ConfigRequest) -> dict[str, str]: @app.get("/v1/models") async def list_models() -> dict[str, list[dict[str, str]]]: """List available models (OpenAI-compatible endpoint).""" - if _state.config.model: + model = _state.gateway._default_config.model + if model: return { "data": [ { - "id": _state.config.model, + "id": model, "object": "model", "owned_by": "its_hub", } @@ -202,7 +179,7 @@ async def chat_completions( metadata = None if not request.return_response_only: metadata = { - "algorithm": _state.config.alg, + "algorithm": result["alg"], "all_responses": result.get("responses"), "response_counts": result.get("response_counts"), "selected_index": result.get("selected_index"), @@ -253,7 +230,7 @@ async def _generate(): its_config = _build_its_config(request, its_budget, its_endpoint, its_api_key) - if not its_config.api_endpoint: + if not (its_config.api_endpoint or _state.gateway._default_config.api_endpoint): yield f"data: {json.dumps({'error': 'Service not configured'})}\n\n" yield "data: [DONE]\n\n" return diff --git a/its_hub/integration/iaas/models.py b/its_hub/integration/iaas/models.py index 1d453792..348a556b 100644 --- a/its_hub/integration/iaas/models.py +++ b/its_hub/integration/iaas/models.py @@ -2,23 +2,18 @@ from typing import Any -from pydantic import BaseModel, Field, field_validator, model_validator +from pydantic import BaseModel, Field, field_validator -from its_hub.api.types import ChatMessage -from its_hub.core.gateway import SUPPORTED_ALGORITHMS +from its_hub.api.types import SUPPORTED_ALGORITHMS, ChatMessage class ConfigRequest(BaseModel): """Configuration request for setting up the IaaS service.""" - provider: str = Field("openai", description="LM provider: 'openai'") endpoint: str = Field(..., description="Language model endpoint URL") api_key: str | None = Field(None, description="API key for the language model") model: str = Field(..., description="Model name identifier") alg: str = Field(..., description="Scaling algorithm to use") - extra_args: dict[str, Any] | None = Field( - None, description="Additional provider-specific arguments" - ) regex_patterns: list[str] | None = Field( None, description="Regex patterns for self-consistency projection function" ) @@ -64,12 +59,6 @@ def validate_algorithm(cls, v): ) return v - @model_validator(mode="after") - def validate_config_requirements(self): - if self.provider == "openai" and not self.api_key: - raise ValueError("api_key is required when using openai provider") - return self - class ChatCompletionRequest(BaseModel): """Chat completion request with inference-time scaling support.""" diff --git a/tests/conftest.py b/tests/conftest.py index ca40e326..4ee6e1e8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,11 +4,12 @@ import socket import threading import time -from http.server import BaseHTTPRequestHandler, HTTPServer +from http.server import BaseHTTPRequestHandler, HTTPServer, ThreadingHTTPServer import pytest from its_hub import AbstractLanguageModel, AbstractOutcomeRewardModel +from tests.mocks.recording_llm import RecordingLLMHandler def find_free_port() -> int: @@ -303,6 +304,7 @@ def vllm_server(): server.shutdown() server_thread.join() + server.server_close() @pytest.fixture(scope="session") @@ -321,6 +323,7 @@ def openai_server(): server.shutdown() server_thread.join() + server.server_close() @pytest.fixture @@ -351,3 +354,20 @@ def mock_process_reward_model(): "ERROR_TRIGGER": "trigger_error", "VLLM_ERROR_TRIGGER": "error", } + + +@pytest.fixture +def llm_server(): + """Controllable upstream LLM server; see RecordingLLMHandler.""" + RecordingLLMHandler.reset() + port = find_free_port() + server = ThreadingHTTPServer(("localhost", port), RecordingLLMHandler) + thread = threading.Thread(target=server.serve_forever) + thread.daemon = True + thread.start() + time.sleep(0.1) + yield f"http://localhost:{port}" + server.shutdown() + thread.join() + server.server_close() + RecordingLLMHandler.reset() diff --git a/tests/mocks/recording_llm.py b/tests/mocks/recording_llm.py new file mode 100644 index 00000000..8f301b89 --- /dev/null +++ b/tests/mocks/recording_llm.py @@ -0,0 +1,92 @@ +"""Controllable upstream LLM server for concurrency / regression tests. + +Lives outside conftest.py: pytest loads conftest.py as plugin module ``conftest`` +(distinct from ``tests.conftest``), so a class defined there would be duplicated. +""" + +import asyncio +import json +import threading +from http.server import BaseHTTPRequestHandler + + +class RecordingLLMHandler(BaseHTTPRequestHandler): + """Upstream LLM stand-in that records requests and can hold them.""" + + received_bodies: list[dict] = [] # noqa: RUF012 - shared across per-request instances + _hold = threading.Event() + _hold.set() + _lock = threading.Lock() + + @classmethod + def reset(cls): + with cls._lock: + cls.received_bodies.clear() + cls._hold.set() + + @classmethod + def hold(cls): + cls._hold.clear() + + @classmethod + def release(cls): + cls._hold.set() + + @classmethod + async def wait_for_bodies(cls, n: int, timeout: float = 5.0) -> None: + """Poll until at least ``n`` request bodies have been recorded. + + Replaces fixed ``asyncio.sleep`` coordination so tests fail fast and + deterministically when expected upstream traffic never arrives. + """ + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while len(cls.received_bodies) < n: + if loop.time() >= deadline: + raise AssertionError( + f"timed out after {timeout}s waiting for {n} recorded " + f"bodies; got {len(cls.received_bodies)}" + ) + await asyncio.sleep(0.01) + + def do_POST(self): + if self.path != "/v1/chat/completions": + self.send_response(404) + self.end_headers() + return + + body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + with type(self)._lock: + type(self).received_bodies.append(body) + + type(self)._hold.wait(timeout=10) + + payload = json.dumps( + { + "model": body.get("model", "unknown"), + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": f"answer from {body.get('model', 'unknown')}", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 10, + "total_tokens": 20, + }, + } + ).encode() + + try: + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + except (BrokenPipeError, ConnectionResetError): + pass diff --git a/tests/test_gateway.py b/tests/test_gateway.py index 4d334e00..251050a2 100644 --- a/tests/test_gateway.py +++ b/tests/test_gateway.py @@ -1,16 +1,21 @@ """Tests for ITSGateway (core/gateway.py).""" +import asyncio from collections import Counter from unittest.mock import AsyncMock, MagicMock, patch import pytest -from its_hub.api.types import GenerationUsage, ITSRequestConfig -from its_hub.core.gateway import SUPPORTED_ALGORITHMS, ITSGateway +from its_hub.api.types import ( + SUPPORTED_ALGORITHMS, + GenerationUsage, + ITSRequestConfigUpdate, +) +from its_hub.core.gateway import ITSGateway +from tests.mocks.recording_llm import RecordingLLMHandler def _make_result(usage=None): - """Build a mock SelfConsistencyResult-like object.""" result = MagicMock() result.the_one = {"role": "assistant", "content": "answer"} result.responses = [ @@ -26,107 +31,163 @@ def _make_result(usage=None): def _make_config(**overrides): defaults = { "budget": 3, + "alg": "self-consistency", "api_endpoint": "http://llm/v1", "model": "gpt-4", "api_key": "sk-test", } defaults.update(overrides) - return ITSRequestConfig(**defaults) + return ITSRequestConfigUpdate(**defaults) MESSAGES = [{"role": "user", "content": "What is 2+2?"}] +def _mock_lm(): + lm = MagicMock() + lm.close = AsyncMock() + return lm + + +def _patch_build_lm(gw): + return patch.object(gw, "_build_lm", return_value=_mock_lm()) + + +def _algo(gw): + return gw._build_algorithm(gw._default_config.resolve()) + + class TestGatewayConstruction: def test_default_init(self): - with patch("its_hub.core.gateway.OpenAICompatibleLanguageModel"): - gw = ITSGateway() - assert gw._algorithm is not None + gw = ITSGateway() assert gw._orchestrator is not None - assert gw._algorithm_name == "SelfConsistency" - - def test_custom_algorithm(self): - algo = MagicMock() - type(algo).__name__ = "CustomAlgo" - gw = ITSGateway(algorithm=algo) - assert gw._algorithm is algo - assert gw._algorithm_name == "CustomAlgo" + # Stored default is partial; system defaults (alg, budget) materialize + # at resolve() time on ITSRequestConfig. + resolved = gw._default_config.merge( + ITSRequestConfigUpdate(api_endpoint="http://x/v1", model="m") + ).resolve() + assert resolved.alg == "self-consistency" + assert resolved.budget == 4 + + def test_custom_default_config(self): + config = ITSRequestConfigUpdate( + alg="beta-self-consistency", + api_endpoint="http://x/v1", + model="m", + api_key="k", + ) + gw = ITSGateway(default_config=config) + assert gw._default_config is config + assert gw._default_config.alg == "beta-self-consistency" - def test_custom_orchestrator_passed_to_default_algorithm(self): + def test_custom_orchestrator_passed_to_algorithm(self): orch = MagicMock() - with patch("its_hub.core.gateway.SelfConsistency") as sc_cls: - ITSGateway(orchestrator=orch) - sc_cls.assert_called_once_with(orchestrator=orch) + gw = ITSGateway(orchestrator=orch) + assert gw._orchestrator is orch + algo = gw._build_algorithm(_make_config().resolve()) + assert algo.orchestrator is orch -class TestLMClientCaching: +class TestLMClientLifecycle: @pytest.mark.asyncio - async def test_creates_new_client(self): - gw = ITSGateway(algorithm=MagicMock()) - with patch("its_hub.core.gateway.OpenAICompatibleLanguageModel") as lm_cls: - lm_cls.return_value = MagicMock() - lm = await gw._get_or_create_lm("http://a/v1", "m1", "key1") - assert lm is lm_cls.return_value - lm_cls.assert_called_once() + async def test_creates_new_lm_per_request(self): + gw = ITSGateway() + algo = MagicMock() + algo.ainfer = AsyncMock(return_value=_make_result()) + with ( + patch.object(gw, "_build_algorithm", return_value=algo), + patch.object(gw, "_build_lm", return_value=_mock_lm()) as build_lm, + ): + await gw.arun_chat_completion(_make_config(), MESSAGES) + await gw.arun_chat_completion(_make_config(), MESSAGES) + assert build_lm.call_count == 2 @pytest.mark.asyncio - async def test_reuses_cached_client(self): - gw = ITSGateway(algorithm=MagicMock()) - with patch("its_hub.core.gateway.OpenAICompatibleLanguageModel") as lm_cls: - lm_cls.return_value = MagicMock() - lm1 = await gw._get_or_create_lm("http://a/v1", "m1", "key1") - lm2 = await gw._get_or_create_lm("http://a/v1", "m1", "key1") - assert lm1 is lm2 - assert lm_cls.call_count == 1 + async def test_closes_lm_after_request(self): + gw = ITSGateway() + algo = MagicMock() + algo.ainfer = AsyncMock(return_value=_make_result()) + lm = _mock_lm() + with ( + patch.object(gw, "_build_algorithm", return_value=algo), + patch.object(gw, "_build_lm", return_value=lm), + ): + await gw.arun_chat_completion(_make_config(), MESSAGES) + lm.close.assert_awaited_once() @pytest.mark.asyncio - async def test_different_endpoint_creates_new_client(self): - gw = ITSGateway(algorithm=MagicMock()) - with patch("its_hub.core.gateway.OpenAICompatibleLanguageModel") as lm_cls: - lm_cls.return_value = MagicMock() - await gw._get_or_create_lm("http://a/v1", "m1", "key1") - lm_cls.return_value = MagicMock() - await gw._get_or_create_lm("http://b/v1", "m1", "key1") - assert lm_cls.call_count == 2 + async def test_closes_lm_even_when_algorithm_raises(self): + gw = ITSGateway() + algo = MagicMock() + algo.ainfer = AsyncMock(side_effect=RuntimeError("boom")) + lm = _mock_lm() + with ( + patch.object(gw, "_build_algorithm", return_value=algo), + patch.object(gw, "_build_lm", return_value=lm), + pytest.raises(RuntimeError, match="boom"), + ): + await gw.arun_chat_completion(_make_config(), MESSAGES) + lm.close.assert_awaited_once() @pytest.mark.asyncio - async def test_different_api_key_creates_new_client(self): - gw = ITSGateway(algorithm=MagicMock()) - with patch("its_hub.core.gateway.OpenAICompatibleLanguageModel") as lm_cls: - lm_cls.return_value = MagicMock() - await gw._get_or_create_lm("http://a/v1", "m1", "key-A") - lm_cls.return_value = MagicMock() - await gw._get_or_create_lm("http://a/v1", "m1", "key-B") - assert lm_cls.call_count == 2 + async def test_concurrent_different_models_do_not_interfere(self, llm_server): + gw = ITSGateway() + + RecordingLLMHandler.hold() + + config_a = _make_config(api_endpoint=f"{llm_server}/v1", model="model-A") + config_b = _make_config(api_endpoint=f"{llm_server}/v1", model="model-B") + task_a = asyncio.create_task(gw.arun_chat_completion(config_a, MESSAGES)) + await RecordingLLMHandler.wait_for_bodies(3) + + task_b = asyncio.create_task(gw.arun_chat_completion(config_b, MESSAGES)) + await RecordingLLMHandler.wait_for_bodies(6) + + RecordingLLMHandler.release() + + result_a = await task_a + result_b = await task_b + + assert result_a["message"]["content"] == "answer from model-A" + assert result_b["message"]["content"] == "answer from model-B" + + +class TestGatewayShutdown: @pytest.mark.asyncio - async def test_ashutdown_closes_clients(self): - gw = ITSGateway(algorithm=MagicMock()) - mock_lm = MagicMock() - mock_lm.close = AsyncMock() - with patch( - "its_hub.core.gateway.OpenAICompatibleLanguageModel", return_value=mock_lm - ): - await gw._get_or_create_lm("http://a/v1", "m1", "key1") - assert len(gw._lm_cache) == 1 - await gw.ashutdown() - mock_lm.close.assert_awaited_once() - assert len(gw._lm_cache) == 0 + async def test_aclose_closes_pooled_connectors(self): + gw = ITSGateway() + c1 = gw._get_connector("http://upstream-a/v1") + c2 = gw._get_connector("http://upstream-b/v1") + assert not c1.closed + assert not c2.closed + assert len(gw._connector_pool) == 2 + + await gw.aclose() + + assert c1.closed + assert c2.closed + assert len(gw._connector_pool) == 0 @pytest.mark.asyncio - async def test_evicts_oldest_when_cache_full(self): - gw = ITSGateway(algorithm=MagicMock(), max_lm_cache_size=2) - mocks = [] - with patch("its_hub.core.gateway.OpenAICompatibleLanguageModel") as lm_cls: - for i in range(3): - mock_lm = MagicMock() - mock_lm.close = AsyncMock() - lm_cls.return_value = mock_lm - mocks.append(mock_lm) - await gw._get_or_create_lm(f"http://host{i}/v1", "m1", "key1") + async def test_aclose_is_idempotent(self): + gw = ITSGateway() + gw._get_connector("http://upstream/v1") + await gw.aclose() + await gw.aclose() # second call must not raise - assert len(gw._lm_cache) == 2 - mocks[0].close.assert_awaited_once() + @pytest.mark.asyncio + async def test_get_connector_rebuilds_after_close(self): + """A request arriving after aclose() gets a fresh connector.""" + gw = ITSGateway() + c1 = gw._get_connector("http://upstream/v1") + await gw.aclose() + assert c1.closed + + c2 = gw._get_connector("http://upstream/v1") + assert not c2.closed + assert c2 is not c1 + await gw.aclose() class TestRunChatCompletion: @@ -135,8 +196,11 @@ async def test_response_only(self): usage = GenerationUsage(prompt_tokens=10, completion_tokens=20, num_calls=3) algo = MagicMock() algo.ainfer = AsyncMock(return_value=_make_result(usage)) - gw = ITSGateway(algorithm=algo) - with patch.object(gw, "_get_or_create_lm", return_value=MagicMock()): + gw = ITSGateway() + with ( + patch.object(gw, "_build_algorithm", return_value=algo), + _patch_build_lm(gw), + ): result = await gw.arun_chat_completion(_make_config(), MESSAGES) assert result["message"] == {"role": "assistant", "content": "answer"} assert result["usage"]["prompt_tokens"] == 10 @@ -149,8 +213,11 @@ async def test_full_result(self): usage = GenerationUsage(prompt_tokens=5, completion_tokens=10, num_calls=2) algo = MagicMock() algo.ainfer = AsyncMock(return_value=_make_result(usage)) - gw = ITSGateway(algorithm=algo) - with patch.object(gw, "_get_or_create_lm", return_value=MagicMock()): + gw = ITSGateway() + with ( + patch.object(gw, "_build_algorithm", return_value=algo), + _patch_build_lm(gw), + ): result = await gw.arun_chat_completion( _make_config(), MESSAGES, return_response_only=False ) @@ -164,9 +231,12 @@ async def test_full_result(self): async def test_tools_forwarded(self): algo = MagicMock() algo.ainfer = AsyncMock(return_value=_make_result()) - gw = ITSGateway(algorithm=algo) + gw = ITSGateway() tools = [{"type": "function", "function": {"name": "f"}}] - with patch.object(gw, "_get_or_create_lm", return_value=MagicMock()): + with ( + patch.object(gw, "_build_algorithm", return_value=algo), + _patch_build_lm(gw), + ): await gw.arun_chat_completion( _make_config(), MESSAGES, tools=tools, tool_choice="auto" ) @@ -176,18 +246,19 @@ async def test_tools_forwarded(self): @pytest.mark.asyncio async def test_missing_model_raises(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() config = _make_config(model=None) - with pytest.raises(ValueError, match="Model must be specified"): + with pytest.raises(ValueError, match="model must be specified"): await gw.arun_chat_completion(config, MESSAGES) @pytest.mark.asyncio async def test_algorithm_exception_propagates(self): algo = MagicMock() algo.ainfer = AsyncMock(side_effect=RuntimeError("boom")) - gw = ITSGateway(algorithm=algo) + gw = ITSGateway() with ( - patch.object(gw, "_get_or_create_lm", return_value=MagicMock()), + patch.object(gw, "_build_algorithm", return_value=algo), + _patch_build_lm(gw), pytest.raises(RuntimeError, match="boom"), ): await gw.arun_chat_completion(_make_config(), MESSAGES) @@ -196,141 +267,190 @@ async def test_algorithm_exception_propagates(self): async def test_usage_empty_when_none(self): algo = MagicMock() algo.ainfer = AsyncMock(return_value=_make_result(usage=None)) - gw = ITSGateway(algorithm=algo) - with patch.object(gw, "_get_or_create_lm", return_value=MagicMock()): + gw = ITSGateway() + with ( + patch.object(gw, "_build_algorithm", return_value=algo), + _patch_build_lm(gw), + ): result = await gw.arun_chat_completion(_make_config(), MESSAGES) assert result["usage"] == {} + @pytest.mark.asyncio + async def test_temperature_forwarded_to_lm(self, llm_server): + gw = ITSGateway() + config = _make_config(api_endpoint=f"{llm_server}/v1", temperature=0.7) + await gw.arun_chat_completion(config, MESSAGES) + assert RecordingLLMHandler.received_bodies[-1].get("temperature") == 0.7 -class TestHashApiKey: - def test_same_key_same_hash(self): - assert ITSGateway._hash_api_key("sk-abc") == ITSGateway._hash_api_key("sk-abc") + @pytest.mark.asyncio + async def test_reconfigure_does_not_swap_in_flight_algorithm(self, llm_server): + """Reconfigure mid-request must not swap the algorithm used by the + in-flight request.""" + gw = ITSGateway() + # `alg` is deliberately omitted from the per-request overlay so the + # algorithm is sourced from the (reconfigurable) service default — + # only then does a mid-flight /configure actually exercise the + # snapshot. With `alg` in the overlay, `result["alg"]` would be + # "self-consistency" regardless of any reconfigure. + base = ITSRequestConfigUpdate( + api_endpoint=f"{llm_server}/v1", + model="m", + ) + gw.configure(base) - def test_different_keys_different_hash(self): - assert ITSGateway._hash_api_key("sk-abc") != ITSGateway._hash_api_key("sk-xyz") + RecordingLLMHandler.hold() - def test_none_key(self): - h = ITSGateway._hash_api_key(None) - assert isinstance(h, str) - assert len(h) == 16 + task = asyncio.create_task(gw.arun_chat_completion(base, MESSAGES)) + await RecordingLLMHandler.wait_for_bodies(1) - def test_hash_length(self): - assert len(ITSGateway._hash_api_key("sk-test")) == 16 + gw.configure( + ITSRequestConfigUpdate( + alg="beta-self-consistency", + ) + ) + + RecordingLLMHandler.release() + result = await task + + assert result["alg"] == "self-consistency" + assert gw._default_config.alg == "beta-self-consistency" class TestConfigure: def test_configure_self_consistency(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() gw.configure( - alg="self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], + _make_config(alg="self-consistency", regex_patterns=[r"\\boxed{([^}]+)}"]), ) - assert gw._algorithm_name == "SelfConsistency" + assert gw._default_config.alg == "self-consistency" + assert type(_algo(gw)).__name__ == "SelfConsistency" def test_configure_defaults_tool_vote_when_omitted(self): - """With neither regex nor tool_vote, the algorithm's DEFAULT_TOOL_VOTE applies.""" - gw = ITSGateway(algorithm=MagicMock()) - gw.configure(alg="self-consistency") - assert gw._algorithm.tool_vote == "tool_hierarchical" + gw = ITSGateway() + gw.configure(_make_config(alg="self-consistency")) + assert _algo(gw).tool_vote == "tool_hierarchical" def test_configure_with_tool_vote(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() gw.configure( - alg="self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], - tool_vote="tool_name", - exclude_tool_args=["timestamp"], + _make_config( + alg="self-consistency", + regex_patterns=[r"\\boxed{([^}]+)}"], + tool_vote="tool_name", + exclude_tool_args=["timestamp"], + ), ) - assert gw._algorithm_name == "SelfConsistency" - - def test_configure_unsupported_algorithm(self): - gw = ITSGateway(algorithm=MagicMock()) + algo = _algo(gw) + assert algo.tool_vote == "tool_name" + assert algo.exclude_args == ("timestamp",) + + def test_configure_tool_vote_persists_when_omitted(self): + """`None` means trickle down, so a /configure omitting tool_vote + cannot clear a previously-set value.""" + gw = ITSGateway() + gw.configure(_make_config(alg="self-consistency", tool_vote="tool_name")) + gw.configure(_make_config(alg="self-consistency")) # no tool_vote + assert _algo(gw).tool_vote == "tool_name" + + def test_overlay_unsupported_algorithm(self): + gw = ITSGateway() with pytest.raises(ValueError, match="not supported"): - gw.configure(alg="beam-search") + gw.configure(_make_config(alg="beam-search")) - def test_configure_invalid_regex(self): - gw = ITSGateway(algorithm=MagicMock()) + def test_overlay_invalid_regex(self): + gw = ITSGateway() with pytest.raises(ValueError, match="Invalid regex pattern"): - gw.configure(alg="self-consistency", regex_patterns=["[invalid("]) + gw.configure( + _make_config(alg="self-consistency", regex_patterns=["[invalid("]), + ) def test_configure_preserves_orchestrator(self): orch = MagicMock() - gw = ITSGateway(algorithm=MagicMock(), orchestrator=orch) - with patch("its_hub.core.gateway.SelfConsistency") as sc_cls: - sc_cls.return_value = MagicMock() - type(sc_cls.return_value).__name__ = "SelfConsistency" - gw.configure( - alg="self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], - ) - sc_cls.assert_called_once() - assert sc_cls.call_args.kwargs["orchestrator"] is orch + gw = ITSGateway(orchestrator=orch) + gw.configure( + _make_config(alg="self-consistency", regex_patterns=[r"\\boxed{([^}]+)}"]), + ) + assert _algo(gw).orchestrator is orch - def test_configure_invalid_tool_vote(self): - gw = ITSGateway(algorithm=MagicMock()) + def test_overlay_invalid_tool_vote(self): + gw = ITSGateway() with pytest.raises(ValueError, match="tool_vote must be one of"): gw.configure( - alg="self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], - tool_vote="invalid", + _make_config( + alg="self-consistency", + regex_patterns=[r"\\boxed{([^}]+)}"], + tool_vote="invalid", + ), ) + def test_overlay_invalid_temperature(self): + gw = ITSGateway() + with pytest.raises(ValueError, match="temperature must be in"): + gw.configure(_make_config(temperature=3.0)) + def test_configure_adaptive_self_consistency(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() gw.configure( - alg="adaptive-self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], + _make_config( + alg="adaptive-self-consistency", + regex_patterns=[r"\\boxed{([^}]+)}"], + ), ) - assert gw._algorithm_name == "AdaptiveSelfConsistency" - # Unset threshold falls back to the class default. - assert gw._algorithm.threshold == pytest.approx(0.75) + algo = _algo(gw) + assert type(algo).__name__ == "AdaptiveSelfConsistency" + assert algo.threshold == pytest.approx(0.75) def test_configure_adaptive_threshold_plumbed(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() gw.configure( - alg="adaptive-self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], - threshold=0.9, + _make_config( + alg="adaptive-self-consistency", + regex_patterns=[r"\\boxed{([^}]+)}"], + threshold=0.9, + ), ) - assert gw._algorithm.threshold == pytest.approx(0.9) + assert _algo(gw).threshold == pytest.approx(0.9) def test_configure_beta_self_consistency(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() gw.configure( - alg="beta-self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], + _make_config( + alg="beta-self-consistency", regex_patterns=[r"\\boxed{([^}]+)}"] + ), ) - assert gw._algorithm_name == "BetaSelfConsistency" - # Unset confidence_threshold falls back to the class default. - assert gw._algorithm.confidence_threshold == pytest.approx(0.95) + algo = _algo(gw) + assert type(algo).__name__ == "BetaSelfConsistency" + assert algo.confidence_threshold == pytest.approx(0.95) def test_configure_beta_confidence_threshold_plumbed(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() gw.configure( - alg="beta-self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], - confidence_threshold=0.8, + _make_config( + alg="beta-self-consistency", + regex_patterns=[r"\\boxed{([^}]+)}"], + confidence_threshold=0.8, + ), ) - assert gw._algorithm.confidence_threshold == pytest.approx(0.8) + assert _algo(gw).confidence_threshold == pytest.approx(0.8) @pytest.mark.parametrize( "alg", ["adaptive-self-consistency", "beta-self-consistency"], ) def test_configure_family_shares_tool_vote(self, alg): - """Adaptive/beta inherit the tool-vote voting surface from SelfConsistency.""" - gw = ITSGateway(algorithm=MagicMock()) - gw.configure(alg=alg, tool_vote="tool_name") - assert gw._algorithm.tool_vote == "tool_name" + gw = ITSGateway() + gw.configure(_make_config(alg=alg, tool_vote="tool_name")) + assert _algo(gw).tool_vote == "tool_name" def test_configure_preserves_orchestrator_for_family(self): orch = MagicMock() - gw = ITSGateway(algorithm=MagicMock(), orchestrator=orch) + gw = ITSGateway(orchestrator=orch) gw.configure( - alg="beta-self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], + _make_config( + alg="beta-self-consistency", regex_patterns=[r"\\boxed{([^}]+)}"] + ), ) - assert gw._algorithm.orchestrator is orch + assert _algo(gw).orchestrator is orch class TestSupportedAlgorithms: @@ -350,7 +470,7 @@ def test_is_frozenset(self): class TestEndpointValidation: @pytest.mark.asyncio async def test_missing_endpoint_raises(self): - gw = ITSGateway(algorithm=MagicMock()) + gw = ITSGateway() config = _make_config(api_endpoint="") with pytest.raises(ValueError, match="api_endpoint"): await gw.arun_chat_completion(config, MESSAGES) diff --git a/tests/test_iaas.py b/tests/test_iaas.py index 124a6763..c1371417 100644 --- a/tests/test_iaas.py +++ b/tests/test_iaas.py @@ -1,19 +1,24 @@ """Tests for the Inference-as-a-Service (IaaS) integration.""" +import asyncio import logging from collections import Counter from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import pytest_asyncio from fastapi.testclient import TestClient +from httpx import ASGITransport -from its_hub.api.types import ChatMessage +from its_hub.api.types import ChatMessage, ITSRequestConfigUpdate from its_hub.integration.iaas.app import _state, app from its_hub.integration.iaas.models import ( ChatCompletionRequest, ConfigRequest, ) from tests.conftest import TEST_CONSTANTS +from tests.mocks.recording_llm import RecordingLLMHandler from tests.mocks.test_data import TestDataFactory @@ -25,12 +30,45 @@ def iaas_client(): _state.reset() +@pytest_asyncio.fixture +async def async_iaas_client(): + """Async client (allows concurrent /configure + /v1/...).""" + _state.reset() + transport = ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as client: + yield client + _state.reset() + + @pytest.fixture(scope="session") def vllm_endpoint(vllm_server): - """Alias the vllm_server fixture for clarity in IaaS tests.""" return vllm_server +def _install_mock_gateway(mock_return=None, side_effect=None, default=None): + """Replace _state.gateway with a mock, preserving _default_config for merging.""" + if default is None: + default = _state.gateway._default_config + mock_gw = MagicMock() + mock_gw._default_config = default + # _lifespan awaits aclose() on shutdown; make the mock awaitable. + mock_gw.aclose = AsyncMock() + if side_effect: + mock_gw.arun_chat_completion = AsyncMock(side_effect=side_effect) + else: + # Mirror resolve(): alg defaults to "self-consistency" when the update + # omits it. We can't call resolve() directly because the header-only + # path has no endpoint/model yet. + patched = {**(mock_return or {}), "alg": default.alg or "self-consistency"} + + async def _arun(*args, **kwargs): + return patched + + mock_gw.arun_chat_completion = AsyncMock(side_effect=_arun) + _state.gateway = mock_gw + return mock_gw + + def _mock_gateway_result(content="answer", usage=None): """Build a dict matching ITSGateway.arun_chat_completion return (response_only=True).""" if usage is None: @@ -135,8 +173,22 @@ def test_configure_creates_gateway(self, iaas_client, vllm_endpoint): response = iaas_client.post("/configure", json=config) assert response.status_code == 200 assert _state.gateway is not None - assert _state.config.model == TEST_CONSTANTS["DEFAULT_MODEL_NAME"] - assert _state.config.endpoint == vllm_endpoint + assert ( + _state.gateway._default_config.model == TEST_CONSTANTS["DEFAULT_MODEL_NAME"] + ) + assert _state.gateway._default_config.api_endpoint == vllm_endpoint + + def test_configure_without_api_key(self, iaas_client, vllm_endpoint): + """A no-auth local endpoint can be configured without an api_key.""" + config = { + "endpoint": vllm_endpoint, + "model": TEST_CONSTANTS["DEFAULT_MODEL_NAME"], + "alg": "self-consistency", + } + response = iaas_client.post("/configure", json=config) + assert response.status_code == 200 + assert _state.gateway is not None + assert _state.gateway._default_config.api_key is None def test_models_endpoint_after_configure(self, iaas_client, vllm_endpoint): config = { @@ -226,11 +278,11 @@ def test_tool_vote_algorithm_usage_verification(self, iaas_client, vllm_endpoint } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 - + _state.gateway._build_algorithm(_state.gateway._default_config.resolve()) mock_sc.assert_called_once() call_args = mock_sc.call_args assert call_args.kwargs["tool_vote"] == "tool_hierarchical" - assert call_args.kwargs["exclude_args"] == ["timestamp", "id"] + assert call_args.kwargs["exclude_args"] == ("timestamp", "id") class TestAdaptiveAndBetaSelfConsistency: @@ -251,7 +303,7 @@ def test_family_basic_configuration(self, iaas_client, vllm_endpoint, alg): response = iaas_client.post("/configure", json=config) assert response.status_code == 200 assert "success" in response.json()["status"] - assert _state.config.alg == alg + assert _state.gateway._default_config.alg == alg def test_adaptive_threshold_forwarded_to_algorithm( self, iaas_client, vllm_endpoint @@ -268,6 +320,7 @@ def test_adaptive_threshold_forwarded_to_algorithm( } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 + _state.gateway._build_algorithm(_state.gateway._default_config.resolve()) mock_alg.assert_called_once() assert mock_alg.call_args.kwargs["threshold"] == 0.9 @@ -286,6 +339,7 @@ def test_beta_confidence_threshold_forwarded_to_algorithm( } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 + _state.gateway._build_algorithm(_state.gateway._default_config.resolve()) mock_alg.assert_called_once() assert mock_alg.call_args.kwargs["confidence_threshold"] == 0.8 @@ -302,6 +356,7 @@ def test_threshold_omitted_uses_algorithm_default(self, iaas_client, vllm_endpoi } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 + _state.gateway._build_algorithm(_state.gateway._default_config.resolve()) assert "threshold" not in mock_alg.call_args.kwargs @pytest.mark.parametrize( @@ -338,8 +393,9 @@ def test_family_tool_vote_forwarded(self, iaas_client, vllm_endpoint): } response = iaas_client.post("/configure", json=config) assert response.status_code == 200 + _state.gateway._build_algorithm(_state.gateway._default_config.resolve()) assert mock_alg.call_args.kwargs["tool_vote"] == "tool_hierarchical" - assert mock_alg.call_args.kwargs["exclude_args"] == ["timestamp"] + assert mock_alg.call_args.kwargs["exclude_args"] == ("timestamp",) class TestChatCompletions: @@ -353,11 +409,9 @@ def test_chat_completion_with_gateway(self, iaas_client, vllm_endpoint): } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock( - return_value=_mock_gateway_result("Tool voting response") + mock_gw = _install_mock_gateway( + mock_return=_mock_gateway_result("Tool voting response") ) - _state.gateway = mock_gw request_data = TestDataFactory.create_chat_completion_request( user_content="What is 2+2?", budget=8 @@ -380,11 +434,7 @@ def test_chat_completion_full_result(self, iaas_client, vllm_endpoint): } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock( - return_value=_mock_gateway_full_result() - ) - _state.gateway = mock_gw + _install_mock_gateway(mock_return=_mock_gateway_full_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) request_data["return_response_only"] = False @@ -409,11 +459,7 @@ def test_chat_completion_metadata_reports_configured_algorithm( } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock( - return_value=_mock_gateway_full_result() - ) - _state.gateway = mock_gw + _install_mock_gateway(mock_return=_mock_gateway_full_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) request_data["return_response_only"] = False @@ -421,6 +467,44 @@ def test_chat_completion_metadata_reports_configured_algorithm( assert response.status_code == 200 assert response.json()["metadata"]["algorithm"] == "beta-self-consistency" + @pytest.mark.asyncio + async def test_metadata_reports_request_time_algorithm( + self, llm_server, async_iaas_client + ): + """A reconfigure during a request must not change the reported algorithm.""" + base = { + "endpoint": f"{llm_server}/v1", + "api_key": "test-key", + "model": "test-model", + "budget": 1, + } + await async_iaas_client.post( + "/configure", json={**base, "alg": "self-consistency"} + ) + + RecordingLLMHandler.hold() + + request_task = asyncio.create_task( + async_iaas_client.post( + "/v1/chat/completions", + json={ + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}], + "return_response_only": False, + }, + ) + ) + await RecordingLLMHandler.wait_for_bodies(1) + + await async_iaas_client.post( + "/configure", json={**base, "alg": "beta-self-consistency"} + ) + + RecordingLLMHandler.release() + resp = await request_task + assert resp.status_code == 200 + assert resp.json()["metadata"]["algorithm"] == "self-consistency" + @pytest.mark.parametrize( "invalid_request", [ @@ -457,16 +541,15 @@ def _parse_sse(self, response): def _configure_and_mock( self, iaas_client, endpoint, mock_return=None, side_effect=None ): - _state.config.endpoint = endpoint - _state.config.model = TEST_CONSTANTS["DEFAULT_MODEL_NAME"] - _state.config.api_key = TEST_CONSTANTS["DEFAULT_API_KEY"] - mock_gw = MagicMock() - if side_effect: - mock_gw.arun_chat_completion = AsyncMock(side_effect=side_effect) - else: - mock_gw.arun_chat_completion = AsyncMock(return_value=mock_return) - _state.gateway = mock_gw - return mock_gw + return _install_mock_gateway( + mock_return=mock_return, + side_effect=side_effect, + default=ITSRequestConfigUpdate( + api_endpoint=endpoint, + model=TEST_CONSTANTS["DEFAULT_MODEL_NAME"], + api_key=TEST_CONSTANTS["DEFAULT_API_KEY"], + ), + ) def test_stream_content_response(self, iaas_client, vllm_endpoint): self._configure_and_mock( @@ -643,9 +726,7 @@ def test_budget_from_header_overrides_body(self, iaas_client, vllm_endpoint): } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock(return_value=_mock_gateway_result()) - _state.gateway = mock_gw + mock_gw = _install_mock_gateway(mock_return=_mock_gateway_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) response = iaas_client.post( @@ -670,9 +751,7 @@ def test_endpoint_from_header_overrides_service_default( } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock(return_value=_mock_gateway_result()) - _state.gateway = mock_gw + mock_gw = _install_mock_gateway(mock_return=_mock_gateway_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) response = iaas_client.post( @@ -695,9 +774,7 @@ def test_api_key_from_header(self, iaas_client, vllm_endpoint): } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock(return_value=_mock_gateway_result()) - _state.gateway = mock_gw + mock_gw = _install_mock_gateway(mock_return=_mock_gateway_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) response = iaas_client.post( @@ -712,9 +789,7 @@ def test_api_key_from_header(self, iaas_client, vllm_endpoint): def test_headers_without_service_config(self, iaas_client): """Headers alone can configure a request — no /configure needed.""" - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock(return_value=_mock_gateway_result()) - _state.gateway = mock_gw + mock_gw = _install_mock_gateway(mock_return=_mock_gateway_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) response = iaas_client.post( @@ -773,9 +848,7 @@ def test_negative_budget_header_rejected(self, iaas_client, vllm_endpoint): } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock(return_value=_mock_gateway_result()) - _state.gateway = mock_gw + _install_mock_gateway(mock_return=_mock_gateway_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) response = iaas_client.post( @@ -796,9 +869,7 @@ def test_over_max_budget_header_rejected(self, iaas_client, vllm_endpoint): } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock(return_value=_mock_gateway_result()) - _state.gateway = mock_gw + _install_mock_gateway(mock_return=_mock_gateway_result()) request_data = TestDataFactory.create_chat_completion_request(budget=4) response = iaas_client.post( @@ -835,11 +906,9 @@ def test_generation_error_is_sanitized(self, iaas_client, vllm_endpoint): } iaas_client.post("/configure", json=config) - mock_gw = MagicMock() - mock_gw.arun_chat_completion = AsyncMock( + _install_mock_gateway( side_effect=RuntimeError("Connection to http://internal:8100 refused") ) - _state.gateway = mock_gw request_data = TestDataFactory.create_chat_completion_request(budget=4) response = iaas_client.post("/v1/chat/completions", json=request_data) @@ -883,14 +952,15 @@ def test_self_consistency_accepts_tool_vote_without_regex(self): assert config.regex_patterns is None assert config.tool_vote == "tool_hierarchical" - def test_openai_requires_api_key_when_omitted(self): - with pytest.raises(ValueError, match="api_key is required"): - ConfigRequest( - endpoint="http://example.com:8000", - model="test-model", - alg="self-consistency", - regex_patterns=[r"\\boxed{([^}]+)}"], - ) + def test_api_key_is_optional_when_omitted(self): + """A no-auth endpoint can be configured without an api_key.""" + config = ConfigRequest( + endpoint="http://example.com:8000", + model="test-model", + alg="self-consistency", + regex_patterns=[r"\\boxed{([^}]+)}"], + ) + assert config.api_key is None class TestExtProcessor: