Skip to content
Open
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
15 changes: 13 additions & 2 deletions stagehand/agent/anthropic_cua.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@


class AnthropicCUAClient(AgentClient):
ALLOWED_TOOL_USE_FIELDS = {"type", "id", "name", "input"}

ANTHROPIC_KEY_MAPPING = {
"return": "Enter",
"enter": "Enter",
Expand Down Expand Up @@ -283,6 +285,15 @@ def _format_initial_messages(
messages.append({"role": "user", "content": user_content})
return messages

def _sanitise_content_block(self, block_data: dict[str, Any]) -> dict[str, Any]:
if block_data.get("type") == "tool_use":
return {
k: v
for k, v in block_data.items()
if k in self.ALLOWED_TOOL_USE_FIELDS
}
return block_data

def _process_provider_response(
self, response: Any # Anthropic API response object
) -> tuple[Optional[AgentAction], Optional[str], bool, list[dict[str, Any]]]:
Expand All @@ -293,10 +304,10 @@ def _process_provider_response(

raw_assistant_content_blocks = []
if hasattr(response, "content") and isinstance(response.content, list):
# Serialize Pydantic models from response.content for history
try:
raw_assistant_content_blocks = [
block.model_dump() for block in response.content
self._sanitise_content_block(block.model_dump())
for block in response.content
]
except Exception as e:
self.logger.error(
Expand Down