diff --git a/python/infinilm/llm/llm.py b/python/infinilm/llm/llm.py index 59d5a1eca..f331ef488 100644 --- a/python/infinilm/llm/llm.py +++ b/python/infinilm/llm/llm.py @@ -35,6 +35,31 @@ logger = logging.getLogger(__name__) +def validate_request_capacity( + config: EngineConfig, prompt_tokens: int, output_tokens: int +) -> None: + """Reject requests that can never fit in the configured KV cache.""" + if config.cache_type == "paged": + capacity = config.num_blocks * config.block_size + else: + capacity = config.max_cache_len + + if prompt_tokens > capacity: + raise ValueError( + f"The maximum context length is {capacity} tokens. " + f"The prompt contains {prompt_tokens} tokens." + ) + + total_tokens = prompt_tokens + output_tokens + if total_tokens > capacity: + available_output_tokens = max(capacity - prompt_tokens, 0) + raise ValueError( + f"The maximum context length is {capacity} tokens. " + f"The prompt contains {prompt_tokens} tokens, so max_tokens must be " + f"at most {available_output_tokens}; received {output_tokens}." + ) + + class LLMEngine: """Low-level LLM engine that handles inference execution.""" @@ -784,6 +809,7 @@ def add_request( request_id: Optional[str] = None, # For server use request_data: Optional[dict] = None, + chat_template_kwargs: Optional[dict] = None, ) -> InferenceRequest: """Add a request to the engine. @@ -829,17 +855,27 @@ def add_request( elif prompt is not None: prompt_token_ids = self.engine.tokenize(prompt) else: - assert messages is not None, ( - "Either messages or prompt/prompt_token_ids must be provided" - ) + assert ( + messages is not None + ), "Either messages or prompt/prompt_token_ids must be provided" - assert apply_chat_template, ( - "apply_chat_template needs to be true for multi-role conversation" - ) + assert ( + apply_chat_template + ), "apply_chat_template needs to be true for multi-role conversation" prompt = self.engine.apply_chat_template( - messages, add_generation_prompt=add_generation_prompt + messages, + add_generation_prompt=add_generation_prompt, + chat_template_kwargs=chat_template_kwargs, + ) + + forced_tool_prefix = ( + request_data.get("_infinilm_forced_tool_prefix") + if request_data + else None ) + if forced_tool_prefix: + prompt += forced_tool_prefix mm_inputs = resolve_multimodal_inputs(messages) @@ -868,6 +904,12 @@ def add_request( sampling_params = sampling_params.clone() sampling_params.max_tokens = self.config.max_tokens + validate_request_capacity( + self.config, + prompt_tokens=len(prompt_token_ids), + output_tokens=sampling_params.max_tokens, + ) + request = InferenceRequest( request_id=request_id, prompt=prompt, @@ -897,6 +939,7 @@ def add_chat_request( request_id: Optional[str] = None, request_data: Optional[dict] = None, add_generation_prompt: bool = True, + chat_template_kwargs: Optional[dict] = None, **kwargs, ) -> InferenceRequest: """Add a chat request to the engine. @@ -918,6 +961,7 @@ def add_chat_request( sampling_params=sampling_params, request_id=request_id, request_data=request_data, + chat_template_kwargs=chat_template_kwargs, ) async def stream_request( diff --git a/python/infinilm/processors/basic_llm_processor.py b/python/infinilm/processors/basic_llm_processor.py index a6fbc33ac..279ef4edb 100644 --- a/python/infinilm/processors/basic_llm_processor.py +++ b/python/infinilm/processors/basic_llm_processor.py @@ -3,7 +3,11 @@ from ..llm.scheduler import SchedulerOutput from ..llm.static_scheduler import StaticSchedulerOutput -from .processor import InfinilmProcessor, register_processor +from .processor import ( + InfinilmProcessor, + normalize_openai_messages, + register_processor, +) @register_processor("default") @@ -44,8 +48,8 @@ def apply_chat_template( **kwargs, ): normalized_conversation = [] - for message in conversation: - if isinstance(message["content"], list): + for message in normalize_openai_messages(conversation): + if isinstance(message.get("content"), list): assert len(message["content"]) == 1, ( "Only one content item supported in list" ) diff --git a/python/infinilm/processors/processor.py b/python/infinilm/processors/processor.py index a2952bc1e..f1a1c386f 100644 --- a/python/infinilm/processors/processor.py +++ b/python/infinilm/processors/processor.py @@ -1,3 +1,6 @@ +import json + + class InfinilmProcessor: def __init__(self, model_dir_path: str): """Initialize the processor with the model directory path.""" @@ -43,6 +46,50 @@ def get_mm_token_index_list( raise NotImplementedError("get_mm_token_index_list is not implemented yet") +def normalize_openai_messages(messages: list[dict]) -> list[dict]: + """Convert OpenAI JSON-string tool arguments to template mappings.""" + normalized = [] + for message in messages: + if not isinstance(message, dict): + normalized.append(message) + continue + + normalized_message = message.copy() + if "content" not in normalized_message and normalized_message.get("tool_calls"): + normalized_message["content"] = None + + tool_calls = normalized_message.get("tool_calls") + if isinstance(tool_calls, list): + normalized_calls = [] + for tool_call in tool_calls: + if not isinstance(tool_call, dict): + normalized_calls.append(tool_call) + continue + normalized_call = tool_call.copy() + function = normalized_call.get("function") + if isinstance(function, dict): + normalized_function = function.copy() + arguments = normalized_function.get("arguments") + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) + except json.JSONDecodeError as exc: + raise ValueError( + "tool call function.arguments must be valid JSON" + ) from exc + if not isinstance(arguments, dict): + raise ValueError( + "tool call function.arguments must decode to an object" + ) + normalized_function["arguments"] = arguments + normalized_call["function"] = normalized_function + normalized_calls.append(normalized_call) + normalized_message["tool_calls"] = normalized_calls + + normalized.append(normalized_message) + return normalized + + # Global registry mapping model_type strings to their Processor classes _PROCESSOR_REGISTRY = {} diff --git a/python/infinilm/processors/qwen3_5_processor.py b/python/infinilm/processors/qwen3_5_processor.py index 6b8fee906..53eb51fd6 100644 --- a/python/infinilm/processors/qwen3_5_processor.py +++ b/python/infinilm/processors/qwen3_5_processor.py @@ -7,7 +7,7 @@ from ..llm.scheduler import SchedulerOutput from ..llm.static_scheduler import StaticSchedulerOutput from .basic_llm_processor import BasicLLMProcessor -from .processor import register_processor +from .processor import normalize_openai_messages, register_processor @register_processor("qwen3_5_moe") @@ -134,8 +134,8 @@ def apply_chat_template( **kwargs, ): normalized_conversation = [] - for message in conversation: - content = message["content"] + for message in normalize_openai_messages(conversation): + content = message.get("content") if not isinstance(content, list): normalized_conversation.append(message) continue @@ -156,9 +156,10 @@ def apply_chat_template( f"Unsupported Qwen3.5 content type: {item_type}" ) - normalized_conversation.append( - {"role": message.get("role", "user"), "content": normalized_content} - ) + normalized_message = message.copy() + normalized_message["role"] = message.get("role", "user") + normalized_message["content"] = normalized_content + normalized_conversation.append(normalized_message) template_owner = ( self.processor if self.processor is not None else self.tokenizer diff --git a/python/infinilm/server/inference_server.py b/python/infinilm/server/inference_server.py index 462c31084..be2d86619 100644 --- a/python/infinilm/server/inference_server.py +++ b/python/infinilm/server/inference_server.py @@ -19,6 +19,17 @@ from infinilm.config import KVTransferConfig from infinilm.llm import AsyncLLMEngine, FinishReason, SamplingParams from infinilm.moe_config import configure_moe_ep_backend +from infinilm.server.openai_protocol import ( + ToolCallStreamParser, + parse_tool_calls, + strip_reasoning_markers, +) +from infinilm.server.tool_contract import apply_tool_contract +from infinilm.server.tool_constraints import ( + constrain_tools, + direct_write_complete, + forced_write_tool_prefix, +) logger = logging.getLogger(__name__) @@ -27,15 +38,23 @@ def chunk_json( - id_, content=None, role=None, finish_reason=None, model: str = "unknown" + id_, + content=None, + role=None, + tool_calls=None, + finish_reason=None, + model: str = "unknown", + usage=None, ): """Generate JSON chunk for streaming response.""" delta = {} - if content: + if content is not None: delta["content"] = content if role: delta["role"] = role - return { + if tool_calls: + delta["tool_calls"] = tool_calls + payload = { "id": id_, "object": "chat.completion.chunk", "created": int(time.time()), @@ -51,6 +70,9 @@ def chunk_json( } ], } + if usage is not None: + payload["usage"] = usage + return payload def completion_json( @@ -62,8 +84,13 @@ def completion_json( prompt_tokens: int = 0, completion_tokens: int = 0, total_tokens: int = 0, + tool_calls=None, ): """Generate JSON response for non-streaming completion.""" + message = {"role": role, "content": content} + if tool_calls: + message["tool_calls"] = tool_calls + return { "id": id_, "object": "chat.completion", @@ -73,10 +100,7 @@ def completion_json( "choices": [ { "index": 0, - "message": { - "role": role, - "content": content, - }, + "message": message, "logprobs": None, "finish_reason": finish_reason, } @@ -256,13 +280,16 @@ async def chat_completions(request: Request): # logger.debug(f"Received request data: {data}") except Exception as e: logger.error(f"Failed to parse request JSON: {e}") - return JSONResponse(content={"error": "Invalid JSON"}, status_code=400) + return self._error_response("Invalid JSON", status_code=400) + + try: + self._validate_request(data) + except ValueError as exc: + return self._error_response(str(exc), status_code=400) if not data.get("messages"): if not data.get("prompt"): - return JSONResponse( - content={"error": "No message provided"}, status_code=400 - ) + return self._error_response("No message provided", status_code=400) else: data["messages"] = [{"role": "user", "content": data.get("prompt")}] @@ -349,6 +376,99 @@ def _normalize_messages(self, messages: list) -> list: return normalized + @staticmethod + def _error_response(message: str, status_code: int = 500) -> JSONResponse: + error_type = "invalid_request_error" if status_code < 500 else "server_error" + return JSONResponse( + content={ + "error": { + "message": message, + "type": error_type, + "param": None, + "code": None, + } + }, + status_code=status_code, + ) + + @staticmethod + def _validate_request(data: dict) -> None: + if not isinstance(data, dict): + raise ValueError("Request body must be a JSON object") + if "messages" in data and not isinstance(data["messages"], list): + raise ValueError("messages must be an array") + + tools = data.get("tools") + if tools is not None: + if not isinstance(tools, list): + raise ValueError("tools must be an array") + for tool in tools: + if ( + not isinstance(tool, dict) + or tool.get("type") != "function" + or not isinstance(tool.get("function"), dict) + or not tool["function"].get("name") + ): + raise ValueError( + "each tool must be a function with a non-empty name" + ) + + tool_choice = data.get("tool_choice") + if isinstance(tool_choice, str): + if tool_choice not in ("auto", "none", "required"): + raise ValueError("tool_choice must be auto, none, or required") + elif tool_choice is not None and not isinstance(tool_choice, dict): + raise ValueError("tool_choice must be a string or object") + + @staticmethod + def _build_chat_template_kwargs(data: dict) -> dict: + raw_kwargs = data.get("chat_template_kwargs") or {} + if not isinstance(raw_kwargs, dict): + raise ValueError("chat_template_kwargs must be an object") + kwargs = raw_kwargs.copy() + + tools = constrain_tools(data.get("messages", []), data.get("tools") or []) + tool_choice = data.get("tool_choice") + if tools and tool_choice != "none": + if isinstance(tool_choice, dict): + function = tool_choice.get("function") or {} + function_name = function.get("name") + if function_name: + tools = [ + tool + for tool in tools + if tool.get("function", {}).get("name") == function_name + ] + if not tools: + raise ValueError( + f"tool_choice references unknown function {function_name!r}" + ) + kwargs["tools"] = tools + kwargs["tool_choice"] = tool_choice or "auto" + + for key in ("enable_thinking", "reasoning_effort", "preserve_thinking"): + if key in data: + kwargs[key] = data[key] + + thinking = data.get("thinking") + if isinstance(thinking, dict) and "type" in thinking: + kwargs["enable_thinking"] = thinking["type"] != "disabled" + return kwargs + + def _prepare_chat_request(self, data: dict) -> tuple[list, dict, dict]: + messages = data.get("messages", []) + original_tools = data.get("tools") or [] + chat_template_kwargs = self._build_chat_template_kwargs(data) + exposed_tools = chat_template_kwargs.get("tools") or [] + messages = apply_tool_contract(messages, original_tools, exposed_tools) + request_data = dict(data) + if direct_write_complete(messages, original_tools): + request_data["_infinilm_direct_write_complete"] = True + forced_prefix = forced_write_tool_prefix(messages, exposed_tools) + if forced_prefix is not None: + request_data["_infinilm_forced_tool_prefix"] = forced_prefix + return messages, chat_template_kwargs, request_data + def _build_sampling_params(self, data: dict) -> SamplingParams: """Build SamplingParams from request data.""" # Support both: @@ -366,11 +486,16 @@ def pick(key: str, default): return sp.get(key) return default - # Accept common alias - max_tokens = pick("max_tokens", self.max_tokens) + max_tokens = None + for key in ("max_tokens", "max_completion_tokens", "max_new_tokens"): + if key in data and data[key] is not None: + max_tokens = data[key] + break + if key in sp and sp[key] is not None: + max_tokens = sp[key] + break if max_tokens is None: - # Some clients use max_new_tokens - max_tokens = pick("max_new_tokens", self.max_tokens) + max_tokens = self.max_tokens stop = pick("stop", None) if isinstance(stop, str): @@ -391,17 +516,57 @@ async def _stream_chat(self, request_id: str, data: dict, http_request: Request) _abort_reason = FinishReason.CANCELED try: - messages = data.get("messages", []) + messages, chat_template_kwargs, request_data = self._prepare_chat_request( + data + ) sampling_params = self._build_sampling_params(data) + if request_data.get("_infinilm_direct_write_complete"): + role_chunk = chunk_json( + request_id, role="assistant", model=self.model_id + ) + yield f"data: {json.dumps(role_chunk)}\n\n" + complete_chunk = chunk_json( + request_id, + content="已写入文件，已按请求停止。", + finish_reason="stop", + model=self.model_id, + usage={ + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + }, + ) + yield f"data: {json.dumps(complete_chunk)}\n\n" + yield "data: [DONE]\n\n" + return + req = self.engine.add_chat_request( messages=messages, sampling_params=sampling_params, request_id=request_id, - request_data=data, + request_data=request_data, add_generation_prompt=bool(data.get("add_generation_prompt", True)), - chat_template_kwargs=data.get("chat_template_kwargs") or {}, + chat_template_kwargs=chat_template_kwargs, ) + role_chunk = chunk_json(request_id, role="assistant", model=self.model_id) + yield f"data: {json.dumps(role_chunk)}\n\n" + + tool_parser = ( + ToolCallStreamParser( + allowed_tool_names={ + tool["function"]["name"] + for tool in chat_template_kwargs["tools"] + } + ) + if chat_template_kwargs.get("tools") + else None + ) + forced_tool_prefix = request_data.get("_infinilm_forced_tool_prefix") + if tool_parser is not None and forced_tool_prefix: + tool_parser.feed(forced_tool_prefix) + + tool_call_index = 0 async for token_output in self.engine.stream_request( req, @@ -439,24 +604,71 @@ async def _stream_chat(self, request_id: str, data: dict, http_request: Request) ) if not is_eos_token and token_output.token_text: - # Send token - chunk = json.dumps( - chunk_json( + if tool_parser is None: + content_parts, tool_calls = [token_output.token_text], [] + else: + content_parts, tool_calls = tool_parser.feed( + token_output.token_text + ) + for content_part in content_parts: + visible_content = strip_reasoning_markers(content_part) + if not visible_content: + continue + chunk = chunk_json( + request_id, content=content_part, model=self.model_id + ) + yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n" + for tool_call in tool_calls: + delta_call = {"index": tool_call_index, **tool_call} + tool_call_index += 1 + chunk = chunk_json( request_id, - content=token_output.token_text, + tool_calls=[delta_call], model=self.model_id, - ), - ensure_ascii=False, - ) - yield f"data: {chunk}\n\n" + ) + yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n" if token_output.finished: + if tool_parser is not None: + content_parts, tool_calls = tool_parser.finalize() + for content_part in content_parts: + visible_content = strip_reasoning_markers(content_part) + if not visible_content: + continue + chunk = chunk_json( + request_id, content=content_part, model=self.model_id + ) + yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n" + for tool_call in tool_calls: + delta_call = {"index": tool_call_index, **tool_call} + tool_call_index += 1 + chunk = chunk_json( + request_id, + tool_calls=[delta_call], + model=self.model_id, + ) + yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n" finish_reason = self._convert_finish_reason( token_output.finish_reason ) + if tool_parser is not None and tool_parser.has_tool_calls: + finish_reason = "tool_calls" + usage = None + stream_options = data.get("stream_options") + if isinstance(stream_options, dict) and stream_options.get( + "include_usage" + ): + usage = { + "prompt_tokens": req.get_prompt_length(), + "completion_tokens": req.get_num_generated_tokens(), + "total_tokens": req.get_total_length(), + } chunk = json.dumps( chunk_json( - request_id, finish_reason=finish_reason, model=self.model_id + request_id, + finish_reason=finish_reason, + model=self.model_id, + usage=usage, ), ensure_ascii=False, ) @@ -496,16 +708,27 @@ async def _chat(self, request_id: str, data: dict, http_request: Request): _abort_reason = FinishReason.CANCELED try: - messages = data.get("messages", []) + messages, chat_template_kwargs, request_data = self._prepare_chat_request( + data + ) sampling_params = self._build_sampling_params(data) + if request_data.get("_infinilm_direct_write_complete"): + return completion_json( + request_id, + content="已写入文件，已按请求停止。", + role="assistant", + finish_reason="stop", + model=self.model_id, + ) + req = self.engine.add_chat_request( messages=messages, sampling_params=sampling_params, request_id=request_id, - request_data=data, + request_data=request_data, add_generation_prompt=bool(data.get("add_generation_prompt", True)), - chat_template_kwargs=data.get("chat_template_kwargs") or {}, + chat_template_kwargs=chat_template_kwargs, ) # Collect all generated tokens @@ -538,6 +761,22 @@ async def _chat(self, request_id: str, data: dict, http_request: Request): output_text = output_text.strip() finish_reason = self._convert_finish_reason(req.finish_reason) + tool_calls = None + if chat_template_kwargs.get("tools"): + allowed_tool_names = { + tool["function"]["name"] for tool in chat_template_kwargs["tools"] + } + forced_tool_prefix = request_data.get( + "_infinilm_forced_tool_prefix", "" + ) + output_text, tool_calls = parse_tool_calls( + forced_tool_prefix + output_text, + allowed_tool_names=allowed_tool_names, + ) + if output_text is not None: + output_text = strip_reasoning_markers(output_text) + if tool_calls: + finish_reason = "tool_calls" response = completion_json( request_id, @@ -548,6 +787,7 @@ async def _chat(self, request_id: str, data: dict, http_request: Request): prompt_tokens=req.get_prompt_length(), completion_tokens=req.get_num_generated_tokens(), total_tokens=req.get_total_length(), + tool_calls=tool_calls, ) return response @@ -558,7 +798,7 @@ async def _chat(self, request_id: str, data: dict, http_request: Request): except Exception as e: logger.error(f"Chat error for {request_id}: {e}", exc_info=True) _abort_reason = FinishReason.ERROR - return JSONResponse(content={"error": str(e)}, status_code=500) + return self._error_response(str(e), status_code=500) finally: # Unified abort: reason is ERROR if we got here via Exception, else CANCELED. diff --git a/python/infinilm/server/openai_protocol.py b/python/infinilm/server/openai_protocol.py new file mode 100644 index 000000000..08d31ed92 --- /dev/null +++ b/python/infinilm/server/openai_protocol.py @@ -0,0 +1,182 @@ +"""Helpers for translating model tool-call text to OpenAI protocol objects.""" + +import json +import re +import uuid +from typing import Optional + +_TOOL_CALL_OPEN = "" +_TOOL_CALL_CLOSE = "" + + +def strip_reasoning_markers(text: str) -> str: + """Remove model reasoning delimiters from visible chat content.""" + return text.replace("", "").replace("", "") + + +_TOOL_BLOCK_RE = re.compile(r"\s*(.*?)\s*", re.DOTALL) +_FUNCTION_RE = re.compile(r"\n]+)>\s*(.*?)\s*", re.DOTALL) +_PARAMETER_RE = re.compile(r"\n]+)>\s*(.*?)\s*", re.DOTALL) + + +def _new_tool_call(name: str, arguments) -> dict: + if isinstance(arguments, str): + try: + parsed_arguments = json.loads(arguments) + except json.JSONDecodeError: + arguments_json = arguments + else: + arguments_json = json.dumps( + parsed_arguments, ensure_ascii=False, separators=(",", ":") + ) + else: + arguments_json = json.dumps( + arguments, ensure_ascii=False, separators=(",", ":") + ) + + return { + "id": f"call_{uuid.uuid4().hex[:24]}", + "type": "function", + "function": { + "name": name.strip(), + "arguments": arguments_json, + }, + } + + +def _decode_parameter(value: str): + value = value.strip() + try: + return json.loads(value) + except json.JSONDecodeError: + return value + + +def _parse_tool_block(body: str) -> list[dict]: + calls = [] + for function_match in _FUNCTION_RE.finditer(body): + arguments = {} + for parameter_match in _PARAMETER_RE.finditer(function_match.group(2)): + arguments[parameter_match.group(1).strip()] = _decode_parameter( + parameter_match.group(2) + ) + calls.append(_new_tool_call(function_match.group(1), arguments)) + + if calls: + return calls + + try: + payload = json.loads(body.strip()) + except json.JSONDecodeError: + return [] + + payloads = payload if isinstance(payload, list) else [payload] + for item in payloads: + if not isinstance(item, dict): + continue + function = item.get("function", item) + if not isinstance(function, dict) or not function.get("name"): + continue + calls.append(_new_tool_call(function["name"], function.get("arguments", {}))) + return calls + + +def parse_tool_calls( + text: str, allowed_tool_names: Optional[set[str]] = None +) -> tuple[Optional[str], list[dict]]: + """Parse Qwen/Hermes tool-call blocks and return OpenAI tool calls.""" + tool_calls = [] + content_parts = [] + cursor = 0 + saw_tool_block = False + + for match in _TOOL_BLOCK_RE.finditer(text): + parsed = _parse_tool_block(match.group(1)) + if not parsed: + continue + saw_tool_block = True + content_parts.append(text[cursor : match.start()]) + cursor = match.end() + if allowed_tool_names is not None: + parsed = [ + call + for call in parsed + if call["function"]["name"] in allowed_tool_names + ] + tool_calls.extend(parsed) + + if not saw_tool_block: + content = text.strip() + return content or None, [] + + content_parts.append(text[cursor:]) + content = "".join(content_parts).strip() + return content or None, tool_calls + + +class ToolCallStreamParser: + """Incrementally hide XML markers and emit parsed OpenAI tool calls.""" + + def __init__(self, allowed_tool_names: Optional[set[str]] = None): + self._buffer = "" + self.has_tool_calls = False + self.allowed_tool_names = allowed_tool_names + + @staticmethod + def _marker_suffix_length(value: str) -> int: + max_length = min(len(value), len(_TOOL_CALL_OPEN) - 1) + for length in range(max_length, 0, -1): + if _TOOL_CALL_OPEN.startswith(value[-length:]): + return length + return 0 + + def feed(self, chunk: str) -> tuple[list[str], list[dict]]: + self._buffer += chunk + content_parts = [] + tool_calls = [] + + while self._buffer: + start = self._buffer.find(_TOOL_CALL_OPEN) + if start < 0: + keep = self._marker_suffix_length(self._buffer) + safe_length = len(self._buffer) - keep + if safe_length: + content_parts.append(self._buffer[:safe_length]) + self._buffer = self._buffer[safe_length:] + break + + if start: + content_parts.append(self._buffer[:start]) + self._buffer = self._buffer[start:] + + end = self._buffer.find(_TOOL_CALL_CLOSE) + if end < 0: + break + end += len(_TOOL_CALL_CLOSE) + block = self._buffer[:end] + self._buffer = self._buffer[end:] + content, parsed = parse_tool_calls( + block, allowed_tool_names=self.allowed_tool_names + ) + if parsed: + self.has_tool_calls = True + tool_calls.extend(parsed) + elif self.allowed_tool_names is not None and content is None: + # A complete tool block was rejected because its function was + # not exposed for this request. Do not leak its XML markup. + pass + else: + content_parts.append(content if content is not None else block) + + return content_parts, tool_calls + + def finalize(self) -> tuple[list[str], list[dict]]: + if not self._buffer: + return [], [] + content, tool_calls = parse_tool_calls( + self._buffer, allowed_tool_names=self.allowed_tool_names + ) + self._buffer = "" + if tool_calls: + self.has_tool_calls = True + return ([content] if content else []), tool_calls diff --git a/python/infinilm/server/tool_constraints.py b/python/infinilm/server/tool_constraints.py new file mode 100644 index 000000000..b5fe394ce --- /dev/null +++ b/python/infinilm/server/tool_constraints.py @@ -0,0 +1,243 @@ +"""Request-local tool constraints for OpenAI-compatible agent clients.""" + +import json +import re +from collections import Counter +from typing import Optional + +_ALLOW_PATTERNS = ( + re.compile(r"(?:仅|只)(?:允许|能|可)?(?:使用|用)\s*([^；;。\n]+)", re.I), + re.compile(r"\bonly\s+(?:use|allow)\s+([^.;\n]+)", re.I), +) +_BLOCK_PATTERNS = ( + re.compile(r"(?:禁止|不得|禁用|不要使用)\s*([^；;。\n]+)", re.I), + re.compile(r"\b(?:do\s+not\s+use|don't\s+use|forbid|disable)\s+([^.;\n]+)", re.I), +) +_IDENTIFIER_RE = re.compile(r"[A-Za-z][A-Za-z0-9_-]*") +_NAMED_FILE_ZH_RE = re.compile( + r"(?:写入|写到|保存到)\s*[^\s；;。\n]+\.[A-Za-z0-9]+", re.I +) +_NAMED_FILE_EN_RE = re.compile( + r"\b(?:write|save)\b[^.\n]*\b(?:to|into)\s+[^\s.;\n]+\.[A-Za-z0-9]+", + re.I, +) +_STOP_AFTER_WRITE_ZH_RE = re.compile( + r"(?:写完|完成后)[^；;。\n]*(?:立即)?(?:停止|停下)", re.I +) +_STOP_AFTER_WRITE_EN_RE = re.compile( + r"(?:\b(?:stop|wait)\b[^.\n]*\bafter\b[^.\n]*(?:writ|sav)|" + r"\bafter\b[^.\n]*(?:writ|sav)[^.\n]*\b(?:stop|wait)\b)", + re.I, +) +_FAILED_TOOL_RESULT_MARKERS = ( + "tool_error", + "tool execution failed", + '"iserror":true', + '"is_error":true', +) + + +def _normalize_tool_name(value: str) -> str: + return re.sub(r"[-_]", "", value.lower()) + + +def _tool_aliases(name: str) -> set[str]: + normalized = _normalize_tool_name(name) + aliases = {normalized} + if normalized.endswith("file"): + aliases.add(normalized[:-4]) + return aliases + + +def _message_text(message: dict) -> str: + content = message.get("content") + if isinstance(content, str): + return content.strip() + if not isinstance(content, list): + return "" + parts = [] + for item in content: + if isinstance(item, str): + parts.append(item) + elif isinstance(item, dict) and item.get("type") == "text": + text = item.get("text") + if isinstance(text, str): + parts.append(text) + return "\n".join(parts).strip() + + +_SYNTHETIC_USER_RE = re.compile( + r"^(?:Output token limit hit[.]|Your previous response was empty " + r"[(]thinking only, no visible text[)][.])", + re.I, +) + + +def _latest_user_text(messages: list) -> Optional[str]: + for message in reversed(messages): + if not isinstance(message, dict) or message.get("role") != "user": + continue + text = _message_text(message) + if text and not _SYNTHETIC_USER_RE.search(text): + return text + return None + + +def _extract_tool_names(text: str, patterns, tools: list) -> set[str]: + lookup = {} + for tool in tools: + name = tool.get("function", {}).get("name") + if not isinstance(name, str): + continue + for alias in _tool_aliases(name): + lookup[alias] = name + + names = set() + for pattern in patterns: + for match in pattern.finditer(text): + for identifier in _IDENTIFIER_RE.findall(match.group(1)): + name = lookup.get(_normalize_tool_name(identifier)) + if name: + names.add(name) + return names + + +def _is_direct_write_then_stop(text: str) -> bool: + requests_file = bool( + _NAMED_FILE_ZH_RE.search(text) or _NAMED_FILE_EN_RE.search(text) + ) + requests_stop = bool( + _STOP_AFTER_WRITE_ZH_RE.search(text) or _STOP_AFTER_WRITE_EN_RE.search(text) + ) + return requests_file and requests_stop + + +def _canonical_arguments(arguments) -> str: + if isinstance(arguments, str): + try: + arguments = json.loads(arguments) + except json.JSONDecodeError: + return arguments + return json.dumps( + arguments, ensure_ascii=False, sort_keys=True, separators=(",", ":") + ) + + +def _tool_history(messages: list) -> tuple[Counter, dict[str, str]]: + signatures = Counter() + names_by_id = {} + for message in messages: + if not isinstance(message, dict) or message.get("role") != "assistant": + continue + for call in message.get("tool_calls") or []: + if not isinstance(call, dict): + continue + function = call.get("function") or {} + name = function.get("name") + if not isinstance(name, str): + continue + signatures[(name, _canonical_arguments(function.get("arguments", {})))] += 1 + call_id = call.get("id") + if isinstance(call_id, str): + names_by_id[call_id] = name + return signatures, names_by_id + + +def _has_successful_write(messages: list, names_by_id: dict[str, str]) -> bool: + for message in messages: + if not isinstance(message, dict) or message.get("role") != "tool": + continue + if names_by_id.get(message.get("tool_call_id")) != "write_file": + continue + result = _message_text(message).lower().replace(" ", "") + if not any( + marker.replace(" ", "") in result for marker in _FAILED_TOOL_RESULT_MARKERS + ): + return True + return False + + +def constrain_tools(messages: list, tools: list) -> list: + """Filter schemas using explicit user constraints and request history.""" + if not tools: + return tools + text = _latest_user_text(messages) + if not text: + return tools + + allowed = _extract_tool_names(text, _ALLOW_PATTERNS, tools) + blocked = _extract_tool_names(text, _BLOCK_PATTERNS, tools) + constrained = [ + tool + for tool in tools + if (not allowed or tool["function"]["name"] in allowed) + and tool["function"]["name"] not in blocked + ] + + signatures, names_by_id = _tool_history(messages) + repeated_names = {name for (name, _), count in signatures.items() if count >= 2} + constrained = [ + tool for tool in constrained if tool["function"]["name"] not in repeated_names + ] + + if _is_direct_write_then_stop(text): + if _has_successful_write(messages, names_by_id): + return [] + write_tools = [ + tool for tool in constrained if tool["function"]["name"] == "write_file" + ] + if write_tools: + return write_tools + return constrained + + +_NAMED_FILE_TARGET_ZH_RE = re.compile( + r"(?:写入|写到|保存到)\s*([^\s；;。\n]+\.[A-Za-z0-9]+)", re.I +) +_NAMED_FILE_TARGET_EN_RE = re.compile( + r"\b(?:write|save)\b[^.\n]*\b(?:to|into)\s+([^\s.;\n]+\.[A-Za-z0-9]+)", + re.I, +) + + +def forced_write_tool_prefix(messages: list, tools: list) -> Optional[str]: + exposed = { + tool.get("function", {}).get("name") for tool in tools if isinstance(tool, dict) + } + if exposed != {"write_file"}: + return None + + text = _latest_user_text(messages) + if not text or not _is_direct_write_then_stop(text): + return None + + match = _NAMED_FILE_TARGET_ZH_RE.search(text) or _NAMED_FILE_TARGET_EN_RE.search( + text + ) + if not match: + return None + target = match.group(1).strip() + if not target: + return None + + lines = [ + "", + "", + "", + target, + "", + "", + ] + return chr(10).join(lines) + chr(10) + + +def direct_write_complete(messages: list, tools: list) -> bool: + if not tools: + return False + + text = _latest_user_text(messages) + if not text or not _is_direct_write_then_stop(text): + return False + + _, names_by_id = _tool_history(messages) + return _has_successful_write(messages, names_by_id) diff --git a/python/infinilm/server/tool_contract.py b/python/infinilm/server/tool_contract.py new file mode 100644 index 000000000..528aff411 --- /dev/null +++ b/python/infinilm/server/tool_contract.py @@ -0,0 +1,70 @@ +import copy + + +def _append_contract(message: dict, contract: str) -> None: + content = message.get("content") + if isinstance(content, str): + message["content"] = content.rstrip() + chr(10) + chr(10) + contract + elif isinstance(content, list): + content.append({"type": "text", "text": contract}) + else: + message["content"] = contract + + +def apply_tool_contract( + messages: list, original_tools: list, exposed_tools: list +) -> list: + """Make a server-side tool filter explicit in the system context.""" + original_names = { + tool.get("function", {}).get("name") + for tool in original_tools + if isinstance(tool, dict) + } + exposed_names = { + tool.get("function", {}).get("name") + for tool in exposed_tools + if isinstance(tool, dict) + } + if not exposed_names or original_names == exposed_names: + return messages + + quoted_names = ", ".join(f"`{name}`" for name in sorted(exposed_names)) + single_tool_instruction = ( + f" Create the requested deliverable content and call `{next(iter(exposed_names))}` " + "directly." + if len(exposed_names) == 1 + else " When an action is required, choose one tool from this exact set." + ) + contract = ( + "\n" + f"The complete set of tools exposed by the runtime is exactly: {quoted_names}. " + "Do not call, mention, simulate, or substitute any other tool." + f"{single_tool_instruction}\n" + "" + ) + + normalized_messages = copy.deepcopy(messages) + system_message = next( + ( + message + for message in normalized_messages + if isinstance(message, dict) and message.get("role") == "system" + ), + None, + ) + if system_message is None: + normalized_messages.insert(0, {"role": "system", "content": contract}) + else: + _append_contract(system_message, contract) + + user_message = next( + ( + message + for message in reversed(normalized_messages) + if isinstance(message, dict) and message.get("role") == "user" + ), + None, + ) + if user_message is not None: + _append_contract(user_message, contract) + return normalized_messages diff --git a/test/service/test_openai_protocol.py b/test/service/test_openai_protocol.py new file mode 100644 index 000000000..3c675b4e6 --- /dev/null +++ b/test/service/test_openai_protocol.py @@ -0,0 +1,456 @@ +import json + +import pytest +from infinilm.processors.processor import normalize_openai_messages +from infinilm.server.inference_server import InferenceServer, completion_json +from infinilm.server.openai_protocol import ( + ToolCallStreamParser, + parse_tool_calls, + strip_reasoning_markers, +) +from infinilm.server.tool_contract import apply_tool_contract +from infinilm.server.tool_constraints import constrain_tools, forced_write_tool_prefix + + +def test_parse_qwen_tool_call_xml(): + content, calls = parse_tool_calls( + "checking\n\n\n" + "\nBeijing\n\n" + "\n3\n\n" + "\n" + ) + + assert content == "checking" + assert len(calls) == 1 + assert calls[0]["type"] == "function" + assert calls[0]["function"]["name"] == "get_weather" + assert json.loads(calls[0]["function"]["arguments"]) == { + "city": "Beijing", + "days": 3, + } + + +def test_parse_json_tool_call_xml(): + content, calls = parse_tool_calls( + '{"name":"lookup","arguments":{"query":"InfiniLM"}}' + ) + + assert content is None + assert calls[0]["function"]["name"] == "lookup" + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "InfiniLM"} + + +def test_stream_parser_handles_split_markers_without_leaking_xml(): + parser = ToolCallStreamParser() + content_parts = [] + calls = [] + chunks = [ + "I will check.\n\n\n\n", + "PilotDeck\n\n\n", + ] + for chunk in chunks: + content, parsed = parser.feed(chunk) + content_parts.extend(content) + calls.extend(parsed) + content, parsed = parser.finalize() + content_parts.extend(content) + calls.extend(parsed) + + assert "".join(content_parts) == "I will check.\n" + assert parser.has_tool_calls + assert calls[0]["function"]["name"] == "lookup" + assert json.loads(calls[0]["function"]["arguments"]) == {"query": "PilotDeck"} + + +def test_normalize_openai_tool_call_history_arguments(): + messages = [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "weather", + "arguments": '{"city":"Beijing"}', + }, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ] + + normalized = normalize_openai_messages(messages) + + assert normalized[1]["content"] is None + assert normalized[1]["tool_calls"][0]["function"]["arguments"] == { + "city": "Beijing" + } + assert messages[1]["tool_calls"][0]["function"]["arguments"] == ( + '{"city":"Beijing"}' + ) + + +def test_invalid_tool_history_arguments_are_rejected(): + with pytest.raises(ValueError, match="valid JSON"): + normalize_openai_messages( + [ + { + "role": "assistant", + "tool_calls": [{"function": {"name": "bad", "arguments": "{"}}], + } + ] + ) + + +def test_openai_completion_contains_native_tool_calls(): + _, calls = parse_tool_calls("") + response = completion_json("cmpl-test", None, tool_calls=calls) + + assert response["choices"][0]["message"]["content"] is None + assert response["choices"][0]["message"]["tool_calls"] == calls + + +def test_pilotdeck_request_fields_become_template_kwargs(): + tools = [ + { + "type": "function", + "function": {"name": "lookup", "parameters": {"type": "object"}}, + } + ] + kwargs = InferenceServer._build_chat_template_kwargs( + { + "tools": tools, + "tool_choice": "auto", + "reasoning_effort": "low", + "thinking": {"type": "disabled"}, + } + ) + + assert kwargs["tools"] == tools + assert kwargs["tool_choice"] == "auto" + assert kwargs["reasoning_effort"] == "low" + assert kwargs["enable_thinking"] is False + + +def test_parse_multiple_qwen_tool_calls(): + content, calls = parse_tool_calls( + "one" + "\n" + "two" + "" + ) + + assert content is None + assert [call["function"]["name"] for call in calls] == ["lookup", "lookup"] + assert [json.loads(call["function"]["arguments"])["query"] for call in calls] == [ + "one", + "two", + ] + + +def test_forced_tool_choice_filters_template_tools(): + tools = [ + {"type": "function", "function": {"name": "first", "parameters": {}}}, + {"type": "function", "function": {"name": "second", "parameters": {}}}, + ] + kwargs = InferenceServer._build_chat_template_kwargs( + { + "tools": tools, + "tool_choice": { + "type": "function", + "function": {"name": "second"}, + }, + } + ) + + assert [tool["function"]["name"] for tool in kwargs["tools"]] == ["second"] + + +def test_max_completion_tokens_alias(): + server = InferenceServer("/tmp/model", max_tokens=99) + + params = server._build_sampling_params({"max_completion_tokens": 17}) + + assert params.max_tokens == 17 + + +def _tool(name): + return { + "type": "function", + "function": {"name": name, "parameters": {"type": "object"}}, + } + + +def test_pilotdeck_chinese_constraints_filter_tool_schemas(): + tools = [ + _tool("write_file"), + _tool("edit_file"), + _tool("bash"), + _tool("read_file"), + _tool("read_skill"), + _tool("web_search"), + _tool("agent"), + ] + messages = [ + { + "role": "user", + "content": ( + "把大纲写入 outline.md。约束：禁止 read_skill / WebSearch / " + "Agent；仅用 Write/Edit/Bash/Read；写完立即停止。" + ), + } + ] + + assert [tool["function"]["name"] for tool in constrain_tools(messages, tools)] == [ + "write_file" + ] + + +def test_direct_write_then_stop_removes_tools_after_success(): + tools = [_tool("write_file"), _tool("read_file")] + messages = [ + { + "role": "user", + "content": "把结果写入 outline.md。写完 outline.md 后立即停止。", + }, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "write_file", + "arguments": '{"path":"outline.md","content":"ok"}', + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "Successfully wrote outline.md", + }, + ] + + assert constrain_tools(messages, tools) == [] + + +def test_direct_write_then_stop_allows_retry_after_failed_write(): + tools = [_tool("write_file"), _tool("read_file")] + messages = [ + { + "role": "user", + "content": "把结果写入 outline.md。写完 outline.md 后立即停止。", + }, + { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "function": {"name": "write_file", "arguments": "{}"}, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "TOOL_ERROR: permission denied", + }, + ] + + assert [tool["function"]["name"] for tool in constrain_tools(messages, tools)] == [ + "write_file" + ] + + +def test_repeated_identical_calls_remove_the_tool_from_later_rounds(): + tools = [_tool("bash"), _tool("write_file")] + call = { + "id": "call_1", + "function": {"name": "bash", "arguments": '{"command":"pwd"}'}, + } + messages = [ + {"role": "user", "content": "Inspect and write the result."}, + {"role": "assistant", "tool_calls": [call]}, + {"role": "tool", "tool_call_id": "call_1", "content": "ok"}, + {"role": "assistant", "tool_calls": [{**call, "id": "call_2"}]}, + {"role": "tool", "tool_call_id": "call_2", "content": "ok"}, + ] + + assert [tool["function"]["name"] for tool in constrain_tools(messages, tools)] == [ + "write_file" + ] + + +def test_parser_rejects_tool_calls_not_exposed_for_this_request(): + content, calls = parse_tool_calls( + "x" + "", + allowed_tool_names={"write_file"}, + ) + + assert content is None + assert calls == [] + + +def test_filtered_tool_schema_adds_explicit_system_contract(): + messages = [ + {"role": "system", "content": "agent instructions"}, + {"role": "user", "content": "write the result"}, + ] + result = apply_tool_contract( + messages, + [_tool("write_file"), _tool("bash")], + [_tool("write_file")], + ) + + assert "`write_file`" in result[0]["content"] + assert "`bash`" not in result[1]["content"] + assert messages[0]["content"] == "agent instructions" + + +def test_unfiltered_tool_schema_does_not_modify_messages(): + messages = [{"role": "user", "content": "write the result"}] + tools = [_tool("write_file")] + + assert apply_tool_contract(messages, tools, tools) is messages + + +def test_stream_parser_hides_rejected_complete_tool_block(): + parser = ToolCallStreamParser(allowed_tool_names={"write_file"}) + content_parts, calls = parser.feed( + "" + "x" + ) + final_content, final_calls = parser.finalize() + content_parts.extend(final_content) + calls.extend(final_calls) + + assert content_parts == [] + assert calls == [] + + +def test_constraints_survive_pilotdeck_recovery_messages(): + tools = [_tool("write_file"), _tool("read_file"), _tool("bash")] + messages = [ + { + "role": "user", + "content": "把结果写入 outline.md。写完 outline.md 后立即停止。", + }, + {"role": "assistant", "content": None}, + { + "role": "user", + "content": ( + "Your previous response was empty (thinking only, no visible " + "text). Please provide your answer as visible text output." + ), + }, + ] + + assert [tool["function"]["name"] for tool in constrain_tools(messages, tools)] == [ + "write_file" + ] + + +def test_tool_contract_is_repeated_for_current_user_round(): + messages = [ + {"role": "system", "content": "agent instructions"}, + {"role": "user", "content": "write the result"}, + ] + result = apply_tool_contract( + messages, + [_tool("write_file"), _tool("bash")], + [_tool("write_file")], + ) + + assert result[0]["content"].count("") == 1 + assert result[1]["content"].count("") == 1 + + +def test_forced_write_prefix_constrains_generation_to_write_file(): + messages = [{"role": "user", "content": "把结果写入 outline.md。写完后立即停止。"}] + prefix = forced_write_tool_prefix(messages, [_tool("write_file")]) + + assert prefix == ( + "\n\n\n" + "outline.md\n\n\n" + ) + + parser = ToolCallStreamParser(allowed_tool_names={"write_file"}) + parser.feed(prefix) + content_parts, calls = parser.feed( + "outline content\n\n\n" + ) + final_content, final_calls = parser.finalize() + content_parts.extend(final_content) + calls.extend(final_calls) + + assert content_parts == [] + assert calls[0]["function"]["name"] == "write_file" + assert json.loads(calls[0]["function"]["arguments"]) == { + "file_path": "outline.md", + "content": "outline content", + } + + +def test_prepared_request_carries_forced_tool_prefix(): + server = InferenceServer("/tmp/model", skip_load=True) + messages, kwargs, request_data = server._prepare_chat_request( + { + "messages": [ + {"role": "user", "content": "把结果写入 outline.md。写完后立即停止。"} + ], + "tools": [_tool("write_file"), _tool("read_file")], + } + ) + + assert [tool["function"]["name"] for tool in kwargs["tools"]] == ["write_file"] + assert request_data["_infinilm_forced_tool_prefix"].startswith("\n") + + +def test_successful_direct_write_marks_next_round_complete(): + server = InferenceServer("/tmp/model", skip_load=True) + messages, kwargs, request_data = server._prepare_chat_request( + { + "messages": [ + { + "role": "user", + "content": "把结果写入 outline.md。写完后立即停止。", + }, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "function": { + "name": "write_file", + "arguments": "{}", + }, + } + ], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": "Created outline.md.", + }, + ], + "tools": [_tool("write_file"), _tool("read_file")], + } + ) + + assert kwargs == {} + assert request_data["_infinilm_direct_write_complete"] is True + + +def test_strip_reasoning_markers_from_visible_content(): + assert ( + strip_reasoning_markers("hiddenvisible") + == "hiddenvisible" + ) diff --git a/test/service/test_request_capacity.py b/test/service/test_request_capacity.py new file mode 100644 index 000000000..0b43fa4ab --- /dev/null +++ b/test/service/test_request_capacity.py @@ -0,0 +1,40 @@ +import pytest +from infinilm.config.engine_config import EngineConfig +from infinilm.llm.llm import validate_request_capacity + + +def paged_config(num_blocks: int, block_size: int) -> EngineConfig: + return EngineConfig( + model_path="/tmp/model", + cache_type="paged", + num_blocks=num_blocks, + block_size=block_size, + ) + + +def test_pilotdeck_prompt_fits_configured_kv_capacity(): + validate_request_capacity( + paged_config(num_blocks=256, block_size=256), + prompt_tokens=8_348, + output_tokens=512, + ) + + +def test_prompt_larger_than_kv_capacity_is_rejected(): + with pytest.raises( + ValueError, match="maximum context length is 8192.*prompt contains 8348" + ): + validate_request_capacity( + paged_config(num_blocks=128, block_size=64), + prompt_tokens=8_348, + output_tokens=512, + ) + + +def test_output_reservation_larger_than_remaining_capacity_is_rejected(): + with pytest.raises(ValueError, match="max_tokens must be at most 192"): + validate_request_capacity( + paged_config(num_blocks=128, block_size=64), + prompt_tokens=8_000, + output_tokens=512, + )