Skip to content

JuPydanticAgent

Mixes a pydantic-ai agent and a jupyter kernel (through MetaKernel). Gives access to methods that allows the agent to be streamed to the kernel. Support different output types : markdown, html, widgets, raw text, colored text, images, ...

This agent wrapper can be used with any MetaKernel instance; allowing to build complex multi agent applications.

Source code in pydantic_ai_kernel/jupydantic_agent.py
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
class JuPydanticAgent:
    """
    Mixes a pydantic-ai agent and a jupyter kernel (through MetaKernel).
    Gives access to methods that allows the agent to be streamed to the kernel.
    Support different output types : markdown, html, widgets, raw text,
    colored text, images, ...

    This agent wrapper can be used with any MetaKernel instance; allowing
    to build complex multi agent applications.
    """

    def __init__(
        self, agent: Agent[None, Any], kernel: MetaKernel, agent_config: AgentConfig
    ):
        self.agent = agent
        self.agent_history = []
        self.agent_config = agent_config
        self.kernel = kernel
        self.log = self.kernel.log
        self.is_interrupted = False
        self.out_md_stack = []

    async def stream_agent_to_kernel(
        self,
        prompt: Optional[str],
        deferred_tool_results: Optional[DeferredToolResults] = None,
    ) -> DeferredToolRequests | Any:
        """
        Runs the agent once, and streams the output to stdout.
        Returns a DeferredToolRequest to ask for user approval if some tool needs it.
        Else, returns None.

        Can be called several times in the same question if the user gives
        its validation for tool calling.

        Parameters :
        ---
            - prompt (Optional[str]): the prompt given to the agent
            - deferred_tool_result (Optional[DeferredToolResults]) : optional
                tool approval / disapproval of the last agent run.

        Returns :
        ---
            None if all text output was streamed to stdout. A DeferredToolRequests
            object if the user needs to validate some tool calling. Can be the
            Agent output type.
        """
        agent_output = None

        async with self.agent.iter(
            user_prompt=prompt,
            message_history=self.agent_history,
            deferred_tool_results=deferred_tool_results,
        ) as agent_run:
            if self.is_interrupted:
                raise UserInterruptionError("User stopped the execution")
            async for node in agent_run:
                # The inner stream already aborts via UserInterruptionError,
                # but we keep the old guard as a fallback.
                if self.is_interrupted:
                    raise UserInterruptionError("User stopped the execution")

                # self.Print(f"node : {node}")
                if Agent.is_user_prompt_node(node):
                    continue
                elif Agent.is_model_request_node(node):
                    await self.deal_with_model_request_node(node, agent_run.ctx)
                elif Agent.is_call_tools_node(node):
                    self.log.debug(f"Calling tool node : {node}")
                elif Agent.is_end_node(node):
                    # either agent output or DeferredToolRequest
                    agent_output = node.data.output

            self.agent_history = agent_run.all_messages()

        self.last_usage = agent_run.usage

        return agent_output

    async def deal_with_model_request_node(
        self, model_request_node: ModelRequestNode, ctx
    ):
        if self.agent_config is None:
            self.log.warning("No config was found for agent during its request")
            return
        self.log.info(model_request_node.request)
        self.log.info(model_request_node._result)

        # all_parts = model_request_node.request.parts
        # for part in
        # if isinstance(.parts[0], ToolReturnPart):
        #     if isinstance(model_request_node.request.parts[0].content, Widget):
        #         self.Display(model_request_node.request.parts[0].content)
        #         model_request_node.request.parts[0].content = "Widget"
        #         return
        out_md = ""
        think_content = ""
        async with model_request_node.stream(ctx) as request_stream:
            async for event in request_stream:
                if self.is_interrupted:
                    raise UserInterruptionError("Interrupt received while streaming")
                if isinstance(event, PartStartEvent):
                    if (
                        isinstance(event.part, ThinkingPart)
                        and self.agent_config.display_thinking
                    ):
                        self.kernel.Print("========= Thinking =========")
                        self.kernel.Print(event.part.content, end="")

                        if self.use_widget:
                            think_content += event.part.content
                        else:
                            out_md += "> **Thinking** :  \n"
                            out_md += f"> {event.part.content}"

                    if isinstance(event.part, TextPart):
                        self.kernel.Print(event.part.content, end="")
                        out_md += event.part.content
                if isinstance(event, PartDeltaEvent):
                    if (
                        isinstance(event.delta, ThinkingPartDelta)
                        and event.delta.content_delta is not None
                        and self.agent_config.display_thinking
                    ):
                        self.kernel.Print(
                            event.delta.content_delta,
                            end="",
                        )
                        if self.use_widget:
                            think_content += event.delta.content_delta
                        else:
                            out_md += event.delta.content_delta.replace("\n", "  \n> ")

                    if isinstance(event.delta, TextPartDelta):
                        self.kernel.Print(event.delta.content_delta, end="")
                        out_md += event.delta.content_delta
                if isinstance(event, PartEndEvent):
                    if (
                        isinstance(event.part, ThinkingPart)
                        and self.agent_config.display_thinking
                    ):
                        self.kernel.Print("\n========= End Think =========")
                        out_md += "  \n  \n"

        if len(think_content) > 0 and self.use_widget:
            self.out_md_stack.append(
                widgets.Accordion(
                    children=[
                        widgets.HTML(
                            value=think_content,
                            placeholder="",
                            description="",
                        )
                    ],
                    titles=["Though"],
                )
            )
        if len(out_md) > 0:
            # avoid display empty blocks
            self.out_md_stack.append({"text/markdown": out_md})
        if self.formatter == "md":
            # this line could break things, because we bet
            # here that frontends that can't display markdown
            # are also those that can't clear output
            # (for ex. jupyter console).
            # since there is no way to know if a frontend
            # really clears output on clear_output msg, this
            # remains the best option
            # we could also have an environment variable to set
            # if the current frontend implements clear_output
            self.kernel.Display(
                *self.out_md_stack,
                clear_output=True,
            )

    @property
    def formatter(self):
        return self.agent_config.formatter

    @property
    def use_widget(self):
        return self.agent_config.use_widget

    async def run(self, prompt: str, silent: bool = False):
        """
        Makes several runs of the agent, until all tool calls
        are processed. Streams its output to the kernel. All
        the tool calls are dealt with input_message.
        """
        try:
            self.is_interrupted = False
            deferred_tool_result = None
            self.out_md_stack = []
            while True:
                # runs the agent until it returns anything other than a deferred
                # tool request.
                # if it outputs a deferred tool request, ask for user approval,
                # and then re-run the agent.
                prompt_or_tool_result = prompt if deferred_tool_result is None else None
                agent_out = await self.stream_agent_to_kernel(
                    prompt_or_tool_result, deferred_tool_results=deferred_tool_result
                )

                # self.log.info(f"Agent out : {agent_out}")
                # self.log.info(f"Agent history : {self.agent.history_processors}")
                # If an interrupt was received while the agent was running, stop the loop immediately.
                if self.is_interrupted:
                    self.log.info("Execution halted by user interrupt")
                    break

                if isinstance(agent_out, DeferredToolRequests):
                    results = DeferredToolResults()
                    for call in agent_out.approvals:
                        self.log.info(f"Call : {call}")

                        user_validation = self.kernel.raw_input(
                            f"{prompt_user_approval(call)} ([y]/n):"
                        )
                        if user_validation not in ["Y", "yes", 1, "1", "y"]:
                            result = ToolDenied(
                                f"User rejected the tool call : {user_validation}"
                            )
                        else:
                            result = True
                        results.approvals[call.tool_call_id] = result
                    deferred_tool_result = results
                elif isinstance(agent_out, str):
                    break
                elif isinstance(agent_out, BinaryImage):
                    self.kernel.Display({"image/png": agent_out.data})
                    break
                elif isinstance(agent_out, SerializableWidget):
                    self.kernel.Display(agent_out.widget)
                    break
                else:  # here, check for unprocessed tool calls, and remove them
                    last_msg = self.agent_history[-1]
                    # self.log.info(f"Last message : {last_msg}")
                    # self.log.info(f"Last message first part : {last_msg.parts[0]}")
                    # self.log.info(
                    #     f"Last message first part content : {last_msg.parts[0].content}"
                    # )
                    if (
                        isinstance(last_msg, ModelRequest)
                        and isinstance(last_msg.parts[0], ToolReturnPart)
                        and last_msg.parts[0].content
                        == "Tool not executed - a final result was already processed."
                    ):
                        self.agent_history.pop(-1)
                        self.agent_history.pop(-1)
                        self.log.info(
                            f"Popped last elem. Agent history now : {self.agent_history}"
                        )
                        # TODO : force for deferred tool in this case also (else some tool
                        # calls are skipped)
                    break

        except UserInterruptionError:
            self.log.info("User interrupted the loop")
        except Exception:
            self.kernel.Error_display(traceback.format_exc())

run(prompt, silent=False) async

Makes several runs of the agent, until all tool calls are processed. Streams its output to the kernel. All the tool calls are dealt with input_message.

Source code in pydantic_ai_kernel/jupydantic_agent.py
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
async def run(self, prompt: str, silent: bool = False):
    """
    Makes several runs of the agent, until all tool calls
    are processed. Streams its output to the kernel. All
    the tool calls are dealt with input_message.
    """
    try:
        self.is_interrupted = False
        deferred_tool_result = None
        self.out_md_stack = []
        while True:
            # runs the agent until it returns anything other than a deferred
            # tool request.
            # if it outputs a deferred tool request, ask for user approval,
            # and then re-run the agent.
            prompt_or_tool_result = prompt if deferred_tool_result is None else None
            agent_out = await self.stream_agent_to_kernel(
                prompt_or_tool_result, deferred_tool_results=deferred_tool_result
            )

            # self.log.info(f"Agent out : {agent_out}")
            # self.log.info(f"Agent history : {self.agent.history_processors}")
            # If an interrupt was received while the agent was running, stop the loop immediately.
            if self.is_interrupted:
                self.log.info("Execution halted by user interrupt")
                break

            if isinstance(agent_out, DeferredToolRequests):
                results = DeferredToolResults()
                for call in agent_out.approvals:
                    self.log.info(f"Call : {call}")

                    user_validation = self.kernel.raw_input(
                        f"{prompt_user_approval(call)} ([y]/n):"
                    )
                    if user_validation not in ["Y", "yes", 1, "1", "y"]:
                        result = ToolDenied(
                            f"User rejected the tool call : {user_validation}"
                        )
                    else:
                        result = True
                    results.approvals[call.tool_call_id] = result
                deferred_tool_result = results
            elif isinstance(agent_out, str):
                break
            elif isinstance(agent_out, BinaryImage):
                self.kernel.Display({"image/png": agent_out.data})
                break
            elif isinstance(agent_out, SerializableWidget):
                self.kernel.Display(agent_out.widget)
                break
            else:  # here, check for unprocessed tool calls, and remove them
                last_msg = self.agent_history[-1]
                # self.log.info(f"Last message : {last_msg}")
                # self.log.info(f"Last message first part : {last_msg.parts[0]}")
                # self.log.info(
                #     f"Last message first part content : {last_msg.parts[0].content}"
                # )
                if (
                    isinstance(last_msg, ModelRequest)
                    and isinstance(last_msg.parts[0], ToolReturnPart)
                    and last_msg.parts[0].content
                    == "Tool not executed - a final result was already processed."
                ):
                    self.agent_history.pop(-1)
                    self.agent_history.pop(-1)
                    self.log.info(
                        f"Popped last elem. Agent history now : {self.agent_history}"
                    )
                    # TODO : force for deferred tool in this case also (else some tool
                    # calls are skipped)
                break

    except UserInterruptionError:
        self.log.info("User interrupted the loop")
    except Exception:
        self.kernel.Error_display(traceback.format_exc())

stream_agent_to_kernel(prompt, deferred_tool_results=None) async

Runs the agent once, and streams the output to stdout. Returns a DeferredToolRequest to ask for user approval if some tool needs it. Else, returns None.

Can be called several times in the same question if the user gives its validation for tool calling.

Parameters :

- prompt (Optional[str]): the prompt given to the agent
- deferred_tool_result (Optional[DeferredToolResults]) : optional
    tool approval / disapproval of the last agent run.

Returns :

None if all text output was streamed to stdout. A DeferredToolRequests
object if the user needs to validate some tool calling. Can be the
Agent output type.
Source code in pydantic_ai_kernel/jupydantic_agent.py
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
async def stream_agent_to_kernel(
    self,
    prompt: Optional[str],
    deferred_tool_results: Optional[DeferredToolResults] = None,
) -> DeferredToolRequests | Any:
    """
    Runs the agent once, and streams the output to stdout.
    Returns a DeferredToolRequest to ask for user approval if some tool needs it.
    Else, returns None.

    Can be called several times in the same question if the user gives
    its validation for tool calling.

    Parameters :
    ---
        - prompt (Optional[str]): the prompt given to the agent
        - deferred_tool_result (Optional[DeferredToolResults]) : optional
            tool approval / disapproval of the last agent run.

    Returns :
    ---
        None if all text output was streamed to stdout. A DeferredToolRequests
        object if the user needs to validate some tool calling. Can be the
        Agent output type.
    """
    agent_output = None

    async with self.agent.iter(
        user_prompt=prompt,
        message_history=self.agent_history,
        deferred_tool_results=deferred_tool_results,
    ) as agent_run:
        if self.is_interrupted:
            raise UserInterruptionError("User stopped the execution")
        async for node in agent_run:
            # The inner stream already aborts via UserInterruptionError,
            # but we keep the old guard as a fallback.
            if self.is_interrupted:
                raise UserInterruptionError("User stopped the execution")

            # self.Print(f"node : {node}")
            if Agent.is_user_prompt_node(node):
                continue
            elif Agent.is_model_request_node(node):
                await self.deal_with_model_request_node(node, agent_run.ctx)
            elif Agent.is_call_tools_node(node):
                self.log.debug(f"Calling tool node : {node}")
            elif Agent.is_end_node(node):
                # either agent output or DeferredToolRequest
                agent_output = node.data.output

        self.agent_history = agent_run.all_messages()

    self.last_usage = agent_run.usage

    return agent_output