diff --git a/src/aios_kernel/__init__.py b/src/aios_kernel/__init__.py index a9a5302..9093b07 100644 --- a/src/aios_kernel/__init__.py +++ b/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 diff --git a/src/aios_kernel/compute_kernel.py b/src/aios_kernel/compute_kernel.py index 0a9e3a3..080d455 100644 --- a/src/aios_kernel/compute_kernel.py +++ b/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 diff --git a/src/aios_kernel/compute_task.py b/src/aios_kernel/compute_task.py index 6b9d74c..3728cdf 100644 --- a/src/aios_kernel/compute_task.py +++ b/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}" diff --git a/src/aios_kernel/image_2_text_function.py b/src/aios_kernel/image_2_text_function.py new file mode 100644 index 0000000..da84b04 --- /dev/null +++ b/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 + + diff --git a/src/aios_kernel/open_ai_node.py b/src/aios_kernel/open_ai_node.py index 1b686d0..f35fda7 100644 --- a/src/aios_kernel/open_ai_node.py +++ b/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 @@ -89,6 +91,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 diff --git a/src/aios_kernel/workflow_env.py b/src/aios_kernel/workflow_env.py index 94cac90..9201fba 100644 --- a/src/aios_kernel/workflow_env.py +++ b/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): diff --git a/test/env_test.py b/test/env_test.py index 694acbb..5ba2d6b 100644 --- a/test/env_test.py +++ b/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 +from aios_kernel import CalenderEnvironment,WorkflowEnvironment,ComputeKernel,OpenAI_ComputeNode,AIStorage 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()) \ No newline at end of file + # asyncio.run(test_buildin_envs()) + asyncio.run(test_image_to_text()) \ No newline at end of file diff --git a/test/test.png b/test/test.png new file mode 100644 index 0000000..7f6cb7a Binary files /dev/null and b/test/test.png differ