112 lines
3.1 KiB
Python
112 lines
3.1 KiB
Python
|
"""Custom tool decorator for OpenAI custom tools."""
|
||
|
|
|
||
|
|
import inspect
|
||
|
|
from collections.abc import Awaitable, Callable
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from langchain_core.tools import tool
|
||
|
|
|
||
|
|
|
||
|
|
def _make_wrapped_func(func: Callable[..., str]) -> Callable[..., list[dict[str, Any]]]:
|
||
|
|
def wrapped(x: str) -> list[dict[str, Any]]:
|
||
|
|
return [{"type": "custom_tool_call_output", "output": func(x)}]
|
||
|
|
|
||
|
|
return wrapped
|
||
|
|
|
||
|
|
|
||
|
|
def _make_wrapped_coroutine(
|
||
|
|
coroutine: Callable[..., Awaitable[str]],
|
||
|
|
) -> Callable[..., Awaitable[list[dict[str, Any]]]]:
|
||
|
|
async def wrapped(*args: Any, **kwargs: Any) -> list[dict[str, Any]]:
|
||
|
|
result = await coroutine(*args, **kwargs)
|
||
|
|
return [{"type": "custom_tool_call_output", "output": result}]
|
||
|
|
|
||
|
|
return wrapped
|
||
|
|
|
||
|
|
|
||
|
|
def custom_tool(*args: Any, **kwargs: Any) -> Any:
|
||
|
|
"""Decorator to create an OpenAI custom tool.
|
||
|
|
|
||
|
|
Custom tools allow for tools with (potentially long) freeform string inputs.
|
||
|
|
|
||
|
|
See below for an example using LangGraph:
|
||
|
|
|
||
|
|
```python
|
||
|
|
@custom_tool
|
||
|
|
def execute_code(code: str) -> str:
|
||
|
|
\"\"\"Execute python code.\"\"\"
|
||
|
|
return "27"
|
||
|
|
|
||
|
|
|
||
|
|
model = ChatOpenAI(model="gpt-5", output_version="responses/v1")
|
||
|
|
|
||
|
|
agent = create_react_agent(model, [execute_code])
|
||
|
|
|
||
|
|
input_message = {"role": "user", "content": "Use the tool to calculate 3^3."}
|
||
|
|
for step in agent.stream(
|
||
|
|
{"messages": [input_message]},
|
||
|
|
stream_mode="values",
|
||
|
|
):
|
||
|
|
step["messages"][-1].pretty_print()
|
||
|
|
```
|
||
|
|
|
||
|
|
You can also specify a format for a corresponding context-free grammar using the
|
||
|
|
`format` kwarg:
|
||
|
|
|
||
|
|
```python
|
||
|
|
from langchain_openai import ChatOpenAI, custom_tool
|
||
|
|
from langgraph.prebuilt import create_react_agent
|
||
|
|
|
||
|
|
grammar = \"\"\"
|
||
|
|
start: expr
|
||
|
|
expr: term (SP ADD SP term)* -> add
|
||
|
|
| term
|
||
|
|
term: factor (SP MUL SP factor)* -> mul
|
||
|
|
| factor
|
||
|
|
factor: INT
|
||
|
|
SP: " "
|
||
|
|
ADD: "+"
|
||
|
|
MUL: "*"
|
||
|
|
%import common.INT
|
||
|
|
\"\"\"
|
||
|
|
|
||
|
|
format = {"type": "grammar", "syntax": "lark", "definition": grammar}
|
||
|
|
|
||
|
|
# highlight-next-line
|
||
|
|
@custom_tool(format=format)
|
||
|
|
def do_math(input_string: str) -> str:
|
||
|
|
\"\"\"Do a mathematical operation.\"\"\"
|
||
|
|
return "27"
|
||
|
|
|
||
|
|
|
||
|
|
model = ChatOpenAI(model="gpt-5", output_version="responses/v1")
|
||
|
|
|
||
|
|
agent = create_react_agent(model, [do_math])
|
||
|
|
|
||
|
|
input_message = {"role": "user", "content": "Use the tool to calculate 3^3."}
|
||
|
|
for step in agent.stream(
|
||
|
|
{"messages": [input_message]},
|
||
|
|
stream_mode="values",
|
||
|
|
):
|
||
|
|
step["messages"][-1].pretty_print()
|
||
|
|
```
|
||
|
|
"""
|
||
|
|
|
||
|
|
def decorator(func: Callable[..., Any]) -> Any:
|
||
|
|
metadata = {"type": "custom_tool"}
|
||
|
|
if "format" in kwargs:
|
||
|
|
metadata["format"] = kwargs.pop("format")
|
||
|
|
tool_obj = tool(infer_schema=False, **kwargs)(func)
|
||
|
|
tool_obj.metadata = metadata
|
||
|
|
tool_obj.description = func.__doc__
|
||
|
|
if inspect.iscoroutinefunction(func):
|
||
|
|
tool_obj.coroutine = _make_wrapped_coroutine(func)
|
||
|
|
else:
|
||
|
|
tool_obj.func = _make_wrapped_func(func)
|
||
|
|
return tool_obj
|
||
|
|
|
||
|
|
if args and callable(args[0]) and not kwargs:
|
||
|
|
return decorator(args[0])
|
||
|
|
|
||
|
|
return decorator
|