Compila una app de chat híbrida con IA en GKE con Gemma y Gemini

1. Introducción

Descripción general

En este lab, compilarás e implementarás una aplicación de chat de pila completa potenciada por IA en Google Kubernetes Engine (GKE). Esta aplicación "híbrida" demuestra un potente patrón arquitectónico: la capacidad de cambiar sin problemas entre un modelo abierto alojado por el usuario (Gemma 3 12B) que se ejecuta directamente en tu clúster y un servicio de IA administrado (Gemini 2.5 Flash a través de 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   |                        |
|                                            +----------------------+                        |
+--------------------------------------------------------------------------------------------+

Usarás Terraform para aprovisionar la infraestructura, incluido un clúster de GKE Autopilot y una base de datos de Firestore para conservar el historial de sesiones de chat. Luego, completarás el código de la aplicación de Python para controlar conversaciones de varios turnos, interactuar con ambos modelos de IA y, por último, implementar la aplicación final con Cloud Build y Skaffold.

Qué aprenderás

  • Aprovisiona la infraestructura de GKE y Firestore con Terraform.
  • Implementa un modelo de lenguaje grande (Gemma) en GKE Autopilot con manifiestos de Kubernetes.
  • Implementa una interfaz de chat de Gradio en Python que pueda cambiar entre diferentes backends de IA.
  • Usa Firestore para almacenar y recuperar el historial de sesiones de chat.
  • Configurar Workload Identity para otorgar de forma segura a tus cargas de trabajo de GKE acceso a los servicios de Google Cloud (Vertex AI, Firestore)

Requisitos previos

  • Un proyecto de Google Cloud con facturación habilitada.
  • Conocimientos básicos sobre Python, Kubernetes y herramientas de línea de comandos estándar
  • Un token de Hugging Face con acceso a los modelos de Gemma

2. Configuración del proyecto

  1. Si aún no tienes una Cuenta de Google, debes crear una.
    • Usar una cuenta personal en lugar de una cuenta laboral o educativa Es posible que las cuentas laborales y educativas tengan restricciones que te impidan habilitar las APIs necesarias para este lab.
  2. Accede a la consola de Google Cloud.
  3. Habilita la facturación en la consola de Cloud.
    • Completar este lab debería costar menos de USD 1 en recursos de Cloud.
    • Puedes seguir los pasos al final de este lab para borrar recursos y evitar cargos adicionales.
    • Los usuarios nuevos pueden acceder a la prueba gratuita de USD 300.
  4. Crea un proyecto nuevo o elige reutilizar uno existente.

Abre el editor de Cloud Shell

  1. Haz clic en este vínculo para navegar directamente al editor de Cloud Shell.
  2. Si se te solicita autorización en algún momento, haz clic en Autorizar para continuar. Haz clic para autorizar Cloud Shell
  3. Si la terminal no aparece en la parte inferior de la pantalla, ábrela:
    • Haz clic en Ver.
    • Haz clic en Terminal.Abre una terminal nueva en el editor de Cloud Shell
  4. En la terminal, configura tu proyecto con este comando:
    • Formato:
      gcloud config set project [PROJECT_ID]
    • Ejemplo:
      gcloud config set project lab-project-id-example
    • Si no recuerdas el ID de tu proyecto, haz lo siguiente:
      • Puedes enumerar todos los IDs de tus proyectos con el siguiente comando:
        gcloud projects list | awk '/PROJECT_ID/{print $2}'
      Establece el ID del proyecto en la terminal del editor de Cloud Shell
  5. Deberías ver el siguiente mensaje:
    Updated property [core/project].
    Si ves un WARNING y se te pregunta Do you want to continue (Y/n)?, es probable que hayas ingresado el ID del proyecto de forma incorrecta. Presiona n, presiona Enter y vuelve a intentar ejecutar el comando gcloud config set project.

Clona el repositorio

En la terminal de Cloud Shell, clona el repositorio del proyecto y navega al directorio del proyecto:

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

Tómate un momento para explorar la estructura del proyecto:

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

Configura variables de entorno

Configura las variables de entorno para el ID y el número de tu proyecto. Terraform y los comandos posteriores los usarán.

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

Habilita la API de Cloud Resource Manager

Terraform requiere que la API de Cloud Resource Manager esté habilitada para administrar los recursos de tu proyecto, por lo que primero debemos habilitarla. Más adelante, implementaremos nuestra aplicación de chat con Skaffold, que usa Cloud Build para compilar nuestra imagen de contenedor. Ahora habilitaremos la API de Storage y crearemos el bucket necesario para Cloud Build. Usaremos Terraform para habilitar el resto de las APIs necesarias para este proyecto.

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

Crea un bucket de etapa de pruebas de Cloud Build

Skaffold usa Google Cloud Build, que requiere un bucket de Cloud Storage para organizar tu código fuente.

Crea el archivo ahora para asegurarte de que exista:

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

(Si recibes un error que indica que el bucket ya existe, puedes ignorarlo sin problemas).

3. Aprovisiona infraestructura con Terraform

Usaremos Terraform para configurar los recursos necesarios de Google Cloud. Esto garantiza un entorno coherente y reproducible.

  1. Navega al directorio de infraestructura:
    cd infra

Este archivo define las APIs adicionales que necesitaremos para este proyecto: cloudbuild, artifactregistry, container (gke), firestore y aiplatform (vertexai). Consulta el archivo o la siguiente información para ver cómo se habilitan las APIs a través de 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
}

Define el clúster de GKE

Abre infra/main.tf en tu editor. Verás varios comentarios de # TODO. Puedes abrirlo de forma manual o usar este comando para abrir el archivo en el editor:

cloudshell edit main.tf

Primero, debemos definir nuestro clúster de Kubernetes. Usaremos GKE Autopilot, que es ideal para las cargas de trabajo de IA, ya que administra los nodos de forma automática.

Busca # TODO: Create a GKE Autopilot Cluster y agrega el siguiente bloque debajo:

# 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]
}

Observa enable_autopilot = true. Esta sola línea nos evita tener que administrar grupos de nodos, el ajuste de escala automático y el empaquetamiento de nuestras cargas de trabajo de GPU.

Define la base de datos de Firestore

A continuación, necesitamos un lugar para almacenar nuestro historial de chat. Firestore es una base de datos NoSQL sin servidores que se adapta perfectamente a esta necesidad.

Busca # TODO: Create a Firestore Database y agrega lo siguiente:

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]
}

Después de agregar el recurso de la base de datos, busca # TODO: Create an initial Firestore Document y agrega el siguiente bloque. Este recurso crea un documento de marcador de posición inicial en nuestra colección, lo que resulta útil para inicializar la estructura de la base de datos.

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]
}

Define Workload Identity

Por último, debemos configurar la seguridad. Queremos que nuestros Pods de Kubernetes puedan acceder a Vertex AI y Firestore sin que tengamos que administrar secretos ni claves de API. Esto lo hacemos con Workload Identity.

Otorgaremos los roles de IAM necesarios a la cuenta de servicio de Kubernetes (KSA) que usará nuestra app.

Nota: La cuenta de servicio de Kubernetes (gradio-chat-ksa) a la que se hace referencia en estas vinculaciones aún no existe. Se creará más adelante cuando implementemos nuestra aplicación en el clúster. Es perfectamente aceptable (y una práctica común) aprovisionar previamente estas vinculaciones de IAM.

Busca # TODO: Configure Workload Identity IAM bindings y agrega lo siguiente:

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
}

Aplica la configuración

Ahora que definimos nuestra infraestructura, aprovisionémosla.

  1. Primero, deberemos establecer algunas variables para que Terraform las use. Lo haremos con variables de entorno:
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. Inicializa Terraform mediante este comando:
terraform init
  1. Usa terraform plan para obtener una vista previa de los recursos que se crearán.
terraform plan
  1. Aplica la configuración. Cuando se le solicite, escriba yes para confirmar.
terraform apply

Nota: El aprovisionamiento de un clúster de GKE puede tardar entre 10 y 15 minutos. Mientras esperas, puedes revisar el código de la aplicación en la siguiente sección.

  1. Una vez que se complete, configura kubectl para que se comunique con tu clúster nuevo:
gcloud container clusters get-credentials gradio-chat-cluster --region us-central1 --project $TF_VAR_project_id

4. Implementa Gemma autohospedada en GKE

A continuación, implementaremos el modelo Gemma 3 12B directamente en tu clúster de GKE. Esto permite la inferencia de baja latencia y el control completo sobre el entorno de ejecución del modelo.

Configura las credenciales de Hugging Face

Para descargar el modelo de Gemma, tu clúster necesita autenticación con Hugging Face.

  1. Asegúrate de tener un token de Hugging Face.
  2. Crea un Secret de Kubernetes con tu token. Reemplaza [YOUR_HF_TOKEN] por tu token real:
    kubectl create secret generic hf-secret --from-literal=hf_api_token=[YOUR_HF_TOKEN]

Implementa el modelo

Usaremos una implementación estándar de Kubernetes para ejecutar el modelo. El manifiesto se encuentra en deploy/gemma3-12b-deploy.yaml. Puedes abrirlo de forma manual o usar este comando para abrir el archivo en el editor:

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

Dedica un momento a inspeccionar este archivo. Observa la sección resources:

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

Esta es la infraestructura de IA declarativa. Le indicamos a GKE Autopilot que este Pod específico requiere 4 GPU NVIDIA L4. Autopilot encontrará o aprovisionará un nodo que cumpla exactamente con estos requisitos. Si un nodo no está disponible para el aprovisionamiento, se seguirá intentando hasta que haya un nodo que cumpla con los requisitos.

  1. Aplica el manifiesto de la implementación:
    cd ..
kubectl apply -f deploy/gemma3-12b-deploy.yaml
    Esto iniciará el proceso de descarga de los pesos del modelo y el inicio del servidor de inferencia. Por lo general, esto puede tardar varios minutos. Esta implementación de Gemma usa GPUs, que pueden estar sujetas a escasez de disponibilidad. Si las GPUs no están disponibles, el pod de Gemma permanecerá en estado "pendiente" hasta que lo estén, y la consola de Google Cloud mostrará un error como "No se pueden programar pods: La prioridad no es útil para la programación" o "No se pueden programar pods: Los nodos no coincidieron con la afinidad o el selector de nodos del pod". Esto significa que GKE aún no pudo adquirir ninguna GPU para ti. Seguirá intentándolo hasta que pueda adquirir GPUs. Esto podría tardar unos minutos o días, según la disponibilidad de la GPU. Puedes verificar el estado con el siguiente comando:
    kubectl get pods
    Puedes continuar con la implementación de la aplicación incluso si el pod gemma aún no está en funcionamiento. La app de chat se conectará al servicio de Gemma cuando esté disponible. Ten en cuenta que no podrás interactuar con Gemma a través de tu app de chat hasta que el pod gemma muestre el estado Running y 1/1. Pero, mientras tanto, puedes chatear con Gemini.

5. Compila la aplicación de Chat

Ahora, completemos la aplicación de Python. Abre app/app.py en el editor de Cloud Shell. Encontrarás varios bloques # TODO que deben completarse para que la aplicación funcione.

cloudshell edit app/app.py

Paso 1: Procesa el historial de conversaciones

Los LLM requieren que el historial de conversaciones se formatee de manera específica para que comprendan quién dijo qué.

El patrón de "traductor universal": Ten en cuenta que estamos a punto de escribir dos funciones diferentes para procesar el mismo historial de chat. Este es un patrón clave en las aplicaciones de varios modelos.

  • La fuente de verdad (Gradio): Nuestra app mantiene el historial en un formato simple y genérico: [[user_msg1, bot_msg1], ...].
  • Objetivo 1 (Gemma): Necesita que se convierta en una sola cadena sin procesar con tokens especiales específicos.
  • Objetivo 2 (Gemini): Necesita que se convierta en una lista estructurada de objetos de la API.

Al volver a dar formato al historial genérico en el formato objetivo en cada turno, podemos cambiar de modelo sin problemas. Si deseas agregar un modelo diferente más adelante, deberás escribir una nueva función de procesamiento para su formato específico.

Para Gemma (autoalojada)

Comprensión de las plantillas de chat: Cuando alojas tus propios modelos abiertos, por lo general, debes darle formato manualmente a la instrucción en una cadena específica que el modelo se entrenó para reconocer como una conversación. Esto se conoce como "plantilla de chat".

Busca la función process_message_gemma en app.py y reemplázala por el siguiente código:

# 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 Gemini (administrado)

Los servicios administrados suelen preferir los objetos estructurados a las cadenas sin procesar. Necesitamos una función separada para dar formato al historial en objetos types.Content para el SDK de Gemini.

Busca process_message_gemini y reemplázalo por lo siguiente:

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

Paso 2: Llama al modelo Gemma alojado por tu cuenta

Tenemos que enviar nuestra instrucción con formato al servicio de Gemma que se ejecuta en nuestro clúster. Usaremos una solicitud HTTP POST estándar para el nombre de DNS interno del servicio.

Busca la función call_gemma_model y reemplázala por lo siguiente:

# 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()

Paso 3: Llama al modelo de Vertex AI Gemini

Para el modelo administrado, usaremos el SDK de IA generativa de Google. Esto es mucho más simple, ya que controla las llamadas de red por nosotros.

Busca la función call_gemini_model y reemplázala por lo siguiente:

# 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

Paso 4: Implementa la interfaz de inferencia principal

Por último, necesitamos la función principal del orquestador a la que llama Gradio. Esta función debe hacer lo siguiente:

  1. Inicializa el historial si está vacío.
  2. Procesa el mensaje.
  3. Enruta la solicitud al modelo seleccionado (Gemma o Gemini).
  4. Guarda la interacción en Firestore.
  5. Devuelve la respuesta a la IU.

Gradio y administración de estados: El ChatInterface de Gradio controla automáticamente el estado a nivel de la sesión (muestra mensajes en el navegador). Sin embargo, no tiene compatibilidad integrada con bases de datos externas.

Para conservar el historial de chat a largo plazo, usamos un patrón estándar: nos conectamos a la función inference_interface. Cuando aceptamos request: gr.Request como argumento, Gradio nos pasa automáticamente los detalles de la sesión del usuario actual. Usamos esto para crear un documento único de Firestore para cada usuario, lo que garantiza que las conversaciones no se mezclen en un entorno multiusuario.

Busca la función inference_interface y reemplázala por lo siguiente:

# 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. Revisa tu archivo app.py

En este punto, tu aplicación de chat basada en Gradio debería estar lista para implementarse. Asegúrate de que coincida exactamente con el siguiente archivo completo.

Solución de problemas: Si implementas tu aplicación y recibes un error de "Se rechazó la conexión" o "No se puede acceder a este sitio" cuando intentas conectarte a ella, repite los pasos desde este punto. Para ello, copia todo este archivo y pégalo en tu 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. Implementa la aplicación de Chat

Usaremos Skaffold para compilar nuestra imagen de contenedor y, luego, la implementaremos en el clúster. Skaffold es una herramienta de línea de comandos que coordina y automatiza el proceso de compilación, envío e implementación de aplicaciones en Kubernetes. Simplifica el flujo de trabajo de desarrollo, ya que te permite activar todo este proceso con un solo comando, lo que lo hace ideal para iterar en tu aplicación.

Nota: Esto también implementará la cuenta de servicio de Kubernetes que necesitamos para Workload Identity. Puedes ver su definición en el archivo deploy/chat-deploy.yaml. Consulta su definición aquí como referencia:

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

Ejecuta Skaffold para compilar e implementar:

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

Skaffold usará Cloud Build para compilar la imagen del contenedor, enviarla a Artifact Registry creado por Terraform y, luego, aplicar los manifiestos de Kubernetes a tu clúster.

8. Pruebe la aplicación

  1. Espera a que el servicio de la aplicación de chat obtenga una dirección IP externa:
    kubectl get svc gradio-chat-service --watch
    Una vez que EXTERNAL-IP cambie de pending a una dirección IP real, presiona Ctrl+C para dejar de mirar.
  2. Abre un navegador web y ve a http://[EXTERNAL-IP]:7860.
  3. Intenta interactuar con el modelo. De forma predeterminada, la app está configurada para que chatees con tu modelo de Gemma alojado de forma local. Si quieres chatear con Gemini, cambia el modelo en el menú desplegable "Entradas adicionales". Por ejemplo, pídele a la IA: "Cuéntame un chiste de Kubernetes".

Solución de problemas:

  1. Si recibes un error como "No se puede acceder a este sitio" o "[EXTERNAL-IP] rechazó la conexión", es posible que haya ocurrido un problema con tu archivo app.py. Vuelve al paso titulado "Verifica tu archivo app.py" y repite los pasos desde allí.
  2. La IU se establece de forma predeterminada en el modelo "Gemma3 12b it". Si recibes un error de inmediato, es probable que el pod de Gemma aún no esté listo. Nota: Puedes cambiar el menú desplegable a "Gemini" para probar la interacción con la aplicación de chat mientras esperas a que se inicialice Gemma.

Probar Gemma: Asegúrate de que "Gemma3 12b it" esté seleccionado en el menú desplegable y envía un mensaje (p.ej., "Cuéntame un chiste sobre Kubernetes").

Probar Gemini: Cambia el menú desplegable a "Gemini" y haz otra pregunta (p.ej., "¿Cuál es la diferencia entre un pod y un nodo?").

Historial de verificación: Una vez que hayas chateado correctamente con un modelo (Gemma o Gemini) en la app de chat, consulta tu base de datos "chat-app-db" en Firestore para ver los registros de chat. Si pudiste chatear con ambos modelos, observa que el historial de conversaciones se mantiene incluso cuando cambias de modelo.

9. Llega más lejos

Ahora que tienes una aplicación de chat híbrida que funciona, considera estos desafíos para profundizar tu comprensión:

  1. Arquetipo personalizado: Intenta modificar las funciones process_message_gemma y process_message_gemini para incluir una "instrucción del sistema" al principio. Por ejemplo, dile a los modelos "Eres un asistente pirata útil" y observa cómo cambian sus respuestas.
  2. Identidad de usuario persistente: Actualmente, la aplicación genera un UUID aleatorio nuevo para cada sesión. ¿Cómo integrarías un sistema de autenticación real (como el Acceso con Google) para que un usuario pueda ver su historial de conversaciones anterior en diferentes dispositivos?
  3. Experimentación con el modelo: Intenta cambiar el control deslizante temperature en la IU. ¿Cómo afecta una temperatura alta (cerca de 1.0) la creatividad y la precisión de las respuestas en comparación con una temperatura baja (cerca de 0.1)?

10. Conclusión

¡Felicitaciones! Compilaste correctamente una aplicación híbrida basada en IA. Aprendió a hacer lo siguiente:

  • Usa Terraform para la infraestructura como código en Google Cloud.
  • Aloja tus propios LLM de código abierto en GKE para tener un control total.
  • Integra servicios de IA administrados, como Vertex AI, para obtener flexibilidad.
  • Compila una aplicación con estado usando Firestore para la persistencia.
  • Protege tus cargas de trabajo con Workload Identity.

Limpieza

Para evitar que se generen cargos, destruye los recursos que creaste:

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