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

zrb/task/llm/prompt.py

@@ -2,18 +2,153 @@ 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",
+                            conversation_history.long_term_note,
+                        ),
+                        make_markdown_section(
+                            "📝 Contextual Note",
+                            conversation_history.contextual_note,
+                        ),
+                        make_markdown_section(
+                            "📄 Apendixes",
+                            apendixes,
+                        ),
+                    ]
+                ),
+            ),
+            make_markdown_section(
+                "💬 PAST CONVERSATION",
+                "\n".join(
+                    [
+                        make_markdown_section(
+                            "Narrative Summary",
+                            conversation_history.past_conversation_summary,
+                        ),
+                        make_markdown_section(
+                            "Past Transcript",
+                            conversation_history.past_conversation_transcript,
+                        ),
+                    ]
+                ),
+            ),
+        ]
+    )
+
+
+def _get_persona(
     ctx: AnyContext,
     persona_attr: StrAttr | None,
     render_persona: bool,
@@ -30,7 +165,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 +182,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 +199,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_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_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("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 +262,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 +316,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]