Add system message

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>

README.md

@@ -31,27 +31,31 @@ input:
 {context}
 """
 
+sys_message = """Your name is TowheeChat."""
+
 p = (
     pipe.input('question', 'doc', 'history')
     .map('doc', 'doc', lambda x: x[:2000])
-    .map(('question', 'doc', 'history'), 'prompt', ops.prompt.template(temp, ['question', 'context']))
+    .map(('question', 'doc', 'history'), 'prompt', ops.prompt.template(temp, ['question', 'context'], sys_message))
     .map('prompt', 'answer', ops.LLM.OpenAI())
     .output('answer')
 )
 
-an1 = p('Tell me something about Towhee', towhee_docs, []).get()[0]
+an1 = p('Who are you?', [], []).get()[0]
 print(an1)
 
-an2 = p('How to use it', towhee_docs, [('Tell me something about Towhee', an1)]).get()[0]
+an2 = p('Tell me something about Towhee', towhee_docs, []).get()[0]
 print(an2)
 
+an3 = p('How to use it', towhee_docs, [('Tell me something about Towhee', an2)]).get()[0]
+print(an3)
 ```
 
 ## Factory Constructor
 
 Create the operator via the following factory method:
 
-***ops.prompt.template(temp, keys)***
+***ops.prompt.template(temp, keys, sys_msg)***
 
 <br />
 

__init__.py

@@ -1,5 +1,5 @@
 from typing import List
 from .template_prompt import TemplatePrompt
 
-def template(temp: str, keys: List[str]):
-    return TemplatePrompt(temp, keys)
+def template(*args, **kwargs):
+    return TemplatePrompt(*args, **kwargs)

template_prompt.py

@@ -4,12 +4,17 @@ from towhee.operator import PyOperator
 
 
 class TemplatePrompt(PyOperator):
-    def __init__(self, temp: str, keys: List[str]):
+    def __init__(self, temp: str, keys: List[str], sys_msg: str = None):
         super().__init__()
         self._template = temp
         self._keys = keys
+        self._sys_msg = sys_msg
 
     def __call__(self, *args) -> List[Dict[str, str]]:
+        if self._sys_msg:
+            system_message = [{'system': self._sys_msg}]
+        else:
+            system_message = []
         if len(self._keys) == len(args):
             history = []
         else:
@@ -20,9 +25,9 @@ class TemplatePrompt(PyOperator):
         ret = [{'question': prompt_str}]
 
         if not history:
-            return ret
+            return system_message + ret
         else:
             history_data = []
             for item in history:
                 history_data.append({'question': item[0], 'answer': item[1]})
-            return history_data + ret
+            return system_message + history_data + ret