From 5994b85401a68cd23c7e7e20b86bf91c586e6f39 Mon Sep 17 00:00:00 2001 From: Prithvi Kannan Date: Wed, 18 Dec 2024 14:30:03 -0800 Subject: [PATCH] trace type Signed-off-by: Prithvi Kannan --- integrations/langchain/src/databricks_langchain/genie.py | 2 +- src/databricks_ai_bridge/genie.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/integrations/langchain/src/databricks_langchain/genie.py b/integrations/langchain/src/databricks_langchain/genie.py index a0a7c7e..87598ee 100644 --- a/integrations/langchain/src/databricks_langchain/genie.py +++ b/integrations/langchain/src/databricks_langchain/genie.py @@ -35,7 +35,7 @@ def _query_genie_as_agent(input, genie_space_id, genie_agent_name): return {"messages": [AIMessage(content="")]} -@mlflow.trace(type="AGENT") +@mlflow.trace(span_type="AGENT") 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/src/databricks_ai_bridge/genie.py b/src/databricks_ai_bridge/genie.py index f659eab..b69b6b8 100644 --- a/src/databricks_ai_bridge/genie.py +++ b/src/databricks_ai_bridge/genie.py @@ -26,7 +26,7 @@ class GenieResponse: description: Optional[str] = "" -@mlflow.trace(type="PARSER") +@mlflow.trace(span_type="PARSER") def _parse_query_result(resp) -> Union[str, pd.DataFrame]: columns = resp["manifest"]["schema"]["columns"] header = [str(col["name"]) for col in columns]