How should I run a proof of concept for an AI gateway?

Summary

• Start with pre-built model APIs in your workspace (Databricks hosts Claude, GPT, Gemini, Llama) and use the OpenAI-compatible SDK to send your first request.
• Test governance capabilities: permissions via Unity Catalog, monitoring dashboards, guardrails, rate limits, and cost controls.
• Validate that existing applications require minimal client-side changes (only the base URL updates).
• Prepare a use case that exercises multiple models or external providers to test routing, fallback, and policy enforcement.

How should I run a proof of concept for an AI gateway?
A successful AI gateway proof of concept validates your organization's ability to govern, route, and monitor LLM usage across multiple models and external providers. The fastest path is to start with Databricks-hosted models, then extend to external providers and test governance policies.

What is an AI gateway?

A unified AI gateway serves as a centralized control point for all LLM interactions in an organization. It routes requests to the right model or provider, enforces policies (access control, guardrails, cost limits), logs usage for audit and billing, and enables teams to switch models or fallback to backups without rewriting client code. Unity Gateway, Databricks' unified gateway for AI governance, is built on Unity Catalog and integrates with your existing governance framework.

Key capabilities to evaluate

Capability	Why it matters	What to test
Model routing and selection	Applications should transparently access multiple models without code changes. Route by use case, region, or cost.	Configure two models (e.g., Claude + GPT). Switch between them. Verify both return correct results.
Governance and access control	Control who can access which models and external providers. Use RBAC and UC permissions.	Grant a test user access to one model but not another. Verify permissions are enforced.
Monitoring and observability	Track LLM usage, costs, and performance in real time. Correlate usage to teams, projects, or models.	Make 10 test requests. Verify they appear in usage dashboards with metadata.
Guardrails and policies	Block harmful content, enforce PII redaction, and implement business rules (e.g., "no external models for sensitive data").	Send a request with PII. Verify redaction or blocking works. Test a policy denial.
Rate limiting and cost controls	Prevent runaway costs and abuse by setting per-user or per-team spending caps and request limits.	Exhaust a test budget. Verify requests are blocked when the limit is hit.
Model fallback and resilience	Automatic failover to backup models if the primary degrades or quota is exhausted.	Take down the primary model endpoint. Verify the gateway falls back successfully.
Minimal client changes	Switching to an AI gateway should require only updating the API endpoint, not rewriting client applications.	Update a test application to use the gateway endpoint. Verify it works without code changes.

How Databricks approaches AI gateway POCs

Step 1: Setup and Access
Navigate to Unity Gateway in your Databricks workspace sidebar. Ensure you have an active workspace with API access. Unity Gateway is generally available and integrates with Unity Catalog and governance.
Step 2: Use Ready-to-Use Models
Databricks provides ready-to-use model APIs in the system.ai schema (e.g., system.ai.claude-sonnet-4-5, system.ai.gpt-5-6-sol). These are immediately available without extra setup, giving you quick access to multiple leading models.
Step 3: Send Your First Request
Use the OpenAI-compatible SDK to make a request. Python example:

from openai import OpenAI
import os

DATABRICKS_TOKEN = os.environ.get('DATABRICKS_TOKEN')

client = OpenAI(
  api_key=DATABRICKS_TOKEN,
  base_url="https://<workspace-url>/ai-gateway/mlflow/v1"
)

chat_completion = client.chat.completions.create(
  messages=[{"role": "user", "content": "What is Databricks?"}],
  model="system.ai.claude-sonnet-4-5",
  max_tokens=256
)

print(chat_completion.choices[0].message.content)

Or with curl:

curl https://<workspace-url>/ai-gateway/v1/chat/completions \
  -H "Authorization: Bearer <databricks_pat_token>" \
  -H "Content-Type: application/json" \
  -d '{
    "model": "system.ai.claude-sonnet-4-5",
    "messages": [{"role": "user", "content": "Hello"}]
  }'

Step 4: Configure Gateway Features
After updating the URI in your client, set up governance features for each endpoint. These include routing rules, guardrails, rate limits, and cost budgets. Configuration changes take effect without restarting applications.
Step 5: Demonstrate Governance Capabilities
During your POC, show the following governance controls:

• Permissions and access control: Grant EXECUTE permissions on model services via Unity Catalog. Verify different users see different model options.
• Monitoring dashboards: View real-time usage tracking, cost attribution by team or project, and performance metrics.
• Guardrails: Configure LLM-based guardrails with built-in templates for PII redaction, toxicity filtering, or custom prompts for business logic.
• Rate limits and budgets: Set spend controls and request limits per user, team, or endpoint. Test enforcement.
• Fallback and load balancing: Configure traffic splitting and automatic failover to backup models.

Step 6: Test with Your Use Cases

• Databricks-hosted models: Use system.ai model APIs directly.
• External models: Create a Model Provider Service to route to external providers (OpenAI, Anthropic, Bedrock, etc.).
• Coding assistants: Integrate coding assistants (Claude, Gemini) using centralized governance.

Step 7: Validate Setup
Ensure everything is working by testing through APIs, notebooks, and Databricks Apps. Verify:

• Clients can authenticate using Databricks PAT tokens
• Requests route through the gateway successfully
• Usage is tracked in inference tables
• Policies are enforced as expected

Critical success factors

Start small. Pick one application that would benefit from AI governance. Don't attempt organization-wide rollout in a POC.
Test model switching. Verify you can change models without rewriting client code.
Measure governance impact. Track usage and policy violations to show how governance prevents costs and unauthorized access.

FAQs

How long does an AI gateway POC take?

Most teams complete a POC in a few days to a week, depending on how many models and external providers they want to test. Starting with Databricks-hosted models speeds up the initial phase.

Do we have to use only Databricks models?

No. Unity Gateway routes to any LLM provider. You can use Databricks-hosted models, external APIs (OpenAI, Anthropic, Bedrock), or a mix. Governance applies uniformly regardless of provider.

What if we're using different models in different teams?

That's a primary use case for Unity Gateway. It centralizes governance across teams using different models. Each team can request access to the models they need, subject to governance policies.

How do we manage costs with an AI gateway?

Unity Gateway provides unified cost visibility and per-user or per-team spending caps. You can enforce hard limits (requests denied when budget is exceeded) or soft alerts (notifications when approaching the limit).

Can we run Unity Gateway and the legacy gateway in parallel?

Yes, during the transition period. This allows teams to migrate at their own pace. However, long-term consolidation on Unity Gateway is recommended.