How to use `predict_stream` in Databricks deployed ChatAgent.

jcastro — Tue, 15 Apr 2025 21:10:55 GMT

Hello,

We are working on a solution to have a streaming chat agent. Our agent works fine in the Playground but when we try to call it as an api under /invocations, the streaming does not work.

Re: How to use `predict_stream` in Databricks deployed ChatAgent.

Vinay_M_R — Tue, 17 Jun 2025 13:22:05 GMT

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-author-agents
https://docs.databricks.com/gcp/en/generative-ai/agent-framework/author-agent#what-if-i-already-have-an-agent

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

topic How to use `predict_stream` in Databricks deployed ChatAgent. in Generative AI

How to use `predict_stream` in Databricks deployed ChatAgent.

Re: How to use `predict_stream` in Databricks deployed ChatAgent.