Add-ons

"""
Hello World Addon

A simple example add-on with two tools. Functions with preceding docstrings containing YAML definitions (including `tool_name`) are recognized as tools.

## Connection Scheme
Add-ons can optionally specify connection objects for external integrations. The required properties depend on the integration type, for example:

- AWS: region, aws_access_key_id, aws_secret_access_key
- REST with Bearer Auth: bearer_token, server_url
- REST with Basic Auth: username, password, server_url

The scheme is defined in YAML below, after the triple-dash (---). The UI uses this specification to render "Add Connection" fields.
---
cxn_scheme: hello_scheme
cxn_fields:
    bearer_token:
    server_url:
"""

from typing import TextIO, Dict, Any
from saturn.cell_result import CellResult
from saturn.cell_cxn import CellCxn
import numpy as np

"""
Fibonacci for fun
"""
def fibonacci(n):
    """
    Returns the nth Fibonacci number using matrix exponentiation.
    F(0) = 0, F(1) = 1
    """
    if n < 0:
        raise ValueError("n must be a non-negative integer")
    if n == 0:
        return 0
    if n == 1:
        return 1

    # Transformation matrix for Fibonacci
    F = np.array([[1, 1],
                  [1, 0]], dtype=object)

    # Compute the (n-1)th power of the matrix
    F_n = np.linalg.matrix_power(F, n - 1)

    # The top left element gives F(n)
    return F_n[0, 0]

"""
Hello world add-on tool specification
- Everything after the --- must be valid YAML
- The tool_name must match the Python function name
- The function signature must match this example
---
tool_name: hello_world_addon_tool
tool_param_map:
    person_name:
        type: string
    favorite_number:
        type: number
"""
def hello_world_addon_tool(cell_cxn: CellCxn,
                    tool_param_map: Dict[str, str],
                    debugOut: TextIO) -> CellResult:

    nm = tool_param_map.get("person_name", "stranger")
    num = tool_param_map.get("favorite_number", 42 )

    # numpy dependency
    fib = fibonacci( int(num) % 100 )

    # create a cell result with what you want to output as primary result of cell
    cr = CellResult(f'Hello {nm}, your favorite number is {num}, fib({num}%1000)={fib}, dummy server_url={cell_cxn.get("server_url")}')

    # you can also set extra information as part of cell results that may be
    # helpful for subsequent cells to use or just general debugging information
    cr.fib=fib

    return cr

"""
Another example tool with different tool parameters
---
tool_name: hello_world_goodbye_tool
tool_param_map:
    goodbye_msg:
        type: string
"""
def hello_world_goodbye_tool(cell_cxn: CellCxn,
                    tool_param_map: Dict[str, str],
                    debugOut: TextIO) -> CellResult:

    goodbye = tool_param_map.get("goodbye_msg", "Adios!!!")

    cr = CellResult( f'{goodbye}, nice meeting you and your secrets, dummy server_url={cell_cxn.get("server_url")}' )
    return cr