add trace

Signed-off-by: Prithvi Kannan <[email protected]>
databricks · Dec 18, 2024 · 4478ccf · 4478ccf
1 parent f815b62
commit 4478ccf
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/databricks_ai_bridge/genie.py b/src/databricks_ai_bridge/genie.py
@@ -157,6 +157,7 @@ def poll_result():
 
         return poll_result()
 
+    @mlflow.trace()
     def ask_question(self, question):
         resp = self.start_conversation(question)
         return self.poll_for_result(resp["conversation_id"], resp["message_id"])