Skip to content

Commit

Permalink
Enable creating Tools from any Runnable (#11177)
Browse files Browse the repository at this point in the history
<!-- Thank you for contributing to LangChain!

Replace this entire comment with:
  - **Description:** a description of the change, 
  - **Issue:** the issue # it fixes (if applicable),
  - **Dependencies:** any dependencies required for this change,
- **Tag maintainer:** for a quicker response, tag the relevant
maintainer (see below),
- **Twitter handle:** we announce bigger features on Twitter. If your PR
gets announced, and you'd like a mention, we'll gladly shout you out!

Please make sure your PR is passing linting and testing before
submitting. Run `make format`, `make lint` and `make test` to check this
locally.

See contribution guidelines for more information on how to write/run
tests, lint, etc:

https://github.com/hwchase17/langchain/blob/master/.github/CONTRIBUTING.md

If you're adding a new integration, please include:
1. a test for the integration, preferably unit tests that do not rely on
network access,
2. an example notebook showing its use. It lives in `docs/extras`
directory.

If no one reviews your PR within a few days, please @-mention one of
@baskaryan, @eyurtsev, @hwchase17.
 -->
  • Loading branch information
nfcampos authored Sep 29, 2023
2 parents 61b5942 + e35ea56 commit ca5293b
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 5 deletions.
37 changes: 32 additions & 5 deletions libs/langchain/langchain/tools/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -734,7 +734,7 @@ def add(a: int, b: int) -> int:


def tool(
*args: Union[str, Callable],
*args: Union[str, Callable, Runnable],
return_direct: bool = False,
args_schema: Optional[Type[BaseModel]] = None,
infer_schema: bool = True,
Expand Down Expand Up @@ -769,21 +769,46 @@ def search_api(query: str) -> str:
"""

def _make_with_name(tool_name: str) -> Callable:
def _make_tool(dec_func: Callable) -> BaseTool:
if inspect.iscoroutinefunction(dec_func):
def _make_tool(dec_func: Union[Callable, Runnable]) -> BaseTool:
if isinstance(dec_func, Runnable):
runnable = dec_func

if runnable.input_schema.schema().get("type") != "object":
raise ValueError("Runnable must have an object schema.")

async def ainvoke_wrapper(
callbacks: Optional[Callbacks] = None, **kwargs: Any
) -> Any:
return await runnable.ainvoke(kwargs, {"callbacks": callbacks})

def invoke_wrapper(
callbacks: Optional[Callbacks] = None, **kwargs: Any
) -> Any:
return runnable.invoke(kwargs, {"callbacks": callbacks})

coroutine = ainvoke_wrapper
func = invoke_wrapper
schema: Optional[Type[BaseModel]] = runnable.input_schema
description = repr(runnable)
elif inspect.iscoroutinefunction(dec_func):
coroutine = dec_func
func = None
schema = args_schema
description = None
else:
coroutine = None
func = dec_func
schema = args_schema
description = None

if infer_schema or args_schema is not None:
return StructuredTool.from_function(
func,
coroutine,
name=tool_name,
description=description,
return_direct=return_direct,
args_schema=args_schema,
args_schema=schema,
infer_schema=infer_schema,
)
# If someone doesn't want a schema applied, we must treat it as
Expand All @@ -803,7 +828,9 @@ def _make_tool(dec_func: Callable) -> BaseTool:

return _make_tool

if len(args) == 1 and isinstance(args[0], str):
if len(args) == 2 and isinstance(args[0], str) and isinstance(args[1], Runnable):
return _make_with_name(args[0])(args[1])
elif 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])
Expand Down
30 changes: 30 additions & 0 deletions libs/langchain/tests/unit_tests/schema/runnable/test_runnable.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
RunnableSequence,
RunnableWithFallbacks,
)
from langchain.tools.base import BaseTool, tool
from langchain.tools.json.tool import JsonListKeysTool, JsonSpec


Expand Down Expand Up @@ -2779,3 +2780,32 @@ async def af(x: int) -> int:
" b: RunnableLambda(...)\n"
" }"
), "repr where code string contains multiple lambdas gives up"


@pytest.mark.asyncio
async def test_tool_from_runnable() -> None:
prompt = (
SystemMessagePromptTemplate.from_template("You are a nice assistant.")
+ "{question}"
)
llm = FakeStreamingListLLM(responses=["foo-lish"])

chain = prompt | llm | StrOutputParser()

chain_tool = tool("chain_tool", chain)

assert isinstance(chain_tool, BaseTool)
assert chain_tool.name == "chain_tool"
assert chain_tool.run({"question": "What up"}) == chain.invoke(
{"question": "What up"}
)
assert await chain_tool.arun({"question": "What up"}) == await chain.ainvoke(
{"question": "What up"}
)
assert chain_tool.description.endswith(repr(chain))
assert chain_tool.args_schema.schema() == chain.input_schema.schema()
assert chain_tool.args_schema.schema() == {
"properties": {"question": {"title": "Question"}},
"title": "PromptInput",
"type": "object",
}

0 comments on commit ca5293b

Please sign in to comment.