Function _socket_name_mapping

burr/integrations/haystack.py:291–318

Map socket names to a single socket name. In Haystack, components communicate via sockets. A socket called "embedding" in one component can be renamed to "query_embedding" when passed to another component. In Burr, there is a single state object so we need a mapping to resolve

(sockets_connections: list[tuple[str, str]])

def _socket_name_mapping(sockets_connections: list[tuple[str, str]]) -> dict[str, str]:
    """Map socket names to a single socket name.

    In Haystack, components communicate via sockets. A socket called
    "embedding" in one component can be renamed to "query_embedding" when
    passed to another component.

    In Burr, there is a single state object so we need a mapping to resolve
    that `embedding` and `query_embedding` point to the same value. This function
    creates a mapping {socket_name: state_field} to rename sockets when creating
    the Burr `Graph`.
    """
    all_connections: dict[str, set[str]] = {}
    for from_, to in sockets_connections:
        if from_ not in all_connections:
            all_connections[from_] = {from_}
        all_connections[from_].add(to)

        if to not in all_connections:
            all_connections[to] = {to}
        all_connections[to].add(from_)

    reduced_mapping: dict[str, str] = {}
    for key, values in all_connections.items():
        unique_name = min(values)
        reduced_mapping[key] = unique_name

    return reduced_mapping


def _connected_inputs(pipeline) -> dict[str, list[str]]: