Merge pull request #108 from wugren/MVP

AgentMsg support image and video

src/aios/__init__.py

@@ -27,6 +27,7 @@ from .storage.storage import ResourceLocation,AIStorage,UserConfig,UserConfigIte
 from .net import *
 from .knowledge import *
 from .package_manager import *
+from .utils import *
 
 
-AIOS_Version = "0.5.2, build 2023-11-30"
+AIOS_Version = "0.5.2, build 2023-11-30"

src/aios/agent/agent.py

@@ -27,6 +27,7 @@ from ..environment.workspace_env import WorkspaceEnvironment
 from ..storage.storage import AIStorage
 
 from ..knowledge import *
+from ..utils import video_utils, image_utils
 
 logger = logging.getLogger(__name__)
 
@@ -423,11 +424,39 @@ class AIAgent(BaseAIAgent):
     async def _create_openai_thread(self) -> str:
         return None
 
+    def check_and_to_base64(self, image_path: str) -> str:
+        if image_utils.is_file(image_path):
+            return image_utils.to_base64(image_path, (1024, 1024))
+        else:
+            return image_path
+
     async def _process_msg(self,msg:AgentMsg,workspace = None) -> AgentMsg:
         msg_prompt = AgentPrompt()
         if msg.msg_type == AgentMsgType.TYPE_GROUPMSG:
             need_process = False
-            msg_prompt.messages = [{"role":"user","content":f"{msg.sender}:{msg.body}"}]
+            if msg.is_image_msg():
+                image_prompt, images = msg.get_image_body()
+                if image_prompt is None:
+                    content = [[{"type": "text", "text": f"{msg.sender}'s message"}]]
+                    content.extend([{"type": "image_url", "image_url": {"url": self.check_and_to_base64(image)}} for image in images])
+                    msg_prompt.messages = [{"role": "user", "content": content}]
+                else:
+                    content = [{"type": "text", "text": f"{msg.sender}:{image_prompt}"}]
+                    content.extend([{"type": "image_url", "image_url": {"url": self.check_and_to_base64(image)}} for image in images])
+                    msg_prompt.messages = [{"role": "user", "content": content}]
+            elif msg.is_video_msg():
+                video_prompt, video = msg.get_video_body()
+                frames = video_utils.extract_frames(video, (1024, 1024))
+                if video_prompt is None:
+                    content = [{"type": "text", "text": f"{msg.sender}'s message"}]
+                    content.extend([{"type": "image_url", "image_url": {"url": frame}} for frame in frames])
+                    msg_prompt.messages = [{"role": "user", "content": content}]
+                else:
+                    content = [{"type": "text", "text": f"{msg.sender}:{video_prompt}"}]
+                    content.extend([{"type": "image_url", "image_url": {"url": frame}} for frame in frames])
+                    msg_prompt.messages = [{"role": "user", "content": content}]
+            else:
+                msg_prompt.messages = [{"role":"user","content":f"{msg.sender}:{msg.body}"}]
             session_topic = msg.target + "#" + msg.topic
             chatsession = AIChatSession.get_session(self.agent_id,session_topic,self.chat_db)
 
@@ -441,7 +470,25 @@ class AIAgent(BaseAIAgent):
                 resp_msg = msg.create_group_resp_msg(self.agent_id,"")
                 return resp_msg
         else:
-            msg_prompt.messages = [{"role":"user","content":msg.body}]
+            if msg.is_image_msg():
+                image_prompt, images = msg.get_image_body()
+                if image_prompt is None:
+                    msg_prompt.messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": self.check_and_to_base64(image)}} for image in images]}]
+                else:
+                    content = [{"type": "text", "text": image_prompt}]
+                    content.extend([{"type": "image_url", "image_url": {"url": self.check_and_to_base64(image)}} for image in images])
+                    msg_prompt.messages = [{"role": "user", "content": content}]
+            elif msg.is_video_msg():
+                video_prompt, video = msg.get_video_body()
+                frames = video_utils.extract_frames(video, (1024, 1024))
+                if video_prompt is None:
+                    msg_prompt.messages = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": frame}} for frame in frames]}]
+                else:
+                    content = [{"type": "text", "text": video_prompt}]
+                    content.extend([{"type": "image_url", "image_url": {"url": frame}} for frame in frames])
+                    msg_prompt.messages = [{"role": "user", "content": content}]
+            else:
+                msg_prompt.messages = [{"role":"user","content":msg.body}]
             session_topic = msg.get_sender() + "#" + msg.topic
             chatsession = AIChatSession.get_session(self.agent_id,session_topic,self.chat_db)
             if self.enable_thread:

src/aios/agent/agent_base.py

@@ -9,7 +9,7 @@ import time
 import re
 import shlex
 import json
-from typing import List
+from typing import List, Tuple
 
 from .ai_function import FunctionItem, AIFunction
 from ..proto.agent_msg import AgentMsg, AgentMsgType
@@ -410,6 +410,10 @@ class BaseAIAgent(abc.ABC):
     def get_max_token_size(self) -> int:
         pass
 
+    @abstractmethod
+    async def _process_msg(self,msg:AgentMsg,workspace = None) -> AgentMsg:
+        pass
+
     @classmethod
     def get_inner_functions(cls, env:Environment) -> (dict,int):
         if env is None:
@@ -445,10 +449,29 @@ class BaseAIAgent(abc.ABC):
         #logger.debug(f"Agent {self.agent_id} do llm token static system:{system_prompt_len},function:{function_token_len},history:{history_token_len},input:{input_len}, totoal prompt:{system_prompt_len + function_token_len + history_token_len} ")
         if inner_functions is None and env is not None:
             inner_functions,_ = BaseAIAgent.get_inner_functions(env)
+
+        model_name = self.get_llm_model_name()
+        if org_msg.is_video_msg() or org_msg.is_image_msg():
+            if model_name.startswith("gpt-4"):
+                model_name = "gpt-4-vision-preview"
         if is_json_resp:
-            task_result:ComputeTaskResult = await ComputeKernel.get_instance().do_llm_completion(prompt,resp_mode="json",mode_name=self.get_llm_model_name(),max_token=self.get_max_token_size(),inner_functions=inner_functions,timeout=None)
+            task_result: ComputeTaskResult = await (ComputeKernel.get_instance()
+            .do_llm_completion(
+                prompt,
+                resp_mode="json",
+                mode_name=model_name,
+                max_token=self.get_max_token_size(),
+                inner_functions=inner_functions,
+                timeout=None))
         else:
-            task_result:ComputeTaskResult = await ComputeKernel.get_instance().do_llm_completion(prompt,resp_mode="text",mode_name=self.get_llm_model_name(),max_token=self.get_max_token_size(),inner_functions=inner_functions,timeout=None)
+            task_result: ComputeTaskResult = await (ComputeKernel.get_instance()
+            .do_llm_completion(
+                prompt,
+                resp_mode="text",
+                mode_name=model_name,
+                max_token=self.get_max_token_size(),
+                inner_functions=inner_functions,
+                timeout=None))
         if task_result.result_code != ComputeTaskResultCode.OK:
             logger.error(f"_do_llm_complection llm compute error:{task_result.error_str}")
             #error_resp = msg.create_error_resp(task_result.error_str)
@@ -478,7 +501,6 @@ class BaseAIAgent(abc.ABC):
         stack_limit = 5
     ) -> ComputeTaskResult:
         from ..frame.compute_kernel import ComputeKernel
-        
         arguments = None
         try:
             func_name = inner_func_call_node.get("name")

src/aios/knowledge/data/writer.py

@@ -19,10 +19,10 @@ def _join_docs(docs: List[str], separator: str) -> Optional[str]:
         return text
 
 def _merge_splits(
-        splits: Iterable[str], 
-        separator: str, 
-        chunk_size: int, 
-        chunk_overlap: int,  
+        splits: Iterable[str],
+        separator: str,
+        chunk_size: int,
+        chunk_overlap: int,
         length_function: Callable[[str], int]
     ) -> List[str]:
     # We now want to combine these smaller pieces into medium size
@@ -86,11 +86,11 @@ def _split_text_with_regex(
     return [s for s in splits if s != ""]
 
 
-def _split_text(
-        text: str, 
-        separators: List[str], 
-        chunk_size: int, 
-        chunk_overlap: int, 
+def split_text(
+        text: str,
+        separators: List[str],
+        chunk_size: int,
+        chunk_overlap: int,
         length_function: Callable[[str], int]
     ) -> List[str]:
 
@@ -127,7 +127,7 @@ def _split_text(
             if not new_separators:
                 final_chunks.append(s)
             else:
-                other_info = _split_text(s, new_separators, chunk_size, chunk_overlap, length_function)
+                other_info = split_text(s, new_separators, chunk_size, chunk_overlap, length_function)
                 final_chunks.extend(other_info)
     if _good_splits:
         merged_text = _merge_splits(_good_splits, _separator, chunk_size, chunk_overlap, length_function)
@@ -153,7 +153,7 @@ class ChunkListWriter:
                 chunk = file.read(chunk_size)
                 if not chunk:
                     break
-                
+
                 chunk_len = len(chunk)
                 chunk_id = ChunkID.hash_data(chunk)
                 chunk_list.append(chunk_id)
@@ -176,14 +176,14 @@ class ChunkListWriter:
 
         file_hash = HashValue(hash_obj.digest())
         # print(f"calc file hash: {file_path}, {file_hash}")
-        
+
         return ChunkList(chunk_list, file_hash)
 
     def create_chunk_list_from_text(
-        self, 
-        text: str, 
-        chunk_size: int = 4000, 
-        chunk_overlap: int = 200, 
+        self,
+        text: str,
+        chunk_size: int = 4000,
+        chunk_overlap: int = 200,
         separators: str = ["\n\n", "\n", " ", ""]
     ) -> ChunkList:
         enc = tiktoken.encoding_for_model("gpt-3.5-turbo")
@@ -196,8 +196,8 @@ class ChunkListWriter:
                     disallowed_special="all",
                 )
             )
-        
-        text_list = _split_text(text, separators, chunk_size, chunk_overlap, length_function)
+
+        text_list = split_text(text, separators, chunk_size, chunk_overlap, length_function)
         chunk_list = []
         hash_obj = hashlib.sha256()
 
@@ -211,4 +211,4 @@ class ChunkListWriter:
             self.chunk_store.put_chunk(chunk_id, chunk_bytes)
 
         hash = HashValue(hash_obj.digest())
-        return ChunkList(chunk_list, hash)
+        return ChunkList(chunk_list, hash)

src/aios/proto/agent_msg.py

@@ -1,7 +1,10 @@
+import json
 import logging
+import shlex
 import uuid
 from enum import Enum
 import time
+from typing import Tuple, List
 
 logger = logging.getLogger(__name__)
 
@@ -35,8 +38,6 @@ class AgentMsgStatus(Enum):
 # 逻辑上的同一个Message在同一个session中看到的msgid相同
 #       在不同的session中看到的msgid不同
 
-
-
 class AgentMsg:
     def __init__(self,msg_type=AgentMsgType.TYPE_MSG) -> None:
         self.msg_id = "msg#" + uuid.uuid4().hex
@@ -136,14 +137,79 @@ class AgentMsg:
 
         return resp_msg
 
-    def set(self,sender:str,target:str,body:str,topic:str=None) -> None:
+    def set(self,sender:str,target:str,body:str,topic:str=None,body_mime:str=None) -> None:
         self.sender = sender
         self.target = target
         self.body = body
+        self.body_mime = body_mime
         self.create_time = time.time()
         if topic:
             self.topic = topic
 
+    @staticmethod
+    def create_image_body(images: [str], prompt: str = None):
+        return json.dumps({"images": images, "prompt": prompt})
+
+    @staticmethod
+    def parse_image_body(image_body: str) -> Tuple[str, List[str]]:
+        body = json.loads(image_body)
+        return body.get("prompt"), body.get("images")
+
+    @staticmethod
+    def create_video_body(video: str, prompt: str = None):
+        return json.dumps({"video": video, "prompt": prompt})
+
+    @staticmethod
+    def parse_video_body(video_body: str) -> Tuple[str, str]:
+        body = json.loads(video_body)
+        return body.get("prompt"), body.get("video")
+
+    def set_image(self, sender: str, target: str, image_format: str, images: [str], prompt: str = None, topic: str = None):
+        self.sender = sender
+        self.target = target
+        self.create_time = time.time()
+        self.body_mime = f"image/{image_format}"
+        self.body = self.create_image_body(images, prompt)
+        if topic:
+            self.topic = topic
+
+    def is_image_msg(self) -> bool:
+        if self.body_mime is None:
+            return False
+        if self.body_mime.startswith("image/"):
+            return True
+        return False
+
+    def get_image_body(self) -> Tuple[str, List[str]]:
+        if self.body_mime is None:
+            return None
+        if self.body_mime.startswith("image/"):
+            return self.parse_image_body(self.body)
+        return None
+
+    def set_video(self, sender: str, target: str, video_format: str, video: str, prompt: str = None, topic: str = None):
+        self.sender = sender
+        self.target = target
+        self.create_time = time.time()
+        self.body_mime = f"video/{video_format}"
+        self.body = self.create_video_body(video, prompt)
+        if topic:
+            self.topic = topic
+
+    def get_video_body(self) -> Tuple[str, str]:
+        if self.body_mime is None:
+            return None
+        if self.body_mime.startswith("video/"):
+            return self.parse_video_body(self.body)
+        return None
+
+    def is_video_msg(self) -> bool:
+        if self.body_mime is None:
+            return False
+        if self.body_mime.startswith("video/"):
+            return True
+        return False
+
     def get_msg_id(self) -> str:
         return self.msg_id
 
@@ -164,4 +230,4 @@ class AgentMsg:
         str_list = shlex.split(func_string)
         func_name = str_list[0]
         params = str_list[1:]
-        return func_name, params
+        return func_name, params

src/aios/utils/__init__.py

@@ -0,0 +1,2 @@
+from . import image_utils
+from . import video_utils

src/aios/utils/image_utils.py

@@ -0,0 +1,40 @@
+import base64
+import os.path
+from typing import Tuple
+
+import cv2
+
+
+def to_base64(image_path: str, resize: Tuple[int, int] = None) -> str:
+    """Convert image to base64."""
+    ext = os.path.splitext(image_path)[1][1:]
+    if resize is None:
+        with open(image_path, "rb") as image_file:
+            base64_image = base64.b64encode(image_file.read()).decode("utf-8")
+            return f"data:image/{ext};base64,{base64_image}"
+    else:
+        dest_width, dest_height = resize
+        img = cv2.imread(image_path)
+        width, height = img.shape[:2]
+        if width > dest_width or height > dest_height:
+            width_rate = dest_width / width
+            height_rate = dest_height / height
+            rate = min(width_rate, height_rate)
+            dest_width = int(width * rate)
+            dest_height = int(height * rate)
+            img = cv2.resize(img, (dest_width, dest_height), interpolation=cv2.INTER_AREA)
+        _, buf = cv2.imencode(f".{ext}", img)
+        base64_image = base64.b64encode(buf).decode("utf-8")
+        return f"data:image/{ext};base64,{base64_image}"
+
+
+def is_file(image_path: str) -> bool:
+    return os.path.isfile(image_path)
+
+
+def is_base64(image_path: str) -> bool:
+    return image_path.startswith("data:image/")
+
+
+def is_url(image_path: str) -> bool:
+    return image_path.startswith("http://") or image_path.startswith("https://")

src/aios/utils/video_utils.py

@@ -0,0 +1,122 @@
+import base64
+from typing import List, Tuple
+
+import cv2
+import numpy as np
+
+
+def precess_image(image):
+    '''
+    Graying and GaussianBlur
+    :param image: The image matrix,np.array
+    :return: The processed image matrix,np.array
+    '''
+    gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    gray_image = cv2.GaussianBlur(gray_image, (3, 3), 0)
+    return gray_image
+
+
+def abs_diff(pre_image, curr_image):
+    '''
+    Calculate absolute difference between pre_image and curr_image
+    :param pre_image:The image in past frame,np.array
+    :param curr_image:The image in current frame,np.array
+    :return:
+    '''
+    gray_pre_image = precess_image(pre_image)
+    gray_curr_image = precess_image(curr_image)
+    diff = cv2.absdiff(gray_pre_image, gray_curr_image)
+    res, diff = cv2.threshold(diff, 0, 255, cv2.THRESH_BINARY+cv2.THRESH_OTSU)
+    cnt_diff = np.sum(np.sum(diff))
+    return cnt_diff
+
+
+def exponential_smoothing(alpha, s):
+    '''
+    Primary exponential smoothing
+    :param alpha:  Smoothing factor,num
+    :param s:      List of data,list
+    :return:       List of data after smoothing,list
+    '''
+    s_temp = [s[0]]
+    print(s_temp)
+    for i in range(1, len(s), 1):
+        s_temp.append(alpha * s[i - 1] + (1 - alpha) * s_temp[i - 1])
+    return s_temp
+
+
+def extract_frames(video_path: str, resize: Tuple[int, int] = None, smooth=False, alpha=0.07, window=25) -> List[str]:
+    """Extract frames from video."""
+    frames = []
+    vidcap = cv2.VideoCapture(video_path)
+    diff = []
+    frm = 0
+    pre_image = np.array([])
+    cur_image = np.array([])
+
+    while True:
+        frm = frm + 1
+        success, image = vidcap.read()
+        if not success:
+            break
+
+        if frm == 1:
+            pre_image = image
+            cur_image = image
+        else:
+            pre_image = cur_image
+            cur_image = image
+
+        diff.append(abs_diff(pre_image, cur_image))
+
+    if smooth:
+        diff = exponential_smoothing(alpha, diff)
+
+    diff = np.array(diff)
+    mean = np.mean(diff)
+    dev = np.std(diff)
+    diff = (diff - mean) / dev
+
+    idx = []
+    for i, d in enumerate(diff):
+        ub = len(diff) - 1
+        lb = 0
+        if not i - window // 2 < lb:
+            lb = i - window // 2
+        if not i + window // 2 > ub:
+            ub = i + window // 2
+
+        comp_window = diff[lb: ub]
+        if d >= max(comp_window):
+            idx.append(i)
+
+    tmp = np.array(idx)
+    tmp = tmp + 1
+    idx = set(tmp.tolist())
+    vidcap.release()
+
+    vidcap = cv2.VideoCapture(video_path)
+    i = 0
+    frm = 0
+    while vidcap.isOpened() and i < 10:
+        frm = frm + 1
+        success, image = vidcap.read()
+        if not success:
+            break
+        if frm not in idx:
+            continue
+        if resize is not None:
+            dest_width, dest_height = resize
+            width, height = image.shape[:2]
+            if width > dest_width or height > dest_height:
+                width_rate = dest_width / width
+                height_rate = dest_height / height
+                rate = min(width_rate, height_rate)
+                dest_width = int(width * rate)
+                dest_height = int(height * rate)
+                image = cv2.resize(image, (dest_width, dest_height), interpolation=cv2.INTER_AREA)
+        _, buffer = cv2.imencode(".jpg", image)
+        frames.append(f"data:image/jpg;base64,{base64.b64encode(buffer).decode('utf-8')}")
+        i += 1
+    vidcap.release()
+    return frames