refactor: split infer core into subpackages by concern
- Eliminate core/ directory into cache/, runtime/, network/ subpackages plus flat modules
- Split cache.py (647 lines) into cache/{buffer,strategy,pool}.py by layer
- Add explicit ContiguousStrategy, make AllocationStrategy a real ABC
- Move TaskCacheState to cache/strategy.py, drop string forward references
- Rename api/ to network/, server.py to app.py
- Move sample.py into runtime/ alongside executor and graph
- Simplify TaskCacheManager.__init__ to single pool param
- Expose pool.strategy and pool.req_pool as public properties
- Fix KVCache import in attention_backend.py (TYPE_CHECKING guard)
- Fix steady-state decode reading uninitialized position_ids on first step
This commit is contained in:
@@ -0,0 +1,339 @@
|
||||
"""Tool call parsers for extracting structured tool calls from model output.
|
||||
|
||||
Patterned after vLLM's ToolParser abstraction. Each parser knows how to
|
||||
detect and incrementally extract tool calls from raw generated text.
|
||||
|
||||
Subclasses may optionally consume ``token_ids`` for token-level parsing
|
||||
(e.g. Harmony / VLM-style parsers).
|
||||
"""
|
||||
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from astrai.factory import BaseFactory
|
||||
|
||||
|
||||
class BaseToolParser(ABC):
|
||||
"""Abstract tool call parser — one instance per request.
|
||||
|
||||
Maintains streaming state internally so that each call to :meth:`feed`
|
||||
can diff against previously emitted content.
|
||||
|
||||
Args:
|
||||
tools (list of dict, optional): Tool definitions from the request.
|
||||
tool_choice (str): ``"auto"`` / ``"required"`` / ``"none"`` or a named
|
||||
tool choice dict.
|
||||
"""
|
||||
|
||||
def __init__(self, tools: Optional[List[Dict]] = None, tool_choice: str = "auto"):
|
||||
self.tools = tools or []
|
||||
self.tool_choice = tool_choice
|
||||
|
||||
@abstractmethod
|
||||
def feed(
|
||||
self,
|
||||
body: str,
|
||||
current_token_ids: Optional[List[int]] = None,
|
||||
delta_token_ids: Optional[List[int]] = None,
|
||||
) -> List[Dict]:
|
||||
"""Feed the *full* accumulated text each step.
|
||||
|
||||
Returns a list of delta dicts to emit. Each delta is one of:
|
||||
|
||||
- ``{"content": "text"}`` — plain text delta
|
||||
- ``{"tool_calls": [...]}`` — tool-call delta (OpenAI format)
|
||||
|
||||
Returns an empty list when nothing new should be emitted.
|
||||
|
||||
Args:
|
||||
body (str): The complete accumulated generated text so far.
|
||||
current_token_ids (list of int, optional): All token IDs decoded
|
||||
into *body* (cumulative).
|
||||
delta_token_ids (list of int, optional): Only the token IDs for
|
||||
this chunk.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def parse_complete(self, body: str) -> Optional[Dict]:
|
||||
"""Parse the *complete* generated text after generation ends.
|
||||
|
||||
Returns ``None`` when no tool calls were found, otherwise a dict
|
||||
with ``content`` (str or None) and ``tool_calls`` (list of dicts).
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def has_tool_calls(self) -> bool:
|
||||
"""True if the parser detected at least one tool call in the stream."""
|
||||
|
||||
|
||||
class ToolParserFactory(BaseFactory["BaseToolParser"]):
|
||||
pass
|
||||
|
||||
|
||||
_TOOL_CALL_HEAD_RE = re.compile(r'\{\s*"name"\s*:')
|
||||
|
||||
|
||||
def _scan_json(text: str, start: int = 0):
|
||||
"""Scan for a complete JSON object starting at *start*.
|
||||
|
||||
Returns ``(end, complete)`` where *end* is one-past the closing
|
||||
brace (or ``len(text)`` if unclosed), and *complete* is a bool.
|
||||
"""
|
||||
depth = 0
|
||||
in_string = False
|
||||
escape = False
|
||||
for i in range(start, len(text)):
|
||||
c = text[i]
|
||||
if escape:
|
||||
escape = False
|
||||
continue
|
||||
if c == "\\":
|
||||
escape = True
|
||||
continue
|
||||
if c == '"':
|
||||
in_string = not in_string
|
||||
continue
|
||||
if in_string:
|
||||
continue
|
||||
if c == "{":
|
||||
depth += 1
|
||||
elif c == "}":
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
return i + 1, True
|
||||
return len(text), False
|
||||
|
||||
|
||||
def _parse_tool_call_json(json_str: str, complete: bool):
|
||||
"""Extract *name* and *arguments* from a tool-call JSON string.
|
||||
|
||||
Returns ``(name, args, valid)``.
|
||||
"""
|
||||
if complete:
|
||||
try:
|
||||
obj = json.loads(json_str)
|
||||
except json.JSONDecodeError:
|
||||
return None, "", False
|
||||
name = obj.get("name")
|
||||
if not isinstance(name, str) or not name:
|
||||
return None, "", False
|
||||
args = obj.get("arguments")
|
||||
if isinstance(args, dict):
|
||||
if not args:
|
||||
args = ""
|
||||
else:
|
||||
args = json.dumps(args, ensure_ascii=False)
|
||||
args = args[1:-1].rstrip()
|
||||
elif isinstance(args, list):
|
||||
args = json.dumps(args, ensure_ascii=False) if args else ""
|
||||
elif isinstance(args, str):
|
||||
pass
|
||||
else:
|
||||
args = str(args) if args is not None else ""
|
||||
return name, args, True
|
||||
|
||||
name_match = re.search(r'"name"\s*:\s*"([^"]*)"', json_str)
|
||||
if not name_match:
|
||||
return None, "", False
|
||||
name = name_match.group(1)
|
||||
|
||||
args_match = re.search(r'"arguments"\s*:\s*(.*)', json_str, re.DOTALL)
|
||||
if not args_match:
|
||||
return name, "", True
|
||||
|
||||
raw = args_match.group(1).rstrip()
|
||||
if raw.startswith("{"):
|
||||
inner = raw[1:].rstrip()
|
||||
if inner.endswith("}"):
|
||||
inner = inner[:-1].rstrip()
|
||||
raw = inner
|
||||
return name, raw, True
|
||||
|
||||
|
||||
def _find_tool_calls(text: str, start_pos: int = 0):
|
||||
"""Find all complete ``{...}`` tool-call objects in *text*.
|
||||
|
||||
Returns a list of dicts with keys *start*, *end*, *name*, *args*,
|
||||
*complete*.
|
||||
"""
|
||||
results = []
|
||||
pos = start_pos
|
||||
|
||||
while True:
|
||||
brace = text.find("{", pos)
|
||||
if brace == -1:
|
||||
break
|
||||
|
||||
end, complete = _scan_json(text, brace)
|
||||
if not complete:
|
||||
break
|
||||
|
||||
json_str = text[brace:end]
|
||||
|
||||
name, args, valid = _parse_tool_call_json(json_str, complete=True)
|
||||
if not valid or name is None:
|
||||
pos = end
|
||||
continue
|
||||
|
||||
results.append(
|
||||
{
|
||||
"start": brace,
|
||||
"end": end,
|
||||
"name": name,
|
||||
"args": args,
|
||||
"complete": True,
|
||||
}
|
||||
)
|
||||
pos = end
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def _find_partial_tool_call(text: str, start_pos: int = 0):
|
||||
"""Find one incomplete (still-generating) tool-call JSON object."""
|
||||
brace = text.find("{", start_pos)
|
||||
if brace == -1:
|
||||
return None
|
||||
|
||||
json_str = text[brace:]
|
||||
if '"name"' not in json_str:
|
||||
return None
|
||||
|
||||
name, args, valid = _parse_tool_call_json(json_str, complete=False)
|
||||
if not valid or name is None:
|
||||
return None
|
||||
|
||||
return {
|
||||
"start": brace,
|
||||
"name": name,
|
||||
"args": args,
|
||||
"complete": False,
|
||||
}
|
||||
|
||||
|
||||
@ToolParserFactory.register("simple_json")
|
||||
class SimpleJsonToolParser(BaseToolParser):
|
||||
"""Parser for models that output tool calls as plain JSON objects.
|
||||
|
||||
Detects ``{"name": "<func>", "arguments": {...}}`` anywhere in the
|
||||
generated text. Handles single and (non-overlapping) multiple tool
|
||||
calls. Text preceding the first tool call is emitted as plain
|
||||
``content`` deltas.
|
||||
"""
|
||||
|
||||
def __init__(self, tools=None, tool_choice="auto"):
|
||||
super().__init__(tools, tool_choice)
|
||||
self._emitted_content_len = 0
|
||||
self._tc_state: List[Dict] = []
|
||||
self._has_tool_calls = False
|
||||
|
||||
# -------------------------------------------------------------- feed
|
||||
|
||||
def feed(
|
||||
self,
|
||||
body: str,
|
||||
current_token_ids: Optional[List[int]] = None,
|
||||
delta_token_ids: Optional[List[int]] = None,
|
||||
) -> List[Dict]:
|
||||
deltas: List[Dict] = []
|
||||
|
||||
completed = _find_tool_calls(body)
|
||||
|
||||
if not completed:
|
||||
partial = _find_partial_tool_call(body)
|
||||
if not partial:
|
||||
return self._emit_plain_content(body, deltas)
|
||||
all_tcs = [partial]
|
||||
else:
|
||||
all_tcs = completed
|
||||
partial = _find_partial_tool_call(body, completed[-1]["end"])
|
||||
if partial:
|
||||
all_tcs = completed + [partial]
|
||||
|
||||
first_start = all_tcs[0]["start"]
|
||||
if first_start > self._emitted_content_len:
|
||||
content = body[self._emitted_content_len : first_start]
|
||||
self._emitted_content_len = first_start
|
||||
if content:
|
||||
deltas.append({"content": content})
|
||||
|
||||
for i, tc in enumerate(all_tcs):
|
||||
if i >= len(self._tc_state):
|
||||
self._tc_state.append(
|
||||
{
|
||||
"id": f"call_{uuid.uuid4().hex[:12]}",
|
||||
"name_emitted": False,
|
||||
"args_emitted_len": 0,
|
||||
}
|
||||
)
|
||||
self._has_tool_calls = True
|
||||
st = self._tc_state[i]
|
||||
|
||||
if not st["name_emitted"]:
|
||||
st["name_emitted"] = True
|
||||
deltas.append(
|
||||
{
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": i,
|
||||
"id": st["id"],
|
||||
"type": "function",
|
||||
"function": {"name": tc["name"], "arguments": ""},
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
new_args = tc["args"]
|
||||
if len(new_args) > st["args_emitted_len"]:
|
||||
diff = new_args[st["args_emitted_len"] :]
|
||||
st["args_emitted_len"] = len(new_args)
|
||||
deltas.append(
|
||||
{
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": i,
|
||||
"function": {"arguments": diff},
|
||||
}
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
return deltas
|
||||
|
||||
def _emit_plain_content(self, body: str, deltas: List[Dict]) -> List[Dict]:
|
||||
new_content = body[self._emitted_content_len :]
|
||||
if new_content:
|
||||
self._emitted_content_len = len(body)
|
||||
deltas.append({"content": new_content})
|
||||
return deltas
|
||||
|
||||
# -------------------------------------------------------- complete
|
||||
|
||||
def parse_complete(self, body: str) -> Optional[Dict]:
|
||||
completed = _find_tool_calls(body)
|
||||
if not completed:
|
||||
return None
|
||||
|
||||
content = body[: completed[0]["start"]].strip() or None
|
||||
tool_calls = []
|
||||
for i, tc in enumerate(completed):
|
||||
tool_calls.append(
|
||||
{
|
||||
"id": f"call_{uuid.uuid4().hex[:12]}",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc["name"],
|
||||
"arguments": tc["args"],
|
||||
},
|
||||
}
|
||||
)
|
||||
return {"content": content, "tool_calls": tool_calls}
|
||||
|
||||
@property
|
||||
def has_tool_calls(self) -> bool:
|
||||
return self._has_tool_calls
|
||||
Reference in New Issue
Block a user