Merge pull request #76 from SkywardAI/fix/chat

fix/chat_summary use conversational input instead of string

src/kimchima/tests/test_chat.py

@@ -37,11 +37,14 @@ def test_chat_summary(self):
             """
             Test chat_summary method
             """
-
+            con = PipelinesFactory.init_conversation()
+            con.add_message({"role": "user", "content": "Dod you like weather of Melbourne?"})
+            con.add_message({"role": "assistant", "content": "Melbourne is also sunny which is my favourite weather"})
+            con.add_message({"role": "user", "content": "why Melbourne is a good place to travel?"})
             res = Dialog.chat_summary(
                 pipe_con=self.pipe_con,
                 pipe_sum=self.pipe_sum,
-                messages=self.msg,
+                messages=con,
                 prompt=self.prompt,
                 max_length=self.max_length
                 )

src/kimchima/utils/dialog.py

@@ -51,11 +51,10 @@ def chat_summary(clas, *args,**kwargs)-> str:
         response = pipe_con(messages)
 
         logger.info("Finish conversation pipeline")
-
         if prompt is None:
-            return response[0].get('generated_text')
+            return response.messages[-1]["content"]
 
-        raw_response = prompt + response[0].get('generated_text')
+        raw_response = prompt + response.messages[-1]["content"]
 
         if max_length is None:
             max_length = len(raw_response)