Update core/refiner/simple_refiner.py

core/refiner/simple_refiner.py

@@ -1,6 +1,8 @@
 from models import BaseModel
 from .base_refiner import BaseRefiner
 from utils.image_encoder import encode_image
+import asyncio
+
 class SimpleRefiner(BaseRefiner):
     def __init__(self,
                  sys_prompt: str, 
@@ -8,6 +10,30 @@ class SimpleRefiner(BaseRefiner):
         ) -> None:
         BaseRefiner.__init__(self, sys_prompt=sys_prompt, model=model)
     
+    async def refine_async(self, message: str, memory, image_paths=None) -> str:
+        if memory is None:
+            memory = []
+        else:
+            memory = memory.messages[1:]
+        
+        user_context = [{"role": "user", "content": [
+                        {"type": "text", "text": f"{message}"},]}]
+        if image_paths:
+            if not isinstance(image_paths, list):
+                image_paths = [image_paths]
+            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}"},
+                        ]}]
+        else:
+            context = [{"role": "system", "content": self.sys_prompt}] + memory + user_context
+
+        respond_task = asyncio.create_task(self.model.respond_async(context))
+        await respond_task
+        response = respond_task.result()
+        return response
+    
     def refine(self, message: str, memory, image_paths=None) -> str:
         if memory is None:
             memory = []