tensorlakeai / indexify

A realtime serving engine for Data-Intensive Generative AI Applications
https://docs.tensorlake.ai
Apache License 2.0
926 stars 120 forks source link

Fix graph local exec #1014

Closed diptanu closed 2 weeks ago

diptanu commented 2 weeks ago

Context

Allows Router Functions to be the start node of a graph


@indexify_function()
def add_two(x: Sum) -> int:
    return x.val + 2

@indexify_function()
def add_three(x: Sum) -> int:
    return x.val + 3

@indexify_router()
def route_if_even(x: Sum) -> List[Union[add_two, add_three]]:
    print(f"routing input {x}")
    if x.val % 2 == 0:
        return add_three
    else:
        return add_two

graph = Graph(name="test_router", description="test", start_node=route_if_even)
graph.route(route_if_even, [add_two, add_three])
graph = RemoteGraph.deploy(graph)
invocation_id = graph.run(block_until_done=True, x=Sum(val=2))
output = graph.output(invocation_id, "add_three")
self.assertEqual(output, [5])

What

The server didn't care whether the first function was regular function or router. Relaxed some constraints on the SDK to allow router functions as start node.

Testing

Contribution Checklist