test_gemini.py

"""Live diagnostic for the configured Gemini API key and Gemini Live API.

This is intentionally a runnable diagnostic, not a pytest unit test, because it
makes real billable/network API calls.
"""

import argparse
import asyncio
import sys
from pathlib import Path

import async_timeout
from google import genai
from google.genai import types

from gargi_ai.config import get_settings


DEFAULT_PROMPT = (
    "You are Professor AI. Explain why the sky looks blue to a curious "
    "12-year-old in three short paragraphs, then ask one check-for-understanding "
    "question."
)
DEFAULT_LIVE_MODEL = "gemini-3.1-flash-live-preview"


def print_heading(title: str) -> None:
    print(f"\n{'=' * 72}\n{title}\n{'=' * 72}")


def check_standard_response(
    client: genai.Client, model: str, prompt: str
) -> bool:
    print_heading("1. STANDARD GEMINI API TEST")
    print(f"Model : {model}")
    print(f"Prompt: {prompt}\n")

    try:
        response = client.models.generate_content(model=model, contents=prompt)
        answer = (response.text or "").strip()
        if not answer:
            print("FAILED: Gemini returned an empty response.")
            return False

        print("Gemini response:\n")
        print(answer)
        print("\nStatus: PASSED - normal Gemini generation is responding.")
        return True
    except Exception as exc:
        print(f"Status: FAILED - {type(exc).__name__}: {exc}")
        return False


async def check_live_response(
    client: genai.Client,
    model: str,
    output_path: Path | None,
) -> bool:
    print_heading("2. GEMINI LIVE API TEST")
    print(f"Live model: {model}")
    print("Opening a WebSocket session and requesting a short audio greeting...\n")

    audio = bytearray()
    transcript_parts: list[str] = []
    turn_completed = False

    config = types.LiveConnectConfig(
        response_modalities=[types.Modality.AUDIO],
        output_audio_transcription=types.AudioTranscriptionConfig(),
        system_instruction=(
            "You are Professor AI. Keep this diagnostic response to one short sentence."
        ),
    )

    try:
        async with async_timeout.timeout(45):
            async with client.aio.live.connect(model=model, config=config) as session:
                await session.send_realtime_input(
                    text="Say: Hello! Professor AI Live API is working."
                )

                async for response in session.receive():
                    content = response.server_content
                    if not content:
                        continue

                    if content.output_transcription and content.output_transcription.text:
                        transcript_parts.append(content.output_transcription.text)

                    if content.model_turn:
                        for part in content.model_turn.parts or []:
                            if part.inline_data and part.inline_data.data:
                                audio.extend(part.inline_data.data)

                    if content.turn_complete:
                        turn_completed = True
                        break

        transcript = "".join(transcript_parts).strip()
        print(f"WebSocket connected : yes")
        print(f"Turn completed      : {turn_completed}")
        print(f"Audio received      : {len(audio):,} bytes (PCM 24 kHz, 16-bit)")
        print(f"Transcription       : {transcript or '<not returned>'}")

        if output_path and audio:
            output_path.parent.mkdir(parents=True, exist_ok=True)
            output_path.write_bytes(audio)
            print(f"Raw audio saved     : {output_path.resolve()}")

        passed = turn_completed and bool(audio)
        if passed:
            print("\nStatus: PASSED - this key supports Gemini Live API.")
        else:
            print("\nStatus: FAILED - the session opened but no complete audio turn arrived.")
        return passed
    except Exception as exc:
        print(f"Status: FAILED - {type(exc).__name__}: {exc}")
        print(
            "The normal API may still work even if this preview Live model is "
            "unavailable for the key, account, or region."
        )
        return False


async def main() -> int:
    parser = argparse.ArgumentParser(
        description="Test the Gemini API key with normal generation and Live API."
    )
    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
    parser.add_argument("--skip-live", action="store_true")
    parser.add_argument(
        "--live-model",
        default=None,
        help=f"Defaults to GEMINI_LIVE_MODEL or {DEFAULT_LIVE_MODEL}.",
    )
    parser.add_argument(
        "--save-live-audio",
        type=Path,
        metavar="PATH",
        help="Optionally save returned raw 24 kHz PCM audio.",
    )
    args = parser.parse_args()

    settings = get_settings()
    api_key = settings.gemini_api_key.strip()
    if not api_key:
        print("FAILED: GEMINI_API_KEY is missing from .env.")
        return 1

    live_model = (
        args.live_model
        or getattr(settings, "gemini_live_model", None)
        or DEFAULT_LIVE_MODEL
    )
    print("Gemini credentials loaded from .env (key value hidden).")
    print(f"Key length: {len(api_key)} characters")

    client = genai.Client(api_key=api_key)
    try:
        standard_ok = check_standard_response(
            client, settings.gemini_model, args.prompt
        )
        live_ok = True
        if not args.skip_live:
            live_ok = await check_live_response(
                client, live_model, args.save_live_audio
            )

        print_heading("FINAL RESULT")
        print(f"Standard API: {'PASSED' if standard_ok else 'FAILED'}")
        print(
            f"Live API    : "
            f"{'SKIPPED' if args.skip_live else ('PASSED' if live_ok else 'FAILED')}"
        )
        return 0 if standard_ok and live_ok else 1
    finally:
        await client.aio.aclose()
        client.close()


if __name__ == "__main__":
    try:
        raise SystemExit(asyncio.run(main()))
    except KeyboardInterrupt:
        print("\nCancelled.")
        sys.exit(130)