Prompt Engineering for Structured Output with Google AI Studio and Vertex AI

Background

In my previous post, GitLab Code Reviews with Gemini Pro 1.5, I built an automated code reviewer using Google Cloud Functions and Vertex AI.

While that prototype worked remarkably well, I originally relied on standard markdown prompt text and parsed the LLM output using brittle regular expressions to extract JSON blocks:

# The old fragile approach: regex stripping
match = re.search(r"```json\n([\s\S]+?)```", response.text)
if match:
    data = json.loads(match.group(1))

Relying on regex to parse AI output in production is risky. If the model includes preamble commentary or forgets markdown fences, downstream API calls crash.

To make the system truly production-ready, I refactored the pipeline to use Google AI Studio and Vertex AI’s native Structured Output (JSON Mode with OpenAPI schemas).

Prototyping in Google AI Studio

Google AI Studio is an exceptional sandbox for engineering and benchmarking prompts before writing cloud deployment code.

In AI Studio’s Run Settings:

1. Select the model (e.g. gemini-1.5-pro or gemini-1.5-flash).
2. Under Tools, activate Structured Output / JSON Mode.
3. Use the schema editor to define an OpenAPI-compliant specification for the exact payload required.

Once verified in AI Studio, you can translate the schema directly into Python code using the vertexai SDK.

Defining the OpenAPI Response Schema

Rather than asking the LLM in plain prose to “please return JSON”, we enforce an exact JSON Schema at the decoding layer:

output_response_schema = {
    "type": "object",
    "properties": {
        "responses": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "new_line": { "type": "integer" },
                    "old_line": { "type": "integer" },
                    "new_file_path": { "type": "string" },
                    "old_file_path": { "type": "string" },
                    "comment": { "type": "string" },
                    "severity": {
                        "type": "string",
                        "enum": ["minor", "moderate", "major", "critical"]
                    }
                },
                "required": [
                    "new_line",
                    "old_line",
                    "new_file_path",
                    "old_file_path",
                    "comment",
                    "severity"
                ]
            }
        }
    },
    "required": ["responses"]
}

Configuring GenerationConfig in Vertex AI

We pass the schema directly into GenerationConfig with response_mime_type="application/json":

import json
from vertexai.generative_models import GenerativeModel, GenerationConfig

generation_config = GenerationConfig(
    response_mime_type="application/json",
    response_schema=output_response_schema,
    temperature=0.2,
    candidate_count=1
)

model = GenerativeModel(
    model_name="gemini-1.5-pro-001",
    generation_config=generation_config,
    system_instruction=system_instructions
)

response = model.generate_content(prompt_content)

# Zero regex required - output is guaranteed valid JSON matching the schema
parsed_result = json.loads(response.text)

Structuring System Instructions as a List

Another key realization: while the Vertex AI SDK accepts a single large text string for system instructions, passing a structured list of discrete rule blocks produces noticeably more reliable adherence:

system_instructions = [
    "You are a helpful code reviewer. Your mission is to review the code changes for a ROS2 package and provide feedback based on the changes.",
    
    """Diff Header Specification:
The diff header follows the standard Git unified diff format:
@@ -[start line],[number of lines] +[start line],[number of lines] @@
Lines starting with '-' were removed.
Lines starting with '+' were added.""",

    """Review Guidelines:
1. Correctness: Verify memory safety, logic flow, and edge cases.
2. ROS2 Idioms: Ensure proper topic namespacing, node lifecycles, and parameters.
3. Security: Prevent command injection and unvalidated inputs.""",

    """Coordinate Instructions:
- For an added line: set new_line to the line number, and old_line to -1.
- For a removed line: set new_line to -1, and old_line to the line number.
- For an unchanged line: set both new_line and old_line to the same line number."""
]

Takeaways

1. Deterministic Contracts: With native schema enforcement, downstream microservices can safely deserialize the LLM output without try/catch fallback hacks.
2. Speed & Efficiency: Constrained decoding prevents token drift and conversational filler (“Sure, here is your JSON:”), reducing latency and token costs.
3. Reproducibility: Designing schemas in Google AI Studio before deploying to Google Cloud Functions provides a fast feedback loop for prompt tuning.