Inside GitLab AI Gateway: How to Register Next-Gen LLMs with GitLab Duo
As software development platforms race to embed generative AI directly into the developer workflow, GitLab Duo has set itself apart with a decoupled, privacy-first architectural approach: the GitLab AI Gateway.
Rather than hardcoding language models directly into the main GitLab Rails monolith, GitLab routes all AI features (Code Suggestions, Duo Chat, Vulnerability Explanation, and Root Cause Analysis) through a specialized, high-performance Python/FastAPI service: gitlab-org/modelops/applied-ml/code-suggestions/ai-assist.
Recently, I contributed Merge Request !7251 to the GitLab AI Gateway (resolving Issue #3001).
This contribution registered Google’s next-generation Gemini 4 Argon model (gemini_4_argon_vertex) into the AI Gateway’s model registry, provisioning it with a 1,000,000-token context window, 64k output token capacity, and updated pricing indicators across unit primitives.
In this deep dive, I’ll walk through:
- The decoupled architecture of GitLab AI Gateway.
- How model configuration and lifecycles work (
models.ymlandunit_primitives.yml). - Step-by-step instructions for adding a new foundation model to GitLab Duo.
- How to pass all 24 AI Gateway CI jobs, container tests, and linters.
1. Architectural Overview: How GitLab Duo Talks to AI Models
Understanding where the AI Gateway sits in GitLab’s architecture is essential:
+-------------------------------------------------------------+
| Developer IDE / Web UI |
| (VS Code, JetBrains, GitLab Web Workspace) |
+-------------------------------------------------------------+
|
v
+-------------------------------------------------------------+
| GitLab Core (Rails) |
| (Policy, Permissions, Telemetry) |
+-------------------------------------------------------------+
|
v
+-------------------------------------------------------------+
| GitLab AI Gateway |
| (FastAPI / Litellm / Prompt Routing) |
| - models.yml |
| - unit_primitives.yml |
+-------------------------------------------------------------+
/ \
v v
+---------------------------+ +-----------------------------+
| Google Cloud Vertex AI | | Anthropic / AWS |
| (Gemini 4 Argon, 1.5) | | (Claude 3.7 / 3.5) |
+---------------------------+ +-----------------------------+
Why Decouple the AI Gateway?
- Model Agility: Adding, upgrading, or benchmarking models doesn't require restarting or deploying the monolithic Rails app.
- Data Privacy: The AI Gateway enforces strict zero-data-retention guarantees, stripping sensitive corporate telemetry before prompts hit third-party hyperscalers.
- Multi-Provider Fallback: If an upstream provider throttles requests or experiences latency spikes, the gateway can reroute prompts seamlessly.
2. The Model Lifecycle in GitLab AI Gateway
In GitLab AI Gateway, model definitions are primarily configuration-driven and validated by strict Pydantic schemas.
Every model supported by GitLab Duo is governed by two core configuration files:
-
ai_gateway/models/definitions/models.yml: Defines the hardware parameters, context limits, provider endpoints, and cost indicators. -
ai_gateway/models/definitions/unit_primitives.yml: Maps specific GitLab Duo features (such asduo_chat,code_suggestions,explain_vulnerability) to approved model engines.
3. Step-by-Step: Adding Gemini 4 Argon (gemini_4_argon_vertex)
Step A: Defining the Model in models.yml
Under the Vertex AI provider section, we registered the new model entry with its token limits and pricing metadata:
- name: gemini_4_argon_vertex
provider: vertex_ai
model_family: gemini
version: "4.0"
description: "Next-generation multi-modal model engineered for 1M+ context window reasoning and code generation."
engine: "gemini-4-argon"
endpoint: "projects/{project_id}/locations/{location}/publishers/google/models/gemini-4-argon"
parameters:
context_window: 1000000
max_output_tokens: 65536
temperature: 0.2
top_p: 0.95
capabilities:
streaming: true
structured_outputs: true
tool_calling: true
pricing:
input_cost_per_token: 0.00000125
output_cost_per_token: 0.00000500
Why These Parameters Matter:
-
context_window: 1000000: Allows GitLab Duo to ingest entire code repositories and large dependency trees into a single prompt during code reviews. -
max_output_tokens: 65536: Ensures large refactoring tasks don’t get cut off mid-function. -
pricing: Enables accurate enterprise seat billing and token consumption telemetry in GitLab’s usage dashboards.
Step B: Binding to Unit Primitives (unit_primitives.yml)
A model cannot be used in GitLab Duo until it is authorized for specific unit primitives (internal capabilities).
We mapped gemini_4_argon_vertex to Duo Chat and Advanced Code Generation:
duo_chat:
description: "Interactive chat assistant for code review, debugging, and general software queries."
compatible_models:
- claude_3_5_sonnet_vertex
- gemini_1_5_pro_vertex
- gemini_4_argon_vertex # <- Newly registered engine
default_model: gemini_4_argon_vertex
code_generation:
description: "Multi-line code completion and inline method implementation."
compatible_models:
- claude_3_5_sonnet_vertex
- gemini_4_argon_vertex
Step C: Adding Validation Specs & Regression Coverage
The AI Gateway enforces strict schema tests to ensure no malformed configuration reaches production.
We added automated test coverage in tests/models/test_model_registry.py:
import pytest
from ai_gateway.models.registry import ModelRegistry
def test_gemini_4_argon_registration():
registry = ModelRegistry.load_from_definitions()
model = registry.get("gemini_4_argon_vertex")
assert model is not None
assert model.provider == "vertex_ai"
assert model.parameters.context_window == 1_000_000
assert model.parameters.max_output_tokens == 65_536
assert model.capabilities.streaming is True
assert model.pricing.input_cost_per_token > 0
4. Conquering the 24 CI Jobs in GitLab CI/CD
One of the most rewarding aspects of contributing to GitLab AI Gateway is the sheer rigor of its verification pipeline. Merge Request !7251 triggered 24 automated CI jobs, including:
-
poetry-lock-check: Validates package integrity. -
ruff-format&ruff-lint: Ensures compliance with Python PEP standards. -
mypy-typecheck: Enforces strict static type safety across model payloads. -
yamllint: Confirms valid YAML structure across all definition files. -
pytest-unit: Runs thousands of unit tests across prompt compilers and streaming wrappers. -
docker-build-test: Compiles the actual container image deployed to production clusters.
All 24 CI jobs turned green on the first review cycle, allowing the maintainer team to swiftly review and merge the changes into production.
5. What This Unlocks for GitLab Developers
By integrating frontier models with massive context windows into GitLab Duo:
- Full-Repository Grounding: Developers can ask questions about cross-file architectural changes without truncating files.
- Deeper Vulnerability Remediation: When a SAST/DAST security pipeline flags an alert, Duo can analyze both the call site and the library definition to craft an accurate patch.
- Cost Optimization: Granular token metrics allow organizations to balance speed, reasoning depth, and budget seamlessly.
6. How You Can Contribute to GitLab AI Gateway
GitLab's AI Gateway is open source and actively welcoming community contributions:
-
Repository: Check out
gitlab-org/modelops/applied-ml/code-suggestions/ai-assist. -
Find Issues: Look for issues labeled
AI Gateway,model registry, orDuo Chat. - Earn Contributor Points: Every contribution earns points toward the GitLab Contributor Platform for swag, recognition, and hackathon awards.
Links & Reference
- Merge Request !7251: gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/merge_requests/7251
- Resolved Issue #3001: gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/work_items/3001
- GitLab Contributor Profile: contributors.gitlab.com/users/hrlpavan
- GitHub: @hrlpavan
Have questions about GitLab Duo, AI Gateway, or LLM model registries? Leave a comment below!
Top comments (0)