File size: 2,055 Bytes
d5d3bd0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
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
|