Add stability node (Text2Img)

src/aios_kernel/open_ai_node.py

@@ -5,47 +5,49 @@ import asyncio
 from asyncio import Queue
 import logging
 
-from .compute_task import ComputeTask,ComputeTaskResult,ComputeTaskState
+from .compute_task import ComputeTask, ComputeTaskResult, ComputeTaskState, ComputeTaskType
 from .compute_node import ComputeNode
 
 logger = logging.getLogger(__name__)
 
+
 class OpenAI_ComputeNode(ComputeNode):
     _instance = None
+
     def __new__(cls):
         if cls._instance is None:
             cls._instance = super(OpenAI_ComputeNode, cls).__new__(cls)
             cls._instance.is_start = False
         return cls._instance
-    
+
     def __init__(self) -> None:
         super().__init__()
         if self.is_start is True:
             logger.warn("OpenAI_ComputeNode is already start")
             return
-        
+
         self.is_start = True
-        #openai.organization = "org-AoKrOtF2myemvfiFfnsSU8rF" #buckycloud
+        # openai.organization = "org-AoKrOtF2myemvfiFfnsSU8rF" #buckycloud
         self.openai_api_key = ""
         self.node_id = "openai_node"
 
         self.task_queue = Queue()
 
-        if os.getenv("OPENAI_API_KEY") is not None: 
+        if os.getenv("OPENAI_API_KEY") is not None:
             openai.api_key = os.getenv("OPENAI_API_KEY")
         else:
             openai.api_key = self.openai_api_key
-        
+
         self.start()
-    
-    async def push_task(self,task:ComputeTask,proiority:int = 0):
+
+    async def push_task(self, task: ComputeTask, proiority: int = 0):
         logger.info(f"openai_node push task: {task.display()}")
         self.task_queue.put_nowait(task)
-    
-    async def remove_task(self,task_id:str):
+
+    async def remove_task(self, task_id: str):
         pass
-    
-    def _run_task(self,task:ComputeTask):
+
+    def _run_task(self, task: ComputeTask):
         task.state = ComputeTaskState.RUNNING
         mode_name = task.params["model_name"]
         # max_token_size = task.params["max_token_size"]
@@ -57,19 +59,19 @@ class OpenAI_ComputeNode(ComputeNode):
                                             max_tokens=4000,
                                             temperature=1.2)
         logger.info(f"openai response: {resp}")
-        
+
         status_code = resp["choices"][0]["finish_reason"]
         if status_code != "stop":
             task.state = ComputeTaskState.ERROR
-            task.error_str =f"The status code was {status_code}."
+            task.error_str = f"The status code was {status_code}."
             return None
-        
-        result = ComputeTaskResult()    
+
+        result = ComputeTaskResult()
         result.set_from_task(task)
         result.worker_id = self.node_id
         result.result_str = resp["choices"][0]["message"]["content"]
         result.result = resp["choices"][0]["message"]
-        
+
         return result
 
     def start(self):
@@ -82,29 +84,20 @@ class OpenAI_ComputeNode(ComputeNode):
                 if result is not None:
                     task.state = ComputeTaskState.DONE
                     task.result = result
-        
+
         asyncio.create_task(_run_task_loop())
 
     def display(self) -> str:
         return f"OpenAI_ComputeNode: {self.node_id}"
-    
-    def get_task_state(self,task_id:str):
-        pass
 
+    def get_task_state(self, task_id: str):
+        pass
 
     def get_capacity(self):
         pass
 
-
-    def is_support(self,task_type:str) -> bool:
-        return True
-
+    def is_support(self, task_type: ComputeTaskType) -> bool:
+        return task_type == ComputeTaskType.LLM_COMPLETION
 
     def is_local(self) -> bool:
         return False
-
-
-
-    
-    
-