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

zrb/task/llm_task.py

@@ -1,28 +1,15 @@
-import functools
 import inspect
 import json
-import os
-import traceback
 from collections.abc import Callable
+from textwrap import dedent
 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,
-    FunctionToolResultEvent,
-    ModelMessagesTypeAdapter,
-    PartDeltaEvent,
-    PartStartEvent,
-    TextPartDelta,
-    ToolCallPartDelta,
-)
 from pydantic_ai.models import Model
 from pydantic_ai.settings import ModelSettings
 
-from zrb.attr.type import StrAttr, fstring
+from zrb.attr.type import BoolAttr, IntAttr, StrAttr, fstring
 from zrb.context.any_context import AnyContext
 from zrb.context.any_shared_context import AnySharedContext
 from zrb.env.any_env import AnyEnv
@@ -31,12 +18,18 @@ from zrb.llm_config import LLMConfig
 from zrb.llm_config import llm_config as default_llm_config
 from zrb.task.any_task import AnyTask
 from zrb.task.base_task import BaseTask
-from zrb.util.attr import get_attr, get_str_attr
-from zrb.util.cli.style import stylize_faint
-from zrb.util.file import read_file, write_file
+from zrb.task.llm.agent_runner import run_agent_iteration
+from zrb.task.llm.context_enricher import EnrichmentConfig, enrich_context
+from zrb.task.llm.default_context import get_default_context
+from zrb.task.llm.history import ConversationHistoryData, ListOfDict
+from zrb.task.llm.history_summarizer import SummarizationConfig, summarize_history
+from zrb.task.llm.tool_wrapper import wrap_tool
+from zrb.util.attr import get_attr, get_bool_attr, get_int_attr, get_str_attr
+from zrb.util.file import write_file
 from zrb.util.run import run_async
 
-ListOfDict = list[dict[str, Any]]
+# ListOfDict moved to history.py
+# Removed old ConversationHistoryData type alias
 ToolOrCallable = Tool | Callable
 
 
@@ -65,6 +58,12 @@ class LLMTask(BaseTask):
         system_prompt: StrAttr | None = None,
         render_system_prompt: bool = True,
         message: StrAttr | None = None,
+        summarization_prompt: StrAttr | None = None,
+        render_summarization_prompt: bool = True,
+        enrich_context: BoolAttr = False,
+        render_enrich_context: bool = True,
+        context_enrichment_prompt: StrAttr | None = None,
+        render_context_enrichment_prompt: bool = True,
         tools: (
             list[ToolOrCallable] | Callable[[AnySharedContext], list[ToolOrCallable]]
         ) = [],
@@ -72,16 +71,29 @@ class LLMTask(BaseTask):
             list[MCPServer] | Callable[[AnySharedContext], list[MCPServer]]
         ) = [],
         conversation_history: (
-            ListOfDict | Callable[[AnySharedContext], ListOfDict]
-        ) = [],
+            ConversationHistoryData  # Use the new BaseModel
+            | Callable[
+                [AnySharedContext], ConversationHistoryData | dict | list
+            ]  # Allow returning raw dict/list too
+            | dict  # Allow raw dict
+            | list  # Allow raw list (old format)
+        ) = ConversationHistoryData(),  # Default to an empty model instance
         conversation_history_reader: (
-            Callable[[AnySharedContext], ListOfDict] | None
+            Callable[[AnySharedContext], ConversationHistoryData | dict | list | None]
+            | None
+            # Allow returning raw dict/list or None
         ) = None,
         conversation_history_writer: (
-            Callable[[AnySharedContext, ListOfDict], None] | None
+            Callable[[AnySharedContext, ConversationHistoryData], None]
+            | None
+            # Writer expects the model instance
         ) = None,
         conversation_history_file: StrAttr | None = None,
         render_history_file: bool = True,
+        summarize_history: BoolAttr = True,
+        render_summarize_history: bool = True,
+        history_summarization_threshold: IntAttr = 5,  # -1 means no summarization trigger
+        render_history_summarization_threshold: bool = True,
         execute_condition: bool | str | Callable[[AnySharedContext], bool] = True,
         retries: int = 2,
         retry_period: float = 0,
@@ -95,6 +107,9 @@ class LLMTask(BaseTask):
         upstream: list[AnyTask] | AnyTask | None = None,
         fallback: list[AnyTask] | AnyTask | None = None,
         successor: list[AnyTask] | AnyTask | None = None,
+        conversation_context: (
+            dict[str, Any] | Callable[[AnySharedContext], dict[str, Any]] | None
+        ) = None,
     ):
         super().__init__(
             name=name,
@@ -128,6 +143,12 @@ class LLMTask(BaseTask):
         self._system_prompt = system_prompt
         self._render_system_prompt = render_system_prompt
         self._message = message
+        self._summarization_prompt = summarization_prompt
+        self._render_summarization_prompt = render_summarization_prompt
+        self._should_enrich_context = enrich_context
+        self._render_enrich_context = render_enrich_context
+        self._context_enrichment_prompt = context_enrichment_prompt
+        self._render_context_enrichment_prompt = render_context_enrichment_prompt
         self._tools = tools
         self._additional_tools: list[ToolOrCallable] = []
         self._mcp_servers = mcp_servers
@@ -137,7 +158,14 @@ class LLMTask(BaseTask):
         self._conversation_history_writer = conversation_history_writer
         self._conversation_history_file = conversation_history_file
         self._render_history_file = render_history_file
+        self._should_summarize_history = summarize_history
+        self._render_summarize_history = render_summarize_history
+        self._history_summarization_threshold = history_summarization_threshold
+        self._render_history_summarization_threshold = (
+            render_history_summarization_threshold
+        )
         self._max_call_iteration = max_call_iteration
+        self._conversation_context = conversation_context
 
     def add_tool(self, tool: ToolOrCallable):
         self._additional_tools.append(tool)
@@ -145,128 +173,90 @@ class LLMTask(BaseTask):
     def add_mcp_server(self, mcp_server: MCPServer):
         self._additional_mcp_servers.append(mcp_server)
 
+    def set_should_enrich_context(self, enrich_context: bool):
+        self._should_enrich_context = enrich_context
+
+    def set_should_summarize_history(self, summarize_history: bool):
+        self._should_summarize_history = summarize_history
+
+    def set_history_summarization_threshold(self, summarization_threshold: int):
+        self._history_summarization_threshold = summarization_threshold
+
     async def _exec_action(self, ctx: AnyContext) -> Any:
-        history = await self._read_conversation_history(ctx)
-        user_prompt = self._get_message(ctx)
+        history_data: ConversationHistoryData = await self._read_conversation_history(
+            ctx
+        )
+        # Extract history list and conversation context
+        history_list = history_data.history
+        conversation_context = self._get_conversation_context(ctx)
+        # Merge history context without overwriting existing keys
+        for key, value in history_data.context.items():
+            if key not in conversation_context:
+                conversation_context[key] = value
+        # Enrich context based on history (if enabled)
+        if self._get_should_enrich_context(ctx, history_list):
+            conversation_context = await enrich_context(
+                ctx=ctx,
+                config=EnrichmentConfig(
+                    model=self._get_model(ctx),
+                    settings=self._get_model_settings(ctx),
+                    prompt=self._get_context_enrichment_prompt(ctx),
+                ),
+                conversation_context=conversation_context,
+                history_list=history_list,
+            )
+        # Get history handling parameters
+        if self._get_should_summarize_history(ctx, history_list):
+            ctx.log_info("Summarize previous conversation")
+            # Summarize the part to be removed and update context
+            conversation_context = await summarize_history(
+                ctx=ctx,
+                config=SummarizationConfig(
+                    model=self._get_model(ctx),
+                    settings=self._get_model_settings(ctx),
+                    prompt=self._get_summarization_prompt(ctx),
+                ),
+                conversation_context=conversation_context,
+                history_list=history_list,  # Pass the full list for context
+            )
+            # Truncate the history list after summarization
+            history_list = []
+        # Construct user prompt
+        user_prompt = self._get_user_prompt(ctx, conversation_context)
+        # Create and run agent
         agent = self._get_agent(ctx)
         try:
-            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
+            agent_run = await run_agent_iteration(
+                ctx=ctx,
+                agent=agent,
+                user_prompt=user_prompt,
+                history_list=history_list,
+            )
+            if agent_run:
+                new_history_list = json.loads(agent_run.result.all_messages_json())
+                data_to_write = ConversationHistoryData(
+                    context=conversation_context,
+                    history=new_history_list,
+                )
+                await self._write_conversation_history(
+                    ctx, data_to_write
+                )  # Pass the model instance
+                return agent_run.result.data
         except Exception as e:
             ctx.log_error(f"Error in agent execution: {str(e)}")
             raise
 
-    async def _print_node(self, ctx: AnyContext, agent_run: Any, node: Any):
-        if Agent.is_user_prompt_node(node):
-            # A user prompt node => The user has provided input
-            ctx.print(stylize_faint(f">> UserPromptNode: {node.user_prompt}"))
-        elif Agent.is_model_request_node(node):
-            # A model request node => We can stream tokens from the model"s request
-            ctx.print(
-                stylize_faint(">> ModelRequestNode: streaming partial request tokens")
-            )
-            async with node.stream(agent_run.ctx) as request_stream:
-                is_streaming = False
-                async for event in request_stream:
-                    if isinstance(event, PartStartEvent):
-                        if is_streaming:
-                            ctx.print("", plain=True)
-                        ctx.print(
-                            stylize_faint(
-                                f"[Request] Starting part {event.index}: {event.part!r}"
-                            ),
-                        )
-                        is_streaming = False
-                    elif isinstance(event, PartDeltaEvent):
-                        if isinstance(event.delta, TextPartDelta):
-                            ctx.print(
-                                stylize_faint(f"{event.delta.content_delta}"),
-                                end="",
-                                plain=is_streaming,
-                            )
-                        elif isinstance(event.delta, ToolCallPartDelta):
-                            ctx.print(
-                                stylize_faint(f"{event.delta.args_delta}"),
-                                end="",
-                                plain=is_streaming,
-                            )
-                        is_streaming = True
-                    elif isinstance(event, FinalResultEvent):
-                        if is_streaming:
-                            ctx.print("", plain=True)
-                        ctx.print(
-                            stylize_faint(f"[Result] tool_name={event.tool_name}"),
-                        )
-                        is_streaming = False
-                if is_streaming:
-                    ctx.print("", plain=True)
-        elif Agent.is_call_tools_node(node):
-            # A handle-response node => The model returned some data, potentially calls a tool
-            ctx.print(
-                stylize_faint(
-                    ">> CallToolsNode: streaming partial response & tool usage"
-                )
-            )
-            async with node.stream(agent_run.ctx) as handle_stream:
-                async for event in handle_stream:
-                    if isinstance(event, FunctionToolCallEvent):
-                        # Handle empty arguments across different providers
-                        if event.part.args == "" or event.part.args is None:
-                            event.part.args = {}
-                        elif isinstance(
-                            event.part.args, str
-                        ) and event.part.args.strip() in ["null", "{}"]:
-                            # Some providers might send "null" or "{}" as a string
-                            event.part.args = {}
-                        # Handle dummy property if present (from our schema sanitization)
-                        if (
-                            isinstance(event.part.args, dict)
-                            and "_dummy" in event.part.args
-                        ):
-                            del event.part.args["_dummy"]
-                        ctx.print(
-                            stylize_faint(
-                                f"[Tools] The LLM calls tool={event.part.tool_name!r} with args={event.part.args} (tool_call_id={event.part.tool_call_id!r})"  # noqa
-                            )
-                        )
-                    elif isinstance(event, FunctionToolResultEvent):
-                        ctx.print(
-                            stylize_faint(
-                                f"[Tools] Tool call {event.tool_call_id!r} returned => {event.result.content}"  # noqa
-                            )
-                        )
-        elif Agent.is_end_node(node):
-            # Once an End node is reached, the agent run is complete
-            ctx.print(stylize_faint(f"{agent_run.result.data}"))
-
     async def _write_conversation_history(
-        self, ctx: AnyContext, conversations: list[Any]
+        self, ctx: AnyContext, history_data: ConversationHistoryData
     ):
+        # Expects the model instance
         if self._conversation_history_writer is not None:
-            await run_async(self._conversation_history_writer(ctx, conversations))
+            # Pass the model instance directly to the writer
+            await run_async(self._conversation_history_writer(ctx, history_data))
         history_file = self._get_history_file(ctx)
         if history_file != "":
-            write_file(history_file, json.dumps(conversations, indent=2))
+            # Use model_dump_json for serialization
+            write_file(history_file, history_data.model_dump_json(indent=2))
 
     def _get_model_settings(self, ctx: AnyContext) -> ModelSettings | None:
         if callable(self._model_settings):
@@ -282,10 +272,17 @@ class LLMTask(BaseTask):
             self._tools(ctx) if callable(self._tools) else self._tools
         )
         tools_or_callables.extend(self._additional_tools)
-        tools = [
-            tool if isinstance(tool, Tool) else Tool(_wrap_tool(tool), takes_ctx=False)
-            for tool in tools_or_callables
-        ]
+        tools = []
+        for tool_or_callable in tools_or_callables:
+            if isinstance(tool_or_callable, Tool):
+                tools.append(tool_or_callable)
+            else:
+                # Inspect original callable for 'ctx' parameter
+                # This ctx refer to pydantic AI's ctx, not task ctx.
+                original_sig = inspect.signature(tool_or_callable)
+                takes_ctx = "ctx" in original_sig.parameters
+                wrapped_tool = wrap_tool(tool_or_callable)
+                tools.append(Tool(wrapped_tool, takes_ctx=takes_ctx))
         mcp_servers = list(
             self._mcp_servers(ctx) if callable(self._mcp_servers) else self._mcp_servers
         )
@@ -345,22 +342,90 @@ class LLMTask(BaseTask):
             return system_prompt
         return default_llm_config.get_default_system_prompt()
 
-    def _get_message(self, ctx: AnyContext) -> str:
+    def _get_user_prompt(
+        self, ctx: AnyContext, conversation_context: dict[str, Any]
+    ) -> str:
+        user_message = self._get_user_message(ctx)
+        enriched_context = {**get_default_context(user_message), **conversation_context}
+        return dedent(
+            f"""
+            # Context
+            {json.dumps(enriched_context)}
+            # User Message
+            {user_message}
+            """.strip()
+        )
+
+    def _get_user_message(self, ctx: AnyContext) -> str:
         return get_str_attr(ctx, self._message, "How are you?", auto_render=True)
 
-    async def _read_conversation_history(self, ctx: AnyContext) -> ListOfDict:
-        if self._conversation_history_reader is not None:
-            return await run_async(self._conversation_history_reader(ctx))
-        if callable(self._conversation_history):
-            return self._conversation_history(ctx)
+    def _get_summarization_prompt(self, ctx: AnyContext) -> str:
+        summarization_prompt = get_attr(
+            ctx,
+            self._summarization_prompt,
+            None,
+            auto_render=self._render_summarization_prompt,
+        )
+        if summarization_prompt is not None:
+            return summarization_prompt
+        return default_llm_config.get_default_summarization_prompt()
+
+    def _get_should_enrich_context(
+        self, ctx: AnyContext, history_list: ListOfDict
+    ) -> bool:
+        if len(history_list) == 0:
+            return False
+        return get_bool_attr(
+            ctx,
+            self._should_enrich_context,
+            True,  # Default to True if not specified
+            auto_render=self._render_enrich_context,
+        )
+
+    def _get_context_enrichment_prompt(self, ctx: AnyContext) -> str:
+        context_enrichment_prompt = get_attr(
+            ctx,
+            self._context_enrichment_prompt,
+            None,
+            auto_render=self._render_context_enrichment_prompt,
+        )
+        if context_enrichment_prompt is not None:
+            return context_enrichment_prompt
+        return default_llm_config.get_default_context_enrichment_prompt()
+
+    async def _read_conversation_history(
+        self, ctx: AnyContext
+    ) -> ConversationHistoryData:  # Returns the model instance
+        """Reads conversation history from reader, file, or attribute, with validation."""
         history_file = self._get_history_file(ctx)
-        if (
-            len(self._conversation_history) == 0
-            and history_file != ""
-            and os.path.isfile(history_file)
-        ):
-            return json.loads(read_file(history_file))
-        return self._conversation_history
+        # Priority 1 & 2: Reader and File (handled by ConversationHistoryData)
+        history_data = await ConversationHistoryData.read_from_sources(
+            ctx=ctx,
+            reader=self._conversation_history_reader,
+            file_path=history_file,
+        )
+        if history_data:
+            return history_data
+        # Priority 3: Callable or direct conversation_history attribute
+        raw_data_attr: Any = None
+        if callable(self._conversation_history):
+            try:
+                raw_data_attr = await run_async(self._conversation_history(ctx))
+            except Exception as e:
+                ctx.log_warning(
+                    f"Error executing callable conversation_history attribute: {e}. "
+                    "Ignoring."
+                )
+        if raw_data_attr is None:
+            raw_data_attr = self._conversation_history
+        if raw_data_attr:
+            history_data = ConversationHistoryData.parse_and_validate(
+                ctx, raw_data_attr, "attribute"
+            )
+            if history_data:
+                return history_data
+        # Fallback: Return default value
+        return ConversationHistoryData()
 
     def _get_history_file(self, ctx: AnyContext) -> str:
         return get_str_attr(
@@ -370,103 +435,71 @@ class LLMTask(BaseTask):
             auto_render=self._render_history_file,
         )
 
-
-def _wrap_tool(func):
-    sig = inspect.signature(func)
-    if len(sig.parameters) == 0:
-
-        @functools.wraps(func)
-        async def wrapper(_dummy=None):
-            try:
-                return await run_async(func())
-            except Exception as e:
-                # Optionally, you can include more details from traceback if needed.
-                error_details = traceback.format_exc()
-                return f"Error: {e}\nDetails: {error_details}"
-
-        new_sig = inspect.Signature(
-            parameters=[
-                inspect.Parameter(
-                    "_dummy", inspect.Parameter.POSITIONAL_OR_KEYWORD, default=None
-                )
-            ]
+    def _get_should_summarize_history(
+        self, ctx: AnyContext, history_list: ListOfDict
+    ) -> bool:
+        history_len = len(history_list)
+        if history_len == 0:
+            return False
+        summarization_threshold = self._get_history_summarization_threshold(ctx)
+        if summarization_threshold == -1:
+            return False
+        if summarization_threshold > history_len:
+            return False
+        return get_bool_attr(
+            ctx,
+            self._should_summarize_history,
+            False,
+            auto_render=self._render_summarize_history,
         )
-        # Override the wrapper's signature so introspection yields a non-empty schema.
-        wrapper.__signature__ = new_sig
-        return wrapper
-    else:
-
-        @functools.wraps(func)
-        async def wrapper(*args, **kwargs):
-            try:
-                return await run_async(func(*args, **kwargs))
-            except Exception as e:
-                # Optionally, you can include more details from traceback if needed.
-                error_details = traceback.format_exc()
-                return f"Error: {e}\nDetails: {error_details}"
 
-        return wrapper
-
-
-def _extract_api_error_details(error: APIError) -> str:
-    """Extract detailed error information from an APIError."""
-    details = f"{error.message}"
-    # Try to parse the error body as JSON
-    if error.body:
+    def _get_history_summarization_threshold(self, ctx: AnyContext) -> int:
+        # Use get_int_attr with -1 as default (no limit)
         try:
-            if isinstance(error.body, str):
-                body_json = json.loads(error.body)
-            elif isinstance(error.body, bytes):
-                body_json = json.loads(error.body.decode("utf-8"))
-            else:
-                body_json = error.body
-            # Extract error details from the JSON structure
-            if isinstance(body_json, dict):
-                if "error" in body_json:
-                    error_obj = body_json["error"]
-                    if isinstance(error_obj, dict):
-                        if "message" in error_obj:
-                            details += f"\nProvider message: {error_obj['message']}"
-                        if "code" in error_obj:
-                            details += f"\nError code: {error_obj['code']}"
-                        if "status" in error_obj:
-                            details += f"\nStatus: {error_obj['status']}"
-                # Check for metadata that might contain provider-specific information
-                if "metadata" in body_json and isinstance(body_json["metadata"], dict):
-                    metadata = body_json["metadata"]
-                    if "provider_name" in metadata:
-                        details += f"\nProvider: {metadata['provider_name']}"
-                    if "raw" in metadata:
-                        try:
-                            raw_json = json.loads(metadata["raw"])
-                            if "error" in raw_json and isinstance(
-                                raw_json["error"], dict
-                            ):
-                                raw_error = raw_json["error"]
-                                if "message" in raw_error:
-                                    details += (
-                                        f"\nRaw error message: {raw_error['message']}"
-                                    )
-                        except (KeyError, TypeError, ValueError):
-                            # If we can't parse the raw JSON, just include it as is
-                            details += f"\nRaw error data: {metadata['raw']}"
-        except json.JSONDecodeError:
-            # If we can't parse the JSON, include the raw body
-            details += f"\nRaw error body: {error.body}"
-        except Exception as e:
-            # Catch any other exceptions during parsing
-            details += f"\nError parsing error body: {str(e)}"
-    # Include request information if available
-    if hasattr(error, "request") and error.request:
-        if hasattr(error.request, "method") and hasattr(error.request, "url"):
-            details += f"\nRequest: {error.request.method} {error.request.url}"
-        # Include a truncated version of the request content if available
-        if hasattr(error.request, "content") and error.request.content:
-            content = error.request.content
-            if isinstance(content, bytes):
+            return get_int_attr(
+                ctx,
+                self._history_summarization_threshold,
+                -1,
+                auto_render=self._render_history_summarization_threshold,
+            )
+        except ValueError as e:
+            ctx.log_warning(
+                f"Could not convert history_summarization_threshold to int: {e}. "
+                "Defaulting to -1 (no threshold)."
+            )
+            return -1
+
+    def _get_conversation_context(self, ctx: AnyContext) -> dict[str, Any]:
+        """
+        Retrieves the conversation context.
+        If a value in the context dict is callable, it executes it with ctx.
+        """
+        raw_context = get_attr(
+            ctx, self._conversation_context, {}, auto_render=False
+        )  # Context usually shouldn't be rendered
+        if not isinstance(raw_context, dict):
+            ctx.log_warning(
+                f"Conversation context resolved to type {type(raw_context)}, "
+                "expected dict. Returning empty context."
+            )
+            return {}
+        # If conversation_context contains callable value, execute them.
+        processed_context: dict[str, Any] = {}
+        for key, value in raw_context.items():
+            if callable(value):
                 try:
-                    content = content.decode("utf-8")
-                except UnicodeDecodeError:
-                    content = str(content)
-            details += f"\nRequest content: {content}"
-    return details
+                    # Check if the callable expects 'ctx'
+                    sig = inspect.signature(value)
+                    if "ctx" in sig.parameters:
+                        processed_context[key] = value(ctx)
+                    else:
+                        processed_context[key] = value()
+                except Exception as e:
+                    ctx.log_warning(
+                        f"Error executing callable for context key '{key}': {e}. "
+                        "Skipping."
+                    )
+                    processed_context[key] = None
+            else:
+                processed_context[key] = value
+        return processed_context

zrb/util/file.py

@@ -3,13 +3,19 @@ import re
 
 
 def read_file(file_path: str, replace_map: dict[str, str] = {}) -> str:
-    with open(file_path, "r", encoding="utf-8") as f:
+    with open(
+        os.path.abspath(os.path.expanduser(file_path)), "r", encoding="utf-8"
+    ) as f:
         content = f.read()
     for key, val in replace_map.items():
         content = content.replace(key, val)
     return content
 
 
+def read_dir(dir_path: str) -> list[str]:
+    return [f for f in os.listdir(os.path.abspath(os.path.expanduser(dir_path)))]
+
+
 def write_file(file_path: str, content: str | list[str]):
     if isinstance(content, list):
         content = "\n".join([line for line in content if line is not None])
@@ -21,5 +27,5 @@ def write_file(file_path: str, content: str | list[str]):
     content = content.rstrip("\n")
     if should_add_eol:
         content += "\n"
-    with open(file_path, "w") as f:
+    with open(os.path.abspath(os.path.expanduser(file_path)), "w") as f:
         f.write(content)

zrb/util/load.py

@@ -23,6 +23,8 @@ def load_file(script_path: str, sys_path_index: int = 0) -> Any | None:
     # Add script dir path to Python path
     os.environ["PYTHONPATH"] = _get_new_python_path(script_dir_path)
     spec = importlib.util.spec_from_file_location(module_name, script_path)
+    if spec is None:
+        return None
     module = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(module)
     return module