gpt4free

import unittest
import time

from g4f.tools.run_tools import ThinkingProcessor, Reasoning


class TestThinkingProcessor(unittest.TestCase):
    def assertEqualReasoning(self, first, second, msg=None):
        if isinstance(first, Reasoning) and isinstance(second, Reasoning):
            if (
                first.status != second.status
                or first.is_thinking != second.is_thinking
                or first.token != second.token
                or first.label != second.label
            ):
                raise self.failureException(msg or f"{first.label} != {second.label}")
        elif first != second:
            raise self.failureException(msg or f"{first} != {second}")

    def test_non_thinking_chunk(self):
        chunk = "This is a regular text."
        expected_time, expected_result = 0, [chunk]
        actual_time, actual_result = ThinkingProcessor.process_thinking_chunk(chunk)
        self.assertEqual(actual_time, expected_time)
        self.assertEqual(actual_result, expected_result)

    def test_thinking_start(self):
        chunk = "Hello <think>World"
        expected_time = time.time()
        expected_result = [
            "Hello ",
            Reasoning(status="🤔 Is thinking...", is_thinking="<think>"),
            Reasoning("World"),
        ]
        actual_time, actual_result = ThinkingProcessor.process_thinking_chunk(chunk)
        self.assertAlmostEqual(actual_time, expected_time, delta=1)
        self.assertEqual(actual_result[0], expected_result[0])
        self.assertEqualReasoning(actual_result[1], expected_result[1])
        self.assertEqualReasoning(actual_result[2], expected_result[2])

    def test_thinking_end(self):
        start_time = time.time()
        chunk = "token</think> content after"
        expected_result = [
            Reasoning("token"),
            Reasoning(status="", is_thinking="</think>"),
            " content after",
        ]
        actual_time, actual_result = ThinkingProcessor.process_thinking_chunk(
            chunk, start_time
        )
        self.assertEqual(actual_time, 0)
        self.assertEqualReasoning(actual_result[0], expected_result[0])
        self.assertEqualReasoning(actual_result[1], expected_result[1])
        self.assertEqual(actual_result[2], expected_result[2])

    def test_thinking_start_and_end(self):
        start_time = time.time()
        chunk = "<think>token</think> content after"
        expected_result = [
            Reasoning(status="🤔 Is thinking...", is_thinking="<think>"),
            Reasoning("token"),
            Reasoning(status="", is_thinking="</think>"),
            " content after",
        ]
        actual_time, actual_result = ThinkingProcessor.process_thinking_chunk(
            chunk, start_time
        )
        self.assertEqual(actual_time, 0)
        self.assertEqualReasoning(actual_result[0], expected_result[0])
        self.assertEqualReasoning(actual_result[1], expected_result[1])
        self.assertEqualReasoning(actual_result[2], expected_result[2])
        self.assertEqual(actual_result[3], expected_result[3])

    def test_ongoing_thinking(self):
        start_time = time.time()
        chunk = "Still thinking..."
        expected_result = [Reasoning("Still thinking...")]
        actual_time, actual_result = ThinkingProcessor.process_thinking_chunk(
            chunk, start_time
        )
        self.assertEqual(actual_time, start_time)
        self.assertEqualReasoning(actual_result[0], expected_result[0])

    def test_chunk_with_text_after_think(self):
        chunk = "Start <think>Middle</think>End"
        expected_time = 0
        expected_result = [
            "Start ",
            Reasoning(status="🤔 Is thinking...", is_thinking="<think>"),
            Reasoning("Middle"),
            Reasoning(status="", is_thinking="</think>"),
            "End",
        ]
        actual_time, actual_result = ThinkingProcessor.process_thinking_chunk(chunk)
        self.assertEqual(actual_time, expected_time)
        for i in range(1, len(expected_result)):
            self.assertEqualReasoning(actual_result[i], expected_result[i])