Feat: agentic search framework (#16859)

### Summary

Agentic search

<img width="1149" height="1575" alt="image"
src="https://github.com/user-attachments/assets/bce9a3e7-0517-4fb2-80a2-5d2a81a4da78"
/>

---------

Co-authored-by: Yingfeng Zhang <yingfeng.zhang@gmail.com>

rag/llm/chat_model.py

@@ -679,8 +679,22 @@ class Base(ABC):
                             args = json_repair.loads(tc.function.arguments)
                         except Exception:
                             args = {}
-                        yield self._verbose_tool_use(tc.function.name, args, "Begin to call...")
+                        yield f"<think>Executing {tc.function.name} with args: {tc.function.arguments}</think>"
                     results = await asyncio.gather(*[_exec_tool(tc) for tc in tcs])
+
+                    # Terminal-tool short-circuit: stream a terminal tool's
+                    # result (already the final answer) and stop the loop.
+                    _terminal = getattr(self, "terminal_tools", None)
+                    if _terminal:
+                        for tc, name, args, result, err in results:
+                            if name in _terminal and not err:
+                                logging.info(f"[ToolLoop] terminal tool {name!r} called; streaming result and stopping")
+                                out = result if isinstance(result, str) else json.dumps(result, ensure_ascii=False)
+                                if out:
+                                    yield out
+                                yield total_tokens
+                                return
+
                     history = self._append_history_batch(history, results)
                     for tc, name, args, result, err in results:
                         yield self._verbose_tool_use(name, args, err if err else result)
@@ -1925,7 +1939,7 @@ class LiteLLMBase(ABC):
             history = deepcopy(hist)
             try:
                 for _ in range(self.max_rounds + 1):
-                    logging.info(f"{self.tools=}")
+                    logging.info(f"HAS TOOL:{len(self.tools)}\n{history=}")
 
                     completion_args = self._construct_completion_args(history=history, stream=False, tools=True, **gen_conf)
                     response = await litellm.acompletion(
@@ -2122,8 +2136,23 @@ class LiteLLMBase(ABC):
                             args = json_repair.loads(tc.function.arguments)
                         except Exception:
                             args = {}
-                        yield self._verbose_tool_use(tc.function.name, args, "Begin to call...")
+                        yield f"<think>Executing {tc.function.name} with args: {tc.function.arguments}</think>"
                     results = await asyncio.gather(*[_exec_tool(tc) for tc in tcs])
+
+                    # Terminal-tool short-circuit: a terminal tool already
+                    # produces the final answer, so stream its result and stop
+                    # instead of feeding it back for another LLM round.
+                    _terminal = getattr(self, "terminal_tools", None)
+                    if _terminal:
+                        for tc, name, args, result, err in results:
+                            if name in _terminal and not err:
+                                logging.info(f"[ToolLoop] terminal tool {name!r} called; streaming result and stopping")
+                                out = result if isinstance(result, str) else json.dumps(result, ensure_ascii=False)
+                                if out:
+                                    yield out
+                                yield total_tokens
+                                return
+
                     history = self._append_history_batch(
                         history,
                         results,

rag/llm/tool_decorator.py

@@ -91,14 +91,33 @@ def _py_type_to_json(py_type: Any) -> dict[str, Any]:
     return {"type": "string"}
 
 
-_PARAM_RE = re.compile(r"^\s*:param\s+(?P<name>\w+)\s*:\s*(?P<desc>.+?)\s*$")
+_PARAM_RE = re.compile(r"^\s*:param\s+(?P<name>\w+)\s*:\s*(?P<desc>.*?)\s*$")
+_GOOGLE_ARGS_HDR_RE = re.compile(r"^(Args|Arguments|Parameters)\s*:\s*$")
+_GOOGLE_SECTION_HDR_RE = re.compile(
+    r"^(Returns?|Yields?|Raises|Notes?|Examples?|Attributes?|Todo|See Also|Warning|Warnings|Tip)\s*:\s*$"
+)
+# Google-style parameter line: leading indent, identifier, optional ``(type)``,
+# then ``: description``. The description can be empty (continuation lines fill it).
+_GOOGLE_PARAM_RE = re.compile(
+    r"^(?P<indent>\s+)(?P<name>\w+)\s*(?:\([^)]*\))?\s*:\s*(?P<desc>.*)$"
+)
 
 
 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.
+    """Pull a function description and per-parameter descriptions out of a docstring.
 
-    Intentionally minimal — Google/NumPy styles are not parsed. Anything
-    before the first ``:param`` line becomes the function description.
+    Recognises two conventions and handles multi-line descriptions in both:
+
+    * **reST / Sphinx**: ``:param name: description`` followed by deeper-indented
+      continuation lines.
+    * **Google**: an ``Args:`` (or ``Arguments:`` / ``Parameters:``) section
+      whose body is ``    name: description`` lines, with deeper-indented
+      continuation lines folded onto the same entry. Other Google sections
+      (``Returns:``, ``Raises:``, ...) terminate the description but are
+      otherwise dropped — they aren't sent to the LLM.
+
+    Both styles can co-exist in one docstring. Anything before the first
+    parameter entry / section header becomes the function description.
     """
     if not docstring:
         return "", {}
@@ -106,12 +125,65 @@ def _parse_param_docs(docstring: str | None) -> tuple[str, dict[str, str]]:
     lines = inspect.cleandoc(docstring).splitlines()
     desc_lines: list[str] = []
     param_docs: dict[str, str] = {}
+    state = "desc"  # "desc" | "rst_param" | "google_args" | "other_section"
+    current_param: str | None = None
+    current_indent = 0
+    after_first_param = False
+
+    def _append_continuation(name: str, text: str) -> None:
+        param_docs[name] = (param_docs[name] + " " + text).strip() if param_docs.get(name) else text
+
     for line in lines:
+        stripped = line.strip()
+        line_indent = len(line) - len(line.lstrip())
+
+        # reST :param: line — works in any state, resets it.
         m = _PARAM_RE.match(line)
         if m:
-            param_docs[m.group("name")] = m.group("desc")
-        elif not param_docs:
+            current_param = m.group("name")
+            current_indent = line_indent
+            param_docs[current_param] = m.group("desc").strip()
+            state = "rst_param"
+            after_first_param = True
+            continue
+
+        # Google section headers.
+        if _GOOGLE_ARGS_HDR_RE.match(stripped):
+            state = "google_args"
+            current_param = None
+            after_first_param = True
+            continue
+        if _GOOGLE_SECTION_HDR_RE.match(stripped):
+            state = "other_section"
+            current_param = None
+            after_first_param = True
+            continue
+
+        # Google ``    name: desc`` entry inside an Args block.
+        if state == "google_args":
+            gm = _GOOGLE_PARAM_RE.match(line)
+            if gm:
+                current_param = gm.group("name")
+                current_indent = line_indent
+                param_docs[current_param] = gm.group("desc").strip()
+                continue
+
+        # Continuation line for the most recent reST or Google param.
+        if state in ("rst_param", "google_args") and current_param and stripped:
+            if line_indent > current_indent:
+                _append_continuation(current_param, stripped)
+                continue
+
+        # Blank line ends the current param's continuation but stays in-state.
+        if not stripped:
+            current_param = None
+            continue
+
+        # Lines outside any param block, before the first param/section,
+        # accumulate as the function description.
+        if not after_first_param:
             desc_lines.append(line)
+
     return "\n".join(desc_lines).strip(), param_docs
 
 
@@ -153,19 +225,58 @@ def _build_openai_schema(fn: Callable[..., Any]) -> dict[str, Any]:
     }
 
 
-def tool(fn: Callable[..., Any]) -> Callable[..., Any]:
+# Sentinel separating "caller did not pass a timeout" from "caller passed None
+# (= run forever)". Plain ``None`` is a legal value for the kwarg.
+_TIMEOUT_UNSET: Any = object()
+
+
+def tool(
+    fn: Callable[..., Any] | None = None,
+    *,
+    timeout: float | int | None = _TIMEOUT_UNSET,
+) -> 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:
+    Usable in two styles:
+
+    * Bare:        ``@tool``               — no per-tool timeout; the session
+                                              falls back to its caller-supplied
+                                              ``request_timeout`` (default 10s).
+    * Parameterised: ``@tool(timeout=60)``  — 60s timeout, overrides the
+                                              session's default for this tool.
+                                              Pass ``timeout=None`` to disable
+                                              the timeout entirely (the tool
+                                              runs until it completes).
+
+    The wrapped callable is the same callable — we only set attributes on it:
 
     * ``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._tool_timeout`` (only when ``timeout=`` was passed) — read by
+      :class:`FunctionToolSession` to override its default timeout for this
+      tool. May be ``None`` to mean "no timeout".
     """
-    fn.openai_schema = _build_openai_schema(fn)  # type: ignore[attr-defined]
-    fn._is_tool = True  # type: ignore[attr-defined]
-    return fn
+
+    def decorate(f: Callable[..., Any]) -> Callable[..., Any]:
+        f.openai_schema = _build_openai_schema(f)  # type: ignore[attr-defined]
+        f._is_tool = True  # type: ignore[attr-defined]
+        if timeout is not _TIMEOUT_UNSET:
+            f._tool_timeout = timeout  # type: ignore[attr-defined]
+        return f
+
+    # ``@tool`` (no parens) — ``fn`` is the function being decorated.
+    if fn is not None:
+        if not callable(fn):
+            raise TypeError(
+                "@tool used incorrectly. Use `@tool` or `@tool(timeout=N)`; "
+                f"got first positional argument of type {type(fn).__name__}."
+            )
+        return decorate(fn)
+
+    # ``@tool(timeout=N)`` — return the decorator that will receive the function.
+    return decorate
 
 
 def is_tool(obj: Any) -> bool:
@@ -188,21 +299,25 @@ class FunctionToolSession:
         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")
+                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:
+    def tool_call(self, name: str, arguments: dict[str, Any], timeout: float | int = 300) -> 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:
+    async def tool_call_async(self, name: str, arguments: dict[str, Any], request_timeout: float | int = 300) -> 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__}")
+            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):
@@ -214,4 +329,9 @@ class FunctionToolSession:
             # 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)
+        # Per-tool timeout set via ``@tool(timeout=N)`` overrides the
+        # session-default. ``None`` is a legal explicit choice meaning
+        # "wait forever" — ``asyncio.wait_for(..., timeout=None)`` handles it.
+        configured = getattr(fn, "_tool_timeout", _TIMEOUT_UNSET)
+        effective_timeout = request_timeout if configured is _TIMEOUT_UNSET else configured
+        return await asyncio.wait_for(coro, timeout=effective_timeout)