Automatically store `BinaryContent`s in message history somewhere and turn them into `FileUrl`s

[DouweM]

Description

BinaryContent returned by the model (as in #2970) or uploaded by a user (as in https://ai.pydantic.dev/input/) is typically going to be too large to store in a DB or send to an OpenTelemetry backend, so it'd be good to have a way to specify a store_blob function that turns BinaryContent into a FileUrl, and have that be in the stored message history.

If the file is stored somewhere the model API can access, like a public bucket, the FileUrl will not need be converted back into a BinaryContent when using the message history, but if that's not the case you'd need a load_blob function as well that is given the FileUrl and can check the URL (e.g. starting with file: or x-custom: or https://my-bucket.s3), and download it if necessary.

The functions should probably take the run context so they can access deps to know where to store the data, e.g. based on the user or a region.

#2970 currently adds support for output_type=BinaryImage, and with this feature it could support output_type=ImageUrl as well, with the binary content returned by the model automatically uploaded to blob storage.

Inspired by https://github.com/nwcell/pydantic-ai-stash by @nwcell.

References

No response

[nwcell]

Sweet!

Store/load blob would work well as a lambda pair. For example, I have an ArtifactService in my deps that is tightly couple w/ my ORM. I'd reccomend that the store/load functions be set to return the morphed part, along with a key of sorts. I ended up stuffing vendor_metadata for this, but I felt dirty doing so (might make sense to have a dedicated attr in the BinaryContent/FileUrl types here). For me, this made sense, because some use cases will want to dedupe assets, but others won't, so allowing the function to control deduping w/ IDs made the most sense.

Helpers on applying store/load functions to a message tree would also be pretty significant, as the looping required to dynamically manage message parts is a bit stressful.

[DouweM]

@nwcell Would it be possible to encode that key into a custom URL?

Helpers on applying store/load functions to a message tree would also be pretty significant, as the looping required to dynamically manage message parts is a bit stressful.

Agreed

[moritzwilksch]

Very interested in this feature! I currently serialize all context after a rollout for later processing and have to strip out all base64-encoded images which is hacky and loses the reference to the original file. As a workaround, I subclassed ImageURL at some point base on how its attributes are being accessed in pydantic_ai, but thats gnarly:

@dataclass
class MyStoragelayerBackedURL(ImageUrl):  # inherits from pydantic_ai.messages.ImageUrl
    file_id: str = field(init=False, repr=False)
    image_loader: ImageLoader

    @property
    def url(self) -> str:
        resolved = f"{self.image_loader.load(self.file_id)}"
        print(f"resolved url: {resolved}")
        return resolved

    @url.setter
    def url(self, value: str) -> None:
        self.file_id = value

Helpers on applying store/load functions to a message tree

Yes, this sounds amazing, I'd want to be able to write my own file loader to load from any storage backend.

[nwcell]

@DouweM using the url as a key could work for some common scenarios, but I'd be worried about tracking which contents are mutated (some images might be and which db item they were mutated too.

I can think of a few other use cases where the url might not be the best identifier. Deduping logic could get messy when cleaning up old artifacts. Signing/salting urls, if the user wants to do some sort of last mile security would be a pain to manage, especially if you are keeping messages in something structured, such as denormalized SQL or a indexed json store (I'm using JSONB).

We could get away from adding g to the core specs through docs, maybe encourage the developer to extend the content objects if they need custom flags?

Also, the more I think about this, it might be a more robust process to have this positioned as a more flexible content mutator, that contains its own filtering logic? Not sure how I feel about mixing filtering & mutation logic...

For example, with the core use case the user might want to mutate all ImageUrls to BinaryContent when sending to the agent. They'd make the lambda check the content type and mutate if true. This would empower the user to have an easy hook to fully control the contents, allowing for maximum use cases.

Bonus benefits,this would support a wider range of use cases that we aren't focused on for image support.

EG.

• The user could inject a signing process to self hosted URLs.
• They'd user could make the decision if they want to have their own unique ID and either cache it into vendor, or extend the target class (making it not a core component)

Thoughts? (Will think more on this... need to hop on a meeting)

[nwcell]

@DouweM where are you thinking these features would live, exactly, from the developers perspective?

In the run agent commands (agent.run_sync(..., store=list[foo], load=list[bar]))? Are there other places you have in mind?

Would it make sense to add in a message list helper as well (mutate_message_contents(messages=list[ModelMessage], mutators=[foo, bar]))? Sounds like @moritzwilksch is on the same page here w/ a message list mutator.

If this is the direction we want to go, then my question is if we want to have filtering logic live in the mutator functions.

• Pro: Minify interfaces for developer & simplify logic. Docs & instruction could guide developers. Maximum control for unforeseen UCs.
• Con: Mutators would touch all contents, all the time... even if we don't want to mutate them.

[DouweM]

@nwcell We already have history_processors on Agent: https://ai.pydantic.dev/message-history/#processing-message-history, which should cover all generic message mutation needs -- let me know if not.

For this feature, I was imagining something like Agent(..., file_store: AbstractFileStore), so you can implement your own subclass with store(BinaryContent) -> BinaryContent | FileUrl and load(FileUrl) -> FileUrl | BinaryContent methods, and we'd have a FileStore(store_func, load_func) for convenience. load_func would be optional if the FileUrl returned by the store_func represents a public URL that doesn't need to be loaded. All of them would take RunContext so you can use deps to determine what to do.

With that, result.all_messages() should already return a list ready for storage in a DB, and I imagine we'd add helpers to apply storing or loading to an existing list[ModelMessage]. Possibly those could be methods on AbstractFileStore.

---

using the url as a key could work for some common scenarios, but I'd be worried about tracking which contents are mutated (some images might be and which db item they were mutated too.

Perhaps naively, I think all of that could be captured in a URI with a custom scheme like x-my-file:, which can have whatever you like in the path or query string. The load function can then check the URI scheme, look up the file, and either load the data or generate a signed link to S3 or something.

But adding a freeform metadata field to BinaryContent and FileUrl to keep track of additional data would be a good option as well, and would save you having to parse the URI. You'd likely still want to use a custom URI scheme to easily tell apart regular URLs from file-store-managed ones.

[nwcell]

I like the idea of a filestore object. It's cleaner to manage & the logic can get a bit complex, especially, since file objects can get a little stateful. Also, a meta attribute would be fantastic. It feels a bit weird using vendor_metadata, since the namespace indicates that it's for core vendor stuff.

In pydantic_ai_stash, my adapter module had way more mandatory skaffolding, but most of it could probably be removed, if we could guarantee that the store/load functions were being applied against all parts.

Some Pseudo code I've been fiddling w/ today, around what the end user would do...

# Example of an image specific file store 
class ImageFileStore(AbstractFileStore):
    # Example targeted storage function (image specific)
    def _store_filter(self, content):
        '''Example of how a developer could manage their own logic for filtering content'''
        return isinstance(content, BinaryContent) and content.is_image

    def store(self, ctx, content):
        '''User defined function to store images to artifact store & return a relevant ImageUrl'''
        # In this case, we're using binary_content's is_image to help us filter.  probably a more elegant way tot do this....
        if not self._store_filter(content):
            return content
        
        artifact_service = ctx.deps.artifact_service
        image_artifact = artifact_service.store_image(content)
        image_url = artifact_service.build_file_url(artifact)

        return image_url

    # Example of how someone could add in a layer to convert stashed assets into a binary, for last leg conversion
    def load(self, ctx, content):
        if not isinstance(content, ImageUrl):
            return content
        artifact_service = ctx.deps.artifact_service

        # user can define their own mapping logic...
        id = artifact_service.get_key(content)
        image = artifact_service.get_image(id)
        binary_image_content = artifact_service.build_image_binary(image)
        return binary_image_content


# Images live in artifacts and are tracked in ImageUrl, but are treated as BinaryContent by agent
deps = Deps(artifact_service=ArtifactService())
agent.run_sync(message, deps=deps, message_history=message_list, file_store=FileStore())

If we do this right, it opens us up to more dynamic use cases, like signed urls.

class SignedUniversalFileStore(AbstractFileStore)
    def store(self, ctx, content):
        '''User defined function to store binary to artifact store & return a relevant FileUrl
        In the docs, we'd probably want to have a more complete solution that exposes more of the artifact service logic
        '''
        if not isinstance(content, (BinaryContent, FileUrl)):
            return content

        artifact_service = ctx.deps.artifact_service
        
        # Check to ensure that we actually want to mutate content.  Custom filter logic would go here
        if isinstance(content, BinaryContent):
            artifact = artifact_service.store_binary(content)
        # If we're a stored artifact, store as unsigned
        elif isinstance(content, FileUrl):
            if artifact_service.in_artifacts(content):
                artifact = artifact_service.find_stored_artifact(content)
            else:
                artifact = artifact_service.store_external_url(content)

        file_url = artifact_service.build_internal_file_url(artifact)
        return file_url

    # We don't want to load artifact-ified  FileUrl, but we DO need to spit out a signed link for this run
    def load(self, ctx, content):
        '''Assets are in S3, but we want to presign URLs immediately before sending'''
        # Check to ensure that we actually want to mutate content.  Custom filter logic would go here
        if not isinstance(content, FileUrl):
            return content
        
        artifact_service = ctx.deps.artifact_service
        if not artifact_service.in_artifacts(content):
            return content

        # Make artifact object from BinaryContent
        artifact = artifact_service.get_artifact(content)
        file_url = artifact_service.build_signed_url(artifact)

        return file_url

As for the history processor, topic... history_processors are really handy, but I'm not sure how helpful they are w/ post processing output. Such as when they are already in some persistent storage structure. Probably out of scope for this particular featureset, but a solid official way to manage message lists outside of the agent would be great, even if it was just a well documented example.

[DouweM]

@nwcell Letting the FileStore also process things that already FileUrls is handy, I hadn't considered that before.

As for the history processor, topic... history_processors are really handy, but I'm not sure how helpful they are w/ post processing output. Such as when they are already in some persistent storage structure. Probably out of scope for this particular featureset, but a solid official way to manage message lists outside of the agent would be great, even if it was just a well documented example.

A separate issue on what you'd like to to make history processors more helpful in post processing would be great!

[nwcell]

Should we have a dedicated couple of filter rules, or force filter logic through the load/unload function? One is cleaner, the other is a simpler interface, but would split things up more.

I can see both, but it probably matters more on how we're handling other similar decisions w/ Pydantic AI's source.

[DouweM]

@nwcell What kind of filter rules are you thinking?

[nwcell]

@DouweM primarily filtering for types, sniffing for things (like the metadata payload), and routing to specific load/unload processes.

Most common UC will probably be blindly converting BinaryContent to and from FileUrl bits.

That said, some users will want to convert FileUrl -> FileUrl (the url signing UC).

More logic will also be needed for the routing. Different asset types might be treated differently. The logic on a stored asset is different from a 3rd party asset. Metadata can also be key (do I already have this stored? Can I de-dupe?)

One way is to just dump the responsibility onto the user and give solid examples on how to check for types (my code example above). Maybe use a generic hook name, so we have the ability to make a more robust base class at a later date.

The other would require at minimum 2 filter functions and possibly further scaffolding for routing logic.

Either way, I'd make sure to provide the user an ability to filter for any content type their heard desires.

[DouweM]

@nwcell Are you thinking about part-mapping in general, or specifically about this file storage feature? I was thinking of starting with store and load as arbitrary (RunContext, BinaryContent | FileUrl) -> (BinaryContent | FileUrl) functions, and we could implement convenience features (or classes for e.g. S3 storage) later if desired.

If you're thinking about part-mapping as a type of history processor, I suggest filing a new issue for that. Generic part-mapping could likely also solve the file storage case, but I think this is a specific enough use case to warrant convenience methods and special documentation etc.

[nwcell]

@DouweM I'm only really thinking of use cases around file storage. Mostly trying to think through the edge cases. If we t0ugh everything behind a content = foo, we should be fine for 95% of UCs.

With that said, I think we've found our answer, which is to not get too fancy & allow for the developer to craft whatever store/load logic they want. A solid example or 2 in the docs around a simple UC and a more complex should be enough guidance.

By part mapping (getting into tools and whatnot), I don't have a ton of opinions in that area yet. I'd probably be looking into a message storage object before that (out of scope for this feature).