diff --git a/src/blackbox/blackbox_factory.py b/src/blackbox/blackbox_factory.py index d7e22b2..60ffb31 100644 --- a/src/blackbox/blackbox_factory.py +++ b/src/blackbox/blackbox_factory.py @@ -10,6 +10,7 @@ from .text_to_audio import TextToAudio from .tesou import Tesou from .fastchat import Fastchat from .g2e import G2E +from .text_and_image import TextAndImage class BlackboxFactory: @@ -25,6 +26,7 @@ class BlackboxFactory: self.fastchat = Fastchat() self.audio_chat = AudioChat(self.asr, self.tesou, self.tts) self.g2e = G2E() + self.text_and_image = TextAndImage() def __call__(self, *args, **kwargs): return self.processing(*args, **kwargs) @@ -52,4 +54,6 @@ class BlackboxFactory: return self.audio_chat if blackbox_name == "g2e": return self.g2e + if blackbox_name == 'text_and_image': + return self.text_and_image raise ValueError("Invalid blockbox type") \ No newline at end of file diff --git a/src/blackbox/g2e.py b/src/blackbox/g2e.py index 90095a1..84f9e90 100755 --- a/src/blackbox/g2e.py +++ b/src/blackbox/g2e.py @@ -22,7 +22,7 @@ class G2E(Blackbox): if context == None: context = [] url = 'http://120.196.116.194:48890/v1' - + #url = 'http://120.196.116.194:48892/v1' background_prompt = '''KOMBUKIKI是一款茶饮料,目标受众 年龄:20-35岁 性别:女性 地点:一线城市、二线城市 职业:精英中产、都市白领 收入水平:中高收入,有一定消费能力 兴趣和爱好:注重健康,有运动习惯 @@ -47,8 +47,11 @@ class G2E(Blackbox): prompt_template = [ - {"role": "system", "content": background_prompt + prompt1}, + {"role": "system", "content": background_prompt + prompt1}, ] + #prompt_template = [ + # {"role": "system", "content": ''}, + #] messages = prompt_template + context + [ @@ -61,8 +64,8 @@ class G2E(Blackbox): api_key='YOUR_API_KEY', base_url=url ) - model_name = client.models.list().data[0].id + print(model_name) response = client.chat.completions.create( model=model_name, messages=messages, diff --git a/src/blackbox/text_and_image.py b/src/blackbox/text_and_image.py new file mode 100644 index 0000000..c63fdf6 --- /dev/null +++ b/src/blackbox/text_and_image.py @@ -0,0 +1,28 @@ +from typing import Any, Coroutine + +from fastapi import Request, Response, status +from fastapi.responses import JSONResponse +from .blackbox import Blackbox + + +class TextAndImage(Blackbox): + + def __call__(self, *args, **kwargs): + return self.processing(*args, **kwargs) + + def processing(self, text: str, image: bytes): + # TODO: implement the processing logic + return "text and image processing..." + + def valid(self, *args, **kwargs) -> bool: + return True + + async def fast_api_handler(self, request: Request) -> Response: + # If the request is a form data, you can use the following code + data = (await request.form()).get("image") + if data is None: + return JSONResponse(content={"error": "image is required"}, status_code=status.HTTP_400_BAD_REQUEST) + imageBytes = await data.read() + text = (await request.form()).get("text") + result = self.processing(text, imageBytes) + return JSONResponse(content={"result": result }, status_code=status.HTTP_200_OK) \ No newline at end of file