diff --git a/README.md b/README.md index 38d7600..50cf051 100644 --- a/README.md +++ b/README.md @@ -178,6 +178,13 @@ prompts, sampling, and output fields. Existing `vllm:` configurations are also accepted by API workers for backward compatibility; endpoint-local `model` values take precedence over the broadcast model name. +Direct vLLM and API workers share JSON validation, output transformations, +refusal detection, and semantic retries. Configure those under +`inference.output_processing`; use `response_format` and `retry_response_format` +for constrained decoding on either backend. See the +[structured JSON recipe](recipes/STRUCTURED_JSON.md) for examples and native +vLLM tuning details. + --- ## dataset formats diff --git a/examples/worker.openai-compatible.yaml b/examples/worker.openai-compatible.yaml index ac63149..a42ec81 100644 --- a/examples/worker.openai-compatible.yaml +++ b/examples/worker.openai-compatible.yaml @@ -4,6 +4,16 @@ worker: token: "replace-with-captionflow-worker-token" name: "my-api-pool" + # Caption policy normally belongs in orchestrator.inference.output_processing + # and works identically with direct vLLM. Optional local overrides: + # output_processing: + # refusal_markers: [] + # validate_json_output: true + # repair_invalid_json_escapes: true + # canonicalize_json_output: true + # normalize_yxyx_bboxes: true + # deduplicate_json_elements: true + # Provider API keys stay in this process. They are never sent to CaptionFlow. openai_compatible: # Keep enough local items ready to fill all discovered request slots. @@ -13,24 +23,9 @@ worker: image_format: "jpeg" image_quality: 90 - # Optional local override for refusal detection. An empty list disables - # text-based refusal detection. This takes precedence over the shared - # orchestrator inference.refusal_markers setting. - # refusal_markers: - # - "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. + # Provider-specific fields applied only to fallback requests. Prefer shared + # inference.retry_response_format for constraints that also apply to direct + # vLLM workers. # retry_extra_body: # response_format: # type: "json_object" diff --git a/recipes/STRUCTURED_JSON.md b/recipes/STRUCTURED_JSON.md index 0ba17fe..3950ce0 100644 --- a/recipes/STRUCTURED_JSON.md +++ b/recipes/STRUCTURED_JSON.md @@ -1,104 +1,140 @@ # 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. +Direct vLLM and OpenAI-compatible workers use the same caption pipeline: +prompt formatting, output validation and transformations, semantic retries, +and stage success/failure accounting. Choose either backend without losing +JSON captioning or recovery behavior. -## Configure the worker +## Shared inference configuration -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: +Put caption policy and decoding constraints in the orchestrator's `inference` +configuration (the legacy `vllm` spelling also works): ```yaml orchestrator: inference: + model: "vision-model" inference_prompts: - >- - Describe the image as one JSON object matching the requested schema. - Return JSON only. + Describe the visible image as one JSON object matching the requested + schema. Use only visible evidence, without speculation. Return JSON only. sampling: temperature: 0.3 max_tokens: 2048 + + output_processing: + validate_json_output: true + canonicalize_json_output: true + # Optional repairs; boxes must use [ymin, xmin, ymax, xmax] order. + repair_invalid_json_escapes: true + normalize_yxyx_bboxes: true + deduplicate_json_elements: true + # Optional: override refusal substrings, or [] to disable detection. + refusal_markers: ["i cannot describe", "unable to provide a caption"] + + response_format: + type: "json_schema" + json_schema: + name: "caption" + strict: true + schema: + type: "object" + additionalProperties: false + required: ["description"] + properties: + description: {type: "string"} + 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_response_format: + type: "json_object" ``` -`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: +An empty caption, configured refusal match, invalid JSON, or provider content +rejection can receive one semantic retry. Unrelated HTTP/transport failures +use the endpoint's request retry policy, not the semantic prompt. Only items +with an accepted output proceed to the next stage; failed items are not +reported as successful empty captions. With multiple prompts, one accepted +output is sufficient, and that item is not retried. + +`output_processing`, `response_format`, `retry_response_format`, `retry_prompt`, +`retry_sampling`, and `retry_without_image` can also be set per stage. Stage +output policy merges over inference defaults. Set `retry_without_image: true` +only when deliberately retrying from text/metadata; it omits the image on both +backends. Prompt, output-policy, and retry changes apply on config reload +without reloading native model weights. + +`validate_json_output` checks JSON syntax and rejects non-finite numbers. It +does **not** validate against JSON Schema. `response_format` requests constrained +decoding from the backend; support for specific schemas depends on the model, +provider, and vLLM version. Supported envelope types are `text`, `json_object`, +and `json_schema`; a retry inherits the primary constraint unless overridden. +Use `retry_response_format: {type: text}` to request unconstrained retry decoding. + +The optional transforms require `validate_json_output: true`: + +- `repair_invalid_json_escapes` escapes only invalid backslashes, preserving + existing valid escape pairs. It cannot repair truncated JSON. +- `canonicalize_json_output` emits compact JSON with literal Unicode. +- `normalize_yxyx_bboxes` sorts inverted coordinate pairs. It does not infer + boxes, clamp their range, or assess localization quality. +- `deduplicate_json_elements` removes exact duplicates from arrays named + `elements`, including nested arrays of that name. + +## Direct vLLM worker + +Use a normal GPU worker config and launch with `caption-flow worker --vllm`. +The worker translates `response_format` into vLLM's native structured-output +parameters, including the older guided-decoding API. It retains native tensor +parallelism, memory/cache configuration, token checks, and image resize recovery. + +Native `sampling` and `retry_sampling` are passed through to `SamplingParams`, +including vLLM-specific controls such as `top_k`, `min_p`, `seed`, and +`repetition_penalty`. Retry settings overlay primary settings without mutating +the cached primary parameters. Use fields supported by your installed vLLM. +Do not also configure native `structured_outputs`/`guided_decoding` when using +the shared `response_format` envelope. + +## OpenAI-compatible worker + +Launch with `caption-flow worker --openai-compatible`. Endpoint credentials, +request concurrency, image encoding, and provider-specific request extensions +remain local: ```yaml worker: + server: "ws://orchestrator.example:8765" + token: "replace-with-captionflow-worker-token" openai_compatible: - retry_extra_body: - response_format: - type: "json_object" + 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 ``` -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`. +This adapter forwards shared `response_format` and standard OpenAI sampling +fields to chat completions. Non-standard provider fields belong in endpoint +`extra_body`; retry-only extensions belong in `openai_compatible.retry_extra_body`. +Shared response formats override endpoint defaults; local `retry_extra_body` +overrides the shared retry body. Neither extension can replace `model` or +`messages`. + +## Compatibility and local overrides + +Either worker can override shared policy using `worker.output_processing`. +Precedence is inference defaults, stage policy, then local worker policy. +The legacy `inference.refusal_markers` is still accepted, with +`inference.output_processing.refusal_markers` taking precedence. + +Existing JSON/refusal flags under `worker.openai_compatible` remain accepted +as compatibility aliases for local output policy. New configs should use +shared inference policy; explicit `worker.output_processing` overrides those +legacy aliases. Existing endpoint `extra_body.response_format` and +`retry_extra_body.response_format` continue working for API workers. diff --git a/src/caption_flow/models.py b/src/caption_flow/models.py index 1478631..3fe4584 100644 --- a/src/caption_flow/models.py +++ b/src/caption_flow/models.py @@ -179,12 +179,13 @@ class ProcessingStage: dtype: Optional[str] = None gpu_memory_utilization: Optional[float] = None - # Optional semantic fallback for empty/refused caption responses. API - # workers can omit the image on the fallback request when an upstream - # provider rejects the image before inference. + # Caption policy and recovery apply to every inference backend. retry_prompt: Optional[str] = None retry_without_image: bool = False retry_sampling: Optional[Dict[str, Any]] = None + output_processing: Optional[Dict[str, Any]] = None + response_format: Optional[Dict[str, Any]] = None + retry_response_format: Optional[Dict[str, Any]] = None @dataclass diff --git a/src/caption_flow/utils/image_processor.py b/src/caption_flow/utils/image_processor.py index e42d801..371d439 100644 --- a/src/caption_flow/utils/image_processor.py +++ b/src/caption_flow/utils/image_processor.py @@ -40,13 +40,18 @@ def prepare_for_inference(item: ProcessingItem) -> Image.Image: if item.image is not None: image = item.image - item.metadata["image_width"], item.metadata["image_height"] = image.size - item.metadata["image_format"] = image.format or "unknown" + for name, value in ( + ("image_width", image.width), + ("image_height", image.height), + ("image_format", image.format or "unknown"), + ): + if item.metadata.get(name) is None: + item.metadata[name] = value # item.image = None return image - item.image = None image = ImageProcessor.decode_image_data(item.image_data) + item.image = image item.image_data = b"" item.metadata["image_format"] = image.format or "unknown" item.metadata["image_width"], item.metadata["image_height"] = image.size diff --git a/src/caption_flow/utils/output_policy.py b/src/caption_flow/utils/output_policy.py new file mode 100644 index 0000000..3808924 --- /dev/null +++ b/src/caption_flow/utils/output_policy.py @@ -0,0 +1,203 @@ +"""Caption acceptance and JSON transformations, independent of inference backends.""" + +import json +import logging +import math +from dataclasses import dataclass, fields +from typing import Any, Mapping, Optional + +logger = logging.getLogger(__name__) + +DEFAULT_REFUSAL_MARKERS = ( + "i'm sorry", + "i’m sorry", + "i cannot", + "i can't", + "i can’t", + "unable to provide", + "unable to describe", + "cannot provide", + "can't provide", + "can’t provide", + "cannot assist", + "can't assist", + "can’t assist", +) + + +def validate_response_format(value: Any) -> None: + """Validate the shared constrained-decoding envelope before either backend runs.""" + if value is None: + return + if not isinstance(value, dict) or value.get("type") not in { + "text", + "json_object", + "json_schema", + }: + raise ValueError( + "response_format must be a mapping with type text, json_object or json_schema" + ) + if value["type"] == "json_schema": + schema = value.get("json_schema") + if not isinstance(schema, dict) or not isinstance(schema.get("schema"), dict): + raise ValueError("response_format.json_schema.schema must be a mapping") + + +def normalize_refusal_markers(value: Any) -> tuple[str, ...]: + if not isinstance(value, list) or any( + not isinstance(marker, str) or not marker.strip() for marker in value + ): + raise ValueError("refusal_markers must be a list of non-empty strings") + return tuple(dict.fromkeys(marker.strip().casefold() for marker in value)) + + +@dataclass(frozen=True) +class OutputPolicy: + """Shared interpretation of a model's caption output.""" + + validate_json_output: bool = False + repair_invalid_json_escapes: bool = False + canonicalize_json_output: bool = False + normalize_yxyx_bboxes: bool = False + deduplicate_json_elements: bool = False + refusal_markers: tuple[str, ...] = DEFAULT_REFUSAL_MARKERS + + @classmethod + def from_config(cls, config: Mapping[str, Any], *, partial: bool = False) -> "OutputPolicy": + options = dict(config) + known = {field.name for field in fields(cls)} + if unknown := options.keys() - known: + raise ValueError(f"Unknown output_processing option(s): {', '.join(sorted(unknown))}") + for name, value in options.items(): + if name != "refusal_markers" and not isinstance(value, bool): + raise ValueError(f"output_processing.{name} must be a boolean") + if "refusal_markers" in options: + options["refusal_markers"] = normalize_refusal_markers(options["refusal_markers"]) + policy = cls(**options) + check_dependencies = not partial or "validate_json_output" in options + if ( + check_dependencies + and not policy.validate_json_output + and any( + getattr(policy, name) + for name in known - {"validate_json_output", "refusal_markers"} + ) + ): + raise ValueError( + "output_processing.validate_json_output must be enabled for JSON transforms" + ) + return policy + + def is_refusal(self, text: str) -> bool: + return any(marker in text.strip().casefold() for marker in self.refusal_markers) + + @staticmethod + def clean(text: str) -> str: + if not text: + return "" + for token in ("<|end|>", "<|endoftext|>", "<|im_end|>"): + text = text.split(token, 1)[0] + return text.strip() + + def process(self, text: str) -> Optional[str]: + text = self.clean(text) + if not text or self.is_refusal(text): + return None + if not self.validate_json_output: + return text + try: + parsed = self._loads(text) + except ValueError as error: + if not self.repair_invalid_json_escapes: + logger.warning("Rejecting invalid JSON output: %s", error) + return None + repaired = self._repair_json_escapes(text) + try: + parsed = self._loads(repaired) + except ValueError: + logger.warning("Rejecting invalid JSON output: %s", error) + return None + text = repaired + 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=(",", ":"), allow_nan=False) + return text + + @staticmethod + def _loads(text: str) -> Any: + def reject_constant(value: str) -> None: + raise ValueError(f"Invalid JSON numeric constant: {value}") + + def finite_float(value: str) -> float: + result = float(value) + if not math.isfinite(result): + raise ValueError("JSON number exceeds finite float range") + return result + + return json.loads(text, parse_constant=reject_constant, parse_float=finite_float) + + @classmethod + def _deduplicate_element_arrays(cls, value: Any) -> None: + if isinstance(value, dict): + # Normalize children first so recursively equivalent elements deduplicate. + for child in value.values(): + cls._deduplicate_element_arrays(child) + elements = value.get("elements") + if isinstance(elements, list): + unique, seen = [], set() + for element in elements: + fingerprint = json.dumps(element, sort_keys=True, ensure_ascii=False) + if fingerprint not in seen: + seen.add(fingerprint) + unique.append(element) + value["elements"] = unique + 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(v) is int for v 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: + """Preserve valid escape pairs, escaping only invalid backslashes.""" + repaired = [] + index = 0 + while index < len(text): + char = text[index] + if char != "\\": + repaired.append(char) + index += 1 + continue + following = text[index + 1 : index + 2] + digits = text[index + 2 : index + 6] + valid_unicode = ( + following == "u" + and len(digits) == 4 + and all(digit in "0123456789abcdefABCDEF" for digit in digits) + ) + if following and (following in '"\\/bfnrt' or valid_unicode): + size = 6 if valid_unicode else 2 + repaired.append(text[index : index + size]) + index += size + else: + repaired.append("\\\\") + index += 1 + return "".join(repaired) diff --git a/src/caption_flow/utils/vllm_config.py b/src/caption_flow/utils/vllm_config.py index b6a8a4c..02e8226 100644 --- a/src/caption_flow/utils/vllm_config.py +++ b/src/caption_flow/utils/vllm_config.py @@ -4,9 +4,43 @@ from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Tuple +from .output_policy import validate_response_format + logger = logging.getLogger(__name__) +def create_native_sampling_params(sampling: Dict[str, Any], response_format=None): + """Preserve native vLLM tuning and translate a shared output constraint.""" + from vllm import SamplingParams + + validate_response_format(response_format) + if response_format is not None and any( + sampling.get(key) is not None for key in ("structured_outputs", "guided_decoding") + ): + raise ValueError("Use response_format or native structured-output sampling, not both") + params = { + "temperature": 0.7, + "top_p": 0.95, + "max_tokens": 256, + "stop": ["<|end|>", "<|endoftext|>", "<|im_end|>"], + "repetition_penalty": 1.05, + "skip_special_tokens": True, + **sampling, + } + if response_format and response_format.get("type") != "text": + from vllm import sampling_params + + if response_format["type"] == "json_schema": + constraint = {"json": response_format["json_schema"]["schema"]} + elif response_format["type"] == "json_object": + constraint = {"json_object": True} + if hasattr(sampling_params, "StructuredOutputsParams"): + params["structured_outputs"] = sampling_params.StructuredOutputsParams(**constraint) + else: + params["guided_decoding"] = sampling_params.GuidedDecodingParams(**constraint) + return SamplingParams(**params) + + @dataclass class VLLMConfigChange: """Represents changes between vLLM configurations.""" @@ -45,6 +79,9 @@ class VLLMConfigManager: "retry_prompt", "retry_sampling", "retry_without_image", + "output_processing", + "response_format", + "retry_response_format", } def __init__(self): @@ -97,17 +134,8 @@ def analyze_config_change( def create_sampling_params(self, vllm_config: Dict[str, Any]): """Create SamplingParams from config.""" - from vllm import SamplingParams - - sampling_config = vllm_config.get("sampling", {}) - - params = SamplingParams( - temperature=sampling_config.get("temperature", 0.7), - top_p=sampling_config.get("top_p", 0.95), - max_tokens=sampling_config.get("max_tokens", 256), - stop=sampling_config.get("stop", ["<|end|>", "<|endoftext|>", "<|im_end|>"]), - repetition_penalty=sampling_config.get("repetition_penalty", 1.05), - skip_special_tokens=sampling_config.get("skip_special_tokens", True), + params = create_native_sampling_params( + vllm_config.get("sampling", {}), vllm_config.get("response_format") ) self.current_sampling_params = params diff --git a/src/caption_flow/workers/caption.py b/src/caption_flow/workers/caption.py index ee996fa..6b9282f 100644 --- a/src/caption_flow/workers/caption.py +++ b/src/caption_flow/workers/caption.py @@ -1,13 +1,10 @@ """Caption worker with processor abstraction for distributed captioning.""" -import os - -os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0" - import asyncio import inspect import json import logging +import os import time from collections import defaultdict, deque from dataclasses import dataclass @@ -30,9 +27,11 @@ WorkUnit, ) from ..utils.image_processor import ImageProcessor +from ..utils.output_policy import OutputPolicy, validate_response_format from ..utils.prompt_template import PromptTemplateManager -from ..utils.vllm_config import VLLMConfigManager +from ..utils.vllm_config import VLLMConfigManager, create_native_sampling_params from .base import BaseWorker +from .pipeline import StageRequest, run_caption_pipeline, stage_context logger = logging.getLogger(__name__) logger.setLevel(os.environ.get("CAPTIONFLOW_LOG_LEVEL", "INFO").upper()) @@ -159,24 +158,21 @@ def load_model(self, model_name: str, stage: ProcessingStage, base_config: Dict[ self.models[model_name] = LLM(**self._filter_engine_args(vllm_params, EngineArgs)) logger.info(f"Model {model_name} loaded successfully") - def create_sampling_params(self, stage: ProcessingStage, base_sampling: Dict[str, Any]): + def create_sampling_params( + self, stage: ProcessingStage, base_sampling: Dict[str, Any], retry: bool = False + ): """Create sampling params for a stage.""" - from vllm import SamplingParams - sampling_config = base_sampling.copy() if stage.sampling: sampling_config.update(stage.sampling) - - params = SamplingParams( - temperature=sampling_config.get("temperature", 0.7), - top_p=sampling_config.get("top_p", 0.95), - max_tokens=sampling_config.get("max_tokens", 256), - stop=sampling_config.get("stop", ["<|end|>", "<|endoftext|>", "<|im_end|>"]), - repetition_penalty=sampling_config.get("repetition_penalty", 1.05), - skip_special_tokens=sampling_config.get("skip_special_tokens", True), - ) - - self.sampling_params[stage.name] = params + if retry and stage.retry_sampling: + sampling_config.update(stage.retry_sampling) + response_format = stage.response_format + if retry and stage.retry_response_format is not None: + response_format = stage.retry_response_format + params = create_native_sampling_params(sampling_config, response_format) + if not retry: + self.sampling_params[stage.name] = params return params def get_model_for_stage(self, stage_name: str, model_name: str) -> Tuple[Any, Any, Any, Any]: @@ -219,6 +215,10 @@ class CaptionWorker(BaseWorker): def __init__(self, config: Dict[str, Any]): super().__init__(config) + self.output_processing_overrides = dict(config.get("output_processing", {})) + # Local overrides may inherit validation from the orchestrator. Check + # their types now and validate the merged policy when stages arrive. + OutputPolicy.from_config(self.output_processing_overrides, partial=True) # Processor configuration self.processor_type = None @@ -509,6 +509,25 @@ def _parse_stages_config(self, vllm_config: Dict[str, Any]) -> List[ProcessingSt if default_retry_sampling is not None and not isinstance(default_retry_sampling, dict): raise ValueError("retry_sampling must be a mapping") + def output_options(stage_config: Dict[str, Any]) -> Dict[str, Any]: + options = dict(vllm_config.get("output_processing", {})) + if "refusal_markers" in vllm_config: + options.setdefault("refusal_markers", vllm_config["refusal_markers"]) + options.update(stage_config.get("output_processing", {})) + OutputPolicy.from_config({**options, **self.output_processing_overrides}) + result = { + "output_processing": options, + "response_format": stage_config.get( + "response_format", vllm_config.get("response_format") + ), + "retry_response_format": stage_config.get( + "retry_response_format", vllm_config.get("retry_response_format") + ), + } + validate_response_format(result["response_format"]) + validate_response_format(result["retry_response_format"]) + return result + if not stages_config: # Backward compatibility return [ @@ -523,6 +542,7 @@ def _parse_stages_config(self, vllm_config: Dict[str, Any]) -> List[ProcessingSt retry_sampling=( dict(default_retry_sampling) if default_retry_sampling is not None else None ), + **output_options({}), ) ] @@ -551,6 +571,7 @@ def _parse_stages_config(self, vllm_config: Dict[str, Any]) -> List[ProcessingSt ) ), retry_sampling=dict(retry_sampling) if retry_sampling is not None else None, + **output_options(stage_cfg), ) stages.append(stage) @@ -591,6 +612,7 @@ def _setup_vllm(self): if not self.vllm_config: raise RuntimeError("vLLM config not received") + os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0") os.environ["CUDA_VISIBLE_DEVICES"] = str(self.gpu_id) # Initialize model manager @@ -637,34 +659,41 @@ def _handle_vllm_config_update(self, new_config: Dict[str, Any]) -> bool: # Check if mock mode changed old_mock_mode = self.mock_mode - self.mock_mode = new_config.get("mock_results", False) + new_mock_mode = new_config.get("mock_results", False) - if old_mock_mode != self.mock_mode: - logger.info(f"Mock mode changed from {old_mock_mode} to {self.mock_mode}") + if old_mock_mode != new_mock_mode: + logger.info(f"Mock mode changed from {old_mock_mode} to {new_mock_mode}") # Parse new stages new_stages = self._parse_stages_config(new_config) + new_order = self._topological_sort_stages(new_stages) # Check if stages changed significantly - stages_changed = len(new_stages) != len(self.stages) + stages_changed = len(new_stages) != len( + self.stages + ) or self.vllm_config_manager.should_reload_vllm(self.vllm_config, new_config) if not stages_changed: for old, new in zip(self.stages, new_stages, strict=False): if ( old.name != new.name or old.model != new.model - or old.prompts != new.prompts - or old.output_field != new.output_field + or old.tensor_parallel_size != new.tensor_parallel_size + or old.max_model_len != new.max_model_len + or old.dtype != new.dtype + or old.gpu_memory_utilization != new.gpu_memory_utilization ): stages_changed = True break - if stages_changed or old_mock_mode != self.mock_mode: + if stages_changed or old_mock_mode != new_mock_mode: logger.info("Configuration changed significantly") old_config = self.vllm_config + old_stages, old_order = self.stages, self.stage_order + self.mock_mode = new_mock_mode self.vllm_config = new_config self.stages = new_stages - self.stage_order = self._topological_sort_stages(self.stages) + self.stage_order = new_order if not self.mock_mode: try: @@ -675,9 +704,9 @@ def _handle_vllm_config_update(self, new_config: Dict[str, Any]) -> bool: except Exception as e: logger.error(f"Failed to reload vLLM: {e}") # Restore previous state + self.mock_mode = old_mock_mode self.vllm_config = old_config - self.stages = self._parse_stages_config(old_config) - self.stage_order = self._topological_sort_stages(self.stages) + self.stages, self.stage_order = old_stages, old_order # Attempt to restore previous models try: self._setup_vllm() @@ -699,9 +728,11 @@ def _handle_vllm_config_update(self, new_config: Dict[str, Any]) -> bool: if not self.mock_mode: logger.info("Updating sampling parameters without model reload") base_sampling = new_config.get("sampling", {}) - for stage in self.stages: + for stage in new_stages: self.model_manager.create_sampling_params(stage, base_sampling) self.vllm_config = new_config + self.stages = new_stages + self.stage_order = new_order return True def _processing_thread(self): @@ -981,14 +1012,7 @@ def _validate_and_split_batch( ) # Test with first prompt # Build context - context = item.metadata.copy() - for prev_stage_name, stage_result in item.stage_results.items(): - for i, output in enumerate(stage_result.outputs): - context[f"{prev_stage_name}_output_{i}"] = output - if len(stage_result.outputs) == 1: - context[stage_result.output_field] = stage_result.outputs[0] - else: - context[stage_result.output_field] = stage_result.outputs + context = stage_context(item) logger.debug(f"Validation context for {item.item_key}: {context}") # Format test prompt @@ -1069,157 +1093,112 @@ def _resize_image_for_tokens( return new_item - def _process_batch_multi_stage( - self, batch: List[ProcessingItem], max_attempts: int = 3 - ) -> List[Tuple[ProcessingItem, Dict]]: - """Process a batch through all stages with token validation.""" - results = [] + def _output_policy(self, stage: Optional[ProcessingStage] = None) -> OutputPolicy: + config = self.vllm_config or {} + options = dict(config.get("output_processing", {})) + if "refusal_markers" in config: + options.setdefault("refusal_markers", config["refusal_markers"]) + if stage and stage.output_processing: + options.update(stage.output_processing) + options.update(self.output_processing_overrides) + return OutputPolicy.from_config(options) - # Get max model length from config - max_model_len = self.vllm_config.get("max_model_len", 16384) + @property + def refusal_markers(self): + return self._output_policy().refusal_markers - # Process each stage in order - for stage_name in self.stage_order: - stage = next(s for s in self.stages if s.name == stage_name) - logger.debug(f"Processing batch through stage: {stage_name}") + def _is_refusal_text(self, text: str) -> bool: + return self._output_policy().is_refusal(text) - # Check if model manager is properly initialized - if not self.model_manager: - logger.error("Model manager not initialized") - self.items_failed += len(batch) - return [] + def _validated_output(self, text: str) -> Optional[str]: + return self._output_policy().process(text) - # Get model components - try: - llm, processor, tokenizer, sampling_params = self.model_manager.get_model_for_stage( - stage_name, stage.model - ) - except KeyError as e: - logger.error(f"Model not found during batch processing: {e}") - self.items_failed += len(batch) - return [] - - # Validate batch before processing - processable_batch, too_long_items = self._validate_and_split_batch( - batch, stage, processor, tokenizer, sampling_params, max_model_len - ) - - # Handle items that are too long - for item in too_long_items: - logger.warning(f"Item {item.item_key} exceeds token limit, attempting resize") + def _process_batch_multi_stage( + self, batch: List[ProcessingItem] + ) -> List[Tuple[ProcessingItem, Dict]]: + """Use the shared caption pipeline with this worker's generation backend.""" + try: + results, failed = run_caption_pipeline(self, batch, self.stages, self.stage_order) + finally: + self._finish_caption_batch() + self.items_failed += failed + self.items_processed += len(results) + return results - # Try resizing the image - resized_item = self._resize_image_for_tokens(item, target_ratio=0.7) + def _finish_caption_batch(self) -> None: + """Release any backend-local request buffers after success or failure.""" - # Re-validate - resized_processable, still_too_long = self._validate_and_split_batch( - [resized_item], stage, processor, tokenizer, sampling_params, max_model_len + def _prepare_stage_batch(self, batch: list, stage: ProcessingStage) -> list: + """Preserve native token validation and image resize recovery.""" + if self.model_manager is None: + logger.error("Model manager not initialized") + return [] + try: + _, processor, tokenizer, sampling = self.model_manager.get_model_for_stage( + stage.name, stage.model + ) + except KeyError as error: + logger.error("Model not found during batch processing: %s", error) + return [] + max_length = stage.max_model_len or self.vllm_config.get("max_model_len", 16384) + processable, too_long = self._validate_and_split_batch( + batch, stage, processor, tokenizer, sampling, max_length + ) + for item in too_long: + for ratio in (0.7, 0.5): + resized = self._resize_image_for_tokens(item, target_ratio=ratio) + accepted, _ = self._validate_and_split_batch( + [resized], stage, processor, tokenizer, sampling, max_length ) - - if resized_processable: - processable_batch.extend(resized_processable) - logger.info(f"Successfully resized {item.item_key} for processing") - else: - # Try even smaller - resized_item = self._resize_image_for_tokens(item, target_ratio=0.5) - resized_processable, still_too_long = self._validate_and_split_batch( - [resized_item], stage, processor, tokenizer, sampling_params, max_model_len - ) - - if resized_processable: - processable_batch.extend(resized_processable) - logger.info(f"Successfully resized {item.item_key} to 50% for processing") - else: - logger.error(f"Item {item.item_key} still too long after resize, skipping") - self.items_failed += 1 - - # Send error result - stage_result = StageResult( - stage_name=stage_name, - output_field=stage.output_field, - outputs=[], - error="Image too large even after resizing", - ) - item.stage_results[stage_name] = stage_result - - self.result_queue.put( - { - "item": item, - "outputs": {}, - "processing_time_ms": 0.0, - "error": f"Failed stage {stage_name}: token limit exceeded", - } - ) - - # Process the validated batch - if processable_batch: - # Build requests for processable items - requests = [] - for item in processable_batch: - converted_img = ImageProcessor.prepare_for_inference(item) - template_manager = PromptTemplateManager(stage.prompts) - - # Build context - context = item.metadata.copy() - for prev_stage_name, stage_result in item.stage_results.items(): - for i, output in enumerate(stage_result.outputs): - context[f"{prev_stage_name}_output_{i}"] = output - if len(stage_result.outputs) == 1: - context[stage_result.output_field] = stage_result.outputs[0] - else: - context[stage_result.output_field] = stage_result.outputs - - # Format prompts - formatted_prompts = template_manager.format_all(context) - - # Build requests - for prompt in formatted_prompts: - req = self._build_vllm_input(converted_img, prompt, processor, tokenizer) - requests.append(req) - - # Run inference - outputs = llm.generate(requests, sampling_params) - - # Process outputs - for idx, item in enumerate(processable_batch): - base_idx = idx * len(stage.prompts) - stage_outputs = [] - - for j in range(len(stage.prompts)): - if base_idx + j < len(outputs) and outputs[base_idx + j].outputs: - original_output = outputs[base_idx + j].outputs[0].text - cleaned_output = self._clean_output(original_output) - if cleaned_output: - stage_outputs.append(cleaned_output) - - if stage_outputs: - stage_result = StageResult( - stage_name=stage_name, - output_field=stage.output_field, - outputs=stage_outputs, - ) - item.stage_results[stage_name] = stage_result - else: - logger.error(f"No outputs for {item.item_key} in stage {stage_name}") - self.items_failed += 1 - - # Update batch for next stage - batch = processable_batch - - # Convert to results - for item in batch: - # Aggregate outputs by field - outputs_by_field = defaultdict(list) - for stage_result in item.stage_results.values(): - outputs_by_field[stage_result.output_field].extend(stage_result.outputs) - - results.append((item, dict(outputs_by_field))) - self.items_processed += 1 - - return results + if accepted: + processable.extend(accepted) + break + else: + self.result_queue.put( + { + "item": item, + "outputs": {}, + "processing_time_ms": 0.0, + "error": f"Failed stage {stage.name}: token limit exceeded", + } + ) + return processable + + def _generate_stage( + self, stage: ProcessingStage, requests: list[StageRequest], retry: bool + ) -> list: + """Translate one generation batch into native vLLM inputs.""" + llm, processor, tokenizer, sampling = self.model_manager.get_model_for_stage( + stage.name, stage.model + ) + if retry: + sampling = self.model_manager.create_sampling_params( + stage, self.vllm_config.get("sampling", {}), retry=True + ) + inputs = [ + self._build_vllm_input( + None + if retry and stage.retry_without_image + else ImageProcessor.prepare_for_inference(request.item), + request.prompt, + processor, + tokenizer, + ) + for request in requests + ] + outputs = llm.generate(inputs, sampling) + if len(outputs) != len(requests): + raise RuntimeError("vLLM returned a different number of results than requests") + return [output.outputs[0].text if output.outputs else "" for output in outputs] def _build_vllm_input(self, image: Image.Image, prompt: str, processor, tokenizer) -> Dict: """Build vLLM input.""" + if image is None: + messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}] + prompt_text = processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + return {"prompt_token_ids": tokenizer(prompt_text, add_special_tokens=False).input_ids} try: from qwen_vl_utils import process_vision_info @@ -1251,14 +1230,7 @@ def _build_vllm_input(self, image: Image.Image, prompt: str, processor, tokenize def _clean_output(self, text: str) -> str: """Clean model output.""" - if not text: - return "" - - for token in ["<|end|>", "<|endoftext|>", "<|im_end|>", "I'm sorry", "I cannot"]: - if token in text: - text = text.split(token)[0] - - return text.strip() + return self._validated_output(text) or "" def _get_heartbeat_data(self) -> Dict[str, Any]: """Get heartbeat data.""" diff --git a/src/caption_flow/workers/openai_compatible.py b/src/caption_flow/workers/openai_compatible.py index b38a727..55da4fa 100644 --- a/src/caption_flow/workers/openai_compatible.py +++ b/src/caption_flow/workers/openai_compatible.py @@ -10,12 +10,11 @@ import base64 import email.utils import io -import json import logging import math import os import time -from collections import defaultdict, deque +from collections import deque from dataclasses import dataclass, field from datetime import datetime, timezone from threading import Lock, Thread @@ -25,10 +24,11 @@ from PIL import Image from .. import __version__ -from ..models import ProcessingStage, StageResult +from ..models import ProcessingStage from ..utils.image_processor import ImageProcessor -from ..utils.prompt_template import PromptTemplateManager +from ..utils.output_policy import OutputPolicy from .caption import CaptionWorker, ProcessingItem +from .pipeline import CaptionRejectedError, StageRequest logger = logging.getLogger(__name__) logger.setLevel(os.environ.get("CAPTIONFLOW_LOG_LEVEL", "INFO").upper()) @@ -569,21 +569,6 @@ class OpenAICompatibleWorker(CaptionWorker): "presence_penalty", "seed", } - _DEFAULT_REFUSAL_MARKERS = ( - "i'm sorry", - "i’m sorry", - "i cannot", - "i can't", - "i can’t", - "unable to provide", - "unable to describe", - "cannot provide", - "can't provide", - "can’t provide", - "cannot assist", - "can't assist", - "can’t assist", - ) _POLICY_ERROR_MARKERS = ( "contentfilter", "content filter", @@ -596,12 +581,23 @@ class OpenAICompatibleWorker(CaptionWorker): ) def __init__(self, config: Dict[str, Any]): # noqa: C901 - super().__init__(config) raw_config = config.get("openai_compatible") if raw_config is True: raw_config = {} if not isinstance(raw_config, dict): raise ValueError("openai_compatible worker configuration must be a mapping") + # Compatibility adapter for configs written before output policy was shared. + legacy_output = { + name: raw_config[name] + for name in OutputPolicy.__dataclass_fields__ + if name in raw_config + } + super().__init__( + { + **config, + "output_processing": {**legacy_output, **config.get("output_processing", {})}, + } + ) self.api_config = raw_config raw_endpoints = raw_config.get("endpoints") @@ -617,6 +613,7 @@ def __init__(self, config: Dict[str, Any]): # noqa: C901 ] self.endpoint_pool = AdaptiveEndpointPool(endpoints) self.api_loop: Optional[asyncio.AbstractEventLoop] = None + self._stage_image_urls: Dict[int, Optional[str]] = {} self.system_prompt = raw_config.get("system_prompt") self.include_image = bool(raw_config.get("include_image", True)) self.image_detail = raw_config.get("image_detail", "auto") @@ -624,52 +621,6 @@ def __init__(self, config: Dict[str, Any]): # noqa: C901 self.image_quality = int(raw_config.get("image_quality", 90)) configured_dimension = int(raw_config.get("max_image_dimension", 0) or 0) self.max_image_dimension = configured_dimension if configured_dimension > 0 else None - self._local_refusal_markers = ( - self._normalize_refusal_markers(raw_config["refusal_markers"]) - if "refusal_markers" in raw_config - else None - ) - self.refusal_markers = ( - self._local_refusal_markers - 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") @@ -711,24 +662,19 @@ def _apply_local_overrides(self) -> None: if batch_size is None: batch_size = sum(state.config.max_concurrency for state in self.endpoint_pool.states) self.vllm_config["batch_size"] = max(1, int(batch_size)) - if self._local_refusal_markers is not None: - self.refusal_markers = self._local_refusal_markers - elif "refusal_markers" in self.vllm_config: - self.refusal_markers = self._normalize_refusal_markers( - self.vllm_config["refusal_markers"] - ) - else: - self.refusal_markers = self._DEFAULT_REFUSAL_MARKERS + self._output_policy() def _handle_vllm_config_update(self, new_config: Dict[str, Any]) -> bool: """Apply shared prompt/stage changes without touching local credentials.""" if not new_config: return True + new_stages = self._parse_stages_config(new_config) + new_order = self._topological_sort_stages(new_stages) self.vllm_config = dict(new_config) self._apply_local_overrides() self.mock_mode = bool(self.vllm_config.get("mock_results", False)) - self.stages = self._parse_stages_config(self.vllm_config) - self.stage_order = self._topological_sort_stages(self.stages) + self.stages = new_stages + self.stage_order = new_order return True def _processing_thread(self): @@ -746,285 +692,60 @@ def _processing_thread(self): loop.close() self.api_loop = None - def _process_batch_multi_stage( # noqa: C901 - self, batch: List[ProcessingItem], max_attempts: int = 3 - ) -> List[Tuple[ProcessingItem, Dict]]: - del max_attempts # Retries are controlled per endpoint. + def _prepare_stage_batch(self, batch: list, stage: ProcessingStage) -> list: if not self.api_loop: raise RuntimeError("OpenAI-compatible endpoint loop is not ready") - - active_batch = list(batch) - image_urls: Dict[int, Optional[str]] = {} if self.include_image: - image_urls = self._image_data_urls(active_batch) - - for stage_name in self.stage_order: - stage = next(stage for stage in self.stages if stage.name == stage_name) - requests: List[ChatRequest] = [] - owners: List[ProcessingItem] = [] - sampling = self._sampling_for_stage(stage) - - for item in active_batch: - context = self._stage_context(item) - - for prompt in PromptTemplateManager(stage.prompts).format_all(context): - requests.append( - ChatRequest( - prompt=prompt, - image_data_url=image_urls.get(id(item)), - requested_model=stage.model, - parameters=sampling, - system_prompt=self.system_prompt, - image_detail=self.image_detail, - ) - ) - owners.append(item) - - api_started = time.monotonic() - responses = self.api_loop.run_until_complete(self.endpoint_pool.run_many(requests)) - outputs_by_item: Dict[int, List[str]] = defaultdict(list) - retryable_item_ids = set() - for owner, response in zip(owners, responses, strict=True): - if isinstance(response, Exception): - logger.error( - "API request failed for item %s in stage %s: %s", - owner.item_key, - stage_name, - response, - ) - if self._is_semantic_caption_failure(response): - retryable_item_ids.add(id(owner)) - continue - 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)) - - retry_items = [ - item - for item in active_batch - if id(item) in retryable_item_ids - and not outputs_by_item.get(id(item)) - and stage.retry_prompt - ] - 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)[ - 0 - ] - retry_requests.append( - ChatRequest( - prompt=retry_prompt, - image_data_url=( - None if stage.retry_without_image else image_urls.get(id(item)) - ), - requested_model=stage.model, - parameters=retry_sampling, - extra_body=self.retry_extra_body, - system_prompt=self.system_prompt, - image_detail=self.image_detail, - ) - ) - - logger.info( - "Retrying %d empty or refused caption(s) in stage %s%s", - len(retry_items), - stage_name, - " without images" if stage.retry_without_image else "", - ) - retry_responses = self.api_loop.run_until_complete( - self.endpoint_pool.run_many(retry_requests) - ) - for item, response in zip(retry_items, retry_responses, strict=True): - if isinstance(response, Exception): - logger.error( - "Fallback API request failed for item %s in stage %s: %s", - item.item_key, - stage_name, - response, - ) - continue - cleaned = self._validated_output(self._clean_output(response)) - if cleaned is not None: - outputs_by_item[id(item)].append(cleaned) - else: - logger.error( - "Fallback returned no usable output for %s in stage %s", - item.item_key, - stage_name, - ) - - logger.info( - "API stage %s completed %d request(s) in %.3f seconds", - stage_name, - len(requests) + len(retry_items), - time.monotonic() - api_started, + missing = [item for item in batch if id(item) not in self._stage_image_urls] + if missing: + self._stage_image_urls.update(self._image_data_urls(missing)) + return batch + + def _finish_caption_batch(self) -> None: + self._stage_image_urls.clear() + + def _generate_stage( + self, stage: ProcessingStage, requests: list[StageRequest], retry: bool + ) -> list: + """Encode requests and classify provider errors for the shared pipeline.""" + sampling = self._sampling_for_stage(stage) + if retry and stage.retry_sampling: + sampling.update( + { + key: value + for key, value in stage.retry_sampling.items() + if key in self._SUPPORTED_SAMPLING_KEYS and value is not None + } ) - - next_batch = [] - for item in active_batch: - outputs = outputs_by_item.get(id(item), []) - if outputs: - item.stage_results[stage_name] = StageResult( - stage_name=stage_name, - output_field=stage.output_field, - outputs=outputs, - ) - next_batch.append(item) - else: - logger.error("No outputs for %s in stage %s", item.item_key, stage_name) - self.items_failed += 1 - active_batch = next_batch - if not active_batch: - break - - results = [] - for item in active_batch: - outputs_by_field: Dict[str, List[str]] = defaultdict(list) - for stage_result in item.stage_results.values(): - outputs_by_field[stage_result.output_field].extend(stage_result.outputs) - results.append((item, dict(outputs_by_field))) - self.items_processed += 1 - return results - - @staticmethod - def _normalize_refusal_markers(value: Any) -> Tuple[str, ...]: - if not isinstance(value, list) or any( - not isinstance(marker, str) or not marker.strip() for marker in value - ): - raise ValueError("refusal_markers must be a list of non-empty strings") - return tuple(dict.fromkeys(marker.strip().casefold() for marker in value)) - - 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 + response_format = stage.response_format + if retry and stage.retry_response_format is not None: + response_format = stage.retry_response_format + extra_body = {"response_format": response_format} if response_format else {} + if retry: + extra_body.update(self.retry_extra_body) + calls = [ + ChatRequest( + prompt=request.prompt, + image_data_url=( + None + if retry and stage.retry_without_image + else self._stage_image_urls.get(id(request.item)) + ), + requested_model=stage.model, + parameters=sampling, + extra_body=extra_body, + system_prompt=self.system_prompt, + image_detail=self.image_detail, ) - if following in '"\\/bfnrt' or valid_unicode: - repaired.append(char) - else: - repaired.append("\\\\") - index += 1 - return "".join(repaired) + for request in requests + ] + responses = self.api_loop.run_until_complete(self.endpoint_pool.run_many(calls)) + return [ + CaptionRejectedError(str(response)) + if isinstance(response, Exception) and self._is_semantic_caption_failure(response) + else response + for response in responses + ] @classmethod def _is_semantic_caption_failure(cls, error: Exception) -> bool: @@ -1036,17 +757,6 @@ def _is_semantic_caption_failure(cls, error: Exception) -> bool: return "message content" in message return False - @staticmethod - def _stage_context(item: ProcessingItem) -> Dict[str, Any]: - context = dict(item.metadata) - for previous_name, result in item.stage_results.items(): - for index, output in enumerate(result.outputs): - context[f"{previous_name}_output_{index}"] = output - context[result.output_field] = ( - result.outputs[0] if len(result.outputs) == 1 else result.outputs - ) - return context - def _sampling_for_stage(self, stage: ProcessingStage) -> Dict[str, Any]: sampling = dict(self.vllm_config.get("sampling", {})) if stage.sampling: diff --git a/src/caption_flow/workers/pipeline.py b/src/caption_flow/workers/pipeline.py new file mode 100644 index 0000000..89350bb --- /dev/null +++ b/src/caption_flow/workers/pipeline.py @@ -0,0 +1,127 @@ +"""Backend-independent stage execution and semantic caption recovery.""" + +import logging +from collections import defaultdict +from dataclasses import dataclass +from typing import Any, Protocol + +from ..models import ProcessingStage, StageResult +from ..utils.output_policy import OutputPolicy +from ..utils.prompt_template import PromptTemplateManager + +logger = logging.getLogger(__name__) + + +class CaptionRejectedError(ValueError): + """A backend rejected content before returning a caption; semantic retry is allowed.""" + + +@dataclass(frozen=True) +class StageRequest: + item: Any + prompt: str + + +def stage_context(item: Any) -> dict: + context = dict(item.metadata) + for name, result in item.stage_results.items(): + for index, output in enumerate(result.outputs): + context[f"{name}_output_{index}"] = output + context[result.output_field] = ( + result.outputs[0] if len(result.outputs) == 1 else result.outputs + ) + return context + + +class CaptionBackend(Protocol): + def _prepare_stage_batch(self, batch: list, stage: ProcessingStage) -> list: ... + def _generate_stage( + self, stage: ProcessingStage, requests: list[StageRequest], retry: bool + ) -> list: ... + def _output_policy(self, stage: ProcessingStage) -> OutputPolicy: ... + + +def _collect_outputs( + backend: CaptionBackend, + stage: ProcessingStage, + policy: OutputPolicy, + requests: list[StageRequest], + *, + retry: bool, +) -> tuple[dict, set]: + outputs = defaultdict(list) + retryable = set() + if not requests: + return outputs, retryable + responses = backend._generate_stage(stage, requests, retry) + for request, response in zip(requests, responses, strict=True): + key = id(request.item) + if isinstance(response, Exception): + logger.error( + "Generation failed for %s in %s: %s", request.item.item_key, stage.name, response + ) + if isinstance(response, CaptionRejectedError): + retryable.add(key) + continue + cleaned = policy.process(response) + if cleaned is None: + retryable.add(key) + else: + outputs[key].append(cleaned) + return outputs, retryable + + +def run_caption_pipeline( + backend: CaptionBackend, batch: list, stages: list[ProcessingStage], order: list[str] +) -> tuple[list, int]: + """Run the same acceptance, retry and stage-progression rules for every backend.""" + active = list(batch) + failed = 0 + stage_map = {stage.name: stage for stage in stages} + for name in order: + if not active: + break + stage = stage_map[name] + prepared = backend._prepare_stage_batch(active, stage) + failed += len(active) - len(prepared) + active = prepared + if not active: + break + policy = backend._output_policy(stage) + templates = PromptTemplateManager(stage.prompts) + requests = [ + StageRequest(item, prompt) + for item in active + for prompt in templates.format_all(stage_context(item)) + ] + outputs, retryable = _collect_outputs(backend, stage, policy, requests, retry=False) + if stage.retry_prompt: + retries = [ + StageRequest( + item, + PromptTemplateManager([stage.retry_prompt]).format_all(stage_context(item))[0], + ) + for item in active + if id(item) in retryable and not outputs[id(item)] + ] + recovered, _ = _collect_outputs(backend, stage, policy, retries, retry=True) + outputs.update(recovered) + + successful = [] + for item in active: + if outputs[id(item)]: + item.stage_results[name] = StageResult(name, stage.output_field, outputs[id(item)]) + successful.append(item) + else: + item.stage_results.pop(name, None) + failed += 1 + active = successful + + results = [] + for item in active: + fields = defaultdict(list) + for name in order: + result = item.stage_results[name] + fields[result.output_field].extend(result.outputs) + results.append((item, dict(fields))) + return results, failed diff --git a/tests/test_caption_pipeline.py b/tests/test_caption_pipeline.py new file mode 100644 index 0000000..4ec4ad8 --- /dev/null +++ b/tests/test_caption_pipeline.py @@ -0,0 +1,413 @@ +"""The same caption contract through the native and HTTP generation adapters.""" + +import asyncio +import json +from types import SimpleNamespace +from unittest.mock import AsyncMock, Mock + +import pytest +from PIL import Image + +from caption_flow.utils.image_processor import ImageProcessor +from caption_flow.utils.output_policy import OutputPolicy, validate_response_format +from caption_flow.utils.vllm_config import create_native_sampling_params +from caption_flow.workers.caption import CaptionWorker, MultiStageVLLMManager, ProcessingItem +from caption_flow.workers.openai_compatible import OpenAICompatibleWorker + + +def item(index=0): + return ProcessingItem( + unit_id="unit", + job_id=f"job-{index}", + chunk_id="chunk", + item_key=f"image-{index}", + item_index=index, + image=Image.new("RGB", (32, 16)), + image_data=b"", + metadata={"source": f"original-{index}"}, + ) + + +@pytest.fixture(params=["native", "api"]) +def backend(request, monkeypatch): + import vllm + + monkeypatch.setattr(vllm, "SamplingParams", SimpleNamespace) + monkeypatch.setattr( + vllm, + "sampling_params", + SimpleNamespace(StructuredOutputsParams=SimpleNamespace), + raising=False, + ) + config = {"server": "ws://localhost:8765", "token": "test", "batch_image_processing": False} + if request.param == "api": + monkeypatch.setenv("CAPTIONFLOW_TEST_KEY", "test") + worker = OpenAICompatibleWorker( + { + **config, + "openai_compatible": {"api_key_env": "CAPTIONFLOW_TEST_KEY"}, + } + ) + worker.api_loop = asyncio.new_event_loop() + worker._image_data_urls = Mock(side_effect=lambda batch: {id(i): "image" for i in batch}) + else: + worker = CaptionWorker(config) + worker.model_manager = MultiStageVLLMManager(0) + worker.model_manager.models["vision"] = Mock() + worker.model_manager.processors["vision"] = Mock() + worker.model_manager.tokenizers["vision"] = Mock() + worker._validate_and_split_batch = Mock(side_effect=lambda batch, *args: (batch, [])) + worker._build_vllm_input = Mock( + side_effect=lambda image, prompt, *args: {"image": image, "prompt": prompt} + ) + + def configure(**settings): + shared = {"model": "vision", "inference_prompts": ["Describe"], **settings} + worker.vllm_config = shared + worker.stages = worker._parse_stages_config(shared) + worker.stage_order = worker._topological_sort_stages(worker.stages) + if request.param == "native": + for stage in worker.stages: + worker.model_manager.create_sampling_params(stage, shared.get("sampling", {})) + + def respond(*batches): + if request.param == "api": + generator = AsyncMock(side_effect=list(batches)) + worker.endpoint_pool.run_many = generator + else: + generator = Mock( + side_effect=[ + [SimpleNamespace(outputs=[SimpleNamespace(text=text)]) for text in batch] + for batch in batches + ] + ) + worker.model_manager.models["vision"].generate = generator + return generator + + yield SimpleNamespace(worker=worker, kind=request.param, configure=configure, respond=respond) + if request.param == "api": + worker.api_loop.close() + + +def test_both_backends_share_the_stage_runner(): + assert ( + OpenAICompatibleWorker._process_batch_multi_stage + is CaptionWorker._process_batch_multi_stage + ) + + +def test_json_transformations_and_retry_constraints_work_on_both_backends(backend): + backend.configure( + output_processing={ + "validate_json_output": True, + "canonicalize_json_output": True, + "normalize_yxyx_bboxes": True, + "deduplicate_json_elements": True, + }, + sampling={"max_tokens": 500, "top_k": 20, "min_p": 0.1}, + response_format={ + "type": "json_schema", + "json_schema": {"name": "caption", "schema": {"type": "object"}}, + }, + retry_prompt="Retry JSON: {column:source}", + retry_without_image=True, + retry_sampling={"max_tokens": 1000, "top_k": 10}, + retry_response_format={"type": "json_object"}, + ) + generator = backend.respond( + ['{"truncated":', '{"elements": [{"bbox":[8,9,1,2]}, {"bbox":[8,9,1,2]}]}'], + ['{"caption":"recovered"}'], + ) + first, second = item(), item(1) + results = backend.worker._process_batch_multi_stage([first, second]) + assert results == [ + (first, {"captions": ['{"caption":"recovered"}']}), + (second, {"captions": ['{"elements":[{"bbox":[1,2,8,9]}]}']}), + ] + assert backend.worker.items_processed == 2 + assert backend.worker.items_failed == 0 + assert generator.call_count == 2 + primary, retry = generator.call_args_list + assert len(retry.args[0]) == 1 + if backend.kind == "native": + assert primary.args[1].top_k == 20 + assert primary.args[1].min_p == 0.1 + assert primary.args[1].structured_outputs.json == {"type": "object"} + assert retry.args[1].max_tokens == 1000 + assert retry.args[1].top_k == 10 + assert retry.args[1].structured_outputs.json_object is True + assert retry.args[0][0] == {"image": None, "prompt": "Retry JSON: original-0"} + # Retry overrides must never leak back into subsequent primary batches. + assert backend.worker.model_manager.sampling_params["default"] is primary.args[1] + else: + call = retry.args[0][0] + assert call.prompt == "Retry JSON: original-0" + assert call.image_data_url is None + assert call.parameters == {"max_tokens": 1000} + assert call.extra_body == {"response_format": {"type": "json_object"}} + assert primary.args[0][0].extra_body["response_format"]["type"] == "json_schema" + + +@pytest.mark.parametrize("bad_caption", ["", "I cannot describe it.", '{"truncated":']) +def test_failed_items_never_reach_the_next_stage_or_count_as_success(backend, bad_caption): + backend.configure( + output_processing={"validate_json_output": True}, + stages=[ + {"name": "first", "prompts": ["JSON"], "output_field": "raw", "retry_prompt": "Retry"}, + { + "name": "second", + "prompts": ["Expand {column:first_output_0}"], + "requires": ["first"], + "output_field": "caption", + "output_processing": {"validate_json_output": False}, + }, + ], + ) + generator = backend.respond([bad_caption, '{"caption":"ok"}'], [bad_caption], ["Expanded"]) + failed, accepted = item(), item(1) + results = backend.worker._process_batch_multi_stage([failed, accepted]) + assert results == [(accepted, {"raw": ['{"caption":"ok"}'], "caption": ["Expanded"]})] + assert failed.stage_results == {} + assert backend.worker.items_processed == 1 + assert backend.worker.items_failed == 1 + requests = generator.call_args_list[2].args[0] + assert len(requests) == 1 + prompt = requests[0]["prompt"] if backend.kind == "native" else requests[0].prompt + assert prompt == 'Expand {"caption":"ok"}' + if backend.kind == "api": + backend.worker._image_data_urls.assert_called_once() + assert backend.worker._stage_image_urls == {} + + +def test_configurable_refusal_policy_preserves_literal_ocr(backend): + backend.configure(output_processing={"refusal_markers": []}) + text = 'A sign reads "I cannot stay".' + backend.respond([text]) + image_item = item() + assert backend.worker._process_batch_multi_stage([image_item]) == [ + (image_item, {"captions": [text]}) + ] + assert backend.worker._clean_output(text) == text + + +def test_one_valid_prompt_is_enough_and_retry_is_not_redundant(backend): + backend.configure(inference_prompts=["First", "Second"], retry_prompt="Retry") + generator = backend.respond(["", "Accepted"]) + image_item = item() + assert backend.worker._process_batch_multi_stage([image_item]) == [ + (image_item, {"captions": ["Accepted"]}) + ] + assert generator.call_count == 1 + + +def test_multiple_outputs_flow_into_the_next_stage_context(backend): + backend.configure( + stages=[ + {"name": "first", "prompts": ["First", "Second"], "output_field": "draft"}, + {"name": "second", "prompts": ["Combine {column:draft}"], "requires": ["first"]}, + ] + ) + generator = backend.respond(["One", "Two"], ["Combined"]) + image_item = item() + assert backend.worker._process_batch_multi_stage([image_item]) == [ + (image_item, {"draft": ["One", "Two"], "captions": ["Combined"]}) + ] + calls = generator.call_args_list[1].args[0] + assert ( + calls[0]["prompt"] if backend.kind == "native" else calls[0].prompt + ) == "Combine One, Two" + + +def test_native_generation_rejects_mismatched_response_count(backend): + if backend.kind != "native": + pytest.skip("Native batch shape validation") + backend.configure() + backend.respond([]) + with pytest.raises(RuntimeError, match="different number of results"): + backend.worker._process_batch_multi_stage([item()]) + + +def test_native_resize_recovery_retains_limits_and_original_dimensions(monkeypatch): + worker = CaptionWorker( + {"server": "ws://localhost:8765", "token": "test", "batch_image_processing": False} + ) + worker.vllm_config = {"model": "vision", "max_model_len": 512} + stage = worker._parse_stages_config(worker.vllm_config)[0] + stage.max_model_len = 1024 + worker.model_manager = Mock() + worker.model_manager.get_model_for_stage.return_value = (Mock(), Mock(), Mock(), Mock()) + image_item = item() + ImageProcessor.prepare_for_inference(image_item) + checks = [] + + def validate(batch, *args): + checks.append(args[-1]) + if batch[0].image.width > 16: + return [], batch + return batch, [] + + monkeypatch.setattr(worker, "_validate_and_split_batch", validate) + resized = worker._prepare_stage_batch([image_item], stage) + assert len(resized) == 1 + assert resized[0].image.width == 16 + assert resized[0].metadata["image_width"] == 32 + assert checks == [1024, 1024, 1024] + monkeypatch.setattr(worker, "_validate_and_split_batch", lambda batch, *args: ([], batch)) + assert worker._prepare_stage_batch([image_item], stage) == [] + assert "token limit exceeded" in worker.result_queue.get_nowait()["error"] + + +def test_local_partial_policy_inherits_shared_validation(backend): + backend.worker.output_processing_overrides = {"canonicalize_json_output": True} + OutputPolicy.from_config(backend.worker.output_processing_overrides, partial=True) + backend.configure(output_processing={"validate_json_output": True}) + policy = backend.worker._output_policy(backend.worker.stages[0]) + assert policy.process('{ "caption" : "ok" }') == '{"caption":"ok"}' + with pytest.raises(ValueError, match="must be enabled"): + backend.configure() + + +def test_hot_reload_applies_policy_prompts_and_retry_without_reloading_models(backend): + backend.configure() + worker = backend.worker + worker._setup_vllm = Mock() + updated = { + **worker.vllm_config, + "inference_prompts": ["New prompt"], + "retry_prompt": "New retry", + "retry_sampling": {"max_tokens": 1500}, + "output_processing": {"validate_json_output": True}, + "retry_response_format": {"type": "json_object"}, + } + assert worker._handle_vllm_config_update(updated) + worker._setup_vllm.assert_not_called() + assert worker.stages[0].prompts == ["New prompt"] + assert worker.stages[0].retry_prompt == "New retry" + assert worker.stages[0].retry_sampling == {"max_tokens": 1500} + assert worker._output_policy(worker.stages[0]).validate_json_output + original_config, original_stages = worker.vllm_config, worker.stages + with pytest.raises(ValueError, match="schema must be a mapping"): + worker._handle_vllm_config_update({**updated, "response_format": {"type": "json_schema"}}) + assert worker.vllm_config is original_config + assert worker.stages is original_stages + + +def test_transport_failure_does_not_trigger_semantic_retry(backend): + if backend.kind != "api": + pytest.skip("Provider transport classification is specific to HTTP") + backend.configure(retry_prompt="Retry") + generator = backend.respond([RuntimeError("transport failed")]) + assert backend.worker._process_batch_multi_stage([item()]) == [] + assert generator.call_count == 1 + assert backend.worker.items_failed == 1 + + +@pytest.mark.parametrize("modern", [True, False]) +def test_native_constraints_support_both_vllm_apis(monkeypatch, modern): + import vllm + + constraint_class = "StructuredOutputsParams" if modern else "GuidedDecodingParams" + monkeypatch.setattr( + vllm, + "sampling_params", + SimpleNamespace(**{constraint_class: SimpleNamespace}), + raising=False, + ) + monkeypatch.setattr(vllm, "SamplingParams", SimpleNamespace) + result = create_native_sampling_params( + {"top_k": 12, "seed": 42, "max_tokens": 30}, {"type": "json_object"} + ) + assert result.top_k == 12 + assert result.seed == 42 + assert result.max_tokens == 30 + assert getattr(result, "structured_outputs" if modern else "guided_decoding").json_object + plain = create_native_sampling_params({}, {"type": "text"}) + assert not hasattr(plain, "structured_outputs") + assert not hasattr(plain, "guided_decoding") + with pytest.raises(ValueError, match="not both"): + create_native_sampling_params({"guided_decoding": {"json_object": True}}, {"type": "text"}) + + +@pytest.mark.parametrize( + "value", + [ + [], + {}, + {"type": "unknown"}, + {"type": "json_schema"}, + {"type": "json_schema", "json_schema": {"schema": []}}, + ], +) +def test_response_format_rejects_invalid_configuration(value): + with pytest.raises(ValueError, match="response_format"): + validate_response_format(value) + + +@pytest.mark.parametrize( + "text", + ['{"x":NaN}', '{"x":Infinity}', '{"x":-Infinity}', '{"x":1e999}', '{"unfinished":', "", None], +) +def test_shared_policy_rejects_unusable_json(text): + assert OutputPolicy(validate_json_output=True).process(text) is None + + +def test_escape_repair_preserves_existing_valid_backslashes_and_unicode(): + policy = OutputPolicy(validate_json_output=True, repair_invalid_json_escapes=True) + output = policy.process(r'{"valid":"\\\\","unicode":"\u00e9","bad":"\q"}') + assert json.loads(output) == {"valid": "\\\\", "unicode": "é", "bad": r"\q"} + assert policy.process(r'{"truncated":"\q"') is None + + +def test_recursive_normalization_deduplicates_after_children_and_preserves_other_arrays(): + policy = OutputPolicy( + validate_json_output=True, normalize_yxyx_bboxes=True, deduplicate_json_elements=True + ) + value = { + "elements": [{"elements": [1, 1]}, {"elements": [1]}], + "other": [1, 1], + "nested": [{"bbox": [True, 1, 2, 3]}, {"bbox": [4, 3, 2, 1]}], + } + assert json.loads(policy.process(json.dumps(value))) == { + "elements": [{"elements": [1]}], + "other": [1, 1], + "nested": [{"bbox": [True, 1, 2, 3]}, {"bbox": [2, 1, 4, 3]}], + } + + +def test_output_policy_rejects_unknown_options(): + with pytest.raises(ValueError, match="Unknown output_processing"): + OutputPolicy.from_config({"unknown": True}) + + +def test_native_decode_can_be_reused_and_original_dimensions_survive_resize(monkeypatch): + image_item = item() + image_item.image = None + image_item.image_data = b"encoded" + decode = Mock(return_value=Image.new("RGB", (640, 480))) + monkeypatch.setattr(ImageProcessor, "decode_image_data", decode) + first = ImageProcessor.prepare_for_inference(image_item) + assert ImageProcessor.prepare_for_inference(image_item) is first + decode.assert_called_once_with(b"encoded") + assert image_item.image_data == b"" + image_item.image = first.resize((320, 240)) + ImageProcessor.prepare_for_inference(image_item) + assert image_item.metadata["image_width"] == 640 + assert image_item.metadata["image_height"] == 480 + image_item.metadata["image_width"] = None + ImageProcessor.prepare_for_inference(image_item) + assert image_item.metadata["image_width"] == 320 + + +def test_native_text_only_retry_uses_chat_template_without_multimodal_payload(): + worker = CaptionWorker( + {"server": "ws://localhost:8765", "token": "test", "batch_image_processing": False} + ) + processor = Mock() + processor.apply_chat_template.return_value = "formatted" + tokenizer = Mock(return_value=SimpleNamespace(input_ids=[1, 2, 3])) + assert worker._build_vllm_input(None, "Retry", processor, tokenizer) == { + "prompt_token_ids": [1, 2, 3] + } + assert processor.apply_chat_template.call_args.args[0] == [ + {"role": "user", "content": [{"type": "text", "text": "Retry"}]} + ] diff --git a/tests/test_openai_compatible_worker.py b/tests/test_openai_compatible_worker.py index cf002eb..d31546e 100644 --- a/tests/test_openai_compatible_worker.py +++ b/tests/test_openai_compatible_worker.py @@ -640,8 +640,9 @@ def test_worker_validates_config_and_applies_shared_updates(): assert worker.mock_mode is True assert worker.stage_order == ["default"] + assert worker._process_batch_multi_stage([]) == [] with pytest.raises(RuntimeError, match="loop is not ready"): - worker._process_batch_multi_stage([]) + worker._prepare_stage_batch([], worker.stages[0]) def test_worker_configures_refusal_markers_from_shared_or_local_config():