toolbox.py

from llama_index.core.tools import FunctionTool
from llama_index.tools.duckduckgo import DuckDuckGoSearchToolSpec

import pint
import sympy as sp


class _Math:

    @staticmethod
    def symbolic_calc(expression: str) -> str:
        """
        Evaluates complex mathematical expressions using SymPy.
        
        Args:
            expression: Mathematical expression as string
        
        Returns:
            Result of the calculation
        """
        print(f"-> symbolic_math_calculator tool used (input: {expression}) !")
        try:
            result = sp.sympify(expression)
            print(f"-> (output: {result}) !")
            return str(result)
        except Exception as e:
            return f"Error in calculation: {str(e)}"

    @staticmethod
    def unit_converter(value: float, from_unit: str, to_unit: str) -> str:
        """
        Converts values between different units of measurement.
        
        Args:
            value: The numerical value to convert
            from_unit: The source unit (e.g., 'meter', 'kg', 'celsius')
            to_unit: The target unit (e.g., 'feet', 'pound', 'fahrenheit')
        
        Returns:
            The converted value with appropriate units
        """
        print(f"-> unit_converter tool used (inputs: [value: {value}, from_unit: {from_unit}, to_unit: {to_unit}]) !")
        try:
            # Create unit registry
            ureg = pint.UnitRegistry()
            
            # Create quantity with source unit
            quantity = value * ureg(from_unit)
            
            # Convert to target unit
            result = quantity.to(to_unit)
            
            answer = f"{value} {from_unit} = {result.magnitude} {to_unit}"
            print(f"-> (output: {result}) !")
            return answer
        except Exception as e:
            return f"Error in unit conversion: {str(e)}"


class _MathToolbox:
    symbolic_calc = FunctionTool.from_defaults(
        name="symbolic_calc",
        description="Evaluates complex mathematical expressions using SymPy",
        fn=_Math.symbolic_calc
    )
    unit_converter = FunctionTool.from_defaults(
        name="unit_converter",
        description="Converts values between different units of measurement",
        fn=_Math.unit_converter
    )


class _WebSearchToolbox:
    duck_duck_go_tools = DuckDuckGoSearchToolSpec().to_tool_list()


class _Encryption:
    
    @staticmethod
    def ascii_encode(text: str) -> str:
        """
        Converts each character in a string to its ASCII value.
        
        Args:
            text: The text to encode
            
        Returns:
            Space-separated ASCII values
        """
        print(f"-> ascii_encode tool used (input: {text[:30]}...) !")
        try:
            ascii_values = [str(ord(char)) for char in text]
            result = " ".join(ascii_values)
            return result
        except Exception as e:
            return f"Error in ASCII encoding: {str(e)}"
    
    @staticmethod
    def ascii_decode(text: str) -> str:
        """
        Converts space-separated ASCII values back to characters.
        
        Args:
            text: Space-separated ASCII values
            
        Returns:
            Decoded string
        """
        print(f"-> ascii_decode tool used (input: {text[:30]}...) !")
        try:
            ascii_values = text.split()
            result = "".join([chr(int(value)) for value in ascii_values])
            return result
        except Exception as e:
            return f"Error in ASCII decoding: {str(e)}"

    @staticmethod
    def base64_encode(text: str) -> str:
        """
        Encodes a string to base64.
        
        Args:
            text: The text to encode
            
        Returns:
            Base64 encoded string
        """
        print(f"-> base64_encode tool used (input: {text[:30]}...) !")
        import base64
        try:
            encoded_bytes = base64.b64encode(text.encode('utf-8'))
            encoded_text = encoded_bytes.decode('utf-8')
            return encoded_text
        except Exception as e:
            return f"Error in base64 encoding: {str(e)}"
    
    @staticmethod
    def base64_decode(encoded_text: str) -> str:
        """
        Decodes a base64 string to plain text.
        
        Args:
            encoded_text: The base64 encoded text
            
        Returns:
            Decoded string
        """
        print(f"-> base64_decode tool used (input: {encoded_text[:30]}...) !")
        import base64
        try:
            decoded_bytes = base64.b64decode(encoded_text)
            decoded_text = decoded_bytes.decode('utf-8')
            return decoded_text
        except Exception as e:
            return f"Error in base64 decoding: {str(e)}"
    
    @staticmethod
    def caesar_cipher_encode(text: str, shift: int) -> str:
        """
        Encodes text using Caesar cipher with specified shift.
        
        Args:
            text: The text to encode
            shift: Number of positions to shift each character
            
        Returns:
            Caesar cipher encoded string
        """
        print(f"-> caesar_cipher_encode tool used (input: {text[:30]}..., shift: {shift}) !")
        result = ""
        try:
            for char in text:
                if char.isalpha():
                    ascii_offset = ord('a') if char.islower() else ord('A')
                    encoded_char = chr((ord(char) - ascii_offset + shift) % 26 + ascii_offset)
                    result += encoded_char
                else:
                    result += char
            return result
        except Exception as e:
            return f"Error in Caesar cipher encoding: {str(e)}"
    
    @classmethod
    def caesar_cipher_decode(cls, encoded_text: str, shift: int) -> str:
        """
        Decodes Caesar cipher text with specified shift.
        
        Args:
            encoded_text: The encoded text
            shift: Number of positions the text was shifted
            
        Returns:
            Decoded string
        """
        print(f"-> caesar_cipher_decode tool used (input: {encoded_text[:30]}..., shift: {shift}) !")
        # To decode, we shift in the opposite direction
        return cls.caesar_cipher_encode(encoded_text, -shift)
    
    @classmethod
    def caesar_cipher_brute_force(cls, text: str) -> str:
        """
        Performs a brute force attack on a Caesar cipher by trying all 26 shifts.
        
        Args:
            text: The Caesar cipher encoded text
            
        Returns:
            All possible decoding results with their respective shifts
        """
        print(f"-> caesar_cipher_brute_force tool used (input: {text[:30]}...) !")
        results = []
        
        # Try all 26 possible shifts for English alphabet
        for shift in range(26):
            decoded = cls.caesar_cipher_decode(text, shift)
            results.append(f"Shift {shift}: {decoded}")
        
        output = "\n".join(results)
        return output

    @staticmethod
    def reverse_string(text: str) -> str:
        """
        Reverses a string.
        
        Args:
            text: The text to reverse
            
        Returns:
            Reversed string
        """
        print(f"-> reverse_string tool used (input: {text[:30]}...) !")
        reversed_text = text[::-1]
        return reversed_text


class _EncryptionToolbox:
    ascii_encode = FunctionTool.from_defaults(
        name="ascii_encode",
        description="Convert each character in a string to its ASCII value",
        fn=_Encryption.ascii_encode
    )
    ascii_decode = FunctionTool.from_defaults(
        name="ascii_decode",
        description="Convert space-separated ASCII values back to characters",
        fn=_Encryption.ascii_decode
    )
    base64_encode = FunctionTool.from_defaults(
        name="base64_encode",
        description="Encode a string to base64",
        fn=_Encryption.base64_encode
    )
    base64_decode = FunctionTool.from_defaults(
        name="base64_decode",
        description="Decode a base64 string to plain text",
        fn=_Encryption.base64_decode
    )
    caesar_cipher_encode = FunctionTool.from_defaults(
        name="caesar_cipher_encode",
        description="Encode a string using Caesar cipher with specified shift",
        fn=_Encryption.caesar_cipher_encode
    )
    caesar_cipher_decode = FunctionTool.from_defaults(
        name="caesar_cipher_decode",
        description="Decode a Caesar cipher string with specified shift",
        fn=_Encryption.caesar_cipher_decode
    )
    caesar_cipher_brute_force = FunctionTool.from_defaults(
        name="caesar_cipher_brute_force",
        description="Try all 26 possible shifts to decode a Caesar cipher text",
        fn=_Encryption.caesar_cipher_brute_force
    )
    reverse_string = FunctionTool.from_defaults(
        name="reverse_string",
        description="Reverse a string",
        fn=_Encryption.reverse_string
    )


class Toolbox:
    math = _MathToolbox()
    web_search = _WebSearchToolbox()
    encryption = _EncryptionToolbox()