From 776d8c1f0fe87a7270f6cc5e46488d93ecff74ce Mon Sep 17 00:00:00 2001 From: Prithvi Kannan Date: Thu, 24 Oct 2024 10:06:24 -0700 Subject: [PATCH] rename Signed-off-by: Prithvi Kannan --- integrations/langchain/src/databricks_langchain/genie.py | 2 +- integrations/langchain/tests/test_genie.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/integrations/langchain/src/databricks_langchain/genie.py b/integrations/langchain/src/databricks_langchain/genie.py index 0528bfe..153c2df 100644 --- a/integrations/langchain/src/databricks_langchain/genie.py +++ b/integrations/langchain/src/databricks_langchain/genie.py @@ -32,7 +32,7 @@ def _query_genie_as_agent(input, genie_space_id, genie_agent_name): return {"messages": [AIMessage(content="")]} -def create_genie_agent(genie_space_id, genie_agent_name="Genie"): +def GenieAgent(genie_space_id, genie_agent_name="Genie", description=""): """Create a genie agent that can be used to query the API""" from functools import partial diff --git a/integrations/langchain/tests/test_genie.py b/integrations/langchain/tests/test_genie.py index 9e41bd3..70c6c28 100644 --- a/integrations/langchain/tests/test_genie.py +++ b/integrations/langchain/tests/test_genie.py @@ -3,9 +3,9 @@ from langchain_core.messages import AIMessage from databricks_langchain.genie import ( + GenieAgent, _concat_messages_array, _query_genie_as_agent, - create_genie_agent, ) @@ -64,7 +64,7 @@ def test_query_genie_as_agent(MockGenie): def test_create_genie_agent(MockRunnableLambda): mock_runnable = MockRunnableLambda.return_value - agent = create_genie_agent("space-id", "Genie") + agent = GenieAgent("space-id", "Genie") assert agent == mock_runnable # Check that the partial function is created with the correct arguments