61 lines
1.6 KiB
Python
61 lines
1.6 KiB
Python
|
|
||
|
|
from enum import Enum
|
||
|
|
import uuid
|
||
|
|
import time
|
||
|
|
|
||
|
|
class ComputeTaskState(Enum):
|
||
|
|
DONE = 0
|
||
|
|
INIT = 1
|
||
|
|
RUNNING = 2
|
||
|
|
ERROR = 3
|
||
|
|
PENDING = 4
|
||
|
|
|
||
|
|
|
||
|
|
class ComputeTask:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self.task_type = "llm_completion"
|
||
|
|
self.create_time = None
|
||
|
|
|
||
|
|
self.task_id:str = None
|
||
|
|
self.callchain_id:str = None
|
||
|
|
self.params:dict = {}
|
||
|
|
self.refers:dict = None
|
||
|
|
self.pading_data:bytearray = None
|
||
|
|
|
||
|
|
self.state = ComputeTaskState.INIT
|
||
|
|
self.result = None
|
||
|
|
self.error_str = None
|
||
|
|
|
||
|
|
def set_llm_params(self,prompts,model_name,max_token_size,callchain_id = None):
|
||
|
|
self.task_type = "llm_completion"
|
||
|
|
self.create_time = time.time()
|
||
|
|
self.task_id = uuid.uuid4().hex
|
||
|
|
self.callchain_id = callchain_id
|
||
|
|
self.params["prompts"] = prompts.messages
|
||
|
|
if model_name is not None:
|
||
|
|
self.params["model_name"] = model_name
|
||
|
|
else:
|
||
|
|
self.params["model_name"] = "gpt-4-0613"
|
||
|
|
self.params["max_token_size"] = max_token_size
|
||
|
|
|
||
|
|
def display(self) -> str:
|
||
|
|
return f"ComputeTask: {self.task_id} {self.task_type} {self.state}"
|
||
|
|
|
||
|
|
|
||
|
|
class ComputeTaskResult:
|
||
|
|
def __init__(self) -> None:
|
||
|
|
self.create_time = None
|
||
|
|
self.task_id:str = None
|
||
|
|
self.callchain_id:str = None
|
||
|
|
self.worker_id:str = None
|
||
|
|
self.result_code:int = 0
|
||
|
|
self.result_str:str = None
|
||
|
|
|
||
|
|
self.result:dict = {}
|
||
|
|
self.result_refers:dict = None
|
||
|
|
self.pading_data:bytearray = None
|
||
|
|
|
||
|
|
def set_from_task(self,task:ComputeTask):
|
||
|
|
self.task_id = task.task_id
|
||
|
|
self.callchain_id = task.callchain_id
|