zrb 1.15.3__py3-none-any.whl → 1.21.29__py3-none-any.whl

zrb/task/llm/prompt.py

@@ -5,20 +5,135 @@ 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.context.any_shared_context import AnySharedContext
 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,
@@ -35,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,
@@ -52,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,
@@ -69,164 +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: StrListAttr | None,
-    render_modes: bool,
+    workflows_attr: StrListAttr | None,
+    render_workflows: bool,
 ) -> list[str]:
-    """Gets the modes, prioritizing task-specific, then default."""
-    raw_modes = get_str_list_attr(
+    """Gets the workflows, prioritizing task-specific, then default."""
+    raw_workflows = get_str_list_attr(
         ctx,
-        [] if modes_attr is None else 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().lower() 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: StrListAttr | 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:
-    builtin_workflow_dir = os.path.join(os.path.dirname(__file__), "default_workflow")
-    modes = set(get_modes(ctx, modes_attr, render_modes))
-
-    # Get user-defined workflows
-    workflows = {
-        workflow_name.strip().lower(): content
-        for workflow_name, content in llm_context_config.get_workflows().items()
-        if workflow_name.strip().lower() 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 available builtin workflow names from the file system
-    available_builtin_workflow_names = set()
-    try:
-        for filename in os.listdir(builtin_workflow_dir):
-            if filename.endswith(".md"):
-                available_builtin_workflow_names.add(filename[:-3].lower())
-    except FileNotFoundError:
-        # Handle case where the directory might not exist
-        ctx.log_error(
-            f"Warning: Default workflow directory not found at {builtin_workflow_dir}"
-        )
-    except Exception as e:
-        # Catch other potential errors during directory listing
-        ctx.log_error(f"Error listing default workflows: {e}")
-
-    # Determine which builtin workflows are requested and not already loaded
-    requested_builtin_workflow_names = [
-        workflow_name
-        for workflow_name in available_builtin_workflow_names
-        if workflow_name in modes and workflow_name not in workflows
-    ]
-
-    # Add builtin-workflows if requested
-    if len(requested_builtin_workflow_names) > 0:
-        for workflow_name in requested_builtin_workflow_names:
-            workflow_file_path = os.path.join(
-                builtin_workflow_dir, f"{workflow_name}.md"
-            )
-            try:
-                with open(workflow_file_path, "r") as f:
-                    workflows[workflow_name] = f.read()
-            except FileNotFoundError:
-                ctx.log_error(
-                    f"Warning: Builtin workflow file not found: {workflow_file_path}"
-                )
-            except Exception as e:
-                ctx.log_error(f"Error reading builtin workflow {workflow_name}: {e}")
-
     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]+)"
     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 = ""
@@ -239,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(
@@ -309,7 +305,7 @@ def get_summarization_system_prompt(
 
 def get_attachments(
     ctx: AnyContext,
-    attachment: "UserContent | list[UserContent] | Callable[[AnySharedContext], UserContent | list[UserContent]] | None" = None,  # noqa
+    attachment: "UserContent | list[UserContent] | Callable[[AnyContext], UserContent | list[UserContent]] | None" = None,  # noqa
 ) -> "list[UserContent]":
     if attachment is None:
         return []

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"]