ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat : 推理层增加 vLLM 风格工具调用解析

Browse Source 
- 新增 BaseToolParser 抽象基类，定义 feed/parse_complete 流式接口
- 新增 SimpleJsonToolParser，解析 {"name":"...","arguments":{...}} 格式
- 新增 ToolParserFactory，基于 BaseFactory 实现可插拔注册
- 集成 parser 到 OpenAIResponseBuilder，支持流式/非流式工具调用
- 扩展 ChatMessage 和 ChatCompletionRequest，增加 tools/tool_choice 字段
- 重构 format_chunk 接口，传入累积文本支持全量重新解析
- 新增 74 个单元测试，覆盖扫描/查找/流式解析/完整解析/工厂
This commit is contained in:
ViperEkura
2026-06-06 08:54:10 +08:00
parent 986be957ec
commit 52aa4d01d5
9 changed files with 1154 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/inference/__init__.py
+10
View File
@@ -11,12 +11,17 @@ Layers:
from astrai.inference.api import (
AnthropicMessage,
BaseToolParser,
ChatCompletionRequest,
ChatMessage,
FunctionDef,
GenContext,
MessagesRequest,
ProtocolHandler,
SimpleJsonToolParser,
StopChecker,
ToolDef,
ToolParserFactory,
get_app,
run_server,
)
@@ -74,10 +79,15 @@ __all__ = [
"ProtocolHandler",
"StopChecker",
"GenContext",
"BaseToolParser",
"SimpleJsonToolParser",
"ToolParserFactory",
"OpenAIResponseBuilder",
"AnthropicResponseBuilder",
"ChatMessage",
"ChatCompletionRequest",
"FunctionDef",
"ToolDef",
"AnthropicMessage",
"MessagesRequest",
"get_app",
astrai/inference/api/__init__.py
+13 -1
View File
@@ -1,4 +1,4 @@
"""Inference API: protocol handler, stop checker, and FastAPI server.
"""Inference API: protocol handler, stop checker, tool parsers, and FastAPI server.
``app`` is no longer a module-level global. Use :func:`get_app` to access the
lazy singleton FastAPI instance.
@@ -9,18 +9,30 @@ from astrai.inference.api.server import (
AnthropicMessage,
ChatCompletionRequest,
ChatMessage,
FunctionDef,
MessagesRequest,
ToolDef,
get_app,
run_server,
)
from astrai.inference.api.tool_parser import (
BaseToolParser,
SimpleJsonToolParser,
ToolParserFactory,
)
__all__ = [
"ProtocolHandler",
"StopChecker",
"GenContext",
"BaseToolParser",
"SimpleJsonToolParser",
"ToolParserFactory",
"AnthropicMessage",
"ChatCompletionRequest",
"ChatMessage",
"FunctionDef",
"ToolDef",
"MessagesRequest",
"get_app",
"run_server",
astrai/inference/api/anthropic.py
+11 -9
View File
@@ -73,15 +73,17 @@ class AnthropicResponseBuilder(ResponseBuilder):
),
]
def format_chunk(self, token: str) -> str:
return sse_event(
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": token},
},
event="content_block_delta",
)
def format_chunk(self, token: str, body: str) -> List[str]:
return [
sse_event(
{
"type": "content_block_delta",
"index": 0,
"delta": {"type": "text_delta", "text": token},
},
event="content_block_delta",
)
]
def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
events: List[str] = []
astrai/inference/api/openai.py
+146 -12
View File
@@ -3,7 +3,7 @@
import logging
import time
import uuid
from typing import Any, Dict, List, Tuple
from typing import Any, Dict, List, Optional, Tuple, Union
from pydantic import BaseModel
@@ -13,6 +13,7 @@ from astrai.inference.api.protocol import (
StopInfo,
sse_event,
)
from astrai.inference.api.tool_parser import BaseToolParser, ToolParserFactory
from astrai.inference.engine import InferenceEngine
logger = logging.getLogger(__name__)
@@ -26,12 +27,37 @@ _UNSUPPORTED_PARAMS = (
)
def _resolve_tool_choice(
request: BaseModel,
) -> Union[str, Dict[str, Any]]:
tc = getattr(request, "tool_choice", None)
if tc is None:
return "auto"
if isinstance(tc, str):
return tc
if isinstance(tc, dict):
return tc
return "auto"
def _resolve_tools(request: BaseModel) -> Optional[List[Dict[str, Any]]]:
raw = getattr(request, "tools", None)
if not raw:
return None
if isinstance(raw, list):
return [t.model_dump() if hasattr(t, "model_dump") else t for t in raw]
return None
class OpenAIResponseBuilder(ResponseBuilder):
def prepare(
self, request: BaseModel, engine: InferenceEngine
) -> Tuple[str, GenContext, List[str]]:
messages = [{"role": m.role, "content": m.content} for m in request.messages]
prompt = engine.tokenizer.apply_chat_template(messages, tokenize=False)
tools = _resolve_tools(request)
prompt = engine.tokenizer.apply_chat_template(
messages, tokenize=False, tools=tools or []
)
self._resp_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
self._model = request.model
@@ -42,17 +68,20 @@ class OpenAIResponseBuilder(ResponseBuilder):
default = fields[param].default if param in fields else None
if value is not None and value != default:
logger.warning(
"ChatCompletionRequest param '%s'=%r is not supported and will be ignored",
param,
value,
)
if value is not None and value != default:
logger.warning(
"ChatCompletionRequest param '%s'=%r is not supported and will be ignored",
"ChatCompletionRequest param '%s'=%r is not supported"
" and will be ignored",
param,
value,
)
self._parser: Optional[BaseToolParser] = None
if tools:
tool_choice = _resolve_tool_choice(request)
self._parser = ToolParserFactory.create(
"simple_json", tools=tools, tool_choice=tool_choice
)
self._content_started = False
ctx = GenContext(
resp_id=self._resp_id,
created=int(time.time()),
@@ -84,7 +113,77 @@ class OpenAIResponseBuilder(ResponseBuilder):
)
]
def format_chunk(self, token: str) -> str:
def format_chunk(self, token: str, body: str) -> List[str]:
if self._parser is not None:
return self._format_tool_chunk(body)
return [
sse_event(
{
"id": self._resp_id,
"object": "chat.completion.chunk",
"created": 0,
"model": self._model,
"choices": [
{
"index": 0,
"delta": {"content": token},
"finish_reason": None,
}
],
}
)
]
def _format_tool_chunk(self, body: str) -> List[str]:
deltas = self._parser.feed(body)
events: List[str] = []
for d in deltas:
if "content" in d:
if not self._content_started:
events.append(self._role_chunk())
self._content_started = True
events.append(
sse_event(
{
"id": self._resp_id,
"object": "chat.completion.chunk",
"created": 0,
"model": self._model,
"choices": [
{
"index": 0,
"delta": {"content": d["content"]},
"finish_reason": None,
}
],
}
)
)
elif "tool_calls" in d:
if not self._content_started:
events.append(self._role_chunk())
self._content_started = True
events.append(
sse_event(
{
"id": self._resp_id,
"object": "chat.completion.chunk",
"created": 0,
"model": self._model,
"choices": [
{
"index": 0,
"delta": {"tool_calls": d["tool_calls"]},
"finish_reason": None,
}
],
}
)
)
return events
def _role_chunk(self) -> str:
return sse_event(
{
"id": self._resp_id,
@@ -92,12 +191,19 @@ class OpenAIResponseBuilder(ResponseBuilder):
"created": 0,
"model": self._model,
"choices": [
{"index": 0, "delta": {"content": token}, "finish_reason": None}
{
"index": 0,
"delta": {"role": "assistant"},
"finish_reason": None,
}
],
}
)
def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
finish_reason = "stop"
if self._parser is not None and self._parser.has_tool_calls:
finish_reason = "tool_calls"
return [
sse_event(
{
@@ -105,7 +211,9 @@ class OpenAIResponseBuilder(ResponseBuilder):
"object": "chat.completion.chunk",
"created": ctx.created,
"model": self._model,
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
"choices": [
{"index": 0, "delta": {}, "finish_reason": finish_reason}
],
}
),
sse_event(
@@ -120,6 +228,32 @@ class OpenAIResponseBuilder(ResponseBuilder):
def format_response(
self, ctx: GenContext, content: str, stop: StopInfo
) -> Dict[str, Any]:
if self._parser is not None:
parsed = self._parser.parse_complete(content)
if parsed and parsed.get("tool_calls"):
return {
"id": self._resp_id,
"object": "chat.completion",
"created": ctx.created,
"model": self._model,
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": parsed.get("content"),
"tool_calls": parsed["tool_calls"],
},
"finish_reason": "tool_calls",
}
],
"usage": {
"prompt_tokens": ctx.prompt_tokens,
"completion_tokens": ctx.completion_tokens,
"total_tokens": ctx.prompt_tokens + ctx.completion_tokens,
},
}
return {
"id": self._resp_id,
"object": "chat.completion",
astrai/inference/api/protocol.py
+9 -3
View File
@@ -78,8 +78,13 @@ class ResponseBuilder(ABC):
"""SSE events that open the stream."""
@abstractmethod
def format_chunk(self, token: str) -> str:
"""SSE event for a single generated token."""
def format_chunk(self, token: str, body: str) -> List[str]:
"""SSE events for a single generated token.
Receives the current token and the full accumulated *body* so
that tool-call parsers can re-parse the complete text each step.
Returns a list of SSE event strings (may be empty).
"""
@abstractmethod
def format_stream_end(self, ctx: GenContext, stop: StopInfo) -> List[str]:
@@ -145,7 +150,8 @@ class ProtocolHandler:
break
ctx.completion_tokens += 1
yield self.builder.format_chunk(token)
for event in self.builder.format_chunk(token, body):
yield event
yielded += token
stop = StopInfo(matched=matched, body=body, yielded=yielded)
astrai/inference/api/server.py
+16 -1
View File
@@ -32,7 +32,20 @@ _app_instance: Optional[FastAPI] = None
class ChatMessage(BaseModel):
role: str
content: str
content: Optional[str] = None
tool_calls: Optional[List[Dict[str, Any]]] = None
tool_call_id: Optional[str] = None
class FunctionDef(BaseModel):
name: str
description: Optional[str] = None
parameters: Optional[Dict[str, Any]] = None
class ToolDef(BaseModel):
type: str = "function"
function: FunctionDef
class ChatCompletionRequest(BaseModel):
@@ -51,6 +64,8 @@ class ChatCompletionRequest(BaseModel):
frequency_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
logit_bias: Optional[Dict[int, float]] = None
user: Optional[str] = None
tools: Optional[List[ToolDef]] = None
tool_choice: Optional[Union[str, Dict[str, Any]]] = "auto"
class AnthropicMessage(BaseModel):
astrai/inference/api/tool_parser.py
+300
View File
@@ -0,0 +1,300 @@
"""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.
"""
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.
"""
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) -> 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.
"""
@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"]):
@classmethod
def _validate_component(cls, component_cls: type):
if not issubclass(component_cls, BaseToolParser):
raise TypeError(
f"{component_cls.__name__} must inherit from BaseToolParser"
)
_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)``.
"""
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 complete and raw.endswith("}"):
raw = raw[:-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]
if not _TOOL_CALL_HEAD_RE.search(json_str):
pos = end
continue
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 not _TOOL_CALL_HEAD_RE.search(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) -> 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