mirror of
https://github.com/mudler/LocalAI.git
synced 2026-10-28 09:08:54 -04:00
pos-ei-don
5392eb813e
With a reasoning parser configured, the streaming path sent every raw text delta as ChatDelta.content. LocalAI uses the backend's ChatDeltas as they are, so with thinking enabled the reasoning arrived as delta.content, and the final chunk then repeated the whole answer and the reasoning a second time. Non-streaming requests were split correctly. Use vLLM's composed parser (ParserManager.get_parser, vLLM >= 0.30), the one vLLM's own OpenAI server streams with: one instance per request, built with the request's chat_template_kwargs (enable_thinking), fed each delta via parse_delta. Reasoning, content and tool calls go out in their own fields as they are produced; the final chunk carries only token counts and logprobs. If the parser raises mid-stream, the rest of the output is buffered and the final chunk sends what the full-text parse has beyond what was already streamed (suffixes per field and per tool-call index). Where the full parse disagrees with what went out, nothing is guessed and the mismatch is logged. Only active when a reasoning_parser is set; without one, or on a vLLM without ParserManager, streaming is unchanged. Signed-off-by: pos-ei-don <1822533+pos-ei-don@users.noreply.github.com>
834 lines
35 KiB
Python
834 lines
35 KiB
Python
import unittest
|
|
import subprocess
|
|
import time
|
|
import backend_pb2
|
|
import backend_pb2_grpc
|
|
|
|
import grpc
|
|
|
|
import unittest
|
|
import subprocess
|
|
import time
|
|
import grpc
|
|
import backend_pb2_grpc
|
|
import backend_pb2
|
|
|
|
class TestBackendServicer(unittest.TestCase):
|
|
"""
|
|
TestBackendServicer is the class that tests the gRPC service.
|
|
|
|
This class contains methods to test the startup and shutdown of the gRPC service.
|
|
"""
|
|
def setUp(self):
|
|
self.service = subprocess.Popen(["python", "backend.py", "--addr", "localhost:50051"])
|
|
time.sleep(10)
|
|
|
|
def tearDown(self) -> None:
|
|
self.service.terminate()
|
|
self.service.wait()
|
|
|
|
def test_server_startup(self):
|
|
try:
|
|
self.setUp()
|
|
with grpc.insecure_channel("localhost:50051") as channel:
|
|
stub = backend_pb2_grpc.BackendStub(channel)
|
|
response = stub.Health(backend_pb2.HealthMessage())
|
|
self.assertEqual(response.message, b'OK')
|
|
except Exception as err:
|
|
print(err)
|
|
self.fail("Server failed to start")
|
|
finally:
|
|
self.tearDown()
|
|
def test_load_model(self):
|
|
"""
|
|
This method tests if the model is loaded successfully
|
|
"""
|
|
try:
|
|
self.setUp()
|
|
with grpc.insecure_channel("localhost:50051") as channel:
|
|
stub = backend_pb2_grpc.BackendStub(channel)
|
|
response = stub.LoadModel(backend_pb2.ModelOptions(Model="facebook/opt-125m"))
|
|
self.assertTrue(response.success)
|
|
self.assertEqual(response.message, "Model loaded successfully")
|
|
except Exception as err:
|
|
print(err)
|
|
self.fail("LoadModel service failed")
|
|
finally:
|
|
self.tearDown()
|
|
|
|
def test_text(self):
|
|
"""
|
|
This method tests if the embeddings are generated successfully
|
|
"""
|
|
try:
|
|
self.setUp()
|
|
with grpc.insecure_channel("localhost:50051") as channel:
|
|
stub = backend_pb2_grpc.BackendStub(channel)
|
|
response = stub.LoadModel(backend_pb2.ModelOptions(Model="facebook/opt-125m"))
|
|
self.assertTrue(response.success)
|
|
req = backend_pb2.PredictOptions(Prompt="The capital of France is")
|
|
resp = stub.Predict(req)
|
|
self.assertIsNotNone(resp.message)
|
|
except Exception as err:
|
|
print(err)
|
|
self.fail("text service failed")
|
|
finally:
|
|
self.tearDown()
|
|
|
|
def test_sampling_params(self):
|
|
"""
|
|
This method tests if all sampling parameters are correctly processed
|
|
NOTE: this does NOT test for correctness, just that we received a compatible response
|
|
"""
|
|
try:
|
|
self.setUp()
|
|
with grpc.insecure_channel("localhost:50051") as channel:
|
|
stub = backend_pb2_grpc.BackendStub(channel)
|
|
response = stub.LoadModel(backend_pb2.ModelOptions(Model="facebook/opt-125m"))
|
|
self.assertTrue(response.success)
|
|
|
|
req = backend_pb2.PredictOptions(
|
|
Prompt="The capital of France is",
|
|
TopP=0.8,
|
|
Tokens=50,
|
|
Temperature=0.7,
|
|
TopK=40,
|
|
PresencePenalty=0.1,
|
|
FrequencyPenalty=0.2,
|
|
RepetitionPenalty=1.1,
|
|
MinP=0.05,
|
|
Seed=42,
|
|
StopPrompts=["\n"],
|
|
StopTokenIds=[50256],
|
|
BadWords=["badword"],
|
|
IncludeStopStrInOutput=True,
|
|
IgnoreEOS=True,
|
|
MinTokens=5,
|
|
Logprobs=5,
|
|
PromptLogprobs=5,
|
|
SkipSpecialTokens=True,
|
|
SpacesBetweenSpecialTokens=True,
|
|
TruncatePromptTokens=10,
|
|
GuidedDecoding=True,
|
|
N=2,
|
|
)
|
|
resp = stub.Predict(req)
|
|
self.assertIsNotNone(resp.message)
|
|
self.assertIsNotNone(resp.logprobs)
|
|
except Exception as err:
|
|
print(err)
|
|
self.fail("sampling params service failed")
|
|
finally:
|
|
self.tearDown()
|
|
|
|
def test_explicit_zero_temperature_and_seed_are_preserved(self):
|
|
"""Temperature=0 is greedy decoding and 0 is a valid seed — neither is
|
|
an unset value. A dropped seed turns a reproducible request random."""
|
|
import sys, os
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from backend import BackendServicer
|
|
|
|
servicer = BackendServicer()
|
|
request = backend_pb2.PredictOptions(Prompt="hello", Temperature=0, Seed=0)
|
|
sampling_params = servicer._build_sampling_params(request)
|
|
self.assertEqual(sampling_params.temperature, 0)
|
|
self.assertEqual(sampling_params.seed, 0)
|
|
# Other protobuf-default scalar fields must remain filtered.
|
|
self.assertEqual(sampling_params.top_p, 0.9)
|
|
|
|
|
|
def test_messages_to_dicts(self):
|
|
"""
|
|
Tests _messages_to_dicts conversion of proto Messages to dicts.
|
|
"""
|
|
import sys, os
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from backend import BackendServicer
|
|
servicer = BackendServicer()
|
|
msgs = [
|
|
backend_pb2.Message(role="user", content="hello"),
|
|
backend_pb2.Message(
|
|
role="assistant",
|
|
content="",
|
|
tool_calls='[{"id":"call_1","type":"function","function":{"name":"foo","arguments":"{}"}}]',
|
|
reasoning_content="thinking...",
|
|
),
|
|
backend_pb2.Message(role="tool", content="result", name="foo", tool_call_id="call_1"),
|
|
]
|
|
result = servicer._messages_to_dicts(msgs)
|
|
self.assertEqual(len(result), 3)
|
|
self.assertEqual(result[0], {"role": "user", "content": "hello"})
|
|
self.assertEqual(result[1]["reasoning_content"], "thinking...")
|
|
self.assertIsInstance(result[1]["tool_calls"], list)
|
|
self.assertEqual(result[1]["tool_calls"][0]["id"], "call_1")
|
|
self.assertEqual(result[2]["tool_call_id"], "call_1")
|
|
self.assertEqual(result[2]["name"], "foo")
|
|
|
|
def test_parse_options(self):
|
|
"""
|
|
Tests _parse_options correctly parses key:value strings.
|
|
"""
|
|
import sys, os
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from backend import BackendServicer
|
|
servicer = BackendServicer()
|
|
opts = servicer._parse_options([
|
|
"tool_parser:hermes",
|
|
"reasoning_parser:deepseek_r1",
|
|
"invalid_no_colon",
|
|
"key_with_colons:a:b:c",
|
|
])
|
|
self.assertEqual(opts["tool_parser"], "hermes")
|
|
self.assertEqual(opts["reasoning_parser"], "deepseek_r1")
|
|
self.assertEqual(opts["key_with_colons"], "a:b:c")
|
|
self.assertNotIn("invalid_no_colon", opts)
|
|
|
|
def test_apply_engine_args_known_keys(self):
|
|
"""
|
|
Tests _apply_engine_args overlays user-supplied JSON onto AsyncEngineArgs.
|
|
"""
|
|
import sys, os, json as _json
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from backend import BackendServicer
|
|
from vllm.engine.arg_utils import AsyncEngineArgs
|
|
|
|
servicer = BackendServicer()
|
|
base = AsyncEngineArgs(model="facebook/opt-125m")
|
|
extras = _json.dumps({
|
|
"trust_remote_code": True,
|
|
"max_num_seqs": 32,
|
|
})
|
|
out = servicer._apply_engine_args(base, extras)
|
|
self.assertTrue(out.trust_remote_code)
|
|
self.assertEqual(out.max_num_seqs, 32)
|
|
# untouched fields preserved
|
|
self.assertEqual(out.model, "facebook/opt-125m")
|
|
|
|
def test_apply_engine_args_unknown_key_raises(self):
|
|
"""
|
|
Tests _apply_engine_args rejects unknown keys with a helpful suggestion.
|
|
"""
|
|
import sys, os, json as _json
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from backend import BackendServicer
|
|
from vllm.engine.arg_utils import AsyncEngineArgs
|
|
|
|
servicer = BackendServicer()
|
|
base = AsyncEngineArgs(model="facebook/opt-125m")
|
|
with self.assertRaises(ValueError) as ctx:
|
|
servicer._apply_engine_args(base, _json.dumps({"trustremotecode": True}))
|
|
self.assertIn("trustremotecode", str(ctx.exception))
|
|
# close-match hint for the typo
|
|
self.assertIn("trust_remote_code", str(ctx.exception))
|
|
|
|
def test_apply_engine_args_empty_passthrough(self):
|
|
"""
|
|
Tests that empty engine_args returns the base unchanged.
|
|
"""
|
|
import sys, os
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from backend import BackendServicer
|
|
from vllm.engine.arg_utils import AsyncEngineArgs
|
|
|
|
servicer = BackendServicer()
|
|
base = AsyncEngineArgs(model="facebook/opt-125m")
|
|
self.assertIs(servicer._apply_engine_args(base, ""), base)
|
|
self.assertIs(servicer._apply_engine_args(base, None), base)
|
|
|
|
def test_tokenize_string(self):
|
|
"""
|
|
Tests the TokenizeString RPC returns valid tokens.
|
|
"""
|
|
try:
|
|
self.setUp()
|
|
with grpc.insecure_channel("localhost:50051") as channel:
|
|
stub = backend_pb2_grpc.BackendStub(channel)
|
|
response = stub.LoadModel(backend_pb2.ModelOptions(Model="facebook/opt-125m"))
|
|
self.assertTrue(response.success)
|
|
resp = stub.TokenizeString(backend_pb2.PredictOptions(Prompt="Hello world"))
|
|
self.assertGreater(resp.length, 0)
|
|
self.assertEqual(len(resp.tokens), resp.length)
|
|
except Exception as err:
|
|
print(err)
|
|
self.fail("TokenizeString service failed")
|
|
finally:
|
|
self.tearDown()
|
|
|
|
def test_free(self):
|
|
"""
|
|
Tests the Free RPC doesn't crash.
|
|
"""
|
|
try:
|
|
self.setUp()
|
|
with grpc.insecure_channel("localhost:50051") as channel:
|
|
stub = backend_pb2_grpc.BackendStub(channel)
|
|
response = stub.LoadModel(backend_pb2.ModelOptions(Model="facebook/opt-125m"))
|
|
self.assertTrue(response.success)
|
|
free_resp = stub.Free(backend_pb2.HealthMessage())
|
|
self.assertTrue(free_resp.success)
|
|
except Exception as err:
|
|
print(err)
|
|
self.fail("Free service failed")
|
|
finally:
|
|
self.tearDown()
|
|
|
|
def test_embedding(self):
|
|
"""
|
|
This method tests if the embeddings are generated successfully
|
|
"""
|
|
try:
|
|
self.setUp()
|
|
with grpc.insecure_channel("localhost:50051") as channel:
|
|
stub = backend_pb2_grpc.BackendStub(channel)
|
|
response = stub.LoadModel(backend_pb2.ModelOptions(Model="intfloat/e5-mistral-7b-instruct"))
|
|
self.assertTrue(response.success)
|
|
embedding_request = backend_pb2.PredictOptions(Embeddings="This is a test sentence.")
|
|
embedding_response = stub.Embedding(embedding_request)
|
|
self.assertIsNotNone(embedding_response.embeddings)
|
|
# assert that is a list of floats
|
|
self.assertIsInstance(embedding_response.embeddings, list)
|
|
# assert that the list is not empty
|
|
self.assertTrue(len(embedding_response.embeddings) > 0)
|
|
except Exception as err:
|
|
print(err)
|
|
self.fail("Embedding service failed")
|
|
finally:
|
|
self.tearDown()
|
|
|
|
|
|
class TestStreamingToolParser(unittest.TestCase):
|
|
"""
|
|
Server-less unit tests for the streaming + tool-parser machinery in
|
|
BackendServicer._predict. These tests instantiate BackendServicer
|
|
directly and mock the vLLM engine + tool parser, so they do not need
|
|
a GPU, a model, or a running gRPC server. Kept in a separate class to
|
|
avoid the parent setUp() which spawns a subprocess.
|
|
|
|
Covers #582 (follow-up to #10346):
|
|
1. Markup-leak prevention with a non-streaming parser (buffer fallback)
|
|
2. No content duplication on the plain-text path with the buffer fallback
|
|
3. Native streaming progressive plain-text emission
|
|
4. Native streaming structured tool_call, no markup leak
|
|
5. Parser exception → graceful fallback to buffer, still no markup
|
|
6. No-tool-parser regression: unchanged per-delta content stream
|
|
"""
|
|
|
|
@staticmethod
|
|
def _make_generate(chunks):
|
|
"""Build a fake vLLM engine.generate that yields cumulative chunks."""
|
|
from types import SimpleNamespace
|
|
async def gen(*a, **k):
|
|
for i, t in enumerate(chunks):
|
|
yield SimpleNamespace(
|
|
outputs=[SimpleNamespace(
|
|
text=t,
|
|
token_ids=list(range(i + 1)),
|
|
logprobs=None,
|
|
)],
|
|
prompt_token_ids=[0],
|
|
)
|
|
return lambda *a, **k: gen()
|
|
|
|
@staticmethod
|
|
def _collect(servicer, req):
|
|
import asyncio
|
|
async def run():
|
|
return [r async for r in servicer._predict(req, None, streaming=True)]
|
|
return asyncio.run(run())
|
|
|
|
def _new_servicer(self):
|
|
import sys, os
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from backend import BackendServicer
|
|
s = BackendServicer()
|
|
s.reasoning_parser_cls = None
|
|
s.tool_parser_cls = None
|
|
s.tokenizer = None
|
|
return s
|
|
|
|
# ── Case 1+2: parser without streaming method → buffer fallback ──
|
|
def test_buffer_path_no_markup_no_duplication(self):
|
|
from types import SimpleNamespace
|
|
|
|
def parser_cls(called, content_text, calls):
|
|
class _P:
|
|
def __init__(self, tokenizer, tools=None):
|
|
pass
|
|
# NOTE: NO extract_tool_calls_streaming → takes the buffer path
|
|
def extract_tool_calls(self, c, request=None):
|
|
return SimpleNamespace(
|
|
tools_called=called, content=content_text, tool_calls=calls,
|
|
)
|
|
return _P
|
|
|
|
tools_json = '[{"type":"function","function":{"name":"calc","parameters":{}}}]'
|
|
|
|
# Tool-call case: no raw markup in any delta.content
|
|
s = self._new_servicer()
|
|
s.llm = SimpleNamespace(generate=self._make_generate([
|
|
'<tool_call>\n{"name": "calc"',
|
|
'<tool_call>\n{"name": "calc", "arguments": {"x": 1}}\n</tool_call>',
|
|
]))
|
|
call = SimpleNamespace(id="call_1",
|
|
function=SimpleNamespace(name="calc", arguments='{"x": 1}'))
|
|
s.tool_parser_cls = parser_cls(True, "", [call])
|
|
req = backend_pb2.PredictOptions(Prompt="x", Tools=tools_json)
|
|
replies = self._collect(s, req)
|
|
contents = [cd.content for r in replies for cd in r.chat_deltas if cd.content]
|
|
self.assertFalse(
|
|
any("<tool_call" in c for c in contents),
|
|
f"markup leaked: {contents!r}",
|
|
)
|
|
names = [tc.name for r in replies for cd in r.chat_deltas for tc in cd.tool_calls]
|
|
self.assertIn("calc", names, "tool_call missing from final chunk")
|
|
|
|
# Plain-text-with-tools case: full content delivered exactly once
|
|
s2 = self._new_servicer()
|
|
s2.llm = SimpleNamespace(generate=self._make_generate([
|
|
"The capital ",
|
|
"The capital of France is Paris.",
|
|
]))
|
|
s2.tool_parser_cls = parser_cls(False, "", [])
|
|
req2 = backend_pb2.PredictOptions(Prompt="x", Tools=tools_json)
|
|
joined = "".join(
|
|
cd.content for r in self._collect(s2, req2)
|
|
for cd in r.chat_deltas if cd.content
|
|
)
|
|
self.assertEqual(
|
|
joined.count("The capital of France is Paris."), 1,
|
|
f"buffered content duplicated: {joined!r}",
|
|
)
|
|
|
|
# ── Case 3: native streaming, progressive plain text ──
|
|
def test_native_streaming_progressive_plain_text(self):
|
|
from types import SimpleNamespace
|
|
|
|
class _DeltaMsg:
|
|
def __init__(self, content=None, reasoning=None, tool_calls=None):
|
|
self.content = content
|
|
self.reasoning = reasoning
|
|
self.tool_calls = tool_calls or []
|
|
|
|
class StreamingParser:
|
|
def __init__(self, tokenizer, tools=None):
|
|
pass
|
|
def extract_tool_calls(self, c, request=None):
|
|
# Should NOT be called when native streaming runs successfully.
|
|
raise AssertionError("extract_tool_calls invoked on native-streaming path")
|
|
def extract_tool_calls_streaming(
|
|
self, previous_text, current_text, delta_text,
|
|
previous_token_ids, current_token_ids, delta_token_ids, request,
|
|
):
|
|
if not delta_text:
|
|
return None
|
|
return _DeltaMsg(content=delta_text)
|
|
|
|
s = self._new_servicer()
|
|
s.llm = SimpleNamespace(generate=self._make_generate([
|
|
"Paris ",
|
|
"Paris is ",
|
|
"Paris is the capital of France.",
|
|
]))
|
|
s.tool_parser_cls = StreamingParser
|
|
req = backend_pb2.PredictOptions(
|
|
Prompt="x",
|
|
Tools='[{"type":"function","function":{"name":"calc","parameters":{}}}]',
|
|
)
|
|
replies = self._collect(s, req)
|
|
|
|
intermediate_content = [
|
|
cd.content for r in replies[:-1] for cd in r.chat_deltas if cd.content
|
|
]
|
|
self.assertTrue(
|
|
len(intermediate_content) > 0,
|
|
"Plain-text response not streamed progressively (native streaming inactive?)",
|
|
)
|
|
assembled = "".join(
|
|
cd.content for r in replies for cd in r.chat_deltas if cd.content
|
|
)
|
|
self.assertEqual(
|
|
assembled, "Paris is the capital of France.",
|
|
f"Assembled content wrong: {assembled!r}",
|
|
)
|
|
|
|
# ── Case 4: native streaming, structured tool_call, no markup ──
|
|
def test_native_streaming_tool_call_no_markup_leak(self):
|
|
from types import SimpleNamespace
|
|
|
|
class _DeltaMsg:
|
|
def __init__(self, content=None, reasoning=None, tool_calls=None):
|
|
self.content = content
|
|
self.reasoning = reasoning
|
|
self.tool_calls = tool_calls or []
|
|
|
|
class _ToolCallStreamer:
|
|
def __init__(self, tokenizer, tools=None):
|
|
self._emitted = False
|
|
def extract_tool_calls(self, c, request=None):
|
|
raise AssertionError("extract_tool_calls invoked on native-streaming path")
|
|
def extract_tool_calls_streaming(
|
|
self, previous_text, current_text, delta_text,
|
|
previous_token_ids, current_token_ids, delta_token_ids, request,
|
|
):
|
|
if "</tool_call>" in current_text and not self._emitted:
|
|
self._emitted = True
|
|
fn = SimpleNamespace(name="calc", arguments='{"x": 1}')
|
|
tc = SimpleNamespace(id="call_1", type="function", index=0, function=fn)
|
|
return _DeltaMsg(tool_calls=[tc])
|
|
return None
|
|
|
|
s = self._new_servicer()
|
|
s.llm = SimpleNamespace(generate=self._make_generate([
|
|
'<tool_call>\n',
|
|
'<tool_call>\n{"name": "calc"',
|
|
'<tool_call>\n{"name": "calc", "arguments": {"x": 1}}\n</tool_call>',
|
|
]))
|
|
s.tool_parser_cls = _ToolCallStreamer
|
|
req = backend_pb2.PredictOptions(
|
|
Prompt="x",
|
|
Tools='[{"type":"function","function":{"name":"calc","parameters":{}}}]',
|
|
)
|
|
replies = self._collect(s, req)
|
|
|
|
contents = [cd.content for r in replies for cd in r.chat_deltas if cd.content]
|
|
self.assertFalse(
|
|
any("<tool_call" in c or "</tool_call>" in c for c in contents),
|
|
f"markup leaked as content: {contents!r}",
|
|
)
|
|
names = [tc.name for r in replies for cd in r.chat_deltas for tc in cd.tool_calls if tc.name]
|
|
args = [tc.arguments for r in replies for cd in r.chat_deltas for tc in cd.tool_calls if tc.arguments]
|
|
self.assertIn("calc", names, f"tool_call name missing; got {names!r}")
|
|
self.assertIn('{"x": 1}', args, f"tool_call args missing; got {args!r}")
|
|
|
|
# ── Case 5: parser exception → fallback to buffer, no leak ──
|
|
def test_native_streaming_parser_exception_falls_back_to_buffer(self):
|
|
from types import SimpleNamespace
|
|
call = SimpleNamespace(id="call_1",
|
|
function=SimpleNamespace(name="calc", arguments='{"x": 1}'))
|
|
|
|
class _BrokenStreamer:
|
|
def __init__(self, tokenizer, tools=None):
|
|
pass
|
|
def extract_tool_calls(self, c, request=None):
|
|
return SimpleNamespace(tools_called=True, content="", tool_calls=[call])
|
|
def extract_tool_calls_streaming(self, *a, **kw):
|
|
raise RuntimeError("simulated parser bug")
|
|
|
|
s = self._new_servicer()
|
|
s.llm = SimpleNamespace(generate=self._make_generate([
|
|
'<tool_call>\n{"name": "calc"',
|
|
'<tool_call>\n{"name": "calc", "arguments": {"x": 1}}\n</tool_call>',
|
|
]))
|
|
s.tool_parser_cls = _BrokenStreamer
|
|
req = backend_pb2.PredictOptions(
|
|
Prompt="x",
|
|
Tools='[{"type":"function","function":{"name":"calc","parameters":{}}}]',
|
|
)
|
|
replies = self._collect(s, req)
|
|
|
|
contents = [cd.content for r in replies for cd in r.chat_deltas if cd.content]
|
|
self.assertFalse(
|
|
any("<tool_call" in c for c in contents),
|
|
f"markup leaked after parser exception: {contents!r}",
|
|
)
|
|
names = [tc.name for r in replies for cd in r.chat_deltas for tc in cd.tool_calls]
|
|
self.assertIn("calc", names, "tool_call missing from final chunk after fallback")
|
|
|
|
# ── Case 6: no tool parser → unchanged per-delta content stream ──
|
|
def test_no_tool_parser_unchanged_per_delta_stream(self):
|
|
from types import SimpleNamespace
|
|
s = self._new_servicer() # tool_parser_cls already None
|
|
s.llm = SimpleNamespace(generate=self._make_generate([
|
|
"Hello ", "Hello world", "Hello world!",
|
|
]))
|
|
req = backend_pb2.PredictOptions(Prompt="x", Tools="")
|
|
replies = self._collect(s, req)
|
|
|
|
intermediate = [
|
|
cd.content for r in replies[:-1] for cd in r.chat_deltas if cd.content
|
|
]
|
|
self.assertEqual(
|
|
intermediate, ["Hello ", "world", "!"],
|
|
f"plain streaming changed; got {intermediate!r}",
|
|
)
|
|
|
|
|
|
class TestReasoningSplit(unittest.TestCase):
|
|
"""Server-less tests for BackendServicer._split_reasoning.
|
|
|
|
vLLM's BaseThinkingReasoningParser returns the whole completion as
|
|
reasoning and None as content whenever the end token is missing. Taken
|
|
literally that duplicates a thinking-disabled answer into both fields.
|
|
"""
|
|
|
|
class _Parser:
|
|
start_token = "<think>"
|
|
end_token = "</think>"
|
|
|
|
def _split(self, generated, prompt, reasoning, content):
|
|
import sys, os
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from backend import BackendServicer
|
|
return BackendServicer._split_reasoning(
|
|
self._Parser(), generated, prompt, reasoning, content,
|
|
)
|
|
|
|
def test_thinking_off_is_not_duplicated_into_reasoning(self):
|
|
"""No tags anywhere: the answer is content, and only content."""
|
|
r, c = self._split(
|
|
"391", "user: 17*23?\n<think>\n\n</think>\n\n",
|
|
reasoning="391", content=None,
|
|
)
|
|
self.assertEqual(r, "")
|
|
self.assertEqual(c, "391")
|
|
|
|
def test_prefilled_start_tag_keeps_truncated_reasoning(self):
|
|
"""Prompt left the block open and the end token never arrived
|
|
(budget exhausted): that really is all reasoning."""
|
|
r, c = self._split(
|
|
"thinking and thinking", "user: hi\n<think>\n",
|
|
reasoning="thinking and thinking", content=None,
|
|
)
|
|
self.assertEqual(r, "thinking and thinking")
|
|
self.assertEqual(c, "")
|
|
|
|
def test_end_token_present_keeps_parser_split(self):
|
|
r, c = self._split(
|
|
"adding two and two</think>4", "user: hi\n<think>\n",
|
|
reasoning="adding two and two", content="4",
|
|
)
|
|
self.assertEqual(r, "adding two and two")
|
|
self.assertEqual(c, "4")
|
|
|
|
def test_stop_right_after_end_token_yields_empty_content(self):
|
|
"""Content must not fall back to the raw text — that would put the
|
|
reasoning into the answer."""
|
|
r, c = self._split(
|
|
"reasoned</think>", "user: hi\n<think>\n",
|
|
reasoning="reasoned", content=None,
|
|
)
|
|
self.assertEqual(r, "reasoned")
|
|
self.assertEqual(c, "")
|
|
|
|
def test_unknown_token_layout_keeps_previous_behaviour(self):
|
|
class _Bare:
|
|
pass
|
|
import sys, os
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from backend import BackendServicer
|
|
r, c = BackendServicer._split_reasoning(
|
|
_Bare(), "raw", "prompt", "raw", None,
|
|
)
|
|
self.assertEqual(r, "raw")
|
|
self.assertEqual(c, "raw")
|
|
|
|
|
|
class TestReasoningParserConstruction(unittest.TestCase):
|
|
"""The parser must learn whether thinking was on for this request.
|
|
|
|
vLLM's engine-based parsers (Qwen3Parser and friends) read
|
|
chat_template_kwargs["enable_thinking"] and default to True, so a parser
|
|
built without it treats a thinking-disabled completion as pure reasoning.
|
|
"""
|
|
|
|
def _servicer(self):
|
|
import sys, os
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from backend import BackendServicer
|
|
s = BackendServicer()
|
|
s.tokenizer = object()
|
|
return s
|
|
|
|
def test_chat_template_kwargs_are_forwarded(self):
|
|
seen = {}
|
|
|
|
class _Parser:
|
|
def __init__(self, tokenizer, **kwargs):
|
|
seen.update(kwargs)
|
|
|
|
s = self._servicer()
|
|
s.reasoning_parser_cls = _Parser
|
|
s._new_reasoning_parser({"enable_thinking": False})
|
|
self.assertEqual(
|
|
seen.get("chat_template_kwargs"), {"enable_thinking": False},
|
|
)
|
|
|
|
def test_parser_without_the_kwarg_still_builds(self):
|
|
"""Older parsers take only the tokenizer — must not break them."""
|
|
class _Old:
|
|
def __init__(self, tokenizer):
|
|
self.tokenizer = tokenizer
|
|
|
|
s = self._servicer()
|
|
s.reasoning_parser_cls = _Old
|
|
self.assertIsInstance(
|
|
s._new_reasoning_parser({"enable_thinking": False}), _Old,
|
|
)
|
|
|
|
|
|
class TestStreamingReasoningSplit(unittest.TestCase):
|
|
"""With a reasoning parser, streaming must split reasoning from content.
|
|
|
|
The stream used to send each raw text delta as ChatDelta.content, so with
|
|
thinking on the reasoning text arrived as delta.content and the final
|
|
chunk repeated it as reasoning_content. These tests replace the engine
|
|
and vLLM's composed parser (ParserManager.get_parser) with fakes; no GPU,
|
|
model or server needed.
|
|
"""
|
|
|
|
THINK = ["Let me ", "count the ", "primes.", "</think>"]
|
|
ANSWER = ["There are ", "37 primes."]
|
|
|
|
class _DeltaMsg:
|
|
def __init__(self, content=None, reasoning=None, tool_calls=None):
|
|
self.content = content
|
|
self.reasoning = reasoning
|
|
self.tool_calls = tool_calls or []
|
|
|
|
@classmethod
|
|
def _parser_cls(cls, seen=None, fail_after=None):
|
|
delta_msg = cls._DeltaMsg
|
|
|
|
class _Parser:
|
|
"""Stands in for vLLM's DelegatingParser: reasoning until the
|
|
end tag, then content, then a tool call between <tool_call> tags
|
|
(name first, arguments as they arrive)."""
|
|
def __init__(self, tokenizer, tools=None, **kwargs):
|
|
if seen is not None:
|
|
seen.update(kwargs, tools=tools)
|
|
self._state = "reasoning"
|
|
self._calls = 0
|
|
|
|
def parse_delta(self, delta_text, delta_token_ids, request,
|
|
prompt_token_ids=None, *, finished):
|
|
from types import SimpleNamespace
|
|
self._calls += 1
|
|
if fail_after is not None and self._calls > fail_after:
|
|
raise RuntimeError("simulated parser bug")
|
|
if self._state == "reasoning":
|
|
if "</think>" in delta_text:
|
|
self._state = "content"
|
|
before = delta_text.split("</think>", 1)[0]
|
|
return delta_msg(reasoning=before) if before else None
|
|
return delta_msg(reasoning=delta_text)
|
|
if delta_text == "<tool_call>":
|
|
self._state = "tool"
|
|
fn = SimpleNamespace(name="search", arguments="")
|
|
return delta_msg(tool_calls=[SimpleNamespace(index=0, id="call_1", function=fn)])
|
|
if self._state == "tool":
|
|
if delta_text == "</tool_call>":
|
|
return None
|
|
fn = SimpleNamespace(name=None, arguments=delta_text)
|
|
return delta_msg(tool_calls=[SimpleNamespace(index=0, id=None, function=fn)])
|
|
return delta_msg(content=delta_text)
|
|
return _Parser
|
|
|
|
class _ToolParser:
|
|
"""Full-text tool extraction used after a streaming parser failure."""
|
|
def __init__(self, tokenizer, tools=None):
|
|
pass
|
|
|
|
def extract_tool_calls(self, text, request=None):
|
|
from types import SimpleNamespace
|
|
if "<tool_call>" not in text:
|
|
return SimpleNamespace(tools_called=False, content=text, tool_calls=[])
|
|
before, rest = text.split("<tool_call>", 1)
|
|
args = rest.split("</tool_call>", 1)[0]
|
|
call = SimpleNamespace(id="call_1", function=SimpleNamespace(name="search", arguments=args))
|
|
return SimpleNamespace(tools_called=True, content=before, tool_calls=[call])
|
|
|
|
class _ReasoningParser:
|
|
"""Full-text split used after the loop (BaseThinkingReasoningParser shape)."""
|
|
start_token = "<think>"
|
|
end_token = "</think>"
|
|
|
|
def __init__(self, tokenizer, **kwargs):
|
|
pass
|
|
|
|
def extract_reasoning(self, text, request=None):
|
|
if "</think>" in text:
|
|
r, c = text.split("</think>", 1)
|
|
return r, c
|
|
return text, None
|
|
|
|
def _servicer(self, stream_parser_cls, pieces=None):
|
|
import sys, os
|
|
from types import SimpleNamespace
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from backend import BackendServicer
|
|
s = BackendServicer()
|
|
s.tokenizer = None
|
|
s.tool_parser_cls = self._ToolParser
|
|
s.reasoning_parser_cls = self._ReasoningParser
|
|
s.stream_parser_cls = stream_parser_cls
|
|
s.llm = SimpleNamespace(generate=TestStreamingToolParser._make_generate(
|
|
self._cumulative(pieces or (self.THINK + self.ANSWER))))
|
|
return s
|
|
|
|
@staticmethod
|
|
def _cumulative(pieces):
|
|
out, acc = [], ""
|
|
for p in pieces:
|
|
acc += p
|
|
out.append(acc)
|
|
return out
|
|
|
|
@staticmethod
|
|
def _joined(replies, field):
|
|
return "".join(getattr(cd, field) for r in replies for cd in r.chat_deltas)
|
|
|
|
def test_reasoning_is_not_streamed_as_content(self):
|
|
s = self._servicer(self._parser_cls())
|
|
replies = TestStreamingToolParser._collect(s, backend_pb2.PredictOptions(Prompt="x"))
|
|
self.assertEqual(self._joined(replies, "content"), "There are 37 primes.")
|
|
self.assertEqual(self._joined(replies, "reasoning_content"), "Let me count the primes.")
|
|
# Progressive: reasoning and content arrive before the final chunk.
|
|
self.assertTrue(any(cd.reasoning_content for r in replies[:-1] for cd in r.chat_deltas))
|
|
self.assertTrue(any(cd.content for r in replies[:-1] for cd in r.chat_deltas))
|
|
# The final chunk carries no text, otherwise LocalAI repeats it.
|
|
final = replies[-1].chat_deltas[0]
|
|
self.assertEqual((final.content, final.reasoning_content), ("", ""))
|
|
|
|
def test_parser_error_mid_stream_sends_the_rest_once(self):
|
|
s = self._servicer(self._parser_cls(fail_after=2))
|
|
replies = TestStreamingToolParser._collect(s, backend_pb2.PredictOptions(Prompt="x"))
|
|
self.assertEqual(self._joined(replies, "content"), "There are 37 primes.")
|
|
self.assertEqual(self._joined(replies, "reasoning_content"), "Let me count the primes.")
|
|
|
|
def test_without_stream_parser_streams_raw_deltas(self):
|
|
"""vLLM without ParserManager: previous behaviour."""
|
|
s = self._servicer(None)
|
|
replies = TestStreamingToolParser._collect(s, backend_pb2.PredictOptions(Prompt="x"))
|
|
intermediate = "".join(cd.content for r in replies[:-1] for cd in r.chat_deltas)
|
|
self.assertEqual(intermediate, "".join(self.THINK + self.ANSWER))
|
|
|
|
def test_parser_gets_chat_template_kwargs_and_tools(self):
|
|
seen = {}
|
|
s = self._servicer(self._parser_cls(seen=seen))
|
|
tools = '[{"type":"function","function":{"name":"calc","parameters":{}}}]'
|
|
parser, parser_request = s._new_stream_parser(
|
|
backend_pb2.PredictOptions(Prompt="x", Tools=tools),
|
|
{"tokenize": False, "add_generation_prompt": True, "enable_thinking": False,
|
|
"tools": [{}]},
|
|
)
|
|
self.assertIsNotNone(parser)
|
|
self.assertEqual(seen["chat_template_kwargs"], {"enable_thinking": False})
|
|
self.assertEqual([t.function.name for t in seen["tools"]], ["calc"])
|
|
|
|
def test_parser_error_after_partial_tool_arguments(self):
|
|
"""The rest of the arguments follows on the same index; no raw
|
|
tool-call markup reaches content."""
|
|
pieces = self.THINK + ["Checking.", "<tool_call>", '{"q": ', '"primes"}', "</tool_call>"]
|
|
tools = '[{"type":"function","function":{"name":"search","parameters":{}}}]'
|
|
for fail_after in (6, 7): # after the name / after part of the arguments
|
|
with self.subTest(fail_after=fail_after):
|
|
s = self._servicer(self._parser_cls(fail_after=fail_after), pieces)
|
|
replies = TestStreamingToolParser._collect(
|
|
s, backend_pb2.PredictOptions(Prompt="x", Tools=tools))
|
|
self.assertEqual(self._joined(replies, "content"), "Checking.")
|
|
self.assertEqual(self._joined(replies, "reasoning_content"), "Let me count the primes.")
|
|
calls = [tc for r in replies for cd in r.chat_deltas for tc in cd.tool_calls]
|
|
self.assertEqual({tc.index for tc in calls}, {0})
|
|
self.assertEqual([tc.name for tc in calls if tc.name], ["search"])
|
|
self.assertEqual("".join(tc.arguments for tc in calls), '{"q": "primes"}')
|