| 218 | """The agent tool can use plugins from parent runner.""" |
| 219 | |
| 220 | class ModelResponseCapturePlugin(BasePlugin): |
| 221 | |
| 222 | def __init__(self): |
| 223 | super().__init__('plugin') |
| 224 | self.model_responses = {} |
| 225 | |
| 226 | async def after_model_callback( |
| 227 | self, |
| 228 | *, |
| 229 | callback_context: CallbackContext, |
| 230 | llm_response: LlmResponse, |
| 231 | ) -> Optional[LlmResponse]: |
| 232 | response_text = [] |
| 233 | for part in llm_response.content.parts: |
| 234 | if not part.text: |
| 235 | continue |
| 236 | response_text.append(part.text) |
| 237 | if response_text: |
| 238 | if callback_context.agent_name not in self.model_responses: |
| 239 | self.model_responses[callback_context.agent_name] = [] |
| 240 | self.model_responses[callback_context.agent_name].append( |
| 241 | ''.join(response_text) |
| 242 | ) |
| 243 | |
| 244 | mock_model = testing_utils.MockModel.create( |
| 245 | responses=[ |
no outgoing calls