common/org_search_component.py

from typing import List, Tuple, Optional, Any
import os

import gradio as gr

try:
    from base.api_base import BaseAPI
except ImportError:
    from .base.api_base import BaseAPI


class OrgSearch(BaseAPI):

    def __init__(self):
        super().__init__(
            url=f"{os.getenv('CDS_API_URL')}/v1/organization/search",
            headers={"x-api-key": os.getenv('CDS_API_KEY')}
        )

    def __call__(self, name: str, **kwargs):
        is_valid = False

        payload = {
            "names": [{
                "value": name,
                "type": "main"
            }],
            "status": "authorized"
        }

        if kwargs.get("ein"):
            ein = kwargs.get("ein")
            if "-" not in ein:
                ein = f"{ein[:2]}-{ein[2:]}"
            payload["ids"] = [{
                "value": ein,
                "type": "ein"
            }]
            is_valid = True

        if kwargs.get("street") or kwargs.get("city") or kwargs.get("state") or kwargs.get("postal_code"):
            payload["addresses"] = [{
                "street1": kwargs.get("street") or "",
                "city": kwargs.get("city") or "",
                "state": kwargs.get("state") or "",
                "postal_code": kwargs.get("postal_code") or ""
            }]
            is_valid = True

        if not is_valid:
            return None

        result = self.post(payload=payload)
        return result.get("payload", [])


search_org = OrgSearch()


def callback_to_state(event: gr.SelectData, state: gr.State) -> Tuple[List[Any], int]:
    """Handles DataFrame `select` events. Updates the internal state for either the recipient or funder based on
    selection. Also sends along the selected Candid entity ID to the proposal generation tab.

    Parameters
    ----------
    event : gr.SelectData
    state : gr.State

    Returns
    -------
    Tuple[List[Any], int]
        (Updated state, Candid entity ID)
    """

    row, _ = event.index

    if len(state) == 0:
        return [], None

    # the state should be a nested list of lists
    # if the state is a single list with non-list elements then we just want a pass-through
    if all(isinstance(s, list) for s in state):
        return state[row], state[row][0]
    return state, state[0]


def lookup_organization(
    name: str,
    ein: Optional[str] = None,
    # street: Optional[str] = None,
    city: Optional[str] = None,
    state: Optional[str] = None,
    postal: Optional[str] = None,
) -> Tuple[List[List[str]], List[List[str]]]:
    """Performs a simple search using the CDS organization search API. Results are sent to the DataFrame table and also
    populate the state for the recipient information.

    Parameters
    ----------
    name : str
        Org name
    ein : Optional[str], optional
        Org EIN, by default None
    street : Optional[str], optional
        Street address, by default None
    city : Optional[str], optional
        Address city, by default None
    state : Optional[str], optional
        Address state, by default None
    postal : Optional[str], optional
        Address postal code, by default None

    Returns
    -------
    Tuple[List[List[str]], List[List[str]]]
        (recip data, recip data)

    Raises
    ------
    gr.Error
        Raised if not enough information was entered to run a search
    gr.Error
        Raised if no search results were returned
    """

    results = search_org(name=name, ein=ein, city=city, state=state, postal=postal)
    if results is None:
        raise gr.Error("You must provide a name, and either an EIN or an address.")
    if not results:
        raise gr.Error("No organizations could be found. Please refine the search criteria.")

    data = []
    for applicant_data in results:
        address = applicant_data.get("addresses", [{}])[0].get("normalized")
        seal = (applicant_data.get("current_seal", {}) or {}).get("image")

        record = [
            applicant_data.get('candid_entity_id'),
            applicant_data.get('main_sort_name'),
            address
        ]

        if seal:
            record.append(f"![{applicant_data.get('current_seal', {}).get('seal')}]({seal})")
        else:
            record.append("")

        data.append(record)
    return data, data


def render(org_id_element: Optional[gr.Blocks] = None) -> Tuple[gr.Blocks, gr.State]:
    """Main blocks build and render function.

    Parameters
    ----------
    org_id_element : Optional[gr.Blocks], optional
        Callback Gradio element, by default None

    Returns
    -------
    Tuple[gr.Blocks, gr.State]
        (component, selected org state)
    """

    with gr.Blocks() as component:
        org_data = gr.State([])
        selected_org_data = gr.State([])

        with gr.Row():
            with gr.Column(scale=2):
                name = gr.Textbox(label="Name of organization", lines=1)
                ein = gr.Textbox(label="EIN of organization", lines=1)
            with gr.Column(scale=3):
                with gr.Group():
                    with gr.Row():
                        with gr.Column():
                            # street = gr.Textbox(label="Street address", lines=1)
                            city = gr.Textbox(label="City", lines=1)
                        with gr.Column():
                            state = gr.Textbox(label="State/province", lines=1)
                            postal = gr.Textbox(label="Postal code", lines=1)

        search_button = gr.Button("Find organization", variant="primary")
        org_info = gr.DataFrame(
            label="Organizations",
            type="array",
            headers=["Candid ID", "Name", "Address", "Seal"],
            col_count=(4, "fixed"),
            datatype=["number", "str", "str", "markdown"],
            wrap=True,
            column_widths=["20%", "30%", "30%", "20%"]
        )

        if org_id_element is None:
            org_id_element = gr.Textbox(label="Selected Candid entity ID", lines=1)

        # pylint: disable=no-member
        search_button.click(
            fn=lambda name, ein, city, state, postal: lookup_organization(
                name=name,
                ein=ein,
                # street=street,
                city=city,
                state=state,
                postal=postal
            ),
            # inputs=[name, ein, street, city, state, postal],
            inputs=[name, ein, city, state, postal],
            outputs=[org_info, org_data],
            api_name=False,
            show_api=False
        )

        # pylint: disable=no-member
        org_info.select(
            fn=callback_to_state,
            inputs=org_data,
            outputs=[selected_org_data, org_id_element],
            api_name=False,
            show_api=False
        )
    return component, selected_org_data