Stream responses from Vertex AI through Kong AI Gateway using Google Generative AI SDK

echo '
_format_version: "3.0"
plugins:
  - name: ai-proxy-advanced
    service: gemini-service
    config:
      llm_format: gemini
      genai_category: text/generation
      targets:
      - route_type: llm/v1/chat
        logging:
          log_payloads: false
          log_statistics: true
        model:
          provider: gemini
          name: gemini-2.0-flash-exp
          options:
            gemini:
              api_endpoint: "${{ env "DECK_GCP_API_ENDPOINT" }}"
              project_id: "${{ env "DECK_GCP_PROJECT_ID" }}"
              location_id: "${{ env "DECK_GCP_LOCATION_ID" }}"
        auth:
          allow_override: false
          gcp_use_service_account: true
          gcp_service_account_json: "${{ env "DECK_GCP_SERVICE_ACCOUNT_JSON" }}"
' | deck gateway apply -

cat << 'EOF' > vertex_stream.py
#!/usr/bin/env python3
from google import genai
from google.genai.types import HttpOptions
import os
import sys

PROJECT_ID = os.getenv("DECK_GCP_PROJECT_ID")
LOCATION = os.getenv("DECK_GCP_LOCATION_ID")

if not PROJECT_ID:
    print("Error: DECK_GCP_PROJECT_ID environment variable not set")
    sys.exit(1)

def vertex_stream(show_raw=False, show_chunks=False):
    """Stream responses from Vertex AI through Kong Gateway"""

    # Configure client to route through Kong Gateway
    client = genai.Client(
        vertexai=True,
        project=PROJECT_ID,
        location=LOCATION,
        http_options=HttpOptions(
            base_url="http://localhost:8000/gemini",
            api_version="v1"
        )
    )

    try:
        if show_raw:
            print("Streaming with raw output...\n")

        chunk_num = 0
        for chunk in client.models.generate_content_stream(
            model="gemini-2.0-flash-exp",
            contents="Explain quantum entanglement in one paragraph"
        ):
            chunk_num += 1

            if show_chunks:
                print(f"\n--- Chunk {chunk_num} ---")
                if hasattr(chunk, 'candidates') and chunk.candidates:
                    candidate = chunk.candidates[0]
                    if hasattr(candidate, 'finish_reason') and candidate.finish_reason:
                        print(f"Finish reason: {candidate.finish_reason}")
                    if hasattr(chunk, 'usage_metadata') and chunk.usage_metadata:
                        if hasattr(chunk.usage_metadata, 'total_token_count'):
                            print(f"Total tokens: {chunk.usage_metadata.total_token_count}")
                print("Text: ", end="")

            if show_raw:
                print(f"\nChunk {chunk_num}:", chunk)
                print("-" * 80)

            print(chunk.text, end="", flush=True)

            if show_chunks:
                print()

        if not show_chunks:
            print()

    except Exception as e:
        print(f"Error: {e}")

if __name__ == "__main__":
    show_raw = "--raw" in sys.argv
    show_chunks = "--chunks" in sys.argv
    vertex_stream(show_raw, show_chunks)
EOF

python3 vertex_stream.py

Connecting to: http://localhost:8000/gemini/v1/projects/your-project/locations/us-central1/publishers/google/models/gemini-2.0-flash-exp:streamGenerateContent

Quantum entanglement is a bizarre phenomenon where two or more particles become linked together in such a way that they share the same fate, no matter how far apart they are. Measuring the state of one entangled particle instantly influences the state of the other, even across vast distances, seemingly violating the classical concept of locality. This "spooky action at a distance" means knowing the property of one particle immediately reveals the corresponding property of its entangled partner, even before any measurement is made on it directly.

python3 vertex_stream.py --chunks

Connecting to: http://localhost:8000/gemini/v1/projects/your-project/locations/us-central1/publishers/google/models/gemini-2.0-flash-exp:streamGenerateContent

--- Chunk 1 ---
Total tokens: None
Text: Quantum

--- Chunk 2 ---
Total tokens: None
Text:  entanglement is a

--- Chunk 3 ---
Total tokens: None
Text:  bizarre phenomenon where two or more particles become linked together in such a way that they

--- Chunk 4 ---
Total tokens: None
Text:  share the same fate, no matter how far apart they are. Measuring the properties

--- Chunk 5 ---
Total tokens: None
Text:  of one entangled particle instantaneously determines the corresponding properties of the other, even if they're separated by vast distances. This correlation isn't due to some pre-existing hidden

--- Chunk 6 ---
Finish reason: STOP
Total tokens: 100
Text:  information but is instead a fundamental connection arising from their shared quantum state, defying classical intuition about locality and causality.

python3 vertex_stream.py --raw

python3 vertex_stream.py --raw --chunks