wandb/weave

`@weave.op()` does not work with Pydantic `BaseModel` functions or class variables

Closed this issue · 4 comments

@weave.op() failure with BaseModel internal function

The code below fails the the error AttributeError: 'ObjectRecord' object has no attribute 'format'

import weave
from pydantic import BaseModel


class Book(BaseModel):
    title: str
    author: str

    def format(self) -> str:
        return f"{self.title} by {self.author}"


@weave.op()
def my_func(book: Book) -> str:
    return book.format()


weave.init("scratch-test")
book = Book(title="The Name of the Wind", author="Patrick Rothfuss")
print(my_func(book))

@weave.op() failure with BaseModel class variable

The code below fails with the error AttributeError: 'ObjectRecord' object has no attribute 'class_var'

from typing import ClassVar

import weave
from pydantic import BaseModel


class Book(BaseModel):
    title: str
    author: str

    class_var: ClassVar[str] = "class_var"


@weave.op()
def my_func(book: Book) -> str:
    return book.class_var


weave.init("scratch-test")
book = Book(title="The Name of the Wind", author="Patrick Rothfuss")
print(my_func(book))

Would you be open to a PR if I look into fixing this?

Any reference material I should look at before submitting one?

The culprit:

return {k: getattr(obj, k) for k in obj.model_fields}

Looks like this only grabs model fields and doesn't grab any other information, which explains the missing functions and class variables.

Due to how pydantic BaseModel handling is implemented, @weave.op() removes everything except for the fields. This means built-in methods like model_dump are also no longer present.

I will try to update the code to handle porting over everything into the ObjectRecord.

Fixed with #1390

Thanks!!