| from typing import Literal | |
| from chatAPI import chat | |
| from typing import Union | |
| class ChatbotBase: | |
| def __init__(self,system_prompt) -> None: | |
| self.history=[] | |
| self.system_prompt='' | |
| self.set_system_prompt(system_prompt) | |
| def set_system_prompt(self,prompt): | |
| '''set system prompt. MUST BE used after reset_bot()''' | |
| self.system_prompt=prompt | |
| self.add_message(prompt,image=None,role='system') | |
| def reset_bot(self): | |
| self.history=[] | |
| self.system_prompt='' | |
| def add_message(self,message,image:Union[str,list[str]],role:Literal['user','assistant','system']): | |
| '''add new message to history | |
| image: base64 string of image | |
| ''' | |
| if isinstance(image,list): | |
| content=[] | |
| content.append({'type':'text','text':message}) | |
| for img in image: | |
| content.append({ | |
| 'type':'image_url', | |
| "image_url": { | |
| "url": f"data:image/jpeg;base64,{image}", | |
| "detail": "high" | |
| }}) | |
| self.history.append({ | |
| 'role':role, | |
| 'content':content}) | |
| elif image: | |
| self.history.append({ | |
| 'role':role, | |
| 'content':[ | |
| {'type':'text','text':message}, | |
| {'type':'image_url',"image_url": { | |
| "url": f"data:image/jpeg;base64,{image}", | |
| "detail": "high" | |
| }} | |
| ]}) | |
| else: | |
| self.history.append({ | |
| 'role':role, | |
| 'content':message}) | |
| def pop_history(self): | |
| '''remove last history message''' | |
| self.history.pop() | |
| def get_history(self): | |
| return self.history | |
| def generate(self): | |
| '''generate new response and store to history''' | |
| print('code generating') | |
| rsp=chat(self.history) | |
| self.add_message(rsp,image=None,role='assistant') | |
| print(self.history) | |
| return rsp | |