From ff394b4729f4160121f62ad7b8a6fa748f9eec64 Mon Sep 17 00:00:00 2001 From: Moiz Sajid Date: Fri, 15 Mar 2024 02:15:26 +0500 Subject: [PATCH] Fixed linting problems --- .../functions_utils.py | 8 +++++--- .../tests/integration_tests/test_tools.py | 20 ++++++++++++++++--- 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/libs/vertexai/langchain_google_vertexai/functions_utils.py b/libs/vertexai/langchain_google_vertexai/functions_utils.py index 29a103c0..92239dea 100644 --- a/libs/vertexai/langchain_google_vertexai/functions_utils.py +++ b/libs/vertexai/langchain_google_vertexai/functions_utils.py @@ -55,7 +55,7 @@ def _format_tool_to_vertex_function(tool: BaseTool) -> FunctionDescription: def _format_tools_to_vertex_tool( - tools: List[Union[BaseTool, Type[BaseModel]]], + tools: List[Union[BaseTool, Type[BaseModel], dict]], ) -> List[VertexTool]: "Format tool into the Vertex Tool instance." function_declarations = [] @@ -64,12 +64,14 @@ def _format_tools_to_vertex_tool( func = _format_tool_to_vertex_function(tool) elif isinstance(tool, type) and issubclass(tool, BaseModel): func = _format_pydantic_to_vertex_function(tool) - else: + elif isinstance(tool, dict): func = { "name": tool["name"], - "description": tool.get("description"), + "description": tool.pop("description"), "parameters": _get_parameters_from_schema(tool["parameters"]), } + else: + raise ValueError(f"Unsupported tool call type {tool}") function_declarations.append(FunctionDeclaration(**func)) return [VertexTool(function_declarations=function_declarations)] diff --git a/libs/vertexai/tests/integration_tests/test_tools.py b/libs/vertexai/tests/integration_tests/test_tools.py index c31a1a1b..5df76d75 100644 --- a/libs/vertexai/tests/integration_tests/test_tools.py +++ b/libs/vertexai/tests/integration_tests/test_tools.py @@ -95,9 +95,12 @@ def test_tools() -> None: @pytest.mark.extended def test_custom_tool() -> None: - from langchain.agents import AgentExecutor, create_openai_functions_agent, tool + from langchain.agents import AgentExecutor, tool + from langchain.agents.format_scratchpad import ( + format_to_openai_function_messages, + ) - @tool + @tool("search", return_direct=True) def search(query: str) -> str: """Look up things online.""" return "LangChain" @@ -115,8 +118,19 @@ def search(query: str) -> str: MessagesPlaceholder("agent_scratchpad"), ] ) + llm_with_tools = llm.bind(functions=tools) - agent = create_openai_functions_agent(llm, tools, prompt) + agent: Any = ( + { + "input": lambda x: x["input"], + "agent_scratchpad": lambda x: format_to_openai_function_messages( + x["intermediate_steps"] + ), + } + | prompt + | llm_with_tools + | _TestOutputParser() + ) agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True) response = agent_executor.invoke({"input": "What is LangChain?"})