Building a Production LangGraph Agent on Databricks - NorthStar Brand Copilot

[SashankKotta]

Building a Production LangGraph Agent on Databricks: The NorthStar Brand Copilot

How we built and deployed an end-to-end CPG AI agent on Databricks using LangGraph, MCP servers, Lakebase memory, Genie, Vector Search, and MLflow and the best practices we learned along the way.

TL;DR

We built NorthStar Brand Copilot, an AI assistant for brand managers and field-sales reps at a fictional multi-category CPG company (Snacks, Beverages, Personal Care). It's a LangGraph custom agent that routes each question to the right Databricks-native capability:

• Genie (NatualLanguage→SQL) for the numbers — sell-in/sell-out, trade-promotion ROI, inventory, market share
• AI Search(RAG) for the documents — product specs, allergens, consumer reviews, brand guidelines, the promo playbook
• Lakebase (long-term memory) for the decisions, For ex: "remember we cut BOGO (Buy One Get One) at Walgreens," recalled across sessions. This provides a persistent, semantic context layer, ensuring the agent doesn't process queries in isolation but evolves its understanding based on past decisions and specific user history.
• LangGraph for agent orchestration. While simpler frameworks (like standard LangChain chains or monolithic agent patterns) are easier to start with, LangGraph was chosen for its superior ability to handle cyclical, complex workflows and granular state management. It allows the agent to reason, refine its tool-calling strategy iteratively, and maintain persistent state, which is essential for enterprise-grade reliability compared to more rigid, linear alternatives. 
• While LangGraph is ideal for cyclical and stateful workflows, other orchestration frameworks offer different strengths, Below are some other alternatives:
• LangChain (Core/Chains): Best for linear, straightforward pipelines where complex looping or persistent state management is not the primary requirement. It offers a vast ecosystem of integrations.
• LlamaIndex: Excels in data-heavy applications, providing highly optimized query engines and ingestion pipelines for RAG-centric workflows.
• DSPy: Focuses on production reliability through a declarative approach, programmatically optimizing prompts and module assembly to improve performance systematically rather than through manual tuning.

Note: The CPG organization and data used in this blog post are fictional and intended for demonstration purposes only.

Full Source Code Repository

The full source code, including synthetic data generation, setup scripts and deployment automation, is available in the source_code_repo.

The whole thing runs inside a single Databricks App, deployed with Declarative Automation Bundles (DABs), Agent powered by Claude Sonnet 4.5, and observed/evaluated with MLflow. This post walks through the architecture and the practices that make it production-grade.

Architecture Flow:

Why this architecture:

A few years ago, deploying an agent on Databricks meant logging a model to Unity Catalog and starting up a Model Serving endpoint. That still works, but the current recommended pattern is to run the agent inside a Databricks App as an MLflow ResponsesAgent served by AgentServer. We deliberately chose this newer path because it gives you:

• One artifact, one URL, one deploy: The UI, the API (/invocations), and the agent loop all live in the same app.
• Native identity & governance: The app runs as a service principal; every tool call is authorized through Unity Catalog and resource grants, no long-lived tokens baked into code.
• First-class observability: MLflow tracing is wired in with a single autolog line. Agent Evaluation is available too, but it is not automatic, it requires you to define an eval dataset, choose scorers, add the databricks-agents dependency, and set quality thresholds.

Everything governed by Unity Catalog · traces + eval in MLflow

The LangGraph custom agent

The heart of this project is agent_server/agent.py. It's intentionally small, the framework does the heavy lifting. Here's the shape of it:

from databricks_langchain import ChatDatabricks
from langchain.agents import create_agent
from mlflow.genai.agent_server import invoke, stream
mlflow.langchain.autolog()                       # full tracing, one line
@stream()
async def stream_handler(request: ResponsesAgentRequest):
    user_messages = to_chat_completions_input([i.model_dump() for i in request.input])
    messages = {"messages": [{"role": "system", "content": AGENT_INSTRUCTIONS}] + user_messages}
    tools = await _build_tools(sp_workspace_client)        # Genie + Vector Search (MCP) + time
    # ... open Lakebase memory store, add memory tools ...
    agent = create_agent(tools=tools, model=ChatDatabricks(endpoint=MODEL_ENDPOINT))
    async for event in process_agent_astream_events(agent.astream(...)):
        yield event

Best practice 1: Let the model route; describe tools well

We don't hand-code a supervisor graph with hard-wired branches. Instead we give a single create_agent a well-described tool set and a routing-oriented system prompt, and let Claude decide which tool(s) to call. The prompt is explicit about when to use each capability:

"Route quantitative questions to Genie. Route document-based or qualitative questions to Vector Search. For queries requiring both, call Genie first to retrieve numbers, then query AI Search(Formerly Vector Search) for guidance. Always cite sources. Never invent numbers."

This is the single highest-leverage piece of the agent. Tool descriptions and the system prompt are your routing logic. Invest there before you reach for custom graph code.

Gotcha: create_agent does not accept a prompt = argument the way older create_react_agent examples do. Prepend your instructions as a system message in the input state instead.

Best practice 2: Build your tools defensively (i.e. graceful degradation)

Make sure the _build_tools always returns something usable, even if a backend is down or a resource was deleted:

async def _build_tools(workspace_client):
    tools = [get_current_time]
    try:
        mcp_client = init_mcp_client(workspace_client)
        tools.extend(_stringify_tool(t) for t in await mcp_client.get_tools())
    except Exception:
        logger.warning("Failed to fetch MCP tools; continuing.", exc_info=True)
    return tools

This mattered in practice: When the Vector Search index was later deleted / Not available in one workspace, the agent kept working, it just lost the RAG path instead of crashing. An agent that degrades is better than an agent returns a 500 error.

Best practice 3: Normalize tool output for your LLM provider

A subtle but critical fix. Databricks-managed MCP tools (Genie especially) return structured content blocks that include an id field. The Claude endpoint rejects that field in tool_result content with:

tool_result.content.0.text.id: Extra inputs are not permitted

The fix is a thin wrapper that coerces every MCP tool's output to a plain string:

def _stringify_tool(t: StructuredTool) -> StructuredTool:
    async def _wrapped(**kwargs):
        out = await t.ainvoke(kwargs)
        return out if isinstance(out, str) else json.dumps(out, default=str)
    return StructuredTool(
        name=t.name, description=t.description, args_schema=t.args_schema, coroutine=_wrapped
    )

Lesson: when you bridge MCP tools to a specific model endpoint, validate the tool-result schema. Coercing to plain text is a safe default.

Best practice 4: Pin your model and centralize config

The model endpoint, Genie space, AI Search(Formerly Vector Search) catalog/schema, Lakebase instance, and embedding config are all environment variables with sensible defaults, set in databricks.yml / app.yaml. We pin a specific, capable model — databricks-claude-sonnet-4-5, rather than a floating alias. This keeps behavior reproducible across the multiple workspaces if deployed to.

MCP servers: the integration layer

Instead of writing bespoke API clients for Genie and AI Search(Formerly Vectory Search), the agent consumes them as MCP (Model Context Protocol) servers. Databricks exposes managed MCP endpoints for its services, and databricks-langchain gives you a client that turns them into LangChain tools:

from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient
def init_mcp_client(workspace_client):
    host = get_databricks_host_from_env()
    return DatabricksMultiServerMCPClient([
        DatabricksMCPServer(
            name="genie",
            url=f"{host}/api/2.0/mcp/genie/{GENIE_SPACE_ID}",
            workspace_client=workspace_client,
        ),
        DatabricksMCPServer(
            name="vector-search",
        url=f"{host}/api/2.0/mcp/vector-search/{VS_CATALOG}/{VS_SCHEMA}",
            workspace_client=workspace_client,
        ),
    ])
tools = await mcp_client.get_tools()

Why MCP is the right call

• Managed endpoints, no glue code. A Genie space becomes a tool at /api/2.0/mcp/genie/<space_id>; an AI Search index is reached at /api/2.0/mcp/ai-search/<catalog>/<schema>/<index_name>, and the schema-level path /api/2.0/mcp/ai-search/<catalog>/<schema> exposes every index in that schema as a tool. The legacy /api/2.0/mcp/vector-search/ prefix still works for backward compatibility. No SDK wiring per service.
• Identity flows through. The MCP client uses the app's WorkspaceClient, so calls execute as the app's service principal and are governed by the grants you declared.
• Interoperable. The same MCP pattern extends to custom MCP servers (Databricks Apps named mcp-*) and external/UC-connection MCP servers, declare the app as an app resource in databricks.yml and the bundle grants CAN_USE on deployment (requires CLI v0.298.0+).
  

Note: We selected MCP as our preferred standard because it offers a unified, extensible protocol for tool integration. While Genie APIs or standard UC functions could handle specific tasks, MCP provides a future-proof framework. We anticipate expanding the Copilot's capabilities with a broader array of tools and external services, and MCP’s standardized interface significantly simplifies this long-term extensibility.

There are three tool types worth knowing on the platform: 

• Unity Catalog function tools (governed SQL UDFs)
• agent code tools (defined inline, for low-latency REST calls — our get_current_time is one)
• MCP tools (the interoperable path we lean on here).

Agent memory: long-term recall with Lakebase

This is what turns a chatbot into a copilot. NorthStar remembers decisions ("we decided to cut BOGO(Buy One Get One) at Walgreens"), flags action items, and recalls them later, even in a brand-new session.

Databricks recognizes two kinds of memory:

Type	Use case	Backing	Key
Short-term	History within one session	AsyncCheckpointSaver	thread_id
Long-term	Facts that persist across sessions	AsyncDatabricksStore	user_id

NorthStar uses long-term memory on Lakebase (Databricks' managed Postgres) via AsyncDatabricksStore, with semantic search powered by the same databricks-gte-large-en embedding endpoint used elsewhere.

How it's wired

The store is opened as an async context manager and threaded into the agent through RunnableConfig:

async with store_cm as store:
    await store.setup()                  # idempotent; creates store tables on first use
    tools = tools + memory_tools()
    agent = create_agent(tools=tools, model=ChatDatabricks(endpoint=MODEL_ENDPOINT))
    config = {"configurable": {"user_id": user_id, "store": store}}
    async for event in process_agent_astream_events(agent.astream(input=messages, config=config, ...)):
        yield event

Memory is exposed to the model as three tools, returned from a factory function:

• get_user_memory
• save_user_memory
• delete_user_memory 

The factory pattern lets each tool close over nothing but RunnableConfig, from which it pulls the store and user_id. Memories are namespaced per user: ("user_memories", user_id.replace(".", "-")).

Memory best practices we followed

1. Extract user_id explicitly, return None on miss. get_user_id() derives identity from the trusted Databricks Apps OBO context (request.context.user_id) first, and falls back to custom_inputs.user_id only for direct API calls that carry no authenticated user. When both are present they must match — a custom_inputs.user_id that conflicts with the authenticated OBO identity is rejected — so a caller can never read or write another user's memories by supplying someone else's ID. It returns None rather than a default so the caller decides the fallback, never silently write one user's memory under another's namespace.
2. Pass the store via config, not as a parameter. Tools read config.get("configurable", {}).get("store"). This keeps tool signatures clean and LLM-friendly.
3. Validate the LLM's JSON before persisting. save_user_memory json.loads the payload and rejects non-objects, the model will occasionally hand you malformed JSON.
4. Check user_id and store separately, with distinct messages. Easier to debug "no user_id" vs. "store not configured."
5. setup() is idempotent — call it on startup. It creates the store tables on first use; safe to run every time.

Two deployment gotchas, both real:

• The value_from: "database" binding in app.yaml resolves to the Lakebase hostname, not the instance name. AsyncDatabricksStore wants the instance name, so we call resolve_lakebase_instance_name() (it lists DB instances and matches the rw/ro DNS) before constructing the store. Skip this and you get "Unable to resolve Lakebase provisioned instances."
• Direct /invocations calls must pass custom_inputs.user_id explicitly — the chat UI / OBO path supplies it automatically, but raw API calls don't. Because a raw caller could supply any value here, restrict direct /invocations access to trusted callers and treat the OBO context as the authoritative identity (any custom_inputs.user_id that conflicts with it is rejected). Without any identity the agent correctly refuses ("requires user identification").

Just like tool-building, memory degrades gracefully — if databricks-langchain[memory] isn't installed or Lakebase is unreachable, the agent drops the memory tools and keeps answering.

Lakebase: Deep Dive and Freshness

Beyond serving as a persistent key-value store, Lakebase acts as a durable, managed storage layer that allows agents to maintain state beyond short-lived session histories. By leveraging managed Postgres, it provides a consistent, reliable backend that ensures agent memory remains available even when compute clusters spin down or sessions expire, effectively turning stateless LLM calls into a continuous, learning assistant.

Handling Stale Data

As with any long-term memory system, stale or outdated information can degrade agent performance. To maintain relevance, implement the following strategies:

• Timestamps & Versioning: Always include a 'last_updated' timestamp with stored facts. When the agent retrieves memory, the logic should filter for freshness, potentially ignoring facts older than a defined threshold.
• Explicit Update Patterns: Rather than simply appending new memories, utilize the save_user_memory tool to perform 'upserts' (update-or-insert). Overwriting specific keys with the latest information ensures the agent always prioritizes the most current context.
• Periodic Pruning: Implement scheduled cleanup jobs that periodically scan the memory store for expired records or deprecated information, ensuring the Lakebase instance remains performant and focused on relevant data.
• Contextual Validation: When the agent retrieves a memory, prompt it to verify if the information is still applicable based on the current user request, adding an extra layer of human-in-the-loop or logical validation.

Databricks features, end to end

The demo is a tour of the platform's GenAI surface area:

Capability	Databricks feature	Role in the agent
Analytics (NL→SQL)	Genie space over 7 governed Delta tables	Quantitative questions: ROI, sell-through, inventory
Document insights (RAG)	Vector Search (Delta-sync index, databricks-gte-large-en)	Specs, allergens, reviews, playbook, briefs
Long-term memory	Lakebase (managed Postgres) + AsyncDatabricksStore	Persist & recall decisions across sessions
LLM	Foundation Model API — databricks-claude-sonnet-4-5	Reasoning + routing
Governance	Unity Catalog	Every table, index, and function authorized via grants
Hosting	Databricks Apps (FastAPI + custom SPA)	One app: Dashboard tab + Assistant tab
Packaging	Automation Bundles (DABs)	Declarative deploy + resource permissions
Observability	MLflow Tracing (autolog)	Per-request span tree: routing → tool → LLM
Quality	MLflow Agent Evaluation	Scored on Correctness / Relevance / Safety

A nice touch: the app serves a custom two-tab SPA (vanilla JS + Chart.js) from the same FastAPI process — a Dashboard tab backed by a /api/analytics SQL endpoint, and an Assistant tab that calls /invocations. Running the app backend-only (command: ["uv", "run", "start-server"]) means the FastAPI app serves both the UI and the agent — no separate Node/Next.js frontend to build.

Best Practice 5 — Production Safeguards: AI Gateway and Guardrails

When moving from prototype to production, ensuring stability and safety is non-negotiable.

• AI Gateway: We leverage the Model Serving AI Gateway to manage costs and reliability. This provides built-in rate limiting, caching for repeated prompts, and centralized API key management, ensuring our calls to the model are robust and compliant with corporate policies.
• Guardrails: For production deployments, we recommend wrapping your agent in a guardrail layer that inspects inputs and outputs before and after they reach the model. This layer is critical for establishing operational rules, such as screening for unauthorized PII disclosure, enforcing topic alignment, and detecting prompt-injection attempts. These guardrails can be implemented at the application layer or via dedicated safety models and should be treated as a first-class part of your agent's request/response path to ensure consistent reliability and security.

Best Practice 6 — Review Cycles and Continuous Assessment

How do you know it's ready for users? We treat the application as a software product.

• The Review Process: For stakeholder feedback we review the app's captured conversations in MLflow Traces, where reviewers can inspect real interactions for tone, accuracy, and correctness before the production deployment. Note that the Databricks Review App Chat UI is built around a Model Serving endpoint, so it does not directly exercise an agent deployed as a Databricks App, if you want the interactive Review App experience for staging, deploy the agent to a separate Model Serving test endpoint and point the Review App at that.
• Leveraging Assessments: Don't just rely on manual testing. We use MLflow Agent Evaluation to run our "eval set" (a suite of 10+ common CPG business questions). By setting up a recurring job that runs these evaluations against new agent versions, we can catch regressions in correctness or safety before they hit production. Treat these "Judge" scores as probabilistic quality signals — a useful regression check, but a complement to, not a replacement for, deterministic unit tests, security tests, and integration tests. LLM judges are themselves non-deterministic and depend on the datasets, scorers, and thresholds you define.

MLflow tracing & evaluation

mlflow.langchain.autolog() captures the entire agent run. The trick for clean traces: wrap each invocation in a parent span (@mlflow.trace(name=..., span_type="AGENT")), otherwise autolog emits fragmented one-span-per-call traces. With the parent span you get a single unified trace:

northstar_brand_copilot → LangGraph → ChatDatabricks → [Genie query OR Vector Search retrieve] → ChatDatabricks

For quality, mlflow.genai.evaluate(...) runs built-in judge scorers over a curated 10-question CPG eval set:

Safety 1.0 · RelevanceToQuery 0.9 · Correctness 0.8

Gotcha: the GenAI judge scorers require the databricks-agents package. Without it, scorers fail silently — assessments come back None, which looks like passing tests but isn't. Install databricks-agents before evaluating.

MLFlow Trace:

MLFlow Evaluations:

Deployment: Automation Bundles, step by step

Everything is declared in databricks.yml. The app, its run command, every environment variable, and — crucially — every resource grant the service principal needs:

resources:
  apps:
    agent_langgraph:
      name: "northstar-brand-copilot"
      source_code_path: ./
      config:
        command: ["uv", "run", "start-server"]
        env:
          - { name: MODEL_ENDPOINT, value: "databricks-claude-sonnet-4-5" }
          - { name: GENIE_SPACE_ID, value: "01f1..." }
          - { name: LAKEBASE_INSTANCE_NAME, value_from: "database" }
      resources:
        - name: llm          # CAN_QUERY on the Claude endpoint
        - name: embedding    # CAN_QUERY on the embedding endpoint
        - name: genie_space  # CAN_RUN on the Genie space
        - name: vector_index # SELECT on the VS index (uc_securable)
        - name: database     # CAN_CONNECT_AND_CREATE on Lakebase
        - name: warehouse    # CAN_USE on the SQL warehouse
        - name: experiment   # CAN_MANAGE on the MLflow experiment

Note that running the app (databricks bundle run) after deployment is not optional. deploy (databricks bundle deploy) uploads code and reconciles resources; bundle run is what actually restarts the app, Avoiding it would be testing stale code.

The grants, bundle can't express

While the bundle grants the SP access to endpoints like the Genie space, Vector Search index, and Lakebase instance, additional permissions are needed for execution. Specifically, since Genie runs SQL queries under the service principal's identity, the SP must be granted Unity Catalog table access and SQL Warehouse usage. Similarly, Lakebase requires configuring a dedicated Postgres role. Those are applied by deployment/grant_resources.py:

• A. USAGE on the catalog/schema + SELECT on the tables (one statement per SQL Statement API call, it only runs one at a time)
• B. CAN_USE on the SQL warehouse
• C. A Lakebase Postgres role for the SP + grants on the AsyncDatabricksStore memory tables (run as a serverless job, because it needs databricks_ai_bridge)

More gotchas from the trenches:

• Authentication: Deployed Databricks Apps require an OAuth token for API requests; Personal Access Tokens (PATs) will not work.

  • Quick fix: Retrieve your token using: databricks auth token ... | jq -r .access_token

• YAML Configuration Syntax: Under the sql_warehouse resource configuration, the correct property name is id:, not sql_warehouse_id: (a common typo in some documentation templates).

• Terraform State Loss Recovery: If your Terraform state is lost, running databricks bundle deploy will attempt to re-grant permissions on all resources, which fails on the Lakebase database grant.

  • Workaround: Because bundle deploy uploads the source code before it fails during grant reconciliation, you can complete the deployment by running databricks apps deploy <app> --source-code-path <bundle files path> to update the application code directly without triggering resource reconciliation.

Final App:

Figure 1(Dashbaord):

Figure 2(Agent):

Figure 3(Agent Memory: Saved to Lakebase):

Performance Considerations — Latency

Latency is a critical metric for a responsive Copilot experience. Since the agent invokes multiple tools and LLM endpoints sequentially, we prioritize streaming responses. By utilizing streaming responses and optimizing tool execution, we ensure that the initial token latency is kept to a minimum, providing a snappy, responsive interface.

Lessons learned

1. Prefer the App + ResponsesAgent pattern over log-to-UC + Model Serving for new agents, It unifies UI, API, identity, and observability.
2. Your system prompt and tool descriptions are your routing engine. Spend time there before writing graph code.
3. Consume Databricks services as MCP tools — managed endpoints, governed identity, zero glue code.
4. Normalize tool output to your LLM provider's expectations (the Genie id-field fix).
5. Make every external dependency degrade gracefully — tools and memory both.
6. Treat memory as a first-class, per-user, namespaced store with explicit user_id handling and JSON validation.
7. Declare all grants in the bundle, and know which ones it can't express (UC tables, warehouse, Lakebase role).
8. Wrap invocations in a parent MLflow span and install databricks-agents for real evaluation scores.

The result is a governed, observable, reproducible agent that answers real CPG business questions — and a blueprint you can lift for your own domain.