refactor(add_messages): clearer step-by-step structure

Restructure add_messages so fast vs. slow path, REMOVE_ALL handling,
and format application each live in a single numbered section with a
short lead-in comment. No behavior changes; the previous commits'
optimizations are preserved.

- fold the two path branches into one `if pure_append else slow_path`
  so format handling happens at a single exit instead of being
  duplicated between fast and slow paths
- drop the now-redundant `left_seq` local; index `left` directly after
  the coerce step (with a single `cast(list, left)` for the type
  checker)
- use `set.isdisjoint` for the overlap check

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Sydney Runkle
2026-04-24 07:31:03 -04:00
co-authored by Claude Opus 4.7
parent e601988ab3
commit 4f01115fa8
+50 -46
View File
@@ -184,24 +184,25 @@ def add_messages(
```
"""
# coerce to list
# 1. Coerce scalars to lists.
if not isinstance(left, list):
left = [left] # type: ignore[assignment]
if not isinstance(right, list):
right = [right] # type: ignore[assignment]
left = cast(list, left)
# Optimization 1: skip conversion + ID assignment on left when it already
# contains fully-resolved BaseMessage objects (the common case after the
# first call, since add_messages always returns list[BaseMessage] with IDs).
# 2. Normalize `left`. After the first call, `left` is the previous return
# value of `add_messages` — a list of fully-resolved BaseMessages with
# IDs — and needs no work. Fresh user input (dicts, tuples, message
# chunks, BaseMessages without IDs) falls through to full conversion.
left_msgs: list[BaseMessage]
left_seq = cast(list, left)
if (
left_seq
and isinstance(left_seq[0], BaseMessage)
and left_seq[0].id is not None
and not isinstance(left_seq[0], BaseMessageChunk)
left
and isinstance(left[0], BaseMessage)
and not isinstance(left[0], BaseMessageChunk)
and left[0].id is not None
):
left_msgs = left_seq
left_msgs = left
else:
left_msgs = [
message_chunk_to_message(cast(BaseMessageChunk, m))
@@ -211,12 +212,13 @@ def add_messages(
if m.id is None:
m.id = str(uuid.uuid4())
# always normalise right — it's fresh external input
# 3. Normalize `right`always fresh input. Assign missing IDs and detect
# any RemoveMessage sentinels in a single pass.
right_msgs: list[BaseMessage] = [
message_chunk_to_message(cast(BaseMessageChunk, m))
for m in convert_to_messages(right)
]
remove_all_idx = None
remove_all_idx: int | None = None
has_remove = False
for idx, m in enumerate(right_msgs):
if m.id is None:
@@ -226,51 +228,53 @@ def add_messages(
if m.id == REMOVE_ALL_MESSAGES:
remove_all_idx = idx
# 4. REMOVE_ALL_MESSAGES: discard everything up to and including the sentinel.
if remove_all_idx is not None:
return right_msgs[remove_all_idx + 1 :]
# Optimization 2: pure-append fast path — no removals, no ID overlaps with
# left, and no duplicate IDs within right (all imply a dedup/update is needed).
# 5. Decide fast vs. slow path. The fast path (pure append) is only valid
# when `right` has no removals, no intra-right duplicate IDs, and no IDs
# that overlap with `left` — any of those would force the indexed merge
# below to update or dedup.
pure_append = False
if not has_remove:
left_ids = {m.id for m in left_msgs}
right_id_set = {m.id for m in right_msgs}
if len(right_id_set) == len(right_msgs) and not (right_id_set & left_ids):
result = left_msgs + right_msgs
if format == "langchain-openai":
return _format_messages(result)
elif format:
msg = (
f"Unrecognized {format=}. Expected one of 'langchain-openai', None."
)
raise ValueError(msg)
return result
right_ids = {m.id for m in right_msgs}
pure_append = len(right_ids) == len(right_msgs) and right_ids.isdisjoint(
left_ids
)
# slow path: updates or removals present — full indexed merge
merged = left_msgs.copy()
merged_by_id = {m.id: i for i, m in enumerate(merged)}
ids_to_remove = set()
for m in right_msgs:
if (existing_idx := merged_by_id.get(m.id)) is not None:
if isinstance(m, RemoveMessage):
ids_to_remove.add(m.id)
if pure_append:
merged = left_msgs + right_msgs
else:
# 6. Slow path: build id→index map over `left`, then replay `right`.
# In-place replacement for matching IDs, append for new IDs, and a
# deferred removal pass so RemoveMessages can target either side.
merged = left_msgs.copy()
merged_by_id = {m.id: i for i, m in enumerate(merged)}
ids_to_remove = set()
for m in right_msgs:
if (existing_idx := merged_by_id.get(m.id)) is not None:
if isinstance(m, RemoveMessage):
ids_to_remove.add(m.id)
else:
ids_to_remove.discard(m.id)
merged[existing_idx] = m
else:
ids_to_remove.discard(m.id)
merged[existing_idx] = m
else:
if isinstance(m, RemoveMessage):
raise ValueError(
f"Attempting to delete a message with an ID that doesn't exist ('{m.id}')"
)
merged_by_id[m.id] = len(merged)
merged.append(m)
merged = [m for m in merged if m.id not in ids_to_remove]
if isinstance(m, RemoveMessage):
raise ValueError(
f"Attempting to delete a message with an ID that doesn't exist ('{m.id}')"
)
merged_by_id[m.id] = len(merged)
merged.append(m)
merged = [m for m in merged if m.id not in ids_to_remove]
# 7. Apply optional output format.
if format == "langchain-openai":
merged = _format_messages(merged)
elif format:
return _format_messages(merged)
if format:
msg = f"Unrecognized {format=}. Expected one of 'langchain-openai', None."
raise ValueError(msg)
return merged