Google

Google Generative AI (Gemini) integration for the EazyML GenAI framework.

This class wraps the Gemini API using the google.generativeai library, handling model configuration, prompt generation, token counting, and output parsing.

class GoogleGM(model='gemini-2.0-flash-lite', api_key=None, **kwargs)

Bases: GenerativeModel

Google Generative AI (Gemini) integration for the EazyML GenAI framework. This class wraps the Gemini API using the google.generativeai library, handling model configuration, prompt generation, token counting, and output parsing.

Args:

model_name (str, optional): The name of the Gemini model to use. Defaults to “gemini-2.0-flash”.

api_key (str, optional): Google API key for authentication. If not provided, it falls back to the ‘GEMINI_API_KEY’ environment variable.

kwargs: Optional parameters for custom safety settings and generation configuration.

  • show_token_details (bool, optional): Whether to display input and output token counts. Defaults to False.

  • temperature (float, optional): Controls the randomness of the output. Lower values (e.g., 0) make the output more deterministic. Defaults to 0.

  • top_k (int, optional): Considers the top ‘k’ most likely tokens for the next token selection. Defaults to 40.

  • top_p (float, `optional): Nucleus sampling parameter. Considers the smallest set of tokens whose cumulative probability exceeds ‘p’. Defaults to 1.0.

  • candidate_count (int, optional): The number of candidate responses to generate. Defaults to 1.

Example:
from eazyml_genai.components import GoogleGM
# initialize google generative model
google_gm = GoogleGM(model="gemini-2.0-flash",
         api_key=os.getenv('GEMINI_API_KEY'))

# response from generative model given question and retrieved documents
response, input_tokens, output_tokens = google_gm.predict(question=question,
                            payloads=payloads,
                            show_token_details=True
                            )

# parse google response to simple text format
parsed_response = google_gm.parse(response=response)
predict(question, payloads, **kwargs)

Generates a response from the Gemini model based on the provided prompt.

Args:

prompt (str): Input text or prompt for the model.

kwargs (dict):

  • safety_settings (dict, optional): Dictionary defining harm category thresholds.

  • generation_config (GenerationConfig, optional): Configuration object for generation behavior.

  • tool_config (dict, optional): Tool configuration for function calling.

  • tools (list, optional): Tools to be used with function calling.

Returns:

response: The response object containing generated candidates.

parse(response)

Extracts the text content from the first candidate and part of the response.

Args:

response: The response object, assumed to have a structure containing candidates, content, and parts. The exact type of ‘response’ is not specified, but it should behave like a nested list/object as shown in the return description.

Returns:

(str): The text content located at response.candidates[0].content.parts[0].text. Returns the extracted text as a string.

count_tokens(prompt)

Counts the number of tokens in the given prompt.

Args:

prompt (str): The input prompt string.

Returns:

dict: Token usage details.