diff --git a/.gitignore b/.gitignore index f100649..ffd309c 100644 --- a/.gitignore +++ b/.gitignore @@ -30,5 +30,6 @@ site/ build/ dist/ /models/**/* +/.scratchpad/ .scratchpad/ .plans/ \ No newline at end of file diff --git a/config.yaml b/config.yaml index c44fe7c..978a95c 100644 --- a/config.yaml +++ b/config.yaml @@ -14,6 +14,7 @@ models: model_type: llm model_path: /mnt/Ironwolf-4TB/Models/OpenVINO/Qwen3.5/Qwen3.5-0.8B-int8_asym-ov/ device: CPU + tool_call_parser: qwen35 runtime_config: PERFORMANCE_HINT: LATENCY scheduler_config: diff --git a/src/cli/groups/add.py b/src/cli/groups/add.py index 7de4413..01ffad7 100644 --- a/src/cli/groups/add.py +++ b/src/cli/groups/add.py @@ -62,8 +62,13 @@ default=None, type=float, help='Confidence threshold for accepting draft tokens.') +@click.option('--tool-call-parser', + type=click.Choice(['qwen35', 'hermes', 'gemma4']), + required=False, + default=None, + help='Tool-call output format for this model (qwen35 XML, hermes JSON, or gemma4 call syntax). llm/vlm only; required for tool calling.') @click.pass_context -def add(ctx, model_path, model_name, engine, model_type, device, runtime_config, scheduler_config, cache_dir, draft_model_path, draft_device, num_assistant_tokens, assistant_confidence_threshold): +def add(ctx, model_path, model_name, engine, model_type, device, runtime_config, scheduler_config, cache_dir, draft_model_path, draft_device, num_assistant_tokens, assistant_confidence_threshold, tool_call_parser): """- Add a model configuration to the config file.""" # Validate model path @@ -128,6 +133,8 @@ def add(ctx, model_path, model_name, engine, model_type, device, runtime_config, load_config["num_assistant_tokens"] = num_assistant_tokens if assistant_confidence_threshold is not None: load_config["assistant_confidence_threshold"] = assistant_confidence_threshold + if tool_call_parser: + load_config["tool_call_parser"] = tool_call_parser ctx.obj.server_config.save_model_config(model_name, load_config) console.print(f"[green]Model configuration saved:[/green] {model_name}") diff --git a/src/engine/ov_genai/llm.py b/src/engine/ov_genai/llm.py index db06022..dd588ed 100755 --- a/src/engine/ov_genai/llm.py +++ b/src/engine/ov_genai/llm.py @@ -16,7 +16,7 @@ from src.server.schemas.modeling.contract_ovgenai_llm_and_vlm import OVGenAI_GenConfig from src.server.model_registry import ModelRegistry from src.server.schemas.registration import ModelLoadConfig -from src.engine.ov_genai.streamers import ChunkStreamer +from src.engine.ov_genai.streamers import ensure_tool_call_parser, select_streamer from src.server.utils.chat import flatten_messages logger = logging.getLogger(__name__) @@ -83,6 +83,7 @@ async def generate_text(self, gen_config: OVGenAI_GenConfig) -> AsyncIterator[Un Async non-streaming text generation. Yields in order: metrics (dict), new_text (str). """ + ensure_tool_call_parser(gen_config, self.load_config) generation_kwargs = self.create_generation_config(gen_config) # Support pre-encoded input_ids, raw prompts, and chat messages @@ -101,7 +102,13 @@ async def generate_text(self, gen_config: OVGenAI_GenConfig) -> AsyncIterator[Un perf_metrics = result.perf_metrics decoder_tokenizer = self.model.get_tokenizer() - text = decoder_tokenizer.decode(result.tokens)[0] if getattr(result, "tokens", None) else "" + # gemma4 protocol tags are special=True: keep them in the decoded text + # so the route-level parse_generation can split reasoning/tool calls. + keep_special = getattr(gen_config, "tool_call_parser", None) == "gemma4" + text = ( + decoder_tokenizer.decode(result.tokens, skip_special_tokens=not keep_special)[0] + if getattr(result, "tokens", None) else "" + ) metrics_dict = self.collect_metrics(gen_config, perf_metrics) yield metrics_dict @@ -112,10 +119,10 @@ async def generate_stream(self, gen_config: OVGenAI_GenConfig) -> AsyncIterator[ Async streaming text generation. Yields token chunks (str) as they arrive, then metrics (dict), then final new_text (str). """ - + ensure_tool_call_parser(gen_config, self.load_config) generation_kwargs = self.create_generation_config(gen_config) decoder_tokenizer = self.model.get_tokenizer() - streamer = ChunkStreamer(decoder_tokenizer, gen_config) + streamer = select_streamer(decoder_tokenizer, gen_config) # Track active request and streamer for cancellation self._active_request_id = gen_config.request_id diff --git a/src/engine/ov_genai/qwen_tool_parser.py b/src/engine/ov_genai/qwen_tool_parser.py deleted file mode 100644 index 9169110..0000000 --- a/src/engine/ov_genai/qwen_tool_parser.py +++ /dev/null @@ -1,520 +0,0 @@ -"""Qwen3.5 XML tool-call parser + streamer built directly on StreamerBase. - -Avoids TextParserStreamer/IncrementalParser entirely (a live VLMPipeline run -with a Python TextParserStreamer subclass hung deterministically after tool -call completion). Instead, ToolCallStreamer(StreamerBase) does its own -incremental decode (cumulative decode + delta slicing, like OpenArc's -ChunkStreamer) and feeds text deltas through: - - ReasoningSplitter - everything before is reasoning_content - (Qwen3.5 emits no opening tag) - QwenXmlToolCallParser - state machine for the Qwen XML tool format: - - - - - VALUE - - ... - - - -Parallel calls are sequential blocks; text may precede the first -block. The parser strips tool XML from content and emits OpenAI-style -streaming fragments: - - {"index": i, "id": ..., "type": "function", - "function": {"name": "", "arguments": ""}} # call start - {"index": i, "function": {"name": NAME}} # name known - {"index": i, "function": {"arguments": FRAGMENT}} # argument deltas - -Concatenating all `arguments` fragments for an index yields complete JSON. -Queued delta messages look like: - {"content": str} / {"reasoning_content": str} / {"tool_calls": [frag, ...]} -None signals completion. -""" -import asyncio -import itertools -import json -from typing import Any, Dict, List, Optional, Tuple, Union - -import openvino_genai -from openvino_genai import StreamerBase, StreamingStatus - -TOOL_OPEN = "" -TOOL_CLOSE = "" -FUNC_OPEN = " int: - """Length of the longest suffix of `text` that is a proper prefix of a tag. - - Anything before that suffix is safe to process now; the suffix must wait - for more deltas (it could grow into a full tag). - """ - max_hold = 0 - upper = max(len(t) for t in tags) - 1 - for n in range(1, min(len(text), upper) + 1): - suffix = text[-n:] - if any(t != suffix and t.startswith(suffix) for t in tags): - max_hold = n - return max_hold - - -def _json_escape_fragment(s: str) -> str: - """Escape a raw string fragment for embedding inside a JSON string.""" - out = [] - for ch in s: - if ch == '"': - out.append('\\"') - elif ch == "\\": - out.append("\\\\") - elif ch == "\n": - out.append("\\n") - elif ch == "\t": - out.append("\\t") - elif ch == "\r": - out.append("\\r") - elif ord(ch) < 0x20: - out.append(f"\\u{ord(ch):04x}") - else: - out.append(ch) - return "".join(out) - - -class ReasoningSplitter: - """Splits the leading thinking block from regular content. - - Qwen3.5 with enable_thinking=True emits reasoning immediately (no opening - tag) and terminates it with . Everything before the close tag is - reasoning; everything after is content. - """ - - def __init__(self, close_tag: str = THINK_CLOSE, enabled: bool = True): - self._close_tag = close_tag - self._buf = "" - # Only treat the stream as reasoning when the prompt actually opened a - # block (chat template appends '\n' to the prompt when - # enable_thinking=True; with False it pre-closes an empty block and the - # stream is pure content). - self.in_reasoning = enabled - - def feed(self, text: str) -> Tuple[str, str]: - """Returns (reasoning_delta, content_delta).""" - self._buf += text - reasoning_out, content_out = "", "" - while self._buf: - if not self.in_reasoning: - content_out += self._buf - self._buf = "" - break - idx = self._buf.find(self._close_tag) - if idx != -1: - reasoning_out += self._buf[:idx] - self._buf = self._buf[idx + len(self._close_tag):] - self.in_reasoning = False - continue - hold = _holdback_suffix(self._buf, (self._close_tag,)) - reasoning_out += self._buf[: len(self._buf) - hold] - self._buf = self._buf[len(self._buf) - hold:] - break - return reasoning_out, content_out - - -class QwenXmlToolCallParser: - """Incremental parser for Qwen XML tool calls. - - `tools` is the OpenAI-style tools list from the request; used for argument - type coercion (booleans arrive Python-style: True/False). - """ - - def __init__(self, tools: Optional[List[Dict[str, Any]]] = None): - self._param_types: Dict[str, Dict[str, str]] = {} - for tool in tools or []: - fn = tool.get("function", {}) - props = (fn.get("parameters") or {}).get("properties") or {} - self._param_types[fn.get("name", "")] = { - k: (v.get("type") or "string") for k, v in props.items() - } - self.reset() - - def reset(self) -> None: - self._ids = itertools.count() - self._buf = "" - self._state = CONTENT - self._calls: List[Dict[str, Any]] = [] - self._cur: Optional[Dict[str, Any]] = None # current call under construction - self._param_name = "" - self._param_raw = "" # buffered value (non-string types) - self._param_index = 0 # params seen in current call - self._value_started = False # first chunk of current value? - self._pending_newline = False # trailing \n hold-back in values - self.status = StreamingStatus.RUNNING - self.errors: List[str] = [] - - @property - def tool_calls(self) -> List[Dict[str, Any]]: - """All calls seen so far, in final OpenAI non-streaming shape.""" - return list(self._calls) - - def feed(self, delta_text: str) -> Tuple[str, List[Dict[str, Any]]]: - """Consume a text delta. Returns (content_text, tool_call_fragments).""" - if not delta_text: - return "", [] - self._buf += delta_text - - content_out: List[str] = [] - fragments: List[Dict[str, Any]] = [] - - while self._buf: - if not self._step(content_out, fragments): - break # need more input - - return "".join(content_out), fragments - - # -- state machine ---------------------------------------------------------- - - def _step(self, content_out: list, fragments: list) -> bool: - """Advance the machine once. False = wait for more input.""" - buf = self._buf - - if self._state == CONTENT: - idx = buf.find(TOOL_OPEN) - if idx == -1: - hold = _holdback_suffix(buf, (TOOL_OPEN,)) - safe, self._buf = buf[: len(buf) - hold], buf[len(buf) - hold:] - if safe and not self._calls: - content_out.append(safe) - elif safe and safe.strip(): - # After a completed call, real text means the model is - # rambling past its answer -> stop generation. - self.status = StreamingStatus.TOOL_CALL_STOP - # else: trailing whitespace after the final call -> drop - return False # remaining buf is a partial tag prefix; wait - if idx > 0: - pre, self._buf = buf[:idx], buf[idx:] - if not self._calls: - content_out.append(pre) - elif pre.strip(): - self.status = StreamingStatus.TOOL_CALL_STOP - return True - self._buf = buf[len(TOOL_OPEN):] - self._start_call(fragments) - self._state = IN_TOOL_CALL - return True - - if self._state == IN_TOOL_CALL: - return self._expect_tag(buf, (FUNC_OPEN,), IN_FUNC_NAME) - - if self._state == IN_FUNC_NAME: - gt = buf.find(GT) - if gt == -1: - if "<" in buf: # names never contain '<'; bail on garbage - self.errors.append(f"malformed function name: {buf!r}") - self._state = CONTENT - return True - return False - name = buf[:gt].strip() - self._buf = buf[gt + 1:] - self._cur["function"]["name"] = name - fragments.append({ - "index": len(self._calls) - 1, - "function": {"name": name}, - }) - self._state = IN_FUNCTION - return True - - if self._state == IN_FUNCTION: - stripped = buf.lstrip() - if stripped.startswith(PARAM_OPEN): - self._buf = stripped[len(PARAM_OPEN):] - self._state = IN_PARAM_NAME - return True - if stripped.startswith(FUNC_CLOSE): - self._buf = stripped[len(FUNC_CLOSE):] - self._close_call(fragments) - self._state = AFTER_FUNCTION - return True - if not stripped or any( - t.startswith(stripped) and t != stripped - for t in (PARAM_OPEN, FUNC_CLOSE) - ): - return False # whitespace only, or partial tag - self.errors.append(f"unexpected text inside : {stripped[:20]!r}") - self._buf = stripped[1:] - return True - - if self._state == IN_PARAM_NAME: - gt = buf.find(GT) - if gt == -1: - if "<" in buf: - self.errors.append(f"malformed parameter name: {buf!r}") - self._state = IN_FUNCTION - return True - return False - self._param_name = buf[:gt].strip() - self._buf = buf[gt + 1:] - self._param_raw = "" - self._value_started = False - self._pending_newline = False - self._state = IN_PARAM - # open the JSON member - ptype = self._param_type(self._cur["function"]["name"], self._param_name) - sep = "{" if self._param_index == 0 else ", " - self._param_index += 1 - self._frag(fragments, f'{sep}{json.dumps(self._param_name)}: ') - if ptype == "string": - self._frag(fragments, '"') - return True - - if self._state == IN_PARAM: - idx = buf.find(PARAM_CLOSE) - ptype = self._param_type(self._cur["function"]["name"], self._param_name) - if idx == -1: - hold = _holdback_suffix(buf, (PARAM_CLOSE,)) - safe, self._buf = buf[: len(buf) - hold], buf[len(buf) - hold:] - if safe: - self._consume_value(safe, fragments, ptype, final=False) - return False - value, self._buf = buf[:idx], buf[idx + len(PARAM_CLOSE):] - self._consume_value(value, fragments, ptype, final=True) - self._state = IN_FUNCTION - return True - - if self._state == AFTER_FUNCTION: - return self._expect_tag(buf, (TOOL_CLOSE,), CONTENT) - - return False - - def _expect_tag(self, buf: str, tags, next_state: str) -> bool: - stripped = buf.lstrip() - for t in tags: - if stripped.startswith(t): - self._buf = stripped[len(t):] - self._state = next_state - return True - if not stripped or any( - t.startswith(stripped) and t != stripped for t in tags - ): - return False # whitespace only, or partial tag - self.errors.append(f"expected {tags}, got {stripped[:20]!r}") - self._buf = stripped[1:] - return True - - # -- call/fragment helpers ---------------------------------------------------- - - def _frag(self, fragments: list, arguments: str): - fragments.append({ - "index": len(self._calls) - 1, - "function": {"arguments": arguments}, - }) - self._cur["function"]["arguments"] += arguments - - def _start_call(self, fragments: list): - self._cur = { - "id": f"call_{next(self._ids):024x}", - "type": "function", - "function": {"name": "", "arguments": ""}, - } - self._calls.append(self._cur) - self._param_index = 0 - fragments.append({ - "index": len(self._calls) - 1, - "id": self._cur["id"], - "type": "function", - "function": {"name": "", "arguments": ""}, - }) - - def _close_call(self, fragments: list): - self._frag(fragments, "}" if self._param_index > 0 else "{}") - self._cur["done"] = True - - def _param_type(self, func_name: str, param: str) -> str: - return self._param_types.get(func_name, {}).get(param, "string") - - def _consume_value(self, text: str, fragments: list, ptype: str, final: bool): - # drop the single wrapper newline right after - if not self._value_started and text.startswith("\n"): - text = text[1:] - if text: - self._value_started = True - if ptype == "string": - # hold back a trailing newline: it may be the wrapper before - # ; if more value text follows we flush it then. - if self._pending_newline: - text = "\n" + text - self._pending_newline = False - if text.endswith("\n"): - if final: - text = text[:-1] - else: - self._pending_newline = True - text = text[:-1] - if text: - self._param_raw += text - self._frag(fragments, _json_escape_fragment(text)) - if final: - self._frag(fragments, '"') - else: - # buffer short scalar values; coerce + emit once complete - self._param_raw += text - if final: - self._frag(fragments, self._coerce(self._param_raw, ptype)) - - def _coerce(self, raw: str, ptype: str) -> str: - """Serialize a raw XML parameter value as JSON text.""" - v = raw.strip() - if ptype in ("integer", "number"): - return v # model emits digits - if ptype == "boolean": - low = v.lower() - if low in ("true", "1", "yes"): - return "true" - if low in ("false", "0", "no"): - return "false" - return json.dumps(v) # fallback: keep as string - if ptype in ("array", "object"): - try: - return json.dumps(json.loads(v)) - except (json.JSONDecodeError, TypeError): - return json.dumps(v) - return json.dumps(raw.strip("\n")) - - def finalize(self) -> None: - """Call at end of generation to surface truncated structures.""" - if self._cur is not None and not self._cur.get("done"): - self.errors.append("generation ended with unterminated tool call") - if self._state != CONTENT: - self.errors.append(f"generation ended in state {self._state}") - - -class ToolCallStreamer(StreamerBase): - """StreamerBase doing its own incremental decode + tool-call parsing. - - Mirrors ChunkStreamer's contract: decoded/parsed delta messages are put on - `queue` from the generation thread; None signals completion. Return value - controls generation: TOOL_CALL_STOP when the model rambles past a closed - tool call, CANCEL after cancel(). - """ - - def __init__(self, tokenizer, tools: Optional[list] = None, - enable_thinking: bool = True): - super().__init__() - self.tokenizer = tokenizer - self.reasoning = ReasoningSplitter(enabled=enable_thinking) - self.tool_parser = QwenXmlToolCallParser(tools) - self.tokens_cache: List[int] = [] - self.last_print_len = 0 - self._wrapper_ids: Dict[int, str] = {} - for tag in WRAPPER_TAGS: - ids = tokenizer.encode(tag).input_ids.data.tolist()[0] - if len(ids) == 1: - self._wrapper_ids[ids[0]] = tag - self.queue: "asyncio.Queue[Optional[dict]]" = asyncio.Queue() - self._cancelled = asyncio.Event() - try: - self._loop = asyncio.get_running_loop() - except RuntimeError: - self._loop = None # offline use: no loop, put_nowait directly - - def _enqueue(self, msg) -> None: - # write()/end() run on the C++ generation thread; asyncio.Queue is not - # thread-safe, so hop through call_soon_threadsafe when a loop exists. - if self._loop is not None: - self._loop.call_soon_threadsafe(self.queue.put_nowait, msg) - else: - self.queue.put_nowait(msg) - - def _process_delta(self, delta: str) -> None: - reasoning, text = self.reasoning.feed(delta) - content, fragments = self.tool_parser.feed(text) - msg: Dict[str, Any] = {} - if reasoning: - msg["reasoning_content"] = reasoning - if content: - msg["content"] = content - if fragments: - msg["tool_calls"] = fragments - if msg: - self._enqueue(msg) - - def _decode_available(self) -> None: - text = self.tokenizer.decode(self.tokens_cache) - if len(text) > self.last_print_len: - delta = text[self.last_print_len:] - if REPLACEMENT_CHAR in delta: - # partial UTF-8 at the boundary; wait for more tokens - return - self.last_print_len = len(text) - self._process_delta(delta) - - def _flush_text(self) -> None: - """Decode and process everything left in the cache, then reset it. - - Called before an atomically-detected wrapper token: the special token - is a hard boundary, so any held-back partial tag in the text parser - must resolve now (it can't be a prefix of a wrapper tag anymore). - """ - if not self.tokens_cache: - return - text = self.tokenizer.decode(self.tokens_cache) - if len(text) > self.last_print_len: - self._process_delta(text[self.last_print_len:]) - self.tokens_cache = [] - self.last_print_len = 0 - - def write(self, token: Union[int, List[int]]) -> StreamingStatus: - if self._cancelled.is_set(): - self._enqueue(None) - return StreamingStatus.CANCEL - - ids = token if isinstance(token, list) else [token] - for tid in ids: - tag = self._wrapper_ids.get(int(tid)) - if tag is not None: - self._flush_text() - if tag == THINK_OPEN: - # Opening tag is baked into the prompt by the chat - # template; a stray one in the stream is a no-op. - continue - # Atomic delivery: hold-back in the text parser resolves - # immediately, so this can never straddle chunks. - self._process_delta(tag) - else: - self.tokens_cache.append(int(tid)) - self._decode_available() - - return self.tool_parser.status - - def end(self) -> None: - self._flush_text() - self.tool_parser.finalize() - self._enqueue(None) - - def cancel(self) -> None: - self._cancelled.set() - - def is_cancelled(self) -> bool: - return self._cancelled.is_set() diff --git a/src/engine/ov_genai/streamers.py b/src/engine/ov_genai/streamers.py index 64d20ce..4e23596 100644 --- a/src/engine/ov_genai/streamers.py +++ b/src/engine/ov_genai/streamers.py @@ -4,6 +4,8 @@ from openvino_genai import StreamerBase from src.server.schemas.modeling.contract_ovgenai_llm_and_vlm import OVGenAI_GenConfig +from src.engine.ov_genai.tool_parse import gemma4 as gemma4_tool_parse +from src.engine.ov_genai.tool_parse import qwen35 as qwen35_tool_parse class ChunkStreamer(StreamerBase): @@ -79,4 +81,36 @@ def end(self) -> None: chunk = text[self.last_print_len:] if chunk: self._enqueue(chunk) - self._enqueue(None) \ No newline at end of file + self._enqueue(None) + + +def ensure_tool_call_parser(gen_config: OVGenAI_GenConfig, load_config) -> None: + """Fall back to the model's registered tool-call parser when the request + does not carry one (routes set gen_config.tool_call_parser; direct engine + callers such as tests and bench rely on the load-time registration).""" + if getattr(gen_config, "tool_call_parser", None) is None: + parser = getattr(load_config, "tool_call_parser", None) + if parser is not None: + gen_config.tool_call_parser = parser.value + + +def select_streamer(tokenizer, gen_config: OVGenAI_GenConfig) -> StreamerBase: + """Pick the streaming implementation for a generation request. + + Tool-call requests on qwen35 models stream through Qwen35ToolCallStreamer + (token-ID block boundaries, parsed OpenAI deltas). gemma4 requests use + Gemma4ToolCallStreamer whenever tools are requested OR thinking is enabled + (its thought-channel tags are special=True, so the plain text path cannot + split reasoning); everything else uses ChunkStreamer. All of them enqueue + on .text_queue, so consumers are unaffected. + """ + parser_name = getattr(gen_config, "tool_call_parser", None) + if gen_config.tools and parser_name == "qwen35": + return qwen35_tool_parse.Qwen35ToolCallStreamer(tokenizer, gen_config) + if parser_name == "gemma4": + thinking = True + if getattr(gen_config, "chat_template_kwargs", None): + thinking = bool(gen_config.chat_template_kwargs.get("enable_thinking", True)) + if gen_config.tools or thinking: + return gemma4_tool_parse.Gemma4ToolCallStreamer(tokenizer, gen_config) + return ChunkStreamer(tokenizer, gen_config) \ No newline at end of file diff --git a/src/engine/ov_genai/tool_parse/__init__.py b/src/engine/ov_genai/tool_parse/__init__.py new file mode 100644 index 0000000..ecd5708 --- /dev/null +++ b/src/engine/ov_genai/tool_parse/__init__.py @@ -0,0 +1 @@ +"""Tool-call output parsers, selectable per model at load time.""" diff --git a/src/engine/ov_genai/tool_parse/gemma4.py b/src/engine/ov_genai/tool_parse/gemma4.py new file mode 100644 index 0000000..c606140 --- /dev/null +++ b/src/engine/ov_genai/tool_parse/gemma4.py @@ -0,0 +1,846 @@ +"""Gemma 4 tool-call parser (Chimera-X-26B-A4B and other Gemma-4-derived exports). + +Unlike Qwen3.5, every Gemma 4 protocol tag is special=True, so the tags NEVER +appear in decoded text; they exist only as token IDs. Token-ID boundary +detection is mandatory, not just preferred: + + - Gemma4ToolCallStreamer sees raw ids in write() and intercepts the protocol + ids (they would be stripped from its own incremental decode otherwise); + the payload text ('call:name{key:value, ...}') arrives token-less between + boundaries. + - Quote tokens (<|"|>, id 52) are special too, so argument strings arrive + UNQUOTED (city:Paris, not city:"Paris"). Values are emitted as JSON + strings, no schema-driven type coercion (same decision as qwen35). + +Assistant protocol (chat_template.jinja): + reasoning : <|channel>thought\n \n + tool call : <|tool_call>call:{:, ...} + turn end : (eos_token_id 106) + +Consumers: + Gemma4ToolCallStreamer - StreamerBase subclass used by the engine for + streaming requests; does its own incremental decode + and enqueues {"chat_delta": [...]} dicts. Deliberately + avoids TextParserStreamer, which deadlocks when + generate runs on a worker thread under asyncio. + parse_generation() - raw-text splitter for non-streaming requests (the + engine decodes with skip_special_tokens=False for + gemma4-parser models so the tags survive as text). + +Parsed OpenAI streaming fragments (same contract as qwen35): + {"index": i, "id": ..., "type": "function", + "function": {"name": "", "arguments": ""}} # call start + {"index": i, "function": {"name": NAME}} # name known + {"index": i, "function": {"arguments": FRAGMENT}} # argument deltas + +Concatenating all `arguments` fragments for an index yields complete JSON. +""" +import asyncio +import json +import logging +import re +import uuid +from typing import Any, Callable, Dict, List, Optional, Tuple + +from openvino_genai import IncrementalParser, StreamerBase, StreamingStatus + +from src.engine.ov_genai.tool_parse.qwen35 import _is_partial, _json_escape_fragment + +logger = logging.getLogger(__name__) + +# Tag text, used only by the raw-text path (parse_generation). Invisible to +# the token-ID path: all of these are special=True tokens. +TOOL_OPEN = "<|tool_call>" +TOOL_CLOSE = "" +CHANNEL_OPEN = "<|channel>" +CHANNEL_CLOSE = "" +TURN_END = "" +EOS = "" +CALL_PREFIX = "call:" + +# Protocol token IDs (validated against the Chimera-X-26B tokenizer; the +# streamer re-checks them via encode lookups at construction time). +TOOL_OPEN_ID = 48 +TOOL_CLOSE_ID = 49 +CHANNEL_OPEN_ID = 100 +CHANNEL_CLOSE_ID = 101 +TOOL_BOUNDARY_IDS = (TOOL_OPEN_ID, TOOL_CLOSE_ID) +THOUGHT_HEADER = "thought" + +# Parser states +OUTSIDE = "outside" # no open block: text is content (or ramble after a call) +HEAD = "head" # awaiting 'call:' +NAME = "name" # awaiting '{' that ends the function name +ARGS_KEY = "args_key" # awaiting 'key:' (or '}' for a zero-argument call) +VALUE = "value" # inside an argument value (brace-depth tracked) +TAIL = "tail" # payload closed, awaiting the close boundary + +_PAYLOAD_RE = re.compile(r"call:([A-Za-z_][\w]*)\{(.*)\}", re.DOTALL) + + +def wants_engine_stream(tools: Optional[List[Dict[str, Any]]], thinking_enabled: bool) -> bool: + """True when a gemma4 request must use the token-ID engine streamer. + + Thought-channel tags are special=True, so a plain ChunkStreamer text path + cannot split reasoning; the engine streamer is required whenever tools are + requested OR thinking is enabled. With both off the model answers with + plain content (the template injects an empty thought channel) and + ChunkStreamer suffices. + """ + return bool(tools) or bool(thinking_enabled) + + +class Gemma4ChannelSplitter: + """Splits Gemma 4 thought channels from content by token ID. + + All channel tags are special=True, so boundaries are ID-only: a + CHANNEL_OPEN_ID opens a channel, CHANNEL_CLOSE_ID closes it. The visible + 'thought\\n' header text is consumed so it does not leak into content; + channels with other names pass through as content without their wrapper + tags (conservative). Thought text streams out as reasoning deltas with a + one-line tail held back so the newline before the close tag is trimmed. + + The driver contract is exact attribution: each feed() call either carries + text (no protocol ids) or protocol ids (empty text). + """ + + def __init__(self): + self._state = "content" # content | header | thought | opaque + self._header = "" # accumulated channel-name line + self._reasoning = "" # held reasoning tail + + def feed(self, text: str, token_ids) -> Tuple[str, str, List[int]]: + """Returns (reasoning_delta, content_delta, passthrough_token_ids).""" + reason_parts: List[str] = [] + content_parts: List[str] = [] + passthrough: List[int] = [] + for token in (token_ids or []): + tid = int(token) + if tid == CHANNEL_OPEN_ID and self._state == "content": + self._state = "header" + self._header = "" + continue + if tid == CHANNEL_CLOSE_ID and self._state in ("thought", "opaque"): + if self._state == "thought": + tail = self._reasoning.rstrip("\n") + if tail: + reason_parts.append(tail) + self._reasoning = "" + self._state = "content" + continue + passthrough.append(tid) + if text: + if self._state == "header": + self._header += text + if "\n" in self._header: + line, _, rest = self._header.partition("\n") + self._header = "" + if line.strip() == THOUGHT_HEADER: + self._state = "thought" + self._reasoning += rest + else: + self._state = "opaque" + content_parts.append(line + "\n" + rest) + elif self._state == "thought": + # Stream thought text out, keeping a one-line tail so the + # newline before the close tag can be trimmed cleanly. + self._reasoning += text + cut = self._reasoning.rfind("\n") + if cut > 0: + reason_parts.append(self._reasoning[:cut]) + self._reasoning = self._reasoning[cut:] + else: # content | opaque + content_parts.append(text) + return "".join(reason_parts), "".join(content_parts), passthrough + + +class Gemma4ToolCallParser(IncrementalParser): + """IncrementalParser for Gemma 4 tool calls, bound by token IDs. + + Payload grammar: 'call:{:, ...}' where values may be + nested brace objects. Boundary detection is token-ID only (the tag text + never appears in the stream); payload text is streamed through the state + machine and emitted as incremental OpenAI argument fragments: + + 'call:get_weather{city:Paris, unit:{...}}' -> + {"index": 0, "id": ..., "type": "function", + "function": {"name": "", "arguments": ""}} + {"index": 0, "function": {"name": "get_weather"}} + {"index": 0, "function": {"arguments": '{"city": "'}} + {"index": 0, "function": {"arguments": 'Paris"}} + {"index": 0, "function": {"arguments": ', "unit": "'}} + ... + + Values are JSON strings (unquoted in the payload because quote tokens are + special=True). With stop_after_tool_call=True the parser requests + StreamingStatus.TOOL_CALL_STOP right after each complete call; by default + (False) sequential parallel calls are allowed and generation is only + stopped when real text follows a completed call. + """ + + def __init__( + self, + on_fragment: Optional[Callable[[Dict[str, Any]], None]] = None, + stop_after_tool_call: bool = False, + ): + super().__init__() + self.on_fragment = on_fragment + self._stop_after = stop_after_tool_call + self.reset() + + def reset(self) -> None: + self._state = OUTSIDE + self._pending = "" + self._cur: Optional[Dict[str, Any]] = None # call under construction + self._cur_index = -1 + self._calls_seen = 0 # calls started (open boundary seen) + self._calls: List[Dict[str, Any]] = [] # completed calls, OpenAI shape + self._param_index = 0 + self._depth = 0 + self._skip_rest = False + self._stopped = False + self.errors: List[str] = [] + self.status = StreamingStatus.RUNNING + self.set_status(StreamingStatus.RUNNING) + + @property + def completed_calls(self) -> List[Dict[str, Any]]: + """All completed calls, in final OpenAI non-streaming shape.""" + return list(self._calls) + + def finalize(self) -> None: + """Call at end of generation: force-close an unterminated call and + surface malformed structures in `errors`.""" + if self._cur is not None: + self.errors.append("generation ended with unterminated tool call") + if self._state == VALUE: + self._end_value() + self._finish_call({}) + elif self._state != OUTSIDE: + self.errors.append(f"generation ended in state {self._state}") + + # -- IncrementalParser API --------------------------------------------------- + + def parse(self, msg: dict, delta_text: str, delta_tokens=None) -> str: + """Consume one decoded delta. Returns content text for this delta. + + Mutates `msg` with completed calls under msg["tool_calls"][str(index)] + and emits streaming fragments through `on_fragment`. The driver + contract is that deltas either carry text (no protocol ids) or + protocol ids (empty text); if both arrive, text is processed first. + """ + if self._stopped: + return "" + out: List[str] = [] + if delta_text: + self._pending += delta_text + self._machine(out, final=bool(delta_tokens)) + for token in (delta_tokens or []): + tid = int(token) + if tid in TOOL_BOUNDARY_IDS: + self._machine(out, final=True) + self._apply_boundary(tid, msg) + if self._stopped: + self._pending = "" + break + return "".join(out) + + # -- state machine ------------------------------------------------------------ + + def _machine(self, out: List[str], final: bool) -> None: + """Advance the machine over `self._pending`. + + final=True means the segment ends at a boundary: nothing may be held + back waiting for more text. + """ + while True: + if self._skip_rest: + self._pending = "" + return + if self._state == OUTSIDE: + if not self._pending: + return + text, self._pending = self._pending, "" + self._emit_outside(out, text) + return + if self._state == TAIL: + if self._pending.strip(): + self.errors.append( + f"unexpected text after payload: {self._pending[:20]!r}" + ) + self._pending = "" + return + if self._state == HEAD: + stripped = self._pending.lstrip() + self._pending = stripped + if stripped.startswith(CALL_PREFIX): + self._pending = stripped[len(CALL_PREFIX):] + self._state = NAME + continue + if not stripped: + return + if not final and _is_partial(stripped, (CALL_PREFIX,)): + return + self.errors.append( + f"expected 'call:' after tool open, got {stripped[:20]!r}" + ) + self._skip_rest = True + self._pending = "" + return + if self._state == NAME: + brace = self._pending.find("{") + if brace == -1: + if final or "}" in self._pending: + self.errors.append( + f"malformed tool call name: {self._pending[:20]!r}" + ) + self._skip_rest = True + self._pending = "" + return + name = self._pending[:brace].strip() + self._pending = self._pending[brace + 1:] + if not name: + self.errors.append("empty tool call name") + self._skip_rest = True + self._pending = "" + return + if self._cur is not None: + self._cur["function"]["name"] = name + self._frag({"index": self._cur_index, "function": {"name": name}}) + self._state = ARGS_KEY + continue + if self._state == ARGS_KEY: + i = 0 + n = len(self._pending) + while i < n: + ch = self._pending[i] + if ch == ":": + key = self._pending[:i].strip() + self._pending = self._pending[i + 1:] + if not key: + self.errors.append("empty argument key") + self._skip_rest = True + self._pending = "" + return + self._begin_param(key) + self._state = VALUE + self._depth = 0 + break + if ch == "}": + pre = self._pending[:i].strip() + self._pending = self._pending[i + 1:] + if pre: + self.errors.append( + f"trailing text before payload end: {pre[:20]!r}" + ) + self._state = TAIL + break + i += 1 + else: + if final and self._pending.strip(): + self.errors.append( + f"malformed argument list: {self._pending[:20]!r}" + ) + self._skip_rest = True + self._pending = "" + return + continue + if self._state == VALUE: + i = 0 + n = len(self._pending) + while i < n: + ch = self._pending[i] + if ch == "{": + self._depth += 1 + elif ch == "}": + if self._depth == 0: + self._consume_value_chars(self._pending[:i]) + self._pending = self._pending[i + 1:] + self._end_value() + self._state = TAIL + break + self._depth -= 1 + elif ch == "," and self._depth == 0: + self._consume_value_chars(self._pending[:i]) + self._pending = self._pending[i + 1:] + self._end_value() + self._state = ARGS_KEY + break + i += 1 + else: + # No boundary char in this segment; everything is value text. + self._consume_value_chars(self._pending) + self._pending = "" + return + continue + + # -- boundaries ---------------------------------------------------------------- + + def _apply_boundary(self, tag_id: int, msg: Optional[dict] = None) -> None: + if tag_id == TOOL_OPEN_ID: + if self._cur is not None: + self.errors.append("new tool call opened before previous closed") + self._close_open_value() + self._finish_call(msg if msg is not None else {}) + self._start_call() + else: + if self._cur is None: + self.errors.append("tool call close outside a tool block") + elif self._state == TAIL: + self._finish_call(msg) + else: + self.errors.append( + f"tool call closed mid-payload (state {self._state})" + ) + self._close_open_value() + self._finish_call(msg) + + # -- helpers ------------------------------------------------------------------- + + def _emit_outside(self, out: List[str], text: str) -> None: + if not text: + return + if self._calls_seen == 0: + out.append(text) # content before the first call + elif text.strip(): + # Real text after a completed call means the model is rambling + # past its answer -> stop generation. + self.status = StreamingStatus.TOOL_CALL_STOP + self.set_status(self.status) + # else: whitespace between/after calls -> drop + + def _start_call(self) -> None: + self._cur_index = self._calls_seen + self._calls_seen += 1 + self._cur = { + "id": f"call_{uuid.uuid4().hex[:24]}", + "type": "function", + "function": {"name": "", "arguments": ""}, + } + self._param_index = 0 + self._depth = 0 + self._state = HEAD + self._skip_rest = False + self._frag({ + "index": self._cur_index, + "id": self._cur["id"], + "type": "function", + "function": {"name": "", "arguments": ""}, + }) + + def _finish_call(self, msg: Optional[dict]) -> None: + call = self._cur + if call is None: + self._state = OUTSIDE + self._skip_rest = False + return + self._frag({ + "index": self._cur_index, + "function": {"arguments": "}" if self._param_index > 0 else "{}"}, + }) + self._cur = None + self._state = OUTSIDE + self._skip_rest = False + finished = { + "id": call["id"], + "type": "function", + "function": { + "name": call["function"]["name"], + "arguments": call["function"]["arguments"], + }, + } + self._calls.append(finished) + self._param_index = 0 + if msg is not None: + # String-keyed object: JsonContainer::concatenate throws on lists. + msg.setdefault("tool_calls", {})[str(self._cur_index)] = finished + if self._stop_after: + self.status = StreamingStatus.TOOL_CALL_STOP + self.set_status(self.status) + self._stopped = True + + def _frag(self, payload: Dict[str, Any]) -> None: + arguments = payload.get("function", {}).get("arguments") + if arguments is not None and self._cur is not None: + self._cur["function"]["arguments"] += arguments + if self.on_fragment is not None: + self.on_fragment(payload) + + def _begin_param(self, key: str) -> None: + self._param_index += 1 + sep = "{" if self._param_index == 1 else ", " + self._frag({ + "index": self._cur_index, + "function": {"arguments": f"{sep}{json.dumps(key)}: "}, + }) + self._frag({"index": self._cur_index, "function": {"arguments": '"'}}) + + def _consume_value_chars(self, text: str) -> None: + if text: + self._frag({ + "index": self._cur_index, + "function": {"arguments": _json_escape_fragment(text)}, + }) + + def _end_value(self) -> None: + self._frag({"index": self._cur_index, "function": {"arguments": '"'}}) + + def _close_open_value(self) -> None: + if self._state == VALUE: + self._end_value() + + +class Gemma4ToolCallStreamer(StreamerBase): + """Engine streamer for gemma4 tool requests. + + StreamerBase with its own incremental decode (cumulative decode + delta + slicing, like ChunkStreamer): protocol tag tokens are matched by ID and + fed through the channel splitter / tool parser as ID-only events, while + everything else is decoded incrementally and fed as token-less text + deltas. Special=True tags vanish from decode(), which is exactly why they + must be intercepted here. This deliberately avoids TextParserStreamer, + whose parser chain deadlocks when generation runs on a worker thread + under asyncio. + + Parsed OpenAI deltas are enqueued on text_queue (same contract as + ChunkStreamer) wrapped as {"chat_delta": [...]} so the route can + distinguish them from metrics/error dicts. + """ + + def __init__(self, tokenizer, gen_config): + super().__init__() + self._channels = Gemma4ChannelSplitter() + self.tool_parser = Gemma4ToolCallParser() + self.tool_parser.on_fragment = self._collect_fragment + self._fragments: List[Dict[str, Any]] = [] + # Raw tagged-output reconstruction: decoded text deltas plus the tag + # text of every intercepted protocol id. Lets non-streaming callers + # obtain the raw text parse_generation expects without relying on + # pipeline-level skip_special_tokens (not available on this wheel). + self._raw_parts: List[str] = [] + self.tokenizer = tokenizer + self._protocol_ids: Dict[int, int] = {} + self._protocol_tags: Dict[int, str] = {} + for tag, tag_id in ( + (TOOL_OPEN, TOOL_OPEN_ID), + (TOOL_CLOSE, TOOL_CLOSE_ID), + (CHANNEL_OPEN, CHANNEL_OPEN_ID), + (CHANNEL_CLOSE, CHANNEL_CLOSE_ID), + ): + ids = tokenizer.encode(tag).input_ids.data.tolist()[0] + if len(ids) == 1: + if ids[0] != tag_id: + logger.warning( + "gemma4 tag %r encodes to id %d, expected %d", + tag, ids[0], tag_id, + ) + self._protocol_ids[ids[0]] = tag_id + self._protocol_tags[tag_id] = tag + self.tokens_cache: List[int] = [] + self.last_print_len = 0 + self.text_queue: "asyncio.Queue" = asyncio.Queue() + self._cancelled = asyncio.Event() + try: + self._loop = asyncio.get_running_loop() + except RuntimeError: + self._loop = None # offline use: put_nowait directly + + def _collect_fragment(self, fragment: Dict[str, Any]) -> None: + self._fragments.append(fragment) + + def _enqueue(self, item) -> None: + # write()/end() run on the generation thread; asyncio.Queue is not + # thread-safe, so hop through call_soon_threadsafe when a loop exists. + if self._loop is not None: + self._loop.call_soon_threadsafe(self.text_queue.put_nowait, item) + else: + self.text_queue.put_nowait(item) + + def _process_delta(self, text: str, delta_tokens) -> None: + self._raw_parts.append(text) + reason_delta, text_delta, content_ids = self._channels.feed(text, delta_tokens) + content = self.tool_parser.parse({}, text_delta, content_ids) + deltas: List[Dict[str, Any]] = [] + if reason_delta: + deltas.append({"reasoning_content": reason_delta}) + if content: + deltas.append({"content": content}) + if self._fragments: + deltas.append({"tool_calls": self._fragments}) + self._fragments = [] + if deltas: + self._enqueue({"chat_delta": deltas}) + + def _decode_available(self) -> None: + text = self.tokenizer.decode(self.tokens_cache) + if len(text) > self.last_print_len: + delta = text[self.last_print_len:] + if chr(65533) in delta: + # partial UTF-8 at the boundary; wait for more tokens + return + self.last_print_len = len(text) + self._process_delta(delta, []) + + def _flush_text(self) -> None: + """Decode and process everything left in the cache, then reset it. + + Called before an intercepted protocol token: pre-boundary text must be + processed before the boundary itself. + """ + if not self.tokens_cache: + return + text = self.tokenizer.decode(self.tokens_cache) + if len(text) > self.last_print_len: + self._process_delta(text[self.last_print_len:], []) + self.tokens_cache = [] + self.last_print_len = 0 + + def write(self, token) -> StreamingStatus: + if self._cancelled.is_set(): + self._enqueue(None) + return StreamingStatus.CANCEL + ids = token if isinstance(token, list) else [token] + for tid in ids: + protocol_id = self._protocol_ids.get(int(tid)) + if protocol_id is not None: + self._flush_text() + self._raw_parts.append(self._protocol_tags[protocol_id]) + self._process_delta("", [protocol_id]) + else: + self.tokens_cache.append(int(tid)) + self._decode_available() + return self.tool_parser.status + + @property + def raw_text(self) -> str: + """Reconstructed raw tagged output (protocol tags as literal text).""" + return "".join(self._raw_parts) + + def end(self) -> None: + self._flush_text() + self.tool_parser.finalize() + if self.tool_parser.errors: + logger.warning("gemma4 tool parser errors: %s", self.tool_parser.errors) + self._enqueue(None) + + def cancel(self) -> None: + self._cancelled.set() + + def is_cancelled(self) -> bool: + return self._cancelled.is_set() + + +# -- raw-text path (non-streaming) ------------------------------------------------ + + +def _split_args(args_text: str) -> Dict[str, str]: + """Split 'key:value,key:{nested:{...}},...' at top level (brace-depth 0). + + Values stay raw strings (unquoted in the payload; quote tokens are + special=True, so no type information survives anyway). + """ + parts: List[str] = [] + depth = 0 + current = "" + for ch in args_text: + if ch == "{": + depth += 1 + elif ch == "}": + depth -= 1 + if ch == "," and depth == 0: + parts.append(current) + current = "" + else: + current += ch + if current: + parts.append(current) + arguments: Dict[str, str] = {} + for part in parts: + if ":" not in part: + continue + key, value = part.split(":", 1) + arguments[key.strip()] = value.strip() + return arguments + + +def _parse_payload(payload: str) -> Optional[Dict[str, Any]]: + """Parse a complete 'call:{...}' payload into an OpenAI tool call.""" + match = _PAYLOAD_RE.fullmatch(payload.strip()) + if not match: + return None + return { + "id": f"call_{uuid.uuid4().hex[:24]}", + "type": "function", + "function": { + "name": match.group(1), + "arguments": json.dumps(_split_args(match.group(2))), + }, + } + + +def _split_channels_text(text: str) -> Tuple[str, str]: + """Extract thought channels from raw tagged text. + + Returns (reasoning, remainder). Thought-channel bodies are rstripped of + the newline before the close tag; channels with other names pass through + as content without their wrapper tags; unterminated channels are kept + verbatim in the remainder. + """ + reasoning_parts: List[str] = [] + out: List[str] = [] + pos = 0 + while True: + i = text.find(CHANNEL_OPEN, pos) + if i == -1: + out.append(text[pos:]) + break + out.append(text[pos:i]) + j = text.find(CHANNEL_CLOSE, i + len(CHANNEL_OPEN)) + if j == -1: + out.append(text[i:]) + break + block = text[i + len(CHANNEL_OPEN):j] + header, sep, body = block.partition("\n") + if sep and header.strip() == THOUGHT_HEADER: + reasoning_parts.append(body.rstrip("\n")) + else: + out.append(block) + pos = j + len(CHANNEL_CLOSE) + return "".join(reasoning_parts), "".join(out) + + +def parse_generation( + text: str, + tools: Optional[List[Dict[str, Any]]] = None, # accepted; unused + enable_thinking: bool = True, # accepted for dispatch symmetry; unused +) -> tuple[str, str, Optional[List[Dict[str, Any]]]]: + """Split raw tagged model output into (reasoning, content, tool_calls). + + Expects text decoded with skip_special_tokens=False (the engine does this + for gemma4-parser models) so <|channel>/<|tool_call> tags survive as + literal text. + """ + if CHANNEL_OPEN not in text and TOOL_OPEN not in text: + return "", text, None + cleaned = text.replace(TURN_END, "").replace(EOS, "") + reasoning, rest = _split_channels_text(cleaned) + calls: List[Dict[str, Any]] = [] + content_parts: List[str] = [] + pos = 0 + while True: + i = rest.find(TOOL_OPEN, pos) + if i == -1: + content_parts.append(rest[pos:]) + break + content_parts.append(rest[pos:i]) + j = rest.find(TOOL_CLOSE, i + len(TOOL_OPEN)) + if j == -1: + # Unterminated block: keep the raw text as content. + content_parts.append(rest[i:]) + break + payload = rest[i + len(TOOL_OPEN):j] + call = _parse_payload(payload) + if call is not None: + calls.append(call) + else: + logger.debug("gemma4 malformed tool call payload: %r", payload[:80]) + pos = j + len(TOOL_CLOSE) + return reasoning, "".join(content_parts).strip(), calls or None + + +if __name__ == "__main__": + # Live smoke test for the engine path (Gemma4ToolCallStreamer) on the + # Chimera-X-26B VLM: reasoning channel + parallel tool calls in one stream. + # python -m src.engine.ov_genai.tool_parse.gemma4 [DEVICE] [MODEL_PATH] + import sys + + import openvino_genai as ov + + DEVICE = sys.argv[1] if len(sys.argv) > 1 and not sys.argv[1].startswith("-") else "GPU.0" + MODEL_PATH = ( + sys.argv[2] if len(sys.argv) > 2 else + "/mnt/Ironwolf-4TB/Models/OpenVINO/Gemma/Chimera-X-26B-A4B-int4-ov/" + ) + SMOKE_TOOLS = [{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string", "description": "City name"}}, + "required": ["city"], + }, + }, + }] + + pipe = ov.VLMPipeline(MODEL_PATH, DEVICE) + tokenizer = pipe.get_tokenizer() + + from types import SimpleNamespace + gen_config = SimpleNamespace( + tools=SMOKE_TOOLS, chat_template_kwargs={"enable_thinking": True} + ) + + streamer = Gemma4ToolCallStreamer(tokenizer, gen_config) + + history = ov.ChatHistory([ + {"role": "user", "content": "What's the weather in Tokyo and Paris? Use the tool for both cities at once."} + ]) + history.set_tools(SMOKE_TOOLS) + history.set_extra_context({"enable_thinking": True}) + + config = ov.GenerationConfig() + config.max_new_tokens = 1024 + + pipe.generate(history, generation_config=config, streamer=streamer) + + deltas: List[Dict[str, Any]] = [] + while not streamer.text_queue.empty(): + item = streamer.text_queue.get_nowait() + if item is None: + break + if isinstance(item, dict) and "chat_delta" in item: + deltas.extend(item["chat_delta"]) + else: + print("queue item:", item) + + print("=" * 60) + print("PARSED DELTAS") + print("=" * 60) + reasoning = "".join(d.get("reasoning_content", "") for d in deltas) + content = "".join(d.get("content", "") for d in deltas) + tool_frags = [f for d in deltas for f in d.get("tool_calls", [])] + print("reasoning:", repr(reasoning)) + print("content:", repr(content)) + print("fragments:") + for f in tool_frags: + print(" ", f) + args = "".join( + f["function"]["arguments"] + for f in tool_frags + if "arguments" in f.get("function", {}) + ) + names = [f["function"]["name"] for f in tool_frags if f.get("function", {}).get("name")] + print("names:", names) + print("arguments:", repr(args)) + try: + print("arguments JSON:", json.loads(args) if args else None) + except json.JSONDecodeError as exc: + print("arguments JSON ERROR:", exc) + print("parser status:", streamer.tool_parser.get_status()) + print("parser errors:", streamer.tool_parser.errors) + print("=" * 60) + per_call: Dict[str, str] = {} + for f in tool_frags: + fn = f.get("function", {}) + if "arguments" in fn: + per_call[str(f["index"])] = per_call.get(str(f["index"]), "") + fn["arguments"] + ok = True + for index, text in sorted(per_call.items()): + try: + print(f"call {index} args:", json.loads(text)) + except json.JSONDecodeError as exc: + ok = False + print(f"call {index} args INVALID: {text!r} ({exc})") + print("PARALLEL OK" if ok and len(per_call) >= 2 else "NOTE: fewer than 2 calls emitted") diff --git a/src/engine/ov_genai/tool_parse/hermes.py b/src/engine/ov_genai/tool_parse/hermes.py new file mode 100644 index 0000000..e2cec92 --- /dev/null +++ b/src/engine/ov_genai/tool_parse/hermes.py @@ -0,0 +1,229 @@ +"""Hermes JSON tool-call parser: + + + {"name": NAME, "arguments": {...}} + + +Kept fully separate from the Qwen3.5 XML parser (qwen35.py). Only the +ReasoningSplitter and the partial-tag hold-back helper are shared (reasoning +splitting is not tool parsing). +""" +import json +import uuid +from typing import Any, Dict, List, Optional + +from src.engine.ov_genai.tool_parse.qwen35 import ReasoningSplitter, _holdback_suffix + +TOOL_OPEN = "" +TOOL_CLOSE = "" + + +def _extract_hermes_tool_call_payloads(text: str) -> List[str]: + payloads: List[str] = [] + cursor = 0 + + while True: + start = text.find(TOOL_OPEN, cursor) + if start < 0: + break + + payload_start = start + len(TOOL_OPEN) + end = text.find(TOOL_CLOSE, payload_start) + if end < 0: + payload = text[payload_start:].strip() + if payload: + payloads.append(payload) + break + + payload = text[payload_start:end].strip() + if payload: + payloads.append(payload) + + cursor = end + len(TOOL_CLOSE) + + return payloads + + +def _format_tool_call_arguments(arguments: Any) -> str: + if isinstance(arguments, str): + try: + return json.dumps(json.loads(arguments)) + except json.JSONDecodeError: + return arguments + return json.dumps(arguments) + + +def _payload_to_tool_call(payload: str) -> Optional[Dict[str, Any]]: + try: + data = json.loads(payload) + except json.JSONDecodeError: + return None + if not (isinstance(data, dict) and "name" in data and "arguments" in data): + return None + return { + "id": f"call_{uuid.uuid4().hex[:24]}", + "type": "function", + "function": { + "name": str(data.get("name", "")), + "arguments": _format_tool_call_arguments(data.get("arguments", {})), + }, + } + + +def parse_hermes_tool_calls(text: str) -> Optional[List[Dict[str, Any]]]: + tool_calls: List[Dict[str, Any]] = [ + tc + for payload in _extract_hermes_tool_call_payloads(text) + if (tc := _payload_to_tool_call(payload)) is not None + ] + return tool_calls if tool_calls else None + + +def parse_generation( + text: str, + tools: Optional[List[Dict[str, Any]]] = None, + enable_thinking: bool = True, +) -> tuple[str, str, Optional[List[Dict[str, Any]]]]: + """Split model output into (reasoning, content, tool_calls). + + `tools` is accepted for dispatch symmetry with qwen35; hermes payloads are + self-describing so it is unused. + """ + reasoning, remainder = ReasoningSplitter(enabled="" in text).feed(text) + if remainder.startswith(""): + remainder = remainder[len("") :] + + hermes = parse_hermes_tool_calls(remainder) + if hermes: + content = remainder + for payload in _extract_hermes_tool_call_payloads(remainder): + content = content.replace(f"{TOOL_OPEN}{payload}{TOOL_CLOSE}", "") + content = content.replace(f"{TOOL_OPEN}\n{payload}\n{TOOL_CLOSE}", "") + return reasoning, content.strip(), hermes + + return reasoning, remainder, None + + +THINK_OPEN = "" + + +class HermesStreamParser: + """Incremental hermes parser for streaming. + + Plain content streams live (with a hold-back on a partial + prefix); JSON inside a call is buffered until (or stream end + for unterminated payloads), then emitted as the two-fragment sequence: + call-start-with-name, then arguments payload. + + Thinking is decided lazily from the stream start: hermes models that think + emit a full ... block, unlike qwen35 (whose chat template + pre-opens the block). When enable_thinking is set, the first characters + decide: a leading enables reasoning splitting, anything else is + plain content from the start. + """ + + def __init__( + self, + tools: Optional[List[Dict[str, Any]]] = None, + enable_thinking: bool = True, + ): + self._undecided = enable_thinking + self._pending = "" + self._reasoning = ReasoningSplitter(enabled=False) + self._buf = "" + self._in_tool = False + self._index = 0 + + def feed(self, text: str) -> List[Dict[str, Any]]: + out: List[Dict[str, Any]] = [] + if self._undecided: + self._pending += text + stripped = self._pending.lstrip() + if len(stripped) < len(THINK_OPEN) and THINK_OPEN.startswith(stripped): + return [] # still could grow into a tag + self._undecided = False + if stripped.startswith(THINK_OPEN): + self._reasoning.in_reasoning = True + text = stripped[len(THINK_OPEN):] + else: + text = self._pending + self._pending = "" + if not text: + return out + reasoning, content = self._reasoning.feed(text) + if reasoning: + out.append({"reasoning_content": reasoning}) + if content: + self._buf += content + out.extend(self._drain(final=False)) + return out + + def finish(self) -> List[Dict[str, Any]]: + out: List[Dict[str, Any]] = [] + if self._undecided: + self._undecided = False + text = self._pending + self._pending = "" + if text: + self._buf += text + return out + self._drain(final=True) + + def _drain(self, final: bool) -> List[Dict[str, Any]]: + out: List[Dict[str, Any]] = [] + while self._buf: + if not self._in_tool: + idx = self._buf.find(TOOL_OPEN) + if idx == -1: + hold = 0 if final else _holdback_suffix(self._buf, (TOOL_OPEN,)) + safe, self._buf = self._buf[: len(self._buf) - hold], self._buf[len(self._buf) - hold:] + if safe: + out.append({"content": safe}) + break + if idx > 0: + out.append({"content": self._buf[:idx]}) + self._buf = self._buf[idx + len(TOOL_OPEN):] + self._in_tool = True + else: + end = self._buf.find(TOOL_CLOSE) + if end == -1: + if final: + self._emit_tool_calls(self._buf, out) + self._buf = "" + break + payload = self._buf[:end] + self._buf = self._buf[end + len(TOOL_CLOSE):] + self._in_tool = False + self._emit_tool_calls(payload, out) + return out + + def _emit_tool_calls(self, payload: str, out: List[Dict[str, Any]]) -> None: + tc = _payload_to_tool_call(payload.strip()) + if tc is None: + return + idx = self._index + self._index += 1 + out.append( + { + "tool_calls": [ + { + "index": idx, + "id": tc["id"], + "type": tc["type"], + "function": { + "name": tc["function"]["name"], + "arguments": "", + }, + } + ] + } + ) + out.append( + { + "tool_calls": [ + { + "index": idx, + "function": {"arguments": tc["function"]["arguments"]}, + } + ] + } + ) diff --git a/src/engine/ov_genai/tool_parse/qwen35.py b/src/engine/ov_genai/tool_parse/qwen35.py new file mode 100644 index 0000000..1df7523 --- /dev/null +++ b/src/engine/ov_genai/tool_parse/qwen35.py @@ -0,0 +1,834 @@ +"""Qwen3.5 XML tool-call parser built on openvino_genai IncrementalParser. + +Tool-call block boundaries are detected by TOKEN ID, not text matching: + - a block opens only in a delta whose delta_tokens contains TOOL_OPEN_ID + - a block closes only in a delta whose delta_tokens contains TOOL_CLOSE_ID + - the boundary tags are single special=False tokens in the Qwen3.5 vocab, so + their text always arrives in the same delta as their ID (even with + skip_special_tokens=True, which only strips true specials like <|im_end|>) + - the tag text is only used to slice the captured payload, never to decide + that a boundary exists + +Two consumers: + Qwen35ToolCallStreamer - StreamerBase subclass used by the engine for + streaming tool requests; does its own incremental + decode, matches boundary tags by token ID, and + enqueues {"chat_delta": [...]} dicts. Deliberately + avoids TextParserStreamer, which deadlocks when + generate runs on a worker thread under asyncio. + Qwen35StreamParser - text-delta facade over ChunkStreamer output; feeds + ReasoningSplitter + QwenXMLToolParser, synthesizing + boundary token IDs from the tag text. + +Parsed OpenAI streaming fragments: + {"index": i, "id": ..., "type": "function", + "function": {"name": "", "arguments": ""}} # call start + {"index": i, "function": {"name": NAME}} # name known + {"index": i, "function": {"arguments": FRAGMENT}} # argument deltas + +Concatenating all `arguments` fragments for an index yields complete JSON. +""" +import asyncio +import json +import logging +import uuid +from typing import Any, Callable, Dict, List, Optional, Tuple + +from openvino_genai import IncrementalParser, StreamerBase, StreamingStatus + +logger = logging.getLogger(__name__) + +# Qwen3.5 XML format. Boundary tags are single special=False tokens; the inner +# tags (" that ends the function name +BODY = "body" # inside : params or +PARAM_NAME = "param_name" # awaiting ">" that ends the parameter name +PARAM = "param" # inside a parameter value +TAIL = "tail" # after , awaiting the close boundary + + +def _holdback_suffix(text: str, tags) -> int: + """Length of the longest suffix of `text` that is a proper prefix of a tag. + + Anything before that suffix is safe to process now; the suffix must wait + for more deltas (it could grow into a full tag). + """ + max_hold = 0 + upper = max(len(t) for t in tags) - 1 + for n in range(1, min(len(text), upper) + 1): + suffix = text[-n:] + if any(t != suffix and t.startswith(suffix) for t in tags): + max_hold = n + return max_hold + + +def _is_partial(text: str, tags) -> bool: + """True if `text` is a proper prefix of one of `tags`.""" + return any(t != text and t.startswith(text) for t in tags) + + +def _json_escape_fragment(s: str) -> str: + """Escape a raw string fragment for embedding inside a JSON string.""" + out = [] + for ch in s: + if ch == '"': + out.append('\\"') + elif ch == "\\": + out.append("\\\\") + elif ch == "\n": + out.append("\\n") + elif ch == "\t": + out.append("\\t") + elif ch == "\r": + out.append("\\r") + elif ord(ch) < 0x20: + out.append(f"\\u{ord(ch):04x}") + else: + out.append(ch) + return "".join(out) + + +class ReasoningSplitter: + """Splits the leading thinking block from regular content. + + Qwen3.5 with enable_thinking=True emits reasoning immediately (no opening + think tag) and terminates it with the close tag. Everything before the + close tag is reasoning; everything after is content. + """ + + def __init__(self, close_tag: str = THINK_CLOSE, enabled: bool = True): + self._close_tag = close_tag + self._buf = "" + # Only treat the stream as reasoning when the prompt actually opened a + # think block (chat template appends it when enable_thinking=True; with + # False it pre-closes an empty block and the stream is pure content). + self.in_reasoning = enabled + + def feed(self, text: str) -> Tuple[str, str]: + """Returns (reasoning_delta, content_delta).""" + self._buf += text + reasoning_out, content_out = "", "" + while self._buf: + if not self.in_reasoning: + content_out += self._buf + self._buf = "" + break + idx = self._buf.find(self._close_tag) + if idx != -1: + reasoning_out += self._buf[:idx] + self._buf = self._buf[idx + len(self._close_tag):] + self.in_reasoning = False + continue + hold = _holdback_suffix(self._buf, (self._close_tag,)) + reasoning_out += self._buf[: len(self._buf) - hold] + self._buf = self._buf[len(self._buf) - hold:] + break + return reasoning_out, content_out + + +class _TextTagSynthesizer: + """Stateful text adapter upholding the parser's boundary invariant. + + Buffers text and emits (chunk, delta_tokens) pairs such that every + complete boundary tag arrives as its own chunk carrying its token ID, and + no chunk ends with a partial boundary tag. Used by text-only feeders + (ChunkStreamer path, parse_generation) where no real token IDs exist. + """ + + def __init__(self): + self._buf = "" + + def feed(self, text: str) -> List[Tuple[str, Optional[List[int]]]]: + self._buf += text + out: List[Tuple[str, Optional[List[int]]]] = [] + while self._buf: + events = [] + for tag, tag_id in ((TOOL_OPEN, TOOL_OPEN_ID), (TOOL_CLOSE, TOOL_CLOSE_ID)): + pos = self._buf.find(tag) + if pos != -1: + events.append((pos, tag, tag_id)) + if events: + pos, tag, tag_id = min(events) + if pos > 0: + out.append((self._buf[:pos], None)) + self._buf = self._buf[pos:] + out.append((tag, [tag_id])) + self._buf = self._buf[len(tag):] + continue + hold = _holdback_suffix(self._buf, (TOOL_OPEN, TOOL_CLOSE)) + keep = len(self._buf) - hold + if keep > 0: + out.append((self._buf[:keep], None)) + self._buf = self._buf[keep:] + break + return out + + def flush(self) -> str: + """Release the held tail at end of stream (it never became a tag).""" + chunk, self._buf = self._buf, "" + return chunk + + +class QwenXMLToolParser(IncrementalParser): + """IncrementalParser for Qwen3.5 XML tool calls, bound by token IDs. + + Boundary detection is token-ID only: + - A block opens only in a delta whose delta_tokens contains TOOL_OPEN_ID. + - A block closes only in a delta whose delta_tokens contains TOOL_CLOSE_ID. + - Boundary tag text is sliced out of the accumulated text; text without + token info drains as ordinary content. + + Parameter values are emitted as JSON strings (no schema-driven type + coercion). `on_fragment` receives OpenAI streaming fragments as they are + produced. With + stop_after_tool_call=True the parser requests StreamingStatus.TOOL_CALL_STOP + right after each complete call; by default (False) sequential parallel + calls are allowed and generation is only stopped when real text follows a + completed call. + """ + + def __init__( + self, + on_fragment: Optional[Callable[[Dict[str, Any]], None]] = None, + stop_after_tool_call: bool = False, + ): + super().__init__() + self.on_fragment = on_fragment + self._stop_after = stop_after_tool_call + self.reset() + + def reset(self) -> None: + self._state = OUTSIDE + self._pending = "" + self._cur: Optional[Dict[str, Any]] = None # call under construction + self._cur_index = -1 + self._calls_seen = 0 # calls started (open boundary seen) + self._calls: List[Dict[str, Any]] = [] # completed calls, OpenAI shape + self._param_name = "" + self._param_raw = "" + self._param_index = 0 + self._value_started = False + self._pending_newline = False + self._skip_rest = False + self._stopped = False + self.errors: List[str] = [] + self.status = StreamingStatus.RUNNING + self.set_status(StreamingStatus.RUNNING) + + @property + def completed_calls(self) -> List[Dict[str, Any]]: + """All completed calls, in final OpenAI non-streaming shape.""" + return list(self._calls) + + def finalize(self) -> None: + """Call at end of generation: force-close an unterminated call and + surface malformed structures in `errors`.""" + if self._cur is not None: + self.errors.append("generation ended with unterminated tool call") + self._finish_call({}) + elif self._state not in (OUTSIDE,): + self.errors.append(f"generation ended in state {self._state}") + + # -- IncrementalParser API --------------------------------------------------- + + def parse(self, msg: dict, delta_text: str, delta_tokens=None) -> str: + """Consume one decoded delta. Returns content text for this delta. + + Mutates `msg` with completed calls under msg["tool_calls"][str(index)] + and emits streaming fragments through `on_fragment`. + """ + if self._stopped: + return "" + self._pending += delta_text + remaining: Dict[int, int] = {} + if delta_tokens: + for token in delta_tokens: + tid = int(token) + if tid in BOUNDARY_IDS: + remaining[tid] = remaining.get(tid, 0) + 1 + + out: List[str] = [] + while True: + # Earliest boundary event: ID still available AND tag text in pending + events = [] + for tag_id, tag in ((TOOL_OPEN_ID, TOOL_OPEN), (TOOL_CLOSE_ID, TOOL_CLOSE)): + if remaining.get(tag_id): + pos = self._pending.find(tag) + if pos != -1: + events.append((pos, tag_id, tag)) + if not events: + self._machine(out, final=False) + break + pos, tag_id, tag = min(events) + segment = self._pending[:pos] + remainder = self._pending[pos + len(tag):] + # The machine must see exactly the text before the boundary; + # final=True guarantees it drains the segment completely. + self._pending = segment + if segment: + self._machine(out, final=True) + self._pending = remainder + remaining[tag_id] -= 1 + self._apply_boundary(tag_id, msg) + if self._stop_after and self.status == StreamingStatus.TOOL_CALL_STOP: + self._pending = "" + break + return "".join(out) + + # -- state machine ------------------------------------------------------------ + + def _machine(self, out: List[str], final: bool) -> None: + """Advance the machine over `self._pending`. + + final=True means the segment ends at a boundary: nothing may be held + back waiting for more text. + """ + while True: + if self._skip_rest: + self._pending = "" + return + if self._state == OUTSIDE: + if not self._pending: + return + text, self._pending = self._pending, "" + self._emit_outside(out, text) + return + if self._state == TAIL: + if self._pending.strip(): + self.errors.append( + f"unexpected text after : {self._pending[:20]!r}" + ) + self._pending = "" + return + if self._state == HEAD: + stripped = self._pending.lstrip() + self._pending = stripped + if stripped.startswith(FUNC_OPEN): + self._pending = stripped[len(FUNC_OPEN):] + self._state = FUNC_NAME + continue + if not stripped: + return + if not final and _is_partial(stripped, (FUNC_OPEN,)): + return + self.errors.append( + f"expected ': {stripped[:20]!r}" + ) + self._pending = stripped[1:] + continue + if self._state == PARAM_NAME: + gt = self._pending.find(GT) + if gt == -1: + if "<" in self._pending or final: + self.errors.append( + f"malformed parameter name: {self._pending[:20]!r}" + ) + self._skip_rest = True + self._pending = "" + return + self._param_name = self._pending[:gt].strip() + self._pending = self._pending[gt + 1:] + self._param_raw = "" + self._value_started = False + self._pending_newline = False + self._state = PARAM + sep = "{" if self._param_index == 0 else ", " + self._param_index += 1 + self._frag({ + "index": self._cur_index, + "function": {"arguments": f"{sep}{json.dumps(self._param_name)}: "}, + }) + self._frag({"index": self._cur_index, "function": {"arguments": '"'}}) + continue + if self._state == PARAM: + idx = self._pending.find(PARAM_CLOSE) + if idx == -1: + if final: + value, self._pending = self._pending, "" + self._consume_value(value, final=True) + self._state = BODY + continue + hold = _holdback_suffix(self._pending, (PARAM_CLOSE,)) + safe = self._pending[: len(self._pending) - hold] + self._pending = self._pending[len(self._pending) - hold:] + if safe: + self._consume_value(safe, final=False) + return + value, self._pending = ( + self._pending[:idx], + self._pending[idx + len(PARAM_CLOSE):], + ) + self._consume_value(value, final=True) + self._state = BODY + continue + return + + def _apply_boundary(self, tag_id: int, msg: Optional[dict] = None) -> None: + if tag_id == TOOL_OPEN_ID: + if self._cur is not None: + self.errors.append("new tool call opened before previous closed") + self._finish_call(msg if msg is not None else {}) + self._start_call() + else: + if self._state == TAIL: + self._finish_call(msg) + elif self._cur is None: + self.errors.append("tool call close outside a tool block") + else: + self.errors.append("tool call closed before ") + self._finish_call(msg) + + # -- helpers ------------------------------------------------------------------- + + def _emit_outside(self, out: List[str], text: str) -> None: + if not text: + return + if self._calls_seen == 0: + out.append(text) # content before the first call + elif text.strip(): + # Real text after a completed call means the model is rambling + # past its answer -> stop generation. + self.status = StreamingStatus.TOOL_CALL_STOP + self.set_status(self.status) + # else: whitespace between/after calls -> drop + + def _start_call(self) -> None: + self._cur_index = self._calls_seen + self._calls_seen += 1 + self._cur = { + "id": f"call_{uuid.uuid4().hex[:24]}", + "type": "function", + "function": {"name": "", "arguments": ""}, + } + self._param_index = 0 + self._state = HEAD + self._skip_rest = False + self._frag({ + "index": self._cur_index, + "id": self._cur["id"], + "type": "function", + "function": {"name": "", "arguments": ""}, + }) + + def _finish_call(self, msg: Optional[dict]) -> None: + call = self._cur + if call is None: + self._state = OUTSIDE + self._skip_rest = False + return + self._frag({ + "index": self._cur_index, + "function": {"arguments": "}" if self._param_index > 0 else "{}"}, + }) + self._cur = None + self._state = OUTSIDE + self._skip_rest = False + finished = { + "id": call["id"], + "type": "function", + "function": { + "name": call["function"]["name"], + "arguments": call["function"]["arguments"], + }, + } + self._calls.append(finished) + self._param_index = 0 + if msg is not None: + # String-keyed object: JsonContainer::concatenate throws on lists. + msg.setdefault("tool_calls", {})[str(self._cur_index)] = finished + if self._stop_after: + self.status = StreamingStatus.TOOL_CALL_STOP + self.set_status(self.status) + self._stopped = True + + def _frag(self, payload: Dict[str, Any]) -> None: + arguments = payload.get("function", {}).get("arguments") + if arguments is not None and self._cur is not None: + self._cur["function"]["arguments"] += arguments + if self.on_fragment is not None: + self.on_fragment(payload) + + def _consume_value(self, text: str, final: bool) -> None: + # drop the single wrapper newline right after + if not self._value_started and text.startswith("\n"): + text = text[1:] + if text: + self._value_started = True + # hold back a trailing newline: it may be the wrapper before + # ; if more value text follows we flush it then. + if self._pending_newline: + text = "\n" + text + self._pending_newline = False + if text.endswith("\n"): + if final: + text = text[:-1] + else: + self._pending_newline = True + text = text[:-1] + if text: + self._param_raw += text + self._frag({ + "index": self._cur_index, + "function": {"arguments": _json_escape_fragment(text)}, + }) + if final: + self._frag({"index": self._cur_index, "function": {"arguments": '"'}}) + + +class Qwen35ToolCallStreamer(StreamerBase): + """Engine streamer for qwen35 tool requests. + + StreamerBase with its own incremental decode (cumulative decode + delta + slicing, like ChunkStreamer): boundary tag tokens are matched by ID and + fed through QwenXMLToolParser with (tag_text, [tag_id]) pairs; everything + else is decoded incrementally and fed as token-less text deltas. This + deliberately avoids TextParserStreamer, whose parser chain deadlocks when + generation runs on a worker thread under asyncio. + + Parsed OpenAI deltas are enqueued on text_queue (same contract as + ChunkStreamer) wrapped as {"chat_delta": [...]} so the route can + distinguish them from metrics/error dicts. + """ + + def __init__(self, tokenizer, gen_config): + super().__init__() + # enable_thinking=True -> the chat template injects the opening think + # tag into the prompt, so the stream starts inside a think block. + # With enable_thinking=False the keep-splitter is disabled entirely. + thinking = True + if getattr(gen_config, "chat_template_kwargs", None): + thinking = bool(gen_config.chat_template_kwargs.get("enable_thinking", True)) + self._reasoning = ReasoningSplitter(enabled=thinking) + self.tool_parser = QwenXMLToolParser() + self.tool_parser.on_fragment = self._collect_fragment + self._fragments: List[Dict[str, Any]] = [] + self.tokenizer = tokenizer + self._wrapper_ids: Dict[int, str] = {} + for tag, tag_id in ((TOOL_OPEN, TOOL_OPEN_ID), (TOOL_CLOSE, TOOL_CLOSE_ID)): + ids = tokenizer.encode(tag).input_ids.data.tolist()[0] + if len(ids) == 1: + self._wrapper_ids[ids[0]] = tag + self.tokens_cache: List[int] = [] + self.last_print_len = 0 + self.text_queue: "asyncio.Queue" = asyncio.Queue() + self._cancelled = asyncio.Event() + try: + self._loop = asyncio.get_running_loop() + except RuntimeError: + self._loop = None # offline use: put_nowait directly + + def _collect_fragment(self, fragment: Dict[str, Any]) -> None: + self._fragments.append(fragment) + + def _enqueue(self, item) -> None: + # write()/end() run on the generation thread; asyncio.Queue is not + # thread-safe, so hop through call_soon_threadsafe when a loop exists. + if self._loop is not None: + self._loop.call_soon_threadsafe(self.text_queue.put_nowait, item) + else: + self.text_queue.put_nowait(item) + + def _process_delta(self, text: str, delta_tokens) -> None: + reason_delta, text_delta = self._reasoning.feed(text) + content = self.tool_parser.parse({}, text_delta, delta_tokens) + deltas: List[Dict[str, Any]] = [] + if reason_delta: + deltas.append({"reasoning_content": reason_delta}) + if content: + deltas.append({"content": content}) + if self._fragments: + deltas.append({"tool_calls": self._fragments}) + self._fragments = [] + if deltas: + self._enqueue({"chat_delta": deltas}) + + def _decode_available(self) -> None: + text = self.tokenizer.decode(self.tokens_cache) + if len(text) > self.last_print_len: + delta = text[self.last_print_len:] + if chr(65533) in delta: + # partial UTF-8 at the boundary; wait for more tokens + return + self.last_print_len = len(text) + self._process_delta(delta, []) + + def _flush_text(self) -> None: + """Decode and process everything left in the cache, then reset it. + + Called before an atomically-detected boundary token: pre-boundary text + must be processed before the boundary itself. + """ + if not self.tokens_cache: + return + text = self.tokenizer.decode(self.tokens_cache) + if len(text) > self.last_print_len: + self._process_delta(text[self.last_print_len:], []) + self.tokens_cache = [] + self.last_print_len = 0 + + def write(self, token) -> StreamingStatus: + if self._cancelled.is_set(): + self._enqueue(None) + return StreamingStatus.CANCEL + ids = token if isinstance(token, list) else [token] + for tid in ids: + tag = self._wrapper_ids.get(int(tid)) + if tag is not None: + self._flush_text() + self._process_delta(tag, [int(tid)]) + else: + self.tokens_cache.append(int(tid)) + self._decode_available() + return self.tool_parser.status + + def end(self) -> None: + self._flush_text() + self.tool_parser.finalize() + if self.tool_parser.errors: + logger.warning("qwen35 tool parser errors: %s", self.tool_parser.errors) + self._enqueue(None) + + def cancel(self) -> None: + self._cancelled.set() + + def is_cancelled(self) -> bool: + return self._cancelled.is_set() + + +class Qwen35StreamParser: + """Streaming facade over text deltas (ChunkStreamer path). + + feed() consumes detokenized text (no token IDs available, so boundary IDs + are synthesized from the tag text) and returns OpenAI delta dicts in order + [{reasoning_content}, {content}, {tool_calls}], skipping empty entries. + finish() finalizes the parser (force-closes truncated calls) and returns []. + """ + + def __init__( + self, + tools: Optional[List[Dict[str, Any]]] = None, # accepted; unused + enable_thinking: bool = True, + ): + self._reasoning = ReasoningSplitter(enabled=enable_thinking) + self._tool_parser = QwenXMLToolParser() + self._synth = _TextTagSynthesizer() + + def feed(self, text: str) -> List[Dict[str, Any]]: + out: List[Dict[str, Any]] = [] + reason_delta, text_delta = self._reasoning.feed(text) + if reason_delta: + out.append({"reasoning_content": reason_delta}) + content_parts: List[str] = [] + fragments: List[Dict[str, Any]] = [] + for chunk, delta_tokens in self._synth.feed(text_delta): + c, f = self._parse_chunk(chunk, delta_tokens) + content_parts.append(c) + fragments.extend(f) + content = "".join(content_parts) + if content: + out.append({"content": content}) + if fragments: + out.append({"tool_calls": fragments}) + return out + + def finish(self) -> List[Dict[str, Any]]: + tail = self._synth.flush() + if tail: + self._parse_chunk(tail, None) + self._tool_parser.finalize() + if self._tool_parser.errors: + logger.warning("qwen35 tool parser errors: %s", self._tool_parser.errors) + return [] + + def _parse_chunk(self, chunk: str, delta_tokens): + fragments: List[Dict[str, Any]] = [] + + def sink(fragment: Dict[str, Any]) -> None: + fragments.append(fragment) + + self._tool_parser.on_fragment = sink + try: + content = self._tool_parser.parse({}, chunk, delta_tokens) + finally: + self._tool_parser.on_fragment = None + return content, fragments + + +def parse_generation( + text: str, + tools: Optional[List[Dict[str, Any]]] = None, # accepted; unused + enable_thinking: bool = True, +) -> tuple[str, str, Optional[List[Dict[str, Any]]]]: + """Split model output into (reasoning, content, tool_calls).""" + reasoning, remainder = ReasoningSplitter(enabled=THINK_CLOSE in text).feed(text) + if remainder.startswith(THINK_OPEN): + remainder = remainder[len(THINK_OPEN):] + if TOOL_OPEN not in remainder: + return reasoning, remainder, None + + parser = QwenXMLToolParser() + synth = _TextTagSynthesizer() + content_parts: List[str] = [] + for chunk, delta_tokens in synth.feed(remainder): + content_parts.append(parser.parse({}, chunk, delta_tokens)) + tail = synth.flush() + if tail: + content_parts.append(parser.parse({}, tail, None)) + parser.finalize() + if parser.errors: + logger.debug("qwen35 tool parser errors: %s", parser.errors) + return reasoning, "".join(content_parts), parser.completed_calls or None + + +if __name__ == "__main__": + # Live smoke test for the engine path (Qwen35ToolCallStreamer), mirroring + # the scratchpad demo: reasoning + tool call in one stream. + # python -m src.engine.ov_genai.tool_parse.qwen35 [DEVICE] [MODEL_PATH] + import sys + + import openvino_genai as ov + + DEVICE = sys.argv[1] if len(sys.argv) > 1 and not sys.argv[1].startswith("-") else "GPU.0" + MODEL_PATH = ( + sys.argv[2] if len(sys.argv) > 2 else + "/mnt/Ironwolf-4TB/Models/OpenVINO/Qwen3.5-2B-int4_sym-ov/" + ) + SMOKE_TOOLS = [{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string", "description": "City name"}}, + "required": ["city"], + }, + }, + }] + + pipe = ov.VLMPipeline(MODEL_PATH, DEVICE) + tokenizer = pipe.get_tokenizer() + + from types import SimpleNamespace + gen_config = SimpleNamespace( + tools=SMOKE_TOOLS, chat_template_kwargs={"enable_thinking": True} + ) + + streamer = Qwen35ToolCallStreamer(tokenizer, gen_config) + + history = ov.ChatHistory([{"role": "user", "content": "What's the weather in Tokyo and Paris? Use the tool for both cities at once."}]) + history.set_tools(SMOKE_TOOLS) + history.set_extra_context({"enable_thinking": True}) + + config = ov.GenerationConfig() + config.max_new_tokens = 512 + + pipe.generate(history, generation_config=config, streamer=streamer) + + deltas: List[Dict[str, Any]] = [] + while not streamer.text_queue.empty(): + item = streamer.text_queue.get_nowait() + if item is None: + break + if isinstance(item, dict) and "chat_delta" in item: + deltas.extend(item["chat_delta"]) + else: + print("queue item:", item) + + print("=" * 60) + print("PARSED DELTAS") + print("=" * 60) + reasoning = "".join(d.get("reasoning_content", "") for d in deltas) + content = "".join(d.get("content", "") for d in deltas) + tool_frags = [f for d in deltas for f in d.get("tool_calls", [])] + print("reasoning:", repr(reasoning)) + print("content:", repr(content)) + print("fragments:") + for f in tool_frags: + print(" ", f) + args = "".join( + f["function"]["arguments"] + for f in tool_frags + if "arguments" in f.get("function", {}) + ) + names = [f["function"]["name"] for f in tool_frags if f.get("function", {}).get("name")] + print("names:", names) + print("arguments:", repr(args)) + try: + print("arguments JSON:", json.loads(args) if args else None) + except json.JSONDecodeError as exc: + print("arguments JSON ERROR:", exc) + print("parser status:", streamer.tool_parser.get_status()) + print("parser errors:", streamer.tool_parser.errors) + print("=" * 60) + per_call: Dict[str, str] = {} + for f in tool_frags: + fn = f.get("function", {}) + if "arguments" in fn: + per_call[str(f["index"])] = per_call.get(str(f["index"]), "") + fn["arguments"] + ok = True + for index, text in sorted(per_call.items()): + try: + print(f"call {index} args:", json.loads(text)) + except json.JSONDecodeError as exc: + ok = False + print(f"call {index} args INVALID: {text!r} ({exc})") + print("PARALLEL OK" if ok and len(per_call) >= 2 else "NOTE: fewer than 2 calls emitted") + diff --git a/src/engine/ov_genai/vlm.py b/src/engine/ov_genai/vlm.py index e337abb..614f1dd 100644 --- a/src/engine/ov_genai/vlm.py +++ b/src/engine/ov_genai/vlm.py @@ -22,7 +22,8 @@ from src.server.utils.resolve_vlm_type import is_qwen3_5_architecture, resolve_vlm_vision_token from src.server.model_registry import ModelRegistry from src.server.schemas.registration import ModelLoadConfig -from src.engine.ov_genai.streamers import ChunkStreamer +from src.engine.ov_genai.streamers import ensure_tool_call_parser, select_streamer +from src.engine.ov_genai.tool_parse.gemma4 import Gemma4ToolCallStreamer logger = logging.getLogger(__name__) @@ -160,20 +161,35 @@ async def generate_text(self, gen_config: OVGenAI_GenConfig) -> AsyncIterator[Un Yields in order: metrics (dict), new_text (str). """ try: + ensure_tool_call_parser(gen_config, self.load_config) generation_kwargs = self.create_generation_config(gen_config) prompt, ov_images = self._resolve_prompt_and_images(gen_config) + # gemma4 non-streaming: generate through the token-ID streamer and + # reconstruct the raw tagged output. Gemma 4 protocol tags are + # special=True and the VLM decode always strips them, so the plain + # result text cannot be parsed for reasoning/tool calls. + gemma4_stream = getattr(gen_config, "tool_call_parser", None) == "gemma4" + streamer = ( + Gemma4ToolCallStreamer(self.model_path.get_tokenizer(), gen_config) + if gemma4_stream else None + ) + result = 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} if streamer is not None else {}), ) perf_metrics = result.perf_metrics - text = result.texts[0] if getattr(result, "texts", None) else "" + if gemma4_stream: + text = streamer.raw_text + else: + text = result.texts[0] if getattr(result, "texts", None) else "" logger.info(f"[{self.load_config.model_name}] Generation completed, generated {len(text)} characters") metrics_dict = self.collect_metrics(gen_config, perf_metrics) @@ -189,10 +205,11 @@ async def generate_stream(self, Async streaming generation for VLM. Yields token chunks (str) as they arrive, then metrics (dict). """ + ensure_tool_call_parser(gen_config, self.load_config) generation_kwargs = self.create_generation_config(gen_config) decoder_tokenizer = self.model_path.get_tokenizer() - streamer = ChunkStreamer(decoder_tokenizer, gen_config) + streamer = select_streamer(decoder_tokenizer, gen_config) # Track active request and streamer for cancellation self._active_request_id = gen_config.request_id diff --git a/src/server/model_registry.py b/src/server/model_registry.py index 7dbbfcf..8a20885 100644 --- a/src/server/model_registry.py +++ b/src/server/model_registry.py @@ -35,6 +35,7 @@ class ModelRecord: engine: str = "" device: str = "" runtime_config: Dict[str, Any] = field(default_factory=dict) + tool_call_parser: Optional[str] = None def registered_models(self) -> dict: @@ -45,6 +46,7 @@ def registered_models(self) -> dict: "engine": self.engine, "device": self.device, "runtime_config": self.runtime_config, + "tool_call_parser": self.tool_call_parser, "status": self.status.value, "time_loaded": self.time_loaded.isoformat(), } @@ -96,6 +98,9 @@ async def register_load(self, loader: ModelLoadConfig) -> str: engine=loader.engine, device=loader.device, runtime_config=loader.runtime_config, + tool_call_parser=( + loader.tool_call_parser.value if loader.tool_call_parser else None + ), status=ModelStatus.LOADING, ) diff --git a/src/server/routes/openai.py b/src/server/routes/openai.py index cd923dd..543b4be 100644 --- a/src/server/routes/openai.py +++ b/src/server/routes/openai.py @@ -31,137 +31,19 @@ OpenArcASRConfig, RerankRequest, ) -from src.engine.ov_genai.qwen_tool_parser import ( - QwenXmlToolCallParser, - ReasoningSplitter, -) +from src.engine.ov_genai.tool_parse import gemma4, hermes, qwen35 logger = logging.getLogger(__name__) router = APIRouter(prefix="/v1") -# ---- tool call helpers ---- - -def _extract_hermes_tool_call_payloads(text: str) -> List[str]: - open_tag = "" - close_tag = "" - payloads: List[str] = [] - cursor = 0 - - while True: - start = text.find(open_tag, cursor) - if start < 0: - break - - payload_start = start + len(open_tag) - end = text.find(close_tag, payload_start) - if end < 0: - payload = text[payload_start:].strip() - if payload: - payloads.append(payload) - break - - payload = text[payload_start:end].strip() - if payload: - payloads.append(payload) - - cursor = end + len(close_tag) - - return payloads - - -def _format_tool_call_arguments(arguments: Any) -> str: - if isinstance(arguments, str): - try: - return json.dumps(json.loads(arguments)) - except json.JSONDecodeError: - return arguments - return json.dumps(arguments) - - -def _finalize_qwen_calls(parser: QwenXmlToolCallParser) -> List[Dict[str, Any]]: - tool_calls: List[Dict[str, Any]] = [] - for call in parser.tool_calls: - fn = call.get("function") or {} - name = fn.get("name") or "" - if not name: - continue - tool_calls.append( - { - "id": call.get("id") or f"call_{uuid.uuid4().hex[:24]}", - "type": "function", - "function": { - "name": name, - "arguments": fn.get("arguments") or "{}", - }, - } - ) - return tool_calls - - -def parse_hermes_tool_calls(text: str) -> Optional[List[Dict[str, Any]]]: - tool_calls: List[Dict[str, Any]] = [] - for payload in _extract_hermes_tool_call_payloads(text): - try: - data = json.loads(payload) - if isinstance(data, dict) and "name" in data and "arguments" in data: - tool_calls.append( - { - "id": f"call_{uuid.uuid4().hex[:24]}", - "type": "function", - "function": { - "name": str(data.get("name", "")), - "arguments": _format_tool_call_arguments( - data.get("arguments", {}) - ), - }, - } - ) - except json.JSONDecodeError: - continue - return tool_calls if tool_calls else None - - -def parse_generation( - text: str, - tools: Optional[List[Dict[str, Any]]] = None, - enable_thinking: bool = True, -) -> tuple[str, str, Optional[List[Dict[str, Any]]]]: - """Split model output into (reasoning, content, tool_calls). - - Prefers Qwen XML ```` tool calls; falls back to Hermes JSON - inside ```` tags. - """ - # Whole-string parse: only split reasoning when is present. - # Otherwise Hermes / plain replies would be swallowed as thinking. - reasoning, remainder = ReasoningSplitter(enabled="" in text).feed(text) - if remainder.startswith(""): - remainder = remainder[len("") :] - - if "{payload}", "") - content = content.replace(f"\n{payload}\n", "") - return reasoning, content.strip(), hermes - - return reasoning, remainder, None - - -def parse_tool_calls( - text: str, tools: Optional[List[Dict[str, Any]]] = None -) -> Optional[List[Dict[str, Any]]]: - _, _, tool_calls = parse_generation(text, tools) - return tool_calls +# Tool-call parser modules keyed by ModelLoadConfig.tool_call_parser value. +_TOOL_PARSERS = { + "qwen35": qwen35, + "hermes": hermes, + "gemma4": gemma4, +} def _prepend_system_instruction(messages: Any, instruction: str) -> Any: @@ -258,6 +140,20 @@ async def openai_chat_completions( try: logger.info(f'"{request.model}" request received') + tool_parser_name = None + async with _registry._lock: + for record in _registry._models.values(): + if record.model_name == request.model: + tool_parser_name = record.tool_call_parser + break + + if tool_parser_name is None and request.tools: + raise ValueError( + f"Model '{request.model}' has no tool_call_parser configured; " + "set one in the model config (e.g. 'openarc add --tool-call-parser qwen35|hermes|gemma4')" + ) + parser_module = _TOOL_PARSERS.get(tool_parser_name) if tool_parser_name else None + messages, tools = _apply_tool_choice( request.messages, request.tools, @@ -284,6 +180,8 @@ async def openai_chat_completions( "presence_penalty": request.presence_penalty, "chat_template_kwargs": chat_template_kwargs, } + if parser_module is not None: + config_kwargs["tool_call_parser"] = tool_parser_name config_kwargs = {k: v for k, v in config_kwargs.items() if v is not None} generation_config = OVGenAI_GenConfig(**config_kwargs) @@ -304,9 +202,26 @@ async def event_stream() -> AsyncIterator[bytes]: metrics_data = None tool_call_sent = False cancel_request_id = None - reasoning = ReasoningSplitter(enabled=thinking_enabled) - tool_parser = QwenXmlToolCallParser(tools) - accumulated_text = "" + stream_parser = None + # qwen35 tool requests and gemma4 requests (tools or thinking; + # gemma4 thought-channel tags are token-ID-only) stream through + # the engine's tool streamers (parsed deltas on the worker + # queue); the text-delta facade only handles no-tool qwen35 + # requests. + engine_tool_stream = ( + parser_module is qwen35 and bool(tools) + ) or ( + parser_module is gemma4 + and gemma4.wants_engine_stream(tools, thinking_enabled) + ) + if parser_module is qwen35 and not engine_tool_stream: + stream_parser = qwen35.Qwen35StreamParser( + tools, enable_thinking=thinking_enabled + ) + elif parser_module is hermes: + stream_parser = hermes.HermesStreamParser( + enable_thinking=thinking_enabled + ) def _chunk(delta: dict) -> bytes: return ( @@ -331,20 +246,23 @@ def _chunk(delta: dict) -> bytes: if isinstance(item, dict): if item.get("error"): raise RuntimeError(item["error"]) + if "chat_delta" in item: + # Parsed deltas from Qwen35ToolCallStreamer + for delta in item["chat_delta"]: + if "tool_calls" in delta: + tool_call_sent = True + yield _chunk(delta) + continue metrics_data = item.get("metrics", item) continue - accumulated_text += item - reason_delta, text_delta = reasoning.feed(item) - content_delta, fragments = tool_parser.feed(text_delta) - - if reason_delta: - yield _chunk({"reasoning_content": reason_delta}) - if content_delta: - yield _chunk({"content": content_delta}) - if fragments: - tool_call_sent = True - yield _chunk({"tool_calls": fragments}) + if stream_parser is not None: + for delta in stream_parser.feed(item): + if "tool_calls" in delta: + tool_call_sent = True + yield _chunk(delta) + else: + yield _chunk({"content": item}) except asyncio.CancelledError: if cancel_request_id: await _workers.infer_cancel(cancel_request_id) @@ -353,39 +271,11 @@ def _chunk(delta: dict) -> bytes: ) raise - tool_parser.finalize() - if not tool_call_sent: - hermes = parse_hermes_tool_calls(accumulated_text) - if hermes: - tool_call_sent = True - for idx, tc in enumerate(hermes): - yield _chunk( - { - "tool_calls": [ - { - "index": idx, - "id": tc["id"], - "type": tc["type"], - "function": { - "name": tc["function"]["name"], - "arguments": "", - }, - } - ] - } - ) - yield _chunk( - { - "tool_calls": [ - { - "index": idx, - "function": { - "arguments": tc["function"]["arguments"] - }, - } - ] - } - ) + if stream_parser is not None: + for delta in stream_parser.finish(): + if "tool_calls" in delta: + tool_call_sent = True + yield _chunk(delta) prompt_tokens = (metrics_data or {}).get("input_token", 0) completion_tokens = (metrics_data or {}).get("new_token", 0) @@ -426,9 +316,12 @@ def _chunk(delta: dict) -> bytes: completion_tokens = metrics.get("new_token", 0) total_tokens = metrics.get("total_token", prompt_tokens + completion_tokens) - reasoning_text, content_text, tool_calls = parse_generation( - text, tools, thinking_enabled - ) + if parser_module is not None: + reasoning_text, content_text, tool_calls = parser_module.parse_generation( + text, tools, thinking_enabled + ) + else: + reasoning_text, content_text, tool_calls = None, text, None message = {"role": "assistant"} finish_reason = "stop" diff --git a/src/server/schemas/modeling/contract_ovgenai_llm_and_vlm.py b/src/server/schemas/modeling/contract_ovgenai_llm_and_vlm.py index 7233603..79fbbd8 100644 --- a/src/server/schemas/modeling/contract_ovgenai_llm_and_vlm.py +++ b/src/server/schemas/modeling/contract_ovgenai_llm_and_vlm.py @@ -65,6 +65,18 @@ class OVGenAI_GenConfig(BaseModel): default=None, description="List of tools/functions available to the model. None by default." ) + tool_call_parser: Optional[str] = Field( + default=None, + description=( + "Name of the server-side tool-call parser registered for the model " + "(e.g. 'qwen35'). Set by /v1/chat/completions; when tools are present " + "and the parser is qwen35, the engine streams with the token-ID " + "Qwen35ToolCallStreamer instead of ChunkStreamer (stream_chunk_tokens " + "does not apply to that path). gemma4 requests use the token-ID " + "Gemma4ToolCallStreamer whenever tools are present or thinking is " + "enabled (its protocol tags are special=True, invisible to text)." + ), + ) request_id: Optional[str] = Field( default=None, description="Request ID for tracking and cancellation." diff --git a/src/server/schemas/registration.py b/src/server/schemas/registration.py index 825ad54..fec3393 100644 --- a/src/server/schemas/registration.py +++ b/src/server/schemas/registration.py @@ -60,6 +60,19 @@ class EngineType(str, Enum): OPENVINO = "openvino" +class ToolCallParser(str, Enum): + """Tool-call output format the model was trained to emit, selected at load time. + + Options: + - qwen35: Qwen3.5 XML format (...) + - hermes: Hermes JSON format ({"name": ..., "arguments": {...}}) + - gemma4: Gemma 4 call syntax (<|tool_call>call:NAME{KEY:VALUE, ...})""" + + HERMES_PARSER = "hermes" + QWEN35_PARSER = "qwen35" + GEMMA4_PARSER = "gemma4" + + class ModelLoadConfig(BaseModel): model_path: str = Field( description=""" @@ -118,6 +131,14 @@ class ModelLoadConfig(BaseModel): default=None, description="Optional OpenVINO scheduler properties.", ) + tool_call_parser: Optional[ToolCallParser] = Field( + default=None, + description=""" + Tool-call parser for this model, selected at load time (llm/vlm only). + + When unset, /chat/completions requests containing tools are rejected + with 400.""", + ) class ModelUnloadConfig(BaseModel): diff --git a/src/server/worker_registry.py b/src/server/worker_registry.py index 82f6bc7..3a3da5a 100644 --- a/src/server/worker_registry.py +++ b/src/server/worker_registry.py @@ -134,7 +134,10 @@ async def infer_llm(packet: WorkerPacket, llm_instance: OVGenAI_LLM) -> WorkerPa try: async for item in llm_instance.generate_type(packet.gen_config): if isinstance(item, dict): - metrics = item + if "chat_delta" in item and packet.stream_queue is not None: + await packet.stream_queue.put(item) + elif "chat_delta" not in item: + metrics = item else: if packet.gen_config.stream: final_text += item @@ -166,7 +169,10 @@ async def infer_vlm(packet: WorkerPacket, vlm_model: OVGenAI_VLM) -> WorkerPacke try: async for item in vlm_model.generate_type(packet.gen_config): if isinstance(item, dict): - metrics = item + if "chat_delta" in item and packet.stream_queue is not None: + await packet.stream_queue.put(item) + elif "chat_delta" not in item: + metrics = item else: if packet.gen_config.stream: final_text += item diff --git a/tests/integration/test_tool_call_parser_integration.py b/tests/integration/test_tool_call_parser_integration.py new file mode 100644 index 0000000..f383e69 --- /dev/null +++ b/tests/integration/test_tool_call_parser_integration.py @@ -0,0 +1,317 @@ +import json +import os +from typing import Any, Dict, List + +import pytest # type: ignore[import] + +from test_model_path import model_path +from src.engine.ov_genai.llm import OVGenAI_LLM +from src.engine.ov_genai.vlm import OVGenAI_VLM +from src.engine.ov_genai.tool_parse import gemma4, hermes, qwen35 +from src.server.schemas.registration import ( + EngineType, + ModelLoadConfig, + ModelType, + ToolCallParser, +) +from src.server.schemas.modeling.contract_ovgenai_llm_and_vlm import OVGenAI_GenConfig + +HERMES_MODEL_PATH = model_path("OpenVINO/Qwen3-0.6B-int8_asym-ov") +QWEN35_MODEL_PATH = model_path("OpenVINO/Qwen3.5-2B-int4_sym-ov") + +TOOLS = [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather for a location", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"}, + "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}, + }, + "required": ["location"], + }, + }, + } +] + +MESSAGES = [ + {"role": "user", "content": "What is the weather in Warsaw in celsius? Use the get_weather tool."}, +] + + +class _DummyRegistry: + async def register_unload(self, model_name: str) -> bool: + return True + + +async def _generate_text(model_dir, model_name: str, tool_call_parser: ToolCallParser, + engine_cls, model_type: ModelType) -> str: + if not model_dir.exists(): + pytest.skip(f"Model path not found: {model_dir}") + + load_config = ModelLoadConfig( + model_path=str(model_dir), + model_name=model_name, + model_type=model_type, + engine=EngineType.OV_GENAI, + device=os.getenv("OPENARC_TEST_DEVICE", "CPU"), + runtime_config={}, + tool_call_parser=tool_call_parser, + ) + llm = engine_cls(load_config) + llm.load_model(load_config) + + try: + gen_config = OVGenAI_GenConfig( + messages=MESSAGES, + tools=TOOLS, + max_tokens=256, + temperature=0.1, + top_k=1, + top_p=1.0, + stream=False, + chat_template_kwargs={"enable_thinking": False}, + ) + outputs = [] + async for item in llm.generate_text(gen_config): + outputs.append(item) + + assert len(outputs) == 2 + _, text = outputs + assert isinstance(text, str) and text.strip() + return text + finally: + await llm.unload_model(_DummyRegistry(), load_config.model_name) + + +@pytest.mark.asyncio +async def test_hermes_tool_call_integration() -> None: + text = await _generate_text( + HERMES_MODEL_PATH, "integration-hermes", ToolCallParser.HERMES_PARSER, + OVGenAI_LLM, ModelType.LLM, + ) + + _, content, tool_calls = hermes.parse_generation(text, TOOLS, enable_thinking=False) + + assert tool_calls, f"Expected hermes tool calls in output: {text!r}" + assert tool_calls[0]["type"] == "function" + assert tool_calls[0]["function"]["name"] == "get_weather" + + arguments = json.loads(tool_calls[0]["function"]["arguments"]) + assert "location" in arguments + assert "" not in content + + +@pytest.mark.asyncio +async def test_qwen35_tool_call_integration() -> None: + text = await _generate_text( + QWEN35_MODEL_PATH, "integration-qwen35", ToolCallParser.QWEN35_PARSER, + OVGenAI_VLM, ModelType.VLM, + ) + + _, content, tool_calls = qwen35.parse_generation(text, TOOLS, enable_thinking=False) + + assert tool_calls, f"Expected qwen35 tool calls in output: {text!r}" + assert tool_calls[0]["type"] == "function" + assert tool_calls[0]["function"]["name"] == "get_weather" + + arguments = json.loads(tool_calls[0]["function"]["arguments"]) + assert "location" in arguments + assert "" not in content + + +async def _collect_streamed_deltas(model_dir, model_name: str, tool_call_parser: ToolCallParser, + engine_cls, model_type: ModelType) -> List[Dict[str, Any]]: + """Stream a tool request through the engine and return parsed chat deltas.""" + if not model_dir.exists(): + pytest.skip(f"Model path not found: {model_dir}") + + load_config = ModelLoadConfig( + model_path=str(model_dir), + model_name=model_name, + model_type=model_type, + engine=EngineType.OV_GENAI, + device=os.getenv("OPENARC_TEST_DEVICE", "CPU"), + runtime_config={}, + tool_call_parser=tool_call_parser, + ) + engine = engine_cls(load_config) + engine.load_model(load_config) + + try: + gen_config = OVGenAI_GenConfig( + messages=MESSAGES, + tools=TOOLS, + max_tokens=256, + temperature=0.1, + top_k=1, + top_p=1.0, + stream=True, + chat_template_kwargs={"enable_thinking": False}, + tool_call_parser="qwen35", + ) + deltas: List[Dict[str, Any]] = [] + async for item in engine.generate_stream(gen_config): + if isinstance(item, dict) and "chat_delta" in item: + deltas.extend(item["chat_delta"]) + return deltas + finally: + await engine.unload_model(_DummyRegistry(), load_config.model_name) + + +@pytest.mark.asyncio +async def test_qwen35_tool_call_streaming_integration() -> None: + deltas = await _collect_streamed_deltas( + QWEN35_MODEL_PATH, "integration-qwen35-stream", ToolCallParser.QWEN35_PARSER, + OVGenAI_VLM, ModelType.VLM, + ) + + tool_frags = [f for d in deltas for f in d.get("tool_calls", [])] + assert tool_frags, f"Expected streamed tool-call fragments: {deltas!r}" + + names = [f["function"]["name"] for f in tool_frags if f.get("function", {}).get("name")] + assert "get_weather" in names + + args = "".join( + f["function"]["arguments"] + for f in tool_frags + if "arguments" in f.get("function", {}) + ) + assert json.loads(args).get("location"), f"Expected location argument: {args!r}" + + assert tool_frags[0].get("id", "").startswith("call_") + assert [f["index"] for f in tool_frags] == sorted(f["index"] for f in tool_frags) + + +def test_qwen35_streamer_offline_write_tokens() -> None: + """Drive Qwen35ToolCallStreamer.write() with token chunks from a real + tokenizer (no model, no GPU): locks in the write/decode/parser contract.""" + if not QWEN35_MODEL_PATH.exists(): + pytest.skip(f"Model path not found: {QWEN35_MODEL_PATH}") + + import openvino_genai as ov + + tokenizer = ov.Tokenizer(str(QWEN35_MODEL_PATH)) + gen_config = OVGenAI_GenConfig( + tools=TOOLS, + chat_template_kwargs={"enable_thinking": True}, + ) + streamer = qwen35.Qwen35ToolCallStreamer(tokenizer, gen_config) + + text = ( + "Checking the weather.\n" + qwen35.THINK_CLOSE + "\n\n" + + qwen35.TOOL_OPEN + "\n\n" + "\nOslo\n\n" + "\n" + qwen35.TOOL_CLOSE + "\n" + ) + token_ids = tokenizer.encode(text).input_ids.data.tolist()[0] + for i in range(0, len(token_ids), 3): + streamer.write(token_ids[i : i + 3]) + streamer.end() + + items = [] + while not streamer.text_queue.empty(): + items.append(streamer.text_queue.get_nowait()) + assert items, "Streamer produced no queue items" + assert items[-1] is None + + deltas: List[Dict[str, Any]] = [] + for item in items[:-1]: + assert "chat_delta" in item, f"Unexpected queue item: {item!r}" + deltas.extend(item["chat_delta"]) + + reasoning = "".join(d.get("reasoning_content", "") for d in deltas) + assert "Checking the weather." in reasoning + + tool_frags = [f for d in deltas for f in d.get("tool_calls", [])] + args = "".join( + f["function"]["arguments"] + for f in tool_frags + if "arguments" in f.get("function", {}) + ) + assert json.loads(args) == {"location": "Oslo"} + assert tool_frags[0]["id"].startswith("call_") + + +# ---- gemma4 (Chimera-X-26B) ---- + +GEMMA4_MODEL_PATH = model_path("OpenVINO/Gemma/Chimera-X-26B-A4B-int4-ov") + + +@pytest.mark.asyncio +@pytest.mark.skipif( + os.getenv("OPENARC_TEST_DEVICE", "CPU").startswith("CPU"), + reason="gemma4 live test needs GPU (26B model; CPU inference is impractical)", +) +async def test_gemma4_tool_call_integration() -> None: + text = await _generate_text( + GEMMA4_MODEL_PATH, "integration-gemma4", ToolCallParser.GEMMA4_PARSER, + OVGenAI_VLM, ModelType.VLM, + ) + + _, content, tool_calls = gemma4.parse_generation(text, TOOLS, enable_thinking=False) + + assert tool_calls, f"Expected gemma4 tool calls in output: {text!r}" + assert tool_calls[0]["type"] == "function" + assert tool_calls[0]["function"]["name"] == "get_weather" + + arguments = json.loads(tool_calls[0]["function"]["arguments"]) + assert "location" in arguments + assert "<|tool_call>" not in content and "<|channel>" not in content + + +def test_gemma4_streamer_offline_write_tokens() -> None: + """Drive Gemma4ToolCallStreamer.write() with token chunks from a real + tokenizer (no model, no GPU): locks in the write/decode/parser contract + for special=True protocol tags (reasoning channel + tool call).""" + if not GEMMA4_MODEL_PATH.exists(): + pytest.skip(f"Model path not found: {GEMMA4_MODEL_PATH}") + + import openvino_genai as ov + + tokenizer = ov.Tokenizer(str(GEMMA4_MODEL_PATH)) + gen_config = OVGenAI_GenConfig( + tools=TOOLS, + chat_template_kwargs={"enable_thinking": True}, + ) + streamer = gemma4.Gemma4ToolCallStreamer(tokenizer, gen_config) + + text = ( + "Checking the weather.\n" + + gemma4.CHANNEL_OPEN + "thought\nLet me check.\n" + gemma4.CHANNEL_CLOSE + "\n" + + gemma4.TOOL_OPEN + "call:get_weather{location:Oslo}" + gemma4.TOOL_CLOSE + "\n" + ) + token_ids = tokenizer.encode(text).input_ids.data.tolist()[0] + for i in range(0, len(token_ids), 3): + streamer.write(token_ids[i : i + 3]) + streamer.end() + + items = [] + while not streamer.text_queue.empty(): + items.append(streamer.text_queue.get_nowait()) + assert items, "Streamer produced no queue items" + assert items[-1] is None + + deltas: List[Dict[str, Any]] = [] + for item in items[:-1]: + assert "chat_delta" in item, f"Unexpected queue item: {item!r}" + deltas.extend(item["chat_delta"]) + + reasoning = "".join(d.get("reasoning_content", "") for d in deltas) + assert "Let me check." in reasoning + + content = "".join(d.get("content", "") for d in deltas) + assert "Checking the weather." in content + assert "<|channel>" not in content and "thought" not in content.split("Checking")[0] + + tool_frags = [f for d in deltas for f in d.get("tool_calls", [])] + args = "".join( + f["function"]["arguments"] + for f in tool_frags + if "arguments" in f.get("function", {}) + ) + assert json.loads(args) == {"location": "Oslo"} + assert tool_frags[0]["id"].startswith("call_") diff --git a/tests/unit/test_tool_call_parser_unit.py b/tests/unit/test_tool_call_parser_unit.py index 13d48fc..44a14d1 100644 --- a/tests/unit/test_tool_call_parser_unit.py +++ b/tests/unit/test_tool_call_parser_unit.py @@ -1,10 +1,13 @@ import json -from typing import Any, AsyncIterator, Dict, List +from types import SimpleNamespace +from typing import Any, AsyncIterator, Dict, List, Optional import pytest # type: ignore[import] +from fastapi import HTTPException from fastapi.responses import StreamingResponse import src.server.routes.openai as openai_routes +from src.engine.ov_genai.tool_parse import gemma4, hermes, qwen35 from src.server.schemas.requests_openai import OpenAIChatCompletionRequest from src.server.utils.chat import flatten_messages, normalize_tool_calls_for_template @@ -61,6 +64,26 @@ async def is_disconnected(self) -> bool: return False +class _FakeRegistry: + """Minimal registry stand-in: one record with a configurable parser.""" + + def __init__(self, tool_call_parser: Optional[str]) -> None: + class _Lock: + async def __aenter__(self) -> "_Lock": + return self + + async def __aexit__(self, *exc: Any) -> bool: + return False + + self._lock = _Lock() + self._models = { + "fake-id": SimpleNamespace( + model_name="demo-model", + tool_call_parser=tool_call_parser, + ) + } + + def _extract_sse_payloads(chunks: List[bytes]) -> List[str]: payloads: List[str] = [] for chunk in chunks: @@ -70,14 +93,17 @@ def _extract_sse_payloads(chunks: List[bytes]) -> List[str]: return payloads -def test_parse_tool_calls_supports_hermes_tool_call_tags() -> None: +# ---- hermes parser unit tests ---- + + +def test_parse_generation_supports_hermes_tool_call_tags() -> None: text = ( "" '{"name":"search","arguments":{"query":"OpenVINO"}}' "" ) - tool_calls = openai_routes.parse_tool_calls(text) + _, _, tool_calls = hermes.parse_generation(text) assert tool_calls is not None assert len(tool_calls) == 1 @@ -86,10 +112,10 @@ def test_parse_tool_calls_supports_hermes_tool_call_tags() -> None: assert json.loads(tool_calls[0]["function"]["arguments"]) == {"query": "OpenVINO"} -def test_parse_tool_calls_supports_missing_closing_tag_until_eos() -> None: +def test_parse_generation_supports_missing_closing_tag_until_eos() -> None: text = '{"name":"search","arguments":{"query":"vLLM"}}' - tool_calls = openai_routes.parse_tool_calls(text) + _, _, tool_calls = hermes.parse_generation(text) assert tool_calls is not None assert len(tool_calls) == 1 @@ -97,98 +123,38 @@ def test_parse_tool_calls_supports_missing_closing_tag_until_eos() -> None: assert json.loads(tool_calls[0]["function"]["arguments"]) == {"query": "vLLM"} -def test_parse_tool_calls_rejects_plain_json_without_tool_call_tags() -> None: +def test_parse_generation_rejects_plain_json_without_tool_call_tags() -> None: text = '{"name":"search","arguments":{"query":"legacy"}}' - tool_calls = openai_routes.parse_tool_calls(text) + _, _, tool_calls = hermes.parse_generation(text) assert tool_calls is None -@pytest.mark.asyncio -async def test_openai_chat_completions_non_streaming_tool_calls(monkeypatch: pytest.MonkeyPatch) -> None: - class _Workers: - async def generate(self, model_name: str, generation_config: Any) -> Dict[str, Any]: - return { - "text": ( - "" - '{"name":"search","arguments":{"query":"OpenArc"}}' - "" - ), - "metrics": {"input_token": 4, "new_token": 6, "total_token": 10}, - } - - monkeypatch.setattr(openai_routes, "_workers", _Workers()) - - request = OpenAIChatCompletionRequest( - model="demo-model", - messages=[{"role": "user", "content": "Find OpenArc docs"}], - stream=False, - ) - - response = await openai_routes.openai_chat_completions(request, _DummyRequest()) - - choice = response["choices"][0] - assert choice["finish_reason"] == "tool_calls" - assert choice["message"]["content"] is None - assert len(choice["message"]["tool_calls"]) == 1 - assert choice["message"]["tool_calls"][0]["function"]["name"] == "search" - assert json.loads(choice["message"]["tool_calls"][0]["function"]["arguments"]) == { +def test_hermes_stream_parser_incremental() -> None: + parser = hermes.HermesStreamParser(enable_thinking=False) + deltas: List[Dict[str, Any]] = [] + for chunk in ( + "The answer", + " is.{"name":"search","arguments":{"query":"OpenArc"}}', + "", + ): + deltas.extend(parser.feed(chunk)) + deltas.extend(parser.finish()) + + content = "".join(d.get("content", "") for d in deltas) + assert content == "The answer is." + + tool_deltas = [d for d in deltas if "tool_calls" in d] + assert len(tool_deltas) == 2 + assert tool_deltas[0]["tool_calls"][0]["function"]["name"] == "search" + assert json.loads(tool_deltas[1]["tool_calls"][0]["function"]["arguments"]) == { "query": "OpenArc" } -@pytest.mark.asyncio -async def test_openai_chat_completions_streaming_hermes_tool_call(monkeypatch: pytest.MonkeyPatch) -> None: - class _Workers: - async def stream_generate(self, model_name: str, generation_config: Any) -> AsyncIterator[Any]: - yield "{"name":"search","arguments":{"query":"OpenArc"}}' - yield "" - yield {"metrics": {"input_token": 2, "new_token": 3, "total_token": 5}} - - async def infer_cancel(self, request_id: str) -> None: - return None - - monkeypatch.setattr(openai_routes, "_workers", _Workers()) - - request = OpenAIChatCompletionRequest( - model="demo-model", - messages=[{"role": "user", "content": "Find OpenArc docs"}], - stream=True, - ) - - response = await openai_routes.openai_chat_completions(request, _DummyRequest()) - assert isinstance(response, StreamingResponse) - - chunks: List[bytes] = [] - async for chunk in response.body_iterator: - chunks.append(chunk) - - payloads = _extract_sse_payloads(chunks) - assert payloads[-1] == "[DONE]" - - json_payloads = [json.loads(p) for p in payloads if p != "[DONE]"] - - content_deltas = [ - payload - for payload in json_payloads - if payload["choices"][0]["delta"].get("content") - ] - assert content_deltas == [] - - tool_deltas = [ - payload - for payload in json_payloads - if payload["choices"][0]["delta"].get("tool_calls") - ] - assert len(tool_deltas) >= 2 - assert tool_deltas[0]["choices"][0]["delta"]["tool_calls"][0]["function"]["name"] == "search" - assert json.loads( - tool_deltas[1]["choices"][0]["delta"]["tool_calls"][0]["function"]["arguments"] - ) == {"query": "OpenArc"} - - assert json_payloads[-1]["choices"][0]["finish_reason"] == "tool_calls" +# ---- qwen35 parser unit tests ---- QWEN_TOOLS = [ @@ -227,38 +193,42 @@ async def infer_cancel(self, request_id: str) -> None: QWEN_SINGLE = ( - "The user asked for weather.\n\n\n" - "\n\n" + "The user asked for weather.\n" + qwen35.THINK_CLOSE + "\n\n" + + qwen35.TOOL_OPEN + "\n\n" "\nWarsaw\n\n" "\ncelsius\n\n" - "\n\n" + "\n" + qwen35.TOOL_CLOSE + "\n" ) QWEN_PARALLEL = ( - "Need two tools.\n\n\n" - "\n\n" + "Need two tools.\n" + qwen35.THINK_CLOSE + "\n\n" + + qwen35.TOOL_OPEN + "\n\n" "\nWarsaw\n\n" "\ncelsius\n\n" - "\n\n" - "\n\n" + "\n" + qwen35.TOOL_CLOSE + "\n" + + qwen35.TOOL_OPEN + "\n\n" "\ncall mom\n\n" "\n3\n\n" "\nTrue\n\n" - "\n\n" + "\n" + qwen35.TOOL_CLOSE + "\n" ) QWEN_TYPED = ( - "Reminder time.\n\n\n" - "\n\n" + "Reminder time.\n" + qwen35.THINK_CLOSE + "\n\n" + + qwen35.TOOL_OPEN + "\n\n" "\nwater the plants\n\n" "\n5\n\n" "\nFalse\n\n" - "\n\n" + "\n" + qwen35.TOOL_CLOSE + "\n" ) -def test_parse_tool_calls_supports_qwen_xml() -> None: - tool_calls = openai_routes.parse_tool_calls(QWEN_SINGLE, QWEN_TOOLS) +def _qwen_tool_calls(text: str) -> Optional[List[Dict[str, Any]]]: + return qwen35.parse_generation(text, QWEN_TOOLS)[2] + + +def test_parse_generation_supports_qwen_xml() -> None: + tool_calls = _qwen_tool_calls(QWEN_SINGLE) assert tool_calls is not None assert len(tool_calls) == 1 @@ -269,8 +239,8 @@ def test_parse_tool_calls_supports_qwen_xml() -> None: } -def test_parse_tool_calls_supports_qwen_xml_parallel_and_bools() -> None: - tool_calls = openai_routes.parse_tool_calls(QWEN_PARALLEL, QWEN_TOOLS) +def test_parse_generation_supports_qwen_xml_parallel_and_bools() -> None: + tool_calls = _qwen_tool_calls(QWEN_PARALLEL) assert tool_calls is not None assert [c["function"]["name"] for c in tool_calls] == [ @@ -279,28 +249,177 @@ def test_parse_tool_calls_supports_qwen_xml_parallel_and_bools() -> None: ] assert json.loads(tool_calls[1]["function"]["arguments"]) == { "task": "call mom", - "days": 3, - "urgent": True, + "days": "3", + "urgent": "True", } def test_parse_generation_strips_thinking_and_xml() -> None: - reasoning, content, tool_calls = openai_routes.parse_generation( - QWEN_SINGLE, QWEN_TOOLS - ) + reasoning, content, tool_calls = qwen35.parse_generation(QWEN_SINGLE, QWEN_TOOLS) assert "weather" in reasoning - assert "" not in content + assert qwen35.TOOL_OPEN not in content assert " None: - tool_calls = openai_routes.parse_tool_calls(QWEN_TYPED, QWEN_TOOLS) +def test_parse_generation_qwen_typed_params() -> None: + tool_calls = _qwen_tool_calls(QWEN_TYPED) + assert tool_calls is not None args = json.loads(tool_calls[0]["function"]["arguments"]) - assert args["days"] == 5 - assert args["urgent"] is False + assert args == {"task": "water the plants", "days": "5", "urgent": "False"} + + +def test_parse_generation_without_think_close_is_pure_content() -> None: + reasoning, content, tool_calls = qwen35.parse_generation("just an answer", QWEN_TOOLS) + assert reasoning == "" + assert content == "just an answer" + assert tool_calls is None + + +def _feed_tokenized(text: str, **parser_kwargs): + """Feed text through QwenXMLToolParser with synthesized boundary IDs.""" + parser = qwen35.QwenXMLToolParser(**parser_kwargs) + synth = qwen35._TextTagSynthesizer() + content: List[str] = [] + for chunk, delta_tokens in synth.feed(text): + content.append(parser.parse({}, chunk, delta_tokens)) + tail = synth.flush() + if tail: + content.append(parser.parse({}, tail, None)) + return parser, "".join(content) + + +def test_boundary_requires_token_id() -> None: + parser = qwen35.QwenXMLToolParser() + out = parser.parse({}, "abc" + qwen35.TOOL_OPEN + "def", None) + assert out == "abc" + qwen35.TOOL_OPEN + "def" + assert parser.completed_calls == [] + + +def test_boundary_fires_on_token_id_and_slices_tag() -> None: + parser = qwen35.QwenXMLToolParser() + out = parser.parse({}, "abc" + qwen35.TOOL_OPEN + "def", [qwen35.TOOL_OPEN_ID]) + assert out == "abc" + assert parser._calls_seen == 1 + + +def test_grouped_delta_with_both_boundary_ids() -> None: + payload = ( + qwen35.TOOL_OPEN + "\n\n" + "\nWarsaw\n\n" + "\n" + qwen35.TOOL_CLOSE + ) + parser = qwen35.QwenXMLToolParser() + out = parser.parse({}, payload, [qwen35.TOOL_OPEN_ID, qwen35.TOOL_CLOSE_ID]) + assert out == "" + assert len(parser.completed_calls) == 1 + assert json.loads(parser.completed_calls[0]["function"]["arguments"]) == { + "location": "Warsaw" + } + + +def test_fragment_stream_shape() -> None: + parser, _ = _feed_tokenized(QWEN_SINGLE) + fragments: List[Dict[str, Any]] = [] + parser2 = qwen35.QwenXMLToolParser(on_fragment=fragments.append) + synth = qwen35._TextTagSynthesizer() + for chunk, delta_tokens in synth.feed(QWEN_SINGLE): + parser2.parse({}, chunk, delta_tokens) + + assert fragments[0]["function"] == {"name": "", "arguments": ""} + assert fragments[0]["id"].startswith("call_") + assert fragments[1]["function"] == {"name": "get_weather"} + args = "".join( + f["function"]["arguments"] + for f in fragments + if "arguments" in f.get("function", {}) + ) + assert json.loads(args) == {"location": "Warsaw", "unit": "celsius"} + assert [f["index"] for f in fragments] == [0] * len(fragments) + + +def test_ramble_after_call_sets_tool_call_stop() -> None: + parser, content = _feed_tokenized(QWEN_SINGLE + " oops extra text") + assert parser.status == qwen35.StreamingStatus.TOOL_CALL_STOP + assert parser.get_status() == qwen35.StreamingStatus.TOOL_CALL_STOP + assert "oops" not in content + + +def test_sequential_parallel_calls_allowed_by_default() -> None: + parser, content = _feed_tokenized(QWEN_PARALLEL) + assert parser.status == qwen35.StreamingStatus.RUNNING + assert [c["function"]["name"] for c in parser.completed_calls] == [ + "get_weather", + "set_reminder", + ] + + +def test_stop_after_tool_call_stops_first_call() -> None: + parser, _ = _feed_tokenized(QWEN_PARALLEL, stop_after_tool_call=True) + assert [c["function"]["name"] for c in parser.completed_calls] == ["get_weather"] + assert parser.status == qwen35.StreamingStatus.TOOL_CALL_STOP + + +def test_malformed_block_records_error_and_valid_args() -> None: + text = ( + qwen35.TOOL_OPEN + "\ngarbage here\n" + qwen35.TOOL_CLOSE + "\n" + ) + parser, _ = _feed_tokenized(text) + assert any("expected" in e for e in parser.errors) + assert parser.completed_calls[0]["function"]["arguments"] == "{}" + + +def test_finalize_force_closes_unterminated_call() -> None: + text = ( + qwen35.TOOL_OPEN + "\n\n" + "\nWarsaw\n\n" + "\n" + ) + parser, _ = _feed_tokenized(text) + parser.finalize() + assert any("unterminated" in e for e in parser.errors) + assert json.loads(parser.completed_calls[0]["function"]["arguments"]) == { + "location": "Warsaw" + } + + +def test_parse_mutates_msg_with_index_keyed_calls() -> None: + parser = qwen35.QwenXMLToolParser() + synth = qwen35._TextTagSynthesizer() + msg: Dict[str, Any] = {} + for chunk, delta_tokens in synth.feed(QWEN_PARALLEL): + parser.parse(msg, chunk, delta_tokens) + assert msg["tool_calls"]["0"]["function"]["name"] == "get_weather" + assert msg["tool_calls"]["1"]["function"]["name"] == "set_reminder" + + +def test_stream_parser_facade_split_boundary_tag() -> None: + text = ( + "Answer. " + qwen35.TOOL_OPEN + "\n\n" + "\nOslo\n\n" + "\n" + qwen35.TOOL_CLOSE + ) + cut = text.index(qwen35.TOOL_OPEN) + 3 # inside the open tag + parser = qwen35.Qwen35StreamParser(QWEN_TOOLS, enable_thinking=False) + deltas: List[Dict[str, Any]] = [] + deltas.extend(parser.feed(text[:cut])) + deltas.extend(parser.feed(text[cut:])) + deltas.extend(parser.finish()) + + content = "".join(d.get("content", "") for d in deltas) + assert content == "Answer. " + tool_frags = [f for d in deltas for f in d.get("tool_calls", [])] + args = "".join( + f["function"]["arguments"] + for f in tool_frags + if "arguments" in f.get("function", {}) + ) + assert json.loads(args) == {"location": "Oslo"} + + +# ---- _apply_tool_choice ---- def test_apply_tool_choice_none_hides_tools() -> None: @@ -342,6 +461,97 @@ def test_apply_named_tool_choice_rejects_unknown_tool() -> None: ) +# ---- route tests ---- + + +@pytest.mark.asyncio +async def test_openai_chat_completions_non_streaming_tool_calls(monkeypatch: pytest.MonkeyPatch) -> None: + class _Workers: + async def generate(self, model_name: str, generation_config: Any) -> Dict[str, Any]: + return { + "text": ( + "" + '{"name":"search","arguments":{"query":"OpenArc"}}' + "" + ), + "metrics": {"input_token": 4, "new_token": 6, "total_token": 10}, + } + + monkeypatch.setattr(openai_routes, "_workers", _Workers()) + monkeypatch.setattr(openai_routes, "_registry", _FakeRegistry("hermes")) + + request = OpenAIChatCompletionRequest( + model="demo-model", + messages=[{"role": "user", "content": "Find OpenArc docs"}], + stream=False, + ) + + response = await openai_routes.openai_chat_completions(request, _DummyRequest()) + + choice = response["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["content"] is None + assert len(choice["message"]["tool_calls"]) == 1 + assert choice["message"]["tool_calls"][0]["function"]["name"] == "search" + assert json.loads(choice["message"]["tool_calls"][0]["function"]["arguments"]) == { + "query": "OpenArc" + } + + +@pytest.mark.asyncio +async def test_openai_chat_completions_streaming_hermes_tool_call(monkeypatch: pytest.MonkeyPatch) -> None: + class _Workers: + async def stream_generate(self, model_name: str, generation_config: Any) -> AsyncIterator[Any]: + yield "{"name":"search","arguments":{"query":"OpenArc"}}' + yield "" + yield {"metrics": {"input_token": 2, "new_token": 3, "total_token": 5}} + + async def infer_cancel(self, request_id: str) -> None: + return None + + monkeypatch.setattr(openai_routes, "_workers", _Workers()) + monkeypatch.setattr(openai_routes, "_registry", _FakeRegistry("hermes")) + + request = OpenAIChatCompletionRequest( + model="demo-model", + messages=[{"role": "user", "content": "Find OpenArc docs"}], + stream=True, + ) + + response = await openai_routes.openai_chat_completions(request, _DummyRequest()) + assert isinstance(response, StreamingResponse) + + chunks: List[bytes] = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + + payloads = _extract_sse_payloads(chunks) + assert payloads[-1] == "[DONE]" + + json_payloads = [json.loads(p) for p in payloads if p != "[DONE]"] + + content_deltas = [ + payload + for payload in json_payloads + if payload["choices"][0]["delta"].get("content") + ] + assert content_deltas == [] + + tool_deltas = [ + payload + for payload in json_payloads + if payload["choices"][0]["delta"].get("tool_calls") + ] + assert len(tool_deltas) >= 2 + assert tool_deltas[0]["choices"][0]["delta"]["tool_calls"][0]["function"]["name"] == "search" + assert json.loads( + tool_deltas[1]["choices"][0]["delta"]["tool_calls"][0]["function"]["arguments"] + ) == {"query": "OpenArc"} + + assert json_payloads[-1]["choices"][0]["finish_reason"] == "tool_calls" + + @pytest.mark.asyncio async def test_openai_chat_completions_non_streaming_qwen_xml( monkeypatch: pytest.MonkeyPatch, @@ -354,6 +564,7 @@ async def generate(self, model_name: str, generation_config: Any) -> Dict[str, A } monkeypatch.setattr(openai_routes, "_workers", _Workers()) + monkeypatch.setattr(openai_routes, "_registry", _FakeRegistry("qwen35")) request = OpenAIChatCompletionRequest( model="demo-model", @@ -374,17 +585,42 @@ async def generate(self, model_name: str, generation_config: Any) -> Dict[str, A async def test_openai_chat_completions_streaming_qwen_xml( monkeypatch: pytest.MonkeyPatch, ) -> None: + # Tools + qwen35 parser -> the engine streams parsed deltas; the fake + # worker yields the Qwen35ToolCallStreamer queue items directly. + seen_configs: List[Any] = [] + class _Workers: async def stream_generate(self, model_name: str, generation_config: Any) -> AsyncIterator[Any]: - text = QWEN_SINGLE - for i in range(0, len(text), 7): - yield text[i : i + 7] + seen_configs.append(generation_config) + yield {"chat_delta": [{"reasoning_content": "thinking"}]} + yield { + "chat_delta": [ + { + "tool_calls": [ + { + "index": 0, + "id": "call_abc123", + "type": "function", + "function": {"name": "", "arguments": ""}, + } + ] + }, + {"tool_calls": [{"index": 0, "function": {"name": "get_weather"}}]}, + { + "tool_calls": [ + {"index": 0, "function": {"arguments": '{"location": '}}, + {"index": 0, "function": {"arguments": '"Warsaw"}'}}, + ] + }, + ] + } yield {"metrics": {"input_token": 2, "new_token": 3, "total_token": 5}} async def infer_cancel(self, request_id: str) -> None: return None monkeypatch.setattr(openai_routes, "_workers", _Workers()) + monkeypatch.setattr(openai_routes, "_registry", _FakeRegistry("qwen35")) request = OpenAIChatCompletionRequest( model="demo-model", @@ -398,7 +634,13 @@ async def infer_cancel(self, request_id: str) -> None: async for chunk in response.body_iterator: chunks.append(chunk) + assert seen_configs and seen_configs[0].tool_call_parser == "qwen35" + payloads = [json.loads(p) for p in _extract_sse_payloads(chunks) if p != "[DONE]"] + reasoning = "".join( + p["choices"][0]["delta"].get("reasoning_content", "") for p in payloads + ) + assert reasoning == "thinking" names = [] args = "" for payload in payloads: @@ -409,5 +651,594 @@ async def infer_cancel(self, request_id: str) -> None: if fn.get("arguments"): args += fn["arguments"] assert "get_weather" in names - assert json.loads(args) == {"location": "Warsaw", "unit": "celsius"} + assert json.loads(args) == {"location": "Warsaw"} + assert payloads[-1]["choices"][0]["finish_reason"] == "tool_calls" + + +@pytest.mark.asyncio +async def test_openai_chat_completions_streaming_qwen35_no_tools_uses_text_facade( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # No tools -> ChunkStreamer text path; the facade does reasoning-only + # splitting plus the tool-XML guard. + class _Workers: + async def stream_generate(self, model_name: str, generation_config: Any) -> AsyncIterator[Any]: + text = " pondering" + qwen35.THINK_CLOSE + "Just an answer." + for i in range(0, len(text), 5): + yield text[i : i + 5] + yield {"metrics": {"input_token": 2, "new_token": 3, "total_token": 5}} + + async def infer_cancel(self, request_id: str) -> None: + return None + + monkeypatch.setattr(openai_routes, "_workers", _Workers()) + monkeypatch.setattr(openai_routes, "_registry", _FakeRegistry("qwen35")) + + request = OpenAIChatCompletionRequest( + model="demo-model", + messages=[{"role": "user", "content": "Hi"}], + stream=True, + ) + + response = await openai_routes.openai_chat_completions(request, _DummyRequest()) + chunks: List[bytes] = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + + payloads = [json.loads(p) for p in _extract_sse_payloads(chunks) if p != "[DONE]"] + reasoning = "".join( + p["choices"][0]["delta"].get("reasoning_content", "") for p in payloads[:-1] + ) + content = "".join( + p["choices"][0]["delta"].get("content", "") for p in payloads[:-1] + ) + assert reasoning == " pondering" + assert content == "Just an answer." + assert payloads[-1]["choices"][0]["finish_reason"] == "stop" + + +def test_select_streamer_picks_tool_streamer_for_qwen35_tools( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from src.engine.ov_genai import streamers as streamers_mod + + sentinels = {"tool": object(), "chunk": object()} + + class _FakeToolStreamer: + def __init__(self, tokenizer, gen_config): + self.args = (tokenizer, gen_config) + + class _FakeChunkStreamer: + def __init__(self, tokenizer, gen_config): + self.args = (tokenizer, gen_config) + + monkeypatch.setattr(streamers_mod, "ChunkStreamer", _FakeChunkStreamer) + monkeypatch.setattr( + streamers_mod.qwen35_tool_parse, "Qwen35ToolCallStreamer", _FakeToolStreamer + ) + + tokenizer = object() + tool_config = SimpleNamespace( + tools=[{"type": "function"}], tool_call_parser="qwen35", chat_template_kwargs={} + ) + picked = streamers_mod.select_streamer(tokenizer, tool_config) + assert isinstance(picked, _FakeToolStreamer) + assert picked.args[0] is tokenizer + + plain_config = SimpleNamespace(tools=None, tool_call_parser="qwen35", chat_template_kwargs={}) + assert isinstance( + streamers_mod.select_streamer(tokenizer, plain_config), _FakeChunkStreamer + ) + + hermes_config = SimpleNamespace( + tools=[{"type": "function"}], tool_call_parser="hermes", chat_template_kwargs={} + ) + assert isinstance( + streamers_mod.select_streamer(tokenizer, hermes_config), _FakeChunkStreamer + ) + + +# ---- unset tool_call_parser ---- + + +@pytest.mark.asyncio +async def test_tools_request_rejected_without_parser(monkeypatch: pytest.MonkeyPatch) -> None: + class _Workers: + async def generate(self, model_name: str, generation_config: Any) -> Dict[str, Any]: + return {"text": "hello", "metrics": {}} + + monkeypatch.setattr(openai_routes, "_workers", _Workers()) + monkeypatch.setattr(openai_routes, "_registry", _FakeRegistry(None)) + + request = OpenAIChatCompletionRequest( + model="demo-model", + messages=[{"role": "user", "content": "Hi"}], + tools=QWEN_TOOLS, + stream=False, + ) + + with pytest.raises(HTTPException) as exc_info: + await openai_routes.openai_chat_completions(request, _DummyRequest()) + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_plain_request_passthrough_without_parser(monkeypatch: pytest.MonkeyPatch) -> None: + raw_text = "Plain answer. not parsed." + + class _Workers: + async def generate(self, model_name: str, generation_config: Any) -> Dict[str, Any]: + return { + "text": raw_text, + "metrics": {"input_token": 2, "new_token": 3, "total_token": 5}, + } + + monkeypatch.setattr(openai_routes, "_workers", _Workers()) + monkeypatch.setattr(openai_routes, "_registry", _FakeRegistry(None)) + + request = OpenAIChatCompletionRequest( + model="demo-model", + messages=[{"role": "user", "content": "Hi"}], + stream=False, + ) + + response = await openai_routes.openai_chat_completions(request, _DummyRequest()) + choice = response["choices"][0] + assert choice["finish_reason"] == "stop" + assert choice["message"]["content"] == raw_text + assert "reasoning_content" not in choice["message"] + + +@pytest.mark.asyncio +async def test_streaming_passthrough_without_parser(monkeypatch: pytest.MonkeyPatch) -> None: + class _Workers: + async def stream_generate(self, model_name: str, generation_config: Any) -> AsyncIterator[Any]: + yield "chunk one. " + yield "chunk two." + yield {"metrics": {"input_token": 2, "new_token": 3, "total_token": 5}} + + async def infer_cancel(self, request_id: str) -> None: + return None + + monkeypatch.setattr(openai_routes, "_workers", _Workers()) + monkeypatch.setattr(openai_routes, "_registry", _FakeRegistry(None)) + + request = OpenAIChatCompletionRequest( + model="demo-model", + messages=[{"role": "user", "content": "Hi"}], + stream=True, + ) + + response = await openai_routes.openai_chat_completions(request, _DummyRequest()) + chunks: List[bytes] = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + + payloads = [json.loads(p) for p in _extract_sse_payloads(chunks) if p != "[DONE]"] + content = "".join( + p["choices"][0]["delta"].get("content", "") for p in payloads[:-1] + ) + assert content == "chunk one. chunk two." + assert payloads[-1]["choices"][0]["finish_reason"] == "stop" + + +# ---- gemma4 tool-call parser ---- + + +GEMMA_TOOLS = [{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get current weather for a city.", + "parameters": { + "type": "object", + "properties": {"city": {"type": "string", "description": "City name"}}, + "required": ["city"], + }, + }, +}] + + +def _feed_gemma(spec, **parser_kwargs): + """Drive Gemma4ToolCallParser with driver-contract deltas. + + spec: sequence of (delta_text, delta_tokens) pairs; text deltas never + carry protocol ids (the streamer flushes text before boundary tokens). + """ + parser = gemma4.Gemma4ToolCallParser(**parser_kwargs) + content: List[str] = [] + for text, ids in spec: + content.append(parser.parse({}, text, ids)) + return parser, "".join(content) + + +def test_gemma_boundary_requires_token_id() -> None: + # Payload text without the open-boundary ID is plain content. + parser, content = _feed_gemma([("call:get_weather{city:Paris}", None)]) + assert content == "call:get_weather{city:Paris}" + assert parser.completed_calls == [] + + +def test_gemma_boundary_fires_on_token_id() -> None: + parser, content = _feed_gemma([ + ("", [gemma4.TOOL_OPEN_ID]), + ("call:get_weather{city:Paris}", None), + ("", [gemma4.TOOL_CLOSE_ID]), + ]) + assert content == "" + assert len(parser.completed_calls) == 1 + assert json.loads(parser.completed_calls[0]["function"]["arguments"]) == { + "city": "Paris" + } + + +def test_gemma_fragment_stream_shape() -> None: + fragments: List[Dict[str, Any]] = [] + parser = gemma4.Gemma4ToolCallParser(on_fragment=fragments.append) + for text, ids in [ + ("", [gemma4.TOOL_OPEN_ID]), + ("call:get_weather{", None), + ("city:Paris", None), + (", ", None), + ("unit:celsius", None), + ("}", None), + ("", [gemma4.TOOL_CLOSE_ID]), + ]: + parser.parse({}, text, ids) + + assert fragments[0]["function"] == {"name": "", "arguments": ""} + assert fragments[0]["id"].startswith("call_") + assert fragments[1]["function"] == {"name": "get_weather"} + args = "".join( + f["function"]["arguments"] + for f in fragments + if "arguments" in f.get("function", {}) + ) + assert json.loads(args) == {"city": "Paris", "unit": "celsius"} + assert [f["index"] for f in fragments] == [0] * len(fragments) + + +def test_gemma_sequential_parallel_calls_allowed_by_default() -> None: + parser, content = _feed_gemma([ + ("", [gemma4.TOOL_OPEN_ID]), + ("call:get_weather{city:Tokyo}", None), + ("", [gemma4.TOOL_CLOSE_ID]), + ("\n\n", None), + ("", [gemma4.TOOL_OPEN_ID]), + ("call:get_weather{city:Paris}", None), + ("", [gemma4.TOOL_CLOSE_ID]), + ]) + assert parser.status == gemma4.StreamingStatus.RUNNING + assert [c["function"]["name"] for c in parser.completed_calls] == [ + "get_weather", "get_weather", + ] + assert json.loads(parser.completed_calls[0]["function"]["arguments"]) == {"city": "Tokyo"} + assert json.loads(parser.completed_calls[1]["function"]["arguments"]) == {"city": "Paris"} + assert content == "" + + +def test_gemma_ramble_after_call_sets_tool_call_stop() -> None: + parser, content = _feed_gemma([ + ("", [gemma4.TOOL_OPEN_ID]), + ("call:get_weather{city:Paris}", None), + ("", [gemma4.TOOL_CLOSE_ID]), + (" oops extra text", None), + ]) + assert parser.status == gemma4.StreamingStatus.TOOL_CALL_STOP + assert parser.get_status() == gemma4.StreamingStatus.TOOL_CALL_STOP + assert "oops" not in content + + +def test_gemma_stop_after_tool_call_stops_first_call() -> None: + parser, _ = _feed_gemma([ + ("", [gemma4.TOOL_OPEN_ID]), + ("call:first{a:1}", None), + ("", [gemma4.TOOL_CLOSE_ID]), + ("", [gemma4.TOOL_OPEN_ID]), + ("call:second{b:2}", None), + ("", [gemma4.TOOL_CLOSE_ID]), + ], stop_after_tool_call=True) + assert [c["function"]["name"] for c in parser.completed_calls] == ["first"] + assert parser.status == gemma4.StreamingStatus.TOOL_CALL_STOP + + +def test_gemma_malformed_payload_records_error_and_valid_args() -> None: + parser, _ = _feed_gemma([ + ("", [gemma4.TOOL_OPEN_ID]), + ("garbage here", None), + ("", [gemma4.TOOL_CLOSE_ID]), + ]) + assert any("expected 'call:'" in e for e in parser.errors) + assert parser.completed_calls[0]["function"]["arguments"] == "{}" + + +def test_gemma_finalize_force_closes_unterminated_call() -> None: + parser, _ = _feed_gemma([ + ("", [gemma4.TOOL_OPEN_ID]), + ("call:get_weather{city:Oslo", None), + ]) + parser.finalize() + assert any("unterminated" in e for e in parser.errors) + assert json.loads(parser.completed_calls[0]["function"]["arguments"]) == { + "city": "Oslo" + } + + +def test_gemma_parse_mutates_msg_with_index_keyed_calls() -> None: + parser = gemma4.Gemma4ToolCallParser() + msg: Dict[str, Any] = {} + for text, ids in [ + ("", [gemma4.TOOL_OPEN_ID]), + ("call:get_weather{city:Tokyo}", None), + ("", [gemma4.TOOL_CLOSE_ID]), + ("", [gemma4.TOOL_OPEN_ID]), + ("call:set_alarm{when:7am}", None), + ("", [gemma4.TOOL_CLOSE_ID]), + ]: + parser.parse(msg, text, ids) + assert msg["tool_calls"]["0"]["function"]["name"] == "get_weather" + assert msg["tool_calls"]["1"]["function"]["name"] == "set_alarm" + + +def test_gemma_nested_brace_value_and_zero_args() -> None: + parser, _ = _feed_gemma([ + ("", [gemma4.TOOL_OPEN_ID]), + ("call:foo{opts:{a:1, b:{c:2}}, note:hi}", None), + ("", [gemma4.TOOL_CLOSE_ID]), + ("", [gemma4.TOOL_OPEN_ID]), + ("call:noargs{}", None), + ("", [gemma4.TOOL_CLOSE_ID]), + ]) + assert json.loads(parser.completed_calls[0]["function"]["arguments"]) == { + "opts": "{a:1, b:{c:2}}", + "note": "hi", + } + assert parser.completed_calls[1]["function"]["arguments"] == "{}" + + +def test_gemma_channel_splitter_thought_channel() -> None: + splitter = gemma4.Gemma4ChannelSplitter() + deltas = [ + ("Hello ", []), + ("", [gemma4.CHANNEL_OPEN_ID]), + ("thought\nLet me think.\nMore thought.\n", []), + ("", [gemma4.CHANNEL_CLOSE_ID]), + ("Final answer", []), + ] + reason, content = "", "" + for text, ids in deltas: + r, c, passthrough = splitter.feed(text, ids) + reason += r + content += c + assert passthrough == [] or all(i not in (100, 101) for i in passthrough) + assert reason == "Let me think.\nMore thought." + assert content == "Hello Final answer" + + +def test_gemma_channel_splitter_opaque_channel_passthrough() -> None: + splitter = gemma4.Gemma4ChannelSplitter() + deltas = [ + ("", [gemma4.CHANNEL_OPEN_ID]), + ("summary\nstuff\n", []), + ("", [gemma4.CHANNEL_CLOSE_ID]), + ("tail", []), + ] + reason, content = "", "" + for text, ids in deltas: + r, c, _ = splitter.feed(text, ids) + reason += r + content += c + assert reason == "" + assert content == "summary\nstuff\ntail" + + +def test_gemma_channel_splitter_header_across_deltas() -> None: + splitter = gemma4.Gemma4ChannelSplitter() + r1, c1, _ = splitter.feed("", [gemma4.CHANNEL_OPEN_ID]) + r2, c2, _ = splitter.feed("thou", []) # header split mid-name + r3, c3, _ = splitter.feed("ght\nbody", []) + r4, c4, _ = splitter.feed("", [gemma4.CHANNEL_CLOSE_ID]) + assert "".join([r1, r2, r3, r4]) == "body" + assert "".join([c1, c2, c3, c4]) == "" + + +def test_gemma_parse_generation_reasoning_and_parallel_calls() -> None: + raw = ( + "<|channel>thought\nstep one\nstep two\n\n\nHello\n\n" + "<|tool_call>call:get_weather{city:Tokyo}\n" + "<|tool_call>call:get_weather{city:Paris, when:today}" + ) + reasoning, content, calls = gemma4.parse_generation(raw) + assert reasoning == "step one\nstep two" + assert content == "Hello" + assert [c["function"]["name"] for c in calls] == ["get_weather", "get_weather"] + assert json.loads(calls[0]["function"]["arguments"]) == {"city": "Tokyo"} + assert json.loads(calls[1]["function"]["arguments"]) == {"city": "Paris", "when": "today"} + assert all(c["id"].startswith("call_") for c in calls) + + +def test_gemma_parse_generation_passthrough_without_tags() -> None: + assert gemma4.parse_generation("plain answer") == ("", "plain answer", None) + + +def test_gemma_parse_generation_unterminated_tool_block() -> None: + raw = "Intro <|tool_call>call:get_weather{city:Oslo}" + reasoning, content, calls = gemma4.parse_generation(raw) + assert reasoning == "" + assert calls is None + assert "call:get_weather{city:Oslo}" in content + + +def test_gemma_wants_engine_stream() -> None: + assert gemma4.wants_engine_stream([{"type": "function"}], False) is True + assert gemma4.wants_engine_stream(None, True) is True + assert gemma4.wants_engine_stream(None, False) is False + + +def test_select_streamer_gemma4_engine_stream(monkeypatch: pytest.MonkeyPatch) -> None: + from src.engine.ov_genai import streamers as streamers_mod + + class _FakeGemmaStreamer: + def __init__(self, tokenizer, gen_config): + self.args = (tokenizer, gen_config) + + class _FakeChunkStreamer: + def __init__(self, tokenizer, gen_config): + self.args = (tokenizer, gen_config) + + monkeypatch.setattr(streamers_mod, "ChunkStreamer", _FakeChunkStreamer) + monkeypatch.setattr( + streamers_mod.gemma4_tool_parse, "Gemma4ToolCallStreamer", _FakeGemmaStreamer + ) + + tokenizer = object() + # Tools present -> engine streamer even with thinking off. + tools_cfg = SimpleNamespace( + tools=[{"type": "function"}], + tool_call_parser="gemma4", + chat_template_kwargs={"enable_thinking": False}, + ) + assert isinstance( + streamers_mod.select_streamer(tokenizer, tools_cfg), _FakeGemmaStreamer + ) + + # Thinking on, no tools -> engine streamer (channel tags are ID-only). + think_cfg = SimpleNamespace( + tools=None, + tool_call_parser="gemma4", + chat_template_kwargs={"enable_thinking": True}, + ) + assert isinstance( + streamers_mod.select_streamer(tokenizer, think_cfg), _FakeGemmaStreamer + ) + + # Default (no chat_template_kwargs) counts as thinking enabled. + default_cfg = SimpleNamespace( + tools=None, tool_call_parser="gemma4", chat_template_kwargs={} + ) + assert isinstance( + streamers_mod.select_streamer(tokenizer, default_cfg), _FakeGemmaStreamer + ) + + # Thinking explicitly off, no tools -> plain ChunkStreamer. + off_cfg = SimpleNamespace( + tools=None, + tool_call_parser="gemma4", + chat_template_kwargs={"enable_thinking": False}, + ) + assert isinstance( + streamers_mod.select_streamer(tokenizer, off_cfg), _FakeChunkStreamer + ) + + +# ---- gemma4 routes ---- + + +@pytest.mark.asyncio +async def test_openai_chat_completions_non_streaming_gemma4( + monkeypatch: pytest.MonkeyPatch, +) -> None: + raw = ( + "<|channel>thought\npondering\n\n" + "<|tool_call>call:get_weather{city:Warsaw}" + ) + + class _Workers: + async def generate(self, model_name: str, generation_config: Any) -> Dict[str, Any]: + return { + "text": raw, + "metrics": {"input_token": 4, "new_token": 6, "total_token": 10}, + } + + monkeypatch.setattr(openai_routes, "_workers", _Workers()) + monkeypatch.setattr(openai_routes, "_registry", _FakeRegistry("gemma4")) + + request = OpenAIChatCompletionRequest( + model="demo-model", + messages=[{"role": "user", "content": "Weather in Warsaw?"}], + tools=GEMMA_TOOLS, + stream=False, + ) + + response = await openai_routes.openai_chat_completions(request, _DummyRequest()) + choice = response["choices"][0] + assert choice["finish_reason"] == "tool_calls" + assert choice["message"]["tool_calls"][0]["function"]["name"] == "get_weather" + assert json.loads(choice["message"]["tool_calls"][0]["function"]["arguments"]) == { + "city": "Warsaw" + } + assert choice["message"].get("reasoning_content") == "pondering" + assert choice["message"].get("content") in (None, "") + + +@pytest.mark.asyncio +async def test_openai_chat_completions_streaming_gemma4_engine_stream( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Tools + gemma4 parser -> the engine streams parsed deltas via chat_delta. + seen_configs: List[Any] = [] + + class _Workers: + async def stream_generate(self, model_name: str, generation_config: Any) -> AsyncIterator[Any]: + seen_configs.append(generation_config) + yield {"chat_delta": [{"reasoning_content": "pondering"}]} + yield { + "chat_delta": [ + { + "tool_calls": [ + { + "index": 0, + "id": "call_abc123", + "type": "function", + "function": {"name": "", "arguments": ""}, + } + ] + }, + {"tool_calls": [{"index": 0, "function": {"name": "get_weather"}}]}, + { + "tool_calls": [ + {"index": 0, "function": {"arguments": '{"city": "'}}, + {"index": 0, "function": {"arguments": 'Warsaw"}'}}, + ] + }, + ] + } + yield {"metrics": {"input_token": 2, "new_token": 3, "total_token": 5}} + + async def infer_cancel(self, request_id: str) -> None: + return None + + monkeypatch.setattr(openai_routes, "_workers", _Workers()) + monkeypatch.setattr(openai_routes, "_registry", _FakeRegistry("gemma4")) + + request = OpenAIChatCompletionRequest( + model="demo-model", + messages=[{"role": "user", "content": "Weather in Warsaw?"}], + tools=GEMMA_TOOLS, + stream=True, + ) + + response = await openai_routes.openai_chat_completions(request, _DummyRequest()) + chunks: List[bytes] = [] + async for chunk in response.body_iterator: + chunks.append(chunk) + + assert seen_configs and seen_configs[0].tool_call_parser == "gemma4" + + payloads = [json.loads(p) for p in _extract_sse_payloads(chunks) if p != "[DONE]"] + reasoning = "".join( + p["choices"][0]["delta"].get("reasoning_content", "") for p in payloads + ) + assert reasoning == "pondering" + names = [] + args = "" + for payload in payloads: + for frag in payload["choices"][0]["delta"].get("tool_calls") or []: + fn = frag.get("function") or {} + if fn.get("name"): + names.append(fn["name"]) + if fn.get("arguments"): + args += fn["arguments"] + assert "get_weather" in names + assert json.loads(args) == {"city": "Warsaw"} assert payloads[-1]["choices"][0]["finish_reason"] == "tool_calls"