Skip to content

refactor(core): split MOT internal state into _call / _gen sub-objects #1216

Description

@ajbozarth

Sub-issue of #909. Design discussed in #1197 Area 3.

Summary

The seventeen private attributes on ModelOutputThunk (post-#1181) split naturally into two groups already enforced implicitly by __deepcopy__. This issue makes that split structural with two internal sub-objects: mot._call (originating-call data, preserved on copy) and mot._gen (in-flight machinery, reset on copy).

Motivation

ModelOutputThunk.__init__ currently sets seventeen flat private attributes. __copy__ and __deepcopy__ both enumerate them field-by-field, deciding per-field whether to preserve or reset. Adding a new field today means three coordinated edits and a maintainer reading inspection comments to know which group it belongs to. The split is real; the code just doesn't show it.

Design

@dataclass
class _CallInfo:
    """Originating-call data. Preserved across copies because retries and
    sampling routinely need it."""
    action: Component | CBlock | None = None
    context: list[Component | CBlock] | None = None
    model_options: dict[str, Any] | None = None
    generation_id: str | None = None  # added by #1181

@dataclass
class _GenerationState:
    """In-flight computation machinery. Reset to a fresh empty instance on
    __copy__ / __deepcopy__ — a copied MOT is a distinct (non-generating)
    object and must not share queues, tasks, or thread signals."""
    queue: asyncio.Queue = field(default_factory=lambda: asyncio.Queue(maxsize=20))
    chunk_size: int = 3
    first_chunk_received: bool = False
    generate: asyncio.Task[None] | None = None
    generate_type: GenerateType = GenerateType.NONE
    generate_extra: asyncio.Task[Any] | None = None
    cancel_hook: Callable[[], None] | None = None
    process: Callable[[ModelOutputThunk, Any], Coroutine] | None = None
    post_process: Callable[[ModelOutputThunk], Coroutine] | None = None
    on_computed: Callable[[ModelOutputThunk], Coroutine] | None = None
    start: datetime.datetime | None = None

_computed and _cancelled stay flat — they back is_computed() and the public cancelled property and need to remain cheap reads.

_generate_log is not in _CallInfo — it stays accessible at its current path because Area 5 (#1191) promotes it to public mot.generate_log.

__copy__ / __deepcopy__ collapse to:

copied._call = copy(self._call)              # preserved
# _gen left as fresh _GenerationState() from __init__ — reset on copy

Disallow copying an uncomputed MOT

Both __copy__ and __deepcopy__ raise RuntimeError if self._computed is False. The contract is already implicitly "copies are post-generation" (the docstring says "A copied ModelOutputThunk cannot be used for generation"); raising makes it explicit. Audit shows nothing in mellea/ copies an in-progress MOT today, so blast radius in our code is zero.

Migration

All mot._async_queue, mot._action, etc. accesses inside mellea/ migrate to mot._gen.queue, mot._call.action. Wide internal rename, confined to mellea/ and test/. None of the renamed attributes are read outside mellea/ today, so no aliases needed.

Appendix items folded in

From #1197's Appendix A:

  • Item 8 — Add a "deepcopy of a generating MOT raises" / "deepcopy of a computed MOT does not share _gen state" test in test/core/test_base.py. The current preserve-vs-reset rule is enforceable only by inspection; this PR makes it structural and a test should pin it down.

Coordination with #1013

Phase 2 streaming chunking (#1013) will likely add a parsed-repr stream buffer adjacent to the current text queue. With _gen in place, that's a new field on _GenerationState; without it, another flat private attr to remember in three copy methods. Recommendation: do not block this issue on #1013 — that issue could remain a placeholder for months. Land Area 3 first; #1013 inherits the cleaner shape.

ComputedModelOutputThunk interaction

ComputedModelOutputThunk reassigns __class__ in-place — this issue does not change that contract. _call and _gen are inherited like any other attribute; _gen on a ComputedModelOutputThunk is just dormant (no in-flight task).

Acceptance

  • _CallInfo and _GenerationState defined; MOT stores _call and _gen
  • __copy__ / __deepcopy__ reduce to a deep-copy of _call + fresh _GenerationState
  • __copy__ / __deepcopy__ raise RuntimeError on uncomputed MOT
  • All internal access sites migrated (mellea/ + test/)
  • New test: deepcopy of a generating MOT raises
  • New test: deepcopy of a computed MOT does not share _gen state with the original

Sequencing

Should land after #1181 (which adds _generation_id to the _CallInfo field set). Independent of Area 1 (#1215), Area 4 (#1217), Area 5 (#1191) — order them however release timing prefers.

Metadata

Metadata

Assignees

Labels

area/stdlibCore abstractions: Context, MOT, SamplingStrategy, formatters, serializationp2Medium/low: minor bugs, niche features, polish, docs, tests, cleanup. Scoped, lower urgency.refactor

Type

No type
No fields configured for issues without a type.

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions