Subscriber with Pydantic BaseModel argument treated as a single field's type rather than the whole body's type

[jonathansick]

I'm new to FastStream and getting started with adding a FastStream-based (aiokafka) Kafka consumer into a FastAPI app. I have another FastStream app that's publishing messages serialized to JSON from a Pydantic BaseModel.

In my consumer app, I can successfully do the following where I treat the message body as a string and then manually deserialize into the Pydantic model:

"""Kafka router and consumers."""

from __future__ import annotations

from faststream.kafka.fastapi import KafkaRouter
from faststream.security import BaseSecurity
from rubin.squarebot.models.kafka import SquarebotSlackMessageValue
from structlog import get_logger

from ..config import config

__all__ = ["kafka_router", "handle_slack_message"]


kafka_security = BaseSecurity(ssl_context=config.kafka.ssl_context)
kafka_router = KafkaRouter(
    config.kafka.bootstrap_servers, security=kafka_security
)


@kafka_router.subscriber(
    config.message_channels_topic,
    group_id=config.consumer_group_id,
)
async def handle_slack_message(
    msg: str,
) -> None:
    """Handle a Slack message."""
    logger = get_logger(__name__)
    message = SquarebotSlackMessageValue.model_validate_json(msg)
    logger.info("Slack message text", text=message.text)

(the SquarebotSlackMessageValue Pydantic model is over here but its just a fairly straightforward BaseModel)

But I understand I should be able to annotate the subscriber's argument with the Pydantic model type and have it automatically deserialize into the SquarebotSlackMessageValue, like this:

@kafka_router.subscriber(
    config.message_channels_topic,
    group_id=config.consumer_group_id,
)
async def handle_slack_message(
    msg: SquarebotSlackMessageValue,
) -> None:
    """Handle a Slack message."""
    logger = get_logger(__name__)
    logger.debug("Received a Slack message", text=msg.text)

When I do this, though, I get a Pydantic validation error saying that the "msg" field is missing the contents of the SquarebotSlackMessageValue model. Essentially it's treating msg as a single key in the Kafka message body, rather than representing the whole JSON structure of the message body:

2024-02-15 21:40:21,768 ERROR    - lsst.square-events.squarebot.slack.message.channels | unfurlbot | 4-17080332 - RequestValidationError: [{'type': 'missing', 'loc': ('body', 'type'), 'msg': 'Field required', 'input': {'msg': '{"type":"message","channel":"D058L2LN1BN", [truncated....]

Are there some things I should try to make sure that FastStream is treating the subscriber function argument that's annotated as a Pydantic BaseModel type as representing the whole message body, and not just a field?

This discussion as touches on this topic, but my reading is that a single argument annotated as the Pydantic model's type should do what I want.

I'm running in Python 3.12.2 with these key library versions:

aiokafka==0.10.0
fastapi==0.109.2
faststream[kafka]==0.4.2
pydantic==2.6.1

Thanks for looking at this. FastStream is really cool and I've been overall quite happy with how it simplifies setting up Kafka app within our FastAPI apps. 😄

[Lancetnik]

With such annotation, your message body should be SquarebotSlackMessageValue type, not { msg: SquarebotSlackMessageValue }

To consume this one, u should create extra Pydantic model like this

class Data(BaseModel):
    msg: SquarebotSlackMessageValue

And use it as payload annotation

[jonathansick]

Thanks for this reply. I did some more digging on the issue and realized that I had a couple issues on the producer side.

In my producer, I was first manually serializing the Pydantic object to JSON, and then also mis-spelled the content type header in upper case, i.e.:

await self.publisher.publish(
    message=value.model_dump_json(),
    key=key.to_key_bytes(),
    headers={"Content-Type": "application/json"},  # should be "content-type"
)

But I realized I can let FastStream serialize to JSON (and I can probably let it set the content-type too):

await self.publisher.publish(
    message=value,  # a BaseModel type
    key=key.to_key_bytes(),
    headers={"content-type": "application/json"},
)

With this, my subscriber can now use that BaseModel type (SquarebotSlackMessageValue) as the type on the argument and get the whole message value as a parsed and validated Pydantic model.

[Lancetnik]

Yes, FastStream is able to publish raw pydantic models (as well as regular python types) and sets content-type header automatically too