| from .BaseLLM import BaseLLM | |
| from google import genai | |
| import os | |
| import time | |
| class Gemini(BaseLLM): | |
| def __init__(self, model="gemini-2.0-flash"): | |
| super(Gemini, self).__init__() | |
| self.model_name = model | |
| self.messages = [] | |
| def initialize_message(self): | |
| self.messages = [] | |
| def ai_message(self, payload): | |
| self.messages.append(payload) | |
| def system_message(self, payload): | |
| self.messages.append(payload) | |
| def user_message(self, payload): | |
| self.messages.append(payload) | |
| def get_response(self,temperature = 0.8): | |
| response = genai.Client(api_key=os.getenv("GEMINI_API_KEY")).models.generate_content( | |
| model=self.model_name, contents="".join(self.messages), temperature = temperature | |
| ) | |
| return response.text | |
| def chat(self,text): | |
| self.initialize_message() | |
| self.user_message(text) | |
| response = self.get_response() | |
| return response | |
| def print_prompt(self): | |
| for message in self.messages: | |
| print(message) |