Simplify AI agent orchestration with Lakebase Postgres

Simplify AI agent orchestration with Lakebase Postgres | Databricks Blog Skip to main content

Summary

Scale-Ready Task Queue on Postgres: A deep dive into the patterns that turn a pair of Lakebase tables into a durable, concurrent, crash-resilient queue for long-running agent tasks — no broker, cache, or scheduler required.

Fully Databricks-Native Architecture: A reference design that stitches Lakebase, Databricks Apps, Lakeflow Jobs, MLflow, and Unity Catalog Volumes into an end-to-end pipeline for agentic document parsing, with no external infrastructure to operate.

Real-Time Observability and Invariants: A deep dive into using Postgres LISTEN/NOTIFY triggers paired with Server-Sent Events (SSE) to build a low-latency operator dashboard that automatically tracks costs and tasks with zero overhead.

Introduction Traditionally, auditing is a tedious process that often requires detailed document review and information extraction. To accelerate this process, CLA (CliftonLarsonAllen LLP), a leading professional services firm with a growing global presence, worked with Databricks Forward Deployed Engineering team to build and productionize an agentic auditing solution. Jointly, we developed a document processing application that reduces extraction time from hours to minutes with no compromise in quality. The application is built entirely on Databricks, using Lakebase Postgres , Databricks Apps , Lakeflow Jobs , MLflow , and Unity Catalog Volumes . In this blog, we focus on one key component of that system, the Lakebase-powered orchestration layer. The orchestration layer is responsible for coordinating long-running tasks, managing retries, attributing cost, and providing real-time visibility. Lakebase and Databricks Apps, we eliminate the need for separate infrastructure for queueing, orchestration, and observability. Lakebase also makes this architecture practical at scale by separating storage from compute. Unlike traditional Postgres deployments, compute can scale with demand while storage remains durable and independent. Together, these capabilities make Lakebase a practical foundation for a simpler, scalable orchestration pattern for long-running agentic workloads on Databricks. Orchestration Challenges for Agentic Workloads Document parsing is a very common, high-volume agentic workload. Companies across industries need to convert large volumes of contracts, invoices, financial filings, and other documents into structured data. Running this at scale surfaces five distinct distributed-systems problems: Unpredictable per-task latency: A two-page invoice may be processed in seconds, while a two-hundred-page contract can take several minutes, making it difficult to predict how long any individual task will run. Rate-limit-aware throttling: LLM and vision model endpoints limit the number of requests and tokens they can process over a given period. Sending hundreds of tasks at once can overwhelm those limits, trigger throttling, and lead to repeated retries. The orchestrator must proactively limit in-flight work (by concurrent task count, by token budget, or both) rather than rely on reactive retry alone. Workload prioritization: Urgent submissions should not be delayed behind large bulk batches. Per-task priority ensures that higher-priority work (interactive submissions, premium-tier requests, operator-initiated reprocesses) is dispatched first. Cost attribution per task: Finance teams need to attribute spend to specific tasks, customers, and agents, broken down by AI token usage and compute consumption. Real-time progress visibility: Users uploading hundreds of documents need a live progress view.

Many organizations combine several specialized systems for orchestration and observability. Each system brings its own infrastructure, authentication, monitoring, and operational requirements, along with the work required to integrate them. For long-running, independent agentic tasks, that overhead is disproportionate to the actual scheduling complexity. The Databricks-native solution we developed meets all of the above requirements with Lakebase as the foundation. Solution Architecture

The full application stack consists exclusively of Databricks services: Web Application (Databricks Apps). A FastAPI-based user interface where users upload PDFs (stored in Unity Catalog Volumes) and submit parse requests. Requests are written directly to the Lakebase task table. Lakebase. An autoscaling Postgres database hosting the orchestrator's relational state across related tables: tasks (documents to be parsed, holding status, lease information, and the structured result), and task_attempts (one row per execution attempt, capturing the Databricks Job run ID, MLflow trace ID, and per-attempt cost metadata). Lakebase serves as the single source of truth for the orchestrator state. Orchestrator (Databricks Apps). A long-running worker daemon and operator dashboard. The daemon dequeues tasks from Lakebase, dispatches them to the AI Agents layer, and writes results back. The dashboard reads the same tables to surface real-time status. AI Agents (Lakeflow Jobs). Lakeflow Jobs execute the parsing work. Each Job reads a PDF from Unity Catalog Volumes, processes it through Intelligent Document Processing and vision/LLM calls, stores parsed output in Lakebase, and acknowledges the orchestrator via webhook. MLflow Tracing captures execution details such as model calls, token usage, latency, and cost metadata.

Data flows between components as follows. The Web App writes PDFs to Unity Catalog Volumes and parses requests to Lakebase. The Orchestrator dequeues tasks from Lakebase and dispatches Databricks Jobs to the AI Agents layer. The AI Agents process documents, write results back to Lakebase, and call back to the Orchestrator with status updates. Due to these built-in capabilities on Databricks, we did not need to rely on external message brokers (Kafka, Redis), separate schedulers (Airflow, Temporal), or dedicated caching layers. Task Queue Implementation The task queue is backed by two Postgres tables in Lakebase. The tasks table holds one row per logical unit of work, recording the task's current status, lease information, agent assignment, parent extraction, and final result. The task_attempts table holds one row per execution attempt, capturing the Databricks Job run ID, MLflow trace ID, and per-attempt cost metadata. The parent-child relationship supports retries (a single task may...