Create llm.py

RAG/llm.py

@@ -0,0 +1,67 @@
+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