chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,34 @@
|
||||
"""Tool-authoring primitives for omnigent.
|
||||
|
||||
Use the :func:`tool` decorator to mark a module-level Python function
|
||||
as a tool the agent can call. The decorator derives the LLM-facing
|
||||
JSON schema from the function's signature and Google-style docstring;
|
||||
the caller just writes Python::
|
||||
|
||||
from omnigent_client import tool
|
||||
|
||||
@tool
|
||||
def get_current_time() -> dict[str, str]:
|
||||
\"\"\"Return the current UTC time as ISO-8601.\"\"\"
|
||||
return {"now": datetime.now(timezone.utc).isoformat()}
|
||||
|
||||
Pass decorated functions as the ``tools=`` argument to
|
||||
:meth:`OmnigentClient.query` or :meth:`Session.query`.
|
||||
|
||||
Server-side runtime (``omnigent.tools.local``) also consumes this
|
||||
decorator to load ``@tool``-decorated functions bundled inside agent
|
||||
images, so the same decorator powers both authoring and runtime.
|
||||
"""
|
||||
|
||||
from ._decorator import TOOL_MARKER_ATTR, ToolMetadata, get_tool_metadata, tool
|
||||
from ._handler import build_tool_handler
|
||||
from ._state import ToolState
|
||||
|
||||
__all__ = [
|
||||
"TOOL_MARKER_ATTR",
|
||||
"ToolMetadata",
|
||||
"ToolState",
|
||||
"build_tool_handler",
|
||||
"get_tool_metadata",
|
||||
"tool",
|
||||
]
|
||||
@@ -0,0 +1,238 @@
|
||||
"""
|
||||
The ``@tool`` decorator and its metadata.
|
||||
|
||||
A ``@tool``-decorated module-level function is the authoring
|
||||
contract for custom Python tools in omnigent. The decorator:
|
||||
|
||||
1. Validates that the target is a module-level ``def`` or
|
||||
``async def`` (rejects class methods, lambdas, and nested
|
||||
functions; see :func:`_validate_decorator_target`).
|
||||
2. Derives the function-calling JSON schema from the signature
|
||||
and Google-style docstring (see :mod:`._schema`).
|
||||
3. Attaches metadata to the function via the
|
||||
:data:`TOOL_MARKER_ATTR` attribute so the framework can
|
||||
discover decorated functions by scanning a module's namespace.
|
||||
|
||||
The decorator is intentionally pure metadata: it returns the
|
||||
original function unwrapped. The framework's executor handles
|
||||
sync vs async execution (wrapping plain ``def`` bodies in
|
||||
``asyncio.to_thread`` so they don't block the event loop).
|
||||
|
||||
**Decorator stacking**: ``@tool`` should be the outermost
|
||||
decorator. Inner decorators (``@retry``, ``@cache``, etc.) must
|
||||
use ``functools.wraps`` so the signature and docstring survive
|
||||
to the schema-derivation pass; otherwise the schema will reflect
|
||||
the wrapper, not the underlying function.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, ParamSpec, TypeVar, overload
|
||||
|
||||
from ._schema import build_function_schema
|
||||
|
||||
# Marker attribute name. The framework's loader scans
|
||||
# ``module.__dict__`` for objects carrying this attribute to
|
||||
# enumerate the tools a Python file exports.
|
||||
TOOL_MARKER_ATTR = "_omnigent_tool_metadata"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ToolMetadata:
|
||||
"""
|
||||
Metadata attached to a ``@tool``-decorated function.
|
||||
|
||||
Read by the framework loader and by the subprocess runner.
|
||||
Read-only by convention — re-decoration or hand-editing is
|
||||
not supported.
|
||||
|
||||
:param name: The tool name as the LLM sees it. Derived from
|
||||
the function's ``__name__``, e.g. ``"word_count"``.
|
||||
:param description: Human-readable description, derived from
|
||||
the function's docstring's leading paragraph.
|
||||
:param json_schema: The function-calling JSON schema for the
|
||||
tool's parameters, in the OpenAI function-calling shape
|
||||
(an ``object`` schema with ``properties`` / ``required``).
|
||||
Already strict-mode-normalized if ``strict=True`` was
|
||||
passed to the decorator.
|
||||
:param strict: Whether the schema was normalized to strict
|
||||
mode. Stored so the executor side can stay consistent with
|
||||
the LLM's expectations.
|
||||
:param return_annotation: The function's declared return type,
|
||||
used by the executor to deserialize the return value via
|
||||
``pydantic.TypeAdapter``. ``None`` if the function has no
|
||||
return annotation (executor falls back to
|
||||
``json.dumps(value, default=str)``).
|
||||
"""
|
||||
|
||||
name: str
|
||||
description: str
|
||||
json_schema: dict[str, Any]
|
||||
strict: bool
|
||||
return_annotation: type[Any] | None
|
||||
|
||||
|
||||
P = ParamSpec("P")
|
||||
R = TypeVar("R")
|
||||
|
||||
|
||||
@overload
|
||||
def tool(fn: Callable[P, R]) -> Callable[P, R]: ...
|
||||
|
||||
|
||||
@overload
|
||||
def tool(
|
||||
*,
|
||||
strict: bool = ...,
|
||||
) -> Callable[[Callable[P, R]], Callable[P, R]]: ...
|
||||
|
||||
|
||||
def tool(
|
||||
fn: Callable[P, R] | None = None,
|
||||
*,
|
||||
strict: bool = True,
|
||||
) -> Callable[P, R] | Callable[[Callable[P, R]], Callable[P, R]]:
|
||||
"""
|
||||
Mark a module-level function as an omnigent tool.
|
||||
|
||||
The decorator infers the LLM-facing schema from the function's
|
||||
type hints and Google-style docstring, then attaches the
|
||||
derived metadata via the :data:`TOOL_MARKER_ATTR` attribute.
|
||||
The framework's loader scans modules for this marker to
|
||||
register tools at agent-image load time.
|
||||
|
||||
Every ``@tool``-decorated function is sync from the framework's
|
||||
perspective. Authors who want a tool dispatched as background
|
||||
work do not annotate the tool — the LLM picks per call site
|
||||
via ``sys_call_async(tool=..., args=...)`` (see
|
||||
``omnigent/tools/builtins/async_inbox.py``). The author-time
|
||||
``@tool(synchronous=False)`` decoration was removed once
|
||||
``sys_call_async`` shipped — keeping both surfaces would double
|
||||
the dispatch paths without adding capability.
|
||||
|
||||
Usage::
|
||||
|
||||
@tool
|
||||
async def word_count(text: str) -> dict[str, int]:
|
||||
\"\"\"Count words, characters, and lines.\"\"\"
|
||||
...
|
||||
|
||||
Restrictions:
|
||||
- Target must be a module-level ``def`` or ``async def``.
|
||||
Class methods, lambdas, and nested functions are rejected
|
||||
with ``TypeError``.
|
||||
- ``@tool`` should be the outermost decorator; inner
|
||||
decorators must use ``functools.wraps`` for schema derivation
|
||||
to see the underlying signature and docstring.
|
||||
|
||||
:param fn: When used as bare ``@tool`` (no parens), the
|
||||
decorated function. ``None`` when used as
|
||||
``@tool(strict=False)``.
|
||||
:param strict: If ``True`` (default), the derived schema is
|
||||
normalized to strict mode (``additionalProperties: false``
|
||||
on objects, all properties required). Authors who hit a
|
||||
schema strict mode breaks can opt out with
|
||||
``@tool(strict=False)``.
|
||||
:returns: The original function with the tool metadata
|
||||
attached (when called with ``fn``), or a decorator
|
||||
function (when called with keyword args).
|
||||
"""
|
||||
|
||||
def wrap(target: Callable[P, R]) -> Callable[P, R]:
|
||||
_validate_decorator_target(target)
|
||||
schema_result = build_function_schema(target, strict=strict)
|
||||
metadata = ToolMetadata(
|
||||
name=target.__name__,
|
||||
description=schema_result.description,
|
||||
json_schema=schema_result.parameters_json_schema,
|
||||
strict=strict,
|
||||
return_annotation=schema_result.return_annotation,
|
||||
)
|
||||
# Attach metadata via setattr (the dynamic attribute name
|
||||
# is intentional — it's the framework's discovery contract,
|
||||
# see TOOL_MARKER_ATTR).
|
||||
setattr(target, TOOL_MARKER_ATTR, metadata)
|
||||
return target
|
||||
|
||||
if fn is None:
|
||||
return wrap
|
||||
return wrap(fn)
|
||||
|
||||
|
||||
def _validate_decorator_target(target: Any) -> None:
|
||||
"""
|
||||
Reject decorator application to anything other than a
|
||||
module-level ``def`` or ``async def``.
|
||||
|
||||
Class methods would include ``self`` / ``cls`` in the schema,
|
||||
which the LLM has no way to fill. Lambdas have no name and
|
||||
no docstring. Nested functions can close over enclosing scope
|
||||
that doesn't survive subprocess invocation.
|
||||
|
||||
:param target: The object the decorator was applied to.
|
||||
:raises TypeError: If ``target`` is not a module-level
|
||||
function. The message names the offending construct so
|
||||
agent authors get an actionable error.
|
||||
"""
|
||||
if not callable(target):
|
||||
raise TypeError(f"@tool can only be applied to functions, got {type(target).__name__}.")
|
||||
|
||||
if isinstance(target, (staticmethod, classmethod)):
|
||||
raise TypeError(
|
||||
"@tool cannot be applied to staticmethod or classmethod. "
|
||||
"Define the tool as a module-level function instead — "
|
||||
"the framework has no way to bind 'self' or 'cls' from "
|
||||
"an LLM-supplied argument set."
|
||||
)
|
||||
|
||||
if not inspect.isfunction(target):
|
||||
# Covers callables, methods of class instances, etc.
|
||||
raise TypeError(
|
||||
f"@tool requires a plain Python function, got "
|
||||
f"{type(target).__name__}. Define the tool as a "
|
||||
f"module-level def or async def."
|
||||
)
|
||||
|
||||
if target.__name__ == "<lambda>":
|
||||
raise TypeError(
|
||||
"@tool cannot be applied to a lambda. Lambdas have no "
|
||||
"name or docstring; the LLM has nothing to call. Define "
|
||||
"the tool with `def` or `async def` instead."
|
||||
)
|
||||
|
||||
# Nested-function detection: __qualname__ contains a dot path
|
||||
# (e.g. "outer.<locals>.inner") for any function defined inside
|
||||
# another function or class body.
|
||||
qualname = target.__qualname__
|
||||
if qualname != target.__name__:
|
||||
# Allow class-level methods if someone re-binds them at module
|
||||
# scope (rare); the staticmethod/classmethod check above is
|
||||
# the primary defense. Otherwise reject as nested.
|
||||
raise TypeError(
|
||||
f"@tool cannot be applied to nested functions or methods "
|
||||
f"({qualname!r}). Define the tool at module scope so the "
|
||||
f"framework's subprocess runner can re-import it cleanly. "
|
||||
f"State that needs to persist across invocations belongs "
|
||||
f"in module-level globals, not closure variables."
|
||||
)
|
||||
|
||||
|
||||
def get_tool_metadata(obj: Any) -> ToolMetadata | None:
|
||||
"""
|
||||
Return the :class:`ToolMetadata` for an object if it is a
|
||||
``@tool``-decorated function, else ``None``.
|
||||
|
||||
Used by the framework loader to filter a module's namespace
|
||||
down to the decorated functions it should register.
|
||||
|
||||
:param obj: Any value pulled from a module's namespace.
|
||||
:returns: The attached metadata, or ``None`` if the object
|
||||
was not produced by ``@tool``.
|
||||
"""
|
||||
metadata = getattr(obj, TOOL_MARKER_ATTR, None)
|
||||
if isinstance(metadata, ToolMetadata):
|
||||
return metadata
|
||||
return None
|
||||
@@ -0,0 +1,189 @@
|
||||
"""
|
||||
Google-style docstring parsing for ``@tool``-decorated functions.
|
||||
|
||||
Extracts the function description (everything before the first
|
||||
section header) and per-parameter descriptions (from the
|
||||
``Args:`` / ``Arguments:`` / ``Parameters:`` section).
|
||||
|
||||
Used by the schema-derivation logic to populate the function-calling
|
||||
JSON schema's ``description`` and ``properties[name].description``
|
||||
fields.
|
||||
|
||||
We don't depend on a third-party docstring library because the
|
||||
Google-style format is simple and our needs are narrow. NumPy and
|
||||
Sphinx styles are intentionally not supported — authors should
|
||||
use Google style or the explicit ``Annotated[T, Field(description=...)]``
|
||||
form for per-param descriptions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from dataclasses import dataclass
|
||||
|
||||
# Recognized section headers that terminate the description and
|
||||
# the args section. Case-sensitive (Google convention).
|
||||
_SECTION_HEADERS = (
|
||||
"Args:",
|
||||
"Arguments:",
|
||||
"Parameters:",
|
||||
"Returns:",
|
||||
"Return:",
|
||||
"Yields:",
|
||||
"Yield:",
|
||||
"Raises:",
|
||||
"Raise:",
|
||||
"Note:",
|
||||
"Notes:",
|
||||
"Example:",
|
||||
"Examples:",
|
||||
"See Also:",
|
||||
"Warning:",
|
||||
"Warnings:",
|
||||
"Attributes:",
|
||||
)
|
||||
|
||||
_ARGS_HEADERS = ("Args:", "Arguments:", "Parameters:")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ParsedDocstring:
|
||||
"""
|
||||
Result of parsing a Google-style docstring.
|
||||
|
||||
:param description: The function-level description, taken from the
|
||||
text preceding the first section header. Whitespace-trimmed.
|
||||
:param param_descriptions: Mapping from parameter name to its
|
||||
description, extracted from the ``Args:`` / ``Arguments:`` /
|
||||
``Parameters:`` section. Empty if no such section exists.
|
||||
"""
|
||||
|
||||
description: str
|
||||
param_descriptions: dict[str, str]
|
||||
|
||||
|
||||
def parse_google_docstring(doc: str) -> ParsedDocstring:
|
||||
"""
|
||||
Parse a Google-style docstring into description and per-param docs.
|
||||
|
||||
Recognizes ``Args:`` / ``Arguments:`` / ``Parameters:`` as the
|
||||
parameter-list section header. Within that section, lines like
|
||||
`` name: description`` (or `` name (type): description``)
|
||||
are parsed as parameter entries; subsequent more-indented lines
|
||||
are treated as continuations of the current parameter's
|
||||
description. The args section ends at the next recognized
|
||||
section header.
|
||||
|
||||
:param doc: The raw docstring text (typically from
|
||||
``fn.__doc__``). May be ``None``-equivalent (empty string).
|
||||
:returns: A :class:`ParsedDocstring` with the extracted description
|
||||
and parameter descriptions. Returns empty values rather than
|
||||
raising for malformed input.
|
||||
"""
|
||||
if not doc:
|
||||
return ParsedDocstring(description="", param_descriptions={})
|
||||
|
||||
cleaned = inspect.cleandoc(doc)
|
||||
if not cleaned:
|
||||
return ParsedDocstring(description="", param_descriptions={})
|
||||
|
||||
lines = cleaned.split("\n")
|
||||
|
||||
# Locate the first recognized section header. Everything before
|
||||
# it is the description; the args section (if any) is what
|
||||
# contains parameter docs.
|
||||
description_lines: list[str] = []
|
||||
args_section_lines: list[str] = []
|
||||
in_args_section = False
|
||||
in_other_section = False
|
||||
|
||||
for line in lines:
|
||||
stripped = line.strip()
|
||||
if stripped in _SECTION_HEADERS:
|
||||
in_args_section = stripped in _ARGS_HEADERS
|
||||
in_other_section = not in_args_section
|
||||
continue
|
||||
if in_args_section:
|
||||
args_section_lines.append(line)
|
||||
elif in_other_section:
|
||||
# Skip non-args sections (Returns:, Raises:, etc.)
|
||||
continue
|
||||
else:
|
||||
description_lines.append(line)
|
||||
|
||||
description = "\n".join(description_lines).strip()
|
||||
param_descriptions = _parse_args_lines(args_section_lines)
|
||||
|
||||
return ParsedDocstring(
|
||||
description=description,
|
||||
param_descriptions=param_descriptions,
|
||||
)
|
||||
|
||||
|
||||
def _parse_args_lines(lines: list[str]) -> dict[str, str]:
|
||||
"""
|
||||
Parse the body of an ``Args:`` section into per-param descriptions.
|
||||
|
||||
Param lines have the form `` name: description`` or
|
||||
`` name (type): description`` at the section's base indent.
|
||||
Lines indented further are treated as continuations of the
|
||||
current parameter's description.
|
||||
|
||||
:param lines: The lines following the ``Args:`` header (not
|
||||
including the header itself), up to the next section.
|
||||
:returns: Mapping from parameter name to its (whitespace-collapsed)
|
||||
description.
|
||||
"""
|
||||
# Determine the base indent — the indent of the first non-empty line.
|
||||
base_indent: int | None = None
|
||||
for line in lines:
|
||||
if line.strip():
|
||||
base_indent = len(line) - len(line.lstrip())
|
||||
break
|
||||
|
||||
if base_indent is None:
|
||||
return {}
|
||||
|
||||
param_descriptions: dict[str, str] = {}
|
||||
current_name: str | None = None
|
||||
current_parts: list[str] = []
|
||||
|
||||
for line in lines:
|
||||
if not line.strip():
|
||||
# Blank lines within args section are continuation separators;
|
||||
# they don't terminate a param.
|
||||
continue
|
||||
line_indent = len(line) - len(line.lstrip())
|
||||
|
||||
if line_indent == base_indent:
|
||||
# Start of a new param entry.
|
||||
if current_name is not None:
|
||||
param_descriptions[current_name] = " ".join(current_parts).strip()
|
||||
current_name = None
|
||||
current_parts = []
|
||||
|
||||
stripped = line.strip()
|
||||
if ":" not in stripped:
|
||||
# Malformed entry; skip it.
|
||||
continue
|
||||
|
||||
name_part, _, desc = stripped.partition(":")
|
||||
# Handle "name (type)" form by trimming the parenthetical.
|
||||
if "(" in name_part:
|
||||
name_only = name_part.split("(", 1)[0].strip()
|
||||
else:
|
||||
name_only = name_part.strip()
|
||||
|
||||
if name_only.isidentifier():
|
||||
current_name = name_only
|
||||
current_parts = [desc.strip()]
|
||||
# Else: not a valid param entry; ignore.
|
||||
else:
|
||||
# Continuation line for the current param.
|
||||
if current_name is not None:
|
||||
current_parts.append(line.strip())
|
||||
|
||||
if current_name is not None:
|
||||
param_descriptions[current_name] = " ".join(current_parts).strip()
|
||||
|
||||
return param_descriptions
|
||||
@@ -0,0 +1,114 @@
|
||||
"""Adapter that turns ``@tool``-decorated functions into a ToolHandler.
|
||||
|
||||
The stream-layer ``ToolHandler`` takes a list of OpenAI-shape JSON
|
||||
schemas and a single ``execute`` callable. Users who have written
|
||||
tools with the ``@tool`` decorator (Python functions with type hints
|
||||
and Google-style docstrings) shouldn't have to hand-roll that shape:
|
||||
:func:`build_tool_handler` reads each function's tool metadata and
|
||||
builds the handler for them.
|
||||
|
||||
Dispatch is by tool name. Calling an unknown tool raises — the SDK
|
||||
surfaces the error back to the agent as a tool error.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
import json
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from .._tool_handler import ToolCallInfo, ToolHandler
|
||||
from ._decorator import TOOL_MARKER_ATTR, ToolMetadata
|
||||
|
||||
|
||||
def build_tool_handler(functions: list[Callable[..., Any]]) -> ToolHandler:
|
||||
"""Build a :class:`ToolHandler` from ``@tool``-decorated functions.
|
||||
|
||||
Each function must carry tool metadata attached by the
|
||||
:func:`~omnigent_client.tool` decorator (checked via
|
||||
:data:`TOOL_MARKER_ATTR`). The returned handler exposes the
|
||||
OpenAI-shape schemas the SDK sends to the server, and an
|
||||
``execute`` callable that dispatches incoming tool calls by
|
||||
name.
|
||||
|
||||
:param functions: List of ``@tool``-decorated Python functions,
|
||||
e.g. ``[get_current_time, search_docs]``. Each must be a
|
||||
module-level ``def`` or ``async def`` decorated with
|
||||
``@tool``.
|
||||
:returns: A :class:`ToolHandler` ready to pass as
|
||||
``session.tool_handler`` or via the ``tools=`` keyword on
|
||||
``OmnigentClient.query`` / ``Session.query``.
|
||||
:raises TypeError: If any function is missing the ``@tool``
|
||||
marker (i.e. wasn't decorated).
|
||||
:raises ValueError: If two functions share the same tool name
|
||||
— tool names must be unique per handler.
|
||||
"""
|
||||
if not functions:
|
||||
raise ValueError("build_tool_handler() requires at least one function")
|
||||
|
||||
schemas: list[dict[str, object]] = []
|
||||
funcs_by_name: dict[str, Callable[..., Any]] = {}
|
||||
|
||||
for fn in functions:
|
||||
meta: ToolMetadata | None = getattr(fn, TOOL_MARKER_ATTR, None)
|
||||
if meta is None:
|
||||
raise TypeError(
|
||||
f"{fn.__module__}.{fn.__qualname__} is not decorated with "
|
||||
f"@tool. Decorate it with `from omnigent_client import tool` "
|
||||
f"and apply @tool above the function definition."
|
||||
)
|
||||
if meta.name in funcs_by_name:
|
||||
raise ValueError(
|
||||
f"Duplicate tool name {meta.name!r}: "
|
||||
f"{funcs_by_name[meta.name].__qualname__} and "
|
||||
f"{fn.__qualname__} both export the same name."
|
||||
)
|
||||
funcs_by_name[meta.name] = fn
|
||||
schema: dict[str, object] = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": meta.name,
|
||||
"description": meta.description,
|
||||
"parameters": meta.json_schema,
|
||||
},
|
||||
}
|
||||
schemas.append(schema)
|
||||
|
||||
async def execute(call: ToolCallInfo) -> str:
|
||||
"""Dispatch ``call`` to the matching ``@tool`` function.
|
||||
|
||||
Async functions (``async def``) are awaited on the
|
||||
event loop. Sync functions (``def``) are dispatched to
|
||||
a worker thread via ``asyncio.to_thread`` so blocking
|
||||
calls inside — ``time.sleep``, file I/O, subprocess,
|
||||
``requests`` — don't stall the event loop. Without the
|
||||
thread bounce, several concurrent ``@tool`` invocations
|
||||
(e.g. a parallel fan-out of async client tools) would
|
||||
serialize: each body would block every sibling AND any
|
||||
caller render loop sharing the loop.
|
||||
|
||||
The return value is JSON-serialized unless the function
|
||||
already returned a string (which is passed through).
|
||||
"""
|
||||
fn = funcs_by_name.get(call.name)
|
||||
if fn is None:
|
||||
# The SDK will surface this back to the agent as a tool
|
||||
# error — this typically means the LLM invented a tool
|
||||
# name that wasn't in the schemas we sent.
|
||||
raise KeyError(f"Unknown tool {call.name!r}. Registered: {sorted(funcs_by_name)}")
|
||||
if inspect.iscoroutinefunction(fn):
|
||||
result = await fn(**call.arguments)
|
||||
else:
|
||||
# Sync body — route to a worker thread so it
|
||||
# doesn't block the event loop (see the fan-out
|
||||
# serialization case above).
|
||||
result = await asyncio.to_thread(lambda: fn(**call.arguments))
|
||||
if isinstance(result, str):
|
||||
return result
|
||||
# Pydantic models and dataclasses commonly aren't JSON-ready
|
||||
# out of the box — ``default=str`` handles datetime/UUID/etc.
|
||||
return json.dumps(result, default=str)
|
||||
|
||||
return ToolHandler(schemas=schemas, execute=execute)
|
||||
@@ -0,0 +1,267 @@
|
||||
"""
|
||||
Schema derivation for ``@tool``-decorated functions.
|
||||
|
||||
Given a typed Python function, produce the function-calling JSON
|
||||
schema the LLM sees. The pipeline:
|
||||
|
||||
1. Inspect the signature for parameters, annotations, and defaults.
|
||||
2. Parse the Google-style docstring for description and per-param
|
||||
descriptions (see :mod:`omnigent.tools._docstring`).
|
||||
3. Build a Pydantic model from the parameters via ``create_model``;
|
||||
Pydantic does the heavy lifting for type → schema (primitives,
|
||||
Pydantic models, ``Optional``, ``Literal``, ``Annotated[..., Field]``,
|
||||
etc.).
|
||||
4. Apply strict-mode normalization (see
|
||||
:mod:`omnigent.tools._strict`) when ``strict=True``.
|
||||
|
||||
Permissive types (``Any``, ``object``, missing annotations) are
|
||||
allowed but produce an INFO-level warning so authors can find them.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import logging
|
||||
import typing
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Annotated, Any, get_args, get_origin
|
||||
|
||||
from pydantic import Field, create_model
|
||||
from pydantic.fields import FieldInfo
|
||||
|
||||
from ._docstring import parse_google_docstring
|
||||
from ._state import ToolState
|
||||
from ._strict import ensure_strict_schema
|
||||
|
||||
_logger = logging.getLogger(__name__)
|
||||
|
||||
# Reserved parameter name for framework-injected per-conversation
|
||||
# per-agent tool state. A ``@tool`` function that declares a
|
||||
# parameter with this name receives a live :class:`ToolState` at
|
||||
# call time; the parameter is stripped from the LLM-facing JSON
|
||||
# schema. Convention over configuration — every stateful tool uses
|
||||
# the same identifier.
|
||||
STATE_PARAM_NAME = "tool_state"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FunctionSchemaResult:
|
||||
"""
|
||||
Output of :func:`build_function_schema`.
|
||||
|
||||
:param description: Function-level description, derived from
|
||||
the docstring's leading paragraph(s).
|
||||
:param parameters_json_schema: JSON schema for the function's
|
||||
parameters, in the OpenAI function-calling shape (an
|
||||
``object`` schema with ``properties`` and ``required``).
|
||||
Already normalized to strict mode if requested.
|
||||
:param return_annotation: The function's return type annotation,
|
||||
or ``None`` if no return annotation was provided. Used by
|
||||
the executor to deserialize the tool's return value via
|
||||
``pydantic.TypeAdapter``.
|
||||
"""
|
||||
|
||||
description: str
|
||||
parameters_json_schema: dict[str, Any]
|
||||
return_annotation: type[Any] | None
|
||||
|
||||
|
||||
def build_function_schema(
|
||||
fn: Callable[..., Any],
|
||||
*,
|
||||
strict: bool = True,
|
||||
) -> FunctionSchemaResult:
|
||||
"""
|
||||
Build the function-calling schema for a Python function.
|
||||
|
||||
:param fn: The Python function to derive a schema for. Must be
|
||||
a module-level ``def`` or ``async def`` (the ``@tool``
|
||||
decorator enforces this elsewhere; we do not re-validate
|
||||
here).
|
||||
:param strict: If ``True``, apply strict-mode normalization to
|
||||
the resulting schema (see :mod:`._strict`).
|
||||
:returns: A :class:`FunctionSchemaResult` with description,
|
||||
JSON schema, and return-type annotation.
|
||||
"""
|
||||
sig = inspect.signature(fn)
|
||||
type_hints = typing.get_type_hints(fn, include_extras=True)
|
||||
parsed_doc = parse_google_docstring(fn.__doc__ or "")
|
||||
|
||||
fields: dict[str, tuple[Any, FieldInfo]] = {}
|
||||
for name, param in sig.parameters.items():
|
||||
ann = type_hints.get(name, Any)
|
||||
|
||||
# Framework-injected parameter — reserved by convention:
|
||||
# any parameter named exactly ``tool_state`` is filled by
|
||||
# the runtime with a :class:`ToolState`. Skipped from the
|
||||
# LLM-facing schema; the LLM has no way to supply it.
|
||||
# Enforce the convention: if someone types a param as
|
||||
# ToolState but names it something else, fail loud so they
|
||||
# know the right contract.
|
||||
if name == STATE_PARAM_NAME:
|
||||
if ann is not ToolState and ann is not Any:
|
||||
raise TypeError(
|
||||
f"@tool function {fn.__name__!r} declares parameter "
|
||||
f"'{STATE_PARAM_NAME}' with unexpected type "
|
||||
f"{ann!r}. It must be typed as ToolState "
|
||||
f"(or left unannotated); any other type is a bug."
|
||||
)
|
||||
continue
|
||||
if ann is ToolState:
|
||||
raise TypeError(
|
||||
f"@tool function {fn.__name__!r} types parameter "
|
||||
f"{name!r} as ToolState but the parameter must be named "
|
||||
f"{STATE_PARAM_NAME!r}. Rename it and the framework "
|
||||
f"will inject a live ToolState at call time."
|
||||
)
|
||||
|
||||
_warn_if_permissive(fn.__name__, name, ann)
|
||||
|
||||
doc_desc = parsed_doc.param_descriptions.get(name)
|
||||
default = param.default if param.default is not inspect.Parameter.empty else ...
|
||||
field_info = _build_field_info(ann, default, doc_desc)
|
||||
fields[name] = (ann, field_info)
|
||||
|
||||
if fields:
|
||||
# Pydantic uses the model name when generating $defs refs;
|
||||
# capitalize so it looks reasonable in the schema output.
|
||||
model_name = f"{_pascal_case(fn.__name__)}Args"
|
||||
# mypy can't statically narrow create_model's overload for our
|
||||
# dynamic field dict, but pydantic accepts (Type, FieldInfo)
|
||||
# tuples here — they're the documented field-definition shape.
|
||||
Model = create_model(model_name, **fields) # type: ignore[call-overload]
|
||||
params_schema: dict[str, Any] = Model.model_json_schema()
|
||||
else:
|
||||
# Zero-arg tool: the schema is an empty object.
|
||||
params_schema = {
|
||||
"type": "object",
|
||||
"properties": {},
|
||||
"required": [],
|
||||
}
|
||||
|
||||
if strict:
|
||||
params_schema = ensure_strict_schema(params_schema)
|
||||
|
||||
return_annotation = type_hints.get("return")
|
||||
|
||||
return FunctionSchemaResult(
|
||||
description=parsed_doc.description,
|
||||
parameters_json_schema=params_schema,
|
||||
return_annotation=return_annotation,
|
||||
)
|
||||
|
||||
|
||||
def _build_field_info(
|
||||
annotation: Any,
|
||||
default: Any,
|
||||
doc_description: str | None,
|
||||
) -> FieldInfo:
|
||||
"""
|
||||
Construct a Pydantic ``FieldInfo`` for one parameter.
|
||||
|
||||
Handles three description sources, with this priority:
|
||||
1. An explicit ``Field(description=...)`` in
|
||||
``Annotated[T, Field(description=...)]``.
|
||||
2. A bare string in ``Annotated[T, "desc"]`` (a common shorthand
|
||||
supported by some agent SDKs).
|
||||
3. The docstring entry for this parameter (Google-style ``Args:``).
|
||||
|
||||
:param annotation: The parameter's type annotation, possibly
|
||||
wrapped in ``Annotated[...]``.
|
||||
:param default: The parameter's default value, or ``...`` if
|
||||
the parameter is required.
|
||||
:param doc_description: Description from the docstring's
|
||||
``Args:`` section, or ``None`` if absent.
|
||||
:returns: A ``FieldInfo`` ready to pass to
|
||||
``pydantic.create_model``. The default (if any) and
|
||||
description are baked in at construction time so
|
||||
``model_json_schema`` picks them up correctly.
|
||||
"""
|
||||
# Pull metadata out of Annotated[T, ...] for description discovery.
|
||||
annotated_str_desc: str | None = None
|
||||
annotated_field: FieldInfo | None = None
|
||||
if get_origin(annotation) is Annotated:
|
||||
for extra in get_args(annotation)[1:]:
|
||||
if isinstance(extra, FieldInfo) and annotated_field is None:
|
||||
annotated_field = extra
|
||||
elif isinstance(extra, str) and annotated_str_desc is None:
|
||||
annotated_str_desc = extra
|
||||
|
||||
# Determine the effective description with the priority above.
|
||||
description: str | None = None
|
||||
if annotated_field is not None and annotated_field.description is not None:
|
||||
description = annotated_field.description
|
||||
elif annotated_str_desc is not None:
|
||||
description = annotated_str_desc
|
||||
elif doc_description:
|
||||
description = doc_description
|
||||
|
||||
# Field(default=PydanticUndefined) is the marker for "required";
|
||||
# we map our `...` sentinel to it via PydanticUndefined import.
|
||||
# Easier: build the constructor kwargs and let Pydantic translate
|
||||
# default=... directly (it accepts ``...`` as "required" too).
|
||||
field_kwargs: dict[str, Any] = {}
|
||||
if description is not None:
|
||||
field_kwargs["description"] = description
|
||||
if default is not ...:
|
||||
field_kwargs["default"] = default
|
||||
|
||||
if annotated_field is not None:
|
||||
# Merge: preserve other metadata (gt/lt/min_length/etc.) from
|
||||
# the author's Field, but override description and default.
|
||||
# merge_field_infos's stub return is too loose for mypy; cast.
|
||||
merged: FieldInfo = FieldInfo.merge_field_infos(annotated_field, FieldInfo(**field_kwargs))
|
||||
return merged
|
||||
|
||||
# pydantic.Field stub returns Any (it's polymorphic by default).
|
||||
# We're constructing a fresh FieldInfo; cast to the documented type.
|
||||
field_obj: FieldInfo = Field(**field_kwargs)
|
||||
return field_obj
|
||||
|
||||
|
||||
def _warn_if_permissive(fn_name: str, param_name: str, annotation: Any) -> None:
|
||||
"""
|
||||
Log a warning if a parameter's type provides no validation constraint.
|
||||
|
||||
``Any``, ``object``, and missing annotations all produce a
|
||||
permissive schema (no ``type`` field) that the LLM can fill
|
||||
with arbitrary structure. Useful but easy to write by accident.
|
||||
|
||||
:param fn_name: The decorated function's ``__name__``, for the
|
||||
log message, e.g. ``"process_payload"``.
|
||||
:param param_name: The offending parameter's name, e.g.
|
||||
``"data"``.
|
||||
:param annotation: The annotation as resolved by
|
||||
``typing.get_type_hints``.
|
||||
"""
|
||||
# Strip Annotated[...] so we inspect the underlying type.
|
||||
underlying = annotation
|
||||
if get_origin(underlying) is Annotated:
|
||||
underlying = get_args(underlying)[0]
|
||||
|
||||
if underlying is Any or underlying is object:
|
||||
type_name = (
|
||||
"Any" if underlying is Any else getattr(underlying, "__name__", str(underlying))
|
||||
)
|
||||
_logger.info(
|
||||
"Tool '%s' parameter '%s' has no concrete type annotation "
|
||||
"(resolved to %s); LLM will get a permissive schema with "
|
||||
"no validation.",
|
||||
fn_name,
|
||||
param_name,
|
||||
type_name,
|
||||
)
|
||||
|
||||
|
||||
def _pascal_case(snake: str) -> str:
|
||||
"""
|
||||
Convert a snake_case identifier to PascalCase.
|
||||
|
||||
Used to give the dynamically-created Pydantic model a readable
|
||||
name in schema ``$defs`` references.
|
||||
|
||||
:param snake: A snake_case identifier, e.g. ``"word_count"``.
|
||||
:returns: PascalCase form, e.g. ``"WordCount"``.
|
||||
"""
|
||||
return "".join(part.capitalize() for part in snake.split("_") if part)
|
||||
@@ -0,0 +1,300 @@
|
||||
"""Per-agent ToolState for stateful ``@tool`` functions.
|
||||
|
||||
See ``designs/TOOL_STATE.md`` for the full design. The primitive is
|
||||
a simple key-value store, JSON-serialized, scoped to one
|
||||
(conversation, agent) pair via the storage directory provided by
|
||||
the framework. The ``@tool`` decorator hides ``ToolState``-typed
|
||||
parameters from the LLM-facing schema; the subprocess runner
|
||||
reconstructs a ``ToolState`` from the directory path and injects it
|
||||
when the tool function is called.
|
||||
|
||||
Tool authors see::
|
||||
|
||||
from omnigent_client import tool, ToolState
|
||||
|
||||
@tool
|
||||
def add_task(desc: str, state: ToolState) -> str:
|
||||
with state.transaction("queue") as q:
|
||||
q = q or []
|
||||
q.append({"desc": desc})
|
||||
return f"#{len(q) - 1}"
|
||||
|
||||
and nothing else.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager, suppress
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
import fcntl as _fcntl
|
||||
except ImportError: # pragma: no cover - exercised on native Windows
|
||||
_fcntl = None # type: ignore[assignment]
|
||||
|
||||
if os.name == "nt":
|
||||
import msvcrt as _msvcrt
|
||||
else: # pragma: no cover - exercised on POSIX
|
||||
_msvcrt = None # type: ignore[assignment]
|
||||
|
||||
# Subdirectory segments reserved by the framework — JSON key files
|
||||
# live at ``{root}/{key}.json``. Keys must not contain path
|
||||
# separators; we sanitize eagerly rather than allowing the bug to
|
||||
# surface as directory traversal.
|
||||
_KEY_SUFFIX = ".json"
|
||||
_THREAD_LOCKS_GUARD = threading.Lock()
|
||||
_THREAD_LOCKS: dict[Path, threading.Lock] = {}
|
||||
|
||||
|
||||
class ToolState:
|
||||
"""Per-agent, per-conversation key-value state for ``@tool`` functions.
|
||||
|
||||
Values are JSON-serialized. The keyspace is shared across every
|
||||
tool invoked for the same registered agent within the same
|
||||
conversation. Use :meth:`transaction` for atomic read-modify-write;
|
||||
plain :meth:`get` and :meth:`set` do not serialize concurrent
|
||||
writers on the same key.
|
||||
|
||||
Instances are constructed by the framework. Tool authors receive
|
||||
a ``ToolState`` by declaring a parameter of this type on their
|
||||
``@tool``-decorated function; the decorator strips the parameter
|
||||
from the LLM-facing schema and the subprocess runner injects the
|
||||
live ``ToolState`` at call time.
|
||||
|
||||
:param root: The directory this namespace lives in, e.g.
|
||||
``{workspace}/.tool_state/{agent_id}``. The directory does
|
||||
not need to exist yet; it is created lazily on first write.
|
||||
"""
|
||||
|
||||
def __init__(self, root: Path) -> None:
|
||||
self._root = root
|
||||
|
||||
# ── Primary API ──────────────────────────────────────────
|
||||
|
||||
def get(self, key: str, *, default: Any = None) -> Any:
|
||||
"""Return the stored value at ``key``, or ``default`` if absent.
|
||||
|
||||
:param key: The state key, e.g. ``"queue"``.
|
||||
:param default: Value to return when the key has never been
|
||||
written. ``None`` by default.
|
||||
:returns: The deserialized JSON value, or ``default``.
|
||||
"""
|
||||
path = self._path_for(key)
|
||||
if not path.exists():
|
||||
return default
|
||||
with path.open("r") as f:
|
||||
# Shared lock: allow parallel reads, block concurrent writers
|
||||
# briefly so we see a complete JSON payload.
|
||||
_lock_file(f, exclusive=False)
|
||||
try:
|
||||
return json.loads(f.read() or "null")
|
||||
finally:
|
||||
_unlock_file(f)
|
||||
|
||||
def set(self, key: str, value: Any) -> None:
|
||||
"""Replace (or create) the value at ``key``. JSON-serialized.
|
||||
|
||||
Non-atomic relative to concurrent writers on the same key —
|
||||
use :meth:`transaction` for read-modify-write sequences.
|
||||
|
||||
:param key: The state key.
|
||||
:param value: Any JSON-serializable value.
|
||||
"""
|
||||
path = self._path_for(key)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
# Write through a temp file + rename so a reader never sees
|
||||
# a half-written JSON payload even without a lock.
|
||||
tmp = path.with_suffix(_KEY_SUFFIX + ".tmp")
|
||||
with tmp.open("w") as f:
|
||||
json.dump(value, f)
|
||||
tmp.replace(path)
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
"""Remove ``key``. No-op if absent.
|
||||
|
||||
:param key: The state key to remove.
|
||||
"""
|
||||
path = self._path_for(key)
|
||||
# Idempotent delete — tools commonly don't know whether the
|
||||
# key was ever set.
|
||||
with suppress(FileNotFoundError):
|
||||
path.unlink()
|
||||
|
||||
def keys(self) -> list[str]:
|
||||
"""List all keys currently stored in this namespace.
|
||||
|
||||
:returns: Sorted list of keys, e.g. ``["counter", "queue"]``.
|
||||
Empty list if nothing has been written yet.
|
||||
"""
|
||||
if not self._root.exists():
|
||||
return []
|
||||
return sorted(p.stem for p in self._root.iterdir() if p.suffix == _KEY_SUFFIX)
|
||||
|
||||
def __contains__(self, key: object) -> bool:
|
||||
"""Return whether ``key`` currently exists in this namespace.
|
||||
|
||||
:param key: Candidate key, e.g. ``"queue"``.
|
||||
:returns: ``True`` when ``key`` is a valid stored key, else ``False``.
|
||||
"""
|
||||
if not isinstance(key, str):
|
||||
return False
|
||||
return self._path_for(key).exists()
|
||||
|
||||
@contextmanager
|
||||
def transaction(self, key: str, *, default: Any = None) -> Iterator[Any]:
|
||||
"""Atomic read-modify-write for one key.
|
||||
|
||||
Typical usage — supply a ``default`` so first-time callers
|
||||
get a usable container without a ``None`` check::
|
||||
|
||||
with state.transaction("queue", default=[]) as queue:
|
||||
queue.append(item)
|
||||
# queue is written back on normal exit.
|
||||
|
||||
The yielded value is the current contents, or a fresh
|
||||
``default`` if the key was never set. Mutating the yielded
|
||||
object in place is the expected pattern — the same object
|
||||
is serialized back on exit. Rebinding the local name inside
|
||||
the ``with`` block does NOT propagate (Python closures), so
|
||||
for "replace the value" semantics use :meth:`set` explicitly.
|
||||
|
||||
On a normal exit the yielded object is JSON-serialized and
|
||||
written back. On exception no write happens — the prior
|
||||
value is preserved.
|
||||
|
||||
:param key: The state key to lock + read + write.
|
||||
:param default: Value yielded when the key has never been
|
||||
written. Defaults to ``None``. Pass ``[]`` or ``{}``
|
||||
(or any JSON-serializable value) to skip the absent-key
|
||||
branch in caller code.
|
||||
:yields: The current value at ``key``, or ``default`` if
|
||||
the key has no stored value yet. Mutate in place.
|
||||
"""
|
||||
path = self._path_for(key)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
thread_lock = _get_thread_lock(path)
|
||||
# ``a+`` creates the file if missing and positions at end;
|
||||
# we seek to 0 before read. Opening with ``r+`` would fail
|
||||
# when the file doesn't exist yet, which is a common first-
|
||||
# call case for a tool.
|
||||
with thread_lock, path.open("a+") as f:
|
||||
_lock_file(f, exclusive=True)
|
||||
try:
|
||||
value = _read_transaction_value(f, default)
|
||||
yield value
|
||||
_write_transaction_value(f, value)
|
||||
finally:
|
||||
_unlock_file(f)
|
||||
|
||||
# ── Internals ────────────────────────────────────────────
|
||||
|
||||
def _path_for(self, key: str) -> Path:
|
||||
"""Resolve ``key`` to the on-disk path, rejecting traversal.
|
||||
|
||||
:param key: Caller-supplied key.
|
||||
:returns: The ``{root}/{key}.json`` path.
|
||||
:raises ValueError: If ``key`` is empty, contains a path
|
||||
separator, or starts with a dot (no hidden or
|
||||
escaped paths).
|
||||
"""
|
||||
if not key:
|
||||
raise ValueError("ToolState key must be a non-empty string")
|
||||
if "/" in key or "\\" in key or key.startswith("."):
|
||||
# Rejects traversal and hidden-file sigils. Authors who
|
||||
# really need slashes can encode them (e.g. "a__b") —
|
||||
# we'd rather break loudly than accept quiet bugs.
|
||||
raise ValueError(
|
||||
f"ToolState key {key!r} contains an illegal character. "
|
||||
f"Keys must be plain names (no '/', '\\', or leading '.')."
|
||||
)
|
||||
return self._root / f"{key}{_KEY_SUFFIX}"
|
||||
|
||||
|
||||
def _read_transaction_value(f: Any, default: Any) -> Any:
|
||||
"""Seek to 0 and decode the JSON value under the open file handle.
|
||||
|
||||
Returns ``default`` when the file is empty (first-time use of
|
||||
the key). Factored out of :meth:`ToolState.transaction` so the
|
||||
context manager stays short.
|
||||
|
||||
:param f: Open file handle positioned anywhere; will be seek(0)ed.
|
||||
:param default: Value to return on empty/whitespace content.
|
||||
:returns: Decoded JSON value or ``default``.
|
||||
"""
|
||||
f.seek(0)
|
||||
raw = f.read()
|
||||
if raw.strip():
|
||||
return json.loads(raw)
|
||||
return default
|
||||
|
||||
|
||||
def _write_transaction_value(f: Any, value: Any) -> None:
|
||||
"""Truncate and re-serialize ``value`` as JSON under the file handle.
|
||||
|
||||
Caller must hold the exclusive flock before calling. Factored
|
||||
out of :meth:`ToolState.transaction` so the context manager
|
||||
stays short.
|
||||
|
||||
:param f: Open file handle (must support ``r+``-style truncate).
|
||||
:param value: Any JSON-serializable value to persist.
|
||||
"""
|
||||
f.seek(0)
|
||||
f.truncate()
|
||||
json.dump(value, f)
|
||||
# Flush before releasing the flock held by the caller. Python file objects
|
||||
# buffer writes; if we unlock before flushing, another process can acquire
|
||||
# the lock and read stale on-disk contents, losing the prior update.
|
||||
f.flush()
|
||||
os.fsync(f.fileno())
|
||||
|
||||
|
||||
def _lock_file(f: Any, *, exclusive: bool) -> None:
|
||||
"""Acquire an advisory file lock for ``f``.
|
||||
|
||||
POSIX uses ``fcntl.flock``. Native Windows has no ``fcntl`` module, so it
|
||||
locks one byte with ``msvcrt.locking``; that API is exclusive-only, which
|
||||
is conservative for reads but preserves cross-process serialization.
|
||||
"""
|
||||
if _fcntl is not None:
|
||||
_fcntl.flock(f, _fcntl.LOCK_EX if exclusive else _fcntl.LOCK_SH)
|
||||
return
|
||||
if _msvcrt is None: # pragma: no cover - defensive for unusual platforms
|
||||
return
|
||||
f.seek(0)
|
||||
_msvcrt.locking(f.fileno(), _msvcrt.LK_LOCK, 1)
|
||||
|
||||
|
||||
def _unlock_file(f: Any) -> None:
|
||||
"""Release a lock acquired by :func:`_lock_file`."""
|
||||
if _fcntl is not None:
|
||||
_fcntl.flock(f, _fcntl.LOCK_UN)
|
||||
return
|
||||
if _msvcrt is None: # pragma: no cover - defensive for unusual platforms
|
||||
return
|
||||
f.seek(0)
|
||||
_msvcrt.locking(f.fileno(), _msvcrt.LK_UNLCK, 1)
|
||||
|
||||
|
||||
def _get_thread_lock(path: Path) -> threading.Lock:
|
||||
"""Return the per-key in-process mutex for ``path``.
|
||||
|
||||
``flock`` serializes across processes, but threads in the same
|
||||
process can still interleave on separate file descriptors. This
|
||||
helper layers a per-path ``threading.Lock`` on top so
|
||||
``transaction()`` is atomic under both thread and process
|
||||
contention.
|
||||
|
||||
:param path: The key file path, e.g. ``Path("/tmp/state/queue.json")``.
|
||||
:returns: The shared mutex guarding that path within this process.
|
||||
"""
|
||||
with _THREAD_LOCKS_GUARD:
|
||||
lock = _THREAD_LOCKS.get(path)
|
||||
if lock is None:
|
||||
lock = threading.Lock()
|
||||
_THREAD_LOCKS[path] = lock
|
||||
return lock
|
||||
@@ -0,0 +1,64 @@
|
||||
"""
|
||||
Strict JSON-schema normalization for ``@tool``-derived schemas.
|
||||
|
||||
Strict-mode schemas have two extra constraints beyond ordinary
|
||||
JSON Schema:
|
||||
|
||||
- Every object schema sets ``additionalProperties: false``.
|
||||
- Every property in an object is listed in ``required`` (Python
|
||||
defaults are still applied on the executor side after the LLM
|
||||
emits a value).
|
||||
|
||||
This is the form most major LLM providers accept for
|
||||
function-calling tool schemas without further coercion. Authors
|
||||
who hit a real schema that strict mode breaks can opt out via
|
||||
``@tool(strict=False)``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
|
||||
def ensure_strict_schema(schema: dict[str, Any]) -> dict[str, Any]:
|
||||
"""
|
||||
Recursively normalize a JSON schema to strict-mode rules.
|
||||
|
||||
Returns a new dict (does not mutate the input). Recurses into
|
||||
``properties``, ``items``, ``$defs``, and union variants
|
||||
(``anyOf`` / ``oneOf`` / ``allOf``).
|
||||
|
||||
:param schema: A JSON schema dict (typically produced by
|
||||
:func:`pydantic.BaseModel.model_json_schema`).
|
||||
:returns: A new dict with strict-mode constraints applied.
|
||||
"""
|
||||
if not isinstance(schema, dict):
|
||||
return schema
|
||||
|
||||
out: dict[str, Any] = dict(schema)
|
||||
|
||||
# Recurse into nested $defs FIRST so the recursion sees the
|
||||
# normalized definitions when it walks references.
|
||||
if "$defs" in out:
|
||||
out["$defs"] = {k: ensure_strict_schema(v) for k, v in out["$defs"].items()}
|
||||
|
||||
# Recurse into union variants.
|
||||
for union_key in ("anyOf", "oneOf", "allOf"):
|
||||
if union_key in out and isinstance(out[union_key], list):
|
||||
out[union_key] = [ensure_strict_schema(variant) for variant in out[union_key]]
|
||||
|
||||
# Recurse into array items.
|
||||
if "items" in out:
|
||||
out["items"] = ensure_strict_schema(out["items"])
|
||||
|
||||
if out.get("type") == "object":
|
||||
properties = out.get("properties", {}) or {}
|
||||
out["additionalProperties"] = False
|
||||
out["properties"] = {k: ensure_strict_schema(v) for k, v in properties.items()}
|
||||
# Strict mode requires every property to be listed as required.
|
||||
# The default value (if any) still applies on the executor side
|
||||
# via Pydantic — strict mode only governs what the LLM emits.
|
||||
if properties:
|
||||
out["required"] = list(properties.keys())
|
||||
|
||||
return out
|
||||
Reference in New Issue
Block a user