diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py index 1d9612dff..5edff15c4 100644 --- a/rag/llm/chat_model.py +++ b/rag/llm/chat_model.py @@ -33,6 +33,7 @@ from enum import StrEnum from common.misc_utils import thread_pool_exec from common.token_utils import num_tokens_from_string, total_token_count_from_response from rag.llm import FACTORY_DEFAULT_BASE_URL, LITELLM_PROVIDER_PREFIX, SupportedLiteLLMProvider +from rag.llm.tool_decorator import FunctionToolSession, is_tool from rag.nlp import is_chinese, is_english @@ -350,7 +351,29 @@ class Base(ABC): hist.append({"role": "tool", "tool_call_id": tc.id, "content": content}) return hist - def bind_tools(self, toolcall_session, tools): + def bind_tools(self, toolcall_session=None, tools=None): + """Register tools the LLM can call. + + Two calling styles are accepted: + + * Legacy: ``bind_tools(toolcall_session, tools_schemas)`` where + ``toolcall_session`` implements :class:`ToolCallSession` and + ``tools_schemas`` is a pre-built list of OpenAI function-schema + dicts (used by the agent/dialog layer). + * Decorator: ``bind_tools(tools=[fn1, fn2, ...])`` where each ``fn`` + is decorated with :func:`rag.llm.tool_decorator.tool`. The session + and schemas are derived from the callables automatically. + """ + if tools is None and isinstance(toolcall_session, list): + tools, toolcall_session = toolcall_session, None + + if tools and toolcall_session is None and all(is_tool(t) for t in tools): + session = FunctionToolSession(tools) + self.is_tools = True + self.toolcall_session = session + self.tools = session.schemas + return + if not (toolcall_session and tools): return self.is_tools = True @@ -1565,7 +1588,29 @@ class LiteLLMBase(ABC): hist.append({"role": "tool", "tool_call_id": tc.id, "content": content}) return hist - def bind_tools(self, toolcall_session, tools): + def bind_tools(self, toolcall_session=None, tools=None): + """Register tools the LLM can call. + + Two calling styles are accepted: + + * Legacy: ``bind_tools(toolcall_session, tools_schemas)`` where + ``toolcall_session`` implements :class:`ToolCallSession` and + ``tools_schemas`` is a pre-built list of OpenAI function-schema + dicts (used by the agent/dialog layer). + * Decorator: ``bind_tools(tools=[fn1, fn2, ...])`` where each ``fn`` + is decorated with :func:`rag.llm.tool_decorator.tool`. The session + and schemas are derived from the callables automatically. + """ + if tools is None and isinstance(toolcall_session, list): + tools, toolcall_session = toolcall_session, None + + if tools and toolcall_session is None and all(is_tool(t) for t in tools): + session = FunctionToolSession(tools) + self.is_tools = True + self.toolcall_session = session + self.tools = session.schemas + return + if not (toolcall_session and tools): return self.is_tools = True diff --git a/rag/llm/tool_decorator.py b/rag/llm/tool_decorator.py new file mode 100644 index 000000000..a1029c62d --- /dev/null +++ b/rag/llm/tool_decorator.py @@ -0,0 +1,221 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Lightweight ``@tool`` decorator and matching ``ToolCallSession`` adapter. + +Lets callers register plain Python functions as LLM tools without having to +hand-write the OpenAI function schema or build an MCP-style session:: + + from rag.llm.tool_decorator import tool + + @tool + def get_weather(city: str) -> str: + \"\"\"Get current weather for a city. + + :param city: City name to look up. + \"\"\" + return f"{city}: 21 C, partly cloudy" + + chat_mdl.bind_tools(tools=[get_weather]) + +The decorator introspects the function signature, type hints, and docstring, +attaches an OpenAI-format schema as ``fn.openai_schema``, and marks the +function with ``fn._is_tool = True`` so :meth:`Base.bind_tools` can detect +the new style. +""" + +from __future__ import annotations + +import asyncio +import inspect +import logging +import re +from collections.abc import Mapping +from typing import Any, Callable, Union, get_args, get_origin, get_type_hints + +from common.misc_utils import thread_pool_exec + + +_PY_TO_JSON: dict[type, str] = { + str: "string", + int: "integer", + float: "number", + bool: "boolean", + list: "array", + dict: "object", + type(None): "null", +} + + +def _py_type_to_json(py_type: Any) -> dict[str, Any]: + """Best-effort mapping from a Python annotation to a JSON-schema fragment. + + Handles ``Optional[T]`` / ``T | None`` by unwrapping the non-None branch + and lets the ``required`` list (built from defaults) carry optionality. + Unknown types fall back to ``{"type": "string"}`` so the schema stays + valid even when annotations are missing. + """ + if py_type is inspect.Parameter.empty or py_type is Any: + return {"type": "string"} + + origin = get_origin(py_type) + if origin is Union: + non_none = [a for a in get_args(py_type) if a is not type(None)] + if len(non_none) == 1: + return _py_type_to_json(non_none[0]) + return {"type": "string"} + + if origin in (list, tuple, set, frozenset): + item_args = get_args(py_type) + item_schema = _py_type_to_json(item_args[0]) if item_args else {"type": "string"} + return {"type": "array", "items": item_schema} + + if origin is dict: + return {"type": "object"} + + if isinstance(py_type, type): + return {"type": _PY_TO_JSON.get(py_type, "string")} + + return {"type": "string"} + + +_PARAM_RE = re.compile(r"^\s*:param\s+(?P\w+)\s*:\s*(?P.+?)\s*$") + + +def _parse_param_docs(docstring: str | None) -> tuple[str, dict[str, str]]: + """Pull a short function description and ``:param name:`` lines out of a docstring. + + Intentionally minimal — Google/NumPy styles are not parsed. Anything + before the first ``:param`` line becomes the function description. + """ + if not docstring: + return "", {} + + lines = inspect.cleandoc(docstring).splitlines() + desc_lines: list[str] = [] + param_docs: dict[str, str] = {} + for line in lines: + m = _PARAM_RE.match(line) + if m: + param_docs[m.group("name")] = m.group("desc") + elif not param_docs: + desc_lines.append(line) + return "\n".join(desc_lines).strip(), param_docs + + +def _build_openai_schema(fn: Callable[..., Any]) -> dict[str, Any]: + sig = inspect.signature(fn) + try: + hints = get_type_hints(fn) + except Exception: + hints = {} + + description, param_docs = _parse_param_docs(fn.__doc__) + + properties: dict[str, dict[str, Any]] = {} + required: list[str] = [] + for name, param in sig.parameters.items(): + if name in ("self", "cls") or param.kind in ( + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + ): + continue + schema = _py_type_to_json(hints.get(name, param.annotation)) + if name in param_docs: + schema["description"] = param_docs[name] + properties[name] = schema + if param.default is inspect.Parameter.empty: + required.append(name) + + return { + "type": "function", + "function": { + "name": fn.__name__, + "description": description or fn.__name__, + "parameters": { + "type": "object", + "properties": properties, + "required": required, + }, + }, + } + + +def tool(fn: Callable[..., Any]) -> Callable[..., Any]: + """Mark ``fn`` as an LLM tool and attach an OpenAI-format schema to it. + + The wrapped callable is the same callable — we only set two attributes: + + * ``fn._is_tool = True`` — sentinel so :meth:`Base.bind_tools` can tell a + ``@tool`` callable apart from a raw schema dict. + * ``fn.openai_schema`` — the schema dict passed verbatim to the LLM + provider in the ``tools=[...]`` request field. + """ + fn.openai_schema = _build_openai_schema(fn) # type: ignore[attr-defined] + fn._is_tool = True # type: ignore[attr-defined] + return fn + + +def is_tool(obj: Any) -> bool: + return callable(obj) and getattr(obj, "_is_tool", False) + + +class FunctionToolSession: + """Adapter that lets a list of ``@tool``-decorated callables satisfy the + :class:`common.mcp_tool_call_conn.ToolCallSession` protocol used by the + chat model tool loop (duck-typed, no explicit inheritance to avoid + pulling the MCP client SDK into this module's import graph). + + The chat model only ever calls ``tool_call`` / ``tool_call_async`` with + ``(name, arguments)`` — this class looks the name up in ``tools_map`` and + invokes the callable, awaiting it if it is a coroutine and otherwise + pushing it through ``thread_pool_exec`` so the event loop is not blocked. + """ + + def __init__(self, tools: list[Callable[..., Any]]): + self.tools_map: dict[str, Callable[..., Any]] = {} + for fn in tools: + if not is_tool(fn): + raise TypeError( + f"{getattr(fn, '__name__', fn)!r} is not a @tool-decorated callable" + ) + self.tools_map[fn.openai_schema["function"]["name"]] = fn + + @property + def schemas(self) -> list[dict[str, Any]]: + return [fn.openai_schema for fn in self.tools_map.values()] + + def tool_call(self, name: str, arguments: dict[str, Any], timeout: float | int = 10) -> Any: + return asyncio.run(self.tool_call_async(name, arguments, request_timeout=timeout)) + + async def tool_call_async(self, name: str, arguments: dict[str, Any], request_timeout: float | int = 10) -> Any: + if name not in self.tools_map: + raise KeyError(f"Tool {name!r} is not registered") + if not isinstance(arguments, Mapping): + raise TypeError( + f"Tool arguments for {name} must be an object, got {type(arguments).__name__}" + ) + fn = self.tools_map[name] + logging.info(f"[FunctionTool] invoke name={name} args={str(arguments)[:200]}") + if asyncio.iscoroutinefunction(fn): + coro = fn(**arguments) + else: + # Sync callables run in the thread pool. asyncio.wait_for cancels + # the awaiting task on timeout, but Python cannot interrupt the + # underlying worker thread — the function keeps running in the + # background until it returns. Callers should treat sync tools + # that block on I/O accordingly. + coro = thread_pool_exec(fn, **arguments) + return await asyncio.wait_for(coro, timeout=request_timeout) diff --git a/test/unit_test/rag/llm/test_tool_decorator.py b/test/unit_test/rag/llm/test_tool_decorator.py new file mode 100644 index 000000000..d2b843967 --- /dev/null +++ b/test/unit_test/rag/llm/test_tool_decorator.py @@ -0,0 +1,123 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Smoke tests for the ``@tool`` decorator and ``FunctionToolSession`` adapter. + +Covers the contract that :meth:`rag.llm.chat_model.Base.bind_tools` relies +on: each ``@tool`` callable carries a well-formed OpenAI function schema, +required vs. optional params are derived from defaults, and the session +dispatches both sync and async callables by name. +""" +import asyncio + +import pytest + +from rag.llm.tool_decorator import FunctionToolSession, is_tool, tool + + +@tool +def get_weather(city: str) -> str: + """Get current weather for a city. + + :param city: City name to look up. + """ + return f"{city}: 21 C, partly cloudy" + + +@tool +def add(a: int, b: int = 1) -> int: + """Return the sum of two numbers. + + :param a: first addend + :param b: second addend (defaults to 1) + """ + return a + b + + +@tool +async def echo_async(text: str) -> str: + """Return the input untouched.""" + return text + + +def test_tool_marks_callable_and_attaches_schema(): + assert is_tool(get_weather) + schema = get_weather.openai_schema + assert schema["type"] == "function" + fn = schema["function"] + assert fn["name"] == "get_weather" + assert "current weather" in fn["description"].lower() + params = fn["parameters"] + assert params["type"] == "object" + assert params["properties"]["city"]["type"] == "string" + assert params["properties"]["city"]["description"] == "City name to look up." + assert params["required"] == ["city"] + + +def test_required_vs_optional_from_defaults(): + params = add.openai_schema["function"]["parameters"] + assert params["properties"]["a"]["type"] == "integer" + assert params["properties"]["b"]["type"] == "integer" + assert params["required"] == ["a"] + + +def test_session_dispatches_sync_tool(): + session = FunctionToolSession([get_weather, add]) + result = session.tool_call("get_weather", {"city": "Beijing"}) + assert "Beijing" in result + assert session.tool_call("add", {"a": 2, "b": 3}) == 5 + + +def test_session_dispatches_async_tool(): + session = FunctionToolSession([echo_async]) + result = asyncio.run(session.tool_call_async("echo_async", {"text": "hi"})) + assert result == "hi" + + +def test_session_rejects_unknown_name(): + session = FunctionToolSession([get_weather]) + with pytest.raises(KeyError): + asyncio.run(session.tool_call_async("nope", {})) + + +def test_session_rejects_non_mapping_arguments(): + session = FunctionToolSession([get_weather]) + with pytest.raises(TypeError): + asyncio.run(session.tool_call_async("get_weather", ["Beijing"])) + + +def test_session_rejects_non_tool_callable(): + def plain(x): return x + + with pytest.raises(TypeError): + FunctionToolSession([plain]) + + +def test_schemas_passthrough_for_bind_tools(): + session = FunctionToolSession([get_weather, add]) + schemas = session.schemas + assert [s["function"]["name"] for s in schemas] == ["get_weather", "add"] + + +def test_async_tool_timeout_raises(): + @tool + async def slow() -> str: + """Sleep longer than the timeout.""" + await asyncio.sleep(0.5) + return "done" + + session = FunctionToolSession([slow]) + with pytest.raises(asyncio.TimeoutError): + asyncio.run(session.tool_call_async("slow", {}, request_timeout=0.05))