Swiss Travel Agent Lab

A hands-on lab for building AI agents on Databricks using LangGraph, Unity Catalog functions, MCP (Model Context Protocol), Unity AI Gateway, and Managed Memory. You’ll iteratively create a Swiss travel assistant that queries live train schedules, weather forecasts, and activity recommendations — with tools governed by Unity Catalog, the LLM governed by a Unity AI Gateway model service (routing, rate limits, guardrails, and usage logging defined once and inherited by every lab), and durable per-user memory governed by Unity Catalog.

The labs are designed for a 2–3 hour hands-on session.

Based on the blog series by Alexander Genser:

• Building an AI Agent for Transportation Applications with Databricks
• Advancing Your Agentic AI with MCP — A Travel Planning Use Case on Databricks

Original code: genseral/ai_agent_transport

---

What You’ll Build

A travel planning agent that evolves across four labs:

Lab 1 Architecture

A ReAct-style LangGraph agent that calls Unity Catalog functions to query the Swiss Transport API for live train connections and station departures.

Lab 2 Architecture

The agent evolves into a ResponsesAgent that loads tools from three MCP server types — managed (UC functions as MCP), custom (FastMCP weather server on a Databricks App), and external (Tavily web search via UC Connection). Its LLM continues to route through the Lab 1 Unity AI Gateway model service. Egress note: a Model Serving endpoint can’t reach a Databricks App, so the deployed Lab 2 agent wires only the two workspace-internal MCP servers (managed + external); the custom weather App is built and verified here but wired into the agent in Lab 3 (which runs on Apps, where app-to-app egress works).

Lab 3 Architecture

The agent is packaged as a Databricks App using the OpenAI Agents SDK, wrapped in MLflow ResponsesAgent. It connects to the same three MCP servers from Labs 1–2, with a custom Swiss-themed React + Tailwind CSS chat UI featuring real-time SSE streaming, tool flow visualization, an MCP tool call sidebar (grouped by server type), dark/light mode, rich markdown rendering, and automatic MLflow tracing — all served from a single FastAPI process without a serving endpoint.

Lab 4 Architecture

The Lab 3 app gains durable, user-scoped memory via Databricks Managed Memory (Beta) — a Unity Catalog MEMORY_STORE securable read/written purely over REST, with no database, Vector Search, or embedding endpoint to run. Two primitives share one store. Entries give the agent semantic facts and procedural rules it learns and reuses across conversations (five tools, organized by a layered /memories/semantic|procedural|episodic/ path convention) — this is the deployed app’s durable memory. Conversations are the managed server-side chat thread (threaded via the OpenAI Agents SDK OpenAIConversationsSession); you create and re-attach one in the notebook, where it runs as the signed-in user. In the deployed multi-user app, per-user conversations need an on-behalf-of scope (unity-catalog) that Databricks Apps can’t grant in this beta, so the app relies on entries and degrades gracefully. Every memory is partitioned by an opaque per-user scope set in trusted server code — the isolation boundary. Memory is additive to the three MCP tool types, not a replacement.

The app surfaces memory live in a Memory pane beside the MCP tool-call sidebar — the recall/save interactions each turn, and the current entries with layer + operation badges:

---

Prerequisites

• Databricks workspace with Unity Catalog enabled (any cloud — AWS, Azure, or GCP)
• Serverless compute — required for UC function execution and Databricks Connect
• Foundation Model access — the labs use a Databricks-hosted Claude model (currently databricks-claude-sonnet-4-6) via the Foundation Model API (FMAPI). All foundation models are available out of the box — no endpoint creation needed.
• Unity AI Gateway enabled for your account (Account Console → Previews) — Lab 1 puts the LLM behind a governed model service. You also need CREATE SERVICE on your catalog/schema.
• Databricks Runtime 15.0+ with Python 3.10+
• API key (for Lab 2): Tavily (free tier, card-free) for web search. The weather MCP server uses Open-Meteo, which is keyless — no signup needed.
• Databricks Apps — required for Labs 2 and 3 (deploying custom MCP servers and the agent app)

No prior Databricks agent experience required — the labs are self-contained and guide you step by step.

---

Before You Begin

Complete these steps before opening the first notebook:

1. Clone the repository into your workspace (see Getting Started below)
2. Verify your catalog — Each notebook has catalog and schema widgets at the top (defaulting to workshop and travel_agent_). If your workspace uses a different catalog, set the Catalog widget before running — and use the same catalog/schema across all three labs so they share the resources you create
3. Verify Foundation Model access — All Databricks-hosted foundation models are available out of the box via FMAPI — no endpoint creation needed. Browse available models and gateway services under AI/ML > AI Gateway in the sidebar. In Lab 1 you’ll create a Unity AI Gateway model service (travel_agent_llm) that fronts a Claude model (currently Sonnet 4.6); every lab reaches the LLM through that service’s name rather than naming the model directly
4. Get an API key (needed for Lab 2):
   • Tavily — free tier, card-free (web search)
   • No weather key needed — the weather MCP server uses keyless Open-Meteo
5. Store the Tavily key in Databricks Secrets — Lab 2 walks you through this, but you can do it in advance. See the “One-Time Setup” cell in Lab 2

---

Repository Structure

travel-agent-lab/
├── notebooks/
│   ├── 01-travel-agent-foundations.py     # Lab 1: LangGraph + UC function tools
│   ├── 02-travel-agent-mcp-extension.py   # Lab 2: Three types of MCP servers
│   ├── 03-travel-agent-app.py            # Lab 3: Agent as a Databricks App
│   ├── 04-agent-memory.py                # Lab 4: Durable managed memory
│   ├── langgraph-mcp-tool-calling-agent.py # Reference: ResponsesAgent + MCP
│   └── reference-implementation.md         # Architecture and design notes
├── agent-app/                             # Agent Databricks App (Labs 3–4) — [README](agent-app/README.md)
│   ├── agent_server/                      # Lab 3 app (no memory)
│   │   ├── agent.py                       # Swiss Travel Agent (OpenAI Agents SDK + MCP)
│   │   ├── utils.py                       # Stream processing, tool flow messages, session helpers
│   │   └── start_server.py               # AgentServer bootstrap + static file serving
│   ├── agent_server_memory/               # Lab 4 overlay (copied onto agent_server/ by notebook 04)
│   │   ├── agent.py                       # Lab 3 agent + managed-memory wiring
│   │   └── utils_memory.py               # Managed-memory tools, scope, episodic session
│   ├── frontend/                          # Custom React + Tailwind CSS chat UI
│   │   ├── src/components/               # ChatWindow, MessageBubble, ToolFlowTracker, ToolCallSidebar, etc.
│   │   ├── src/hooks/                    # useChat (streaming), useTheme (dark/light)
│   │   └── package.json                  # React, Vite, Tailwind, react-markdown
│   ├── scripts/
│   │   └── start_app.py                  # Builds frontend + starts backend
│   ├── app.yaml                          # Databricks App config (env vars)
│   ├── pyproject.toml                    # Dependencies and scripts
│   └── .env.example                      # Local dev template
├── mcp-weather-server/                     # Custom MCP server (Lab 2)
│   ├── app.py                              # FastMCP weather server (Open-Meteo)
│   ├── app.yaml                            # Databricks App config
│   └── requirements.txt                    # App dependencies
├── assets/
│   ├── diagrams/                           # Architecture diagrams (PNG)
│   │   ├── lab1-architecture.png
│   │   ├── lab1-react-loop.png
│   │   ├── lab2-agent-architecture.png
│   │   ├── lab2-mcp-architecture.png
│   │   ├── lab3-agent-app-architecture.png
│   │   └── lab4-agent-memory-architecture.png
│   └── screenshots/                        # App UI screenshots
│       └── memory-and-mcp-panes.png        # Lab 4 Memory pane + MCP sidebar
├── scripts/
│   ├── run_cell.py                         # Local testing: persistent kernel runner
│   └── _kernel_process.py                  # Local testing: kernel launcher
├── requirements-dev.txt                    # Local dev dependencies
├── .env.example                            # Template for local secrets
└── README.md

---

Getting Started

1. Clone the repo into your Databricks workspace (recommended)

Create a Git folder in your Databricks workspace and clone this repository directly:

1. In the Databricks sidebar, click Workspace
2. Navigate to your user folder (/Workspace/Users//)
3. Click the dropdown menu, then Create > Git folder
4. Paste the repository URL and click Create Git folder

That’s it — the notebooks are ready to run. No local setup needed.

Alternative: import via CLI

If you prefer not to use Git folders, you can import the notebooks and MCP server code directly:

databricks workspace import-dir notebooks/ /Workspace/Users//travel-agent-lab/notebooks/ --profile  --overwrite
databricks workspace import-dir mcp-weather-server/ /Workspace/Users//travel-agent-lab/mcp-weather-server/ --profile  --overwrite
databricks workspace import-dir agent-app/ /Workspace/Users//travel-agent-lab/agent-app/ --profile  --overwrite
databricks workspace import-dir assets/ /Workspace/Users//travel-agent-lab/assets/ --profile  --overwrite

2. Start with Lab 1

Open 01-travel-agent-foundations.py and run through the cells. The notebook covers:

1. Install dependencies — databricks-langchain, unitycatalog-langchain, mlflow
2. Configure catalog and schema — set up your governed namespace
3. Govern the LLM with Unity AI Gateway — create a travel_agent_llm model service (a UC securable) that fronts a Claude foundation model, with traffic routing, a rate limit, an inference table, and a jailbreak guardrail; set up a dedicated service principal for the agent’s LLM calls. The agent then reaches the LLM via ChatDatabricks(use_ai_gateway=True) pointed at the service — governance defined once, inherited everywhere
4. Create UC functions — Python UDFs that call the Swiss Transport API for train connections and station departures
5. Build a LangGraph agent — ReAct-style tool-calling loop over the governed LLM
6. Test interactively — Ask about train connections between Swiss cities
7. Log with MLflow — Register the agent using the models-from-code pattern
8. Deploy — Serve with agents.deploy() and query the live endpoint

Prerequisite: Lab 1 must be completed first — Lab 2 uses the UC functions created in Lab 1.

3. Continue with Lab 2

Open 02-travel-agent-mcp-extension.py to extend the agent with MCP:

1. Managed MCP — Expose Lab 1’s UC functions as an MCP server (zero code change)
2. Custom MCP — Deploy a FastMCP weather server as a Databricks App with Open-Meteo (keyless)
3. External MCP — Connect to Tavily’s web search via UC Connection
4. Build the unified agent — A ResponsesAgent whose LLM routes through the Lab 1 model service and that loads tools from the MCP servers
5. Deploy — Log, register, and deploy the travel planner (LLM via the model service + the dedicated SP). The served agent wires the two workspace-internal MCP servers; the weather App is deferred to Lab 3 (egress)

Prerequisite: Labs 1 and 2 must be completed first — Lab 3 reuses the MCP servers deployed in Labs 1–2.

4. Complete with Lab 3

Open 03-travel-agent-app.py to package the agent as a Databricks App:

1. Explore the agent app — OpenAI Agents SDK with MCP, wrapped in MLflow ResponsesAgent
2. Configure — Write user-specific MCP config to app.yaml
3. Deploy — Create the app with auto-granted permissions via Databricks SDK
4. Grant app-to-app access — CAN_USE on the custom MCP weather server
5. Test — Built-in chat UI, streaming, and programmatic API test
6. Monitor — Automatic MLflow tracing via AgentServer

Prerequisite: Labs 1–3 must be completed first — Lab 4 adds memory to the deployed Lab 3 app.

5. Add durable memory with Lab 4

Open 04-agent-memory.py to give the agent memory that survives restarts and new sessions:

1. Enable the preview — turn on “Managed Memory for Agents” and confirm API access
2. Create the memory store — one UC MEMORY_STORE (REST) + grants (mind the USE_CATALOG/USE_SCHEMA gotcha)
3. Episodic memory — durable conversations threaded via OpenAIConversationsSession
4. Semantic & procedural memory — five entry tools with a layered /memories/… path convention
5. See it personalize — teach a fact + a rule, then watch a fresh request recall and honour them
6. Per-user isolation — scope is the boundary; one user never sees another’s memories
7. Redeploy, inspect & monitor — memories as governed UC objects, traced in MLflow

---

Cleanup

The workshop creates a few resources that keep costing while they run — the two Databricks Apps run continuous compute, and the serving endpoints bill when active. Remove them when you’re done. Each notebook ends with a commented-out “Cleanup” cell; Lab 3’s cleanup is a full teardown for all three labs (uncomment and run it). Alternatively delete resources per-lab, or from the UI (Compute > Apps and Serving).

The remaining objects have no runtime compute cost and are safe to leave (delete manually from Catalog Explorer if you want a clean slate): UC functions (get_connections, get_station_board), the Tavily UC connection, Databricks Secrets (travel-agent-keys), MLflow experiments, the inference table (travel_agent_llm_payload), the memory store, and the schema itself.

---

Key Concepts

Unity Catalog Functions as Tools

Python UDFs registered in Unity Catalog that wrap external APIs. They become governed, discoverable, and reusable tools — any agent in the workspace can use them with proper permissions.

Unity AI Gateway (governed LLM)

The agent’s LLM is not called directly — it’s fronted by a Unity AI Gateway model service, a new Unity Catalog securable at catalog.schema. (here, travel_agent_llm) that wraps one or more Foundation Model API models behind governed routing / traffic split, rate limits, inference-table logging, guardrails, and usage tracking. Callers invoke it by its three-level name over an OpenAI-compatible API and inherit all of that governance.

The point is define-once / inherit-everywhere: the model choice lives on the service (currently Claude Sonnet 4.6) and can be changed centrally without touching any lab or app code — no caller ever names the underlying model.

• Lab 1 creates and configures the model service (via the Unity Catalog REST API), sets up a dedicated service principal for the LLM hop, and rewires the agent to ChatDatabricks(model=, use_ai_gateway=True).
• Lab 2 routes its ResponsesAgent’s LLM through the same service.
• Lab 3 (the app) uses AsyncDatabricksOpenAI(use_ai_gateway=True) against the same service, and Lab 4 adds managed memory on top.

Governance defined once in Lab 1 is inherited by every subsequent lab. Usage is observable in system.ai_gateway.usage and the inference _payload table.

Model Context Protocol (MCP)

An open standard providing a uniform interface for AI agents to discover and call tools. Instead of custom integrations per data source, MCP gives you a standard list_tools / call_tool interface.

Databricks supports three MCP server types:

Agent Patterns

Lab 1 uses ChatAgent so you learn the fundamentals. Lab 2 upgrades to ResponsesAgent, which is the recommended pattern for new agents and the interface used by Agent Bricks and agents.deploy(). Lab 3 shows that ResponsesAgent is framework-agnostic — same wrapper, different implementation (OpenAI Agents SDK instead of LangGraph).

Agent on Databricks Apps

Lab 3 introduces the “agent as an application” pattern: instead of deploying to a serving endpoint, the agent runs inside a Databricks App with MLflow AgentServer. A custom React + Tailwind CSS frontend provides a Swiss-themed chat interface with real-time tool flow visualization, dark/light mode, and MLflow experiment links — all served from a single FastAPI process. The app connects to the same MCP servers as Labs 1–2, showing that MCP infrastructure is reusable across agent implementations. Because it runs on Apps (not a Model Serving endpoint), app-to-app egress works, so all three tools — trains, weather (Open-Meteo), and web search (Tavily) — run here, including the custom weather MCP that a served agent can’t reach. Its LLM is governed the same way as the other labs: every call routes through the travel_agent_llm Unity AI Gateway model service (via AsyncDatabricksOpenAI(use_ai_gateway=True)), never FMAPI directly.

Managed Memory (durable, governed, zero-ops)

Lab 4 gives the agent durable, user-scoped memory with Databricks Managed Memory (Beta) — a Unity Catalog MEMORY_STORE securable read/written purely over REST, backed by Databricks EntityStore. There’s no database, Vector Search, or embedding endpoint to provision. One store holds two complementary primitives, both partitioned by an opaque per-user scope:

scope is the isolation boundary — set in trusted server code from the verified signed-in user, never chosen by the model, and fail-closed when no identity resolves. Recall happens inside the agent’s turn (list_memories → get_memory); writing is in-agentic-loop, steered by a layered MEMORY_INSTRUCTIONS block. Because memory is threaded through session= (not the responses conversation= param), it’s independent of the model API — the app stays on the chat completions API, which cleanly handles MCP tool results.

Models-from-Code

Labs 1 and 2 use MLflow’s models-from-code pattern: write agent logic to a .py file, then log it with mlflow.pyfunc.log_model(python_model="agent.py"). This captures the code, dependencies, and configuration as a versioned MLflow model that can be deployed to a serving endpoint. (Lab 3 instead runs the agent as a Databricks App, so it doesn’t use agents.deploy().)

---

Multi-User Workshop Support

The labs are designed for multi-user workshops where many participants run concurrently in a shared workspace. All user-specific resources (schemas, functions, apps, endpoints) are automatically prefixed with the username to avoid collisions. See the notebooks for the naming convention.

---

Custom MCP Weather Server

The mcp-weather-server/ directory contains a FastMCP server that wraps the Open-Meteo API. It’s deployed as a Databricks App in Lab 2.

Features:

• Geocoding (geocode_city) and weather forecasts (daily and hourly) for any city, via latitude/longitude
• Keyless — Open-Meteo needs no API key or Databricks Secret for non-commercial/educational use, with a generous 10,000 calls/day tier (so no quota-tracking or caching machinery)
• Swiss-authority Alpine data — sources the MeteoSwiss ICON model (ICON-CH2, ~2 km) directly through Open-Meteo
• Graceful degradation — returns friendly errors on bad input or rate limits
• Attribution — credits “Weather data by Open-Meteo.com” (CC BY 4.0), as the license requires

---

Local Development (Optional)

For contributors who want to test notebooks locally using Databricks Connect:

1. Create a virtual environment: python3 -m venv .venv && source .venv/bin/activate
2. Install dependencies: pip install -r requirements-dev.txt
3. Copy .env.example to .env and fill in your secrets
4. Use scripts/run_cell.py to run notebook cells in a persistent local kernel

Agent App (Lab 3) local testing: Run the agent app locally with cd agent-app && uv run start-app. This builds the custom React frontend (first time only) and starts the server on port 8000. Open http://localhost:8000 in your browser to chat with the agent. Requires Node.js 18+ for the frontend build. Use --no-ui to skip the frontend and test via /invocations directly. For custom MCP weather server auth, set the secret scope and key names in agent-app/.env — see agent-app/README.md for details.

See CLAUDE.md for the full local testing procedure.

---

Common Issues

---

Resources

• Databricks MCP Documentation
• AI governance with Unity AI Gateway
• Foundation Model APIs
• Mosaic AI Agent Framework
• Unity Catalog AI Functions
• LangGraph Documentation
• OpenAI Agents SDK
• Databricks Apps Documentation
• FastMCP — Python MCP Server Framework
• Original Repository — genseral/ai_agent_transport

---

Agentic Coding with Claude Code

This repository includes Claude Code skills (.claude/skills/) — structured knowledge files that teach AI coding assistants how to build, deploy, and operate Databricks agents using the validated patterns from this workshop.

What are skills?

Skills are markdown files that provide Claude Code with domain-specific knowledge extracted from working implementations. When you use Claude Code in this repository, it automatically picks up these skills and can apply the patterns to your requests — creating agents, deploying MCP servers, configuring auth, and more.

Available skills

How to use

1. Install Claude Code — follow the setup guide
2. Open this repo in your terminal: cd travel-agent-lab && claude
3. Ask Claude to build things — it will use the skills automatically:
   • “Create a UC function that calls the OpenWeather API”
   • “Build a ResponsesAgent with MCP tools and deploy it”
   • “Set up a FastMCP server as a Databricks App”
   • “Help me test this notebook locally with Databricks Connect”

The skills are based entirely on validated, working code from this workshop — no hallucinated patterns.

---

License

This project is licensed under the Databricks License. See LICENSE for details.