mirror of
https://github.com/microsoft/autogen.git
synced 2025-07-07 17:11:06 +00:00
95bd514a9a
* add initial global search draft * add graphrag dep * fix local search embedding * linting * add from config constructor * remove draft notebook * update config factory and add docstrings * add graphrag sample * add sample prompts * update readme * update deps * Add API docs * Update python/samples/agentchat_graphrag/requirements.txt * Update python/samples/agentchat_graphrag/requirements.txt * update docstrings with snippet and doc ref * lint * improve set up instructions in docstring * lint * update lock * Update python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_global_search.py Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com> * Update python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_local_search.py Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com> * add unit tests * update lock * update uv lock * add docstring newlines * stubs and typing on graphrag tests * fix docstrings * fix mypy error * + linting and type fixes * type fix graphrag sample * Update python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_global_search.py Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com> * Update python/packages/autogen-ext/src/autogen_ext/tools/graphrag/_local_search.py Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com> * Update python/samples/agentchat_graphrag/requirements.txt Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com> * update overrides * fix docstring client imports * additional docstring fix * add docstring missing import * use openai and fix db path * use console for displaying messages * add model config and gitignore * update readme * lint * Update python/samples/agentchat_graphrag/README.md * Update python/samples/agentchat_graphrag/README.md * Comment remaining azure config --------- Co-authored-by: Leonardo Pinheiro <lpinheiro@microsoft.com> Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>
99 lines
3.7 KiB
Python
99 lines
3.7 KiB
Python
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() -> 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
|
|
|
|
# Test run method
|
|
result = await adapter.run_json({"a": 2, "b": 3}, CancellationToken())
|
|
assert result == 5
|
|
|
|
# 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
|