Hello! So I am using the built-in code when we initiate the databricks custom app using agent. After tweaking and using AI by Gemini and Claude, I was finally able to start the conversation and receive the response! Sharing the code below for anyone to play with. Of course, this code only returns the first round of conversation and subsequently fails in second turn. I may look into this further, but it seems I may have to flag this to Databricks support as the built-in template should work out of the box regardless it is a reasoning or an instruction model.
import json
import logging
from contextlib import AsyncExitStack
from datetime import datetime
from typing import AsyncGenerator
# --- PYDANTIC MONKEY-PATCH FOR REASONING MODELS ---
def _flatten_text_payload(value):
if value is None:
return ""
if isinstance(value, str):
return value
if isinstance(value, list):
parts = []
for chunk in value:
if isinstance(chunk, dict):
if chunk.get("type") in ("text", "output_text"):
parts.append(str(chunk.get("text", "")))
elif "text" in chunk:
parts.append(str(chunk.get("text", "")))
elif isinstance(chunk, str):
parts.append(chunk)
else:
parts.append(str(chunk))
return "".join(parts)
if isinstance(value, dict):
if value.get("type") in ("text", "output_text"):
return str(value.get("text", ""))
return str(value.get("text", str(value)))
return str(value)
def _patch_pydantic_class(cls):
if cls is None:
return
orig_init = cls.__init__
def patched_init(self, *args, **kwargs):
if "text" in kwargs:
kwargs["text"] = _flatten_text_payload(kwargs["text"])
orig_init(self, *args, **kwargs)
cls.__init__ = patched_init
orig_validate = getattr(cls, "model_validate", None)
if orig_validate:
@classmethod
def patched_validate(target_cls, obj, *args, **kwargs):
if isinstance(obj, dict) and "text" in obj:
obj["text"] = _flatten_text_payload(obj["text"])
return orig_validate(obj, *args, **kwargs)
cls.model_validate = patched_validate
try:
from openai.types.responses import ResponseOutputText as OpenAIResponseOutputText
_patch_pydantic_class(OpenAIResponseOutputText)
except ImportError:
pass
try:
from mlflow.types.responses import ResponseOutputText as MLflowResponseOutputText
_patch_pydantic_class(MLflowResponseOutputText)
except ImportError:
pass
# ---------------------------------------------------
import mlflow
from agents import Agent, Runner, function_tool, set_default_openai_api, set_default_openai_client
from agents.tracing import set_trace_processors
from databricks.sdk import WorkspaceClient
from databricks_openai import AsyncDatabricksOpenAI
from databricks_openai.agents import McpServer
from mlflow.genai.agent_server import invoke, stream
from mlflow.types.responses import (
ResponsesAgentRequest,
ResponsesAgentResponse,
ResponsesAgentStreamEvent,
)
from agent_server.utils import build_mcp_url, get_session_id
logger = logging.getLogger(__name__)
set_default_openai_client(AsyncDatabricksOpenAI())
set_default_openai_api("chat_completions")
set_trace_processors([])
mlflow.openai.autolog()
logging.getLogger("mlflow.utils.autologging_utils").setLevel(logging.ERROR)
def _sanitize_item_dict(dump: dict) -> dict:
"""
Repair malformed 'text' leaves (some reasoning models return text as a
nested list/dict instead of a plain string) WITHOUT collapsing the
container shapes (content / summary) that the Agents SDK's
chat_completions converter expects to remain lists of parts.
IMPORTANT: do not flatten `content` or `summary` themselves into plain
strings -- Converter.items_to_messages() in the openai-agents SDK
indexes into these as list-of-parts (e.g. content[0]["text"]) when
reconstructing multi-turn history. Flattening them causes:
TypeError: string indices must be integers, not 'str'
on the *second* turn, once history round-trips back through the API.
"""
# Fix a malformed top-level "text" leaf without touching containers
if isinstance(dump.get("text"), (list, dict)):
dump["text"] = _flatten_text_payload(dump["text"])
# "content" must stay a list-of-parts for the chat_completions converter --
# only repair malformed "text" leaves inside each part
content = dump.get("content")
if isinstance(content, list):
for part in content:
if isinstance(part, dict) and isinstance(part.get("text"), (list, dict)):
part["text"] = _flatten_text_payload(part["text"])
elif content is not None and not isinstance(content, (str, list)):
dump["content"] = _flatten_text_payload(content)
# Reasoning items use "summary": [{"type": "summary_text", "text": ...}]
summary = dump.get("summary")
if isinstance(summary, list):
for part in summary:
if isinstance(part, dict) and isinstance(part.get("text"), (list, dict)):
part["text"] = _flatten_text_payload(part["text"])
return dump
def sanitize_input_messages(input_items):
clean = []
for item in input_items:
dump = item.model_dump() if hasattr(item, "model_dump") else item
if isinstance(dump, dict):
clean.append(_sanitize_item_dict(dump))
else:
clean.append(dump)
return clean
@function_tool
def get_current_time() -> str:
"""Get the current date and time."""
return datetime.now().isoformat()
def create_agent(mcp_servers: list[McpServer] | None = None) -> Agent:
return Agent(
name="Agent",
instructions="You are a helpful assistant.",
model="databricks-qwen35-122b-a10b",
tools=[get_current_time],
mcp_servers=mcp_servers or [],
)
@invoke()
async def invoke_handler(request: ResponsesAgentRequest) -> ResponsesAgentResponse:
if session_id := get_session_id(request):
mlflow.update_current_trace(metadata={"mlflow.trace.session": session_id})
try:
async with AsyncExitStack() as stack:
agent = create_agent()
messages = sanitize_input_messages(request.input)
result = await Runner.run(agent, messages)
cleaned_outputs = []
for item in result.new_items:
input_item = item.to_input_item()
if isinstance(input_item, dict):
input_item = _sanitize_item_dict(input_item)
cleaned_outputs.append(input_item)
return ResponsesAgentResponse(output=cleaned_outputs)
except Exception:
logger.exception("invoke_handler failed")
raise
@stream()
async def stream_handler(
request: ResponsesAgentRequest,
) -> AsyncGenerator[ResponsesAgentStreamEvent, None]:
if session_id := get_session_id(request):
mlflow.update_current_trace(metadata={"mlflow.trace.session": session_id})
try:
async with AsyncExitStack() as stack:
agent = create_agent()
messages = sanitize_input_messages(request.input)
result = Runner.run_streamed(agent, input=messages)
# Directly stream events safely without process_agent_stream_events
async for event in result.stream_events():
delta_val = None
if hasattr(event, "delta"):
delta_val = _flatten_text_payload(event.delta)
elif hasattr(event, "data") and hasattr(event.data, "delta"):
delta_val = _flatten_text_payload(event.data.delta)
elif hasattr(event, "item") and hasattr(event.item, "text"):
delta_val = _flatten_text_payload(event.item.text)
if delta_val:
yield ResponsesAgentStreamEvent(delta=delta_val)
except Exception:
logger.exception("stream_handler failed")
raise
