|
from models import BaseModel |
|
from .base_refiner import BaseRefiner |
|
from utils.image_encoder import encode_image |
|
class SimpleRefiner(BaseRefiner): |
|
def __init__(self, |
|
sys_prompt: str, |
|
model: BaseModel, |
|
) -> None: |
|
BaseRefiner.__init__(self, sys_prompt=sys_prompt, model=model) |
|
|
|
def refine(self, message: str, image_path=None) -> str: |
|
if image_path: |
|
context = [{"role": "system", "content": self.sys_prompt}, {"role": "user", "content": [ |
|
{"type": "text", "text": f"{message}"}, |
|
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encode_image(image_path)}"}}, |
|
]}] |
|
else: |
|
context = [{"role": "system", "content": self.sys_prompt}, {"role": "user", "content": [ |
|
{"type": "text", "text": f"{message}"}, |
|
]}] |
|
response = self.model.respond(context) |
|
|
|
return response |