Add openai text embedding compute task
This commit is contained in:
@@ -119,3 +119,33 @@ class ComputeKernel:
|
|||||||
|
|
||||||
return "error!"
|
return "error!"
|
||||||
|
|
||||||
|
def text_embedding(self,input:str,model_name:Optional[str] = None):
|
||||||
|
task_req = ComputeTask()
|
||||||
|
task_req.set_text_embeding_params(input,model_name)
|
||||||
|
self.run(task_req)
|
||||||
|
return task_req
|
||||||
|
|
||||||
|
async def do_text_embedding(self,input:str,model_name:Optional[str] = None) -> [float]
|
||||||
|
task_req = self.text_embedding(input,model_name)
|
||||||
|
async def check_timer():
|
||||||
|
check_times = 0
|
||||||
|
while True:
|
||||||
|
if task_req.state == ComputeTaskState.DONE:
|
||||||
|
break
|
||||||
|
|
||||||
|
if task_req.state == ComputeTaskState.ERROR:
|
||||||
|
break
|
||||||
|
|
||||||
|
if check_times >= 20:
|
||||||
|
task_req.state = ComputeTaskState.ERROR
|
||||||
|
break
|
||||||
|
|
||||||
|
await asyncio.sleep(0.5)
|
||||||
|
check_times += 1
|
||||||
|
|
||||||
|
await asyncio.create_task(check_timer())
|
||||||
|
if task_req.state == ComputeTaskState.DONE:
|
||||||
|
return task_req.result.result
|
||||||
|
|
||||||
|
return "error!"
|
||||||
|
|
||||||
@@ -28,7 +28,7 @@ class ComputeNode(ABC):
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def is_support(self,task_type:str) -> bool:
|
def is_support(self, task: ComputeTask) -> bool:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
|
|||||||
@@ -38,12 +38,15 @@ class ComputeTask:
|
|||||||
self.params["model_name"] = "gpt-4-0613"
|
self.params["model_name"] = "gpt-4-0613"
|
||||||
self.params["max_token_size"] = max_token_size
|
self.params["max_token_size"] = max_token_size
|
||||||
|
|
||||||
def set_embeding_params(self, model_name, input, callchain_id = None):
|
def set_text_embeding_params(self, input, model_name=None, callchain_id = None):
|
||||||
self.task_type = "embeding"
|
self.task_type = "text_embedding"
|
||||||
self.create_time = time.time()
|
self.create_time = time.time()
|
||||||
self.task_id = uuid.uuid4().hex
|
self.task_id = uuid.uuid4().hex
|
||||||
self.callchain_id = callchain_id
|
self.callchain_id = callchain_id
|
||||||
|
if model_name is not None:
|
||||||
self.params["model_name"] = model_name
|
self.params["model_name"] = model_name
|
||||||
|
else:
|
||||||
|
self.params["model_name"] = "text-embedding-ada-002"
|
||||||
self.params["input"] = input
|
self.params["input"] = input
|
||||||
|
|
||||||
def display(self) -> str:
|
def display(self) -> str:
|
||||||
|
|||||||
@@ -8,3 +8,5 @@ class KnowledgeBase:
|
|||||||
|
|
||||||
async def query(self, prompt: AgentPrompt) -> AgentPrompt:
|
async def query(self, prompt: AgentPrompt) -> AgentPrompt:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
@@ -77,8 +77,21 @@ class OpenAI_ComputeNode(ComputeNode):
|
|||||||
result.result = resp["choices"][0]["message"]
|
result.result = resp["choices"][0]["message"]
|
||||||
|
|
||||||
return result
|
return result
|
||||||
if task.task_type == "embeding":
|
if task.task_type == "text_embedding":
|
||||||
pass
|
model_name = task.params["model_name"]
|
||||||
|
input = task.params["input"]
|
||||||
|
logger.info(f"call openai {model_name} input: {input}")
|
||||||
|
|
||||||
|
resp = openai.Embeding.create(model=model_name,
|
||||||
|
input=input)
|
||||||
|
logger.info(f"openai response: {resp}")
|
||||||
|
|
||||||
|
result = ComputeTaskResult()
|
||||||
|
result.set_from_task(task)
|
||||||
|
result.worker_id = self.node_id
|
||||||
|
result.result = resp["data"][0]["embedding"]
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
def start(self):
|
def start(self):
|
||||||
async def _run_task_loop():
|
async def _run_task_loop():
|
||||||
@@ -104,8 +117,13 @@ class OpenAI_ComputeNode(ComputeNode):
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
def is_support(self,task_type:str) -> bool:
|
def is_support(self, task: ComputeTask) -> bool:
|
||||||
|
if task.task_type == "llm_completion":
|
||||||
return True
|
return True
|
||||||
|
if task.task_type == "text_embedding":
|
||||||
|
if task.params["model_name"] == "text-embedding-ada-002":
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
def is_local(self) -> bool:
|
def is_local(self) -> bool:
|
||||||
|
|||||||
Reference in New Issue
Block a user