[UPDATE] Update

[ghstack-poisoned]

Author: mergennachin

extension/llm/runner/llm_session.h

@@ -8,8 +8,8 @@
 
 // Model-agnostic Engine/Session interfaces. Model-specific execution lives in
 // adapters that implement these (TextLLMSession over TextLLMRunner today;
-// Gemma4Session etc. later); the server and pybind layer depend only on these
-// interfaces, never on a concrete runner.
+// Gemma4Session etc. later); the serving code (HTTP control plane + C++ worker
+// binaries) depends only on these interfaces, never on a concrete runner.
 
 #pragma once
 

extension/llm/runner/test/test_util.cpp

@@ -18,6 +18,7 @@ namespace {
 using ::executorch::aten::ScalarType;
 using ::executorch::extension::make_tensor_ptr;
 using ::executorch::extension::llm::convert_to_bfloat16;
+using ::executorch::extension::llm::stop_safe_prefix_len;
 using ::executorch::extension::llm::utf8_complete_prefix_len;
 
 class ConvertToBFloat16Test : public ::testing::Test {
@@ -86,4 +87,39 @@ TEST(Utf8CompletePrefixLenTest, HandlesAsciiAndMultiByteBoundaries) {
   EXPECT_EQ(utf8_complete_prefix_len("\x80"), 1u);
 }
 
+TEST(StopSafePrefixLenTest, NoStopsEmitsEverything) {
+  bool hit = true;
+  EXPECT_EQ(stop_safe_prefix_len("hello world", {}, hit), 11u);
+  EXPECT_FALSE(hit);
+}
+
+TEST(StopSafePrefixLenTest, StopFoundReturnsEarliestOffsetAndExcludesIt) {
+  bool hit = false;
+  // "STOP" begins at offset 6; emit "Hello " (6 bytes), drop the stop and rest.
+  EXPECT_EQ(stop_safe_prefix_len("Hello STOP there", {"STOP"}, hit), 6u);
+  EXPECT_TRUE(hit);
+  // Earliest of several wins.
+  hit = false;
+  EXPECT_EQ(stop_safe_prefix_len("aXbY", {"Y", "X"}, hit), 1u);
+  EXPECT_TRUE(hit);
+}
+
+TEST(StopSafePrefixLenTest, HoldsBackPossiblePartialStopTail) {
+  bool hit = false;
+  // No full stop yet, but the trailing "ST" could become "STOP": hold back
+  // len("STOP")-1 == 3 bytes, so of "hi ST" (5 bytes) only "hi" (2) is safe.
+  EXPECT_EQ(stop_safe_prefix_len("hi ST", {"STOP"}, hit), 2u);
+  EXPECT_FALSE(hit);
+}
+
+TEST(StopSafePrefixLenTest, HoldBackSnapsToUtf8Boundary) {
+  bool hit = false;
+  // "ab" + "€"(3 bytes). Stop "XX" => hold back 1 byte, which would land inside
+  // the euro sign; snap down so the multi-byte char isn't split.
+  const std::string text = "ab\xe2\x82\xac";
+  const size_t safe = stop_safe_prefix_len(text, {"XX"}, hit);
+  EXPECT_FALSE(hit);
+  EXPECT_EQ(safe, 2u); // only "ab"; the € is held whole
+}
+
 } // namespace

extension/llm/runner/util.h

@@ -13,6 +13,7 @@
 #include <executorch/runtime/platform/compiler.h>
 #include <stdio.h>
 #include <time.h>
+#include <algorithm>
 #include <cctype>
 #include <string>
 #include <vector>
@@ -98,6 +99,55 @@ ET_EXPERIMENTAL size_t inline utf8_complete_prefix_len(const std::string& s) {
   return i;
 }
 
+// How many leading bytes of `text` a streaming consumer may safely emit given a
+// set of `stops` strings, and whether a stop was hit (`stop_hit`).
+//   * If any stop occurs, returns the byte offset of the EARLIEST occurrence and
+//     sets stop_hit=true — text before it is safe; the stop and everything after
+//     are dropped (the stop is excluded from output).
+//   * Otherwise returns the length minus the longest possible partial-stop tail
+//     (max(len(stop))-1 bytes), snapped DOWN to a UTF-8 boundary so a multi-byte
+//     character is never split; stop_hit=false. Holding back that tail lets a
+//     stop that straddles the next piece still be caught.
+// `text` is expected to be complete-UTF-8 (e.g. the assembled output of
+// utf8_complete_prefix_len). Empty `stops` => emit everything, no hold-back.
+ET_EXPERIMENTAL size_t inline stop_safe_prefix_len(
+    const std::string& text,
+    const std::vector<std::string>& stops,
+    bool& stop_hit) {
+  stop_hit = false;
+  if (stops.empty()) {
+    return text.size();
+  }
+  size_t earliest = std::string::npos;
+  size_t max_len = 0;
+  for (const auto& s : stops) {
+    if (s.empty()) {
+      continue;
+    }
+    max_len = std::max(max_len, s.size());
+    const size_t p = text.find(s);
+    if (p != std::string::npos &&
+        (earliest == std::string::npos || p < earliest)) {
+      earliest = p;
+    }
+  }
+  if (earliest != std::string::npos) {
+    stop_hit = true;
+    return earliest;
+  }
+  const size_t hold = max_len > 0 ? max_len - 1 : 0;
+  if (text.size() <= hold) {
+    return 0;
+  }
+  size_t end = text.size() - hold;
+  // Don't cut in the middle of a UTF-8 character: back up over continuation
+  // bytes (10xxxxxx).
+  while (end > 0 && (static_cast<unsigned char>(text[end]) & 0xC0) == 0x80) {
+    --end;
+  }
+  return end;
+}
+
 // ----------------------------------------------------------------------------
 // utilities: time
 

extension/llm/server/python/errors.py

@@ -30,13 +30,26 @@ def body(self) -> dict:
 
 
 class ContextLengthExceeded(APIError):
-    def __init__(self, num_tokens: int, max_context: int):
-        super().__init__(
-            status=400,
-            message=(
+    def __init__(self, num_tokens: int, max_context: int, completion_tokens: int = 0):
+        # completion_tokens > 0: the prompt fits but prompt + requested
+        # max_tokens would run past the window — reject up front rather than
+        # fail (or truncate) mid-generation.
+        if completion_tokens > 0:
+            message = (
+                f"This model's maximum context length is {max_context} tokens. "
+                f"However, you requested {num_tokens + completion_tokens} tokens "
+                f"({num_tokens} in the messages, {completion_tokens} in the "
+                f"completion). Please reduce the length of the messages or "
+                f"completion."
+            )
+        else:
+            message = (
                 f"This model's maximum context length is {max_context} tokens, "
                 f"but the request has {num_tokens} prompt tokens."
-            ),
+            )
+        super().__init__(
+            status=400,
+            message=message,
             err_type="invalid_request_error",
             code="context_length_exceeded",
         )

extension/llm/server/python/tests/test_qwen_tool_parser.py

@@ -10,7 +10,14 @@
 
 from executorch.extension.llm.server.python.tool_parsers import QwenFunctionCallDetector
 
-_TOOLS = {"get_weather", "add"}
+# name -> JSON-schema `parameters` (as the server passes it to the detector).
+_TOOLS = {
+    "get_weather": {"type": "object", "properties": {"city": {"type": "string"}}},
+    "add": {
+        "type": "object",
+        "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}},
+    },
+}
 
 
 def _parse(text, tools=_TOOLS):
@@ -87,3 +94,33 @@ def test_missing_tool_call_wrapper_still_parses():
     r = _parse(text)
     assert len(r.calls) == 1
     assert json.loads(r.calls[0].arguments) == {"city": "Paris"}
+
+
+# Schema-aware coercion: the XML format is stringly-typed, so values must be cast
+# to the declared schema type (the cause of several BFCL function-calling misses).
+def test_boolean_value_coerced_by_schema():
+    tools = {"f": {"properties": {"flag": {"type": "boolean"}}}}
+    # The model writes a non-JSON capitalized "True"; the schema says boolean.
+    text = "<function=f><parameter=flag>True</parameter></function>"
+    r = _parse(text, tools)
+    assert json.loads(r.calls[0].arguments) == {"flag": True}
+
+
+def test_string_schema_keeps_numeric_literal_as_string():
+    tools = {"f": {"properties": {"id": {"type": "string"}}}}
+    # A numeric-looking value the schema declares as a string must NOT become int.
+    text = "<function=f><parameter=id>1234</parameter></function>"
+    r = _parse(text, tools)
+    args = json.loads(r.calls[0].arguments)
+    assert args == {"id": "1234"} and isinstance(args["id"], str)
+
+
+def test_untyped_param_falls_back_to_json_guess():
+    # No declared type -> best-effort JSON guess (so loosely-typed tools still work).
+    tools = {"f": {"properties": {}}}
+    text = (
+        "<function=f><parameter=n>42</parameter>"
+        "<parameter=items>[1, 2]</parameter></function>"
+    )
+    r = _parse(text, tools)
+    assert json.loads(r.calls[0].arguments) == {"n": 42, "items": [1, 2]}

extension/llm/server/python/tool_parsers/hermes.py

@@ -40,14 +40,14 @@ class HermesDetector:
     def __init__(self):
         self._next_index = 0
 
-    def detect_and_parse(self, text: str, tool_names: set[str]) -> ParseResult:
+    def detect_and_parse(self, text: str, tools: dict[str, dict]) -> ParseResult:
         """Return leading text + any complete tool calls. On no call or a parse
         failure, return the original text unchanged (kept visible to the client)."""
         if self.bot_token not in text:
             return ParseResult(normal_text=text)
         normal = text[: text.find(self.bot_token)].strip()
         try:
-            calls = self._parse_calls(text, tool_names)
+            calls = self._parse_calls(text, tools)
         except _UndefinedToolCall as e:
             # Degrade the whole response to visible text so the undefined call
             # isn't silently dropped (and its valid siblings aren't executed in
@@ -61,7 +61,7 @@ def detect_and_parse(self, text: str, tool_names: set[str]) -> ParseResult:
             return ParseResult(normal_text=text)
         return ParseResult(normal_text=normal, calls=calls)
 
-    def _parse_calls(self, text: str, tool_names: set[str]) -> list[ToolCallItem]:
+    def _parse_calls(self, text: str, tools: dict[str, dict]) -> list[ToolCallItem]:
         calls = []
         for raw in _CALL_RE.findall(text):
             if not raw.strip():
@@ -72,15 +72,15 @@ def _parse_calls(self, text: str, tool_names: set[str]) -> list[ToolCallItem]:
                     self._make_item(
                         entry.get("name"),
                         entry.get("arguments", entry.get("parameters")),
-                        tool_names,
+                        tools,
                     )
                 )
         return calls
 
     def _make_item(
-        self, name: Optional[str], arguments: Any, tool_names: set[str]
+        self, name: Optional[str], arguments: Any, tools: dict[str, dict]
     ) -> ToolCallItem:
-        if not name or name not in tool_names:
+        if not name or name not in tools:
             raise _UndefinedToolCall(repr(name))
         item = ToolCallItem(
             tool_index=self._next_index,

extension/llm/server/python/tool_parsers/qwen.py

@@ -41,9 +41,35 @@ class _UndefinedToolCall(Exception):
     response to visible text rather than emitting a partial set (spec)."""
 
 
-def _coerce(value: str) -> Any:
-    """Parameter values are raw text; try JSON so numbers/bools/objects get their
-    natural type, otherwise keep the string."""
+def _coerce(value: str, declared_type: Optional[str]) -> Any:
+    """Cast a raw XML parameter string to the type declared in the tool's JSON
+    schema.
+
+    The Qwen XML format is stringly-typed (`<parameter=k>v</parameter>`), so
+    without the schema we'd have to guess. A bare `json.loads` guess mistypes two
+    common ways: a value the schema wants as a string but that looks numeric
+    (`"1234"`) becomes an int, and a value the schema wants as a bool but that the
+    model didn't write as valid JSON (`True`) stays a string. Coercing to the
+    declared type keeps the emitted OpenAI tool_call schema-valid. Falls back to a
+    JSON guess (then the raw string) when the type is unknown or coercion fails,
+    so untyped/loosely-typed params keep working.
+    """
+    if declared_type == "string":
+        return value
+    if declared_type == "boolean":
+        low = value.strip().lower()
+        if low in ("true", "false"):
+            return low == "true"
+    elif declared_type == "integer":
+        try:
+            return int(value.strip())
+        except (ValueError, TypeError):
+            pass
+    elif declared_type == "number":
+        try:
+            return float(value.strip())
+        except (ValueError, TypeError):
+            pass
     try:
         return json.loads(value)
     except (ValueError, TypeError):
@@ -59,9 +85,13 @@ class QwenFunctionCallDetector:
     def __init__(self):
         self._next_index = 0
 
-    def detect_and_parse(self, text: str, tool_names: set[str]) -> ParseResult:
+    def detect_and_parse(self, text: str, tools: dict[str, dict]) -> ParseResult:
         """Return leading text + any complete tool calls. On no call or a parse
-        failure, return the original text unchanged (kept visible to the client)."""
+        failure, return the original text unchanged (kept visible to the client).
+
+        `tools` maps each defined tool name to its JSON-schema ``parameters``
+        object; the schema is used to coerce stringly-typed XML values to their
+        declared types (and the key set validates names)."""
         first = _FUNCTION_RE.search(text)
         if first is None:
             return ParseResult(normal_text=text)
@@ -72,7 +102,7 @@ def detect_and_parse(self, text: str, tool_names: set[str]) -> ParseResult:
             cut = first.start()
         normal = text[:cut].strip()
         try:
-            calls = self._parse_calls(text, tool_names)
+            calls = self._parse_calls(text, tools)
         except _UndefinedToolCall as e:
             logger.debug("undefined tool %s; returning raw text (no partial calls)", e)
             return ParseResult(normal_text=text)
@@ -83,20 +113,21 @@ def detect_and_parse(self, text: str, tool_names: set[str]) -> ParseResult:
             return ParseResult(normal_text=text)
         return ParseResult(normal_text=normal, calls=calls)
 
-    def _parse_calls(self, text: str, tool_names: set[str]) -> list[ToolCallItem]:
+    def _parse_calls(self, text: str, tools: dict[str, dict]) -> list[ToolCallItem]:
         calls = []
         for name, body in _FUNCTION_RE.findall(text):
+            props = (tools.get(name) or {}).get("properties", {})
             args = {
-                key: _coerce(value.strip())
+                key: _coerce(value.strip(), props.get(key, {}).get("type"))
                 for key, value in _PARAMETER_RE.findall(body)
             }
-            calls.append(self._make_item(name, args, tool_names))
+            calls.append(self._make_item(name, args, tools))
         return calls
 
     def _make_item(
-        self, name: Optional[str], arguments: dict, tool_names: set[str]
+        self, name: Optional[str], arguments: dict, tools: dict[str, dict]
     ) -> ToolCallItem:
-        if not name or name not in tool_names:
+        if not name or name not in tools:
             raise _UndefinedToolCall(repr(name))
         item = ToolCallItem(
             tool_index=self._next_index,