zrb 1.21.37__py3-none-any.whl → 1.21.43__py3-none-any.whl

zrb/task/llm/agent_runner.py

@@ -1,4 +1,9 @@
+import asyncio
 import json
+import os
+import sys
+import termios
+import tty
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
@@ -93,11 +98,17 @@ async def _run_single_agent_iteration(
             message_history=ModelMessagesTypeAdapter.validate_python(history_list),
             usage_limits=UsageLimits(request_limit=None),  # We don't want limit
         ) as agent_run:
+            escape_task = asyncio.create_task(_wait_for_escape(ctx))
             async for node in agent_run:
                 # Each node represents a step in the agent's execution
                 try:
                     await print_node(
-                        _get_plain_printer(ctx), agent_run, node, log_indent_level
+                        _get_plain_printer(ctx),
+                        agent_run,
+                        node,
+                        ctx.is_tty,
+                        log_indent_level,
+                        lambda: escape_task.done(),
                     )
                 except APIError as e:
                     # Extract detailed error information from the response
@@ -108,9 +119,56 @@ async def _run_single_agent_iteration(
                     ctx.log_error(f"Error processing node: {str(e)}")
                     ctx.log_error(f"Error type: {type(e).__name__}")
                     raise
+                if escape_task.done():
+                    break
+            # Clean escape_task
+            if not escape_task.done():
+                try:
+                    escape_task.cancel()
+                    await escape_task
+                except asyncio.CancelledError:
+                    pass
             return agent_run
 
 
+async def _wait_for_escape(ctx: AnyContext) -> None:
+    if not ctx.is_tty:
+        # Wait forever
+        await asyncio.Future()
+        return
+    loop = asyncio.get_event_loop()
+    future = loop.create_future()
+    fd = sys.stdin.fileno()
+    old_settings = termios.tcgetattr(fd)
+    try:
+        tty.setcbreak(fd)
+        loop.add_reader(fd, _create_escape_detector(ctx, future, fd))
+        await future
+    except asyncio.CancelledError:
+        raise
+    finally:
+        loop.remove_reader(fd)
+        termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
+
+
+def _create_escape_detector(
+    ctx: AnyContext, future: asyncio.Future[Any], fd: int | Any
+) -> Callable[[], None]:
+    def on_stdin():
+        try:
+            # Read just one byte
+            ch = os.read(fd, 1)
+            if ch == b"\x1b":
+                ctx.print("\n🚫 Interrupted by user.", plain=True)
+                if not future.done():
+                    future.set_result(None)
+        except Exception as e:
+            if not future.done():
+                future.set_exception(e)
+
+    return on_stdin
+
+
 def _create_print_throttle_notif(ctx: AnyContext) -> Callable[[str], None]:
     def _print_throttle_notif(text: str, *args: Any, **kwargs: Any):
         new_line = kwargs.get("new_line", True)

zrb/task/llm/default_workflow/researching/workflow.md

@@ -9,6 +9,7 @@ Follow this workflow to deliver accurate, well-sourced, and synthesized research
 - **Source Hierarchy:** Prioritize authoritative sources over secondary ones
 - **Synthesis Excellence:** Connect information into coherent narratives
 - **Comprehensive Attribution:** Cite all significant claims and data points
+- **Mandatory Citations:** ALWAYS include a "Sources" section at the end of the response with a list of all URLs used.
 
 # Tool Usage Guideline
 - Use `search_internet` for web research and information gathering
@@ -114,6 +115,7 @@ Follow this workflow to deliver accurate, well-sourced, and synthesized research
 - **Detailed Analysis:** Provide comprehensive information for deep dives
 - **Actionable Insights:** Highlight implications and recommended actions
 - **Further Reading:** Suggest additional resources for interested readers
+- **Sources:** ALWAYS end the response with a "Sources" section listing all URLs used.
 
 # Risk Assessment Guidelines
 

zrb/task/llm/print_node.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 from collections.abc import Callable
 from typing import Any
@@ -7,7 +8,12 @@ from zrb.util.cli.style import stylize_faint
 
 
 async def print_node(
-    print_func: Callable, agent_run: Any, node: Any, log_indent_level: int = 0
+    print_func: Callable,
+    agent_run: Any,
+    node: Any,
+    is_tty: bool,
+    log_indent_level: int = 0,
+    stop_check: Callable[[], bool] | None = None,
 ):
     """Prints the details of an agent execution node using a provided print function."""
     from pydantic_ai import Agent
@@ -31,13 +37,17 @@ async def print_node(
         return
     if Agent.is_model_request_node(node):
         # A model request node => We can stream tokens from the model's request
-        print_func(_format_header("🧠 Processing...", log_indent_level))
+        esc_notif = " (Press esc to cancel)" if is_tty else ""
+        print_func(_format_header(f"🧠 Processing{esc_notif}...", log_indent_level))
         # Reference: https://ai.pydantic.dev/agents/#streaming-all-events-and-output
         try:
             async with node.stream(agent_run.ctx) as request_stream:
                 is_streaming = False
                 is_tool_processing = False
                 async for event in request_stream:
+                    if stop_check and stop_check():
+                        return
+                    await asyncio.sleep(0)
                     if isinstance(event, PartStartEvent) and event.part:
                         if is_streaming:
                             print_func("")
@@ -127,6 +137,9 @@ async def print_node(
         try:
             async with node.stream(agent_run.ctx) as handle_stream:
                 async for event in handle_stream:
+                    if stop_check and stop_check():
+                        return
+                    await asyncio.sleep(0)
                     if isinstance(event, FunctionToolCallEvent):
                         args = _get_event_part_args(event)
                         call_id = event.part.tool_call_id

zrb/task/llm/prompt.py

@@ -10,7 +10,7 @@ from zrb.context.any_context import AnyContext
 from zrb.task.llm.conversation_history_model import ConversationHistory
 from zrb.task.llm.workflow import LLMWorkflow, get_available_workflows
 from zrb.util.attr import get_attr, get_str_attr, get_str_list_attr
-from zrb.util.file import read_dir, read_file_with_line_numbers
+from zrb.util.file import read_dir, read_file, read_file_with_line_numbers
 from zrb.util.markdown import make_markdown_section
 
 if TYPE_CHECKING:
@@ -32,11 +32,8 @@ def get_system_and_user_prompt(
 ) -> tuple[str, str]:
     if conversation_history is None:
         conversation_history = ConversationHistory()
-    new_user_message_prompt, apendixes = _get_user_message_prompt(user_message)
     new_system_prompt = _construct_system_prompt(
         ctx=ctx,
-        user_message=user_message,
-        apendixes=apendixes,
         persona_attr=persona_attr,
         render_persona=render_persona,
         system_prompt_attr=system_prompt_attr,
@@ -47,13 +44,12 @@ def get_system_and_user_prompt(
         render_workflows=render_workflows,
         conversation_history=conversation_history,
     )
+    new_user_message_prompt = _get_user_message_prompt(user_message)
     return new_system_prompt, new_user_message_prompt
 
 
 def _construct_system_prompt(
     ctx: AnyContext,
-    user_message: str,
-    apendixes: str,
     persona_attr: StrAttr | None = None,
     render_persona: bool = False,
     system_prompt_attr: StrAttr | None = None,
@@ -71,6 +67,7 @@ def _construct_system_prompt(
     special_instruction_prompt = _get_special_instruction_prompt(
         ctx, special_instruction_prompt_attr, render_special_instruction_prompt
     )
+    project_instructions = _get_project_instructions()
     available_workflows = get_available_workflows()
     active_workflow_names = set(
         _get_active_workflow_names(ctx, workflows_attr, render_workflows)
@@ -98,6 +95,7 @@ def _construct_system_prompt(
                     ]
                 ),
             ),
+            make_markdown_section("📜 PROJECT INSTRUCTIONS", project_instructions),
             make_markdown_section("🛠️ AVAILABLE WORKFLOWS", inactive_workflow_prompt),
             make_markdown_section(
                 "📚 CONTEXT",
@@ -122,10 +120,6 @@ def _construct_system_prompt(
                             "📝 Contextual Note Content",
                             conversation_history.contextual_note,
                         ),
-                        make_markdown_section(
-                            "📄 Apendixes",
-                            apendixes,
-                        ),
                     ]
                 ),
             ),
@@ -133,21 +127,63 @@ def _construct_system_prompt(
     )
 
 
+def _get_project_instructions() -> str:
+    instructions = []
+    cwd = os.path.abspath(os.getcwd())
+    home = os.path.abspath(os.path.expanduser("~"))
+    search_dirs = []
+    if cwd == home or cwd.startswith(os.path.join(home, "")):
+        current_dir = cwd
+        while True:
+            search_dirs.append(current_dir)
+            if current_dir == home:
+                break
+            parent_dir = os.path.dirname(current_dir)
+            if parent_dir == current_dir:
+                break
+            current_dir = parent_dir
+    else:
+        search_dirs.append(cwd)
+    for file_name in ["AGENTS.md", "CLAUDE.md"]:
+        for dir_path in search_dirs:
+            abs_file_name = os.path.join(dir_path, file_name)
+            if os.path.isfile(abs_file_name):
+                content = read_file(abs_file_name)
+                instructions.append(
+                    make_markdown_section(
+                        f"Instruction from `{abs_file_name}`", content
+                    )
+                )
+                break
+    return "\n".join(instructions)
+
+
+def _get_prompt_attr(
+    ctx: AnyContext,
+    attr: StrAttr | None,
+    render: bool,
+    default: str | None,
+) -> str:
+    """Generic helper to get a prompt attribute, prioritizing task-specific then default."""
+    value = get_attr(
+        ctx,
+        attr,
+        None,
+        auto_render=render,
+    )
+    if value is not None:
+        return value
+    return default or ""
+
+
 def _get_persona(
     ctx: AnyContext,
     persona_attr: StrAttr | None,
     render_persona: bool,
 ) -> str:
-    """Gets the persona, prioritizing task-specific, then default."""
-    persona = get_attr(
-        ctx,
-        persona_attr,
-        None,
-        auto_render=render_persona,
+    return _get_prompt_attr(
+        ctx, persona_attr, render_persona, llm_config.default_persona
     )
-    if persona is not None:
-        return persona
-    return llm_config.default_persona or ""
 
 
 def _get_base_system_prompt(
@@ -155,16 +191,9 @@ def _get_base_system_prompt(
     system_prompt_attr: StrAttr | None,
     render_system_prompt: bool,
 ) -> str:
-    """Gets the base system prompt, prioritizing task-specific, then default."""
-    system_prompt = get_attr(
-        ctx,
-        system_prompt_attr,
-        None,
-        auto_render=render_system_prompt,
+    return _get_prompt_attr(
+        ctx, system_prompt_attr, render_system_prompt, llm_config.default_system_prompt
     )
-    if system_prompt is not None:
-        return system_prompt
-    return llm_config.default_system_prompt or ""
 
 
 def _get_special_instruction_prompt(
@@ -172,16 +201,12 @@ def _get_special_instruction_prompt(
     special_instruction_prompt_attr: StrAttr | None,
     render_spcecial_instruction_prompt: bool,
 ) -> str:
-    """Gets the special instruction prompt, prioritizing task-specific, then default."""
-    special_instruction = get_attr(
+    return _get_prompt_attr(
         ctx,
         special_instruction_prompt_attr,
-        None,
-        auto_render=render_spcecial_instruction_prompt,
+        render_spcecial_instruction_prompt,
+        llm_config.default_special_instruction_prompt,
     )
-    if special_instruction is not None:
-        return special_instruction
-    return llm_config.default_special_instruction_prompt
 
 
 def _get_active_workflow_names(
@@ -229,7 +254,8 @@ def _get_workflow_prompt(
     )
 
 
-def _get_user_message_prompt(user_message: str) -> tuple[str, str]:
+def _get_user_message_prompt(user_message: str) -> str:
+    current_directory = os.getcwd()
     processed_user_message = user_message
     # Match “@” + any non-space/comma sequence that contains at least one “/”
     pattern = r"(?<!\w)@(?=[^,\s]*\/)([^,\?\!\s]+)"
@@ -247,19 +273,19 @@ def _get_user_message_prompt(user_message: str) -> tuple[str, str]:
             ref_type = "directory"
         if content != "":
             # Replace the @-reference in the user message with the placeholder
-            placeholder = f"[Reference {i+1}: `{os.path.basename(ref)}`]"
+            rel_path = os.path.relpath(resource_path, current_directory)
+            placeholder = f"`{rel_path}`"
             processed_user_message = processed_user_message.replace(
                 f"@{ref}", placeholder, 1
             )
             apendix_list.append(
                 make_markdown_section(
-                    f"Content of {placeholder} ({ref_type} path: `{resource_path}`)",
+                    f"{placeholder} {ref_type}",
                     "\n".join(content) if isinstance(content, list) else content,
                     as_code=True,
                 )
             )
     apendixes = "\n".join(apendix_list)
-    current_directory = os.getcwd()
     iso_date = datetime.now(timezone.utc).astimezone().isoformat()
     modified_user_message = make_markdown_section(
         "User Request",
@@ -269,10 +295,14 @@ def _get_user_message_prompt(user_message: str) -> tuple[str, str]:
                 f"- Current Time: {iso_date}",
                 "---",
                 processed_user_message,
+                make_markdown_section(
+                    "📄 Apendixes",
+                    apendixes,
+                ),
             ]
         ),
     )
-    return modified_user_message, apendixes
+    return modified_user_message
 
 
 def get_user_message(

zrb/task/llm/workflow.py

@@ -3,9 +3,13 @@ import os
 from zrb.config.config import CFG
 from zrb.config.llm_context.config import llm_context_config
 from zrb.config.llm_context.workflow import LLMWorkflow
+from zrb.context.any_context import AnyContext
+from zrb.xcom.xcom import Xcom
 
+LLM_LOADED_WORKFLOW_XCOM_NAME = "_llm_loaded_workflow_name"
 
-def load_workflow(workflow_name: str | list[str]) -> str:
+
+def load_workflow(ctx: AnyContext, workflow_name: str | list[str]) -> str:
     """
     Loads and formats one or more workflow documents for LLM consumption.
 
@@ -36,9 +40,17 @@ def load_workflow(workflow_name: str | list[str]) -> str:
                 ]
             )
         )
+    llm_loaded_workflow_xcom = get_llm_loaded_workflow_xcom(ctx)
+    llm_loaded_workflow_xcom.push(names)
     return "\n".join(contents)
 
 
+def get_llm_loaded_workflow_xcom(ctx: AnyContext) -> Xcom:
+    if LLM_LOADED_WORKFLOW_XCOM_NAME not in ctx.xcom:
+        ctx.xcom[LLM_LOADED_WORKFLOW_XCOM_NAME] = Xcom([])
+    return ctx.xcom[LLM_LOADED_WORKFLOW_XCOM_NAME]
+
+
 def get_available_workflows() -> dict[str, LLMWorkflow]:
     available_workflows = {
         workflow_name.strip().lower(): workflow

zrb/task/llm_task.py

@@ -82,7 +82,7 @@ class LLMTask(BaseTask):
         render_system_prompt: bool = False,
         special_instruction_prompt: "Callable[[AnyContext], str | None] | str | None" = None,
         render_special_instruction_prompt: bool = False,
-        workflows: StrListAttr | None = None,
+        workflows: "Callable[[AnyContext], list[str] | None] | list[str] | None" = None,
         render_workflows: bool = True,
         message: StrAttr | None = None,
         attachment: "UserContent | list[UserContent] | Callable[[AnyContext], UserContent | list[UserContent]] | None" = None,  # noqa
@@ -235,7 +235,58 @@ class LLMTask(BaseTask):
         self._yolo_mode = yolo_mode
 
     async def _exec_action(self, ctx: AnyContext) -> Any:
-        # Get dependent configurations first
+        # 1. Get dependent configurations
+        (
+            model_settings,
+            model,
+            yolo_mode,
+            summarization_prompt,
+            user_message,
+            attachments,
+        ) = self._get_llm_config(ctx)
+
+        # 2. Prepare initial state (read history from previous session)
+        conversation_history = await read_conversation_history(
+            ctx=ctx,
+            conversation_history_reader=self._conversation_history_reader,
+            conversation_history_file_attr=self._conversation_history_file,
+            render_history_file=self._render_history_file,
+            conversation_history_attr=self._conversation_history,
+        )
+        inject_conversation_history_notes(conversation_history)
+        inject_subagent_conversation_history_into_ctx(ctx, conversation_history)
+
+        # 3. Get system prompt and user prompt
+        system_prompt, user_prompt = self._get_prompts(
+            ctx, user_message, conversation_history
+        )
+
+        # 4. Get the agent instance
+        ctx.log_info(f"SYSTEM PROMPT:\n{system_prompt}")
+        ctx.log_info(f"USER PROMPT:\n{user_prompt}")
+        agent = self._create_agent(
+            ctx,
+            model,
+            system_prompt,
+            model_settings,
+            yolo_mode,
+            summarization_prompt,
+        )
+
+        # 5. Run the agent iteration
+        result = await self._execute_agent(
+            ctx=ctx,
+            agent=agent,
+            user_prompt=user_prompt,
+            attachments=attachments,
+            conversation_history=conversation_history,
+        )
+
+        # 6. Save history and usage
+        await self._save_history_and_usage(ctx, conversation_history)
+        return result
+
+    def _get_llm_config(self, ctx: AnyContext):
         model_settings = get_model_settings(ctx, self._model_settings)
         model = get_model(
             ctx=ctx,
@@ -258,18 +309,22 @@ class LLMTask(BaseTask):
         )
         user_message = get_user_message(ctx, self._message, self._render_message)
         attachments = get_attachments(ctx, self._attachment)
-        # 1. Prepare initial state (read history from previous session)
-        conversation_history = await read_conversation_history(
-            ctx=ctx,
-            conversation_history_reader=self._conversation_history_reader,
-            conversation_history_file_attr=self._conversation_history_file,
-            render_history_file=self._render_history_file,
-            conversation_history_attr=self._conversation_history,
+        return (
+            model_settings,
+            model,
+            yolo_mode,
+            summarization_prompt,
+            user_message,
+            attachments,
         )
-        inject_conversation_history_notes(conversation_history)
-        inject_subagent_conversation_history_into_ctx(ctx, conversation_history)
-        # 2. Get system prompt and user prompt
-        system_prompt, user_prompt = get_system_and_user_prompt(
+
+    def _get_prompts(
+        self,
+        ctx: AnyContext,
+        user_message: str,
+        conversation_history: ConversationHistory,
+    ):
+        return get_system_and_user_prompt(
             ctx=ctx,
             user_message=user_message,
             persona_attr=self._persona,
@@ -282,7 +337,16 @@ class LLMTask(BaseTask):
             render_workflows=self._render_workflows,
             conversation_history=conversation_history,
         )
-        # 3. Summarization
+
+    def _create_agent(
+        self,
+        ctx: AnyContext,
+        model,
+        system_prompt,
+        model_settings,
+        yolo_mode,
+        summarization_prompt,
+    ):
         small_model = get_model(
             ctx=ctx,
             model_attr=self._small_model,
@@ -298,10 +362,7 @@ class LLMTask(BaseTask):
             self._history_summarization_token_threshold,
             self._render_history_summarization_token_threshold,
         )
-        # 4. Get the agent instance
-        ctx.log_info(f"SYSTEM PROMPT:\n{system_prompt}")
-        ctx.log_info(f"USER PROMPT:\n{user_prompt}")
-        agent = get_agent(
+        return get_agent(
             ctx=ctx,
             model=model,
             rate_limitter=self._rate_limitter,
@@ -319,15 +380,10 @@ class LLMTask(BaseTask):
             summarization_token_threshold=summarization_token_threshold,
             history_processors=[],  # TODO: make this a property
         )
-        # 5. Run the agent iteration and save the results/history
-        result = await self._execute_agent(
-            ctx=ctx,
-            agent=agent,
-            user_prompt=user_prompt,
-            attachments=attachments,
-            conversation_history=conversation_history,
-        )
-        # 6. Write conversation history
+
+    async def _save_history_and_usage(
+        self, ctx: AnyContext, conversation_history: ConversationHistory
+    ):
         conversation_history.subagent_history = (
             extract_subagent_conversation_history_from_ctx(ctx)
         )
@@ -338,7 +394,6 @@ class LLMTask(BaseTask):
             conversation_history_file_attr=self._conversation_history_file,
             render_history_file=self._render_history_file,
         )
-        return result
 
     async def _execute_agent(
         self,

zrb/util/run.py

@@ -5,7 +5,7 @@ from typing import Any
 
 async def run_async(value: Any) -> Any:
     """
-    Run a value asynchronously, awaiting if it's awaitable or running in a thread if not.
+    Run a value asynchronously, awaiting if it's awaitable or returning it directly.
 
     Args:
         value (Any): The value to run. Can be awaitable or not.
@@ -14,7 +14,7 @@ async def run_async(value: Any) -> Any:
         Any: The result of the awaited value or the value itself if not awaitable.
     """
     if isinstance(value, asyncio.Task):
-        return value
+        return await value
     if inspect.isawaitable(value):
         return await value
-    return await asyncio.to_thread(lambda: value)
+    return value

{zrb-1.21.37.dist-info → zrb-1.21.43.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: zrb
-Version: 1.21.37
+Version: 1.21.43
 Summary: Your Automation Powerhouse
 License: AGPL-3.0-or-later
 Keywords: Automation,Task Runner,Code Generator,Monorepo,Low Code