diff --git a/src/celeste/modalities/text/providers/anthropic/client.py b/src/celeste/modalities/text/providers/anthropic/client.py index aae3d1b..b1572ee 100644 --- a/src/celeste/modalities/text/providers/anthropic/client.py +++ b/src/celeste/modalities/text/providers/anthropic/client.py @@ -40,11 +40,9 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self._message_start: dict[str, Any] | None = None self._tool_calls: dict[int, dict[str, Any]] = {} - self._thinking_blocks: list[dict[str, Any]] = [] - self._current_thinking: dict[str, Any] | None = None def _parse_chunk(self, event_data: dict[str, Any]) -> TextChunk | None: - """Parse one SSE event into a typed chunk (captures message_start, tool_use, thinking).""" + """Parse one SSE event into a typed chunk (captures message_start and tool_use).""" event_type = event_data.get("type") if event_type == "message_start": message = event_data.get("message") @@ -53,37 +51,19 @@ def _parse_chunk(self, event_data: dict[str, Any]) -> TextChunk | None: return None if event_type == "content_block_start": block = event_data.get("content_block", {}) - block_type = block.get("type") - if block_type == "tool_use": + if block.get("type") == "tool_use": idx = event_data.get("index", len(self._tool_calls)) self._tool_calls[idx] = { "id": block.get("id", ""), "name": block.get("name", ""), "input_json": "", } - elif block_type == "thinking": - self._current_thinking = { - "type": "thinking", - "thinking": "", - "signature": "", - } - elif block_type == "redacted_thinking": - self._thinking_blocks.append(block) elif event_type == "content_block_delta": delta = event_data.get("delta", {}) - delta_type = delta.get("type") - if delta_type == "input_json_delta": + if delta.get("type") == "input_json_delta": idx = event_data.get("index", -1) if idx in self._tool_calls: self._tool_calls[idx]["input_json"] += delta.get("partial_json", "") - elif delta_type == "thinking_delta" and self._current_thinking is not None: - self._current_thinking["thinking"] += delta.get("thinking", "") - elif delta_type == "signature_delta" and self._current_thinking is not None: - self._current_thinking["signature"] += delta.get("signature", "") - elif event_type == "content_block_stop": - if self._current_thinking is not None: - self._thinking_blocks.append(self._current_thinking) - self._current_thinking = None return super()._parse_chunk(event_data) def _aggregate_event_data(self, chunks: list[TextChunk]) -> list[dict[str, Any]]: @@ -110,8 +90,12 @@ def _aggregate_tool_calls( def _aggregate_signature( self, chunks: list[TextChunk], raw_events: list[dict[str, Any]] ) -> list[dict[str, Any]]: - """Return accumulated thinking blocks for round-trip signature.""" - return self._thinking_blocks + """Return the full native content array when the turn carries thinking blocks.""" + blocks = self._aggregate_content_blocks() + has_thinking = any( + b.get("type") in {"thinking", "redacted_thinking"} for b in blocks + ) + return blocks if has_thinking else [] def _aggregate_grounding( self, chunks: list[TextChunk], raw_events: list[dict[str, Any]] @@ -191,8 +175,15 @@ def content_to_blocks(content: Any) -> list[dict[str, Any]]: pending_tool_results = [] if role == Role.ASSISTANT and (message.tool_calls or message.signature): + sig_blocks = message.signature or [] + if any( + b.get("type") not in {"thinking", "redacted_thinking"} + for b in sig_blocks + ): + # thinking signs block positions — echo the full original turn verbatim + messages.append({"role": "assistant", "content": sig_blocks}) + continue content_blocks: list[dict[str, Any]] = [] - sig_blocks = message.signature if sig_blocks: content_blocks.extend(sig_blocks) if content: @@ -280,21 +271,18 @@ def _parse_content( def _parse_reasoning( self, response_data: dict[str, Any] ) -> tuple[str | None, list[dict[str, Any]]]: - """Parse thinking blocks from Anthropic response.""" + """Parse thinking from Anthropic response (signature = full content array).""" blocks = response_data.get("content", []) - reasoning_parts: list[str] = [] - signature_blocks: list[dict[str, Any]] = [] - for block in blocks: - block_type = block.get("type") - if block_type == "thinking": - thinking_text = block.get("thinking", "") - if thinking_text: - reasoning_parts.append(thinking_text) - signature_blocks.append(block) - elif block_type == "redacted_thinking": - signature_blocks.append(block) + reasoning_parts = [ + b["thinking"] + for b in blocks + if b.get("type") == "thinking" and b.get("thinking") + ] + has_thinking = any( + b.get("type") in {"thinking", "redacted_thinking"} for b in blocks + ) text = "\n".join(reasoning_parts) if reasoning_parts else None - return text, signature_blocks + return text, list(blocks) if has_thinking else [] def _parse_tool_calls(self, response_data: dict[str, Any]) -> list[ToolCall]: """Parse tool calls from Anthropic response.""" diff --git a/src/celeste/providers/anthropic/messages/streaming.py b/src/celeste/providers/anthropic/messages/streaming.py index b29d5e9..e1078a2 100644 --- a/src/celeste/providers/anthropic/messages/streaming.py +++ b/src/celeste/providers/anthropic/messages/streaming.py @@ -32,15 +32,21 @@ def _parse_chunk(self, event_data: dict[str, Any]) -> Any: # noqa: ANN401 block = event_data.get("content_block", {}) block_type = block.get("type") idx = event_data.get("index", len(self._content_blocks)) - if block_type == "server_tool_use": + if block_type in {"server_tool_use", "tool_use"}: self._content_blocks[idx] = { - "type": "server_tool_use", + "type": block_type, "id": block.get("id", ""), "name": block.get("name", ""), "input_json": "", } - elif block_type == "web_search_tool_result": + elif block_type in {"web_search_tool_result", "redacted_thinking"}: self._content_blocks[idx] = block + elif block_type == "thinking": + self._content_blocks[idx] = { + "type": "thinking", + "thinking": block.get("thinking", ""), + "signature": block.get("signature", ""), + } elif block_type == "text": self._content_blocks[idx] = { "type": "text", @@ -53,8 +59,12 @@ def _parse_chunk(self, event_data: dict[str, Any]) -> Any: # noqa: ANN401 idx = event_data.get("index", -1) block = self._content_blocks.get(idx) if delta_type == "input_json_delta": - if block and block.get("type") == "server_tool_use": + if block and block.get("type") in {"server_tool_use", "tool_use"}: block["input_json"] += delta.get("partial_json", "") + elif delta_type in {"thinking_delta", "signature_delta"}: + key = delta_type.removesuffix("_delta") + if block and block.get("type") == "thinking": + block[key] += delta.get(key, "") elif delta_type in {"text_delta", "citations_delta"}: block = self._content_blocks.setdefault( idx, {"type": "text", "text": "", "citations": []} @@ -72,19 +82,21 @@ def _aggregate_content_blocks(self) -> list[dict[str, Any]]: blocks: list[dict[str, Any]] = [] for idx in sorted(self._content_blocks): block = self._content_blocks[idx] - if block.get("type") == "server_tool_use": + if block.get("type") in {"server_tool_use", "tool_use"}: input_data: dict[str, Any] = {} if block["input_json"]: with contextlib.suppress(ValueError, TypeError): input_data = json.loads(block["input_json"]) blocks.append( { - "type": "server_tool_use", + "type": block["type"], "id": block["id"], "name": block["name"], "input": input_data, } ) + elif block.get("type") == "text" and not block.get("citations"): + blocks.append({"type": "text", "text": block.get("text", "")}) else: blocks.append(block) return blocks