fix

Signed-off-by: Prithvi Kannan <[email protected]>

integrations/langchain/src/databricks_langchain/genie.py

@@ -1,5 +1,6 @@
-from databricks_ai_bridge.genie import Genie
 import mlflow
+from databricks_ai_bridge.genie import Genie
+
 
 @mlflow.trace()
 def _concat_messages_array(messages):

pyproject.toml

@@ -14,7 +14,7 @@ dependencies = [
   "pandas",
   "tiktoken>=0.8.0",
   "tabulate",
-  "mlflow",
+  "mlflow-skinny",
 ]
 
 [project.license]

src/databricks_ai_bridge/genie.py

@@ -3,12 +3,11 @@
 from datetime import datetime
 from typing import Union
 
+import mlflow
 import pandas as pd
 import tiktoken
 from databricks.sdk import WorkspaceClient
 
-import mlflow
-
 MAX_TOKENS_OF_DATA = 20000  # max tokens of data in markdown format
 MAX_ITERATIONS = 50  # max times to poll the API when polling for either result or the query results, each iteration is ~1 second, so max latency == 2 * MAX_ITERATIONS
 
@@ -75,6 +74,7 @@ def __init__(self, space_id):
             "Content-Type": "application/json",
         }
 
+    @mlflow.trace()
     def start_conversation(self, content):
         resp = self.genie._api.do(
             "POST",
@@ -84,6 +84,7 @@ def start_conversation(self, content):
         )
         return resp
 
+    @mlflow.trace()
     def create_message(self, conversation_id, content):
         resp = self.genie._api.do(
             "POST",
@@ -93,6 +94,7 @@ def create_message(self, conversation_id, content):
         )
         return resp
 
+    @mlflow.trace()
     def poll_for_result(self, conversation_id, message_id):
         @mlflow.trace()
         def poll_result():
@@ -139,6 +141,7 @@ def poll_result():
                     logging.debug(f"Waiting...: {resp['status']}")
                     time.sleep(5)
 
+        @mlflow.trace()
         def poll_query_results():
             iteration_count = 0
             while iteration_count < MAX_ITERATIONS: