Add image embedding task and test case

This commit is contained in:
liyaxing
2023-09-27 17:24:38 +08:00
parent e13e31aae9
commit 87f82ec9e2
4 changed files with 32 additions and 4 deletions
+14 -1
View File
@@ -2,6 +2,8 @@
from enum import Enum
import uuid
import time
from typing import Union
from knowledge import ObjectID
class ComputeTaskResultCode(Enum):
OK = 0
@@ -61,7 +63,7 @@ class ComputeTask:
if inner_functions is not None:
self.params["inner_functions"] = inner_functions
def set_text_embedding_params(self, input, model_name=None, callchain_id = None):
def set_text_embedding_params(self, input: str, model_name=None, callchain_id = None):
self.task_type = ComputeTaskType.TEXT_EMBEDDING
self.create_time = time.time()
self.task_id = uuid.uuid4().hex
@@ -71,6 +73,17 @@ class ComputeTask:
else:
self.params["model_name"] = "text-embedding-ada-002"
self.params["input"] = input
def set_image_embedding_params(self, input = Union[ObjectID, bytes], model_name=None, callchain_id = None):
self.task_type = ComputeTaskType.IMAGE_EMBEDDING
self.create_time = time.time()
self.task_id = uuid.uuid4().hex
self.callchain_id = callchain_id
if model_name is not None:
self.params["model_name"] = model_name
else:
self.params["model_name"] = None
self.params["input"] = input
def set_text_2_image_params(self, prompt: str, model_name, callchain_id=None):
self.task_type = ComputeTaskType.TEXT_2_IMAGE