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]