zrb 1.13.1__py3-none-any.whl → 1.21.33__py3-none-any.whl

zrb/task/llm/prompt.py

@@ -2,18 +2,138 @@ import os
 import platform
 import re
 from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Callable
 
 from zrb.attr.type import StrAttr, StrListAttr
-from zrb.config.llm_config import llm_config as llm_config
-from zrb.config.llm_context.config import llm_context_config
+from zrb.config.llm_config import llm_config
 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.llm.prompt import make_prompt_section
+from zrb.util.markdown import make_markdown_section
 
+if TYPE_CHECKING:
+    from pydantic_ai.messages import UserContent
 
-def get_persona(
+
+def get_system_and_user_prompt(
+    ctx: AnyContext,
+    user_message: str,
+    persona_attr: StrAttr | None = None,
+    render_persona: bool = False,
+    system_prompt_attr: StrAttr | None = None,
+    render_system_prompt: bool = False,
+    special_instruction_prompt_attr: StrAttr | None = None,
+    render_special_instruction_prompt: bool = False,
+    workflows_attr: StrListAttr | None = None,
+    render_workflows: bool = False,
+    conversation_history: ConversationHistory | None = None,
+) -> 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,
+        render_system_prompt=render_system_prompt,
+        special_instruction_prompt_attr=special_instruction_prompt_attr,
+        render_special_instruction_prompt=render_special_instruction_prompt,
+        workflows_attr=workflows_attr,
+        render_workflows=render_workflows,
+        conversation_history=conversation_history,
+    )
+    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,
+    render_system_prompt: bool = False,
+    special_instruction_prompt_attr: StrAttr | None = None,
+    render_special_instruction_prompt: bool = False,
+    workflows_attr: StrListAttr | None = None,
+    render_workflows: bool = False,
+    conversation_history: ConversationHistory | None = None,
+) -> str:
+    persona = _get_persona(ctx, persona_attr, render_persona)
+    base_system_prompt = _get_base_system_prompt(
+        ctx, system_prompt_attr, render_system_prompt
+    )
+    special_instruction_prompt = _get_special_instruction_prompt(
+        ctx, special_instruction_prompt_attr, render_special_instruction_prompt
+    )
+    available_workflows = get_available_workflows()
+    active_workflow_names = set(
+        _get_active_workflow_names(ctx, workflows_attr, render_workflows)
+    )
+    active_workflow_prompt = _get_workflow_prompt(
+        available_workflows, active_workflow_names, True
+    )
+    inactive_workflow_prompt = _get_workflow_prompt(
+        available_workflows, active_workflow_names, False
+    )
+    if conversation_history is None:
+        conversation_history = ConversationHistory()
+    current_directory = os.getcwd()
+    iso_date = datetime.now(timezone.utc).astimezone().isoformat()
+    return "\n".join(
+        [
+            persona,
+            base_system_prompt,
+            make_markdown_section(
+                "📝 SPECIAL INSTRUCTION",
+                "\n".join(
+                    [
+                        special_instruction_prompt,
+                        active_workflow_prompt,
+                    ]
+                ),
+            ),
+            make_markdown_section("🛠️ AVAILABLE WORKFLOWS", inactive_workflow_prompt),
+            make_markdown_section(
+                "📚 CONTEXT",
+                "\n".join(
+                    [
+                        make_markdown_section(
+                            "ℹ️ System Information",
+                            "\n".join(
+                                [
+                                    f"- OS: {platform.system()} {platform.version()}",
+                                    f"- Python Version: {platform.python_version()}",
+                                    f"- Current Directory: {current_directory}",
+                                    f"- Current Time: {iso_date}",
+                                ]
+                            ),
+                        ),
+                        make_markdown_section(
+                            "🧠 Long Term Note Content",
+                            conversation_history.long_term_note,
+                        ),
+                        make_markdown_section(
+                            "📝 Contextual Note Content",
+                            conversation_history.contextual_note,
+                        ),
+                        make_markdown_section(
+                            "📄 Apendixes",
+                            apendixes,
+                        ),
+                    ]
+                ),
+            ),
+        ]
+    )
+
+
+def _get_persona(
     ctx: AnyContext,
     persona_attr: StrAttr | None,
     render_persona: bool,
@@ -30,7 +150,7 @@ def get_persona(
     return llm_config.default_persona or ""
 
 
-def get_base_system_prompt(
+def _get_base_system_prompt(
     ctx: AnyContext,
     system_prompt_attr: StrAttr | None,
     render_system_prompt: bool,
@@ -47,7 +167,7 @@ def get_base_system_prompt(
     return llm_config.default_system_prompt or ""
 
 
-def get_special_instruction_prompt(
+def _get_special_instruction_prompt(
     ctx: AnyContext,
     special_instruction_prompt_attr: StrAttr | None,
     render_spcecial_instruction_prompt: bool,
@@ -64,138 +184,57 @@ def get_special_instruction_prompt(
     return llm_config.default_special_instruction_prompt
 
 
-def get_modes(
+def _get_active_workflow_names(
     ctx: AnyContext,
-    modes_attr: StrAttr | None,
-    render_modes: bool,
-) -> str:
-    """Gets the modes, prioritizing task-specific, then default."""
-    raw_modes = get_str_list_attr(
+    workflows_attr: StrListAttr | None,
+    render_workflows: bool,
+) -> list[str]:
+    """Gets the workflows, prioritizing task-specific, then default."""
+    raw_workflows = get_str_list_attr(
         ctx,
-        modes_attr,
-        auto_render=render_modes,
+        [] if workflows_attr is None else workflows_attr,
+        auto_render=render_workflows,
     )
-    if raw_modes is None:
-        raw_modes = []
-    modes = [mode.strip() for mode in raw_modes if mode.strip() != ""]
-    if len(modes) > 0:
-        return modes
-    return llm_config.default_modes or []
+    if raw_workflows is not None and len(raw_workflows) > 0:
+        return [w.strip().lower() for w in raw_workflows if w.strip() != ""]
+    return []
 
 
-def get_workflow_prompt(
-    ctx: AnyContext,
-    modes_attr: StrAttr | None,
-    render_modes: bool,
+def _get_workflow_prompt(
+    available_workflows: dict[str, LLMWorkflow],
+    active_workflow_names: list[str] | set[str],
+    select_active_workflow: bool,
 ) -> str:
-    modes = get_modes(ctx, modes_attr, render_modes)
-    # Get user-defined workflows
-    workflows = {
-        workflow_name: content
-        for workflow_name, content in llm_context_config.get_workflows().items()
-        if workflow_name in modes
+    selected_workflows = {
+        workflow_name: available_workflows[workflow_name]
+        for workflow_name in available_workflows
+        if (workflow_name in active_workflow_names) == select_active_workflow
     }
-    # Get requested builtin-workflow names
-    requested_builtin_workflow_names = [
-        workflow_name
-        for workflow_name in ("coding", "copywriting", "researching")
-        if workflow_name in modes and workflow_name not in workflows
-    ]
-    # add builtin-workflows if requested
-    if len(requested_builtin_workflow_names) > 0:
-        dir_path = os.path.dirname(__file__)
-        for workflow_name in requested_builtin_workflow_names:
-            workflow_file_path = os.path.join(
-                dir_path, "default_workflow", f"{workflow_name}.md"
-            )
-            with open(workflow_file_path, "r") as f:
-                workflows[workflow_name] = f.read()
     return "\n".join(
         [
-            make_prompt_section(header.capitalize(), content)
-            for header, content in workflows.items()
-            if header.lower() in modes
-        ]
-    )
-
-
-def get_system_and_user_prompt(
-    ctx: AnyContext,
-    user_message: str,
-    persona_attr: StrAttr | None = None,
-    render_persona: bool = False,
-    system_prompt_attr: StrAttr | None = None,
-    render_system_prompt: bool = False,
-    special_instruction_prompt_attr: StrAttr | None = None,
-    render_special_instruction_prompt: bool = False,
-    modes_attr: StrListAttr | None = None,
-    render_modes: bool = False,
-    conversation_history: ConversationHistory | None = None,
-) -> tuple[str, str]:
-    """Combines persona, base system prompt, and special instructions."""
-    persona = get_persona(ctx, persona_attr, render_persona)
-    base_system_prompt = get_base_system_prompt(
-        ctx, system_prompt_attr, render_system_prompt
-    )
-    special_instruction_prompt = get_special_instruction_prompt(
-        ctx, special_instruction_prompt_attr, render_special_instruction_prompt
-    )
-    workflow_prompt = get_workflow_prompt(ctx, modes_attr, render_modes)
-    if conversation_history is None:
-        conversation_history = ConversationHistory()
-    conversation_context, new_user_message = extract_conversation_context(user_message)
-    new_system_prompt = "\n".join(
-        [
-            make_prompt_section("Persona", persona),
-            make_prompt_section("System Prompt", base_system_prompt),
-            make_prompt_section("Special Instruction", special_instruction_prompt),
-            make_prompt_section("Special Workflows", workflow_prompt),
-            make_prompt_section(
-                "Past Conversation",
-                "\n".join(
-                    [
-                        make_prompt_section(
-                            "Summary",
-                            conversation_history.past_conversation_summary,
-                            as_code=True,
-                        ),
-                        make_prompt_section(
-                            "Last Transcript",
-                            conversation_history.past_conversation_transcript,
-                            as_code=True,
-                        ),
-                    ]
+            make_markdown_section(
+                workflow_name.capitalize(),
+                (
+                    (
+                        "> Workflow status: Automatically Loaded/Activated.\n"
+                        f"> Workflow location: `{workflow.path}`\n"
+                        "{workflow.content}"
+                    )
+                    if select_active_workflow
+                    else f"Workflow name: {workflow_name}\n{workflow.description}"
                 ),
-            ),
-            make_prompt_section(
-                "Notes",
-                "\n".join(
-                    [
-                        make_prompt_section(
-                            "Long Term",
-                            conversation_history.long_term_note,
-                            as_code=True,
-                        ),
-                        make_prompt_section(
-                            "Contextual",
-                            conversation_history.contextual_note,
-                            as_code=True,
-                        ),
-                    ]
-                ),
-            ),
-            make_prompt_section("Conversation Context", conversation_context),
+            )
+            for workflow_name, workflow in selected_workflows.items()
         ]
     )
-    return new_system_prompt, new_user_message
 
 
-def extract_conversation_context(user_message: str) -> tuple[str, str]:
-    modified_user_message = user_message
+def _get_user_message_prompt(user_message: str) -> tuple[str, str]:
+    processed_user_message = user_message
     # Match “@” + any non-space/comma sequence that contains at least one “/”
-    pattern = r"(?<!\w)@(?=[^,\s]*\/)([^,\s]+)"
+    pattern = r"(?<!\w)@(?=[^,\s]*\/)([^,\?\!\s]+)"
     potential_resource_path = re.findall(pattern, user_message)
-    apendixes = []
+    apendix_list = []
     for i, ref in enumerate(potential_resource_path):
         resource_path = os.path.abspath(os.path.expanduser(ref))
         content = ""
@@ -208,44 +247,32 @@ def extract_conversation_context(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)}]"
-            modified_user_message = modified_user_message.replace(
+            placeholder = f"[Reference {i+1}: `{os.path.basename(ref)}`]"
+            processed_user_message = processed_user_message.replace(
                 f"@{ref}", placeholder, 1
             )
-            apendixes.append(
-                make_prompt_section(
-                    f"{placeholder} ({ref_type} path: `{resource_path}`)",
-                    content,
+            apendix_list.append(
+                make_markdown_section(
+                    f"Content of {placeholder} ({ref_type} path: `{resource_path}`)",
+                    "\n".join(content) if isinstance(content, list) else content,
                     as_code=True,
                 )
             )
-    conversation_context = "\n".join(
-        [
-            make_prompt_section("Current OS", platform.system()),
-            make_prompt_section("OS Version", platform.version()),
-            make_prompt_section("Python Version", platform.python_version()),
-        ]
-    )
-    iso_date = datetime.now(timezone.utc).astimezone().isoformat()
+    apendixes = "\n".join(apendix_list)
     current_directory = os.getcwd()
-    modified_user_message = "\n".join(
-        [
-            make_prompt_section("User Message", modified_user_message),
-            make_prompt_section(
-                "Context",
-                "\n".join(
-                    [
-                        make_prompt_section(
-                            "Current working directory", current_directory
-                        ),
-                        make_prompt_section("Current time", iso_date),
-                        make_prompt_section("Apendixes", "\n".join(apendixes)),
-                    ]
-                ),
-            ),
-        ]
+    iso_date = datetime.now(timezone.utc).astimezone().isoformat()
+    modified_user_message = make_markdown_section(
+        "User Request",
+        "\n".join(
+            [
+                f"- Current Directory: {current_directory}",
+                f"- Current Time: {iso_date}",
+                "---",
+                processed_user_message,
+            ]
+        ),
     )
-    return conversation_context, modified_user_message
+    return modified_user_message, apendixes
 
 
 def get_user_message(
@@ -274,3 +301,21 @@ def get_summarization_system_prompt(
     if summarization_prompt is not None:
         return summarization_prompt
     return llm_config.default_summarization_prompt
+
+
+def get_attachments(
+    ctx: AnyContext,
+    attachment: "UserContent | list[UserContent] | Callable[[AnyContext], UserContent | list[UserContent]] | None" = None,  # noqa
+) -> "list[UserContent]":
+    if attachment is None:
+        return []
+    if callable(attachment):
+        result = attachment(ctx)
+        if result is None:
+            return []
+        if isinstance(result, list):
+            return result
+        return [result]
+    if isinstance(attachment, list):
+        return attachment
+    return [attachment]

zrb/task/llm/subagent_conversation_history.py

@@ -0,0 +1,41 @@
+from zrb.context.any_context import AnyContext
+from zrb.task.llm.conversation_history_model import ConversationHistory
+from zrb.task.llm.typing import ListOfDict
+from zrb.xcom.xcom import Xcom
+
+
+def inject_subagent_conversation_history_into_ctx(
+    ctx: AnyContext, conversation_history: ConversationHistory
+):
+    subagent_messages_xcom = _get_global_subagent_history_xcom(ctx)
+    existing_subagent_history = subagent_messages_xcom.get({})
+    subagent_messages_xcom.set(
+        {**existing_subagent_history, **conversation_history.subagent_history}
+    )
+
+
+def extract_subagent_conversation_history_from_ctx(
+    ctx: AnyContext,
+) -> dict[str, ListOfDict]:
+    subagent_messsages_xcom = _get_global_subagent_history_xcom(ctx)
+    return subagent_messsages_xcom.get({})
+
+
+def get_ctx_subagent_history(ctx: AnyContext, subagent_name: str) -> ListOfDict:
+    subagent_history = extract_subagent_conversation_history_from_ctx(ctx)
+    return subagent_history.get(subagent_name, [])
+
+
+def set_ctx_subagent_history(ctx: AnyContext, subagent_name: str, messages: ListOfDict):
+    subagent_history = extract_subagent_conversation_history_from_ctx(ctx)
+    subagent_history[subagent_name] = messages
+    subagent_messages_xcom = _get_global_subagent_history_xcom(ctx)
+    subagent_messages_xcom.set(subagent_history)
+
+
+def _get_global_subagent_history_xcom(ctx: AnyContext) -> Xcom:
+    if "_global_subagents" not in ctx.xcom:
+        ctx.xcom["_global_subagents"] = Xcom([{}])
+    if not isinstance(ctx.xcom["_global_subagents"], Xcom):
+        raise ValueError("ctx.xcom._global_subagents must be an Xcom")
+    return ctx.xcom["_global_subagents"]

zrb/task/llm/tool_confirmation_completer.py

@@ -0,0 +1,41 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from prompt_toolkit.completion import Completer
+
+
+def get_tool_confirmation_completer(
+    options: list[str], meta_dict: dict[str, str]
+) -> "Completer":
+    from prompt_toolkit.completion import Completer, Completion
+
+    class ToolConfirmationCompleter(Completer):
+        """Custom completer for tool confirmation that doesn't auto-complete partial words."""
+
+        def __init__(self, options, meta_dict):
+            self.options = options
+            self.meta_dict = meta_dict
+
+        def get_completions(self, document, complete_event):
+            text = document.text.strip()
+            # 1. Input is empty, OR
+            # 2. Input exactly matches the beginning of an option
+            if text == "":
+                # Show all options when nothing is typed
+                for option in self.options:
+                    yield Completion(
+                        option,
+                        start_position=0,
+                        display_meta=self.meta_dict.get(option, ""),
+                    )
+                return
+            # Only complete if text exactly matches the beginning of an option
+            for option in self.options:
+                if option.startswith(text):
+                    yield Completion(
+                        option,
+                        start_position=-len(text),
+                        display_meta=self.meta_dict.get(option, ""),
+                    )
+
+    return ToolConfirmationCompleter(options, meta_dict)