Gemini LLM Wrapper

This guide demonstrates how to create a wrapper for Google's Gemini Large Language Model (LLM). It allows you to interact with the model via the command line, supporting both single-prompt queries and interactive chat sessions.

Modules Used:

google-generativeai: The official Python SDK for the Gemini API.
python-dotenv: To securely manage the API key.
argparse: To handle command-line arguments.

Prerequisites

Get an API Key: Go to Google AI Studio and create an API key.

Installation

pip install google-generativeai python-dotenv

Setup

Create a .env file in your project directory:

GOOGLE_API_KEY=your_actual_api_key_here

The Code

Save this as gemini_tool.py.

import os
import argparse
import google.generativeai as genai
from dotenv import load_dotenv

# 1. Configuration
load_dotenv()
API_KEY = os.getenv('GOOGLE_API_KEY')

def configure():
    if not API_KEY:
        print("Error: GOOGLE_API_KEY not found in .env file.")
        return False
    genai.configure(api_key=API_KEY)
    return True

def chat_mode(model_name):
    """Interactive chat session with history."""
    try:
        model = genai.GenerativeModel(model_name)
        chat = model.start_chat(history=[])

        print(f"--- Chatting with {model_name} (Type 'quit' to exit) ---")

        while True:
            user_input = input("\nYou: ")
            if user_input.lower() in ['quit', 'exit']:
                print("Exiting chat.")
                break

            if not user_input.strip():
                continue

            print("Gemini: ", end="", flush=True)

            # Stream the response for a better experience
            try:
                response = chat.send_message(user_input, stream=True)
                for chunk in response:
                    print(chunk.text, end="", flush=True)
                print() # Newline after response
            except Exception as e:
                print(f"\nError: {e}")

    except Exception as e:
        print(f"Failed to start chat: {e}")

def single_query(model_name, prompt):
    """Send a single prompt and print the response."""
    try:
        model = genai.GenerativeModel(model_name)
        response = model.generate_content(prompt)
        print(response.text)
    except Exception as e:
        print(f"Error: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Gemini LLM CLI Wrapper")
    parser.add_argument("prompt", nargs="?", help="The prompt to send (leave empty for chat mode)")
    parser.add_argument("--model", default="gemini-1.5-flash", help="Model to use (default: gemini-1.5-flash)")

    args = parser.parse_args()

    if configure():
        if args.prompt:
            single_query(args.model, args.prompt)
        else:
            chat_mode(args.model)

Usage

Interactive Chat: Run without arguments to start a conversation.
```
python gemini_tool.py
```
Single Prompt: Provide a prompt argument for quick answers (useful for piping to other tools).
```
python gemini_tool.py "Explain quantum computing in one sentence"
```
Change Model: Use a different model version if needed.
```
python gemini_tool.py --model gemini-1.5-pro
```

programming/python/python

Created: March 17, 2026 • Last Modified: January 28, 2026

Docs Login