Repository: microsoft/TypeChat
File: python/examples/math/program.py (Python)

Code before:

ResultReference = TypedDict(
    "ResultReference", {"@ref": Annotated[int, Doc("Index of the previous expression in the 'steps' array")]}
)

FunctionCall = TypedDict(
    "FunctionCall",
    {
        "@func": Annotated[str, Doc("Name of the function")],
        "@args": NotRequired[Annotated[list["Expression"], Doc("Arguments for the function, if any")]],
    },
)

JsonValue = str | int | float | bool | None | dict[str, "Expression"] | list["Expression"]
Expression = JsonValue | FunctionCall | ResultReference

JsonProgram = TypedDict("Program", {"@steps": list[FunctionCall]})


async def evaluate_json_program(program: Mapping[str, Any], onCall: Callable[[str, list[Any]], Awaitable[Any]]) -> Any:
    results: list[Any] = []

    async def evaluate_array(array: list[Any]) -> list[Any]:
        return await asyncio.gather(*[evaluate_call(e) for e in array])

    async def evaluate_object(expr: FunctionCall):
        if "@ref" in expr:
            index = expr["@ref"]
            if index < len(results):
                return results[index]

        elif "@func" in expr and "@args" in expr:
            function_name = expr["@func"]
            return await onCall(function_name, await evaluate_array(expr["@args"]))

        elif isinstance(expr, collections.abc.Sequence):
            return await evaluate_array(expr)

        else:
            raise ValueError("This condition should never hit")

    async def evaluate_call(expr: FunctionCall):
        # implement awaitable pattern
        if isinstance(expr, int) or isinstance(expr, float) or isinstance(expr, str):
            return expr
        return await evaluate_object(expr)

    for step in program["@steps"]:
        results.append(await evaluate_call(step))

    if len(results) > 0:

Review comment (suggestion):
would be good to define the return types expected for this function like maybe Union[List[str], ... , None]:

Code after:

ResultReference = TypedDict(
    "ResultReference", {"@ref": Annotated[int, Doc("Index of the previous expression in the 'steps' array")]}
)

FunctionCall = TypedDict(
    "FunctionCall",
    {
        "@func": Annotated[str, Doc("Name of the function")],
        "@args": NotRequired[Annotated[list["Expression"], Doc("Arguments for the function, if any")]],
    },
)

JsonValue = str | int | float | bool | None | dict[str, "Expression"] | list["Expression"]
Expression = JsonValue | FunctionCall | ResultReference

JsonProgram = TypedDict("Program", {"@steps": list[FunctionCall]})


async def evaluate_json_program(program: JsonProgram, onCall: Callable[[str, Sequence[Expression]], Awaitable[Expression]]) -> Expression | Sequence[Expression]:
    results: list[Expression] | Expression = []

    async def evaluate_array(array: Sequence[Expression]) -> Sequence[Expression]:
        return await asyncio.gather(*[evaluate_call(e) for e in array]) # type: ignore

    async def evaluate_object(expr: FunctionCall):
        if "@ref" in expr:
            index = expr["@ref"]
            if index < len(results):
                return results[index]

        elif "@func" in expr and "@args" in expr:
            function_name = expr["@func"]
            return await onCall(function_name, await evaluate_array(expr["@args"]))

        elif isinstance(expr, collections.abc.Sequence):
            return await evaluate_array(expr)

        else:
            raise ValueError("This condition should never hit")

    async def evaluate_call(expr: FunctionCall) -> Expression | Sequence[Expression]:
        if isinstance(expr, int) or isinstance(expr, float) or isinstance(expr, str):
            return expr
        return await evaluate_object(expr)

    for step in program["@steps"]:
        results.append(await evaluate_call(step)) # type: ignore

    if len(results) > 0:
        return results[-1]