Add Sampling Strategies and Requirements to Generative Slots

[jakelorocco]

Generative Slots: Adding Sampling / Requirements

Base Decorator for Async and Sync Functions

We want the decorator to be something of the type:

P = ParamSpec('P')
R = TypeVar('R')

def generative(func: Callable[P,R]) -> GenerativeSlot[P, R]:
    ...

Moreover, we want GenerativeSlots to fit the calling conventions of other session and mfunc function signatures. This means we would expect the decorated function to do one of the below:

1. func(MelleaSession) -> R
2. func(Context, Backend) -> tuple[R, Context]

The issue is that for async functions, R is not the return type but rather Coroutine[Any, Any, R]. This makes sense because async functions return an awaitable that returns R. For example:

async def test(num: int) -> int:
    ...

When passing this function into a GenerativeSlot[P, R], R becomes Coroutine[Any, Any, int]. This becomes an issue when passing Context, Backend as parameters into the function; instead of Coroutine[Any, Any, tuple[int, Context]], python expects the return type to be tuple[Coroutine[Any, Any, int], Context].

Our generative slots cannot correctly implement this function return type since we only have the context after the generation has been completed. This means we can't meet the default return type python expects (tuple[Coroutine[Any, Any, int], Context]).

(Note: This python-expected syntax also leads to clunky interactions with async generative slots. For example,

@generative
async def test(num: int) -> int:
    ...

original_return_type, context = test(Context, Backend, num=1)
original_return_type = await original_return_type # You must individually await the actual value.

)

This alone isn't an issue: we can overload the @generative decorator to give it the correct return type (Coroutine[Any, Any, tuple[int, Context]]), which is what we currently do. This gives us the following behavior:

@generative
async def test(num: int) -> int:
    ...

original_return_type, context = await test(Context, Backend, num=1)

Adding Requirements and Sampling Strategies

To add requirements and sampling strategies at the level of @generative, we have to define a function that returns a decorator. This is because @generative(reqs=..., strategy=...) and @generative have to act differently (the first creates a new decorator with reqs and strategy in its closure; the second is the decorator).

This is where the issue arises: we cannot properly overload / type hint this "meta-decorator" in a way that correctly supports both async and sync genslots. We can't correctly specify the return type in a way that type hinters can correctly infer whether the decorated function is async or not.

Proposed Solution

Have @generative add in parameters for requirements and sampling the same way we do for MelleaSessions and Contexts/Backends now:

@generative
def test(num: int) -> int:
    ...

test(m=session, requirements=["req1"], strategy=RejectionSamplingStrategy())

Requirements would default to None and strategy would defaul to RejectionSamplingStrategy(loop=2) (the same as instructions).

Pros: Simplest and allows clearly specifying these parameters during each call.
Cons: May result in a lot of duplicate code to write out the requirements and sampling strategy.

Alternate Solutions

Add a Second Decorator

We could also create a new decorator @add_requirements_and_sampling that goes with the @generative decorator:

@add_requirements_and_sampling(requirements=["req1"], strategy=RejectionSamplingStrategy(loop=5))
@generative
def test(num: int) -> int:
    ...

test(m=session)

Pros: You get the one-time requirement / sampling strategy definition desired.
Cons:

1. The double decorator pattern is a bit weird for this use case since @generative is always required if you use the new decorator.
2. It is difficult to allow this approach and allow setting these parameters for each function call. There's no way to specify not_given for a single function call so we won't be able to tell if the user wants to explicitly set the strategy / requirements to None, or just didn't provide them.

Change the Return Signature

We could change our return signature to match what Python expects. If we didn't have to overload and type-hint that aspect of the interface, we could correctly modify the @generative decorator to take arguments like @generative(reqs, strategy).

This means

@generative
def test() -> int

will ultimately have a return type of tuple[Coroutine[Any, Any, int], Context].

Pros: We can implement @generative to behave like we originally wanted.
Cons:

1. We have to change our user interface for sessions and mfuncs. For the async version of these functions, users will now have to manually call modeloutputthunk.avalue() or modeloutputthunk.astream().
2. The async generative slot interface actually becomes slightly less intuitive. To access the values, you would have to do something like:

original_return_type, context = test()
original_return_type = await original_return_type

This also means that the interface is different for genslots than the rest of the mfuncs / session functions since we still have to call await mot.avalue() inside the generative slot to be able to perform a pydantic validation.

Don't Return Context

If we don't return the context object when Context, Backend is passed in, we can also overload the @generative decorator to accept parameters.

Pros: We keep the desired implementation behavior.
Cons: Users can no longer get the context values for the results of generative slots if they use the Backend + Context parameters.