zrb 1.5.3__tar.gz → 1.5.4__tar.gz

{zrb-1.5.3 → zrb-1.5.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.5.3
+Version: 1.5.4
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later
@@ -26,7 +26,7 @@ Requires-Dist: openai (>=1.70.0,<2.0.0) ; extra == "rag" or extra == "all"
 Requires-Dist: pdfplumber (>=0.11.6,<0.12.0) ; extra == "rag" or extra == "all"
 Requires-Dist: playwright (>=1.51.0,<2.0.0) ; extra == "playwright" or extra == "all"
 Requires-Dist: psutil (>=7.0.0,<8.0.0)
-Requires-Dist: pydantic-ai (>=0.0.49,<0.0.50)
+Requires-Dist: pydantic-ai (>=0.0.52,<0.0.53)
 Requires-Dist: python-dotenv (>=1.1.0,<2.0.0)
 Requires-Dist: python-jose[cryptography] (>=3.4.0,<4.0.0)
 Requires-Dist: requests (>=2.32.3,<3.0.0)

{zrb-1.5.3 → zrb-1.5.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zrb"
-version = "1.5.3"
+version = "1.5.4"
 description = "Your Automation Powerhouse"
 authors = ["Go Frendi Gunawan <gofrendiasgard@gmail.com>"]
 license = "AGPL-3.0-or-later"
@@ -49,7 +49,7 @@ beautifulsoup4 = "^4.13.3"
 fastapi = {extras = ["standard"], version = "^0.115.12"}
 python-jose = {extras = ["cryptography"], version = "^3.4.0"}
 ulid-py = "^1.1.0"
-pydantic-ai = "^0.0.49"
+pydantic-ai = "^0.0.52"
 openai = {version = "^1.70.0", optional = true}
 playwright = {version = "^1.51.0", optional = true}
 psutil = "^7.0.0"

{zrb-1.5.3 → zrb-1.5.4}/src/zrb/task/base_task.py

@@ -1,6 +1,6 @@
 import asyncio
 import os
-from collections.abc import Callable
+from collections.abc import Callable, Coroutine
 from typing import Any
 
 from zrb.attr.type import BoolAttr, fstring
@@ -64,7 +64,7 @@ class BaseTask(AnyTask):
     def __repr__(self):
         return f"<{self.__class__.__name__} name={self._name}>"
 
-    def __rshift__(self, other: AnyTask | list[AnyTask]) -> AnyTask:
+    def __rshift__(self, other: AnyTask | list[AnyTask]) -> AnyTask | list[AnyTask]:
         try:
             if isinstance(other, AnyTask):
                 other.append_upstream(self)
@@ -123,7 +123,9 @@ class BaseTask(AnyTask):
         return [task_input for task_input in inputs if task_input is not None]
 
     def __combine_inputs(
-        self, inputs: list[AnyInput], other_inputs: list[AnyInput] | AnyInput
+        self,
+        inputs: list[AnyInput],
+        other_inputs: list[AnyInput | None] | AnyInput | None,
     ):
         input_names = [task_input.name for task_input in inputs]
         if isinstance(other_inputs, AnyInput):
@@ -429,7 +431,7 @@ class BaseTask(AnyTask):
             session.get_task_status(self).reset()
             # defer this action
             ctx.log_info("Running")
-            action_coro = asyncio.create_task(
+            action_coro: Coroutine = asyncio.create_task(
                 run_async(self.__exec_action_and_retry(session))
             )
             session.defer_action(self, action_coro)

{zrb-1.5.3 → zrb-1.5.4}/src/zrb/task/llm_task.py

@@ -1,4 +1,3 @@
-import copy
 import functools
 import inspect
 import json
@@ -9,6 +8,7 @@ from typing import Any
 
 from openai import APIError
 from pydantic_ai import Agent, Tool
+from pydantic_ai.mcp import MCPServer
 from pydantic_ai.messages import (
     FinalResultEvent,
     FunctionToolCallEvent,
@@ -54,9 +54,9 @@ class LLMTask(BaseTask):
             Callable[[AnySharedContext], Model | str | fstring] | Model | None
         ) = None,
         render_model: bool = True,
-        model_base_url: StrAttr = None,
+        model_base_url: StrAttr | None = None,
         render_model_base_url: bool = True,
-        model_api_key: StrAttr = None,
+        model_api_key: StrAttr | None = None,
         render_model_api_key: bool = True,
         model_settings: (
             ModelSettings | Callable[[AnySharedContext], ModelSettings] | None
@@ -68,6 +68,9 @@ class LLMTask(BaseTask):
         tools: (
             list[ToolOrCallable] | Callable[[AnySharedContext], list[ToolOrCallable]]
         ) = [],
+        mcp_servers: (
+            list[MCPServer] | Callable[[AnySharedContext], list[MCPServer]]
+        ) = [],
         conversation_history: (
             ListOfDict | Callable[[AnySharedContext], ListOfDict]
         ) = [],
@@ -127,6 +130,8 @@ class LLMTask(BaseTask):
         self._message = message
         self._tools = tools
         self._additional_tools: list[ToolOrCallable] = []
+        self._mcp_servers = mcp_servers
+        self._additional_mcp_servers: list[MCPServer] = []
         self._conversation_history = conversation_history
         self._conversation_history_reader = conversation_history_reader
         self._conversation_history_writer = conversation_history_writer
@@ -137,32 +142,36 @@ class LLMTask(BaseTask):
     def add_tool(self, tool: ToolOrCallable):
         self._additional_tools.append(tool)
 
+    def add_mcp_server(self, mcp_server: MCPServer):
+        self._additional_mcp_servers.append(mcp_server)
+
     async def _exec_action(self, ctx: AnyContext) -> Any:
         history = await self._read_conversation_history(ctx)
         user_prompt = self._get_message(ctx)
         agent = self._get_agent(ctx)
         try:
-            async with agent.iter(
-                user_prompt=user_prompt,
-                message_history=ModelMessagesTypeAdapter.validate_python(history),
-            ) as agent_run:
-                async for node in agent_run:
-                    # Each node represents a step in the agent's execution
-                    # Reference: https://ai.pydantic.dev/agents/#streaming
-                    try:
-                        await self._print_node(ctx, agent_run, node)
-                    except APIError as e:
-                        # Extract detailed error information from the response
-                        error_details = _extract_api_error_details(e)
-                        ctx.log_error(f"API Error: {error_details}")
-                        raise
-                    except Exception as e:
-                        ctx.log_error(f"Error processing node: {str(e)}")
-                        ctx.log_error(f"Error type: {type(e).__name__}")
-                        raise
-                new_history = json.loads(agent_run.result.all_messages_json())
-                await self._write_conversation_history(ctx, new_history)
-                return agent_run.result.data
+            async with agent.run_mcp_servers():
+                async with agent.iter(
+                    user_prompt=user_prompt,
+                    message_history=ModelMessagesTypeAdapter.validate_python(history),
+                ) as agent_run:
+                    async for node in agent_run:
+                        # Each node represents a step in the agent's execution
+                        # Reference: https://ai.pydantic.dev/agents/#streaming
+                        try:
+                            await self._print_node(ctx, agent_run, node)
+                        except APIError as e:
+                            # Extract detailed error information from the response
+                            error_details = _extract_api_error_details(e)
+                            ctx.log_error(f"API Error: {error_details}")
+                            raise
+                        except Exception as e:
+                            ctx.log_error(f"Error processing node: {str(e)}")
+                            ctx.log_error(f"Error type: {type(e).__name__}")
+                            raise
+                    new_history = json.loads(agent_run.result.all_messages_json())
+                    await self._write_conversation_history(ctx, new_history)
+                    return agent_run.result.data
         except Exception as e:
             ctx.log_error(f"Error in agent execution: {str(e)}")
             raise
@@ -277,10 +286,15 @@ class LLMTask(BaseTask):
             tool if isinstance(tool, Tool) else Tool(_wrap_tool(tool), takes_ctx=False)
             for tool in tools_or_callables
         ]
+        mcp_servers = list(
+            self._mcp_servers(ctx) if callable(self._mcp_servers) else self._mcp_servers
+        )
+        mcp_servers.extend(self._additional_mcp_servers)
         return Agent(
             self._get_model(ctx),
             system_prompt=self._get_system_prompt(ctx),
             tools=tools,
+            mcp_servers=mcp_servers,
             model_settings=self._get_model_settings(ctx),
             retries=3,
         )