forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
156 lines
5.7 KiB
Python
156 lines
5.7 KiB
Python
"""Interface for tools."""
|
|
from functools import partial
|
|
from typing import Any, Awaitable, Callable, Dict, Optional, Tuple, Type, Union
|
|
|
|
from pydantic import BaseModel, validator
|
|
|
|
from langchain.tools.base import BaseTool, StructuredTool
|
|
|
|
|
|
class Tool(BaseTool):
|
|
"""Tool that takes in function or coroutine directly."""
|
|
|
|
description: str = ""
|
|
func: Callable[..., str]
|
|
"""The function to run when the tool is called."""
|
|
coroutine: Optional[Callable[..., Awaitable[str]]] = None
|
|
"""The asynchronous version of the function."""
|
|
|
|
@validator("func", pre=True, always=True)
|
|
def validate_func_not_partial(cls, func: Callable) -> Callable:
|
|
"""Check that the function is not a partial."""
|
|
if isinstance(func, partial):
|
|
raise ValueError("Partial functions not yet supported in tools.")
|
|
return func
|
|
|
|
@property
|
|
def args(self) -> dict:
|
|
"""The tool's input arguments."""
|
|
if self.args_schema is not None:
|
|
return self.args_schema.schema()["properties"]
|
|
# For backwards compatibility, if the function signature is ambiguous,
|
|
# assume it takes a single string input.
|
|
return {"tool_input": {"type": "string"}}
|
|
|
|
def _to_args_and_kwargs(self, tool_input: Union[str, Dict]) -> Tuple[Tuple, Dict]:
|
|
"""Convert tool input to pydantic model."""
|
|
args, kwargs = super()._to_args_and_kwargs(tool_input)
|
|
# For backwards compatibility. The tool must be run with a single input
|
|
all_args = list(args) + list(kwargs.values())
|
|
if len(all_args) != 1:
|
|
raise ValueError(
|
|
f"Too many arguments to single-input tool {self.name}."
|
|
f" Args: {all_args}"
|
|
)
|
|
return tuple(all_args), {}
|
|
|
|
def _run(self, *args: Any, **kwargs: Any) -> Any:
|
|
"""Use the tool."""
|
|
return self.func(*args, **kwargs)
|
|
|
|
async def _arun(self, *args: Any, **kwargs: Any) -> Any:
|
|
"""Use the tool asynchronously."""
|
|
if self.coroutine:
|
|
return await self.coroutine(*args, **kwargs)
|
|
raise NotImplementedError("Tool does not support async")
|
|
|
|
# TODO: this is for backwards compatibility, remove in future
|
|
def __init__(
|
|
self, name: str, func: Callable, description: str, **kwargs: Any
|
|
) -> None:
|
|
"""Initialize tool."""
|
|
super(Tool, self).__init__(
|
|
name=name, func=func, description=description, **kwargs
|
|
)
|
|
|
|
|
|
class InvalidTool(BaseTool):
|
|
"""Tool that is run when invalid tool name is encountered by agent."""
|
|
|
|
name = "invalid_tool"
|
|
description = "Called when tool name is invalid."
|
|
|
|
def _run(self, tool_name: str) -> str:
|
|
"""Use the tool."""
|
|
return f"{tool_name} is not a valid tool, try another one."
|
|
|
|
async def _arun(self, tool_name: str) -> str:
|
|
"""Use the tool asynchronously."""
|
|
return f"{tool_name} is not a valid tool, try another one."
|
|
|
|
|
|
def tool(
|
|
*args: Union[str, Callable],
|
|
return_direct: bool = False,
|
|
args_schema: Optional[Type[BaseModel]] = None,
|
|
infer_schema: bool = True,
|
|
) -> Callable:
|
|
"""Make tools out of functions, can be used with or without arguments.
|
|
|
|
Args:
|
|
*args: The arguments to the tool.
|
|
return_direct: Whether to return directly from the tool rather
|
|
than continuing the agent loop.
|
|
args_schema: optional argument schema for user to specify
|
|
infer_schema: Whether to infer the schema of the arguments from
|
|
the function's signature. This also makes the resultant tool
|
|
accept a dictionary input to its `run()` function.
|
|
|
|
Requires:
|
|
- Function must be of type (str) -> str
|
|
- Function must have a docstring
|
|
|
|
Examples:
|
|
.. code-block:: python
|
|
|
|
@tool
|
|
def search_api(query: str) -> str:
|
|
# Searches the API for the query.
|
|
return
|
|
|
|
@tool("search", return_direct=True)
|
|
def search_api(query: str) -> str:
|
|
# Searches the API for the query.
|
|
return
|
|
"""
|
|
|
|
def _make_with_name(tool_name: str) -> Callable:
|
|
def _make_tool(func: Callable) -> BaseTool:
|
|
if infer_schema or args_schema is not None:
|
|
return StructuredTool.from_function(
|
|
func,
|
|
name=tool_name,
|
|
return_direct=return_direct,
|
|
args_schema=args_schema,
|
|
infer_schema=infer_schema,
|
|
)
|
|
# If someone doesn't want a schema applied, we must treat it as
|
|
# a simple string->string function
|
|
assert func.__doc__ is not None, "Function must have a docstring"
|
|
return Tool(
|
|
name=tool_name,
|
|
func=func,
|
|
description=f"{tool_name} tool",
|
|
return_direct=return_direct,
|
|
)
|
|
|
|
return _make_tool
|
|
|
|
if len(args) == 1 and isinstance(args[0], str):
|
|
# if the argument is a string, then we use the string as the tool name
|
|
# Example usage: @tool("search", return_direct=True)
|
|
return _make_with_name(args[0])
|
|
elif len(args) == 1 and callable(args[0]):
|
|
# if the argument is a function, then we use the function name as the tool name
|
|
# Example usage: @tool
|
|
return _make_with_name(args[0].__name__)(args[0])
|
|
elif len(args) == 0:
|
|
# if there are no arguments, then we use the function name as the tool name
|
|
# Example usage: @tool(return_direct=True)
|
|
def _partial(func: Callable[[str], str]) -> BaseTool:
|
|
return _make_with_name(func.__name__)(func)
|
|
|
|
return _partial
|
|
else:
|
|
raise ValueError("Too many arguments for tool decorator")
|