Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
194 changes: 186 additions & 8 deletions python/packages/core/agent_framework/_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -3351,6 +3351,8 @@ def _match_mixed_pause_responses(
host_items_by_occurrence[request.id] = index

matched_content_ids: set[int] = set()
preexisting_response_indexes = {index for index, item in enumerate(items) if item.get("response") is not None}
current_idless_response_indexes: set[int] = set()
for match_idless_host_results in (False, True):
for response in responses:
is_idless_host_result = response.type == "function_result" and response.id is None
Expand Down Expand Up @@ -3391,6 +3393,18 @@ def _match_mixed_pause_responses(
if items[pending_index].get("response") is None
]
if len(unanswered_indexes) == 1:
if any(
pending_index in preexisting_response_indexes
and isinstance(items[pending_index].get("response"), Mapping)
and _same_mixed_pause_response(
cast(Mapping[str, Any], items[pending_index]["response"]),
response.to_dict(),
)
for pending_index in host_items_by_call[response.call_id]
):
raise RuntimeError(
f"Ambiguous id-less Host response for mixed pause call_id {response.call_id!r}."
)
item_index = unanswered_indexes[0]
elif not unanswered_indexes and allow_idless_host_duplicates:
duplicate_indexes = [
Expand All @@ -3404,6 +3418,13 @@ def _match_mixed_pause_responses(
]
if len(duplicate_indexes) == 1:
item_index = duplicate_indexes[0]
elif any(
pending_index in current_idless_response_indexes
for pending_index in host_items_by_call[response.call_id]
):
raise RuntimeError(
f"Conflicting id-less Host responses for mixed pause call_id {response.call_id!r}."
)
else:
continue

Expand All @@ -3418,6 +3439,8 @@ def _match_mixed_pause_responses(
raise RuntimeError(f"Conflicting response for mixed pause occurrence {candidate.id!r}.")
items[item_index]["response"] = candidate_state
matched_content_ids.add(id(response))
if is_idless_host_result:
current_idless_response_indexes.add(item_index)

if any(item.get("response") is None for item in items):
return matched_content_ids, True, [], set()
Expand Down Expand Up @@ -3464,22 +3487,177 @@ def bind_approval_response(response: Content) -> Content | None:
consume=False,
)

responses = [
content
for message in messages
for content in message.contents
if content.type in {"function_approval_response", "function_result"}
def attributed_staged_replay_item_index(
content: Content,
reserved_item_indexes: set[int],
) -> int | None:
if content.type == "function_approval_response":
candidate = bind_approval_response(content)
elif content.type == "function_result":
candidate = content
else:
return None
if candidate is None:
return None

candidate_state = candidate.to_dict()
for item_index, item in enumerate(items):
if item_index in reserved_item_indexes:
continue
stored_response = item.get("response")
if not isinstance(stored_response, Mapping) or not _same_mixed_pause_response(
cast(Mapping[str, Any], stored_response),
candidate_state,
):
continue
request = _content_from_state(item.get("request"))
if request is None:
continue
if item.get("kind") == "approval" and candidate.type == "function_approval_response":
request_ids = {
str(identity)
for identity in (
request.id,
request.function_call.id if request.function_call is not None else None,
)
if identity is not None
}
candidate_ids = {
str(identity)
for identity in (
candidate.additional_properties.get(_APPROVAL_REQUEST_ID_KEY),
candidate.id,
)
if identity is not None
}
if request_ids.intersection(candidate_ids):
return item_index
elif (
item.get("kind") == "host"
and candidate.type == "function_result"
and candidate.call_id == request.call_id
and (candidate.id is None or candidate.id == request.id)
):
return item_index
return None

host_requests = [
request
for item in items
if item.get("kind") == "host"
and item.get("response") is None
and (request := _content_from_state(item.get("request"))) is not None
and request.call_id is not None
]
approval_request_ids = {
str(identity)
for item in items
if item.get("kind") == "approval"
and item.get("response") is None
and (request := _content_from_state(item.get("request"))) is not None
for identity in (
request.id,
request.function_call.id if request.function_call is not None else None,
)
if identity is not None
}
host_call_ids = {request.call_id for request in host_requests if request.call_id is not None}
host_occurrences = {
(request.call_id, request.id)
for request in host_requests
if request.call_id is not None and request.id is not None
}
approval_anchor: int | None = None
host_anchor: int | None = None
idless_host_anchor: int | None = None
indexed_contents: list[tuple[int, Content]] = []
content_positions: dict[int, tuple[int, int]] = {}
current_input_indexes: set[int] = set()
content_index = -1
for message_index, message in enumerate(messages):
message_start = content_index + 1
is_current_input = "_attribution" not in message.additional_properties
for message_content_index, content in enumerate(message.contents):
content_index += 1
indexed_contents.append((content_index, content))
content_positions[content_index] = (message_index, message_content_index)
if not is_current_input:
continue
current_input_indexes.add(content_index)
if content.type == "function_approval_response" and approval_anchor is None:
bound_approval = bind_approval_response(content)
if bound_approval is not None and approval_request_ids.intersection(
str(identity)
for identity in (
bound_approval.additional_properties.get(_APPROVAL_REQUEST_ID_KEY),
bound_approval.id,
)
if identity is not None
):
approval_anchor = message_start
elif (
content.type == "function_result"
and content.call_id is not None
and content.id is not None
and (content.call_id, content.id) in host_occurrences
and host_anchor is None
):
host_anchor = message_start
elif (
content.type == "function_result"
and content.id is None
and content.call_id in host_call_ids
and idless_host_anchor is None
):
idless_host_anchor = message_start
response_start = min(
(anchor for anchor in (approval_anchor, host_anchor, idless_host_anchor) if anchor is not None),
default=None,
)
attributed_contents = [
(message_index, content_index, content)
for message_index, message in enumerate(messages)
if "_attribution" in message.additional_properties
for content_index, content in enumerate(message.contents)
]
attributed_replay_occurrences: set[tuple[int, int]] = set()
reserved_replay_item_indexes: set[int] = set()
for match_idless_host_results in (False, True):
for message_index, content_index, content in attributed_contents:
is_idless_host_result = content.type == "function_result" and content.id is None
if is_idless_host_result != match_idless_host_results:
continue
item_index = attributed_staged_replay_item_index(content, reserved_replay_item_indexes)
if item_index is not None:
reserved_replay_item_indexes.add(item_index)
attributed_replay_occurrences.add((message_index, content_index))
response_entries = [
(index, content)
for index, content in indexed_contents
if response_start is not None
and index >= response_start
and index in current_input_indexes
and content.type in {"function_approval_response", "function_result"}
]
responses = [content for _, content in response_entries]
matched_content_ids, incomplete, ordered_responses, host_result_ids = _match_mixed_pause_responses(
items,
responses,
approval_response_binder=bind_approval_response,
)
matched_current_occurrences = {
content_positions[index] for index, content in response_entries if id(content) in matched_content_ids
}

if matched_content_ids:
if matched_current_occurrences or attributed_replay_occurrences:
filtered_messages: list[Message] = []
for message in messages:
message.contents = [content for content in message.contents if id(content) not in matched_content_ids]
for message_index, message in enumerate(messages):
message.contents = [
content
for content_index, content in enumerate(message.contents)
if (message_index, content_index) not in matched_current_occurrences
and (message_index, content_index) not in attributed_replay_occurrences
]
if message.contents:
filtered_messages.append(message)
messages[:] = filtered_messages
Expand Down
Loading
Loading