Add back description param to GenieAgent (#123)

Signed-off-by: Bryan Qiu <[email protected]>

Author: bbqiu

integrations/langchain/src/databricks_langchain/genie.py

@@ -40,9 +40,10 @@ def _query_genie_as_agent(input, genie: Genie, genie_agent_name):
 def GenieAgent(
     genie_space_id,
     genie_agent_name: str = "Genie",
+    description: str = "",
     client: Optional["WorkspaceClient"] = None,
 ):
-    """Create a genie agent that can be used to query the API"""
+    """Create a genie agent that can be used to query the API. If a description is not provided, the description of the genie space will be used."""
     if not genie_space_id:
         raise ValueError("genie_space_id is required to create a GenieAgent")
 
@@ -61,5 +62,5 @@ def GenieAgent(
 
     runnable = RunnableLambda(partial_genie_agent)
     runnable.name = genie_agent_name
-    runnable.description = genie.description
+    runnable.description = description or genie.description
     return runnable

integrations/langchain/tests/unit_tests/test_genie.py

@@ -72,12 +72,31 @@ def test_create_genie_agent(MockRunnableLambda, MockWorkspaceClient):
         title="Sales Space",
         description="description",
     )
-    MockWorkspaceClient.genie.get_space.return_value = mock_space
+    mock_client = MockWorkspaceClient.return_value
+    mock_client.genie.get_space.return_value = mock_space
 
-    agent = GenieAgent("space-id", "Genie", MockWorkspaceClient)
+    agent = GenieAgent("space-id", "Genie", client=mock_client)
     assert agent.description == "description"
 
-    MockWorkspaceClient.genie.get_space.assert_called_once()
+    mock_client.genie.get_space.assert_called_once()
+    assert agent == MockRunnableLambda.return_value
+
+
+@patch("databricks.sdk.WorkspaceClient")
+@patch("langchain_core.runnables.RunnableLambda")
+def test_create_genie_agent_with_description(MockRunnableLambda, MockWorkspaceClient):
+    mock_space = GenieSpace(
+        space_id="space-id",
+        title="Sales Space",
+        description=None,
+    )
+    mock_client = MockWorkspaceClient.return_value
+    mock_client.genie.get_space.return_value = mock_space
+
+    agent = GenieAgent("space-id", "Genie", "this is a description", client=mock_client)
+    assert agent.description == "this is a description"
+
+    mock_client.genie.get_space.assert_called_once()
     assert agent == MockRunnableLambda.return_value