Modify the Agent

import mlflow
from databricks_openai import AsyncDatabricksOpenAI
from agents import set_default_openai_api, set_default_openai_client, Agent
from agents.tracing import set_trace_processors

# Set up async client (recommended for agent servers)
set_default_openai_client(AsyncDatabricksOpenAI())
set_default_openai_api("chat_completions")

# Use MLflow for tracing (disables SDK's built-in tracing)
set_trace_processors([])
mlflow.openai.autolog()

from databricks_openai.agents import McpServer

# UC Functions
uc_server = McpServer(
    url=f"{host}/api/2.0/mcp/functions/{catalog}/{schema}",
    name="uc functions",
)

# Genie Space
genie_server = McpServer(
    url=f"{host}/api/2.0/mcp/genie/{space_id}",
    name="genie space",
)

# Vector Search
vector_server = McpServer(
    url=f"{host}/api/2.0/mcp/vector-search/{catalog}/{schema}/{index}",
    name="vector search",
)

# Add to agent
agent = Agent(
    name="my agent",
    instructions="You are a helpful agent.",
    model="databricks-claude-3-7-sonnet",
    mcp_servers=[uc_server, genie_server, vector_server],
)

agent = Agent(
    name="my agent",
    model="databricks-claude-3-7-sonnet",  # Change here
    ...
)

agent = Agent(
    name="my agent",
    instructions="""You are a helpful data analyst assistant.

    You have access to:
    - Company sales data via Genie
    - Product documentation via vector search

    Always cite your sources when answering questions.""",
    ...
)

from agents import Runner

# Non-streaming
messages = [{"role": "user", "content": "hi"}]
result = await Runner.run(agent, messages)

# Streaming
result = Runner.run_streamed(agent, input=messages)
async for event in result.stream_events():
    # Process stream events
    pass

from agent_server.utils import process_agent_stream_events

result = Runner.run_streamed(agent, input=messages)
async for event in process_agent_stream_events(result.stream_events()):
    yield event  # Yields ResponsesAgentStreamEvent objects