Modify the Agent

import mlflow
from databricks.sdk import WorkspaceClient
from databricks_langchain import ChatDatabricks, DatabricksMCPServer, DatabricksMultiServerMCPClient
from langchain.agents import create_agent

# Enable autologging for tracing
mlflow.langchain.autolog()

# Initialize workspace client
workspace_client = WorkspaceClient()

from databricks_langchain import ChatDatabricks

llm = ChatDatabricks(
    endpoint="databricks-claude-3-7-sonnet",  # or databricks-meta-llama-3-1-70b-instruct
    temperature=0,
    max_tokens=500,
)

# For Responses API agents:
llm = ChatDatabricks(endpoint="my-agent-endpoint", use_responses_api=True)

from databricks_langchain import DatabricksEmbeddings

embeddings = DatabricksEmbeddings(endpoint="databricks-bge-large-en")
vector = embeddings.embed_query("The meaning of life is 42")
vectors = embeddings.embed_documents(["doc1", "doc2"])

from databricks_langchain import DatabricksVectorSearch

# Delta-sync index with Databricks-managed embeddings
vs = DatabricksVectorSearch(index_name="catalog.schema.index_name")

# Direct-access or self-managed embeddings
vs = DatabricksVectorSearch(
    index_name="catalog.schema.index_name",
    embedding=embeddings,
    text_column="content",
)

docs = vs.similarity_search("query", k=5)

from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient

client = DatabricksMultiServerMCPClient([
    DatabricksMCPServer(
        name="system-ai",
        url=f"{host}/api/2.0/mcp/functions/system/ai",
    )
])
tools = await client.get_tools()

server = DatabricksMCPServer.from_uc_function(
    catalog="main",
    schema="tools",
    function_name="send_email",  # Optional - omit for all functions in schema
    name="email-server",
    timeout=30.0,
    handle_tool_error=True,
)

server = DatabricksMCPServer.from_vector_search(
    catalog="main",
    schema="embeddings",
    index_name="product_docs",  # Optional - omit for all indexes in schema
    name="docs-search",
    timeout=30.0,
)

DatabricksMCPServer(
    name="genie",
    url=f"{host_name}/api/2.0/mcp/genie/01f0515f6739169283ef2c39b7329700",
)

from databricks_langchain import MCPServer

server = MCPServer(
    name="external-server",
    url="https://other-server.com/mcp",
    headers={"X-API-Key": "secret"},
    timeout=15.0,
)

from langchain.agents import create_agent

# Create agent - ONLY accepts tools and model, NO prompt/instructions parameter
agent = create_agent(tools=tools, model=llm)

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

# Streaming
async for event in agent.astream(input=messages, stream_mode=["updates", "messages"]):
    # Process stream events
    pass

from agent_server.utils import process_agent_astream_events

async for event in process_agent_astream_events(
    agent.astream(input=messages, stream_mode=["updates", "messages"])
):
    yield event  # Yields ResponsesAgentStreamEvent objects

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

@stream()
async def streaming(request: ResponsesAgentRequest) -> AsyncGenerator[ResponsesAgentStreamEvent, None]:
    agent = await init_agent()
    # Prepend system instructions to user messages
    user_messages = to_chat_completions_input([i.model_dump() for i in request.input])
    messages = {"messages": [{"role": "system", "content": AGENT_INSTRUCTIONS}] + user_messages}

    async for event in process_agent_astream_events(
        agent.astream(input=messages, stream_mode=["updates", "messages"])
    ):
        yield event

# WRONG - will cause "unexpected keyword argument" error
agent = create_agent(tools=tools, model=llm, prompt=AGENT_INSTRUCTIONS)

# CORRECT - add instructions via messages
messages = {"messages": [{"role": "system", "content": AGENT_INSTRUCTIONS}] + user_messages}