Update core/refiner/simple_refiner.py

core/refiner/simple_refiner.py

@@ -8,20 +8,23 @@ class SimpleRefiner(BaseRefiner):
         ) -> None:
         BaseRefiner.__init__(self, sys_prompt=sys_prompt, model=model)
     
-    def refine(self, message: str, memory, image_path=None) -> str:
+    def refine(self, message: str, memory, image_paths=None) -> str:
         if memory is None:
             memory = []
         else:
             memory = memory.messages[1:]
-        if image_path:
+        if image_paths:
+            if not isinstance(image_paths, list):
+                image_paths = [image_paths]
+            user_context = [{"role": "user", "content": [
+                            {"type": "text", "text": f"{message}"},]}]
+            for image_path in image_paths:
+                user_context.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image_path.name)}"}})
             context = [{"role": "system", "content": self.sys_prompt}] + memory + [{"role": "user", "content": [
                             {"type": "text", "text": f"{message}"},
-                            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image_path)}"}},
                         ]}]
         else:
-            context = [{"role": "system", "content": self.sys_prompt}] + memory + [{"role": "user", "content": [
-                            {"type": "text", "text": f"{message}"},
-                        ]}]
+            context = [{"role": "system", "content": self.sys_prompt}] + memory + user_context
         response = self.model.respond(context)
 
         return response