Spaces:
Sleeping
Sleeping
| import google.generativeai as genai | |
| from transformers.agents.llm_engine import MessageRole, get_clean_message_list | |
| from typing import List, Dict | |
| # ... (gemini_role_conversions from your task1.py) | |
| # Role conversion mapping for Gemini | |
| gemini_role_conversions = { | |
| MessageRole.ASSISTANT: "model", | |
| MessageRole.USER: "user", | |
| MessageRole.SYSTEM: "user", # Gemini doesn't have a system role, prepend to user | |
| MessageRole.TOOL_RESPONSE: "user", | |
| } | |
| class GeminiLLM: #renamed | |
| def __init__(self, gemini_key: str, model_name="gemini-2.0-flash-exp"): | |
| genai.configure(api_key=gemini_key) | |
| self.model = genai.GenerativeModel(model_name) | |
| def format_messages(self, messages: List[Dict]) -> List[Dict]: #renamed | |
| # ... (Your _convert_messages logic from task1.py) | |
| cleaned_messages = get_clean_message_list(messages, role_conversions=gemini_role_conversions) | |
| # Handle system messages by prepending to first user message | |
| formatted_messages = [] | |
| system_content = "" | |
| for msg in cleaned_messages: | |
| if msg["role"] == "user" and msg.get("content", "").startswith("System:"): | |
| system_content += msg["content"].replace("System:", "").strip() + "\n" | |
| else: | |
| if system_content and msg["role"] == "user": | |
| msg["content"] = f"{system_content}\n{msg['content']}" | |
| system_content = "" | |
| formatted_messages.append({ | |
| "role": msg["role"], | |
| "parts": [msg["content"]] | |
| }) | |
| return formatted_messages | |
| def generate_response(self, messages: List[Dict], stop_sequences: List[str] = None) -> str: #renamed | |
| # ... (Your Gemini generation logic from task1.py) | |
| """ | |
| Generate a response using the Gemini model | |
| Args: | |
| messages (List[Dict]): List of message dictionaries | |
| stop_sequences (List[str], optional): List of sequences to stop generation | |
| Returns: | |
| str: Generated response | |
| """ | |
| formatted_messages = self.format_messages(messages) | |
| # Create chat session | |
| chat = self.model.start_chat(history=formatted_messages) | |
| # Generate response with safety settings and parameters | |
| response = chat.send_message( | |
| formatted_messages[-1]["parts"][0], | |
| generation_config={ | |
| "temperature": 0, | |
| "max_output_tokens": 4096, | |
| "stop_sequences": stop_sequences if stop_sequences else [] | |
| } | |
| ) | |
| return response.text |