app.py

import os
import time
import google.generativeai as genai
from google.generativeai.types import Content, Part, GenerateContentConfig

def get_client():
    api_key = "AIzaSyCZVL-HKieiZ8-uOjEvfpxU1eqTr2NzySs"
    if not api_key:
        raise ValueError("API key is not set.")
    return genai.Client(api_key=api_key)

def generate(user_input):
    client = get_client()
    model = "gemini-2.0-flash"

    contents = [
        types.Content(
            role="user",
            parts=[types.Part.from_text(text="hello")],
        ),
        types.Content(
            role="model",
            parts=[types.Part.from_text(text="""Hey, you are very helpful assistant and talk like Luke Skywalker""")],
        ),
        types.Content(
            role="user",
            parts=[types.Part.from_text(text=user_input)],
        ),
    ]

    generate_content_config = types.GenerateContentConfig(
        response_mime_type="text/plain",
    )

    start_time = time.time()
    for chunk in client.models.generate_content_stream(
        model=model,
        contents=contents,
        config=generate_content_config,
    ):
        print(chunk.text, end="", flush=True)
    end_time = time.time()

    latency = end_time - start_time
    print(f"\nLatency: {latency:.2f} seconds")

if __name__ == "__main__":
    user_input = input("Enter your prompt: ")
    generate(user_input)