[#169] Fix for stripping vision-tags that were sent by clients withou…

src/engine/ov_genai/vlm.py

@@ -133,6 +133,36 @@ def prepare_inputs(self,
 
         return tokenized_messages, ov_images
 
+    def _strip_stray_vision_tokens(self, prompt: str, ov_images: List[ov.Tensor]) -> str:
+        """
+        Enforce the vision-tag / image count invariant that OpenVINO's
+        inputs_embedder checks (vision_sequence.size() == n_visions).
+
+        prepare_inputs inserts exactly one native vision token per decoded
+        image, so a mismatch can only appear when the prompt *text* itself
+        carries a stray vision placeholder while no image is supplied. That is
+        exactly what happens when a VLM is driven with plain text whose
+        conversation mentions the model's own vision token (e.g. the model is
+        asked to read/analyze source code that references its own token). When
+        no image is provided we strip every stray token so the native vision
+        tag count (0) matches the provided image count (0); otherwise OpenVINO
+        aborts with "The number of native vision tags must match the number of
+        provided images/videos". Idempotent: a no-op when images exist or the
+        token is absent.
+        """
+        if ov_images or not self.vision_token:
+            return prompt
+        token_str = self._vision_token_for_index(0)
+        if not token_str or token_str not in prompt:
+            return prompt
+        stray_count = prompt.count(token_str)
+        logger.warning(
+            f"[{self.load_config.model_name}] prompt contains "
+            f"{stray_count} native vision token(s) but no image was provided; "
+            "stripping token(s) from the input, solves bug found in PR #169"
+        )
+        return prompt.replace(token_str, " ")
+
     def _resolve_prompt_and_images(
         self, gen_config: OVGenAI_GenConfig
     ) -> Tuple[str, List[ov.Tensor]]:
@@ -141,10 +171,13 @@ def _resolve_prompt_and_images(
         """
         if gen_config.input_ids:
             prompt = self.tokenizer.decode(gen_config.input_ids, skip_special_tokens=False)
-            return prompt, []
-        if gen_config.prompt:
-            return gen_config.prompt, []
-        return self.prepare_inputs(gen_config.messages, gen_config.tools, gen_config.chat_template_kwargs)
+            images: List[ov.Tensor] = []
+        elif gen_config.prompt:
+            prompt = gen_config.prompt
+            images = []
+        else:
+            prompt, images = self.prepare_inputs(gen_config.messages, gen_config.tools, gen_config.chat_template_kwargs)
+        return self._strip_stray_vision_tokens(prompt, images), images
 
     def generate_type(self, gen_config: OVGenAI_GenConfig):
         """
@@ -224,13 +257,24 @@ async def generate_stream(self,
         prompt, ov_images = self._resolve_prompt_and_images(gen_config)
 
         async def _run_generation():
-            return await asyncio.to_thread(
-                self.model_path.generate,
-                prompt=prompt,
-                **({'images': ov_images} if len(ov_images) > 0 else {}),
-                generation_config=generation_kwargs,
-                streamer=streamer,
-            )
+            try:
+                return await asyncio.to_thread(
+                    self.model_path.generate,
+                    prompt=prompt,
+                    **({'images': ov_images} if len(ov_images) > 0 else {}),
+                    generation_config=generation_kwargs,
+                    streamer=streamer,
+                )
+            except Exception:
+                # The streamer's end() (the thing that enqueues the None EOF
+                # sentinel) is only invoked on a *successful* finish. A generation
+                # that raises -- e.g. OpenVINO's "native vision tag count must
+                # match image count" check, which fires during input embedding
+                # before any token is emitted -- therefore enqueues no EOF. Push
+                # one here so the drain loop in generate_stream terminates and
+                # can surface the error instead of waiting on the queue forever.
+                streamer.text_queue.put_nowait(None)
+                raise
 
         gen_task = asyncio.create_task(_run_generation())
 
@@ -240,15 +284,22 @@ async def _run_generation():
                 if chunk is None:
                     break
                 yield chunk
+            # Stream fully drained: now await the generation task so that any
+            # error it raised is re-raised to the caller. Previously this await
+            # lived in the `finally` block, where a failing task masked the
+            # exception (and the trailing `yield metrics` never ran), leaving
+            # the HTTP client waiting for a stream that would never end.
+            result = await gen_task
         finally:
             # Clear active request tracking
             self._active_request_id = None
             self._active_streamer = None
-
-            result = await gen_task
-            perf_metrics = result.perf_metrics
-            metrics = self.collect_metrics(gen_config, perf_metrics)
-            yield metrics
+
+        # Reached only on the happy path: the stream drained and generation
+        # succeeded. Emit the metrics last.
+        perf_metrics = result.perf_metrics
+        metrics = self.collect_metrics(gen_config, perf_metrics)
+        yield metrics
 
     async def cancel(self, request_id: str) -> bool:
         """

tests/unit/test_ov_genai_vlm_unit.py

@@ -208,3 +208,48 @@ async def _run():
     assert vlm.vision_token is None
     registry.register_unload.assert_called_once_with("model-name")
     gc_mock.assert_called_once()
+
+
+def test_strip_stray_vision_token_no_image(load_config: ModelLoadConfig) -> None:
+    # Regression (vision-issue.log): a VLM driven with plain text whose
+    # conversation contains the model's own native vision token (e.g. being
+    # asked to "analyze the code" of a model that mentions its own token) must
+    # not trip OpenVINO's "vision_sequence.size() == n_visions" invariant.
+    vlm = OVGenAI_VLM(load_config)
+    vlm.vision_token = "<" + "img" + ">"
+    tok = vlm.vision_token
+
+    stripped = vlm._strip_stray_vision_tokens("A " + tok + " B", [])
+    assert tok not in stripped
+    assert vlm._strip_stray_vision_tokens(stripped, []) == stripped  # idempotent
+
+
+def test_strip_stray_vision_token_kept_when_image_present(load_config: ModelLoadConfig) -> None:
+    # When an image is actually provided, the tag is expected and must stay.
+    vlm = OVGenAI_VLM(load_config)
+    vlm.vision_token = "<" + "img" + ">"
+    tok = vlm.vision_token
+    assert vlm._strip_stray_vision_tokens(tok, [object()]) == tok
+
+
+def test_resolve_prompt_strips_stray_token_for_text_only_chat(
+    load_config: ModelLoadConfig,
+) -> None:
+    # Text-only messages -> no images -> any stray token in the templated
+    # prompt is removed before it reaches the pipeline.
+    vlm = OVGenAI_VLM(load_config)
+    vlm.vision_token = "<" + "img" + ">"
+    tok = vlm.vision_token
+
+    templated_prompt = "system text " + tok + " user asks to analyze the code"
+    tokenizer_mock = MagicMock()
+    tokenizer_mock.apply_chat_template.return_value = templated_prompt
+    vlm.tokenizer = tokenizer_mock
+
+    config = OVGenAI_GenConfig(
+        messages=[{"role": "user", "content": "analyze the code"}],
+        stream=True,
+    )
+    prompt, images = vlm._resolve_prompt_and_images(config)
+    assert images == []
+    assert tok not in prompt