Add a callback to easily track data

[ekaj2]

Description

Right now we're overriding the extract and call methods of the OpenAIExtractor and OpenAICall classes to handle things like storing in a database or tracking usage. Would be nice to get a standardized way to do this and manage serialization automatically.

class SavedOpenAICall(OpenAICall):
    def call(self, **kwargs: Any) -> OpenAICallResponse:
        response = super().call(**kwargs)
        prompt_dump = self.dump() | response.dump()
        label = self.__class__.__name__
        // process data
        return response


T = TypeVar("T", bound=ExtractedType)


def make_serializable(val):
    """
    Attempt to JSON serialize an object. Convert to string if not directly serializable.
    """
    try:
        json.dumps(val)
        return val
    except (TypeError, ValueError):
        return str(val)


def serialize_dict(d):
    """
    Recursively process a dictionary, making all values JSON serializable.
    """
    for key, value in d.items():
        if isinstance(value, dict):
            d[key] = serialize_dict(value)
        else:
            d[key] = make_serializable(value)
    return d


class SavedOpenAIExtractor(OpenAIExtractor[T]):
    def extract(self, retries: int = 0, **kwargs: Any) -> T:
        response = super().extract(retries, **kwargs)
        prompt_dump = serialize_dict(self.dump() | {"response": response})
        label = self.__class__.__name__
        // process data
        return response

Ref (conversation w/ William): https://mirascope-community.slack.com/archives/D0736LQ50E9/p1715651888713359

[willbakst]

My thoughts here:

1. Likely want to set this at the class level when defining a subclass. This will ensure (as in the example) that all calls use the same serialization method set at the class level.
2. If you want every single call to use the same serialization, this would still likely require defining your own subclass as you already have (e.g. SerializingOpenAIExtractor); however, instead of overwriting a specific method, you'll just set the class variable and we'll call the method internally for all methods so you don't have to worry about that.

My first instinct would be something like this:

# this would match the finalized callback type
def my_serializer(...):
    ....

class SerializingOpenAICall(OpenAICall):
    serialization_fn = my_serializer

One benefit of doing it this way is that you can create different classes for different serialization schemes + lock the scheme to the particular call, which is likely necessary if you're serializing to a DB or something.

Need to further flesh this out and play around with things, but I like the feature and general direction. Thanks for the request!

[willbakst]

After further thought and discussion, we're thinking that this feature is actually quite similar to something else we're already working on wrt. ops tooling. Our ops integrations first wrap each method and then (optionally) wrap the LLM client calls themselves internally.

I think that the method wrapper functionality would actually cover this callback feature (and the option to add additional LLM client specific wrappers would provide even more customizability if desired).

For reference, this would end up looking something like this:

# we will include documentation for how to write this function
# this is currently what needs the most additional thought for optimal DX
def saving_fn(...):
    ...
    
def with_saving(cls):
    return wrap_mirascope_class_functions(cls, saving_fn)
    # ^ you would import this (function name tbd) helper utility from mirascope

@with_saving
class SavingCall(OpenAICall):
    """A base call that wraps methods and saves input/output."""

class MyCall(SavingCall):
    ...

There are still aspects of this that need to be further fleshed out (namely wrt. DX), but I like this direction.

[willbakst]

This is now released in v0.14.0! 🎉

@ekaj2 you can find documentation for how to use this change for callbacks here

If you have any questions / run into any issues / the docs aren't clear let us know!

[ekaj2]

Amazing, thanks Will I'll check it out soon