File size: 6,049 Bytes
ad9da03 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 |
import logging
import os
from typing import Dict, Iterator, List, Optional, Union
import requests
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
# Setup logging
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
# Gemini API Configuration
API_KEY_ENV_VAR = "GOOGLE_API_KEY"
BASE_URL = "https://generativelanguage.googleapis.com/v1beta/models/"
DEFAULT_MODEL_ID = "gemini-2.0-flash"
def _get_api_key() -> Optional[str]:
"""
Retrieves the Google API key from environment variables.
Returns:
Optional[str]: The API key if found, otherwise None.
"""
api_key = os.getenv(API_KEY_ENV_VAR)
if not api_key:
logger.error(f"API key not found. Set the variable '{API_KEY_ENV_VAR}'.")
return api_key
def _format_payload_for_gemini(
messages: List[Dict], temperature: float, max_tokens: int
) -> Optional[Dict]:
"""
Formats the message history and configuration into a valid payload for the Gemini REST API.
This function performs two critical tasks:
1. Separates the 'system' instruction from the main conversation history.
2. Consolidates consecutive 'user' messages into a single block to comply with
the Gemini API's requirement of alternating 'user' and 'model' roles.
Args:
messages (List[Dict]): A list of message dictionaries, potentially including a 'system' role.
temperature (float): The generation temperature.
max_tokens (int): The maximum number of tokens to generate.
Returns:
Optional[Dict]: A fully formed payload dictionary ready for the API, or None if the
conversation history is empty.
"""
system_instruction = None
conversation_history = []
for msg in messages:
if msg.get("role") == "system":
system_instruction = {"parts": [{"text": msg.get("content", "")}]}
else:
conversation_history.append(msg)
if not conversation_history:
return None
consolidated_contents = []
current_block = None
for msg in conversation_history:
role = "model" if msg.get("role") == "assistant" else "user"
content = msg.get("content", "")
if current_block and current_block["role"] == "user" and role == "user":
current_block["parts"][0]["text"] += "\n" + content
else:
if current_block:
consolidated_contents.append(current_block)
current_block = {"role": role, "parts": [{"text": content}]}
if current_block:
consolidated_contents.append(current_block)
payload = {
"contents": consolidated_contents,
"safetySettings": [
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
],
"generationConfig": {"temperature": temperature, "maxOutputTokens": max_tokens},
}
if system_instruction:
payload["system_instruction"] = system_instruction
return payload
def call_gemini_api(
messages: List[Dict], stream: bool = False, temperature: float = 0.7, max_tokens: int = 2048
) -> Union[Iterator[str], str]:
"""
Calls the Google Gemini REST API with the provided messages and parameters.
This is the main public function of the module. It handles API key retrieval,
payload formatting, making the HTTP request, and processing the response.
Args:
messages (List[Dict]): The list of messages forming the conversation context.
stream (bool): If True, streams the response. (Currently not implemented).
temperature (float): The generation temperature (creativity).
max_tokens (int): The maximum number of tokens for the response.
Returns:
Union[Iterator[str], str]: An iterator of response chunks if streaming, or a single
response string otherwise. Returns an error string on failure.
"""
api_key = _get_api_key()
if not api_key:
error_msg = "Error: Google API key not configured."
return iter([error_msg]) if stream else error_msg
payload = _format_payload_for_gemini(messages, temperature, max_tokens)
if not payload or not payload.get("contents"):
error_msg = "Error: Conversation is empty or malformed after processing."
return iter([error_msg]) if stream else error_msg
stream_param = "streamGenerateContent" if stream else "generateContent"
request_url = f"{BASE_URL}{DEFAULT_MODEL_ID}:{stream_param}?key={api_key}"
headers = {"Content-Type": "application/json"}
try:
response = requests.post(
request_url, headers=headers, json=payload, stream=stream, timeout=180
)
response.raise_for_status()
if stream:
# TODO: Implement robust stream processing logic here.
pass
else:
data = response.json()
# Safely access nested keys to prevent KeyErrors
if data.get("candidates") and data["candidates"][0].get("content", {}).get("parts"):
return data["candidates"][0]["content"]["parts"][0]["text"]
else:
logger.warning(
f"Gemini's response does not contain 'candidates'. Full response: {data}"
)
return "[BLOCKED OR EMPTY RESPONSE]"
except requests.exceptions.HTTPError as e:
err_msg = f"API HTTP Error ({e.response.status_code}): {e.response.text[:500]}"
logger.error(err_msg, exc_info=False)
return f"Error: {err_msg}"
except Exception as e:
logger.exception("Unexpected error while calling Gemini API:")
return f"Error: {e}"
|