🐛 fix(custom.py): fix JsonAgent.from_toolkit_and_llm method to handle both list and JsonToolkit input

🔥 chore(test_graph.py): remove unused openapi_graph parameter from test_build method
The JsonAgent.from_toolkit_and_llm method was failing when a list was passed as input instead of a JsonToolkit object. The fix now handles both cases. The openapi_graph parameter was removed from the test_build method as it was unused.
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-05-31 11:44:12 -03:00
commit 25d77eaf87
2 changed files with 2 additions and 3 deletions

View file

@ -69,7 +69,7 @@ class JsonAgent(CustomAgentExecutor):
@classmethod
def from_toolkit_and_llm(cls, toolkit: JsonToolkit, llm: BaseLanguageModel):
tools = toolkit.get_tools()
tools = toolkit if isinstance(toolkit, list) else toolkit.get_tools()
tool_names = {tool.name for tool in tools}
prompt = ZeroShotAgent.create_prompt(
tools,

View file

@ -237,11 +237,10 @@ def test_build_params(basic_graph):
assert "memory" in root.params
def test_build(basic_graph, complex_graph, openapi_graph):
def test_build(basic_graph, complex_graph):
"""Test Node's build method"""
assert_agent_was_built(basic_graph)
assert_agent_was_built(complex_graph)
assert_agent_was_built(openapi_graph)
def assert_agent_was_built(graph):