feat: keep last k images

src/blackbox/vlms.py

@@ -12,13 +12,13 @@ import requests
 import base64
 import copy
 import ast
+import random
+from time import time
 
 import io 
 from PIL import Image
 from lmdeploy.serve.openai.api_client import APIClient
-import io 
-from PIL import Image
-from lmdeploy.serve.openai.api_client import APIClient
+
 
 def is_base64(value) -> bool:
     try:
@@ -50,8 +50,8 @@ class VLMS(Blackbox):
         - ignore_eos (bool): indicator for ignoring eos
         - skip_special_tokens (bool): Whether or not to remove special tokens
             in the decoding. Default to be True."""
-        self.url = vlm_config.url
-
+        self.model_dict = vlm_config.urls
+        self.model_url = None
         self.temperature: float = 0.7
         self.top_p:float = 1
         self.max_tokens: (int |None) = 512
@@ -81,7 +81,7 @@ class VLMS(Blackbox):
         data = args[0]
         return isinstance(data, list)
 
-    def processing(self, prompt:str, images:str | bytes, settings: dict, model_name: Optional[str] = None, user_context: List[dict] = None) -> str:
+    def processing(self, prompt:str | None, images:str | bytes | None, settings: dict, model_name: Optional[str] = None, user_context: List[dict] = None) -> str:
         """
         Args:
             prompt: a string query to the model.
@@ -105,6 +105,9 @@ class VLMS(Blackbox):
         else:
             settings = {}
 
+        if not prompt:
+            prompt = '你是一个辅助机器人，请就此图做一个简短的概括性描述，包括图中的主体物品及状态，不超过50字。' if images else '你好'
+
         # Transform the images into base64 format where openai format need.
         if images:
             if is_base64(images):                               # image as base64 str
@@ -148,7 +151,11 @@ class VLMS(Blackbox):
             #         'content': '图片中主要展示了一只老虎，它正在绿色的草地上休息。草地上有很多可以让人坐下的地方，而且看起来相当茂盛。背景比较模糊，可能是因为老虎的影响，让整个图片的其他部分都变得不太清晰了。'
             #         }
             # ]
-        api_client = APIClient(self.url)
+
+        user_context = self.keep_last_k_images(user_context,k = 1)
+        if self.model_url is None: self.model_url = self._get_model_url(model_name) 
+ 
+        api_client = APIClient(self.model_url)
         # api_client = APIClient("http://10.6.80.91:23333")
         model_name = api_client.available_models[0]
         # Reformat input into openai format to request.
@@ -198,10 +205,28 @@ class VLMS(Blackbox):
             print(item["choices"][0]["message"]['content'])
             responses += item["choices"][0]["message"]['content']
             # total_token_usage += item['usage']['total_tokens']          # 'usage': {'prompt_tokens': *, 'total_tokens': *, 'completion_tokens': *}
-
         user_context = messages + [{'role': 'assistant', 'content': responses}]
+        self.custom_print(user_context)
         return responses, user_context
 
+    def _get_model_url(self,model_name:str | None):
+        available_models = {}
+        for model, url in self.model_dict.items():
+            try:
+                response = requests.get(url,timeout=3)
+                if response.status_code == 200:
+                    available_models[model] = url
+            except Exception as e:
+                # print(e)
+                pass
+        if not available_models: print("There are no available running models and please check your endpoint urls.")
+        if model_name and model_name in available_models: 
+            return available_models[model_name]
+        else:
+            model = random.choice(list(available_models.keys()))
+            print(f"No such model {model_name}, using {model} instead.") if model_name else print(f"Using random model {model}.")
+            return available_models[model]
+        
     def _into_openai_format(self, context:List[list]) -> List[dict]:
         """
         Convert the data into openai format.
@@ -255,7 +280,35 @@ class VLMS(Blackbox):
             
         return user_context
     
+    def keep_last_k_images(self, user_context: list, k:int=2):
+        count = 0
+        result =[]
+        for item in user_context[::-1]:
+            if item['role'] == 'user' and len(item['content']) > 1:
+                for idx, info in enumerate(item['content']):
+                    if info['type'] in ('image_url','image') and count >= k:
+                        item['content'].pop(idx)
+                        # item['content'].insert(idx, {'type': 'text', 'text': '<IMAGE>'})
+                    elif info['type'] in ('image_url','image') and count < k:
+                        count += 1
+                    else:
+                        continue
+            result.append(item)
+        return result[::-1]
 
+
+    def custom_print(self, user_context: list):
+        result = []
+        for item in user_context:
+            if item['role'] == 'user':
+                for idx, info in enumerate(item['content']):
+                    if info['type'] in ('image_url','image'):
+                        item['content'].pop(idx)
+                        item['content'].insert(idx, {'type': 'image', 'image': '##<IMAGE>##'})
+                    else:
+                        continue
+            result.append(item)
+        print(result)
     async def fast_api_handler(self, request: Request) -> Response:
         ## TODO: add support for multiple images and support image in form-data format
         json_request = True
@@ -291,8 +344,8 @@ class VLMS(Blackbox):
         if prompt is None:
             return JSONResponse(content={'error': "Question is required"}, status_code=status.HTTP_400_BAD_REQUEST)
 
-        if model_name is None or model_name.isspace():
-            model_name = "Qwen-VL-Chat"
+        # if model_name is None or model_name.isspace():
+        #     model_name = "Qwen-VL-Chat"
 
         response, history = self.processing(prompt, img_data,settings, model_name,user_context=user_context)
         # jsonresp = str(JSONResponse(content={"response": self.processing(prompt, img_data, model_name)}).body, "utf-8")

src/configuration.py

@@ -153,4 +153,4 @@ class VLMConf():
 
     @inject
     def __init__(self, config: Configuration) -> None:
-        self.url = config.get("vlms.url")
+        self.urls = config.get("vlms.urls")