Skip to content

Conversation

@jakelorocco
Copy link
Contributor

@jakelorocco jakelorocco commented Sep 30, 2025

Added top-level async functions and async functions to MelleaSessions. Also added a short section to the tutorial. This PR does not support async generative slots; I have it working, but I am trying to find a way for type checkers to be happy. This PR supports async generative slots with a happy type checker.

There is some code duplication between sync and asynchronous versions of each top level function. Unfortunately, some of that seems to be inevitable since all functions end up calling aact (either directly or through an asyncio.run call). We can remove this code duplication, but that would require parameterizing some these functions and introducing room for user errors. I thought it best to leave the duplication for now and in the future we can potentially break out some of the logic to separate functions (like creating an instruction from the component parts, etc...).

Added a warning when calling async functions with a non-Simple context.

Added a session.clone() method to make interacting with contexts easier on a session level. (See feedback below)

Key concerns:

  1. ModelOutputThunk values: We await the mot's value inside the async functions. This means that users get a fully computed mot from aact. This was to maintain consistency with sync functions and sampling strategies.
    a. Sync functions- act requires that we await the mot.avalue() call internally since the mot is being called from a non-async function that cannot await
    b. Sampling strategies- sampling strategies (along with generative slots) require the mot to be fully computed before validating (or for generative slots, applying the pydantic validation). This means that mots returned by aact when using a sampling strategy must be already be computed. If we allowed users to await mot.avalue() this would cause different behavior when sampling strategies are provided regardless of if the SamplingResults or the mot is returned.
  2. Contexts: we set m.ctx = ctx only after we've already awaited the mot's value. In many ways this is good, it means that we will never linearize a context with an uncomputed mot in it. However, it also means there's slightly wonky behavior if using non-Simple contexts with async functions:
out1 = m.ainstruct("1")
out2 = m.ainstruct("2")

_ = await asyncio.gather(out2, out1)

print(out1) # nothing in context
print(out2) # nothing in context for this generation
print(m.ctx) # out2ctx -> root; we lost the ctx for out1

# VS.
out1 = await m.ainstruct("1")
out2 = await m.ainstruct("2")
print(out1) # empty context
print(out2) # out1 in context
print(m.ctx) # out2ctx -> out1ctx -> root

@mergify
Copy link

mergify bot commented Sep 30, 2025

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Enforce conventional commit

Wonderful, this rule succeeded.

Make sure that we follow https://www.conventionalcommits.org/en/v1.0.0/

  • title ~= ^(fix|feat|docs|style|refactor|perf|test|build|ci|chore|revert|release)(?:\(.+\))?:

@jakelorocco jakelorocco marked this pull request as ready for review September 30, 2025 20:56
@HendrikStrobelt
Copy link
Contributor

Looks good.

Couple of thoughts:

  • Would we need a session.clone() function to run mulitple calls on the same ChatContext?
  • Can you think of an example to show the async instruct call streaming out to console? (Mabye fine with MOT not streaming but seeing how generation and validation are happening)

@jakelorocco
Copy link
Contributor Author

jakelorocco commented Oct 1, 2025

Looks good.

Couple of thoughts:

  • Would we need a session.clone() function to run multiple calls on the same ChatContext?

That's a good point. A session.clone() would create a new session that points to the same backend and context but would prevent conflicts with updating the session's ctx when multiple async functions are running. I will implement that. Done.

  • Can you think of an example to show the async instruct call streaming out to console? (Mabye fine with MOT not streaming but seeing how generation and validation are happening)

We currently don't have a great way to do this. In the initial proposal, we talked about attaching a logging callback for these things but it hasn't been implemented yet. Let's talk about this in the mellea meeting today.

@jakelorocco
Copy link
Contributor Author

jakelorocco commented Oct 1, 2025

Added async generative slots. Note, the async nature of the functions don't interfere with the schema generation:

schema for sync genslot:
{'properties': {'result': {'description': 'Result of test_sync', 'title': 'Result', 'type': 'integer'}}, 'required': ['result'], 'title': 'TestSyncResponse', 'type': 'object'}

schema for async genslot:
{'properties': {'result': {'description': 'Result of test_async', 'title': 'Result', 'type': 'boolean'}}, 'required': ['result'], 'title': 'TestAsyncResponse', 'type': 'object'}

In other words, the function doesn't have some wonky return type like Coroutine[Any, Any, int] when used as a pydantic schema.

@jakelorocco
Copy link
Contributor Author

After chatting with Nathan, removed references to sessions' backend stacks since we aren't using them.

@avinash2692 avinash2692 self-requested a review October 6, 2025 19:23
Copy link
Contributor

@avinash2692 avinash2692 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@avinash2692 avinash2692 merged commit 689e1a9 into main Oct 6, 2025
4 checks passed
@jakelorocco jakelorocco deleted the jal/top-level-async branch October 6, 2025 19:25
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants