diff --git a/examples/server/server.py b/examples/server/server.py
index ec83fe6e3..28fc8f4eb 100644
--- a/examples/server/server.py
+++ b/examples/server/server.py
@@ -4344,6 +4344,11 @@ def capture(match: re.Match[str]) -> str:
@staticmethod
def _regex_literal_prefix(pattern: str) -> str:
+ literal, _ = ResponseParser._regex_literal_prefix_and_remainder(pattern)
+ return literal
+
+ @staticmethod
+ def _regex_literal_prefix_and_remainder(pattern: str) -> Tuple[str, str]:
literal: List[str] = []
index = 0
while index < len(pattern):
@@ -4365,7 +4370,181 @@ def _regex_literal_prefix(pattern: str) -> str:
break
literal.append(char)
index += 1
- return "".join(literal)
+ return "".join(literal), pattern[index:]
+
+ @staticmethod
+ def _find_regex_group_end(pattern: str, start: int) -> int:
+ depth = 0
+ escaped = False
+ in_character_class = False
+ for index in range(start, len(pattern)):
+ char = pattern[index]
+ if escaped:
+ escaped = False
+ continue
+ if char == "\\":
+ escaped = True
+ continue
+ if char == "[":
+ in_character_class = True
+ continue
+ if char == "]" and in_character_class:
+ in_character_class = False
+ continue
+ if in_character_class:
+ continue
+ if char == "(":
+ depth += 1
+ continue
+ if char == ")":
+ depth -= 1
+ if depth == 0:
+ return index
+ return -1
+
+ @classmethod
+ def _consume_optional_literal_prefix(
+ cls,
+ pattern: str,
+ ) -> Optional[Tuple[str, str]]:
+ if not pattern.startswith("(?:"):
+ return None
+ group_end = cls._find_regex_group_end(pattern, 0)
+ if group_end < 0 or group_end + 1 >= len(pattern) or pattern[group_end + 1] != "?":
+ return None
+ literal, remainder = cls._regex_literal_prefix_and_remainder(pattern[3:group_end])
+ if not literal or remainder:
+ return None
+ return literal, pattern[group_end + 2 :]
+
+ @staticmethod
+ def _split_regex_alternatives(pattern: str) -> List[str]:
+ alternatives: List[str] = []
+ start = 0
+ depth = 0
+ escaped = False
+ in_character_class = False
+ for index, char in enumerate(pattern):
+ if escaped:
+ escaped = False
+ continue
+ if char == "\\":
+ escaped = True
+ continue
+ if char == "[":
+ in_character_class = True
+ continue
+ if char == "]" and in_character_class:
+ in_character_class = False
+ continue
+ if in_character_class:
+ continue
+ if char == "(":
+ depth += 1
+ continue
+ if char == ")":
+ depth -= 1
+ continue
+ if char == "|" and depth == 0:
+ alternatives.append(pattern[start:index])
+ start = index + 1
+ alternatives.append(pattern[start:])
+ return alternatives
+
+ @classmethod
+ def _regex_lookahead_literal_specs(cls, pattern: str) -> List[Tuple[str, bool]]:
+ if not pattern.startswith("(?="):
+ return []
+ group_end = cls._find_regex_group_end(pattern, 0)
+ if group_end < 0:
+ return []
+ literals: List[Tuple[str, bool]] = []
+ for alternative in cls._split_regex_alternatives(pattern[3:group_end]):
+ strip_leading_whitespace = False
+ while alternative.startswith(r"\s*"):
+ strip_leading_whitespace = True
+ alternative = alternative[3:]
+ if alternative == "$":
+ continue
+ if alternative.endswith("$"):
+ alternative = alternative[:-1]
+ literal, _ = cls._regex_literal_prefix_and_remainder(alternative)
+ if literal:
+ literals.append((literal, strip_leading_whitespace))
+ return literals
+
+ @classmethod
+ def _regex_capture_parts(
+ cls,
+ pattern: str,
+ ) -> Optional[Tuple[str, str]]:
+ normalized = pattern.lstrip("^")
+ captures = [
+ (index, token)
+ for token in ("(.*?)", "(.*)")
+ if (index := normalized.find(token)) >= 0
+ ]
+ if not captures:
+ return None
+ capture_index, capture_token = min(captures, key=lambda item: item[0])
+ return normalized[:capture_index], normalized[capture_index + len(capture_token) :]
+
+ @classmethod
+ def _regex_capture_end_literal_specs(cls, pattern: str) -> List[Tuple[str, bool]]:
+ capture_parts = cls._regex_capture_parts(pattern)
+ if capture_parts is None:
+ return []
+ _, suffix_pattern = capture_parts
+ literal_specs = cls._regex_lookahead_literal_specs(suffix_pattern)
+ if literal_specs:
+ return literal_specs
+ literal, _ = cls._regex_literal_prefix_and_remainder(suffix_pattern)
+ return [(literal, False)] if literal else []
+
+ @classmethod
+ def _regex_capture_end_literals(cls, pattern: str) -> List[str]:
+ return [literal for literal, _ in cls._regex_capture_end_literal_specs(pattern)]
+
+ @classmethod
+ def _regex_leading_capture(
+ cls,
+ *,
+ field_name: str,
+ field_regex: str,
+ content_regex: Optional[str],
+ ) -> Optional[Dict[str, Any]]:
+ capture_parts = cls._regex_capture_parts(field_regex)
+ if capture_parts is None:
+ return None
+ prefix_pattern, _ = capture_parts
+ prefix_pattern = prefix_pattern.lstrip("^")
+ optional_prefix = cls._consume_optional_literal_prefix(prefix_pattern)
+ if optional_prefix is not None:
+ prefix_pattern = optional_prefix[1]
+ implicit_at_start = False
+ optional_capture_start = cls._consume_optional_literal_prefix(prefix_pattern)
+ if optional_capture_start is not None:
+ capture_start, prefix_pattern = optional_capture_start
+ implicit_at_start = True
+ else:
+ capture_start, prefix_pattern = cls._regex_literal_prefix_and_remainder(prefix_pattern)
+ if not capture_start or prefix_pattern:
+ return None
+ end_literals = cls._regex_capture_end_literals(field_regex)
+ if not end_literals:
+ return None
+ capture_end = end_literals[0]
+ strip_after = False
+ if isinstance(content_regex, str):
+ escaped_end = re.escape(capture_end)
+ strip_after = bool(re.search(escaped_end + r"\\s\*", content_regex))
+ return {
+ "field": field_name,
+ "start": capture_start,
+ "end": capture_end,
+ "strip_after": strip_after,
+ "implicit_at_start": implicit_at_start,
+ }
@staticmethod
def _literal_suffix_prefix_length(text: str, literal: str) -> int:
@@ -4702,7 +4881,14 @@ def _compile_tagged_message_plan(
return None
iterator = cls._compile_iterator_pattern(iterator_pattern)
if iterator is None:
- return None
+ iterator_capture = cls._compile_iterator_block_pattern(iterator_pattern)
+ if (
+ not isinstance(iterator_capture, dict)
+ or not iterator_capture["start"]
+ or iterator_capture["allow_eof"]
+ ):
+ return None
+ iterator = (iterator_capture["start"], iterator_capture["end"])
items_schema = tool_calls_schema.get("items")
if not isinstance(items_schema, dict):
return None
@@ -4715,10 +4901,11 @@ def _compile_tagged_message_plan(
if isinstance(content_schema, dict)
else None
)
- object_regex = schema.get("x-regex") if isinstance(schema.get("x-regex"), str) else None
assistant_prefix: Optional[str] = None
- if isinstance(content_regex, str) and r"<\|im_start\|>assistant\n" in content_regex:
- assistant_prefix = "<|im_start|>assistant\n"
+ if isinstance(content_regex, str):
+ optional_prefix = cls._consume_optional_literal_prefix(content_regex.lstrip("^"))
+ if optional_prefix is not None:
+ assistant_prefix = optional_prefix[0]
leading_capture: Optional[Dict[str, Any]] = None
for field_name, value_schema in properties.items():
if not isinstance(value_schema, dict):
@@ -4726,43 +4913,31 @@ def _compile_tagged_message_plan(
field_regex = value_schema.get("x-regex")
if not isinstance(field_regex, str):
continue
- if "\\n" in field_regex and "" in field_regex:
- leading_capture = {
- "field": field_name,
- "start": "\n",
- "end": "",
- "strip_after": True,
- "implicit_at_start": "(?:\\n)?" in field_regex,
- }
+ if field_name == "content":
+ continue
+ capture = cls._regex_leading_capture(
+ field_name=field_name,
+ field_regex=field_regex,
+ content_regex=content_regex,
+ )
+ if capture is not None:
+ leading_capture = capture
break
- if leading_capture is None and isinstance(object_regex, str):
- if (
- "(?P" in object_regex
- and r"<\|channel\>thought\n" in object_regex
- and r"\" in object_regex
- ):
- leading_capture = {
- "field": "thinking",
- "start": "<|channel>thought\n",
- "end": "",
- "strip_after": False,
- "implicit_at_start": False,
- }
end_markers: List[str] = []
+ content_end_marker_specs: List[Tuple[str, bool]] = []
iterator_start, iterator_end = iterator
if "content" in properties:
end_markers.append(iterator_start)
- if isinstance(content_regex, str) and r"<\|im_end\|>" in content_regex:
- end_markers.append("<|im_end|>")
- if isinstance(object_regex, str) and r"" in object_regex:
- end_markers.append("")
+ if isinstance(content_regex, str):
+ content_end_marker_specs = cls._regex_capture_end_literal_specs(content_regex)
+ end_markers.extend(literal for literal, _ in content_end_marker_specs)
if not end_markers and iterator_start:
end_markers.append(iterator_start)
- trim_before_iterator = (
- isinstance(content_regex, str)
- and r"\s*\n" in content_regex
+ deduped_end_markers = tuple(dict.fromkeys(end_markers))
+ trim_before_iterator = any(
+ literal == iterator_start and strip_leading_whitespace
+ for literal, strip_leading_whitespace in content_end_marker_specs
)
- end_marker_tuple = tuple(end_markers)
direct_deltas = item_plan["kind"] == "tagged-parameters"
direct_init = (
(
@@ -4773,8 +4948,8 @@ def _compile_tagged_message_plan(
bool(leading_capture.get("strip_after")) if leading_capture is not None else False,
bool(leading_capture.get("implicit_at_start")) if leading_capture is not None else False,
trim_before_iterator,
- end_marker_tuple,
- tuple(marker for marker in end_marker_tuple if marker != iterator_start),
+ deduped_end_markers,
+ tuple(marker for marker in deduped_end_markers if marker != iterator_start),
iterator_start,
iterator_end,
item_plan["function_start"],
@@ -4793,9 +4968,9 @@ def _compile_tagged_message_plan(
"assistant_prefix": assistant_prefix,
"leading_capture": leading_capture,
"content_field": "content" if "content" in properties else None,
- "content_end_markers": end_marker_tuple,
+ "content_end_markers": deduped_end_markers,
"trim_before_iterator": trim_before_iterator,
- "stop_markers": tuple(marker for marker in end_marker_tuple if marker != iterator_start),
+ "stop_markers": tuple(marker for marker in deduped_end_markers if marker != iterator_start),
"direct_init": direct_init,
"iterator": {
"start": iterator_start,
@@ -7192,19 +7367,17 @@ def consume_completion_chunk(
logprobs=logprobs,
leading_delta=role_delta,
)
- if self._stream_state_complete():
- return self._chunk_payloads(
- chunk_id=chunk_id,
- created=created,
- model=model,
- deltas=stream_deltas,
- finish_reason=(
- "tool_calls" if self._direct.saw_tool_calls else finish_reason
- ),
- logprobs=logprobs,
- leading_delta=role_delta,
- )
- self._stream_failed = True
+ return self._chunk_payloads(
+ chunk_id=chunk_id,
+ created=created,
+ model=model,
+ deltas=stream_deltas,
+ finish_reason=(
+ "tool_calls" if self._direct.saw_tool_calls else finish_reason
+ ),
+ logprobs=logprobs,
+ leading_delta=role_delta,
+ )
elif self._stream_plan["kind"] == "segment-message":
if role_delta is not None:
stream_deltas = [role_delta, *stream_deltas]
@@ -7219,18 +7392,16 @@ def consume_completion_chunk(
finish_reason=None,
logprobs=logprobs,
)
- if self._stream_state_complete():
- return self._chunk_payloads(
- chunk_id=chunk_id,
- created=created,
- model=model,
- deltas=stream_deltas,
- finish_reason=(
- "tool_calls" if self._message.get("tool_calls") else finish_reason
- ),
- logprobs=logprobs,
- )
- self._stream_failed = True
+ return self._chunk_payloads(
+ chunk_id=chunk_id,
+ created=created,
+ model=model,
+ deltas=stream_deltas,
+ finish_reason=(
+ "tool_calls" if self._message.get("tool_calls") else finish_reason
+ ),
+ logprobs=logprobs,
+ )
else:
previous_message = self._message
partial_deltas: List[Dict[str, Any]] = []
@@ -7238,7 +7409,7 @@ def consume_completion_chunk(
parsed = cast(Dict[str, Any], self._stream_state.parsed)
message = self._parsed_chat_message(
parsed=parsed,
- partial=finish_reason is None,
+ partial=finish_reason is None or not self._stream_state_complete(),
)
if finish_reason is None:
if role_delta is not None:
@@ -7253,22 +7424,20 @@ def consume_completion_chunk(
finish_reason=None,
logprobs=logprobs,
)
- if self._stream_state_complete():
- if role_delta is not None:
- partial_deltas.append(role_delta)
- partial_deltas.extend(self._message_deltas(previous_message, message))
- self._message = message
- return self._chunk_payloads(
- chunk_id=chunk_id,
- created=created,
- model=model,
- deltas=partial_deltas,
- finish_reason=(
- "tool_calls" if message.get("tool_calls") else finish_reason
- ),
- logprobs=logprobs,
- )
- self._stream_failed = True
+ if role_delta is not None:
+ partial_deltas.append(role_delta)
+ partial_deltas.extend(self._message_deltas(previous_message, message))
+ self._message = message
+ return self._chunk_payloads(
+ chunk_id=chunk_id,
+ created=created,
+ model=model,
+ deltas=partial_deltas,
+ finish_reason=(
+ "tool_calls" if message.get("tool_calls") else finish_reason
+ ),
+ logprobs=logprobs,
+ )
else:
self._stream_failed = True