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/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",