Modify the Agent

Main File

agent_server/agent.py - Agent logic, model selection, instructions, MCP servers

Key Files

SDK Setup

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()

Adding MCP Servers

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],
)

After adding MCP servers: Grant permissions in databricks.yml (see add-tools skill)

Changing the Model

Available models (check workspace for current list):

• databricks-claude-3-7-sonnet
• databricks-claude-3-5-sonnet
• databricks-meta-llama-3-3-70b-instruct

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

Note: Some workspaces require granting the app access to the serving endpoint in databricks.yml. See the add-tools skill and examples/serving-endpoint.yaml.

Changing Instructions

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.""",
    ...
)

Running the Agent

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

Converting to Responses API format: Use processagentstreamevents() from agentserver/utils.py to convert streaming output to Responses API compatible format:

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

External Resources

1. databricks-openai SDK
2. Agent examples
3. Agent Framework docs
4. Adding tools
5. OpenAI Agents SDK
6. Responses API

Next Steps

• Discover available tools: see discover-tools skill
• Grant resource permissions: see add-tools skill
• Test locally: see run-locally skill
• Deploy: see deploy skill