Merge pull request #106 from glen0125/MVP

Fix the dall e invocation process

rootfs/agents/David/agent.toml

@@ -19,4 +19,6 @@ For example, if a user says, "I want a painting that depicts a sunset on a summe
 
 Then you will combine these pieces of information into an accurate prompt to pass to the "paint" function, for example: "During a summer sunset, children play on the beach with sailboats visible in the distance, and the sky displays a gradient of orange and purple."
 
+After successfully creating the image, please clearly state the file path where the image is saved in your response. For instance, you might say, "The image has been successfully created and saved at the following path: /path/to/the/image.png."
+
 Please always ensure that your creations respect the user's intentions and that you infuse your unique artistic style and creativity into your work."""

src/aios/agent/agent_base.py

@@ -477,6 +477,8 @@ class BaseAIAgent(abc.ABC):
         org_msg:AgentMsg,
         stack_limit = 5
     ) -> ComputeTaskResult:
+        from ..frame.compute_kernel import ComputeKernel
+        
         arguments = None
         try:
             func_name = inner_func_call_node.get("name")

src/aios/environment/workflow_env.py

@@ -319,9 +319,9 @@ class PaintEnvironment(Environment):
 
 
     async def _paint(self, prompt, model_name = None, negative_prompt = None) -> str:
-        err, result = await ComputeKernel.get_instance().do_text_2_image(prompt, model_name, negative_prompt)
+        result = await ComputeKernel.get_instance().do_text_2_image(prompt, model_name, negative_prompt)
-        if err is not None:
+        if result.result_code == ComputeTaskResultCode.ERROR:
-            return f"exec paint failed. err:{err}"
+            return f"exec paint failed. err:{result.error_str}"
         else:
             return f'exec paint OK, saved as a local file, path is: {result.result["file"]}'
 

src/aios/frame/compute_kernel.py

@@ -240,9 +240,9 @@ class ComputeKernel:
 
     async def do_text_2_image(self, prompt:str, model_name:Optional[str] = None, negative_prompt = None) -> ComputeTaskResult:
         task = self.text_2_image(prompt,model_name, negative_prompt)
-        task = await self._wait_task(task)
+        task_result = await self._wait_task(task)
 
-        return task.result
+        return task_result
         # if task_req.state == ComputeTaskState.DONE:
         #     return None, task_result
 

src/component/openai_node/dall_e_compute_node.py

@@ -14,13 +14,13 @@ from aios import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskTy
 logger = logging.getLogger(__name__)
 
 
-class DallE_ComputeNode(ComputeNode):
+class DallEComputeNode(ComputeNode):
     _instance = None
 
     @classmethod
     def get_instance(cls):
         if cls._instance is None:
-            cls._instance = DallE_ComputeNode()
+            cls._instance = DallEComputeNode()
         return cls._instance
 
     @classmethod

src/service/aios_shell/aios_shell.py

@@ -157,6 +157,13 @@ class AIOS_Shell:
         openai_tts_node = OpenAITTSComputeNode.get_instance()
         ComputeKernel.get_instance().add_compute_node(openai_tts_node)
 
+        dall_e_node = DallEComputeNode.get_instance()
+        if await dall_e_node.initial() is not True:
+            logger.error("dall-e node initial failed!")
+        else:
+            await dall_e_node.start()
+            ComputeKernel.get_instance().add_compute_node(dall_e_node)
+
         llama_nodes = ComputeNodeConfig.get_instance().initial()
         for llama_node in llama_nodes:
             llama_node.start()

test/test_dall_e_node.py

@@ -8,7 +8,7 @@ import logging
 import pytest
 directory = os.path.dirname(__file__)
 sys.path.append(directory + '/../src')
-from aios_kernel.dall_e_compute_node import DallE_ComputeNode
+from aios_kernel.dall_e_compute_node import DallEComputeNode
 from aios_kernel.compute_task import ComputeTaskType, ComputeTask, ComputeTaskState
 
 os.environ["TEXT2IMG_OUTPUT_DIR"] = "./"
@@ -16,7 +16,7 @@ os.environ["openai_api_key"] = ""
 
 @pytest.mark.asyncio
 async def test_dall_e_node(propmt, model):
-    node = DallE_ComputeNode.get_instance()
+    node = DallEComputeNode.get_instance()
     if await node.initial() is not True:
         print("node initial failed!")
         return