Hello @jcastro I wish you a wonderful day ahead!
Databricks recommends the MLflow ChatAgent interface to author production-grade agents.
These ChatAgent provides advanced agent capabilities like Streaming output to enable interactive user experiences by streaming output in smaller chunks.
If you already have an agent built with LangChain, LangGraph, or a similar framework, you donโt need to rewrite your agent to use it on Databricks. Instead, just wrap your existing agent with the MLflow ChatAgent interface.
If your agent supports streaming, you can use predict_stream, the following is a simplified template for converting your agent:
from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import ChatAgentMessage, ChatAgentResponse, ChatAgentChunk
import uuid
class MyWrappedAgent(ChatAgent):
def __init__(self, agent):
self.agent = agent
def predict_stream(self, messages, context=None, custom_inputs=None):
# If your agent supports streaming
for chunk in self.agent.stream(...):
yield ChatAgentChunk(delta=ChatAgentMessage(role="assistant", content=chunk, id=str(uuid.uuid4())))
For more details about this, I am sharing below documentation for your reference:
https://docs.databricks.com/gcp/en/generative-ai/agent-framework/author-agent#-use-chatagent-to-auth...
https://docs.databricks.com/gcp/en/generative-ai/agent-framework/author-agent#what-if-i-already-have...
In the below URL the following notebooks show how to author streaming and non-streaming ChatAgents using the popular libraries OpenAI, LangGraph, and AutoGen:
https://docs.databricks.com/gcp/en/generative-ai/agent-framework/author-agent#chatagent-examples
