Merge pull request #96 from alexsunxl/MVP

Add the gpt4 vision feature in AIOS

src/aios_kernel/__init__.py

@@ -18,6 +18,7 @@ from .email_tunnel import EmailTunnel
 from .storage import ResourceLocation,AIStorage,UserConfig,UserConfigItem
 from .contact_manager import ContactManager,Contact,FamilyMember  
 from .text_to_speech_function import TextToSpeechFunction
+from .image_2_text_function import Image2TextFunction
 from .workspace_env import ShellEnvironment
 from .local_stability_node import Local_Stability_ComputeNode
 from .stability_node import Stability_ComputeNode

src/aios_kernel/compute_kernel.py

@@ -228,5 +228,13 @@ class ComputeKernel:
         # if task_req.state == ComputeTaskState.DONE:
         #     return None, task_result
 
-        # return task_req.error_str, None
+    def image_2_text(self, image_path: str, prompt:str, model_name:Optional[str] = None, negative_prompt = None):
+        task = ComputeTask()
+        task.set_image_2_text_params(image_path,prompt,model_name, negative_prompt)
+        self.run(task)
+        return task
+    async def do_image_2_text(self, image_path: str, prompt:str, model_name:Optional[str] = None, negative_prompt = None) -> ComputeTaskResult:
+        task = self.image_2_text(image_path,prompt, model_name, negative_prompt)
+        task = await self._wait_task(task)
+        return task.result
 

src/aios_kernel/compute_task.py

@@ -24,6 +24,7 @@ class ComputeTaskType(Enum):
     NONE = "None"
     LLM_COMPLETION = "llm_completion"
     TEXT_2_IMAGE = "text_2_image"
+    IMAGE_2_TEXT = "image_2_text"
     IMAGE_2_IMAGE = "image_2_image"
     VOICE_2_TEXT = "voice_2_text"
     TEXT_2_VOICE = "text_2_voice"
@@ -98,6 +99,23 @@ class ComputeTask:
         else:
             self.params["model_name"] = "v1-5-pruned-emaonly"
 
+    def set_image_2_text_params(self, image_path: str, prompt: str, model_name, negative_prompt="", callchain_id=None):
+        self.task_type = ComputeTaskType.IMAGE_2_TEXT
+        self.create_time = time.time()
+        self.task_id = uuid.uuid4().hex
+        self.callchain_id = callchain_id
+        self.params["image_path"] = image_path
+        if prompt == '':
+            self.params["prompt"] = "What's in this image?"
+        else:
+            self.params["prompt"] = prompt
+        self.params["negative_prompt"] = negative_prompt
+        if model_name is not None:
+            self.params["model_name"] = model_name
+        else:
+            self.params["model_name"] = "gpt-4-vision-preview"
+
+
     def display(self) -> str:
         return f"ComputeTask: {self.task_id} {self.task_type} {self.state}"
 

src/aios_kernel/image_2_text_function.py

@@ -0,0 +1,46 @@
+import logging
+from typing import Dict
+
+from aios_kernel import ComputeKernel
+from aios_kernel.ai_function import AIFunction
+
+logger = logging.getLogger(__name__)
+
+class Image2TextFunction(AIFunction):
+    def __init__(self):
+        self.func_id = "image_2_text"
+        self.description = "According to the input image file address, return the description of the image content"
+        logger.info(f"init Image2TextFunction")
+
+    def get_name(self) -> str:
+        return self.func_id
+
+    def get_description(self) -> str:
+        return self.description
+
+    def get_parameters(self) -> Dict:
+        return {
+        }
+
+    async def execute(self, **kwargs) -> str:
+        logger.info(f"execute image_2_text function: {kwargs}")
+        image_path = kwargs.get("image_path")
+        data = await ComputeKernel.get_instance().do_image_2_text(image_path, '')
+        try:
+            result = data['message']['choices'][0]['message']['content']
+        except (KeyError, TypeError, IndexError):
+            logger.error(f"image_2_text error: {data}")
+            result = ""
+        return result
+
+
+    def is_local(self) -> bool:
+        return False
+
+    def is_in_zone(self) -> bool:
+        return True
+
+    def is_ready_only(self) -> bool:
+        return False
+
+

src/aios_kernel/open_ai_node.py

@@ -6,6 +6,8 @@ from asyncio import Queue
 import logging
 import json
 import aiohttp
+import base64
+import requests
 
 from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType,ComputeTaskResultCode
 from .compute_node import ComputeNode
@@ -90,6 +92,50 @@ class OpenAI_ComputeNode(ComputeNode):
         # result["message"] = result_msg
         return result
     
+    def _image_2_text(self, task: ComputeTask):
+        logger.info('openai image_2_text')
+        # 本地图片处理
+        def encode_image(image_path):
+            with open(image_path, "rb") as image_file:
+                return base64.b64encode(image_file.read()).decode('utf-8')
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {self.openai_api_key }"
+        }
+        model_name = task.params["model_name"]
+        base64_image = encode_image(task.params["image_path"])
+        payload = {
+            "model": model_name,
+            "messages": [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": task.params["prompt"]
+                        },
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/jpeg;base64,{base64_image}"
+                            }
+                        }
+                    ]
+                }
+            ],
+            "max_tokens": 300
+        }
+        logger.info('openai send image_2_text request ')
+        # openai 的库的Vision只支持传图片的url地址。本地图片得用request
+        response = requests.post("https://api.openai.com/v1/chat/completions", headers=headers, json=payload)
+        if response.status_code == 200:
+            logger.info('openai image_2_text success')
+            return response.json()
+        else:
+            logger.error('openai image_2_text error')
+            logger.error(response.json())
+            return None
+
     async def _run_task(self, task: ComputeTask):
         task.state = ComputeTaskState.RUNNING
         
@@ -133,6 +179,12 @@ class OpenAI_ComputeNode(ComputeNode):
                 result.worker_id = self.node_id
                 result.result_str = resp["data"][0]["embedding"]
 
+                return result
+            case ComputeTaskType.IMAGE_2_TEXT:
+                result.result_code = ComputeTaskResultCode.OK
+                result.worker_id = self.node_id
+                # result.result_str = resp["data"][0]["image_2_text"]
+                result.result["message"] = self._image_2_text(task)
                 return result
             case ComputeTaskType.LLM_COMPLETION:
                 mode_name = task.params["model_name"]
@@ -238,6 +290,10 @@ class OpenAI_ComputeNode(ComputeNode):
             if model_name.startswith("gpt-"):
                 return True
         
+        if task.task_type == ComputeTaskType.IMAGE_2_TEXT:
+            model_name : str = task.params["model_name"]
+            if model_name.startswith("gpt-4"):
+                return True
         #if task.task_type == ComputeTaskType.TEXT_EMBEDDING:
         #    if task.params["model_name"] == "text-embedding-ada-002":
         #        return True

src/aios_kernel/workflow_env.py

@@ -9,6 +9,7 @@ import logging
 from typing import Optional
 
 from .text_to_speech_function import TextToSpeechFunction
+from .image_2_text_function import Image2TextFunction
 from .compute_kernel import ComputeKernel, ComputeTaskResultCode
 from .environment import Environment,EnvironmentEvent
 from .ai_function import SimpleAIFunction
@@ -344,6 +345,7 @@ class WorkflowEnvironment(Environment):
         self.local = threading.local()
         self.table_name = "WorkflowEnv_" + env_id
         self.add_ai_function(TextToSpeechFunction())
+        self.add_ai_function(Image2TextFunction())
 
 
     def _get_conn(self):

test/env_test.py

@@ -1,11 +1,9 @@
 import asyncio
 import os
 import sys
-
 directory = os.path.dirname(__file__)
 sys.path.append(directory + '/../src')
-
+from aios_kernel import CalenderEnvironment,WorkflowEnvironment,ComputeKernel,OpenAI_ComputeNode,AIStorage
-from aios_kernel import CalenderEnvironment,WorkflowEnvironment
 
 
 async def test_buildin_envs():
@@ -22,6 +20,29 @@ async def test_buildin_envs():
 
     await asyncio.sleep(10)
 
+async def test_image_to_text():
+    # init the compute kernel and add the compute node
+    open_ai_node = OpenAI_ComputeNode.get_instance()
+    if await open_ai_node.initial() is not True:
+        print("openai node initial failed!")
+        return False
+    ComputeKernel.get_instance().add_compute_node(open_ai_node)
+    w_env = WorkflowEnvironment("workflow",os.path.abspath(directory + "/../rootfs/workflow_env.db"))
+    assert w_env.functions['image_2_text'] is not None
+    await ComputeKernel.get_instance().start()
+    fn = w_env.get_ai_function('image_2_text')
+    image_path = os.path.abspath(directory + "/test.png")
+    arguments = {
+        'image_path': image_path
+    }
+
+    # execute the ai function
+    result = await fn.execute(**arguments)
+    assert result is not ""
+    print(result)
+
+    await asyncio.sleep(10)
+
 
 if __name__ == "__main__":
     #test_rstr = "abc is {abc}"
@@ -29,4 +50,5 @@ if __name__ == "__main__":
     #new_str = test_rstr.format_map(values)
     #print(new_str)
 
-    asyncio.run(test_buildin_envs())
+    # asyncio.run(test_buildin_envs())
+    asyncio.run(test_image_to_text())