Update return msg

PingJun-You · Mar 23, 2023 · b8861a4 · b8861a4
1 parent a595464
commit b8861a4
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 11 deletions.
diff --git a/main.py b/main.py
@@ -100,15 +100,21 @@ def handle_text_message(event):
                     if not is_successful:
                         raise Exception(error_message)
                     youtube_transcript_reader = YoutubeTranscriptReader(user_model, os.getenv('OPENAI_MODEL_ENGINE'))
-                    role, response = youtube_transcript_reader.summarize(chunks)
+                    is_successful, response, error_message = youtube_transcript_reader.summarize(chunks)
+                    if not is_successful:
+                        raise Exception(error_message)
+                    role, response = get_role_and_content(response)
                     msg = TextSendMessage(text=response)
                 else:
                     chunks = website.get_content_from_url(url)
                     if len(chunks) == 0:
                         msg = TextSendMessage(text='無法撈取此網站文字')
                     else:
                         website_reader = WebsiteReader(user_model, os.getenv('OPENAI_MODEL_ENGINE'))
-                        role, response = website_reader.summarize(chunks)
+                        is_successful, response, error_message = website_reader.summarize(chunks)
+                        if not is_successful:
+                            raise Exception(error_message)
+                        role, response = get_role_and_content(response)
                         msg = TextSendMessage(text=response)
             else:
                 is_successful, response, error_message = user_model.chat_completions(memory.get(user_id), os.getenv('OPENAI_MODEL_ENGINE'))

diff --git a/src/service/website.py b/src/service/website.py
@@ -50,8 +50,7 @@ def __init__(self, model=None, model_engine=None):
         self.model_engine = model_engine
 
     def send_msg(self, msg):
-        role, content = self.model.chat_completions(msg, self.model_engine)
-        return role, content
+        return self.model.chat_completions(msg, self.model_engine)
 
     def summarize(self, chunks):
         text = '\n'.join(chunks)[:self.text_length_limit]
@@ -60,5 +59,4 @@ def summarize(self, chunks):
         }, {
             "role": "user", "content": self.message_format.format(text)
         }]
-        role, response = self.send_msg(msgs)
-        return role, response
+        return self.send_msg(msgs)
diff --git a/src/service/youtube.py b/src/service/youtube.py
@@ -1,6 +1,7 @@
 import math
 import os
 import re
+from src.utils import get_role_and_content
 
 from youtube_transcript_api import YouTubeTranscriptApi, NoTranscriptFound, TranscriptsDisabled
 
@@ -51,8 +52,7 @@ def __init__(self, model=None, model_engine=None):
         self.model_engine = model_engine
 
     def send_msg(self, msg):
-        role, content = self.model.chat_completions(msg, self.model_engine)
-        return role, content
+        return self.model.chat_completions(msg, self.model_engine)
 
     def summarize(self, chunks):
         summary_msg = []
@@ -63,7 +63,8 @@ def summarize(self, chunks):
                 }, {
                     "role": "user", "content": self.part_message_format.format(i, chunk, i)
                 }]
-                _, content = self.send_msg(msgs)
+                _, response, _ = self.send_msg(msgs)
+                _, content = get_role_and_content(response)
                 summary_msg.append(content)
             text = '\n'.join(summary_msg)
             msgs = [{
@@ -78,5 +79,4 @@ def summarize(self, chunks):
             }, {
                 'role': 'user', 'content': self.single_message_format.format(text)
             }]
-        role, response = self.send_msg(msgs)
-        return role, response
+        return self.send_msg(msgs)