updated FrameController

src/controller/mainFrameController.py

@@ -1,3 +1,4 @@
+from collections import deque
 import customtkinter as ctk
 from views.mainScreen import MainFrame
 from models.data import TextData
@@ -5,6 +6,14 @@ from Ai.interence import VeraMindInference
 from utils.database.database import FakeNewsChecker
 from models.provider import Provider
 from collections import Counter
+from Ai.llm import ArticleRater
+
+BAD_WORDS = ["FAKE", "SATIRE", "Fake", "fake"]
+GOOD_WORDS = ["REAL", "real", "Real"]
+BAD_COLOR = "#ff8080"
+GOOD_COLOR = "#80ff8f"
+WORDS = BAD_WORDS + GOOD_WORDS
+
 
 class MainFrameController:
     """
@@ -22,6 +31,7 @@ class MainFrameController:
         self.model_inference = VeraMindInference('VeraMind-Mini')
         self.db = FakeNewsChecker()
         self.update_provider_list()
+        self.rater = ArticleRater()
 
     def get_text_data(self) -> TextData:
         """
@@ -35,16 +45,63 @@ class MainFrameController:
             text_data.text = self.frame.input_textbox.get("0.0", "end")
         return text_data
 
-    def press_check_button(self) -> None:
-        """
-        Handle the 'Check' button press event.
-        Processes the input, makes a prediction, updates the database, and displays the result.
-        """
-        text_data = self.get_text_data()
-        text_data = self._predict(text_data)
-        self._add_to_db(text_data)
-        self._update_output(text_data.get_output())
+    def press_check_button(self):
+        text_data = self.get_textdata()
+        print(text_data.text)
+        self._predict(text_data)
+        self.frame.output_textbox.configure(state="normal")
+        self.frame.output_textbox.delete("0.0", "end")
+
+        response_stream = self.rater.get_response(text_data.text, text_data.result, float(f"{text_data.confidence * 100:.2f}"))
+
+        highlight_buffer = deque(maxlen=5)
+
+        for chunk in response_stream:
+            # Display the chunk immediately
+            self.frame.output_textbox.insert("end", chunk)
+            self.frame.output_textbox.see("end")
+            self.frame.update_idletasks()
+
+            # Add to highlight buffer
+            highlight_buffer.append(chunk)
+
+            # Process highlighting when buffer is full
+            if len(highlight_buffer) == 5:
+                self.process_highlighting(highlight_buffer)
+
+        # Process any remaining chunks in the buffer
+        if highlight_buffer:
+            self.process_highlighting(highlight_buffer)
+
+        self.frame.output_textbox.configure(state="disabled")
         self.update_provider_list()
+    
+    def process_highlighting(self, highlight_buffer):
+        start_index = self.frame.output_textbox.index(f"end-{sum(len(c) for c in highlight_buffer)}c")
+        end_index = self.frame.output_textbox.index("end")
+        self.highlight_words(start_index, end_index)
+
+        # Keep overlap of 2 chunks
+        highlight_buffer = deque(list(highlight_buffer)[-2:], maxlen=5)
+
+    def highlight_words(self, start_index, end_index):
+        content = self.frame.output_textbox.get(start_index, end_index)
+
+        for word in WORDS:
+            start = 0
+            while True:
+                pos = content.find(word, start)
+                if pos == -1:
+                    break
+                word_start = f"{start_index}+{pos}c"
+                word_end = f"{word_start}+{len(word)}c"
+                tag_name = f"{word.lower()}_color"
+                self.frame.output_textbox.tag_add(tag_name, word_start, word_end)
+                if word in BAD_WORDS:
+                    self.frame.output_textbox.tag_config(tag_name, foreground=BAD_COLOR)
+                elif word in GOOD_WORDS:
+                    self.frame.output_textbox.tag_config(tag_name, foreground=GOOD_COLOR)
+                start = pos + len(word)
 
     def _predict(self, text_data: TextData) -> TextData:
         """
@@ -57,8 +114,6 @@ class MainFrameController:
         text_data.confidence = result["confidence"]
         text_data.result = result["result"]
         text_data.is_fake_news = result["is_fake"]
-        print(f"Prediction: {text_data.result}")
-        print(f"Confidence: {text_data.confidence}")
         return text_data
 
     def _add_to_db(self, text_data: TextData) -> None:
@@ -119,8 +174,6 @@ class MainFrameController:
             count_label = ctk.CTkLabel(provider_frame, text=str(provider.get_fake_percentage())+"%")
             count_label.pack(side="right", padx=5)
 
-
-
     def _update_output(self, output: str) -> None:
         """
         Update the output text box with the result.
@@ -130,6 +183,4 @@ class MainFrameController:
         self.frame.output_textbox.configure(state="normal")
         self.frame.output_textbox.delete("0.0", "end")
         self.frame.output_textbox.insert("0.0", output)
-        self.frame.output_textbox.configure(state="disabled")
-
-    
+        self.frame.output_textbox.configure(state="disabled")