diff --git a/api/src/main/models/base.py b/api/src/main/models/base.py index 9307d620..e428e3c9 100644 --- a/api/src/main/models/base.py +++ b/api/src/main/models/base.py @@ -1,4 +1,5 @@ import logging +import re from re import finditer, DOTALL from dataclasses import dataclass from typing import Generator, Optional, List, Dict, Union @@ -172,6 +173,21 @@ def adaptive_special_tag_buffering(outs, wait_tokens_for=6): else: yield wd + tag_pattern = re.compile("(" + "|".join(re.escape(t) for t in ( + self.special_tags.REASONING, self.special_tags.REASONING_END, + self.special_tags.TOOLCALL, self.special_tags.TOOLCALL_END, + )) + ")") + + def split_special_tags(outs): + """ + Give every special tag a frame of its own. The app matches ``/`` exactly, and a + streamer that emits at word boundaries glues tags to whitespace or text (`'\\n'`). + """ + for wd in outs: + for piece in tag_pattern.split(wd): + if piece: + yield piece + initial_operation = True function_called = True while function_called: @@ -206,7 +222,7 @@ def adaptive_special_tag_buffering(outs, wait_tokens_for=6): ) generation_kwargs.update(kwargs) outputs = self.parse_tool_calling( - adaptive_special_tag_buffering(self.runtime(**generation_kwargs)), + split_special_tags(adaptive_special_tag_buffering(self.runtime(**generation_kwargs))), chat_history=chat_history, tools=tools, stream=stream, diff --git a/api/tests/test_think_frames.py b/api/tests/test_think_frames.py new file mode 100644 index 00000000..71952ec5 --- /dev/null +++ b/api/tests/test_think_frames.py @@ -0,0 +1,44 @@ +""" +Reasoning and tool tags reach the app as frames of their own (SPEC S1.4, #112): the app matches +`` and `` exactly, and the one-at-a-time path's streamer glues them to whitespace +(`'\n'`, `'\n\n'`). +""" +from api.src.main.models.base import BaseModel +from api.src.main.models.config import ChatHistory +from api.src.main.utils import FunctionCalling + + +class ScriptedEngine: + def __init__(self, chunks): + self.chunks = chunks + + def __call__(self, messages, **kwargs): + yield from self.chunks + + +class Plain(BaseModel): + model_id = "test" + supported_tools = FunctionCalling.DISABLED + + +def frames(chunks): + model = Plain(engine=ScriptedEngine(chunks)) + return [f for f in model.chat(ChatHistory(), "Hi", tools=[]) if f] + + +def test_think_tags_glued_to_whitespace_become_their_own_frames(): + out = frames(["\n", "The user ", "greets me.", "\n\n\n", "Hello", "!"]) + + assert "" in out + assert "" in out + assert "".join(out) == "\nThe user greets me.\n\n\nHello!" + + +def test_text_and_tag_in_one_chunk_are_split(): + out = frames(["ok", "x", "done"]) + + assert out == ["ok", "", "x", "", "done"] + + +def test_text_without_tags_is_untouched(): + assert frames(["Hello", " world"]) == ["Hello", " world"] diff --git a/docs/SPEC.md b/docs/SPEC.md index 95c82198..e4a24448 100644 --- a/docs/SPEC.md +++ b/docs/SPEC.md @@ -91,7 +91,8 @@ and S3.4 also record the maintainer's decisions of the same day ([`serving/engin The server then sends one text frame per generated chunk, then ``, then closes. An unknown session closes the socket with code `1008`. Reasoning arrives between `` and `` -frames; tool activity arrives as ` … ` frames (S1.6). Generation runs off +frames; tool activity arrives as ` … ` frames (S1.6). Each tag is a frame of its own, +even when the model's streamer emits it glued to whitespace or text (#112). Generation runs off the event loop; if the client disconnects mid-stream, generation stops and the model is free for the next request (#35).