zrb 1.5.3__py3-none-any.whl → 1.5.4__py3-none-any.whl

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/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,
         )

{zrb-1.5.3.dist-info → zrb-1.5.4.dist-info}/METADATA

@@ -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.dist-info → zrb-1.5.4.dist-info}/RECORD

@@ -296,11 +296,11 @@ zrb/session_state_logger/file_session_state_logger.py,sha256=1ue7-Bcwg4wlLn2G_7A
 zrb/session_state_logger/session_state_logger_factory.py,sha256=wXf2DVmeRmx399MFYYty6uNcPZMcf7iayHBYCLGlhfc,189
 zrb/task/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/task/any_task.py,sha256=9rCdKe-Sayr34Han9AsbhRxFpkbk6Rteg1DOyETulwQ,4917
-zrb/task/base_task.py,sha256=SQRf37bylS586KwyW0eYDe9JZ5Hl18FP8kScHae6y3A,21251
+zrb/task/base_task.py,sha256=H1D2KyJ9qK0GIPJ4kYyvfRe7fseJibmLVGjj8iRuLs4,21320
 zrb/task/base_trigger.py,sha256=jC722rDvodaBLeNaFghkTyv1u0QXrK6BLZUUqcmBJ7Q,4581
 zrb/task/cmd_task.py,sha256=pUKRSR4DZKjbmluB6vi7cxqyhxOLfJ2czSpYeQbiDvo,10705
 zrb/task/http_check.py,sha256=Gf5rOB2Se2EdizuN9rp65HpGmfZkGc-clIAlHmPVehs,2565
-zrb/task/llm_task.py,sha256=di9lYvcwd3uCPe9KzbDsXNy2ZJT49DmIIuYUS_bd8DY,19942
+zrb/task/llm_task.py,sha256=L5VFRJBgcYuucBtF3fTemjD0E-8fzA40HHeCm8T56CQ,20659
 zrb/task/make_task.py,sha256=PD3b_aYazthS8LHeJsLAhwKDEgdurQZpymJDKeN60u0,2265
 zrb/task/rsync_task.py,sha256=GSL9144bmp6F0EckT6m-2a1xG25AzrrWYzH4k3SVUKM,6370
 zrb/task/scaffolder.py,sha256=rME18w1HJUHXgi9eTYXx_T2G4JdqDYzBoNOkdOOo5-o,6806
@@ -341,7 +341,7 @@ zrb/util/string/name.py,sha256=8picJfUBXNpdh64GNaHv3om23QHhUZux7DguFLrXHp8,1163
 zrb/util/todo.py,sha256=1nDdwPc22oFoK_1ZTXyf3638Bg6sqE2yp_U4_-frHoc,16015
 zrb/xcom/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/xcom/xcom.py,sha256=o79rxR9wphnShrcIushA0Qt71d_p3ZTxjNf7x9hJB78,1571
-zrb-1.5.3.dist-info/METADATA,sha256=didD-AcEFyUlwOfXEyLCebJewXMt8hLRmLL1zp5HDZE,8470
-zrb-1.5.3.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-zrb-1.5.3.dist-info/entry_points.txt,sha256=-Pg3ElWPfnaSM-XvXqCxEAa-wfVI6BEgcs386s8C8v8,46
-zrb-1.5.3.dist-info/RECORD,,
+zrb-1.5.4.dist-info/METADATA,sha256=OeBW0SVw6MmYpkFACsrADLiJctonRhH9ZVQPKzgrjAI,8470
+zrb-1.5.4.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+zrb-1.5.4.dist-info/entry_points.txt,sha256=-Pg3ElWPfnaSM-XvXqCxEAa-wfVI6BEgcs386s8C8v8,46
+zrb-1.5.4.dist-info/RECORD,,