Hi @moemedina. No, I didn't.

I'm considering using ChatModel/ChatAgent class to wrap the graph and be able to move on. However, the MLflow documentation is still referring to ChatModel where Chat Agent is the latest recommendation:

• MLflow ChatModel Doc: https://mlflow.org/docs/latest/llms/chat-model-intro/ 
• Databricks Doc: https://docs.databricks.com/aws/en/generative-ai/agent-framework/author-agent 
• ChatAgent API Doc: https://mlflow.org/docs/latest/api_reference/python_api/mlflow.pyfunc.html#mlflow.pyfunc.ChatAgent