Deserialize return_value of task result back to Pydantic model?

[tnychn]

Since I couldn't find any docs about Taskiq's Pydantic integration/behavior,

Suppose I have the following task:

@broker.task
async def analyze_audio(
    audio_b64: str,
    transcript: Transcript,
    context: Annotated[Context, TaskiqDepends()],
) -> Analysis | None:
    audio = base64.b64decode(audio_b64)
    return await context.state.pipeline(audio, transcript)

where both Transcript and Analysis are Pydantic models.

When retrieving the result, broker.result_backend.get_result returns TaskiqResult and TaskiqResult.return_value is a dict.
So I often need to do an extra step of model_validate on it.

task_result = await broker.result_backend.get_result(str(id))
value = task_result.return_value  # dict or None
if value is None:
    return None
return Analysis.model_validate(value)

Taskiq can handle Pydantic models for task arguments. I expect it can handle that for task return value as well.

Is it possible to support this Taskiq? This seems to be a trivial use case and the inconsistent behavior causes confusion.

[adiberk]

Just curious - what do you do usually with the backend result? I am just curious about the use case. If whatever task you offloaded finished, what would you need to do with the finished result?

[tnychn]

Just curious - what do you do usually with the backend result? I am just curious about the use case. If whatever task you offloaded finished, what would you need to do with the finished result?

In my example, I need to retrieve the audio analysis result after the task is finished and return it through my FastAPI server to the client for display.

I've thought of handling storing the result in the task function itself but that seems to be a waste of Taskiq's capability and beats the purpose of a result backend.

[s3rius]

We have support for Pydantic models, and the returned types are parsed when you try to access them.

taskiq/taskiq/task.py

Line 66 in 7ee894f

res.return_value = parse_obj_as(

However, for this code to be triggered, you need to specify the type for AsyncTaskiqTask. We cannot transfer pickled Python objects as that is considered insecure. Instead, we send JSON objects and store information about types on receiving side so it can infer type correctly.

Can you provide a minimal reproducible example demonstrating this behavior? If you use .wait_result(), it should work as expected tho.

[s3rius]

Here's a solution that might work for you.

You can see that we have few type: ignore annotations in this solution. That's because type information is getting stripped internally, so if we want this thing, we have to add these annotations.

async def get_my_result(
    task: AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType],
    task_id: str,
) -> _ReturnType:
    res = await AsyncTaskiqTask(
        task_id=task_id,
        result_backend=task.broker.result_backend,
        return_type=task.return_type,  # type: ignore
    ).wait_result()

    return res.return_value  # type: ignore

Here's an example how to use it:

import asyncio
from typing import ParamSpec, TypeVar
from pydantic import BaseModel

from taskiq.serializers import JSONSerializer
from taskiq_redis import RedisAsyncResultBackend, RedisStreamBroker
from taskiq import AsyncTaskiqDecoratedTask, AsyncTaskiqTask


_FuncParams = ParamSpec("_FuncParams")
_ReturnType = TypeVar("_ReturnType")


broker = RedisStreamBroker(
    url="redis://localhost/0",  # type: ignore
    unacknowledged_batch_size=3,
    xread_count=2,
    maxlen=2_000,
).with_result_backend(
    RedisAsyncResultBackend(
        serializer=JSONSerializer(),
        redis_url="redis://localhost/1",  # type: ignore
        result_ex_time=60 * 5,  # 5 minutes
    )
)


class InModel(BaseModel):
    arg1: str


class OutModel(BaseModel):
    arg2: str


@broker.task
async def testtask(_in: InModel) -> OutModel:
    return OutModel(arg2=_in.arg1)


async def get_my_result(
    task: AsyncTaskiqDecoratedTask[_FuncParams, _ReturnType],
    task_id: str,
) -> _ReturnType:
    res = await AsyncTaskiqTask(
        task_id=task_id,
        result_backend=task.broker.result_backend,
        return_type=task.return_type,  # type: ignore
    ).wait_result()

    return res.return_value  # type: ignore


async def main():
    await broker.startup()

    sent = await testtask.kiq(_in=InModel(arg1="test"))

    a = await get_my_result(testtask, sent.task_id)
    print(a)

    await broker.shutdown()


if __name__ == "__main__":
    asyncio.run(main())

[tnychn]

Your example works but it seems that the result type is incorrectly inferred:

For some reason, the type checker sees the result type as wrapped within a CoroutineType even though it's in fact an OutModel instance at runtime.