From 0f4c2ef8051d5be28cc8412ce925e910ea230c3b Mon Sep 17 00:00:00 2001 From: Yuki Watanabe <31463517+B-Step62@users.noreply.github.com> Date: Wed, 20 Nov 2024 09:32:01 +0900 Subject: [PATCH] Trying to un-flake agent test (#39) Signed-off-by: B-Step62 --- libs/databricks/tests/integration_tests/test_chat_models.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/libs/databricks/tests/integration_tests/test_chat_models.py b/libs/databricks/tests/integration_tests/test_chat_models.py index d2fe498..729d496 100644 --- a/libs/databricks/tests/integration_tests/test_chat_models.py +++ b/libs/databricks/tests/integration_tests/test_chat_models.py @@ -352,7 +352,7 @@ def test_chat_databricks_agent_executor(): temperature=0, max_tokens=100, ) - tools = [add, multiply] + tools = [add] prompt = ChatPromptTemplate.from_messages( [ ("system", "You are a helpful assistant"), @@ -364,8 +364,8 @@ def test_chat_databricks_agent_executor(): agent = create_tool_calling_agent(model, tools, prompt) agent_executor = AgentExecutor(agent=agent, tools=tools) - response = agent_executor.invoke({"input": "What is (10 + 5) * 3?"}) - assert "45" in response["output"] + response = agent_executor.invoke({"input": "What is 1 + 1?"}) + assert "2" in response["output"] def test_chat_databricks_langgraph():