From 117b19621156abaae16d8111babaf7683b7e1ba9 Mon Sep 17 00:00:00 2001 From: bghira Date: Tue, 8 Sep 2026 14:09:49 -0600 Subject: [PATCH] feat: validate and recover structured captions --- examples/worker.openai-compatible.yaml | 20 ++ recipes/STRUCTURED_JSON.md | 104 ++++++++++ src/caption_flow/models.py | 1 + src/caption_flow/utils/vllm_config.py | 1 + src/caption_flow/workers/caption.py | 10 + src/caption_flow/workers/openai_compatible.py | 186 ++++++++++++++++- tests/test_openai_compatible_worker.py | 190 +++++++++++++++++- tests/test_vllm_config.py | 1 + tests/test_worker_caption.py | 23 +++ 9 files changed, 529 insertions(+), 7 deletions(-) create mode 100644 recipes/STRUCTURED_JSON.md diff --git a/examples/worker.openai-compatible.yaml b/examples/worker.openai-compatible.yaml index f39d232..ac63149 100644 --- a/examples/worker.openai-compatible.yaml +++ b/examples/worker.openai-compatible.yaml @@ -20,6 +20,21 @@ worker: # - "i cannot describe" # - "unable to provide a caption" + # Structured-output validation is optional. When enabled, invalid JSON is + # treated like an empty/refused result and can use inference.retry_prompt. + # The remaining transforms require validate_json_output: true. + # validate_json_output: true + # repair_invalid_json_escapes: true + # canonicalize_json_output: true + # normalize_yxyx_bboxes: true + # deduplicate_json_elements: true + + # Provider-specific fields applied only to fallback requests. This is + # useful when the retry needs a stricter or smaller response schema. + # retry_extra_body: + # response_format: + # type: "json_object" + endpoints: - name: "zai-coding-plan" base_url: "https://api.z.ai/api/coding/paas/v4" @@ -33,6 +48,11 @@ worker: # requests_per_minute: 600 timeout_seconds: 120 max_retries: 6 + # Provider-specific fields shared by primary and fallback requests can + # be placed here, including response_format or chat-template options. + # extra_body: + # response_format: + # type: "json_object" # Add more accounts or providers to aggregate their capacity. Model names # can differ between endpoints. diff --git a/recipes/STRUCTURED_JSON.md b/recipes/STRUCTURED_JSON.md new file mode 100644 index 0000000..0ba17fe --- /dev/null +++ b/recipes/STRUCTURED_JSON.md @@ -0,0 +1,104 @@ +# Structured JSON captions + +CaptionFlow's OpenAI-compatible worker can request provider-constrained JSON, +reject malformed responses, and retry failed items with a smaller schema or +different sampling settings. This works for Ideogram-style captions and other +JSON caption formats without coupling CaptionFlow to one schema. + +## Configure the worker + +Provider-specific request fields remain local to the worker. Put the primary +structured-output constraint in an endpoint's `extra_body` and enable local +validation on the worker: + +```yaml +worker: + server: "ws://orchestrator.example:8765" + token: "replace-with-captionflow-worker-token" + + openai_compatible: + validate_json_output: true + canonicalize_json_output: true + + # Optional repairs for common model failures. Bounding boxes are expected + # in [ymin, xmin, ymax, xmax] order. + repair_invalid_json_escapes: true + normalize_yxyx_bboxes: true + deduplicate_json_elements: true + + endpoints: + - name: "local-vllm" + base_url: "http://inference.example:8000/v1" + api_key_env: "VLLM_API_KEY" + model: "vision-model" + initial_concurrency: 16 + max_concurrency: 16 + extra_body: + response_format: + type: "json_schema" + json_schema: + name: "caption" + strict: true + schema: + type: "object" + additionalProperties: false + required: ["description"] + properties: + description: {type: "string"} +``` + +`validate_json_output` checks strict JSON syntax, including rejecting +non-standard `NaN` and `Infinity` values. It does not itself implement JSON +Schema validation; use the endpoint's `response_format` when the provider +supports constrained decoding. The optional transforms require +`validate_json_output: true`. + +`repair_invalid_json_escapes` only escapes backslashes that cannot begin a +valid JSON escape. `normalize_yxyx_bboxes` corrects inverted coordinate pairs; +it does not infer boxes, clamp their range, or assess localization quality. +`deduplicate_json_elements` removes exact duplicate objects from arrays named +`elements`. + +## Configure recovery + +The orchestrator owns the shared prompt and standard sampling settings. A +response that is empty, refused, rejected by a provider safety filter, or +invalid JSON can receive one semantic retry: + +```yaml +orchestrator: + inference: + inference_prompts: + - >- + Describe the image as one JSON object matching the requested schema. + Return JSON only. + sampling: + temperature: 0.3 + max_tokens: 2048 + retry_prompt: >- + Try again. Return exactly one valid JSON object with a concise, + visually grounded description and no surrounding prose. + retry_sampling: + temperature: 0.1 + max_tokens: 1024 +``` + +`retry_sampling` accepts the standard OpenAI sampling fields supported by the +worker. It can be configured globally under `inference` or per stage. Changes +received from the orchestrator apply without restarting the worker. + +Some providers need a different constrained-decoding body for the retry. Set +that locally so provider details and credentials never pass through the +orchestrator: + +```yaml +worker: + openai_compatible: + retry_extra_body: + response_format: + type: "json_object" +``` + +Retry request fields override the endpoint's primary `extra_body`, except for +`model` and `messages`, which CaptionFlow reserves. Non-standard sampling such +as vLLM's `repetition_penalty` can also be placed in `retry_extra_body`. diff --git a/src/caption_flow/models.py b/src/caption_flow/models.py index b0bf00d..1478631 100644 --- a/src/caption_flow/models.py +++ b/src/caption_flow/models.py @@ -184,6 +184,7 @@ class ProcessingStage: # provider rejects the image before inference. retry_prompt: Optional[str] = None retry_without_image: bool = False + retry_sampling: Optional[Dict[str, Any]] = None @dataclass diff --git a/src/caption_flow/utils/vllm_config.py b/src/caption_flow/utils/vllm_config.py index c818767..b6a8a4c 100644 --- a/src/caption_flow/utils/vllm_config.py +++ b/src/caption_flow/utils/vllm_config.py @@ -43,6 +43,7 @@ class VLLMConfigManager: "inference_prompts", "refusal_markers", "retry_prompt", + "retry_sampling", "retry_without_image", } diff --git a/src/caption_flow/workers/caption.py b/src/caption_flow/workers/caption.py index d027ff7..ee996fa 100644 --- a/src/caption_flow/workers/caption.py +++ b/src/caption_flow/workers/caption.py @@ -505,6 +505,9 @@ async def _execute_post_hook(self): def _parse_stages_config(self, vllm_config: Dict[str, Any]) -> List[ProcessingStage]: """Parse stages configuration from vLLM config.""" stages_config = vllm_config.get("stages", []) + default_retry_sampling = vllm_config.get("retry_sampling") + if default_retry_sampling is not None and not isinstance(default_retry_sampling, dict): + raise ValueError("retry_sampling must be a mapping") if not stages_config: # Backward compatibility @@ -517,12 +520,18 @@ def _parse_stages_config(self, vllm_config: Dict[str, Any]) -> List[ProcessingSt requires=[], retry_prompt=vllm_config.get("retry_prompt"), retry_without_image=bool(vllm_config.get("retry_without_image", False)), + retry_sampling=( + dict(default_retry_sampling) if default_retry_sampling is not None else None + ), ) ] # Parse stages stages = [] for stage_cfg in stages_config: + retry_sampling = stage_cfg.get("retry_sampling", default_retry_sampling) + if retry_sampling is not None and not isinstance(retry_sampling, dict): + raise ValueError(f"Stage '{stage_cfg['name']}' retry_sampling must be a mapping") stage = ProcessingStage( name=stage_cfg["name"], model=stage_cfg.get("model", vllm_config.get("model")), @@ -541,6 +550,7 @@ def _parse_stages_config(self, vllm_config: Dict[str, Any]) -> List[ProcessingSt vllm_config.get("retry_without_image", False), ) ), + retry_sampling=dict(retry_sampling) if retry_sampling is not None else None, ) stages.append(stage) diff --git a/src/caption_flow/workers/openai_compatible.py b/src/caption_flow/workers/openai_compatible.py index 2bd850d..b38a727 100644 --- a/src/caption_flow/workers/openai_compatible.py +++ b/src/caption_flow/workers/openai_compatible.py @@ -10,6 +10,7 @@ import base64 import email.utils import io +import json import logging import math import os @@ -189,6 +190,7 @@ class ChatRequest: image_data_url: Optional[str] requested_model: Optional[str] parameters: Dict[str, Any] = field(default_factory=dict) + extra_body: Dict[str, Any] = field(default_factory=dict) system_prompt: Optional[str] = None image_detail: Optional[str] = None @@ -460,6 +462,7 @@ async def _request_once(self, state: EndpointState, request: ChatRequest) -> str "messages": messages, **request.parameters, **config.extra_body, + **request.extra_body, } headers = { "Authorization": f"Bearer {config.api_key}", @@ -592,7 +595,7 @@ class OpenAICompatibleWorker(CaptionWorker): "unsafe or sensitive", ) - def __init__(self, config: Dict[str, Any]): + def __init__(self, config: Dict[str, Any]): # noqa: C901 super().__init__(config) raw_config = config.get("openai_compatible") if raw_config is True: @@ -631,6 +634,52 @@ def __init__(self, config: Dict[str, Any]): if self._local_refusal_markers is not None else self._DEFAULT_REFUSAL_MARKERS ) + validate_json_output = raw_config.get("validate_json_output", False) + if not isinstance(validate_json_output, bool): + raise ValueError("openai_compatible.validate_json_output must be a boolean") + self.validate_json_output = validate_json_output + repair_json_escapes = raw_config.get("repair_invalid_json_escapes", False) + if not isinstance(repair_json_escapes, bool): + raise ValueError("openai_compatible.repair_invalid_json_escapes must be a boolean") + self.repair_invalid_json_escapes = repair_json_escapes + canonicalize_json_output = raw_config.get("canonicalize_json_output", False) + if not isinstance(canonicalize_json_output, bool): + raise ValueError("openai_compatible.canonicalize_json_output must be a boolean") + self.canonicalize_json_output = canonicalize_json_output + normalize_yxyx_bboxes = raw_config.get("normalize_yxyx_bboxes", False) + if not isinstance(normalize_yxyx_bboxes, bool): + raise ValueError("openai_compatible.normalize_yxyx_bboxes must be a boolean") + self.normalize_yxyx_bboxes = normalize_yxyx_bboxes + deduplicate_json_elements = raw_config.get("deduplicate_json_elements", False) + if not isinstance(deduplicate_json_elements, bool): + raise ValueError("openai_compatible.deduplicate_json_elements must be a boolean") + self.deduplicate_json_elements = deduplicate_json_elements + json_transforms = { + "repair_invalid_json_escapes": self.repair_invalid_json_escapes, + "canonicalize_json_output": self.canonicalize_json_output, + "normalize_yxyx_bboxes": self.normalize_yxyx_bboxes, + "deduplicate_json_elements": self.deduplicate_json_elements, + } + enabled_without_validation = [ + name + for name, enabled in json_transforms.items() + if enabled and not validate_json_output + ] + if enabled_without_validation: + raise ValueError( + "openai_compatible.validate_json_output must be enabled when using: " + + ", ".join(enabled_without_validation) + ) + retry_extra_body = raw_config.get("retry_extra_body", {}) + if not isinstance(retry_extra_body, dict): + raise ValueError("openai_compatible.retry_extra_body must be a mapping") + reserved_retry_keys = {"model", "messages"}.intersection(retry_extra_body) + if reserved_retry_keys: + raise ValueError( + "openai_compatible.retry_extra_body cannot override: " + + ", ".join(sorted(reserved_retry_keys)) + ) + self.retry_extra_body = dict(retry_extra_body) async def _pre_start(self): """Fetch shared stage settings, then start the API processing thread.""" @@ -746,8 +795,8 @@ def _process_batch_multi_stage( # noqa: C901 if self._is_semantic_caption_failure(response): retryable_item_ids.add(id(owner)) continue - cleaned = self._clean_output(response) - if cleaned and not self._is_refusal_text(cleaned): + cleaned = self._validated_output(self._clean_output(response)) + if cleaned is not None: outputs_by_item[id(owner)].append(cleaned) else: retryable_item_ids.add(id(owner)) @@ -761,6 +810,15 @@ def _process_batch_multi_stage( # noqa: C901 ] if retry_items: retry_requests = [] + retry_sampling = dict(sampling) + if stage.retry_sampling: + retry_sampling.update( + { + key: value + for key, value in stage.retry_sampling.items() + if key in self._SUPPORTED_SAMPLING_KEYS and value is not None + } + ) for item in retry_items: context = self._stage_context(item) retry_prompt = PromptTemplateManager([stage.retry_prompt]).format_all(context)[ @@ -773,7 +831,8 @@ def _process_batch_multi_stage( # noqa: C901 None if stage.retry_without_image else image_urls.get(id(item)) ), requested_model=stage.model, - parameters=sampling, + parameters=retry_sampling, + extra_body=self.retry_extra_body, system_prompt=self.system_prompt, image_detail=self.image_detail, ) @@ -797,8 +856,8 @@ def _process_batch_multi_stage( # noqa: C901 response, ) continue - cleaned = self._clean_output(response) - if cleaned and not self._is_refusal_text(cleaned): + cleaned = self._validated_output(self._clean_output(response)) + if cleaned is not None: outputs_by_item[id(item)].append(cleaned) else: logger.error( @@ -852,6 +911,121 @@ def _is_refusal_text(self, text: str) -> bool: normalized = text.strip().casefold() return any(marker in normalized for marker in self.refusal_markers) + def _validated_output(self, text: str) -> Optional[str]: + if not text or self._is_refusal_text(text): + return None + if not self.validate_json_output: + return text + + parsed_output = self._parse_json_output(text) + if parsed_output is None: + return None + parsed, text = parsed_output + + if self.normalize_yxyx_bboxes: + self._normalize_yxyx_bbox_values(parsed) + if self.deduplicate_json_elements: + self._deduplicate_element_arrays(parsed) + if ( + self.canonicalize_json_output + or self.normalize_yxyx_bboxes + or self.deduplicate_json_elements + ): + return json.dumps(parsed, ensure_ascii=False, separators=(",", ":")) + return text + + def _parse_json_output(self, text: str) -> Optional[Tuple[Any, str]]: + try: + return self._strict_json_loads(text), text + except (TypeError, ValueError) as error: + if self.repair_invalid_json_escapes: + repaired = self._repair_json_escapes(text) + if repaired != text: + try: + parsed = self._strict_json_loads(repaired) + except (TypeError, ValueError): + pass + else: + logger.warning("Repaired invalid JSON escape sequence: %s", error) + return parsed, repaired + logger.warning("Rejecting invalid JSON output: %s", error) + return None + + @staticmethod + def _strict_json_loads(text: str) -> Any: + def reject_constant(value: str) -> None: + raise ValueError(f"Invalid JSON numeric constant: {value}") + + return json.loads(text, parse_constant=reject_constant) + + @classmethod + def _deduplicate_element_arrays(cls, value: Any) -> None: + if isinstance(value, dict): + elements = value.get("elements") + if isinstance(elements, list): + unique = [] + seen = set() + for element in elements: + fingerprint = json.dumps( + element, ensure_ascii=False, sort_keys=True, separators=(",", ":") + ) + if fingerprint not in seen: + seen.add(fingerprint) + unique.append(element) + value["elements"] = unique + for child in value.values(): + cls._deduplicate_element_arrays(child) + elif isinstance(value, list): + for child in value: + cls._deduplicate_element_arrays(child) + + @classmethod + def _normalize_yxyx_bbox_values(cls, value: Any) -> None: + if isinstance(value, dict): + bbox = value.get("bbox") + if ( + isinstance(bbox, list) + and len(bbox) == 4 + and all(type(coordinate) is int for coordinate in bbox) + ): + ymin, xmin, ymax, xmax = bbox + value["bbox"] = [ + min(ymin, ymax), + min(xmin, xmax), + max(ymin, ymax), + max(xmin, xmax), + ] + for child in value.values(): + cls._normalize_yxyx_bbox_values(child) + elif isinstance(value, list): + for child in value: + cls._normalize_yxyx_bbox_values(child) + + @staticmethod + def _repair_json_escapes(text: str) -> str: + """Escape only backslashes that cannot begin a valid JSON escape.""" + repaired: List[str] = [] + index = 0 + while index < len(text): + char = text[index] + if char != "\\": + repaired.append(char) + index += 1 + continue + + following = text[index + 1] if index + 1 < len(text) else "" + valid_unicode = ( + following == "u" + and all(digit in "0123456789abcdefABCDEF" for digit in text[index + 2 : index + 6]) + and len(text[index + 2 : index + 6]) == 4 + ) + if following in '"\\/bfnrt' or valid_unicode: + repaired.append(char) + else: + repaired.append("\\\\") + index += 1 + return "".join(repaired) + @classmethod def _is_semantic_caption_failure(cls, error: Exception) -> bool: if isinstance(error, EndpointRequestError): diff --git a/tests/test_openai_compatible_worker.py b/tests/test_openai_compatible_worker.py index ad471f0..cf002eb 100644 --- a/tests/test_openai_compatible_worker.py +++ b/tests/test_openai_compatible_worker.py @@ -3,6 +3,7 @@ import asyncio import base64 import io +import json import os from unittest.mock import AsyncMock, patch @@ -317,6 +318,10 @@ def post(self, url, **kwargs): image_data_url="data:image/png;base64,AAAA", requested_model="shared-model", parameters={"max_tokens": 100}, + extra_body={ + "reasoning_effort": "high", + "response_format": {"type": "json_object"}, + }, system_prompt="Be precise", image_detail="low", ), @@ -329,7 +334,8 @@ def post(self, url, **kwargs): assert kwargs["headers"]["User-Agent"].startswith("CaptionFlow/") assert kwargs["json"]["model"] == "provider-model" assert kwargs["json"]["max_tokens"] == 100 - assert kwargs["json"]["reasoning_effort"] == "low" + assert kwargs["json"]["reasoning_effort"] == "high" + assert kwargs["json"]["response_format"] == {"type": "json_object"} assert kwargs["json"]["messages"][0] == {"role": "system", "content": "Be precise"} assert kwargs["json"]["messages"][1]["content"][0]["image_url"]["detail"] == "low" @@ -591,6 +597,28 @@ def test_worker_validates_config_and_applies_shared_updates(): } ) + with pytest.raises(ValueError, match="retry_extra_body must be a mapping"): + OpenAICompatibleWorker( + { + **base, + "openai_compatible": { + "api_key_env": "CAPTIONFLOW_TEST_API_KEY", + "retry_extra_body": [], + }, + } + ) + + with pytest.raises(ValueError, match="retry_extra_body cannot override: model"): + OpenAICompatibleWorker( + { + **base, + "openai_compatible": { + "api_key_env": "CAPTIONFLOW_TEST_API_KEY", + "retry_extra_body": {"model": "other"}, + }, + } + ) + config = { **base, "openai_compatible": { @@ -678,6 +706,117 @@ def test_worker_configures_refusal_markers_from_shared_or_local_config(): ) +def test_worker_can_require_valid_json_output(): + config = { + "server": "ws://localhost:8765", + "token": "orchestrator-token", + "openai_compatible": { + "api_key_env": "CAPTIONFLOW_TEST_API_KEY", + "model": "vision", + "validate_json_output": True, + }, + } + with patch.dict(os.environ, {"CAPTIONFLOW_TEST_API_KEY": "secret"}): + worker = OpenAICompatibleWorker(config) + + assert worker._validated_output('{"caption": "A red square."}') is not None + assert worker._validated_output('{"caption": "truncated"') is None + assert worker._validated_output('{"confidence": NaN}') is None + assert worker._validated_output("I'm sorry, but I cannot describe it.") is None + + config["openai_compatible"]["validate_json_output"] = "true" + with ( + patch.dict(os.environ, {"CAPTIONFLOW_TEST_API_KEY": "secret"}), + pytest.raises(ValueError, match="validate_json_output must be a boolean"), + ): + OpenAICompatibleWorker(config) + + +def test_worker_can_repair_only_invalid_json_escapes(): + config = { + "server": "ws://localhost:8765", + "token": "orchestrator-token", + "openai_compatible": { + "api_key_env": "CAPTIONFLOW_TEST_API_KEY", + "model": "vision", + "validate_json_output": True, + "repair_invalid_json_escapes": True, + }, + } + with patch.dict(os.environ, {"CAPTIONFLOW_TEST_API_KEY": "secret"}): + worker = OpenAICompatibleWorker(config) + + malformed = r'{"text": "C:\users\path and \u12Z4", "quote": "ok\nline"}' + repaired = worker._validated_output(malformed) + assert repaired is not None + assert json.loads(repaired) == { + "text": r"C:\users\path and \u12Z4", + "quote": "ok\nline", + } + assert worker._validated_output('{"text": "still truncated"') is None + assert worker._validated_output(r'{"text": "bad\q"') is None + + config["openai_compatible"]["repair_invalid_json_escapes"] = 1 + with ( + patch.dict(os.environ, {"CAPTIONFLOW_TEST_API_KEY": "secret"}), + pytest.raises(ValueError, match="repair_invalid_json_escapes must be a boolean"), + ): + OpenAICompatibleWorker(config) + + +def test_worker_can_canonicalize_json_and_normalize_yxyx_boxes(): + config = { + "server": "ws://localhost:8765", + "token": "orchestrator-token", + "openai_compatible": { + "api_key_env": "CAPTIONFLOW_TEST_API_KEY", + "model": "vision", + "validate_json_output": True, + "canonicalize_json_output": True, + "normalize_yxyx_bboxes": True, + "deduplicate_json_elements": True, + }, + } + with patch.dict(os.environ, {"CAPTIONFLOW_TEST_API_KEY": "secret"}): + worker = OpenAICompatibleWorker(config) + + output = worker._validated_output( + '{"description":"caf\\u00e9","elements":[' + '{"bbox":[900,800,100,200]},{"bbox":[900,800,100,200]}]}' + ) + assert output == '{"description":"café","elements":[{"bbox":[100,200,900,800]}]}' + + for option in ( + "repair_invalid_json_escapes", + "canonicalize_json_output", + "normalize_yxyx_bboxes", + "deduplicate_json_elements", + ): + invalid = { + **config, + "openai_compatible": {**config["openai_compatible"], option: "true"}, + } + with ( + patch.dict(os.environ, {"CAPTIONFLOW_TEST_API_KEY": "secret"}), + pytest.raises(ValueError, match=f"{option} must be a boolean"), + ): + OpenAICompatibleWorker(invalid) + + without_validation = { + **config, + "openai_compatible": { + **config["openai_compatible"], + "validate_json_output": False, + option: True, + }, + } + with ( + patch.dict(os.environ, {"CAPTIONFLOW_TEST_API_KEY": "secret"}), + pytest.raises(ValueError, match="validate_json_output must be enabled"), + ): + OpenAICompatibleWorker(without_validation) + + def test_worker_runs_shared_caption_stage_through_endpoint_pool(): worker_config = { "server": "ws://localhost:8765", @@ -837,6 +976,55 @@ def test_worker_retries_empty_refusal_but_not_unrelated_request_error(): ) +def test_worker_retries_invalid_json_output(): + worker_config = { + "server": "ws://localhost:8765", + "token": "orchestrator-token", + "openai_compatible": { + "api_key_env": "CAPTIONFLOW_TEST_API_KEY", + "model": "vision", + "validate_json_output": True, + "retry_extra_body": {"response_format": {"type": "json_object"}}, + }, + } + with patch.dict(os.environ, {"CAPTIONFLOW_TEST_API_KEY": "provider-secret"}): + worker = OpenAICompatibleWorker(worker_config) + + worker.vllm_config = { + "model": "vision", + "inference_prompts": ["Describe as JSON"], + "retry_prompt": "Retry as valid JSON", + "retry_sampling": {"max_tokens": 4096, "repetition_penalty": 1.1}, + } + worker.stages = worker._parse_stages_config(worker.vllm_config) + worker.stage_order = worker._topological_sort_stages(worker.stages) + worker.endpoint_pool.run_many = AsyncMock( + side_effect=[['{"caption": "truncated"'], ['{"caption": "A red square."}']] + ) + worker.api_loop = asyncio.new_event_loop() + item = ProcessingItem( + unit_id="unit", + job_id="job", + chunk_id="chunk", + item_key="image.png", + item_index=0, + image=Image.new("RGB", (2, 2), color="red"), + image_data=b"", + metadata={}, + ) + + try: + results = worker._process_batch_multi_stage([item]) + finally: + worker.api_loop.close() + + assert results == [(item, {"captions": ['{"caption": "A red square."}']})] + assert worker.endpoint_pool.run_many.await_count == 2 + retry_request = worker.endpoint_pool.run_many.await_args_list[1].args[0][0] + assert retry_request.parameters == {"max_tokens": 4096} + assert retry_request.extra_body == {"response_format": {"type": "json_object"}} + + def test_cli_selects_openai_compatible_worker(): from caption_flow.cli import main diff --git a/tests/test_vllm_config.py b/tests/test_vllm_config.py index a9cac4e..e4cae5f 100644 --- a/tests/test_vllm_config.py +++ b/tests/test_vllm_config.py @@ -58,6 +58,7 @@ def test_class_constants(self): assert "sampling" in VLLMConfigManager.RUNTIME_UPDATEABLE_FIELDS assert "refusal_markers" in VLLMConfigManager.RUNTIME_UPDATEABLE_FIELDS assert "retry_prompt" in VLLMConfigManager.RUNTIME_UPDATEABLE_FIELDS + assert "retry_sampling" in VLLMConfigManager.RUNTIME_UPDATEABLE_FIELDS assert "retry_without_image" in VLLMConfigManager.RUNTIME_UPDATEABLE_FIELDS diff --git a/tests/test_worker_caption.py b/tests/test_worker_caption.py index 16dec60..6aa236d 100644 --- a/tests/test_worker_caption.py +++ b/tests/test_worker_caption.py @@ -257,6 +257,7 @@ def test_parse_stages_backward_compatibility(self, caption_worker): "inference_prompts": ["Test prompt"], "retry_prompt": "Fallback prompt", "retry_without_image": True, + "retry_sampling": {"max_tokens": 2048}, } stages = caption_worker._parse_stages_config(old_config) @@ -267,6 +268,28 @@ def test_parse_stages_backward_compatibility(self, caption_worker): assert stages[0].output_field == "captions" assert stages[0].retry_prompt == "Fallback prompt" assert stages[0].retry_without_image is True + assert stages[0].retry_sampling == {"max_tokens": 2048} + + old_config["retry_sampling"] = [] + with pytest.raises(ValueError, match="retry_sampling must be a mapping"): + caption_worker._parse_stages_config(old_config) + + def test_parse_stage_retry_sampling_override(self, caption_worker): + """Stage retry sampling overrides the shared fallback settings.""" + config = { + "retry_sampling": {"max_tokens": 1024}, + "stages": [ + {"name": "caption", "prompts": ["Caption"], "retry_sampling": {"top_p": 0.8}} + ], + } + + stage = caption_worker._parse_stages_config(config)[0] + + assert stage.retry_sampling == {"top_p": 0.8} + + config["stages"][0]["retry_sampling"] = "invalid" + with pytest.raises(ValueError, match="Stage 'caption' retry_sampling must be a mapping"): + caption_worker._parse_stages_config(config) def test_topological_sort_stages(self, caption_worker): """Test dependency sorting of stages."""