Skip to main content

Google Vertex AI Integration

HatiData runs natively on Google Cloud Platform and integrates with Vertex AI Agent Builder to provide persistent memory, governed SQL access, and chain-of-thought logging for Gemini-powered agents. The data plane deploys on Cloud Run or GKE, keeping your data within your GCP project while the control plane handles auth, billing, and policy distribution.

Prerequisites​

  • A GCP project with Vertex AI API enabled
  • gcloud CLI authenticated
  • Python 3.10+
  • A running HatiData proxy (local or GCP-deployed)
pip install hatidata google-cloud-aiplatform
export GOOGLE_CLOUD_PROJECT="your-project-id"
export HATIDATA_API_KEY="hd_live_your_api_key"
export HATIDATA_HOST="localhost"  # or your Cloud Run URL

Gemini Agent with Memory​

Build a Vertex AI agent that uses HatiData for persistent memory and context retrieval.

import os
import vertexai
from vertexai.generative_models import GenerativeModel, Part
from hatidata import HatiDataClient
from hatidata.memory import MemoryClient

vertexai.init(project=os.environ["GOOGLE_CLOUD_PROJECT"])

hd = HatiDataClient(
    host=os.environ["HATIDATA_HOST"],
    port=5439,
    api_key=os.environ["HATIDATA_API_KEY"],
)
memory = MemoryClient(hd)

AGENT_ID = "vertex-ai-agent"

def retrieve_context(query: str, top_k: int = 5) -> str:
    """Search HatiData memory for relevant context."""
    results = memory.search(
        agent_id=AGENT_ID,
        query=query,
        top_k=top_k,
        min_score=0.7,
    )
    if not results:
        return ""
    lines = [f"- {r.content}" for r in results]
    return "Relevant context:\n" + "\n".join(lines)

def ask_agent(query: str, session_id: str = "default") -> str:
    context = retrieve_context(query)

    model = GenerativeModel(
        "gemini-1.5-pro",
        system_instruction=(
            "You are a knowledgeable assistant. Use the provided context "
            "to give accurate, personalized responses.\n\n" + context
        ),
    )

    response = model.generate_content(query)
    answer = response.text

    # Persist the interaction
    memory.store(
        agent_id=AGENT_ID,
        content=f"Q: {query}\nA: {answer[:300]}",
        metadata={
            "session_id": session_id,
            "model": "gemini-1.5-pro",
            "type": "interaction",
        },
    )

    return answer

result = ask_agent("What were our top-performing products last quarter?")
print(result)

Function Calling with HatiData Tools​

Vertex AI supports function calling, allowing Gemini to invoke HatiData queries as tools.

from vertexai.generative_models import (
    GenerativeModel,
    FunctionDeclaration,
    Tool,
)

query_func = FunctionDeclaration(
    name="query_warehouse",
    description="Execute a SQL query against the data warehouse and return results.",
    parameters={
        "type": "object",
        "properties": {
            "sql": {
                "type": "string",
                "description": "The SQL query to execute",
            }
        },
        "required": ["sql"],
    },
)

search_memory_func = FunctionDeclaration(
    name="search_memory",
    description="Search agent memory for relevant past interactions and facts.",
    parameters={
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": "The search query",
            }
        },
        "required": ["query"],
    },
)

hatidata_tools = Tool(function_declarations=[query_func, search_memory_func])

model = GenerativeModel(
    "gemini-1.5-pro",
    tools=[hatidata_tools],
    system_instruction="You are a data analyst with access to a SQL warehouse.",
)

def handle_function_call(name: str, args: dict) -> str:
    if name == "query_warehouse":
        rows = hd.query(args["sql"])
        return str(rows[:20])
    elif name == "search_memory":
        results = memory.search(
            agent_id=AGENT_ID,
            query=args["query"],
            top_k=5,
        )
        return "\n".join(r.content for r in results) or "No results found."
    return "Unknown function"

chat = model.start_chat()
response = chat.send_message("What are the revenue trends for this quarter?")

# Handle function calls in the response
for candidate in response.candidates:
    for part in candidate.content.parts:
        if part.function_call:
            result = handle_function_call(
                part.function_call.name,
                dict(part.function_call.args),
            )
            response = chat.send_message(
                Part.from_function_response(
                    name=part.function_call.name,
                    response={"result": result},
                )
            )

GCP-Native Deployment​

Deploy HatiData alongside your Vertex AI agents on GCP using Terraform.

Cloud Run Deployment​

The HatiData proxy runs as a Cloud Run service in your GCP project:

# Deploy using HatiData's Terraform modules
cd terraform/gcp
terraform init
terraform plan -var-file="environments/production.tfvars"
terraform apply -var-file="environments/production.tfvars"

Key Terraform variables for GCP:

# environments/production.tfvars
project_id        = "your-project-id"
region            = "us-central1"
enable_cloud_tier = true
cloud_provider    = "gcp"

# Cloud Run settings
cloud_run_cpu    = "4"
cloud_run_memory = "8Gi"

# Cloud SQL for control plane
cloudsql_tier          = "db-custom-2-7680"
cloudsql_disk_size_gb  = 50

VPC Network Configuration​

Keep the HatiData proxy and your Vertex AI agents in the same VPC for minimal latency:

# Vertex AI and HatiData share the same VPC
resource "google_compute_network" "main" {
  name                    = "agent-network"
  auto_create_subnetworks = false
}

resource "google_compute_subnetwork" "agents" {
  name          = "agent-subnet"
  ip_cidr_range = "10.0.0.0/24"
  network       = google_compute_network.main.id
  region        = "us-central1"
}

Private Service Connect​

For production deployments, use Private Service Connect to keep all traffic on the GCP backbone:

Vertex AI Agent  -->  PSC Endpoint  -->  HatiData Proxy (Cloud Run)
                      (private IP)       (port 5439)

See PrivateLink & VPC for detailed PSC setup instructions.

Workload Identity Federation​

Authenticate the HatiData proxy to GCP services without static credentials:

# The proxy automatically uses Workload Identity when running on GKE/Cloud Run
# No manual credential management needed

# For local development, use application default credentials
import google.auth
credentials, project = google.auth.default()

HatiData's GCP federation provider validates Workload Identity tokens and maps them to HatiData organizations, eliminating the need for long-lived API keys in GCP-native deployments.

Stay in the loop

Product updates, engineering deep-dives, and agent-native insights. No spam.