logo
Ernie
repo-copy-icon

copied

Browse Source

Allow system message

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
Jael Gu 2 years ago
parent
commit
048f1cecd0
  1. 23
      ernie_chat.py

23
ernie_chat.py

@ -63,13 +63,24 @@ class ErnieChat(PyOperator):
new_messages.append(m)
else:
for k, v in m.items():
if k == 'question':
new_m = {'role': 'user', 'content': v}
elif k == 'answer':
new_m = {'role': 'assistant', 'content': v}
if ('role' and 'content' in m) and (m['role'] in ['system', 'assistant', 'user']):
if m['role'] == 'system':
new_messages.append(m)
new_messages.append({'role': 'assistant', 'content': 'OK.'})
else:
new_messages.append(m)
else:
raise KeyError('Invalid message key: only accept key value from ["question", "answer"].')
new_messages.append(new_m)
for k, v in m.items():
if k == 'question':
new_ms = [{'role': 'user', 'content': v}]
elif k == 'answer':
new_ms = [{'role': 'assistant', 'content': v}]
elif k == 'system':
new_ms = [{'role': 'user', 'content': v}, {'role': 'assistant', 'content': 'OK.'}]
else:
raise KeyError(
'Invalid message key: only accept key value from ["question", "answer"].')
new_messages += new_ms
return new_messages
def stream_output(self, response):

Loading…
Cancel
Save