Browse Source
Allow system message
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
1 changed files with
17 additions and
6 deletions
-
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): |
|
|
|
|
