rename

Signed-off-by: Prithvi Kannan <[email protected]>
databricks · Oct 24, 2024 · 776d8c1 · 776d8c1
1 parent 580c928
commit 776d8c1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git 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
@@ -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