File size: 6,331 Bytes
58d133e |
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 |
import google.generativeai as genai
class GoogleGeminiWrapper:
def __init__(self, api_key: str):
"""
Initialize the GoogleGeminiWrapper with the API key.
:param api_key: Your Google Gemini API key.
"""
self.api_key = api_key
genai.configure(api_key=self.api_key)
self.conversation_history = [] # For the new chat method
self.chat_session = None # To store the chat session for gemini
def ask(self, prompt: str, model: str = "gemini-2.0-flash", max_tokens: int = 150, temperature: float = 0.7) -> str:
"""
Send a prompt to the Google Gemini model and get a response (single turn).
:param prompt: The input prompt to send to the model.
:param model: The model to use (default is "gemini-pro").
:param max_tokens: The maximum number of tokens to include in the response.
:param temperature: Sampling temperature (higher values mean more randomness).
:return: The response from the model as a string.
"""
try:
generation_config = {
"temperature": temperature,
"max_output_tokens": max_tokens,
}
model_instance = genai.GenerativeModel(model_name=model, generation_config=generation_config)
response = model_instance.generate_content(prompt)
return response.text.strip()
except Exception as e:
return f"An error occurred: {e}"
def start_chat_session(self, model: str = "gemini-2.0-flash", temperature: float = 0.7, max_tokens: int = 150):
"""
Starts a new chat session or continues an existing one.
"""
generation_config = {
"temperature": temperature,
"max_output_tokens": max_tokens,
}
model_instance = genai.GenerativeModel(model_name=model, generation_config=generation_config)
# For Gemini, conversation history is managed by the chat object itself.
# We re-initialize the chat session if one doesn't exist or if we want to start fresh.
# If you want to persist history across calls to `chat` without explicitly calling reset,
# you might initialize `self.chat_session` in `__init__` or when `chat` is first called.
self.chat_session = model_instance.start_chat(history=self.conversation_history)
def chat(self, prompt: str, model: str = "gemini-2.0-flash", max_tokens: int = 150, temperature: float = 0.7) -> str:
"""
Send a prompt to the Google Gemini model, maintaining conversation history for context.
:param prompt: The input prompt to send to the model.
:param model: The model to use (default is "gemini-pro").
:param max_tokens: The maximum number of tokens to include in the response.
:param temperature: Sampling temperature (higher values mean more randomness).
:return: The response from the model as a string.
"""
try:
if self.chat_session is None:
self.start_chat_session(model=model, temperature=temperature, max_tokens=max_tokens)
response = self.chat_session.send_message(prompt)
assistant_response = response.text.strip()
# Gemini's chat session object updates its history internally.
# We can optionally also store it in our self.conversation_history if needed for other purposes
# or if we want to be able to reconstruct the chat session later.
# For simplicity here, we rely on the chat_session's internal history.
# To manually track:
# self.conversation_history.append({"role": "user", "parts": [prompt]})
# self.conversation_history.append({"role": "model", "parts": [assistant_response]})
return assistant_response
except Exception as e:
# Reset chat session on error to avoid issues with subsequent calls
self.chat_session = None
return f"An error occurred: {e}"
def reset_conversation(self):
"""
Reset the conversation history and the chat session.
"""
self.conversation_history = []
self.chat_session = None # Crucial for Gemini to start a fresh chat
def list_available_models(self):
"""
Lists available Gemini models.
:return: A list of available models.
"""
try:
print("Available Gemini Models:")
for m in genai.list_models():
if 'generateContent' in m.supported_generation_methods:
print(m.name)
return [m.name for m in genai.list_models() if 'generateContent' in m.supported_generation_methods]
except Exception as e:
return f"An error occurred while listing models: {e}"
# Example usage (uncomment to test):
if __name__ == "__main__":
api_key = "AIzaSyBisxoehBz8UF0i9kX42f1V3jp-9RNq04g" # Replace with your actual key
wrapper = GoogleGeminiWrapper(api_key)
# Example 0: List available models
# print("\nListing available models...")
# available_models = wrapper.list_available_models()
# The function already prints, but you can use the returned list if needed
# print(available_models)
#
# Example 1: Simple one-off question
response_ask = wrapper.ask("What is the largest planet in our solar system?")
print(f"Ask response: {response_ask}")
#
# # Example 2: Conversation with history
# print("\nStarting chat conversation...")
# response1 = wrapper.chat("Hi, my name is Alex.")
# print(f"Chat response 1: {response1}")
#
# response2 = wrapper.chat("What is my name?")
# print(f"Chat response 2: {response2}") # Should remember "Alex"
#
# response3 = wrapper.chat("What was the first thing I asked you in this chat?")
# print(f"Chat response 3: {response3}")
#
# # Reset conversation history
# wrapper.reset_conversation()
# print("\nConversation reset.")
#
# response4 = wrapper.chat("Do you remember my name?")
# print(f"Chat response 4 (after reset): {response4}") # Should not remember "Alex"
|