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
66 changes: 27 additions & 39 deletions src/celeste/modalities/text/providers/anthropic/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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]]:
Expand All @@ -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]]
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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."""
Expand Down
24 changes: 18 additions & 6 deletions src/celeste/providers/anthropic/messages/streaming.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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": []}
Expand All @@ -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
Expand Down
Loading