|
25 | 25 | PartStartEvent, |
26 | 26 | TextPart, |
27 | 27 | TextPartDelta, |
| 28 | + ThinkingPart, |
| 29 | + ThinkingPartDelta, |
28 | 30 | ToolCallPart, |
29 | 31 | ToolCallPartDelta, |
30 | 32 | ) |
@@ -86,8 +88,7 @@ def handle_text_delta( |
86 | 88 | A `PartStartEvent` if a new part was created, or a `PartDeltaEvent` if an existing part was updated. |
87 | 89 |
|
88 | 90 | Raises: |
89 | | - UnexpectedModelBehavior: If attempting to apply text content to a part that is |
90 | | - not a TextPart. |
| 91 | + UnexpectedModelBehavior: If attempting to apply text content to a part that is not a TextPart. |
91 | 92 | """ |
92 | 93 | existing_text_part_and_index: tuple[TextPart, int] | None = None |
93 | 94 |
|
@@ -122,6 +123,77 @@ def handle_text_delta( |
122 | 123 | self._parts[part_index] = part_delta.apply(existing_text_part) |
123 | 124 | return PartDeltaEvent(index=part_index, delta=part_delta) |
124 | 125 |
|
| 126 | + def handle_thinking_delta( |
| 127 | + self, |
| 128 | + *, |
| 129 | + vendor_part_id: Hashable | None, |
| 130 | + content: str | None = None, |
| 131 | + signature: str | None = None, |
| 132 | + ) -> ModelResponseStreamEvent: |
| 133 | + """Handle incoming thinking content, creating or updating a ThinkingPart in the manager as appropriate. |
| 134 | +
|
| 135 | + When `vendor_part_id` is None, the latest part is updated if it exists and is a ThinkingPart; |
| 136 | + otherwise, a new ThinkingPart is created. When a non-None ID is specified, the ThinkingPart corresponding |
| 137 | + to that vendor ID is either created or updated. |
| 138 | +
|
| 139 | + Args: |
| 140 | + vendor_part_id: The ID the vendor uses to identify this piece |
| 141 | + of thinking. If None, a new part will be created unless the latest part is already |
| 142 | + a ThinkingPart. |
| 143 | + content: The thinking content to append to the appropriate ThinkingPart. |
| 144 | + signature: An optional signature for the thinking content. |
| 145 | +
|
| 146 | + Returns: |
| 147 | + A `PartStartEvent` if a new part was created, or a `PartDeltaEvent` if an existing part was updated. |
| 148 | +
|
| 149 | + Raises: |
| 150 | + UnexpectedModelBehavior: If attempting to apply a thinking delta to a part that is not a ThinkingPart. |
| 151 | + """ |
| 152 | + existing_thinking_part_and_index: tuple[ThinkingPart, int] | None = None |
| 153 | + |
| 154 | + if vendor_part_id is None: |
| 155 | + # If the vendor_part_id is None, check if the latest part is a ThinkingPart to update |
| 156 | + if self._parts: |
| 157 | + part_index = len(self._parts) - 1 |
| 158 | + latest_part = self._parts[part_index] |
| 159 | + if isinstance(latest_part, ThinkingPart): # pragma: no branch |
| 160 | + existing_thinking_part_and_index = latest_part, part_index |
| 161 | + else: |
| 162 | + # Otherwise, attempt to look up an existing ThinkingPart by vendor_part_id |
| 163 | + part_index = self._vendor_id_to_part_index.get(vendor_part_id) |
| 164 | + if part_index is not None: |
| 165 | + existing_part = self._parts[part_index] |
| 166 | + if not isinstance(existing_part, ThinkingPart): |
| 167 | + raise UnexpectedModelBehavior(f'Cannot apply a thinking delta to {existing_part=}') |
| 168 | + existing_thinking_part_and_index = existing_part, part_index |
| 169 | + |
| 170 | + if existing_thinking_part_and_index is None: |
| 171 | + if content is not None: |
| 172 | + # There is no existing thinking part that should be updated, so create a new one |
| 173 | + new_part_index = len(self._parts) |
| 174 | + part = ThinkingPart(content=content, signature=signature) |
| 175 | + if vendor_part_id is not None: # pragma: no branch |
| 176 | + self._vendor_id_to_part_index[vendor_part_id] = new_part_index |
| 177 | + self._parts.append(part) |
| 178 | + return PartStartEvent(index=new_part_index, part=part) |
| 179 | + else: |
| 180 | + raise UnexpectedModelBehavior('Cannot create a ThinkingPart with no content') |
| 181 | + else: |
| 182 | + if content is not None: |
| 183 | + # Update the existing ThinkingPart with the new content delta |
| 184 | + existing_thinking_part, part_index = existing_thinking_part_and_index |
| 185 | + part_delta = ThinkingPartDelta(content_delta=content) |
| 186 | + self._parts[part_index] = part_delta.apply(existing_thinking_part) |
| 187 | + return PartDeltaEvent(index=part_index, delta=part_delta) |
| 188 | + elif signature is not None: |
| 189 | + # Update the existing ThinkingPart with the new signature delta |
| 190 | + existing_thinking_part, part_index = existing_thinking_part_and_index |
| 191 | + part_delta = ThinkingPartDelta(signature_delta=signature) |
| 192 | + self._parts[part_index] = part_delta.apply(existing_thinking_part) |
| 193 | + return PartDeltaEvent(index=part_index, delta=part_delta) |
| 194 | + else: |
| 195 | + raise UnexpectedModelBehavior('Cannot update a ThinkingPart with no content or signature') |
| 196 | + |
125 | 197 | def handle_tool_call_delta( |
126 | 198 | self, |
127 | 199 | *, |
|
0 commit comments