Implement compose_message_openai function

Flagro · Dec 9, 2024 · 71db4b6 · 71db4b6
1 parent 0be0825
commit 71db4b6
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/bot/rp_bot/ai_agent/models_toolkit.py b/bot/rp_bot/ai_agent/models_toolkit.py
@@ -91,13 +91,19 @@ def get_price(
             + image_generation_pixels * output_pixel_price
         )
 
+    @staticmethod
+    def compose_message_openai(
+        message_text: str, role: Literal["user", "system"] = "user"
+    ) -> Dict[str, str]:
+        return {"role": role, "content": message_text}
+
     @staticmethod
     def compose_messages_openai(
         user_input: str, system_prompt: str
     ) -> List[Dict[str, str]]:
         return [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": user_input},
+            ModelsToolkit.compose_message_openai(system_prompt, role="system"),
+            ModelsToolkit.compose_message_openai(user_input, role="user"),
         ]
 
     async def get_response(self, question: str) -> str: