autogen/python/packages/autogen-ext/tests/tools/test_langchain_tools.py at ac750c46745867995ebff597d2846aedacbc6904 · microsoft/autogen · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
import logging
from typing import Optional, Type, cast

import pytest
from autogen_core import CancellationToken
from autogen_core.tools import Tool
from autogen_ext.tools.langchain import LangChainToolAdapter  # type: ignore
from langchain_core.callbacks.manager import AsyncCallbackManagerForToolRun, CallbackManagerForToolRun
from langchain_core.tools import BaseTool as LangChainTool
from langchain_core.tools import tool  # pyright: ignore
from pydantic import BaseModel, Field


@tool  # type: ignore
def add(a: int, b: int) -> int:
    """Add two numbers"""
    return a + b


class CalculatorInput(BaseModel):
    a: int = Field(description="first number")
    b: int = Field(description="second number")


class CustomCalculatorTool(LangChainTool):
    name: str = "Calculator"
    description: str = "useful for when you need to answer questions about math"
    args_schema: Type[BaseModel] = CalculatorInput
    return_direct: bool = True

    def _run(self, a: int, b: int, run_manager: Optional[CallbackManagerForToolRun] = None) -> int:
        """Use the tool."""
        return a * b

    async def _arun(
        self,
        a: int,
        b: int,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> int:
        """Use the tool asynchronously."""
        return self._run(a, b, run_manager=run_manager.get_sync() if run_manager else None)


@pytest.mark.asyncio
async def test_langchain_tool_adapter(caplog: pytest.LogCaptureFixture) -> None:
    # Create a LangChain tool
    langchain_tool = add  # type: ignore

    # Create an adapter
    adapter = cast(Tool, LangChainToolAdapter(langchain_tool))  # type: ignore

    # Test schema generation
    schema = adapter.schema

    assert schema["name"] == "add"
    assert "description" in schema
    assert schema["description"] == "Add two numbers"
    assert "parameters" in schema
    assert schema["parameters"]["type"] == "object"
    assert "properties" in schema["parameters"]
    assert "a" in schema["parameters"]["properties"]
    assert "b" in schema["parameters"]["properties"]
    assert schema["parameters"]["properties"]["a"]["type"] == "integer"
    assert schema["parameters"]["properties"]["b"]["type"] == "integer"
    assert "required" in schema["parameters"]
    assert set(schema["parameters"]["required"]) == {"a", "b"}
    assert len(schema["parameters"]["properties"]) == 2

    # Check log.
    with caplog.at_level(logging.INFO):
        # Test run method
        result = await adapter.run_json({"a": 2, "b": 3}, CancellationToken())
        assert result == 5
        assert str(result) in caplog.text

    # Test that the adapter's run method can be called multiple times
    result = await adapter.run_json({"a": 5, "b": 7}, CancellationToken())
    assert result == 12

    # Test CustomCalculatorTool
    custom_langchain_tool = CustomCalculatorTool()
    custom_adapter = LangChainToolAdapter(custom_langchain_tool)  # type: ignore

    # Test schema generation for CustomCalculatorTool
    custom_schema = custom_adapter.schema

    assert custom_schema["name"] == "Calculator"
    assert custom_schema["description"] == "useful for when you need to answer questions about math"  # type: ignore
    assert "parameters" in custom_schema
    assert custom_schema["parameters"]["type"] == "object"
    assert "properties" in custom_schema["parameters"]
    assert "a" in custom_schema["parameters"]["properties"]
    assert "b" in custom_schema["parameters"]["properties"]
    assert custom_schema["parameters"]["properties"]["a"]["type"] == "integer"
    assert custom_schema["parameters"]["properties"]["b"]["type"] == "integer"
    assert "required" in custom_schema["parameters"]
    assert set(custom_schema["parameters"]["required"]) == {"a", "b"}

    # Test run method for CustomCalculatorTool
    custom_result = await custom_adapter.run_json({"a": 3, "b": 4}, CancellationToken())
    assert custom_result == 12


class NoSchemaTool(LangChainTool):
    name: str = "NoSchema"
    description: str = "a tool without an explicit args schema"

    def _run(self, a: int, b: int, run_manager: Optional[CallbackManagerForToolRun] = None) -> int:
        return a + b

    async def _arun(
        self,
        a: int,
        b: int,
        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
    ) -> int:
        return a + b


@pytest.mark.asyncio
async def test_langchain_tool_adapter_skips_run_manager() -> None:
    # Tools without an explicit args_schema get their args inferred from the
    # callable's signature. LangChain injects a ``run_manager`` into ``_run``
    # which is not a user-facing input and cannot be turned into a pydantic
    # schema; the adapter must skip it (see #6385).
    tool = NoSchemaTool()
    adapter = LangChainToolAdapter(tool)  # type: ignore

    schema = adapter.schema
    assert schema["name"] == "NoSchema"
    props = schema["parameters"]["properties"]
    assert set(props.keys()) == {"a", "b"}
    assert set(schema["parameters"]["required"]) == {"a", "b"}

    result = await adapter.run_json({"a": 2, "b": 3}, CancellationToken())
    assert result == 5