Merge pull request #32 from BoardWare-Genius/ivan

fix: to support the vlm compatible with non streaming output

src/blackbox/vlms.py

@@ -58,7 +58,6 @@ class VLMS(Blackbox):
         - skip_special_tokens (bool): Whether or not to remove special tokens
             in the decoding. Default to be True."""
         self.model_dict = vlm_config.urls
-        # self.model_url = None
         self.available_models = {}
         self.temperature: float = 0.7
         self.top_p:float = 1
@@ -87,7 +86,6 @@ class VLMS(Blackbox):
                 if response.status_code == 200:
                     self.available_models[model] = url
             except Exception as e:
-                # print(e)
                 pass
     def __call__(self, *args, **kwargs):
         return self.processing(*args, **kwargs)
@@ -96,7 +94,7 @@ class VLMS(Blackbox):
         data = args[0]
         return isinstance(data, list)
 
-    def processing(self, prompt:str | None, images:str | bytes | None, 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, user_context: List[dict] = None) -> str:
         """
         Args:
             prompt: a string query to the model.
@@ -130,8 +128,7 @@ class VLMS(Blackbox):
             prompt = '你是一个辅助机器人，请就此图做一个简短的概括性描述，包括图中的主体物品及状态，不超过50字。' if images else '你好'
 
         # Transform the images into base64 format where openai url)
-        # print(self.config['vlm_model_name'])
-        # print(self.available_models)format need.
+
         if images:
             if is_base64(images):                               # image as base64 str
                 images_data = images    
@@ -144,38 +141,13 @@ class VLMS(Blackbox):
                 images_data = str(base64.b64encode(res.content),'utf-8')
         else:
             images_data = None
-        ## AutoLoad Model
-        # url = 'http://10.6.80.87:8000/' + model_name + '/'
-        # data_input = {'model': model_name, 'prompt': prompt, 'img_data': images_data}
-        # data = requests.post(url, json=data_input)
-        # return data.text
 
-        # 'https://raw.githubusercontent.com/open-mmlab/mmdeploy/main/tests/data/tiger.jpeg'
-        ## Lmdeploy
-        # if not user_context:
-        #     user_context = []
 
         ## Predefine user_context only for testing
         # user_context = [{'role':'user','content':'你好，我叫康康，你是谁？'}, {'role': 'assistant', 'content': '你好！很高兴为你提供帮助。'}]
-            # user_context = [{
-            #         'role': 'user',
-            #         'content': [{
-            #             'type': 'text',
-            #             'text': '图中有什么，请描述一下',
-            #         }, {
-            #             'type': 'image_url',
-            #             'image_url': {
-            #                 'url': 'https://raw.githubusercontent.com/open-mmlab/mmdeploy/main/tests/data/tiger.jpeg'
-            #             },
-            #         }]
-            #     },{
-            #         'role': 'assistant',
-            #         'content': '图片中主要展示了一只老虎，它正在绿色的草地上休息。草地上有很多可以让人坐下的地方，而且看起来相当茂盛。背景比较模糊，可能是因为老虎的影响，让整个图片的其他部分都变得不太清晰了。'
-            #         }
-            # ]
+
         if not user_context and config['system_prompt']: user_context = [{'role':'system','content': config['system_prompt']}] 
-        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) 
+        user_context = self.keep_last_k_images(user_context,k = 2)
 
         # Reformat input into openai format to request.
         if images_data:
@@ -188,8 +160,6 @@ class VLMS(Blackbox):
                         'type': 'image_url',
                         'image_url': {      # Image two
                             'url': 
-                            # 'https://raw.githubusercontent.com/open-mmlab/mmdeploy/main/tests/data/tiger.jpeg'
-                            # './val_data/image_5.jpg'
                             f"data:image/jpeg;base64,{images_data}",
                         },
                     # },{                 # Image one
@@ -213,9 +183,7 @@ class VLMS(Blackbox):
         responses = ''
         total_token_usage = 0   # which can be used to count the cost of a query 
         model_url = self._get_model_url(config['vlm_model_name'])
-        # print(model_url)
-        # print(self.config['vlm_model_name'])
-        # print(self.available_models)
+
         if config['lmdeploy_infer']:
             api_client = APIClient(model_url)
             model_name = api_client.available_models[0]
@@ -225,7 +193,6 @@ class VLMS(Blackbox):
                                                 # session_id=,
                                                 )):
                 # Stream output
-                # print(item["choices"][0]["delta"]['content'],end='\n')
                 yield item["choices"][0]["delta"]['content']
                 responses += item["choices"][0]["delta"]['content']
 
@@ -234,7 +201,6 @@ class VLMS(Blackbox):
                 # total_token_usage += item['usage']['total_tokens']          # 'usage': {'prompt_tokens': *, 'total_tokens': *, 'completion_tokens': *}
         else:
             api_key = "EMPTY_API_KEY"
-            # print(model_url+'/v1')
             api_client = OpenAI(api_key=api_key, base_url=model_url+'/v1')
             model_name = api_client.models.list().data[0].id
             for item in api_client.chat.completions.create(
@@ -256,7 +222,7 @@ class VLMS(Blackbox):
         
         user_context = messages + [{'role': 'assistant', 'content': responses}]
         self.custom_print(user_context)
-        # return responses, user_context
+        # return responses
 
     def _get_model_url(self,model_name:str | None):
         if not self.available_models: print("There are no available running models and please check your endpoint urls.")
@@ -336,7 +302,6 @@ class VLMS(Blackbox):
             result.append(item)
         return result[::-1]
 
-
     def custom_print(self, user_context: list):
         result = []
         for item in user_context:
@@ -354,7 +319,6 @@ class VLMS(Blackbox):
         json_request = True
         try:
             content_type = request.headers.get('content-type', '')
-            print(content_type)
             if content_type == 'application/json':
                 data = await request.json()
             elif 'multipart/form-data' in content_type:
@@ -367,9 +331,10 @@ class VLMS(Blackbox):
         except Exception as e:
             return JSONResponse(content={"error": "json parse error"}, status_code=status.HTTP_400_BAD_REQUEST)    
         
-        model_name = data.get("model_name")
+
         prompt = data.get("prompt")
         settings: dict = data.get('settings') 
+
         context = data.get("context")
         if not context:
             user_context = []
@@ -388,14 +353,12 @@ 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"
-        # response,_ = self.processing(prompt, img_data,settings, model_name,user_context=user_context)
-
-        # return StreamingResponse(self.processing(prompt, img_data,settings, model_name,user_context=user_context), status_code=status.HTTP_200_OK)
-        return EventSourceResponse(self.processing(prompt, img_data,settings, model_name,user_context=user_context), status_code=status.HTTP_200_OK)
-    
+        streaming_output = str(settings.get('stream',False)).strip().lower() == 'true' if settings else False
+        if streaming_output:
+            # return StreamingResponse(self.processing(prompt, img_data,settings, user_context=user_context), status_code=status.HTTP_200_OK)
+            return EventSourceResponse(self.processing(prompt, img_data,settings, user_context=user_context), status_code=status.HTTP_200_OK)
+        else:
             # HTTP JsonResponse
-        response, history = self.processing(prompt, img_data,settings, model_name,user_context=user_context)
-        # return JSONResponse(content={"response": response}, status_code=status.HTTP_200_OK)
+            output = self.processing(prompt, img_data,settings, user_context=user_context)
+            response = ''.join([res for res in output])
+            return JSONResponse(content={"response": response}, status_code=status.HTTP_200_OK)