Generative AI Google Brain

search-gemini-enterprise-gemini-enterprise-prompt-management

4 min read

Gemini Enterprise custom agent with prompt management

Source notebook

Repo path: search/gemini-enterprise/gemini_enterprise_prompt_management.ipynb · Open on GitHub · advanced

Builds a Gemini Enterprise ADK SQL agent using Vertex AI Prompt Management and schema file context.

Summary

This notebook teaches how to connect Vertex AI Prompt Management, ADK, Agent Engine, and Gemini Enterprise for a custom SQL-generation agent. It creates a managed prompt, retrieves the prompt by ID in a before-agent callback, appends uploaded schema DDL from a text file, and tests the agent locally with google.genai Content and Part objects. It then deploys the ADK app to Agent Engine with pinned requirements and registers it in Gemini Enterprise using Discovery Engine v1alpha REST calls with server-side OAuth.

Key code patterns

Initialize Vertex AI and load schema

vertexai.init(
    project=PROJECT_ID,
    location=LOCATION,
    staging_bucket=STAGING_BUCKET,
)
try:
    with open(SCHEMA) as f_handle:
        schema = f_handle.read()
except FileNotFoundError:
    schema = ''

Sets project, region, staging bucket, and local schema text used as file context.

Create managed prompt version

sql_query_gen_prompt_obj = Prompt(
    prompt_name='SQL Generator',
    prompt_data=sql_query_gen_prompt,
    model_name='gemini-2.0-flash-001',
)
agent_prompt = prompts.create_version(prompt=sql_query_gen_prompt_obj)
print('Prompt id = ', agent_prompt.prompt_id)

Stores the SQL-generation instruction in Vertex AI Prompt Management and returns a prompt ID.

Fetch prompt and append uploaded file

def update_instructions_add_schema(callback_context):
    sql_query_gen_prompt = prompts.get(prompt_id=PID).prompt_data
    callback_context._invocation_context.agent.instruction = sql_query_gen_prompt
    for part in callback_context.user_content.parts:
        if hasattr(part, 'inline_data') and getattr(part.inline_data, 'mime_type', ''):
            uploaded_file_content = part.inline_data.data.decode('utf-8')
            callback_context._invocation_context.agent.instruction += uploaded_file_content

Refreshes agent instructions from the managed prompt and adds schema DDL from an inline file part.

Configure ADK agent

root_agent = Agent(
    name='SQL_Generation_Agent',
    model='gemini-2.0-flash',
    description='Agent to convert natural language query to SQL based on a given schema DDL.',
    before_agent_callback=update_instructions_add_schema,
    generate_content_config=types.GenerateContentConfig(temperature=0.01),
)

Uses an ADK before-agent callback and low temperature for SQL generation.

Test with file artifact

app = reasoning_engines.AdkApp(agent=root_agent, enable_tracing=True)
session = app.create_session(user_id='u_123')
file_artifact = types.Part.from_bytes(mime_type='text/plain', data=schema.encode('utf-8'))
contents = types.Content(role='user', parts=[types.Part.from_text(text=query), file_artifact])
for event in app.stream_query(user_id='u_123', session_id=session.id, message=contents.model_dump()):
    print(event['content']['parts'][0]['text'])

Exercises the agent locally with a natural language query plus schema file content.

Deploy to Agent Engine

remote_app = agent_engines.create(
    display_name='SQL generator',
    agent_engine=app,
    requirements=[
        'google-adk (==1.5.0)',
        'google-genai (==1.24.0)',
        'pydantic (==2.11.7)',
        'google-cloud-aiplatform (==1.101.0)',
    ],
)
remote_app.resource_name

Publishes the ADK app to Agent Engine with explicit runtime dependencies.

Register in Gemini Enterprise

curl -X POST .../authorizations?authorizationId=sqlgen2 -d '{serverSideOauth2: {...}}'
curl -X POST .../assistants/default_assistant/agents -d '{
  displayName: "DDL SQL Generator",
  adk_agent_definition: {
    provisioned_reasoning_engine: {reasoning_engine: "projects/.../reasoningEngines/..."},
    authorizations: ["projects/.../authorizations/sqlgen2"]
  }
}'

Connects the deployed reasoning engine to a Gemini Enterprise assistant through Discovery Engine endpoints.

Models & APIs used

  • Models: gemini-2.0-flash-001, gemini-2.0-flash
  • APIs / services: Vertex AI, Vertex AI Prompt Management, Agent Engine, Gemini Enterprise, Discovery Engine API, Cloud Storage
  • SDKs / libraries: vertexai, google-adk, google-genai, google-cloud-aiplatform, pydantic

When to use this

Use this when a Gemini Enterprise custom agent needs centrally managed prompts plus uploaded file context before Agent Engine deployment.

Gotchas & caveats

  • Notebook uses placeholders for PROJECT_ID, PROJECT_NUMBER, LOCATION, STAGING_BUCKET, SCHEMA, PID, OAuth client fields, Gemini Enterprise engine ID, and reasoning engine ID.
  • SCHEMA must point to a readable UTF-8 text file; FileNotFoundError leaves schema as an empty string.
  • The callback decodes inline file data as UTF-8 text and appends it directly to the prompt instruction.
  • Prompt ID must be copied from prompts.list() or agent_prompt.prompt_id.
  • Agent Engine deployment pins google-adk 1.5.0, google-genai 1.24.0, pydantic 2.11.7, and google-cloud-aiplatform 1.101.0.
  • Gemini Enterprise registration uses gcloud auth print-access-token, X-Goog-User-Project, server-side OAuth client credentials, and discoveryengine.googleapis.com v1alpha endpoints.
  • The registration payload references a provisioned reasoning engine in locations/us-central1 while vertexai.init uses LOCATION.

Best practices

  • Create the managed prompt once and reuse its prompt_id rather than hard-coding instructions in the deployed agent.
  • Fetch the prompt in before_agent_callback so prompt text is centrally managed.
  • Pass schema DDL as a text/plain file part instead of assuming schema context.
  • Use low temperature for SQL generation.
  • Enable tracing in AdkApp for local testing.
  • Pin dependencies for Agent Engine deployment.

Graph View

Backlinks