| 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 |
| |
| |
|
|