solver.py

from typing import List

from itf_agent import IAgent
from toolbox import Toolbox
from args import Args


class Summarizer(IAgent):
    def __init__(self, temperature, max_tokens):
        super().__init__(temperature, max_tokens, "04_summarizer.txt", Args.primary_llm_interface)

    def setup_tools(self) -> List:
        """
        Abstract method to set up the tools.
        This method must be overridden by subclasses to define custom tools this agent can use.
        """
        return []

    def setup_slaves(self) -> List:
        """
        Abstract method to set up the slave agents.
        This method must be overridden by subclasses to define custom sub-agents this agent can use.
        """
        return []


class Researcher(IAgent):
    def __init__(self, temperature, max_tokens):
        super().__init__(temperature, max_tokens, "05_researcher.txt", Args.primary_llm_interface)

    def setup_tools(self) -> List:
        return [Toolbox.web_search.duck_duck_go_tools]

    def setup_slaves(self) -> List:
        return []


class EncryptionExpert(IAgent):
    def __init__(self, temperature, max_tokens):
        super().__init__(temperature, max_tokens, "06_encryption_expert.txt", Args.primary_llm_interface)

    def setup_tools(self) -> List:
        return [
                Toolbox.encryption.base64_encode,
                Toolbox.encryption.base64_decode,
                Toolbox.encryption.caesar_cipher_encode,
                Toolbox.encryption.caesar_cipher_decode,
                Toolbox.encryption.reverse_string
                # TODO: Add more encryption tools
            ]

    def setup_slaves(self) -> List:
        return []


class MathExpert(IAgent):
    def __init__(self, temperature, max_tokens):
        super().__init__(temperature, max_tokens, "07_math_expert.txt", Args.primary_llm_interface)

    def setup_tools(self) -> List:
        return [
            Toolbox.math.symbolic_calc,
            Toolbox.math.unit_converter,
            ]

    def setup_slaves(self) -> List:
        reasoner = Reasoner(self.temperature, self.max_tokens)
        return [reasoner]


class Reasoner(IAgent):
    def __init__(self, temperature, max_tokens):
        super().__init__(temperature, max_tokens, "08_reasoner.txt", Args.primary_llm_interface)

    def setup_tools(self) -> List:
        return []

    def setup_slaves(self) -> List:
        return []


class ImageHandler(IAgent):
    def __init__(self, temperature, max_tokens):
        super().__init__(temperature, max_tokens, "09_image_handler.txt", Args.vlm_interface)

    def setup_tools(self) -> List:
        return []

    def setup_slaves(self) -> List:
        return []


class VideoHandler(IAgent):
    def __init__(self, temperature, max_tokens):
        super().__init__(temperature, max_tokens, "10_video_handler.txt", Args.vlm_interface)

    def setup_tools(self) -> List:
        return []

    def setup_slaves(self) -> List:
        return []


class Solver(IAgent):
    def __init__(self, temperature, max_tokens):
        super().__init__(temperature, max_tokens, "03_solver.txt", Args.primary_llm_interface)

    def setup_tools(self) -> List:
        return []

    def setup_slaves(self) -> List:
        summarizer = Summarizer(self.temperature, self.max_tokens)
        researcher = Researcher(self.temperature, self.max_tokens)
        encryption_expert = EncryptionExpert(self.temperature, self.max_tokens)
        math_expert = MathExpert(self.temperature, self.max_tokens)
        reasoner = Reasoner(self.temperature, self.max_tokens)
        image_handler = ImageHandler(self.temperature, self.max_tokens)
        video_handler = VideoHandler(self.temperature, self.max_tokens)

        return [
            summarizer,
            researcher,
            encryption_expert,
            math_expert,
            reasoner,
            image_handler,
            video_handler
            ]


# if __name__ == "__main__":
#     pass