diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 8184329fb..040a3f46e 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -138,7 +138,7 @@ jobs: strategy: fail-fast: false matrix: - dependency-group: [ "base", "transformers", "peft", "catboost" ] + dependency-group: [ "base", "transformers", "peft", "catboost", "typesafe" ] uses: ./.github/workflows/reusable-test.yaml secrets: inherit with: diff --git a/docs/source/installation.rst b/docs/source/installation.rst index edce5cbf3..488a13d08 100644 --- a/docs/source/installation.rst +++ b/docs/source/installation.rst @@ -45,6 +45,9 @@ The table below lists each published extra, what it is typically used for in Aut * - ``openai`` - OpenAI API clients for OpenAI-backed embedders and the ``zero-shot-llm`` style workflows. - ``pip install "autointent[openai]"`` + * - ``typesafe`` + - TypeSafe API client for the ``description_typesafe`` scorer and the ``zero-shot-typesafe`` preset. + - ``pip install "autointent[typesafe]"`` * - ``dspy`` - DSPy-based augmentation and related generation utilities. - ``pip install "autointent[dspy]"`` diff --git a/docs/source/quickstart.rst b/docs/source/quickstart.rst index be0f632b6..c613cd747 100644 --- a/docs/source/quickstart.rst +++ b/docs/source/quickstart.rst @@ -190,6 +190,7 @@ Available Presets - ``transformers-light``: Transformer models with limited search - ``transformers-no-hpo``: Transformer models without hyperparameter optimization - ``zero-shot-llm``: Zero-shot classification using OpenAI models +- ``zero-shot-typesafe``: Zero-shot classification using TypeSafe's ``jev`` model (requires ``TYPESAFE_API_KEY``) - ``zero-shot-encoders``: Zero-shot classification using transformer models Making Predictions diff --git a/pyproject.toml b/pyproject.toml index 914e859de..f2124634c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,6 +80,9 @@ openai = [ "openai (>=2,<3)", "tiktoken (>=0.7,<1)", ] +typesafe = [ + "typesafe-sdk (>=0.7,<1)", +] vllm = [ "vllm>=0.20.0", ] @@ -316,6 +319,7 @@ module = [ "dspy.evaluate.auto_evaluation", "codecarbon", "catboost", + "typesafe_sdk", "jupytext", "docs_utils.*", ] diff --git a/src/autointent/_deps.py b/src/autointent/_deps.py index f56b2a536..e25efee0b 100644 --- a/src/autointent/_deps.py +++ b/src/autointent/_deps.py @@ -37,6 +37,7 @@ "peft", "sentence-transformers", "transformers", + "typesafe", "vllm", "wandb", ] diff --git a/src/autointent/_dump_tools/unit_dumpers.py b/src/autointent/_dump_tools/unit_dumpers.py index 2b1c201bb..3759ac8a0 100644 --- a/src/autointent/_dump_tools/unit_dumpers.py +++ b/src/autointent/_dump_tools/unit_dumpers.py @@ -3,7 +3,7 @@ import importlib import json import logging -from typing import TYPE_CHECKING, Any, TypeVar +from typing import TYPE_CHECKING, Any, TypeVar, cast import aiofiles import joblib @@ -215,7 +215,8 @@ def dump(obj: PeftModel, path: Path, exists_ok: bool) -> None: ptuning_path = path / "ptuning" ptuning_path.mkdir(parents=True, exist_ok=exists_ok) obj.save_pretrained(str(ptuning_path / "peft")) - obj.base_model.save_pretrained(ptuning_path / "base_model") + base_model = cast("PreTrainedModel", obj.base_model) + base_model.save_pretrained(ptuning_path / "base_model") else: # strategy to save lora models: merge adapters and save as usual hugging face model lora_path = path / "lora" diff --git a/src/autointent/_presets/zero-shot-typesafe.yaml b/src/autointent/_presets/zero-shot-typesafe.yaml new file mode 100644 index 000000000..3427a1b21 --- /dev/null +++ b/src/autointent/_presets/zero-shot-typesafe.yaml @@ -0,0 +1,24 @@ +search_space: + - node_type: scoring + target_metric: scoring_f1 + search_space: + - module_name: description_typesafe + question_type: [choice, noul] + temperature: + low: 0.3 + high: 3.0 + - node_type: decision + target_metric: decision_accuracy + search_space: + - module_name: threshold + thresh: + low: 0.1 + high: 0.9 + - module_name: argmax + - module_name: jinoos + - module_name: tunable + - module_name: adaptive +hpo_config: + sampler: tpe + n_trials: 30 + n_startup_trials: 10 diff --git a/src/autointent/advisor/_estimates/_resource.py b/src/autointent/advisor/_estimates/_resource.py index 5de6a4407..dffcdadb7 100644 --- a/src/autointent/advisor/_estimates/_resource.py +++ b/src/autointent/advisor/_estimates/_resource.py @@ -106,13 +106,14 @@ def _embedder_model_name(embedder: EmbedderConfig) -> str | None: "description_bi", "description_cross", "description_llm", + "description_typesafe", }, ) # Cache-honoring modules whose per-entry estimate already bundles the embedder # forward into `time_hours` (vs. classic linear/catboost which don't). _EMBEDDER_FORWARD_TRANSFORMER_MODULES = frozenset( - {"knn", "mlknn", "retrieval", "description_bi", "description_cross", "description_llm"}, + {"knn", "mlknn", "retrieval", "description_bi", "description_cross", "description_llm", "description_typesafe"}, ) @@ -608,6 +609,7 @@ def _emit_resource_findings( "description_bi", "description_cross", "description_llm", + "description_typesafe", }, ) diff --git a/src/autointent/advisor/_runner.py b/src/autointent/advisor/_runner.py index 55c74cbd8..340d70d32 100644 --- a/src/autointent/advisor/_runner.py +++ b/src/autointent/advisor/_runner.py @@ -274,7 +274,7 @@ def _data_phase( ) # partial descriptions x description scorer - description_modules = {"description_bi", "description_cross", "description_llm"} + description_modules = {"description_bi", "description_cross", "description_llm", "description_typesafe"} has_description = any(e.get("module_name") in description_modules for _, e in _walk_modules(search_space)) if has_description and stats.has_descriptions is False: report.add( diff --git a/src/autointent/advisor/_workflows.py b/src/autointent/advisor/_workflows.py index 5ed6c53ec..9acbd15ee 100644 --- a/src/autointent/advisor/_workflows.py +++ b/src/autointent/advisor/_workflows.py @@ -45,6 +45,7 @@ "transformers-light", "nn-heavy", "zero-shot-llm", + "zero-shot-typesafe", "nn-medium", "classic-heavy", "transformers-no-hpo", diff --git a/src/autointent/custom_types/_types.py b/src/autointent/custom_types/_types.py index f322001d1..4211b9f99 100644 --- a/src/autointent/custom_types/_types.py +++ b/src/autointent/custom_types/_types.py @@ -126,6 +126,7 @@ class Split: "transformers-light", "transformers-no-hpo", "zero-shot-llm", + "zero-shot-typesafe", "zero-shot-encoders", ] """Bundled search-space presets that our library supports. diff --git a/src/autointent/generation/_cache.py b/src/autointent/generation/_cache.py index 87ae93163..899956c48 100644 --- a/src/autointent/generation/_cache.py +++ b/src/autointent/generation/_cache.py @@ -55,18 +55,26 @@ def _remove_cache_entry(path: Path) -> None: shutil.rmtree(path, ignore_errors=True) -class StructuredOutputCache: - """Cache for structured output results.""" +class PydanticDiskCache: + """Key-agnostic memory + disk cache of pydantic models under ``user_cache_dir("autointent")/``. - def __init__(self, use_cache: bool = True, max_workers: int | None = None, batch_size: int = 100) -> None: + Each entry is a directory written by :class:`PydanticModelDumper`. Subclasses (or callers) + decide how keys are derived; this class only stores and loads by key. + """ + + def __init__( + self, subdir: str, use_cache: bool = True, max_workers: int | None = None, batch_size: int = 100 + ) -> None: """Initialize the cache. Args: + subdir: Directory name under the autointent user cache dir that holds this cache's entries. use_cache: Whether to use caching. max_workers: Maximum number of worker threads for parallel loading. If None, uses min(32, os.cpu_count() + 4). batch_size: Number of cache files to process in each batch. """ + self.subdir = subdir self.use_cache = use_cache self._memory_cache: dict[str, BaseModel] = {} self.max_workers = max_workers @@ -75,9 +83,17 @@ def __init__(self, use_cache: bool = True, max_workers: int | None = None, batch if self.use_cache: self._load_existing_cache() + def _cache_dir(self) -> Path: + """Directory holding this cache's entries (resolved at call time so tests can redirect it).""" + return Path(user_cache_dir("autointent")) / self.subdir + + def _entry_path(self, key: str) -> Path: + """On-disk directory of one entry.""" + return self._cache_dir() / key + def _load_existing_cache(self) -> None: """Load all existing cache items from disk into memory.""" - cache_dir = Path(user_cache_dir("autointent")) / "structured_outputs" + cache_dir = self._cache_dir() if not cache_dir.exists(): return @@ -138,34 +154,6 @@ def _load_single_cache_file(self, cache_file: Path) -> tuple[str, BaseModel] | N return None - def _get_cache_key( - self, - messages: list[Message], - output_model: type[T], - generation_params: dict[str, Any], - model_name: str, - base_url: str | None, - ) -> str: - """Generate a cache key for the given parameters. - - Args: - messages: List of messages to send to the model. - output_model: Pydantic model class to parse the response into. - generation_params: Generation parameters. - model_name: Name of the language model that will serve the request. - base_url: Base URL of the API endpoint, or None for the default. - - Returns: - Cache key as a hexadecimal string. - """ - hasher = Hasher() - hasher.update(json.dumps(messages)) - hasher.update(json.dumps(output_model.model_json_schema())) - hasher.update(json.dumps(generation_params)) - hasher.update(model_name) - hasher.update(base_url) - return hasher.hexdigest() - def _check_memory_cache(self, cache_key: str, output_model: type[T]) -> T | None: """Check if the result is available in memory cache. @@ -196,7 +184,7 @@ def _load_from_disk(self, cache_key: str, output_model: type[T]) -> T | None: Returns: Cached result if available and valid, None otherwise. """ - cache_path = _get_structured_output_cache_path(cache_key) + cache_path = self._entry_path(cache_key) if cache_path.exists(): try: @@ -216,17 +204,135 @@ def _load_from_disk(self, cache_key: str, output_model: type[T]) -> T | None: return None - def _save_to_disk(self, cache_key: str, result: T) -> None: + def _save_to_disk(self, cache_key: str, result: BaseModel) -> None: """Save result to disk cache. Args: cache_key: The cache key to use. result: The result to cache. """ - cache_path = _get_structured_output_cache_path(cache_key) + cache_path = self._entry_path(cache_key) cache_path.parent.mkdir(parents=True, exist_ok=True) PydanticModelDumper.dump(result, cache_path, exists_ok=True) + async def _load_from_disk_async(self, cache_key: str, output_model: type[T]) -> T | None: + """Load cached result from disk asynchronously. + + Args: + cache_key: The cache key to look up. + output_model: Pydantic model class to parse the response into. + + Returns: + Cached result if available and valid, None otherwise. + """ + cache_path = self._entry_path(cache_key) + + if cache_path.exists(): + try: + cached_data = await PydanticModelDumper.load_async(cache_path) + + if isinstance(cached_data, output_model): + logger.debug("Using cached structured output from disk for key: %s", cache_key) + # Add to memory cache for future access + self._memory_cache[cache_key] = cached_data + return cached_data + + logger.warning("Cached data type mismatch on disk, removing invalid cache") + _remove_cache_entry(cache_path) + except (ValidationError, ImportError) as e: + logger.warning("Failed to load cached structured output from disk: %s", e) + _remove_cache_entry(cache_path) + + return None + + async def _save_to_disk_async(self, cache_key: str, result: BaseModel) -> None: + """Save result to disk cache asynchronously. + + Args: + cache_key: The cache key to use. + result: The result to cache. + """ + cache_path = self._entry_path(cache_key) + cache_path.parent.mkdir(parents=True, exist_ok=True) + await PydanticModelDumper.dump_async(result, cache_path, exists_ok=True) + + def get_by_key(self, key: str, model_type: type[T]) -> T | None: + """Return the cached model for ``key`` if present and of type ``model_type``.""" + if not self.use_cache: + return None + memory_result = self._check_memory_cache(key, model_type) + if memory_result is not None: + return memory_result + return self._load_from_disk(key, model_type) + + def set_by_key(self, key: str, result: BaseModel) -> None: + """Store ``result`` under ``key`` in memory and on disk.""" + if not self.use_cache: + return + self._memory_cache[key] = result + self._save_to_disk(key, result) + logger.debug("Cached %s for key: %s (memory and disk)", type(result).__name__, key) + + async def get_by_key_async(self, key: str, model_type: type[T]) -> T | None: + """Async variant of :meth:`get_by_key`.""" + if not self.use_cache: + return None + memory_result = self._check_memory_cache(key, model_type) + if memory_result is not None: + return memory_result + return await self._load_from_disk_async(key, model_type) + + async def set_by_key_async(self, key: str, result: BaseModel) -> None: + """Async variant of :meth:`set_by_key`.""" + if not self.use_cache: + return + self._memory_cache[key] = result + await self._save_to_disk_async(key, result) + logger.debug("Cached %s for key: %s (memory and disk)", type(result).__name__, key) + + +class StructuredOutputCache(PydanticDiskCache): + """Cache for structured output results, keyed by prompt, schema, generation params and model.""" + + def __init__(self, use_cache: bool = True, max_workers: int | None = None, batch_size: int = 100) -> None: + """Initialize the cache. + + Args: + use_cache: Whether to use caching. + max_workers: Maximum number of worker threads for parallel loading. + If None, uses min(32, os.cpu_count() + 4). + batch_size: Number of cache files to process in each batch. + """ + super().__init__("structured_outputs", use_cache=use_cache, max_workers=max_workers, batch_size=batch_size) + + def _get_cache_key( + self, + messages: list[Message], + output_model: type[T], + generation_params: dict[str, Any], + model_name: str, + base_url: str | None, + ) -> str: + """Generate a cache key for the given parameters. + + Args: + messages: List of messages to send to the model. + output_model: Pydantic model class to parse the response into. + generation_params: Generation parameters. + model_name: Name of the language model that will serve the request. + base_url: Base URL of the API endpoint, or None for the default. + + Returns: + Cache key as a hexadecimal string. + """ + hasher = Hasher() + hasher.update(json.dumps(messages)) + hasher.update(json.dumps(output_model.model_json_schema())) + hasher.update(json.dumps(generation_params)) + hasher.update(model_name) + hasher.update(base_url) + return hasher.hexdigest() + def get( self, messages: list[Message], @@ -249,16 +355,9 @@ def get( """ if not self.use_cache: return None - - cache_key = self._get_cache_key(messages, output_model, generation_params, model_name, base_url) - - # First check in-memory cache - memory_result = self._check_memory_cache(cache_key, output_model) - if memory_result is not None: - return memory_result - - # Fallback to disk cache - return self._load_from_disk(cache_key, output_model) + return self.get_by_key( + self._get_cache_key(messages, output_model, generation_params, model_name, base_url), output_model + ) def set( self, @@ -281,56 +380,7 @@ def set( """ if not self.use_cache: return - - cache_key = self._get_cache_key(messages, output_model, generation_params, model_name, base_url) - - # Store in memory cache - self._memory_cache[cache_key] = result - - # Store in disk cache - self._save_to_disk(cache_key, result) - logger.debug("Cached structured output for key: %s (memory and disk)", cache_key) - - async def _load_from_disk_async(self, cache_key: str, output_model: type[T]) -> T | None: - """Load cached result from disk asynchronously. - - Args: - cache_key: The cache key to look up. - output_model: Pydantic model class to parse the response into. - - Returns: - Cached result if available and valid, None otherwise. - """ - cache_path = _get_structured_output_cache_path(cache_key) - - if cache_path.exists(): - try: - cached_data = await PydanticModelDumper.load_async(cache_path) - - if isinstance(cached_data, output_model): - logger.debug("Using cached structured output from disk for key: %s", cache_key) - # Add to memory cache for future access - self._memory_cache[cache_key] = cached_data - return cached_data - - logger.warning("Cached data type mismatch on disk, removing invalid cache") - _remove_cache_entry(cache_path) - except (ValidationError, ImportError) as e: - logger.warning("Failed to load cached structured output from disk: %s", e) - _remove_cache_entry(cache_path) - - return None - - async def _save_to_disk_async(self, cache_key: str, result: T) -> None: - """Save result to disk cache asynchronously. - - Args: - cache_key: The cache key to use. - result: The result to cache. - """ - cache_path = _get_structured_output_cache_path(cache_key) - cache_path.parent.mkdir(parents=True, exist_ok=True) - await PydanticModelDumper.dump_async(result, cache_path, exists_ok=True) + self.set_by_key(self._get_cache_key(messages, output_model, generation_params, model_name, base_url), result) async def get_async( self, @@ -354,16 +404,9 @@ async def get_async( """ if not self.use_cache: return None - - cache_key = self._get_cache_key(messages, output_model, generation_params, model_name, base_url) - - # First check in-memory cache - memory_result = self._check_memory_cache(cache_key, output_model) - if memory_result is not None: - return memory_result - - # Fallback to disk cache - return await self._load_from_disk_async(cache_key, output_model) + return await self.get_by_key_async( + self._get_cache_key(messages, output_model, generation_params, model_name, base_url), output_model + ) async def set_async( self, @@ -386,12 +429,6 @@ async def set_async( """ if not self.use_cache: return - - cache_key = self._get_cache_key(messages, output_model, generation_params, model_name, base_url) - - # Store in memory cache - self._memory_cache[cache_key] = result - - # Store in disk cache - await self._save_to_disk_async(cache_key, result) - logger.debug("Cached structured output for key: %s (memory and disk)", cache_key) + await self.set_by_key_async( + self._get_cache_key(messages, output_model, generation_params, model_name, base_url), result + ) diff --git a/src/autointent/modules/__init__.py b/src/autointent/modules/__init__.py index f77d104b1..fe33644de 100644 --- a/src/autointent/modules/__init__.py +++ b/src/autointent/modules/__init__.py @@ -31,6 +31,7 @@ RerankScorer, RNNScorer, SklearnScorer, + TypeSafeDescriptionScorer, ) if TYPE_CHECKING: @@ -59,6 +60,7 @@ def _create_modules_dict(modules: list[type[T]]) -> dict[str, type[T]]: BiEncoderDescriptionScorer, CrossEncoderDescriptionScorer, LLMDescriptionScorer, + TypeSafeDescriptionScorer, RerankScorer, SklearnScorer, MLKnnScorer, diff --git a/src/autointent/modules/scoring/__init__.py b/src/autointent/modules/scoring/__init__.py index 9b08b9bf8..5feefe6b2 100644 --- a/src/autointent/modules/scoring/__init__.py +++ b/src/autointent/modules/scoring/__init__.py @@ -1,6 +1,11 @@ from ._bert import BertScorer from ._catboost import CatBoostScorer, FeaturesType -from ._description import BiEncoderDescriptionScorer, CrossEncoderDescriptionScorer, LLMDescriptionScorer +from ._description import ( + BiEncoderDescriptionScorer, + CrossEncoderDescriptionScorer, + LLMDescriptionScorer, + TypeSafeDescriptionScorer, +) from ._dnnc import DNNCScorer from ._gcn import GCNScorer from ._knn import KNNScorer, RerankScorer @@ -29,4 +34,5 @@ "RNNScorer", "RerankScorer", "SklearnScorer", + "TypeSafeDescriptionScorer", ] diff --git a/src/autointent/modules/scoring/_description/__init__.py b/src/autointent/modules/scoring/_description/__init__.py index 8e43187c3..3573b7957 100644 --- a/src/autointent/modules/scoring/_description/__init__.py +++ b/src/autointent/modules/scoring/_description/__init__.py @@ -1,5 +1,11 @@ from .bi_encoder import BiEncoderDescriptionScorer from .cross_encoder import CrossEncoderDescriptionScorer from .llm_encoder import LLMDescriptionScorer +from .typesafe import TypeSafeDescriptionScorer -__all__ = ["BiEncoderDescriptionScorer", "CrossEncoderDescriptionScorer", "LLMDescriptionScorer"] +__all__ = [ + "BiEncoderDescriptionScorer", + "CrossEncoderDescriptionScorer", + "LLMDescriptionScorer", + "TypeSafeDescriptionScorer", +] diff --git a/src/autointent/modules/scoring/_description/typesafe.py b/src/autointent/modules/scoring/_description/typesafe.py new file mode 100644 index 000000000..f474a632f --- /dev/null +++ b/src/autointent/modules/scoring/_description/typesafe.py @@ -0,0 +1,416 @@ +"""TypeSafeDescriptionScorer: zero-shot intent scoring with TypeSafe System One models (jev).""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +from functools import partial +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal + +import aiometer +import numpy as np +import scipy +from dotenv import load_dotenv +from pydantic import BaseModel, PositiveFloat, PositiveInt + +from autointent import Context +from autointent._deps import require +from autointent._dump_tools import Dumper +from autointent._hash import Hasher +from autointent.generation._cache import PydanticDiskCache + +from .base import BaseDescriptionScorer + +if TYPE_CHECKING: + from numpy.typing import NDArray + + from autointent.configs import CrossEncoderConfig, EmbedderConfig + +logger = logging.getLogger(__name__) + +load_dotenv() + +CACHE_SUBDIR = "typesafe_answers" +CHOICE_KEY = "intent" +DEFAULT_MODEL = "jev-latest" +DEFAULT_MODEL_ENV = "TYPESAFE_DEFAULT_MODEL" +DEFAULT_BASE_URL = "https://api.typesafe.ai" +BASE_URL_ENV = "TYPESAFE_BASE_URL" +MAX_CHOICE_OPTIONS = 255 +"""The TypeSafe API's limit on ``choice`` options; above this, ``question_type='noul'`` must be used instead.""" +_EPS = 1e-6 +_CLIENT_ERROR_STATUS_MIN = 400 +_SERVER_ERROR_STATUS_MIN = 500 +_RETRYABLE_CLIENT_ERROR_STATUSES = frozenset({408, 429}) # request timeout, rate limit: transient, like a 5xx + +QuestionType = Literal["choice", "noul"] +_Result = tuple[list[float], int, int, str | None] | Exception +"""Either ``(probability_row, input_tokens, output_tokens, model)`` or the failure that produced no row.""" + + +def _is_fatal(error: BaseException) -> bool: + """Whether an API error means the request itself is wrong rather than a transient failure. + + A 4xx status other than 408 (timeout) or 429 (rate limit) means a wrong API key, model name, + or malformed question, none of which a retry or a uniform-row fallback would fix or even + surface; degrading it to a uniform row (as ``_ask_one_*`` do for everything else, mirroring + the LLM scorer) would silently hide a misconfiguration behind plausible-looking scores. 408 + and 429 behave like a 5xx (transient) and fall back like any other failure. + """ + status = getattr(error, "status", None) + return ( + isinstance(status, int) + and _CLIENT_ERROR_STATUS_MIN <= status < _SERVER_ERROR_STATUS_MIN + and status not in _RETRYABLE_CLIENT_ERROR_STATUSES + ) + + +class TypeSafeAnswer(BaseModel): + """Per-utterance probabilities in intent-description order, as stored in the disk cache.""" + + probabilities: list[float] + model: str | None = None + + +def build_questions(question_type: QuestionType, descriptions: list[str]) -> dict[str, dict[str, Any]]: + """Build the System One questions for one utterance as plain dicts (the SDK accepts them). + + ``choice`` asks a single question whose options are the descriptions; ``noul`` asks one + yes/no question per description. Option keys are ``intent_{i}`` in description order. + """ + if question_type == "choice": + return { + CHOICE_KEY: { + "type": "choice", + "instructions": "Which intent does `utterance` express?", + "criteria": {f"intent_{i}": description for i, description in enumerate(descriptions)}, + } + } + return { + f"intent_{i}": {"type": "noul", "instructions": f"Does `utterance` express this intent: {description}"} + for i, description in enumerate(descriptions) + } + + +def parse_answers(question_type: QuestionType, answers: Any, n_intents: int) -> list[float]: # noqa: ANN401 + """Read a response's ``answers`` mapping back into a probability row in description order.""" + if question_type == "choice": + probabilities = answers[CHOICE_KEY].probabilities + return [float(probabilities[f"intent_{i}"]) for i in range(n_intents)] + return [float(answers[f"intent_{i}"].noul) for i in range(n_intents)] + + +class TypeSafeDescriptionScorer(BaseDescriptionScorer): + """Zero-shot description scorer backed by TypeSafe's System One model (``jev``). + + Unlike :class:`LLMDescriptionScorer`, which prompts a chat model and buckets its answer, + this scorer asks the TypeSafe API typed questions and gets calibrated probabilities back: + + - ``question_type="choice"`` (multiclass): one ``Choice`` whose options are the intent + descriptions; the answer is a distribution over intents. + - ``question_type="noul"`` (multiclass or multilabel): one ``Noul`` (yes/no) per intent, + all in one request; each answer is the probability that the utterance expresses it. + + Probabilities are turned into ``log(p)`` (choice) or ``logit(p)`` (noul) "similarities", so + the base class's softmax / sigmoid with ``temperature=1`` reproduces the model's own + distribution and ``temperature`` sharpens or flattens it. Answers are cached on disk, so + repeated predictions on the same utterances (e.g. across HPO trials) cost nothing. + + Requires the ``typesafe`` extra (``pip install "autointent[typesafe]"``) and the + ``TYPESAFE_API_KEY`` environment variable. + + Args: + question_type: ``"choice"`` (single question, multiclass only) or ``"noul"`` (one yes/no per intent). + model: TypeSafe model name; ``None`` uses the SDK default (``jev-latest`` or ``TYPESAFE_DEFAULT_MODEL``). + temperature: Temperature for scaling the log/logit similarities (default: 1.0). + max_concurrent: Maximum concurrent requests; ``None`` uses the synchronous client (default: 15). + max_per_second: Rate limit for requests (default: 10). + max_retries: Retries per request, delegated to the SDK's ``RetryPolicy`` (default: 3). + use_cache: Cache answers on disk under the autointent cache dir (default: True). + multilabel: Flag indicating classification task type. + + Example: + -------- + .. code-block:: + + from autointent.modules.scoring import TypeSafeDescriptionScorer + + scorer = TypeSafeDescriptionScorer(question_type="choice") + + descriptions = [ + "User wants to book or reserve transportation like flights, trains, or hotels", + "User wants to cancel an existing booking or reservation", + "User asks about weather conditions or forecasts", + ] + scorer.fit([], [], descriptions) + + probabilities = scorer.predict(["Reserve a hotel room", "Delete my booking"]) + """ + + name = "description_typesafe" + + def __init__( + self, + question_type: QuestionType = "choice", + model: str | None = None, + temperature: PositiveFloat = 1.0, + max_concurrent: PositiveInt | None = 15, + max_per_second: PositiveInt = 10, + max_retries: PositiveInt = 3, + use_cache: bool = True, + multilabel: bool = False, + ) -> None: + super().__init__(temperature=temperature, multilabel=multilabel) + if question_type == "choice" and multilabel: + msg = "question_type='choice' picks a single intent and cannot express multilabel targets; use 'noul'" + raise ValueError(msg) + self.question_type: QuestionType = question_type + self.model = model + self.max_concurrent = max_concurrent + self.max_per_second = max_per_second + self.max_retries = max_retries + self.use_cache = use_cache + + @classmethod + def from_context( + cls, + context: Context, + question_type: QuestionType | None = None, + model: str | None = None, + temperature: PositiveFloat = 1.0, + max_concurrent: PositiveInt | None = 15, + max_per_second: PositiveInt = 10, + max_retries: PositiveInt = 3, + use_cache: bool = True, + ) -> TypeSafeDescriptionScorer: + """Build a scorer from pipeline context, defaulting/coercing ``question_type`` by task type. + + ``question_type=None`` defaults to ``"noul"`` on a multilabel dataset, ``"choice"`` otherwise. An + HPO-sampled ``question_type="choice"`` on a multilabel dataset (e.g. the ``zero-shot-typesafe`` + preset's search space) cannot express multilabel targets either, so it is downgraded to + ``"noul"`` with a warning instead of raising, so a single bad trial does not abort the whole + ``Pipeline.fit()``. An explicit ``TypeSafeDescriptionScorer(question_type="choice", multilabel=True)`` + still raises ``ValueError`` (see the constructor). + """ + multilabel = context.is_multilabel() + if question_type is None: + question_type = "noul" if multilabel else "choice" + if multilabel and question_type == "choice": + logger.warning("question_type='choice' cannot express multilabel targets; using 'noul' for this trial") + question_type = "noul" + return cls( + question_type=question_type, + model=model, + temperature=temperature, + max_concurrent=max_concurrent, + max_per_second=max_per_second, + max_retries=max_retries, + use_cache=use_cache, + multilabel=multilabel, + ) + + def get_implicit_initialization_params(self) -> dict[str, Any]: + return {"multilabel": self._multilabel, "question_type": self.question_type} + + @property + def resolved_model(self) -> str: + """Model name actually requested: explicit ``model``, else the SDK's default.""" + default = os.getenv(DEFAULT_MODEL_ENV, DEFAULT_MODEL) + return self.model or default + + @property + def resolved_base_url(self) -> str: + """Base URL the SDK client targets: ``TYPESAFE_BASE_URL`` if set, else TypeSafe's default.""" + return os.getenv(BASE_URL_ENV, DEFAULT_BASE_URL) + + def _fit_implementation(self, descriptions: list[str]) -> None: + """Reject ``choice`` past the API's option limit, then store descriptions and build runtime objects.""" + if self.question_type == "choice" and len(descriptions) > MAX_CHOICE_OPTIONS: + msg = ( + f"question_type='choice' supports at most {MAX_CHOICE_OPTIONS} options " + f"(got {len(descriptions)} descriptions); use question_type='noul' instead" + ) + raise ValueError(msg) + self._description_texts = descriptions + self._init_runtime() + + def _init_runtime(self) -> None: + """Create questions, SDK clients, the disk cache and the event loop from the stored config.""" + self._questions = build_questions(self.question_type, self._description_texts) + self._client, self._async_client = self._create_clients() + self._cache = PydanticDiskCache(CACHE_SUBDIR, use_cache=self.use_cache) + self._init_event_loop() + + def _create_clients(self) -> tuple[Any, Any]: + """Build ``(sync_client, async_client)``. + + The SDK is imported here, not at module level, so the module loads without the + ``typesafe`` extra; tests patch this method to inject fakes. + """ + require("typesafe") + from typesafe_sdk import AsyncTypeSafeClient, RetryPolicy, TypeSafeClient + + retry = RetryPolicy(max_retries=self.max_retries) + return TypeSafeClient(model=self.model, retry=retry), AsyncTypeSafeClient(model=self.model, retry=retry) + + def _init_event_loop(self) -> None: + if self.max_concurrent is not None: + try: + loop = asyncio.get_event_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + else: + if loop.is_closed(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + self._event_loop = loop + + def _cache_key(self, utterance: str) -> str: + """Hash everything that changes the answer: resolved model, base URL, questions, utterance.""" + hasher = Hasher() + hasher.update(self.resolved_model) + hasher.update(self.resolved_base_url) + hasher.update(json.dumps(self._questions, sort_keys=True)) + hasher.update(utterance) + return hasher.hexdigest() + + def _compute_similarities(self, utterances: list[str]) -> NDArray[np.float64]: + """Query (or load from cache) one probability row per utterance and map it to similarities.""" + if not (hasattr(self, "_description_texts") and hasattr(self, "_client")): + msg = "Scorer is not initialized. Call fit() before predict()." + raise RuntimeError(msg) + + n_intents = len(self._description_texts) + probabilities = np.full((len(utterances), n_intents), 1.0 / n_intents, dtype=np.float64) + + cache_keys = [self._cache_key(utterance) for utterance in utterances] + pending: list[int] = [] + for i, cache_key in enumerate(cache_keys): + cached = self._cache.get_by_key(cache_key, TypeSafeAnswer) + if cached is None: + pending.append(i) + else: + probabilities[i] = cached.probabilities + + input_tokens = output_tokens = 0 + seen_model: str | None = None + if pending: + results = self._ask_many([utterances[i] for i in pending]) + for i, result in zip(pending, results, strict=True): + if isinstance(result, Exception): + logger.warning( + "TypeSafe request failed for utterance %r; using uniform scores: %s", utterances[i], result + ) + continue + row, row_input_tokens, row_output_tokens, model = result + probabilities[i] = row + input_tokens += row_input_tokens + output_tokens += row_output_tokens + if model is not None: + seen_model = model + self._cache.set_by_key(cache_keys[i], TypeSafeAnswer(probabilities=row, model=model)) + + log_message = "TypeSafe predict: %d utterances, %d from cache, %d requests, %d input tokens, %d output tokens" + log_args: list[Any] = [ + len(utterances), + len(utterances) - len(pending), + len(pending), + input_tokens, + output_tokens, + ] + if seen_model is not None: + log_message += ", model %s" + log_args.append(seen_model) + logger.info(log_message, *log_args) + return self._to_similarities(probabilities) + + def _ask_many(self, utterances: list[str]) -> list[_Result]: + """Send one request per utterance: through aiometer when ``max_concurrent`` is set, else sequentially.""" + if self.max_concurrent is None: + return [self._ask_one_sync(utterance) for utterance in utterances] + task = aiometer.run_all( + [partial(self._ask_one_async, utterance) for utterance in utterances], + max_at_once=self.max_concurrent, + max_per_second=self.max_per_second, + ) + return self._event_loop.run_until_complete(task) + + def _ask_one_sync(self, utterance: str) -> _Result: + """Ask one question synchronously; re-raise a fatal API error, else return the row or the failure.""" + try: + response = self._client.system_one(state={"utterance": utterance}, questions=self._questions) + except Exception as e: # reason: any SDK/network failure degrades to a uniform row, like the LLM scorer + if _is_fatal(e): + raise + return e + return self._unpack(response) + + async def _ask_one_async(self, utterance: str) -> _Result: + """Ask one question asynchronously; re-raise a fatal API error, else return the row or the failure.""" + try: + response = await self._async_client.system_one(state={"utterance": utterance}, questions=self._questions) + except Exception as e: # reason: any SDK/network failure degrades to a uniform row, like the LLM scorer + if _is_fatal(e): + raise + return e + return self._unpack(response) + + def _unpack(self, response: Any) -> _Result: # noqa: ANN401 + """Turn an SDK response into ``(row, input_tokens, output_tokens, model)``; a malformed answer is a failure.""" + try: + row = parse_answers(self.question_type, response.answers, len(self._description_texts)) + except (KeyError, AttributeError, TypeError, ValueError) as e: + return e + usage = getattr(response, "usage", None) + model = getattr(response, "model", None) + return ( + row, + int(getattr(usage, "input_tokens", 0) or 0), + int(getattr(usage, "output_tokens", 0) or 0), + model if isinstance(model, str) else None, + ) + + def _to_similarities(self, probabilities: NDArray[np.float64]) -> NDArray[np.float64]: + """``log p`` for choice, ``logit p`` for noul, so the base class's scaling is the identity at T=1.""" + clipped = np.clip(probabilities, _EPS, 1.0 - _EPS) + if self.question_type == "choice": + return np.log(clipped) + return scipy.special.logit(clipped) # type: ignore[no-any-return] + + def clear_cache(self) -> None: + """Drop runtime objects (clients, disk-cache handle, event loop).""" + for attribute in ("_client", "_async_client", "_cache", "_event_loop"): + if hasattr(self, attribute): + delattr(self, attribute) + + def dump(self, path: str) -> None: + """Persist config and descriptions; clients, cache handle and questions are rebuilt on load. + + Simple attributes (``question_type``, ``model``, ``temperature``, concurrency settings, + ``use_cache``, ``_description_texts``) are written by the generic Dumper. The SDK clients + and the disk-cache handle are runtime objects it cannot serialize, so they are detached + for the duration of the dump; ``_questions`` is a dict and is excluded. + """ + detached = { + key: self.__dict__.pop(key) for key in ("_client", "_async_client", "_cache") if key in self.__dict__ + } + try: + Dumper.dump(self, Path(path), exclude=[asyncio.BaseEventLoop, dict]) + finally: + self.__dict__.update(detached) + + @classmethod + def load( + cls, + path: str, + embedder_config: EmbedderConfig | None = None, + cross_encoder_config: CrossEncoderConfig | None = None, + ) -> TypeSafeDescriptionScorer: + """Load a dumped scorer and rebuild its runtime objects (SDK clients, disk cache, event loop).""" + instance = super().load(path=path, embedder_config=embedder_config, cross_encoder_config=cross_encoder_config) + instance._init_runtime() # noqa: SLF001 + return instance diff --git a/tests/_fixtures/mock_typesafe.py b/tests/_fixtures/mock_typesafe.py new file mode 100644 index 000000000..d8d16df3a --- /dev/null +++ b/tests/_fixtures/mock_typesafe.py @@ -0,0 +1,83 @@ +"""Fake TypeSafe clients for tests where the API is incidental. + +``TypeSafeDescriptionScorer`` builds its SDK clients lazily in ``_create_clients`` so the +module imports without the ``typesafe`` extra. The fixture below patches that method to +return fakes whose ``system_one`` answers deterministically from the question dicts, and +redirects the autointent disk cache to ``tmp_path`` so runs never see each other's entries. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import TYPE_CHECKING, Any + +import pytest + +from autointent.modules.scoring._description.typesafe import CHOICE_KEY, TypeSafeDescriptionScorer + +if TYPE_CHECKING: + from pathlib import Path + +BEST_PROBABILITY = 0.7 +NOUL_YES = 0.9 +NOUL_NO = 0.1 + + +def expected_best_index(utterance: str, n_options: int) -> int: + """Mirror `fake_response`'s derivation of the "winning" option when `best_index` is omitted.""" + return sum(map(ord, utterance)) % n_options + + +def fake_response( + questions: dict[str, dict[str, Any]], utterance: str, best_index: int | None = None +) -> SimpleNamespace: + """Deterministic answers: one option gets 0.7 and the rest share 0.3 (choice), or 0.9 vs 0.1 (noul). + + ``best_index`` picks the "winning" option; when omitted it is derived from ``utterance`` (see + `expected_best_index`), so different utterances land on different options and the result order + is observable in tests instead of always being index 0. + """ + answers: dict[str, SimpleNamespace] = {} + if CHOICE_KEY in questions and questions[CHOICE_KEY]["type"] == "choice": + keys = list(questions[CHOICE_KEY]["criteria"]) + index = best_index if best_index is not None else expected_best_index(utterance, len(keys)) + rest = (1.0 - BEST_PROBABILITY) / max(len(keys) - 1, 1) + probabilities = {key: (BEST_PROBABILITY if i == index else rest) for i, key in enumerate(keys)} + answers[CHOICE_KEY] = SimpleNamespace(choice=keys[index], confidence=0.5, probabilities=probabilities) + else: + n_options = len(questions) + index = best_index if best_index is not None else expected_best_index(utterance, n_options) + for i, key in enumerate(questions): + answers[key] = SimpleNamespace(noul=NOUL_YES if i == index else NOUL_NO) + return SimpleNamespace(answers=answers, usage=SimpleNamespace(input_tokens=100, output_tokens=10), model="jev-test") + + +class FakeTypeSafeClient: + """Sync stand-in for ``typesafe_sdk.TypeSafeClient`` recording every utterance it was asked about.""" + + def __init__(self) -> None: + self.calls: list[str] = [] + + def system_one(self, state: dict[str, Any], questions: dict[str, dict[str, Any]], **_: Any) -> SimpleNamespace: + self.calls.append(state["utterance"]) + return fake_response(questions, state["utterance"]) + + +class FakeAsyncTypeSafeClient(FakeTypeSafeClient): + """Async stand-in for ``typesafe_sdk.AsyncTypeSafeClient``.""" + + async def system_one( # mypy's override check is disabled project-wide (disable_error_code) + self, state: dict[str, Any], questions: dict[str, dict[str, Any]], **_: Any + ) -> SimpleNamespace: + return FakeTypeSafeClient.system_one(self, state, questions) + + +@pytest.fixture +def patch_typesafe_scorer_client( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> tuple[FakeTypeSafeClient, FakeAsyncTypeSafeClient]: + """Make every ``TypeSafeDescriptionScorer`` use the fakes and an isolated disk cache.""" + monkeypatch.setattr("autointent.generation._cache.user_cache_dir", lambda *_: str(tmp_path / "cache")) + sync_client, async_client = FakeTypeSafeClient(), FakeAsyncTypeSafeClient() + monkeypatch.setattr(TypeSafeDescriptionScorer, "_create_clients", lambda _self: (sync_client, async_client)) + return sync_client, async_client diff --git a/tests/conftest.py b/tests/conftest.py index 0845a3e40..52947f11f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -350,5 +350,6 @@ def _guarded_api_model_info( mock_generator, patch_llm_scorer_generator, ) +from tests._fixtures.mock_typesafe import patch_typesafe_scorer_client # noqa: E402, F401 from tests._fixtures.opensearch_container import opensearch_container # noqa: E402, F401 from tests._fixtures.respx_openai import respx_openai # noqa: E402, F401 diff --git a/tests/generation/structured_output/test_cache_unit.py b/tests/generation/structured_output/test_cache_unit.py index 5c510fa4b..206fe5870 100644 --- a/tests/generation/structured_output/test_cache_unit.py +++ b/tests/generation/structured_output/test_cache_unit.py @@ -240,3 +240,55 @@ async def test_async_same_identity_still_hits_cache() -> None: hit = await cache.get_async(MESSAGES, CacheModel, PARAMS, model_name="gpt-4o", base_url=None) assert hit == result + + +# --- PydanticDiskCache (key-agnostic base) --------------------------------------------- + + +def test_disk_cache_by_key_roundtrip_memory_and_disk() -> None: + from autointent.generation._cache import PydanticDiskCache + + cache = PydanticDiskCache("unit_test_subdir") + cache.set_by_key("k1", CacheModel(name="a", value=1)) + assert cache.get_by_key("k1", CacheModel) == CacheModel(name="a", value=1) + + fresh = PydanticDiskCache("unit_test_subdir") + fresh._memory_cache.clear() + assert fresh.get_by_key("k1", CacheModel) == CacheModel(name="a", value=1) + + +def test_disk_cache_subdirs_are_isolated() -> None: + from autointent.generation._cache import PydanticDiskCache + + first = PydanticDiskCache("subdir_a") + first.set_by_key("shared-key", CacheModel(name="a", value=1)) + + second = PydanticDiskCache("subdir_b") + assert second.get_by_key("shared-key", CacheModel) is None + assert not second._memory_cache + assert first._entry_path("shared-key").parent.name == "subdir_a" + assert second._entry_path("shared-key").parent.name == "subdir_b" + + +def test_structured_output_cache_uses_structured_outputs_subdir() -> None: + cache = StructuredOutputCache(use_cache=True) + key = cache._get_cache_key(MESSAGES, CacheModel, PARAMS, MODEL_NAME, BASE_URL) + assert cache._entry_path(key) == _get_structured_output_cache_path(key) + + +def test_disk_cache_disabled_is_noop() -> None: + from autointent.generation._cache import PydanticDiskCache + + cache = PydanticDiskCache("unit_test_subdir", use_cache=False) + cache.set_by_key("k1", CacheModel(name="a", value=1)) + assert cache.get_by_key("k1", CacheModel) is None + + +@pytest.mark.asyncio +async def test_disk_cache_async_by_key_roundtrip() -> None: + from autointent.generation._cache import PydanticDiskCache + + cache = PydanticDiskCache("unit_test_subdir") + await cache.set_by_key_async("k2", CacheModel(name="b", value=2)) + cache._memory_cache.clear() + assert await cache.get_by_key_async("k2", CacheModel) == CacheModel(name="b", value=2) diff --git a/tests/modules/scoring/test_description_typesafe.py b/tests/modules/scoring/test_description_typesafe.py new file mode 100644 index 000000000..6e3011056 --- /dev/null +++ b/tests/modules/scoring/test_description_typesafe.py @@ -0,0 +1,386 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +import numpy as np +import pytest + +from autointent import Context +from autointent.configs import DataConfig +from autointent.context.data_handler import DataHandler +from autointent.modules.scoring import TypeSafeDescriptionScorer +from autointent.modules.scoring._description.typesafe import ( + CHOICE_KEY, + MAX_CHOICE_OPTIONS, + build_questions, + parse_answers, +) +from tests._fixtures.mock_typesafe import NOUL_NO, NOUL_YES, expected_best_index, fake_response +from tests._helpers import is_strict_labels + +if TYPE_CHECKING: + from pathlib import Path + + from autointent import Dataset + from autointent.modules.scoring._description.typesafe import QuestionType + from tests._fixtures.mock_typesafe import FakeAsyncTypeSafeClient, FakeTypeSafeClient + +TEST_UTTERANCES = ["What is the balance on my account?", "How do I reset my online banking password?"] + + +def _descriptions(data_handler: DataHandler) -> list[str]: + descriptions = data_handler.intent_descriptions + assert all(d is not None for d in descriptions) + return cast("list[str]", descriptions) + + +def test_build_questions_choice_and_noul() -> None: + descriptions = ["book a hotel", "check the weather"] + choice = build_questions("choice", descriptions) + assert list(choice) == [CHOICE_KEY] + assert choice[CHOICE_KEY]["type"] == "choice" + assert choice[CHOICE_KEY]["criteria"] == {"intent_0": "book a hotel", "intent_1": "check the weather"} + + noul = build_questions("noul", descriptions) + assert list(noul) == ["intent_0", "intent_1"] + assert all( + q["type"] == "noul" and desc in q["instructions"] for q, desc in zip(noul.values(), descriptions, strict=True) + ) + + +def test_parse_answers_orders_by_index() -> None: + questions = build_questions("choice", ["a", "b", "c"]) + response = fake_response(questions, "x", best_index=2) + assert parse_answers("choice", response.answers, 3) == pytest.approx([0.15, 0.15, 0.7]) + + questions = build_questions("noul", ["a", "b", "c"]) + response = fake_response(questions, "x", best_index=1) + assert parse_answers("noul", response.answers, 3) == pytest.approx([0.1, 0.9, 0.1]) + + +def test_choice_multiclass_reproduces_model_distribution_at_unit_temperature( + dataset: Dataset, patch_typesafe_scorer_client: tuple[FakeTypeSafeClient, FakeAsyncTypeSafeClient] +) -> None: + sync_client, _ = patch_typesafe_scorer_client + data_handler = DataHandler(dataset) + descriptions = _descriptions(data_handler) + + labels = data_handler.train_labels(0) + assert is_strict_labels(labels) + + scorer = TypeSafeDescriptionScorer(question_type="choice", temperature=1.0, max_concurrent=None) + scorer.fit(data_handler.train_utterances(0), labels, descriptions) + assert scorer._description_texts == descriptions + + probabilities = scorer.predict(TEST_UTTERANCES) + + assert probabilities.shape == (len(TEST_UTTERANCES), len(descriptions)) + np.testing.assert_allclose(probabilities.sum(axis=1), 1.0) + expected_response = fake_response(build_questions("choice", descriptions), TEST_UTTERANCES[0]) + expected = expected_response.answers[CHOICE_KEY].probabilities + np.testing.assert_allclose(probabilities[0], [expected[f"intent_{i}"] for i in range(len(descriptions))], atol=1e-5) + assert sync_client.calls == TEST_UTTERANCES + + +def test_noul_multilabel_reproduces_model_probabilities_at_unit_temperature( + dataset: Dataset, patch_typesafe_scorer_client: tuple[FakeTypeSafeClient, FakeAsyncTypeSafeClient] +) -> None: + data_handler = DataHandler(dataset.to_multilabel()) + descriptions = _descriptions(data_handler) + + labels = data_handler.train_labels(0) + assert is_strict_labels(labels) + + scorer = TypeSafeDescriptionScorer(question_type="noul", temperature=1.0, max_concurrent=None, multilabel=True) + scorer.fit(data_handler.train_utterances(0), labels, descriptions) + + probabilities = scorer.predict(TEST_UTTERANCES) + + assert probabilities.shape == (len(TEST_UTTERANCES), len(descriptions)) + winner = expected_best_index(TEST_UTTERANCES[0], len(descriptions)) + expected = [NOUL_YES if i == winner else NOUL_NO for i in range(len(descriptions))] + np.testing.assert_allclose(probabilities[0], expected, atol=1e-5) + + +def test_temperature_sharpens_choice_distribution( + dataset: Dataset, patch_typesafe_scorer_client: tuple[FakeTypeSafeClient, FakeAsyncTypeSafeClient] +) -> None: + data_handler = DataHandler(dataset) + descriptions = _descriptions(data_handler) + cold = TypeSafeDescriptionScorer(question_type="choice", temperature=0.5, max_concurrent=None) + cold.fit([], [], descriptions) + probabilities = cold.predict(TEST_UTTERANCES[:1]) + winner = expected_best_index(TEST_UTTERANCES[0], len(descriptions)) + assert probabilities[0, winner] > 0.7 # softmax(log p / 0.5) sharpens the 0.7 winner + + +def test_choice_with_multilabel_is_rejected() -> None: + with pytest.raises(ValueError, match="noul"): + TypeSafeDescriptionScorer(question_type="choice", multilabel=True) + + +def test_from_context_defaults_question_type_by_task(dataset: Dataset) -> None: + context = Context() + context.set_dataset(dataset, DataConfig(scheme="ho")) + assert TypeSafeDescriptionScorer.from_context(context).question_type == "choice" + + multilabel_context = Context() + multilabel_context.set_dataset(dataset.to_multilabel(), DataConfig(scheme="ho")) + scorer = TypeSafeDescriptionScorer.from_context(multilabel_context) + assert scorer.question_type == "noul" + assert scorer.get_implicit_initialization_params() == {"multilabel": True, "question_type": "noul"} + + +def test_from_context_coerces_hpo_sampled_choice_to_noul_on_multilabel( + dataset: Dataset, caplog: pytest.LogCaptureFixture +) -> None: + """An HPO trial can sample `question_type="choice"` from the preset's search space even on a + multilabel dataset; `from_context` must downgrade it to `noul` with a warning instead of + raising, so a single bad trial does not abort `Pipeline.fit()`. + """ + multilabel_context = Context() + multilabel_context.set_dataset(dataset.to_multilabel(), DataConfig(scheme="ho")) + + with caplog.at_level("WARNING"): + scorer = TypeSafeDescriptionScorer.from_context(multilabel_context, question_type="choice") + + assert scorer.question_type == "noul" + assert "question_type='choice'" in caplog.text + assert scorer.get_implicit_initialization_params() == {"multilabel": True, "question_type": "noul"} + + +def test_predict_before_fit_raises() -> None: + with pytest.raises(RuntimeError, match="fit"): + TypeSafeDescriptionScorer(max_concurrent=None).predict(["hello"]) + + +def test_async_path_uses_async_client_and_matches_sync( + dataset: Dataset, patch_typesafe_scorer_client: tuple[FakeTypeSafeClient, FakeAsyncTypeSafeClient] +) -> None: + sync_client, async_client = patch_typesafe_scorer_client + descriptions = _descriptions(DataHandler(dataset)) + + concurrent = TypeSafeDescriptionScorer(question_type="choice", max_concurrent=2, max_per_second=100) + concurrent.fit([], [], descriptions) + concurrent_probabilities = concurrent.predict(TEST_UTTERANCES) + assert async_client.calls == TEST_UTTERANCES + assert sync_client.calls == [] + + sequential = TypeSafeDescriptionScorer(question_type="choice", max_concurrent=None, use_cache=False) + sequential.fit([], [], descriptions) + sequential_probabilities = sequential.predict(TEST_UTTERANCES) + np.testing.assert_allclose(sequential_probabilities, concurrent_probabilities) + + # A swapped result order between utterances would slip past assert_allclose above if it + # happened to swap two utterances with identical distributions; check per-utterance winners. + expected_indices = [expected_best_index(utterance, len(descriptions)) for utterance in TEST_UTTERANCES] + np.testing.assert_array_equal(np.argmax(concurrent_probabilities, axis=1), expected_indices) + np.testing.assert_array_equal(np.argmax(sequential_probabilities, axis=1), expected_indices) + + +@pytest.mark.parametrize("max_concurrent", [None, 2]) +@pytest.mark.parametrize("question_type", ["choice", "noul"]) +def test_failed_request_yields_uniform_row_and_warns( + dataset: Dataset, + patch_typesafe_scorer_client: tuple[FakeTypeSafeClient, FakeAsyncTypeSafeClient], + question_type: QuestionType, + max_concurrent: int | None, + caplog: pytest.LogCaptureFixture, +) -> None: + """The uniform fallback is exercised for both the softmax (choice) and sigmoid (noul) paths.""" + sync_client, async_client = patch_typesafe_scorer_client + descriptions = _descriptions(DataHandler(dataset)) + + def boom(*_: object, **__: object) -> None: + msg = "simulated outage" + raise RuntimeError(msg) + + async def boom_async(*_: object, **__: object) -> None: + boom() + + sync_client.system_one = boom # type: ignore[assignment] + async_client.system_one = boom_async # type: ignore[assignment] + + multilabel = question_type == "noul" + scorer = TypeSafeDescriptionScorer( + question_type=question_type, max_concurrent=max_concurrent, multilabel=multilabel + ) + scorer.fit([], [], descriptions) + with caplog.at_level("WARNING"): + probabilities = scorer.predict(TEST_UTTERANCES) + + np.testing.assert_allclose(probabilities, 1.0 / len(descriptions)) + assert "simulated outage" in caplog.text + + +class FakeAPIError(Exception): + """Stand-in for typesafe_sdk's HTTP error types, which carry `.status` (see `_is_fatal`).""" + + def __init__(self, status: int) -> None: + super().__init__(f"api error, status={status}") + self.status = status + + +@pytest.mark.parametrize("max_concurrent", [None, 2]) +def test_fatal_status_propagates_but_transient_status_falls_back( + dataset: Dataset, + patch_typesafe_scorer_client: tuple[FakeTypeSafeClient, FakeAsyncTypeSafeClient], + max_concurrent: int | None, +) -> None: + """A 401 (bad key/misconfiguration) must not be swallowed into a uniform row; a 503 still is.""" + sync_client, async_client = patch_typesafe_scorer_client + descriptions = _descriptions(DataHandler(dataset)) + + def boom_401(*_: object, **__: object) -> None: + raise FakeAPIError(401) + + async def boom_401_async(*_: object, **__: object) -> None: + boom_401() + + sync_client.system_one = boom_401 # type: ignore[assignment] + async_client.system_one = boom_401_async # type: ignore[assignment] + + fatal_scorer = TypeSafeDescriptionScorer(question_type="choice", max_concurrent=max_concurrent) + fatal_scorer.fit([], [], descriptions) + with pytest.raises(FakeAPIError): + fatal_scorer.predict(TEST_UTTERANCES) + + def boom_503(*_: object, **__: object) -> None: + raise FakeAPIError(503) + + async def boom_503_async(*_: object, **__: object) -> None: + boom_503() + + sync_client.system_one = boom_503 # type: ignore[assignment] + async_client.system_one = boom_503_async # type: ignore[assignment] + + transient_scorer = TypeSafeDescriptionScorer(question_type="choice", max_concurrent=max_concurrent) + transient_scorer.fit([], [], descriptions) + probabilities = transient_scorer.predict(TEST_UTTERANCES) + np.testing.assert_allclose(probabilities, 1.0 / len(descriptions)) + + +def test_fit_rejects_too_many_choice_options() -> None: + descriptions = [f"description {i}" for i in range(MAX_CHOICE_OPTIONS + 1)] + scorer = TypeSafeDescriptionScorer(question_type="choice") + with pytest.raises(ValueError, match="noul"): + scorer.fit([], [], descriptions) + + +def test_cache_hit_skips_the_client( + dataset: Dataset, patch_typesafe_scorer_client: tuple[FakeTypeSafeClient, FakeAsyncTypeSafeClient] +) -> None: + sync_client, _ = patch_typesafe_scorer_client + descriptions = _descriptions(DataHandler(dataset)) + + scorer = TypeSafeDescriptionScorer(question_type="choice", max_concurrent=None) + scorer.fit([], [], descriptions) + first = scorer.predict(TEST_UTTERANCES) + assert len(sync_client.calls) == len(TEST_UTTERANCES) + + second = scorer.predict(TEST_UTTERANCES) + assert len(sync_client.calls) == len(TEST_UTTERANCES) + np.testing.assert_allclose(first, second) + + # A fresh instance (new HPO trial) with the same descriptions also hits the disk cache. + other = TypeSafeDescriptionScorer(question_type="choice", max_concurrent=None, temperature=2.0) + other.fit([], [], descriptions) + other.predict(TEST_UTTERANCES) + assert len(sync_client.calls) == len(TEST_UTTERANCES) + + # A different question_type is a different key. + noul = TypeSafeDescriptionScorer(question_type="noul", max_concurrent=None) + noul.fit([], [], descriptions) + noul.predict(TEST_UTTERANCES) + assert len(sync_client.calls) == 2 * len(TEST_UTTERANCES) + + +def test_cache_key_changes_with_model( + dataset: Dataset, patch_typesafe_scorer_client: tuple[FakeTypeSafeClient, FakeAsyncTypeSafeClient] +) -> None: + """A different `model` must be a different cache key, else a cached answer from one model would + be served to another (the model determines the answer, so the key must cover it). + """ + sync_client, _ = patch_typesafe_scorer_client + descriptions = _descriptions(DataHandler(dataset)) + + scorer = TypeSafeDescriptionScorer(question_type="choice", max_concurrent=None) + scorer.fit([], [], descriptions) + scorer.predict(TEST_UTTERANCES) + assert len(sync_client.calls) == len(TEST_UTTERANCES) + + other_model = TypeSafeDescriptionScorer(question_type="choice", max_concurrent=None, model="jev-preview") + other_model.fit([], [], descriptions) + other_model.predict(TEST_UTTERANCES) + assert len(sync_client.calls) == 2 * len(TEST_UTTERANCES) + + +def test_predict_logs_usage( + dataset: Dataset, + patch_typesafe_scorer_client: tuple[FakeTypeSafeClient, FakeAsyncTypeSafeClient], + caplog: pytest.LogCaptureFixture, +) -> None: + descriptions = _descriptions(DataHandler(dataset)) + scorer = TypeSafeDescriptionScorer(question_type="choice", max_concurrent=None) + scorer.fit([], [], descriptions) + with caplog.at_level("INFO", logger="autointent.modules.scoring._description.typesafe"): + scorer.predict(TEST_UTTERANCES) + assert "2 requests" in caplog.text + assert "200 input tokens" in caplog.text + assert "model jev-test" in caplog.text + + +@pytest.mark.parametrize("max_concurrent", [None, 2]) +def test_dump_and_load_roundtrip( + dataset: Dataset, + patch_typesafe_scorer_client: tuple[FakeTypeSafeClient, FakeAsyncTypeSafeClient], + tmp_path: Path, + max_concurrent: int | None, +) -> None: + descriptions = _descriptions(DataHandler(dataset)) + scorer = TypeSafeDescriptionScorer( + question_type="noul", + model="jev-preview", + temperature=0.7, + max_concurrent=max_concurrent, + max_retries=5, + use_cache=False, + ) + scorer.fit([], [], descriptions) + before = scorer.predict(TEST_UTTERANCES) + + dump_dir = tmp_path / "dump" + scorer.dump(str(dump_dir)) + loaded = TypeSafeDescriptionScorer.load(str(dump_dir)) + + assert loaded.question_type == "noul" + assert loaded.model == "jev-preview" + assert loaded.temperature == pytest.approx(0.7) + assert loaded.max_concurrent == max_concurrent + assert loaded.max_retries == 5 + assert loaded.use_cache is False + assert loaded._description_texts == descriptions + np.testing.assert_allclose(loaded.predict(TEST_UTTERANCES), before) + + # dump() must leave the live instance usable + np.testing.assert_allclose(scorer.predict(TEST_UTTERANCES), before) + + +def test_registered_in_scoring_modules() -> None: + from autointent.modules import SCORING_MODULES + + assert SCORING_MODULES["description_typesafe"] is TypeSafeDescriptionScorer + + +def test_create_clients_builds_real_sdk_clients(monkeypatch: pytest.MonkeyPatch) -> None: + """Exercise the real SDK seam once (unmocked): `_create_clients` must build actual SDK client + instances, not just satisfy the fake used everywhere else in this file. Requires the + ``typesafe`` extra; skipped when it isn't installed. Makes no network call. + """ + typesafe_sdk = pytest.importorskip("typesafe_sdk") + monkeypatch.setenv("TYPESAFE_API_KEY", "dummy-key-for-tests") + + sync_client, async_client = TypeSafeDescriptionScorer(model="jev-1.13.0", max_retries=2)._create_clients() + + assert isinstance(sync_client, typesafe_sdk.TypeSafeClient) + assert isinstance(async_client, typesafe_sdk.AsyncTypeSafeClient) diff --git a/tests/pipeline/test_presets.py b/tests/pipeline/test_presets.py index 737ef3ebd..211796fff 100644 --- a/tests/pipeline/test_presets.py +++ b/tests/pipeline/test_presets.py @@ -14,6 +14,7 @@ from autointent import Dataset from autointent.generation import Generator + from tests._fixtures.mock_typesafe import FakeAsyncTypeSafeClient, FakeTypeSafeClient @pytest.mark.parametrize( @@ -28,10 +29,17 @@ pytest.param("transformers-light", marks=pytest.mark.transformers), pytest.param("transformers-no-hpo", marks=pytest.mark.transformers), "zero-shot-llm", + "zero-shot-typesafe", "zero-shot-encoders", ], ) -def test_presets(dataset: Dataset, preset: str, patch_llm_scorer_generator: Generator, tmp_path: Path) -> None: +def test_presets( + dataset: Dataset, + preset: str, + patch_llm_scorer_generator: Generator, + patch_typesafe_scorer_client: tuple[FakeTypeSafeClient, FakeAsyncTypeSafeClient], + tmp_path: Path, +) -> None: project_dir = tmp_path pipeline_optimizer = Pipeline.from_preset(preset) # type: ignore[arg-type] # reason: parametrize values are runtime strings; mypy can't narrow to the SearchSpacePreset Literal @@ -44,6 +52,29 @@ def test_presets(dataset: Dataset, preset: str, patch_llm_scorer_generator: Gene pipeline_optimizer.fit(dataset, refit_after=False) +def test_zero_shot_typesafe_preset_on_multilabel( + dataset: Dataset, + patch_typesafe_scorer_client: tuple[FakeTypeSafeClient, FakeAsyncTypeSafeClient], + tmp_path: Path, +) -> None: + """The preset's `question_type: [choice, noul]` search space must not crash on a multilabel dataset. + + HPO can sample `choice` for the `description_typesafe` module even though the dataset is + multilabel; `from_context` must downgrade it to `noul` (with a warning) instead of raising, + so the whole `Pipeline.fit()` doesn't abort on the first such trial. + """ + project_dir = tmp_path + + pipeline_optimizer = Pipeline.from_preset("zero-shot-typesafe") + apply_test_models(pipeline_optimizer) + + pipeline_optimizer.set_config(LoggingConfig(project_dir=project_dir, dump_modules=True, clear_ram=True)) + pipeline_optimizer.set_config(DataConfig(scheme="ho")) + pipeline_optimizer.set_config(HPOConfig(timeout=60)) # limit budget time because we want tests to be fast + + pipeline_optimizer.fit(dataset.to_multilabel(), refit_after=False) + + def test_apply_test_models_retargets_pipeline_slots() -> None: from autointent import Pipeline from tests.conftest import (