feat: add stream support to mfuncs and session #395

akihikokuroda · 2026-02-01T19:18:50Z

Misc PR

Type of PR

Bug Fix
New Feature
Documentation
Other

Description

Link to Issue: fix: fix: make asynchronous functions / session functions return uncomputed model output thunks #381

Testing

Tests added to the respective file if code was changed
New code has 100% coverage if code as added
Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)

github-actions · 2026-02-01T19:19:01Z

The PR description has been updated. Please fill out the template for your PR to be reviewed.

mergify · 2026-02-01T19:19:25Z

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)(?:\(.+\))?:

HendrikStrobelt · 2026-02-02T14:13:05Z

Thank you. That looks great. We might still have to create a new issue on how to handle streaming with requirement-checking and sampling strategies. @jakelorocco and @nrfulton

nrfulton · 2026-02-02T17:22:49Z

Thank you. That looks great. We might still have to create a new issue on how to handle streaming with requirement-checking and sampling strategies.

Agreed. While doing this we should think about how much we can stream sampling.

jakelorocco

Will follow up / answer the questions I indicated tomorrow morning. Thank you! This looks good.

jakelorocco · 2026-02-02T21:00:33Z

mellea/stdlib/functional.py

+    await_result: bool = True,
 ) -> tuple[ModelOutputThunk[S], Context] | SamplingResult:


Lets set this to False by default. I think that's likely a better longterm default and the expected behavior. Will confirm with the team tomorrow though. I know this is an api breaking change too, so we will just need to flag it in the release notes.

Others agreed that we should set this to False by default.

This change requires making sure all tests and examples working without failure. It may take some time but it will be done.

jakelorocco · 2026-02-02T21:01:07Z

mellea/stdlib/functional.py

+    await_result: bool = True,
 ) -> tuple[ModelOutputThunk[S], Context] | SamplingResult:


We should be able to add type hinting so that if await_result == True or strategy != None, we return a computed model output thunk.

jakelorocco · 2026-02-02T21:02:07Z

mellea/stdlib/functional.py

+            # ._generate_log should never be None after generation.
+            assert result._generate_log is not None
+            result._generate_log.is_final_result = True
+            generate_logs.append(result._generate_log)


Let me think about this generate_log change. I have to review where they are being used / if they are still being used anywhere.

I believe we can actually get rid of the generate_logs.append(result._generate_log) lines in this function.

jakelorocco · 2026-02-02T21:04:49Z

mellea/core/base.py

+class ComputedModelOutputThunk(ModelOutputThunk[S]):
+    """A `ComputedModelOutputThunk` is a `ModelOutputThunk` that is guaranteed to be computed.
+
+    This subclass provides a clear type distinction between thunks that may need awaiting
+    and those that are already computed. It should be returned from synchronous functions
+    and sampling strategies to indicate that no awaiting is needed.
+
+    Key differences from ModelOutputThunk:
+    - Always initialized with a value (cannot be None)
+    - _computed is always True
+    - Cannot be used for streaming (generation fields are not set)
+    - Provides type safety to indicate "already computed"
+    """
+
+    def __init__(
+        self,
+        value: str,
+        meta: dict[str, Any] | None = None,
+        parsed_repr: S | None = None,
+        tool_calls: dict[str, ModelToolCall] | None = None,
+    ):


I feel like this is a bit clunky to create a computed model output thunk. I don't think we will ever have to create one from scratch. Is it possible to have the constructor just take the regular model output thunk as it's only arg, check that everything is computed, set computed to true, and then just defer all of it's getters to the base model output thunk? Or maybe theres some easier way to do this where we don't have to copy over all of the fields.

jakelorocco · 2026-02-02T21:05:58Z

mellea/stdlib/functional.py


    Returns:
        A (ModelOutputThunk, Context) if `return_sampling_results` is `False`, else returns a `SamplingResult`.
+        Always returns ComputedModelOutputThunk since sync functions must await completion.


Where possible, with the sync functions, can we change the return type hinting to be the computed model output thunk?

jakelorocco · 2026-02-02T21:06:58Z

test/backends/test_tool_argument_validation_README.md

This should be removed.

Signed-off-by: Akihiko Kuroda <akihikokuroda2020@gmail.com>

akihikokuroda force-pushed the stream branch from cfe3031 to 114b35a Compare February 1, 2026 19:20

akihikokuroda marked this pull request as draft February 1, 2026 23:58

akihikokuroda changed the title ~~feat: add stream support~~ [WIP] feat: add stream support Feb 1, 2026

akihikokuroda changed the title ~~[WIP] feat: add stream support~~ feat: add stream support Feb 2, 2026

akihikokuroda marked this pull request as ready for review February 2, 2026 01:13

HendrikStrobelt requested a review from jakelorocco February 2, 2026 14:11

nrfulton changed the title ~~feat: add stream support~~ feat: add stream support to mfuncs and session Feb 2, 2026

akihikokuroda force-pushed the stream branch from b78f882 to 524d2a7 Compare February 2, 2026 19:41

jakelorocco reviewed Feb 2, 2026

View reviewed changes

akihikokuroda added 10 commits February 3, 2026 15:46

add ComputedModelOutputThunk

9081564

Signed-off-by: Akihiko Kuroda <akihikokuroda2020@gmail.com>

additional

146b002

Signed-off-by: Akihiko Kuroda <akihikokuroda2020@gmail.com>

additional

eff2e6f

Signed-off-by: Akihiko Kuroda <akihikokuroda2020@gmail.com>

stream support

8b51454

Signed-off-by: Akihiko Kuroda <akihikokuroda2020@gmail.com>

CI test error fix

3ba29e8

Signed-off-by: Akihiko Kuroda <akihikokuroda2020@gmail.com>

CI test error fix

e46bbb0

Signed-off-by: Akihiko Kuroda <akihikokuroda2020@gmail.com>

CI test error fix

5307bcc

Signed-off-by: Akihiko Kuroda <akihikokuroda2020@gmail.com>

fix examples

352556d

Signed-off-by: Akihiko Kuroda <akihikokuroda2020@gmail.com>

default change

bd96f65

Signed-off-by: Akihiko Kuroda <akihikokuroda2020@gmail.com>

review comments

efd8d6a

Signed-off-by: Akihiko Kuroda <akihikokuroda2020@gmail.com>

akihikokuroda force-pushed the stream branch from 524d2a7 to efd8d6a Compare February 3, 2026 20:47

		await_result: bool = True,
		) -> tuple[ModelOutputThunk[S], Context] \| SamplingResult:

feat: add stream support to mfuncs and session #395

Are you sure you want to change the base?

feat: add stream support to mfuncs and session #395

Uh oh!

Conversation

akihikokuroda commented Feb 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Misc PR

Type of PR

Description

Testing

Uh oh!

github-actions bot commented Feb 1, 2026

Uh oh!

mergify bot commented Feb 1, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merge Protections

🟢 Enforce conventional commit

Uh oh!

HendrikStrobelt commented Feb 2, 2026

Uh oh!

nrfulton commented Feb 2, 2026

Uh oh!

jakelorocco left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

akihikokuroda commented Feb 1, 2026 •

edited

Loading

mergify bot commented Feb 1, 2026 •

edited

Loading