Criar um app de chat de IA híbrida no GKE com Gemma e Gemini

1. Introdução

Visão geral

Neste laboratório, você vai criar e implantar um aplicativo de chat completo com tecnologia de IA no Google Kubernetes Engine (GKE). Esse aplicativo "híbrido" demonstra um padrão arquitetônico poderoso: a capacidade de alternar sem problemas entre um modelo aberto auto-hospedado (Gemma 3 12B) executado diretamente no cluster e um serviço de IA gerenciado (Gemini 2.5 Flash via Vertex AI).

                                   +----------------------+
                                   |   User (Web Browser) |
                                   +-----------+----------+
                                               |
                                               v
+----------------------------------------------+---------------------------------------------+
| Google Cloud Platform                        |                                             |
|   |                                  +-------+-------+                                     |
|   |                                  | Load Balancer |                                     |
|   |                                  +-------+-------+                                     |
|   |                                          v                                             |
|   +------------------------------------------+-----------------------------------------+   |
|   |  Google Kubernetes Engine (GKE)          |                                         |   |
|   |                                          v                                         |   |
|   |                              +-----------+-----------+                             |   |
|   |                              |    Gradio Chat App    |                             |   |
|   |                              +--+-----------------+--+                             |   |
|   |                                 |                 |                                |   |
|   |                   (Self-hosted) |                 | (Managed via SDK)              |   |
|   |                                 v                 |                                |   |
|   |                  +--------------+---+             |                                |   |
|   |                  | Gemma 3 Model    |             |                                |   |
|   |                  | (GPU Node)       |             |                                |   |
|   |                  +------------------+             |                                |   |
|   +---------------------------------------------------|--------------------------------+   |
|                                                       |                                    |
|                                                       v                                    |
|                                            +----------+-----------+                        |
|                                            | Vertex AI (Gemini)   |                        |
|                                            +----------------------+                        |
|                                                       | (Save History)                     |
|                                                       v                                    |
|                                            +----------+-----------+                        |
|                                            | Firestore Database   |                        |
|                                            +----------------------+                        |
+--------------------------------------------------------------------------------------------+

Você vai usar o Terraform para provisionar a infraestrutura, incluindo um cluster do Autopilot do GKE e um banco de dados do Firestore para manter o histórico de sessões de chat. Em seguida, você vai concluir o código do aplicativo Python para processar conversas de várias rodadas, interagir com os dois modelos de IA e implantar o aplicativo final usando o Cloud Build e o Skaffold.

O que você vai aprender

  • Provisionar a infraestrutura do GKE e do Firestore usando o Terraform.
  • Implante um modelo de linguagem grande (Gemma) no Autopilot do GKE usando manifestos do Kubernetes.
  • Implemente uma interface de chat do Gradio em Python que possa alternar entre diferentes back-ends de IA.
  • Use o Firestore para armazenar e recuperar o histórico de sessões de chat.
  • Configure a Identidade da carga de trabalho para conceder acesso seguro às cargas de trabalho do GKE aos serviços do Google Cloud (Vertex AI, Firestore).

Pré-requisitos

  • Ter um projeto do Google Cloud com o faturamento ativado.
  • Conhecimento básico de Python, Kubernetes e ferramentas de linha de comando padrão.
  • Um token do Hugging Face com acesso aos modelos da Gemma.

2. Configuração do projeto

  1. Se você ainda não tiver uma Conta do Google, crie uma.
    • Use uma conta pessoal em vez de uma conta escolar ou de trabalho. As contas escolares e de trabalho podem ter restrições que impedem a ativação das APIs necessárias para este laboratório.
  2. Faça login no Console do Google Cloud.
  3. Ative o faturamento no console do Cloud.
    • A conclusão deste laboratório custa menos de US $1 em recursos do Cloud.
    • Siga as etapas no final deste laboratório para excluir recursos e evitar mais cobranças.
    • Novos usuários podem aproveitar o teste sem custos financeiros de US$300.
  4. Crie um projeto ou reutilize um projeto existente.

Abrir editor do Cloud Shell

  1. Clique neste link para navegar diretamente até o editor do Cloud Shell.
  2. Se for preciso autorizar em algum momento, clique em Autorizar para continuar. Clique para autorizar o Cloud Shell
  3. Se o terminal não aparecer na parte de baixo da tela, abra-o:
    • Clique em Visualizar.
    • Clique em TerminalAbrir um novo terminal no editor do Cloud Shell.
  4. No terminal, defina o projeto com este comando:
    • Formato:
      gcloud config set project [PROJECT_ID]
      
    • Exemplo:
      gcloud config set project lab-project-id-example
      
    • Se você não se lembrar do ID do projeto:
      • Para listar todos os IDs de projeto, use:
        gcloud projects list | awk '/PROJECT_ID/{print $2}'
        
      Definir o ID do projeto no terminal do Editor do Cloud Shell
  5. Você vai receber esta mensagem:
    Updated property [core/project].
    
    Se você vir um WARNING e for perguntado Do you want to continue (Y/n)?, provavelmente inseriu o ID do projeto incorretamente. Pressione n, Enter e tente executar o comando gcloud config set project novamente.

Clone o repositório

No terminal do Cloud Shell, clone o repositório do projeto e navegue até o diretório dele:

git clone https://github.com/GoogleCloudPlatform/devrel-demos.git
cd devrel-demos/containers/gradio-chat-gke

Reserve um momento para explorar a estrutura do projeto:

gradio-chat-gke/
├── app/
   ├── app.py                # Main application logic (you will edit this)
   ├── requirements.txt      # Python dependencies
   └── themes.py             # UI theming
├── deploy/
   ├── chat-deploy.yaml      # Kubernetes deployment for the chat app
   ├── Dockerfile            # Container definition for the chat app
   └── gemma3-12b-deploy.yaml# Kubernetes deployment for Gemma model
├── infra/
   └── main.tf               # Terraform infrastructure definition
└── skaffold.yaml             # Skaffold configuration for building/deploying

Definir variáveis de ambiente

Configure variáveis de ambiente para o ID do projeto e o número do projeto. Eles serão usados pelo Terraform e por comandos subsequentes.

export GOOGLE_CLOUD_PROJECT=$(gcloud config get-value project)
export PROJECT_NUMBER=$(gcloud projects describe $GOOGLE_CLOUD_PROJECT --format="value(projectNumber)")
export REGION=us-central1

Ativar a API Resource Manager

O Terraform exige que a API Resource Manager esteja ativada para gerenciar os recursos do projeto. Portanto, precisamos ativá-la primeiro. Depois, vamos implantar o aplicativo de chat com o Skaffold, que usa o Cloud Build para criar a imagem do contêiner. Vamos ativar a API Storage e criar o bucket necessário para o Cloud Build agora. Vamos usar o próprio Terraform para ativar o restante das APIs necessárias para este projeto.

gcloud services enable cloudresourcemanager.googleapis.com storage-api.googleapis.com

Criar um bucket temporário do Cloud Build

O Skaffold usa o Google Cloud Build, que exige um bucket do Cloud Storage para preparar o código-fonte.

Crie agora para garantir que ele exista:

gcloud storage buckets create gs://${GOOGLE_CLOUD_PROJECT}_cloudbuild

Se você receber um erro informando que o bucket já existe, ignore-o.

3. Provisionar infraestrutura com o Terraform

Vamos usar o Terraform para configurar os recursos necessários do Google Cloud. Isso garante um ambiente consistente e reproduzível.

  1. Navegue até o diretório de infraestrutura:
    cd infra
    

Esse arquivo define as APIs adicionais necessárias para este projeto: cloudbuild, artifactregistry, container (gke), firestore e aiplatform (vertexai). Confira no arquivo ou abaixo como as APIs são ativadas pelo Terraform:

resource "google_project_service" "cloudbuild" {
  service            = "cloudbuild.googleapis.com"
  disable_on_destroy = false
  project            = var.project_id
}

resource "google_project_service" "artifactregistry" {
  service            = "artifactregistry.googleapis.com"
  disable_on_destroy = false
  project            = var.project_id
}

resource "google_project_service" "container" {
  service            = "container.googleapis.com"
  disable_on_destroy = false
  project            = var.project_id
}

resource "google_project_service" "firestore" {
  service            = "firestore.googleapis.com"
  disable_on_destroy = false
  project            = var.project_id
}

resource "google_project_service" "vertexai" {
  service            = "aiplatform.googleapis.com"
  disable_on_destroy = false
  project            = var.project_id
}

Definir o cluster do GKE

Abra infra/main.tf no editor. Você vai encontrar vários comentários # TODO. É possível abrir manualmente ou usar este comando para abrir o arquivo no editor:

cloudshell edit main.tf

Primeiro, precisamos definir nosso cluster do Kubernetes. Vamos usar o GKE Autopilot, que é ideal para cargas de trabalho de IA porque gerencia os nós automaticamente.

Encontre # TODO: Create a GKE Autopilot Cluster e adicione o bloco abaixo dele:

# Create a GKE Autopilot Cluster
resource "google_container_cluster" "primary" {
  name     = var.cluster_name
  location = var.region
  project  = var.project_id

  # Enable Autopilot mode
  enable_autopilot = true

  deletion_protection = false

  # Networking
  network    = "default"
  subnetwork = "projects/${var.project_id}/regions/${var.region}/subnetworks/default"

  # Timeout for cluster creation
  timeouts {
    create = "30m"
    update = "30m"
  }

  depends_on = [google_project_service.container]
}

Observe enable_autopilot = true. Essa única linha evita que gerenciemos pools de nós, escalonamento automático e bin-packing das nossas cargas de trabalho de GPU.

Definir o banco de dados do Firestore

Em seguida, precisamos de um lugar para armazenar nosso histórico de chat. O Firestore é um banco de dados NoSQL sem servidor que atende perfeitamente a essa necessidade.

Encontre # TODO: Create a Firestore Database e adicione:

resource "google_firestore_database" "database" {
  project     = var.project_id
  name        = "chat-app-db"
  location_id = "nam5"
  type        = "FIRESTORE_NATIVE"

  depends_on = [google_project_service.firestore]
}

Depois de adicionar o recurso de banco de dados, encontre # TODO: Create an initial Firestore Document e adicione o seguinte bloco. Esse recurso cria um documento de marcador de posição inicial na nossa coleção, o que é útil para inicializar a estrutura do banco de dados.

resource "google_firestore_document" "initial_document" {
  project     = var.project_id
  collection  = "chat_sessions"
  document_id = "initialize"
  fields = <<EOF
  EOF

  depends_on = [google_firestore_database.database]
}

Definir a Identidade da carga de trabalho

Por fim, precisamos configurar a segurança. Queremos que nossos pods do Kubernetes possam acessar a Vertex AI e o Firestore sem precisar gerenciar segredos ou chaves de API. Fazemos isso com a Identidade da carga de trabalho.

Vamos conceder os papéis do IAM necessários à conta de serviço do Kubernetes (KSA, na sigla em inglês) que nosso app vai usar.

Observação:a conta de serviço do Kubernetes (gradio-chat-ksa) referenciada nessas vinculações ainda não existe. Ele será criado mais tarde, quando implantarmos o aplicativo no cluster. É perfeitamente aceitável (e uma prática comum) pré-provisionar essas vinculações do IAM.

Encontre # TODO: Configure Workload Identity IAM bindings e adicione:

locals {
  ksa_principal = "principal://iam.googleapis.com/projects/${var.project_number}/locations/global/workloadIdentityPools/${var.project_id}.svc.id.goog/subject/ns/default/sa/gradio-chat-ksa"
}

resource "google_project_iam_member" "ksa_token_creator" {
  project = var.project_id
  role    = "roles/iam.serviceAccountTokenCreator"
  member  = local.ksa_principal
}

resource "google_project_iam_member" "ksa_vertex_user" {
  project = var.project_id
  role    = "roles/aiplatform.user"
  member  = local.ksa_principal
}

resource "google_project_iam_member" "ksa_datastore_user" {
  project = var.project_id
  role    = "roles/datastore.user"
  member  = local.ksa_principal
}

Aplicar a configuração

Agora que nossa infraestrutura está definida, vamos provisioná-la.

  1. Primeiro, precisamos definir algumas variáveis para o Terraform usar. Vamos fazer isso usando variáveis de ambiente:
export TF_VAR_project_id=$(gcloud config get-value project)
export TF_VAR_project_number=$(gcloud projects describe $TF_VAR_project_id --format="value(projectNumber)")
export TF_VAR_region="us-central1"
  1. Inicialize o Terraform:
terraform init
  1. Use terraform plan para visualizar os recursos que serão criados.
terraform plan
  1. Aplique a configuração. Quando solicitado, digite yes para confirmar.
terraform apply

Observação: o provisionamento de um cluster do GKE pode levar de 10 a 15 minutos. Enquanto espera, você pode analisar o código do aplicativo na próxima seção.

  1. Depois de concluir, configure kubectl para se comunicar com o novo cluster:
gcloud container clusters get-credentials gradio-chat-cluster --region us-central1 --project $TF_VAR_project_id

4. Implantar o Gemma autohospedado no GKE

Em seguida, vamos implantar o modelo Gemma 3 12B diretamente no cluster do GKE. Isso permite inferência de baixa latência e controle total sobre o ambiente de execução do modelo.

Configurar credenciais do Hugging Face

Para fazer o download do modelo Gemma, seu cluster precisa de autenticação com o Hugging Face.

  1. Verifique se você tem um token do Hugging Face.
  2. Crie um secret do Kubernetes com seu token: substitua [YOUR_HF_TOKEN] pelo token real:
    kubectl create secret generic hf-secret --from-literal=hf_api_token=[YOUR_HF_TOKEN]
    

Implante o modelo

Vamos usar uma implantação padrão do Kubernetes para executar o modelo. O manifesto está localizado em deploy/gemma3-12b-deploy.yaml. É possível abrir manualmente ou usar este comando para abrir o arquivo no editor:

cd ../deploy
cloudshell edit gemma3-12b-deploy.yaml

Reserve um momento para inspecionar esse arquivo. Observe a seção resources:

        resources:
          requests:
            nvidia.com/gpu: 4
      nodeSelector:
        cloud.google.com/gke-accelerator: nvidia-l4

Essa é a infraestrutura de IA declarativa. Estamos informando ao GKE Autopilot que esse pod específico requer quatro GPUs NVIDIA L4. O Autopilot vai encontrar ou provisionar um nó que atenda exatamente a esses requisitos. Se um nó não estiver disponível para provisionamento, ele vai continuar tentando até que um nó que atenda aos requisitos esteja disponível.

  1. Aplique o manifesto de implantação:
    cd ..
    kubectl apply -f deploy/gemma3-12b-deploy.yaml
    
    Isso vai iniciar o processo de download dos pesos do modelo e iniciar o servidor de inferência. Isso geralmente leva alguns minutos. Essa implantação do Gemma usa GPUs, que podem estar sujeitas a falta de disponibilidade. Se as GPUs não estiverem disponíveis, o pod do Gemma vai ficar "pendente" até que estejam. O console do Google Cloud vai mostrar um erro como "Não é possível programar pods: a remoção não é útil para a programação" e/ou "Não é possível programar pods: os nós não corresponderam à afinidade/ao seletor de nós do pod". Isso significa que o GKE ainda não conseguiu adquirir nenhuma GPU para você. Ele vai continuar tentando até conseguir GPUs. Isso pode levar alguns minutos ou dias, dependendo da disponibilidade da GPU. É possível verificar o status com:
    kubectl get pods
    
    Você pode implantar o aplicativo mesmo que o pod gemma ainda não esteja em execução. O app de chat se conectará ao serviço do Gemma sempre que ele estiver disponível. Não será possível interagir com a Gemma pelo app de chat até que o pod gemma mostre o status Running e 1/1. Mas, enquanto isso, você pode conversar com o Gemini.

5. Criar o aplicativo de chat

Agora, vamos concluir o aplicativo Python. Abra app/app.py no editor do Cloud Shell. Você vai encontrar vários blocos # TODO que precisam ser preenchidos para que o aplicativo funcione.

cloudshell edit app/app.py

Etapa 1: processar o histórico de conversas

Os LLMs exigem que o histórico de conversas seja formatado de maneira específica para entender quem disse o quê.

O padrão "Tradutor universal":vamos escrever duas funções diferentes para processar o mesmo histórico de chat. Esse é um padrão fundamental em aplicativos multimodelos.

  • A fonte da verdade (Gradio): nosso app mantém o histórico em um formato simples e genérico: [[user_msg1, bot_msg1], ...].
  • Destino 1 (Gemma): precisa ser convertido em uma única string bruta com tokens especiais específicos.
  • Destino 2 (Gemini): precisa ser convertido em uma lista estruturada de objetos da API.

Ao reformatar o histórico genérico para o formato de destino a cada vez, podemos alternar entre os modelos sem problemas. Para adicionar um modelo diferente depois, você precisaria escrever uma nova função de processamento para o formato específico dele.

Para o Gemma (auto-hospedado)

Entender os modelos de chat:ao hospedar seus próprios modelos abertos, geralmente é necessário formatar manualmente o comando em uma string específica que o modelo foi treinado para reconhecer como uma conversa. Isso é conhecido como "modelo de chat".

Encontre a função process_message_gemma em app.py e substitua pelo código a seguir:

# This function takes a user's message and the conversation history as input.
#   Its job is to format these elements into a single,
#   structured prompt that can be understood by the language model (LLM).
#   This structured format helps the LLM maintain context and generate more relevant responses.
def process_message_gemma(message, history):
    user_prompt_format = "User's Turn:\n>>> {prompt}\n"
    assistant_prompt_format = "Assistant's Turn:\n>>> {prompt}\n"

    history_message = ""
    for user_turn, assistant_turn in history:
        history_message += user_prompt_format.format(prompt=user_turn)
        history_message += assistant_prompt_format.format(prompt=assistant_turn)

    # Format the new user message
    new_user_message = user_prompt_format.format(prompt=message)
    # Create a new aggregated message to be used as a single flat string in a json object sent to the LLM
    aggregated_message = (
        history_message + new_user_message + assistant_prompt_format.format(prompt="")
    )
    return aggregated_message

Para o Gemini (gerenciado)

Os serviços gerenciados geralmente preferem objetos estruturados em vez de strings brutas. Precisamos de uma função separada para formatar o histórico em objetos types.Content para o SDK do Gemini.

Encontre process_message_gemini e substitua por:

def process_message_gemini(message, history):
    contents = []
    for user_turn, model_turn in history:
        contents.append(
            types.Content(role="user", parts=[types.Part.from_text(text=user_turn)])
        )
        contents.append(
            types.Content(role="model", parts=[types.Part.from_text(text=model_turn)])
        )

    contents.append(
        types.Content(role="user", parts=[types.Part.from_text(text=message)])
    )
    return contents

Etapa 2: chamar o modelo Gemma auto-hospedado

Precisamos enviar nosso comando formatado para o serviço Gemma em execução no cluster. Vamos usar uma solicitação HTTP POST padrão para o nome DNS interno do serviço.

Encontre a função call_gemma_model e substitua por:

# Construct the request, send it to Gemma, return the model's response
# aggregated_message = current user message + history
def call_gemma_model(aggregated_message, model_temperature, top_p, max_tokens):
    json_message = {
        "prompt": aggregated_message,
        "temperature": model_temperature,
        "top_p": top_p,
        "max_tokens": max_tokens,
        "stop": ["User's Turn:"],
    }

    # Log what will be sent to the LLM
    print("*** JSON request: " + str(json_message))

    # Send the constructed json with the user prompt to the model and put the model's response in the json_data variable
    json_data = post_request(json_message)

    # The response from the model is a list of predictions. We'll take the first result.
    raw_output = json_data["predictions"][0]

    # The vLLM server returns the full prompt in the response. We need to extract
    # just the newly generated text from the model.
    assistant_turn_marker = "Assistant's Turn:\n>>>"
    marker_pos = raw_output.rfind(assistant_turn_marker)

    if marker_pos != -1:
        output = raw_output[marker_pos + len(assistant_turn_marker) :]
    else:
        output = raw_output

    # Clean up potential over-generation
    stop_marker = "User's Turn:"
    stop_pos = output.lower().find(stop_marker.lower())
    if stop_pos != -1:
        output = output[:stop_pos]

    return output.strip()

Etapa 3: chamar o modelo do Gemini da Vertex AI

Para o modelo gerenciado, vamos usar o SDK da IA generativa do Google. Isso é muito mais simples, já que processa as chamadas de rede para nós.

Encontre a função call_gemini_model e substitua por:

# Send a request to Gemini via the VertexAI API. Return the model's response
# contents = list of types.Content objects
def call_gemini_model(contents, model_temperature, top_p, max_tokens):
    gemini_model = "gemini-2.5-flash"

    response = client.models.generate_content(
        model=gemini_model,
        contents=contents,
        config={
            "temperature": model_temperature,
            "max_output_tokens": max_tokens,
            "top_p": top_p,
        },
    )
    return response.text

Etapa 4: implementar a interface principal de inferência

Por fim, precisamos da função principal de orquestração que o Gradio chama. Essa função precisa:

  1. Inicialize o histórico se ele estiver vazio.
  2. Processa a mensagem.
  3. Encaminhar a solicitação para o modelo selecionado (Gemma ou Gemini).
  4. Salve a interação no Firestore.
  5. Retorne a resposta para a interface.

Gradio e gerenciamento de estado:o ChatInterface do Gradio processa automaticamente o estado no nível da sessão (exibindo mensagens no navegador). No entanto, ele não tem suporte integrado para bancos de dados externos.

Para manter o histórico de chat por um longo período, usamos um padrão padrão: nos conectamos à função inference_interface. Ao aceitar request: gr.Request como argumento, o Gradio transmite automaticamente os detalhes da sessão do usuário atual. Usamos isso para criar um documento exclusivo do Firestore para cada usuário, garantindo que as conversas não se misturem em um ambiente multiusuário.

Encontre a função inference_interface e substitua por:

# This is the primary chat function. Every time a user sends a message, gradio calls this function,
# which sends the user's input to the appropriate AI (as indicated on the user interface), updates
# the chat history for future use during this session, and records the chat history in Firestore.
def inference_interface(
    message,
    history,
    model_name,
    model_temperature,
    top_p,
    max_tokens,
    request: gr.Request,
):

    # set history to empty array
    if history is None:
        history = []

    # Get or create session document
    session_hash = request.session_hash
    doc_id = f"session-{session_hash}"
    doc_ref = db.collection("chat_sessions").document(doc_id)

    # Create the session document if it doesn't exist
    if not doc_ref.get().exists:
        doc_ref.set({"Session start": datetime.datetime.now()})

    # Log info
    print("Model: " + model_name)
    print("LLM Engine: " + llm_engine)
    print("* History: " + str(history))

    # Pass the message and history to the appropriate model, as indicated by the user via the ui
    if model_name == "Gemma3 12b it":
        aggregated_message = process_message_gemma(message, history)
        output = call_gemma_model(
            aggregated_message, model_temperature, top_p, max_tokens
        )

    elif model_name == "Gemini":
        gemini_contents = process_message_gemini(message, history)
        output = call_gemini_model(
            gemini_contents, model_temperature, top_p, max_tokens
        )

    else:
        # Handle the case where no valid model is selected
        output = "Error: Invalid model selected."

    interaction = {"user": message, model_name: output}

    # Log the updated chat history
    print("* History: " + str(history) + " " + str(interaction))

    # Save the updated history to Firestore
    save_chat_history(interaction, doc_ref)

    return output

6. Verificar o arquivo app.py

Neste ponto, seu aplicativo de chat baseado em gradio deve estar pronto para implantação. Verifique se ele corresponde exatamente ao arquivo completo a seguir.

Solução de problemas:se você implantar o aplicativo e receber um erro "conexão recusada" ou "não é possível acessar este site" ao tentar se conectar a ele, repita as etapas a partir deste ponto. Comece copiando todo o arquivo e colando em app.py.

# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import datetime

import google.auth
import google.cloud
import gradio as gr
import requests
import themes
from google import genai

from google.cloud import firestore
from google.genai import types

## Do one-time initialization things

## grab the project id from google auth
_, project = google.auth.default()
print(f"Project: {project}")

# Set initial values for model
llm_engine = "vllm"
host = "http://gemma-service:8000"
context_path = "/generate"

# initialize vertex for interacting with Gemini
client = genai.Client(
    vertexai=True,
    project=project,
    location="global",
)

# Initialize Firestore client
db = firestore.Client(database="chat-app-db")


# This is the primary chat function. Every time a user sends a message, gradio calls this function,
# which sends the user's input to the appropriate AI (as indicated on the user interface), updates
# the chat history for future use during this session, and records the chat history in Firestore.
def inference_interface(
    message,
    history,
    model_name,
    model_temperature,
    top_p,
    max_tokens,
    request: gr.Request,
):

    # set history to empty array
    if history is None:
        history = []

    # Get or create session document
    session_hash = request.session_hash
    doc_id = f"session-{session_hash}"
    doc_ref = db.collection("chat_sessions").document(doc_id)

    # Create the session document if it doesn't exist
    if not doc_ref.get().exists:
        doc_ref.set({"Session start": datetime.datetime.now()})

    # Log info
    print("Model: " + model_name)
    print("LLM Engine: " + llm_engine)
    print("* History: " + str(history))

    # Pass the message and history to the appropriate model, as indicated by the user via the ui
    if model_name == "Gemma3 12b it":
        aggregated_message = process_message_gemma(message, history)
        output = call_gemma_model(
            aggregated_message, model_temperature, top_p, max_tokens
        )

    elif model_name == "Gemini":
        gemini_contents = process_message_gemini(message, history)
        output = call_gemini_model(
            gemini_contents, model_temperature, top_p, max_tokens
        )

    else:
        # Handle the case where no valid model is selected
        output = "Error: Invalid model selected."

    interaction = {"user": message, model_name: output}

    # Log the updated chat history
    print("* History: " + str(history) + " " + str(interaction))

    # Save the updated history to Firestore
    save_chat_history(interaction, doc_ref)

    return output


# Construct the request, send it to Gemma, return the model's response
# aggregated_message = current user message + history
def call_gemma_model(aggregated_message, model_temperature, top_p, max_tokens):
    json_message = {
        "prompt": aggregated_message,
        "temperature": model_temperature,
        "top_p": top_p,
        "max_tokens": max_tokens,
        "stop": ["User's Turn:"],
    }

    # Log what will be sent to the LLM
    print("*** JSON request: " + str(json_message))  # Log the JSON request

    # Send the constructed json with the user prompt to the model and put the model's response in the json_data variable
    json_data = post_request(json_message)

    # The response from the model is a list of predictions.
    # We'll take the first result.
    raw_output = json_data["predictions"][0]

    # The vLLM server returns the full prompt in the response. We need to extract
    # just the newly generated text from the model. The prompt ends with
    # "Assistant's Turn:\n>>>", so we find the last occurrence of that and
    # take everything after it.
    assistant_turn_marker = "Assistant's Turn:\n>>>"
    marker_pos = raw_output.rfind(assistant_turn_marker)

    if marker_pos != -1:
        # Get the text generated by the assistant
        output = raw_output[marker_pos + len(assistant_turn_marker) :]
    else:
        # Fallback in case the marker isn't found
        output = raw_output

    # The model sometimes continues the conversation and includes the next user's turn.
    # The 'stop' parameter is a good hint, but we parse the output as a safeguard.
    stop_marker = "User's Turn:"
    stop_pos = output.lower().find(stop_marker.lower())
    if stop_pos != -1:
        output = output[:stop_pos]

    # The model also sometimes prefixes its response with "Output:". We'll remove this.
    output = output.lstrip()
    prefix_marker = "Output:"
    if output.lower().startswith(prefix_marker.lower()):
        output = output[len(prefix_marker) :]

    return output.strip()


# Send a request to Gemini via the VertexAI API. Return the model's response
# contents = list of types.Content objects
def call_gemini_model(contents, model_temperature, top_p, max_tokens):
    gemini_model = "gemini-2.5-flash"

    response = client.models.generate_content(
        model=gemini_model,
        contents=contents,
        config={
            "temperature": model_temperature,
            "max_output_tokens": max_tokens,
            "top_p": top_p,
        },
    )
    output = response.text  # Extract the generated text
    # Consider handling additional response attributes (safety, usage, etc.)
    return output


def process_message_gemini(message, history):
    contents = []
    for user_turn, model_turn in history:
        contents.append(
            types.Content(role="user", parts=[types.Part.from_text(text=user_turn)])
        )
        contents.append(
            types.Content(role="model", parts=[types.Part.from_text(text=model_turn)])
        )

    contents.append(
        types.Content(role="user", parts=[types.Part.from_text(text=message)])
    )
    return contents


# This function takes a user's message and the conversation history as input.
#   Its job is to format these elements into a single,
#   structured prompt that can be understood by the language model (LLM).
#   This structured format helps the LLM maintain context and generate more relevant responses.
def process_message_gemma(message, history):
    user_prompt_format = "User's Turn:\n>>> {prompt}\n"
    assistant_prompt_format = "Assistant's Turn:\n>>> {prompt}\n"

    history_message = ""
    for user_turn, assistant_turn in history:
        history_message += user_prompt_format.format(prompt=user_turn)
        history_message += assistant_prompt_format.format(prompt=assistant_turn)

    # Format the new user message
    new_user_message = user_prompt_format.format(prompt=message)
    # Create a new aggregated message to be used as a single flat string in a json object sent to the LLM
    aggregated_message = (
        history_message + new_user_message + assistant_prompt_format.format(prompt="")
    )
    return aggregated_message


# Function to save chat history to Firestore
def save_chat_history(interaction, doc_ref):
    timestamp_str = str(datetime.datetime.now())

    # Save the chat history, merging with existing data
    doc_ref.update({timestamp_str: interaction})

    print("Chat history saved successfully!")  # Optional: Log success


# Send the json message to the model and return the model's response. This is used for Gemma but not Gemini. It could also be used for other models.
def post_request(json_message):
    print("*** Request" + str(json_message), flush=True)
    # Set a timeout and check for HTTP errors. This will raise an exception on a bad status code (4xx or 5xx).
    response = requests.post(host + context_path, json=json_message, timeout=60)
    response.raise_for_status()
    json_data = response.json()
    print("*** Output: " + str(json_data), flush=True)
    return json_data


# custom css to hide default footer
css = """
footer {display: none !important;} .gradio-container {min-height: 0px !important;}
"""

# Add a dropdown to select the model to chat with
model_dropdown = gr.Dropdown(
    ["Gemma3 12b it", "Gemini"],
    label="Model",
    info="Select the model you would like to chat with.",
    value="Gemma3 12b it",
)

# Make the model temperature, top_p, and max tokents modifiable via sliders in the GUI
model_temperature = gr.Slider(
    minimum=0.1, maximum=1.0, value=0.9, label="Temperature", render=False
)
top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, label="Top_p", render=False)
max_tokens = gr.Slider(
    minimum=1, maximum=4096, value=1024, label="Max Tokens", render=False
)

# Call gradio to create the chat interface
app = gr.ChatInterface(
    inference_interface,
    additional_inputs=[model_dropdown, model_temperature, top_p, max_tokens],
    theme=themes.google_theme(),
    css=css,
    title="Chat with AI",
)

app.launch(server_name="0.0.0.0", allowed_paths=["images"])

7. Implantar o aplicativo de chat

Vamos usar o Skaffold para criar a imagem do contêiner e implantá-la no cluster. O Skaffold é uma ferramenta de linha de comando que orquestra e automatiza o processo de criação, envio e implantação de aplicativos no Kubernetes. Ele simplifica o fluxo de trabalho de desenvolvimento, permitindo que você acione todo esse processo com um único comando, o que o torna ideal para iterar no aplicativo.

Observação: isso também vai implantar a conta de serviço do Kubernetes necessária para a Identidade da carga de trabalho. É possível conferir a definição dele no arquivo deploy/chat-deploy.yaml. Confira a definição aqui, para referência:

apiVersion: v1
kind: ServiceAccount
metadata:
  name: gradio-chat-ksa

Execute o Skaffold para criar e implantar:

skaffold run --default-repo=us-central1-docker.pkg.dev/$GOOGLE_CLOUD_PROJECT/chat-app-repo

O Skaffold vai usar o Cloud Build para criar a imagem do contêiner, enviá-la para o Artifact Registry criado pelo Terraform e aplicar os manifestos do Kubernetes ao cluster.

8. Teste o aplicativo

  1. Aguarde até que o serviço do aplicativo de chat receba um endereço IP externo:
    kubectl get svc gradio-chat-service --watch
    
    Quando o EXTERNAL-IP mudar de pending para um endereço IP real, pressione Ctrl+C para parar de assistir.
  2. Abra um navegador da Web e acesse http://[EXTERNAL-IP]:7860.
  3. Interaja com o modelo! Por padrão, o app está configurado para você conversar com seu modelo Gemma hospedado localmente. Se quiser conversar com o Gemini, mude o modelo no menu suspenso "Outras entradas". Por exemplo, pergunte à IA: "Conte uma piada sobre o Kubernetes".

Solução de problemas:

  1. Se você receber um erro como "Não é possível acessar este site" ou "[EXTERNAL-IP] recusou a conexão", talvez algo tenha dado errado com o arquivo app.py. Volte para a etapa "Verificar o arquivo app.py" e repita as etapas de lá.
  2. A interface usa o modelo "Gemma3 12b it" por padrão. Se você receber um erro imediatamente, provavelmente é porque o pod do Gemma ainda não está pronto. Dica:você pode mudar o menu suspenso para "Gemini" e testar a interação com o aplicativo de chat enquanto espera a inicialização da Gemma.

Teste a Gemma:verifique se "Gemma3 12b it" está selecionado no menu suspenso e envie uma mensagem (por exemplo, "Conte uma piada sobre o Kubernetes").

Teste o Gemini:mude o menu suspenso para "Gemini" e faça outra pergunta (por exemplo, "Qual é a diferença entre um pod e um nó?").

Verificar o histórico:depois de conversar com um modelo (Gemma ou Gemini) no app de chat, confira o banco de dados "chat-app-db" no Firestore para ver os registros de conversa. Se você conseguiu conversar com os dois modelos, observe que o histórico da conversa é mantido mesmo ao trocar de modelo.

9. Indo mais longe

Agora que você tem um aplicativo de chat híbrido funcionando, considere estes desafios para aprofundar seu entendimento:

  1. Personagem personalizado:tente modificar as funções process_message_gemma e process_message_gemini para incluir um "comando do sistema" no início. Por exemplo, diga aos modelos: "Você é um assistente pirata útil" e veja como isso muda as respostas deles.
  2. Identidade de usuário persistente:no momento, o aplicativo gera um novo UUID aleatório para cada sessão. Como você integraria um sistema de autenticação real (como o Login do Google) para que um usuário pudesse ver o histórico de conversas anteriores em diferentes dispositivos?
  3. Experimentação com modelos:tente mudar o controle deslizante temperature na interface. Como uma temperatura alta (próxima de 1,0) afeta a criatividade e a precisão das respostas em comparação com uma temperatura baixa (próxima de 0,1)?

10. Conclusão

Parabéns! Você criou um aplicativo de IA híbrida. Você aprendeu a:

  • Use o Terraform para infraestrutura como código no Google Cloud.
  • Hospede seus próprios LLMs de código aberto no GKE para ter controle total.
  • Integre serviços de IA gerenciados, como a Vertex AI, para ter mais flexibilidade.
  • Criar um aplicativo com estado usando o Firestore para persistência.
  • Proteja suas cargas de trabalho usando a Identidade da carga de trabalho.

Limpeza

Para evitar cobranças, destrua os recursos criados:

cd infra
terraform destroy -var="project_id=$GOOGLE_CLOUD_PROJECT" -var="project_number=$PROJECT_NUMBER" -var="region=$REGION"