solver.py

from llama_index.core.tools import FunctionTool

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)


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[FunctionTool]:
        return [
                Toolbox.web_search.duckduckgo_text_search,
                Toolbox.web_search.duckduckgo_images_search,
                Toolbox.web_search.duckduckgo_videos_search
            ]


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[FunctionTool]:
        return [
                Toolbox.encryption.ascii_encode,
                Toolbox.encryption.ascii_decode,
                Toolbox.encryption.base64_encode,
                Toolbox.encryption.base64_decode,
                Toolbox.encryption.caesar_cipher_encode,
                Toolbox.encryption.caesar_cipher_decode,
                Toolbox.encryption.caesar_cipher_brute_force,
                Toolbox.encryption.reverse_string,
                Toolbox.math.unit_converter
            ]

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


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[FunctionTool]:
        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)


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

    async def query(self, question: str, has_context=True) -> str:
        return "Image Handler is not available due to maintainance !"


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

    async def query(self, question: str, has_context=True) -> str:
        return "Video Handler is not available due to maintainance !"


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

    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