[BUG] Backlinks cause errors when using documents as part of Pydantic union types

[ldorigo]

Describe the bug
If I have a document containing backlinks; using that document within a pydantic basemodel that uses union types causes an exception. See simple MRO below.

To Reproduce

from typing import Literal
from uuid import uuid4

from beanie import BackLink
from beanie import Document
from beanie import Link
from beanie import init_beanie
from pydantic import UUID4
from pydantic import BaseModel
from pydantic import Field

class Improvement(Document):
    id: UUID4 = Field(default_factory=uuid4)
    issue_id: int
    conversation_messages: list[Link["ConversationMessage"]] = []

class ConversationMessage(Document):
    id: UUID4 = Field(default_factory=uuid4)
    message_type: Literal["message"]  # Placeholder for actual MessageType
    improvements: list[BackLink["Improvement"]] = Field(
        default=[], json_schema_extra={"original_field": "conversation_messages"}
    )

    class Settings:
        is_root = True


class WelcomeMessage(ConversationMessage):
    message_type: Literal["welcome_message"] = "welcome_message"
    contents: str


class EscalationProposalMessage(ConversationMessage):
    message_type: Literal["escalation_proposal"] = "escalation_proposal"


class Answer(BaseModel):
    message: EscalationProposalMessage | WelcomeMessage


await init_beanie(
    client.get_database(),
    document_models=[Improvement, ConversationMessage, WelcomeMessage],
)
t = WelcomeMessage(contents="fwrwerF")
esc = EscalationProposalMessage()
Answer(message=t)

Note that the same happens with a simplified example without inheritance:

l
from uuid import uuid4

from beanie import BackLink
from beanie import Document
from beanie import Link
from beanie import init_beanie
from pydantic import UUID4
from pydantic import BaseModel
from pydantic import Field


class Improvement(Document):
    id: UUID4 = Field(default_factory=uuid4)
    issue_id: int
    conversation_messages: list[Link["WelcomeMessage"] | Link["EscalationProposalMessage"]] = []




class WelcomeMessage(Document):
    message_type: Literal["welcome_message"] = "welcome_message"
    contents: str
    improvements: list[BackLink["Improvement"]] = Field(
        default=[], json_schema_extra={"original_field": "conversation_messages"}
    )


class EscalationProposalMessage(Document):
    message_type: Literal["escalation_proposal"] = "escalation_proposal"
    improvements: list[BackLink["Improvement"]] = Field(
        default=[], json_schema_extra={"original_field": "conversation_messages"}
    )


class Answer(BaseModel):
    message: EscalationProposalMessage | WelcomeMessage


await init_beanie(
    client.get_database(),
    document_models=[Improvement, WelcomeMessage, EscalationProposalMessage],
)
t = WelcomeMessage(contents="fwrwerF")
esc = EscalationProposalMessage()
Answer(message=t)

The error:

{
	"name": "TypeError",
	"message": "'WelcomeMessage' object does not support item assignment",
	"stack": "---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[27], line 46
     44 t = WelcomeMessage(contents=\"fwrwerF\")
     45 esc = EscalationProposalMessage()
---> 46 Answer(message=t)

    [... skipping hidden 1 frame]

File ~/.cache/pypoetry/virtualenvs/learnwise-chat-RkYLlhmr-py3.12/lib/python3.12/site-packages/beanie/odm/documents.py:240, in Document.fill_back_refs(cls, values)
    238 @model_validator(mode=\"before\")
    239 def fill_back_refs(cls, values):
--> 240     return cls._fill_back_refs(values)

File ~/.cache/pypoetry/virtualenvs/learnwise-chat-RkYLlhmr-py3.12/lib/python3.12/site-packages/beanie/odm/documents.py:229, in Document._fill_back_refs(cls, values)
    221             values[field_name] = BackLink[link_info.document_class](
    222                 link_info.document_class
    223             )
    224         if (
    225             link_info.link_type
    226             in [LinkTypes.BACK_LIST, LinkTypes.OPTIONAL_BACK_LIST]
    227             and field_name not in values
    228         ):
--> 229             values[field_name] = [
    230                 BackLink[link_info.document_class](
    231                     link_info.document_class
    232                 )
    233             ]
    234 return values

TypeError: 'WelcomeMessage' object does not support item assignment"
}

Note that the error always happens on the second item in the union type: if we define Answer with message: WelcomeMessage | EscalationProposalMessage, then assigning a WelcomeMessage to messages works, but it breaks when assigning an EscalationProposalMessage.

[github-actions]

This issue is stale because it has been open 30 days with no activity.

[ldorigo]

Not stale still an issue

[github-actions]

This issue is stale because it has been open 30 days with no activity.

[ldorigo]

Not stale still an issue. We're migrating away from beanie because of too many issues like this. Won't un-stale it again, but it just means the bug will remain unadressed if it gets closed automatically next time

[staticxterm]

Hi @ldorigo,
I am unable to reproduce the exact error message you get. Would you mind posting other relevant info, such as the Beanie version and Pydantic versions you were testing this? Also, please provide a fully runnable example.

According to the Relations documentation, you are missing the original_field reference in the ConversationMessage Document model (and for your second example, in the WelcomeMessage and EscalationProposalMessage models). The minimum reproducible example are the Improvement and ConversationMessage Document classes from your first example.