zrb 1.21.9__py3-none-any.whl → 1.21.28__py3-none-any.whl

zrb/task/llm/conversation_history.py

@@ -1,7 +1,6 @@
 import json
 import os
 from collections.abc import Callable
-from copy import deepcopy
 from typing import Any
 
 from zrb.attr.type import StrAttr
@@ -13,6 +12,37 @@ from zrb.util.attr import get_str_attr
 from zrb.util.file import read_file, write_file
 from zrb.util.markdown import make_markdown_section
 from zrb.util.run import run_async
+from zrb.xcom.xcom import Xcom
+
+
+def _get_global_subagent_messages_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"]
+
+
+def inject_subagent_history_into_ctx(
+    ctx: AnyContext, conversation_history: ConversationHistory
+):
+    subagent_messages_xcom = _get_global_subagent_messages_xcom(ctx)
+    existing_subagent_history = subagent_messages_xcom.get({})
+    subagent_messages_xcom.set(
+        {**existing_subagent_history, **conversation_history.subagent_history}
+    )
+
+
+def set_ctx_subagent_history(ctx: AnyContext, subagent_name: str, messages: ListOfDict):
+    subagent_messages_xcom = _get_global_subagent_messages_xcom(ctx)
+    subagent_history = subagent_messages_xcom.get({})
+    subagent_history[subagent_name] = messages
+    subagent_messages_xcom.set(subagent_history)
+
+
+def get_subagent_histories_from_ctx(ctx: AnyContext) -> dict[str, ListOfDict]:
+    subagent_messsages_xcom = _get_global_subagent_messages_xcom(ctx)
+    return subagent_messsages_xcom.get({})
 
 
 def inject_conversation_history_notes(conversation_history: ConversationHistory):
@@ -56,7 +86,10 @@ def get_history_file(
 
 async def _read_from_source(
     ctx: AnyContext,
-    reader: Callable[[AnyContext], dict[str, Any] | list | None] | None,
+    reader: (
+        Callable[[AnyContext], ConversationHistory | dict[str, Any] | list | None]
+        | None
+    ),
     file_path: str | None,
 ) -> "ConversationHistory | None":
     # Priority 1: Reader function
@@ -165,28 +198,6 @@ async def write_conversation_history(
         write_file(history_file, json.dumps(history_data.to_dict(), indent=2))
 
 
-def replace_system_prompt_in_history(
-    history_list: ListOfDict, replacement: str = "<main LLM system prompt>"
-) -> ListOfDict:
-    """
-    Returns a new history list where any part with part_kind 'system-prompt'
-    has its 'content' replaced with the given replacement string.
-    Args:
-        history: List of history items (each item is a dict with a 'parts' list).
-        replacement: The string to use in place of system-prompt content.
-
-    Returns:
-        A deep-copied list of history items with system-prompt content replaced.
-    """
-    new_history = deepcopy(history_list)
-    for item in new_history:
-        parts = item.get("parts", [])
-        for part in parts:
-            if part.get("part_kind") == "system-prompt":
-                part["content"] = replacement
-    return new_history
-
-
 def count_part_in_history_list(history_list: ListOfDict) -> int:
     """Calculates the total number of 'parts' in a history list."""
     history_part_len = 0

zrb/task/llm/conversation_history_model.py

@@ -10,27 +10,24 @@ class ConversationHistory:
 
     def __init__(
         self,
-        past_conversation_summary: str = "",
-        past_conversation_transcript: str = "",
         history: ListOfDict | None = None,
         contextual_note: str | None = None,
         long_term_note: str | None = None,
         project_path: str | None = None,
+        subagent_history: dict[str, ListOfDict] | None = None,
     ):
-        self.past_conversation_transcript = past_conversation_transcript
-        self.past_conversation_summary = past_conversation_summary
         self.history = history if history is not None else []
         self.contextual_note = contextual_note if contextual_note is not None else ""
         self.long_term_note = long_term_note if long_term_note is not None else ""
         self.project_path = project_path if project_path is not None else os.getcwd()
+        self.subagent_history = subagent_history if subagent_history is not None else {}
 
     def to_dict(self) -> dict[str, Any]:
         return {
-            "past_conversation_summary": self.past_conversation_summary,
-            "past_conversation_transcript": self.past_conversation_transcript,
             "history": self.history,
             "contextual_note": self.contextual_note,
             "long_term_note": self.long_term_note,
+            "subagent_history": self.subagent_history,
         }
 
     def model_dump_json(self, indent: int = 2) -> str:
@@ -44,15 +41,11 @@ class ConversationHistory:
             if isinstance(data, cls):
                 return data  # Already a valid instance
             if isinstance(data, dict):
-                # This handles both the new format and the old {'context': ..., 'history': ...}
                 return cls(
-                    past_conversation_summary=data.get("past_conversation_summary", ""),
-                    past_conversation_transcript=data.get(
-                        "past_conversation_transcript", ""
-                    ),
                     history=data.get("history", data.get("messages", [])),
                     contextual_note=data.get("contextual_note", ""),
                     long_term_note=data.get("long_term_note", ""),
+                    subagent_history=data.get("subagent_history", {}),
                 )
             elif isinstance(data, list):
                 # Handle very old format (just a list) - wrap it

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

@@ -1,5 +1,5 @@
 ---
-description: "A comprehensive workflow for software engineering tasks, including writing, modifying, and debugging code, as well as creating new applications."
+description: "A comprehensive workflow for software engineering tasks, including writing, modifying, and debugging code, as well as creating new applications. ALWAYS activate this workflow whenever you need to deal with software engineering tasks."
 ---
 
 This workflow provides a structured approach to software engineering tasks. Adhere to these guidelines to deliver high-quality, idiomatic code that respects the project's existing patterns and conventions.
@@ -35,8 +35,7 @@ Always consider if a more specific workflow is available and appropriate for the
 When requested to perform tasks like fixing bugs, adding features, refactoring, or explaining code, follow this sequence:
 1. **Understand & Strategize:** Think about the user's request and the relevant codebase context. When the task involves **complex refactoring, codebase exploration or system-wide analysis**, your **first and primary tool** must be 'codebase_investigator'. Use it to build a comprehensive understanding of the code, its structure, and dependencies. For **simple, targeted searches** (like finding a specific function name, file path, or variable declaration), you should use 'search_file_content' or 'glob' directly.
 2. **Plan:** Build a coherent and grounded (based on the understanding in step 1) plan for how you intend to resolve the user's task. Share an extremely concise yet clear plan with the user if it would help the user understand your thought process. As part of the plan, you should use an iterative development process that includes writing unit tests to verify your changes. Use output logs or debug statements as part of this process to arrive at a solution.
-3. **Implement:** Use the available tools (e.g., 'replace', 'write_file' 'run_shell_command' ...) to act on the plan, strictly adhering to the project's established conventions (detailed under 'Core 
-Mandates').
+3. **Implement:** Use the available tools (e.g., 'replace_in_file', 'write_to_file' 'run_shell_command' ...) to act on the plan, strictly adhering to the project's established conventions (detailed under 'Core Mandates').
 4. **Verify (Tests):** If applicable and feasible, verify the changes using the project's testing procedures. Identify the correct test commands and frameworks by examining 'README' files, build/package configuration (e.g., 'package.json'), or existing test execution patterns. NEVER assume standard test commands.
 5. **Verify (Standards):** VERY IMPORTANT: After making code changes, execute the project-specific build, linting and type-checking commands (e.g., 'tsc', 'npm run lint', 'ruff check .') that you have identified for this project (or obtained from the user). This ensures code quality and adherence to standards. If unsure about these commands, you can ask the user if they'd like you to run them and if so how to.
 6. **Finalize:** After all verification passes, consider the task complete. Do not remove or revert any changes or created files (like tests). Await the user's next instruction.

zrb/task/llm/file_replacement.py

@@ -0,0 +1,206 @@
+import difflib
+import os
+import shlex
+import subprocess
+import tempfile
+from typing import Any
+
+from zrb.config.config import CFG
+from zrb.task.llm.file_tool_model import FileReplacement
+from zrb.util.file import read_file
+
+
+def is_single_path_replacement(param: Any):
+    if isinstance(param, dict):
+        return _dict_has_exact_keys(
+            param, {"path", "old_text", "new_text"}
+        ) or _dict_has_exact_keys(param, {"path", "old_text", "new_text", "count"})
+    if isinstance(param, list):
+        current_path = None
+        for single_replacement in param:
+            if not is_single_path_replacement(single_replacement):
+                return False
+            if current_path is not None and current_path != single_replacement["path"]:
+                return False
+            current_path = single_replacement["path"]
+        return True
+    return False
+
+
+def _dict_has_exact_keys(dictionary: dict, required_keys: set) -> bool:
+    """
+    Check if a dictionary contains exactly the specified keys.
+    More efficient for large dictionaries.
+    """
+    if len(dictionary) != len(required_keys):
+        return False
+    return all(key in dictionary for key in required_keys)
+
+
+def edit_replacement(
+    replacement: list[FileReplacement] | FileReplacement,
+    diff_edit_command_tpl: str | None = None,
+) -> tuple[list[FileReplacement] | FileReplacement, bool]:
+    # Normalize input to list
+    replacement_list = [replacement] if isinstance(replacement, dict) else replacement
+    if not replacement_list:
+        return replacement, False
+    path = replacement_list[0]["path"]
+    original_content = read_file(path)
+    # Calculate initial proposed content based on AI's suggestion
+    proposed_content = _apply_initial_replacements(original_content, replacement_list)
+    # Open external editor for user modification
+    edited_content = _open_diff_editor(
+        path, original_content, proposed_content, diff_edit_command_tpl
+    )
+    # If content hasn't changed from proposal, return original replacement
+    if edited_content == proposed_content:
+        return replacement, False
+    # Calculate optimized replacements based on user's final edit
+    optimized_replacements = _generate_optimized_replacements(
+        path, original_content, edited_content
+    )
+    return optimized_replacements, True
+
+
+def _apply_initial_replacements(
+    content: str, replacement_list: list[FileReplacement]
+) -> str:
+    new_content = content
+    for single_replacement in replacement_list:
+        old_text = single_replacement["old_text"]
+        new_text = single_replacement["new_text"]
+        count = single_replacement.get("count", -1)
+        new_content = new_content.replace(old_text, new_text, count)
+    return new_content
+
+
+def _open_diff_editor(
+    original_path: str,
+    original_content: str,
+    proposed_content: str,
+    diff_edit_command_tpl: str | None,
+) -> str:
+    if diff_edit_command_tpl is None:
+        diff_edit_command_tpl = CFG.DEFAULT_DIFF_EDIT_COMMAND_TPL
+    _, extension = os.path.splitext(original_path)
+    with tempfile.NamedTemporaryFile(delete=False, suffix=extension) as old_file:
+        old_file_name = old_file.name
+        old_file.write(original_content.encode())
+        old_file.flush()
+        with tempfile.NamedTemporaryFile(delete=False, suffix=extension) as new_file:
+            new_file_name = new_file.name
+            new_file.write(proposed_content.encode())
+            new_file.flush()
+            diff_edit_command = diff_edit_command_tpl.format(
+                old=old_file_name, new=new_file_name
+            )
+            subprocess.call(shlex.split(diff_edit_command))
+            edited_content = read_file(new_file_name)
+    if os.path.exists(old_file_name):
+        os.remove(old_file_name)
+    if os.path.exists(new_file_name):
+        os.remove(new_file_name)
+    return edited_content
+
+
+def _generate_optimized_replacements(
+    path: str, original_content: str, edited_content: str
+) -> list[FileReplacement]:
+    matcher = difflib.SequenceMatcher(None, original_content, edited_content)
+    hunks = _group_opcodes_into_hunks(matcher.get_opcodes())
+    replacements = []
+    for hunk in hunks:
+        replacement = _create_replacement_from_hunk(
+            path, original_content, edited_content, hunk
+        )
+        if replacement:
+            replacements.append(replacement)
+    return replacements
+
+
+def _group_opcodes_into_hunks(opcodes, merge_threshold=200):
+    """
+    Groups opcodes into hunks.
+    'equal' blocks smaller than merge_threshold are treated as context (glue) within a hunk.
+    """
+    hunks = []
+    current_hunk = []
+    for tag, i1, i2, j1, j2 in opcodes:
+        if tag == "equal":
+            if i2 - i1 < merge_threshold:
+                if current_hunk:
+                    current_hunk.append((tag, i1, i2, j1, j2))
+            else:
+                if current_hunk:
+                    hunks.append(current_hunk)
+                    current_hunk = []
+        else:
+            current_hunk.append((tag, i1, i2, j1, j2))
+    if current_hunk:
+        hunks.append(current_hunk)
+    return hunks
+
+
+def _create_replacement_from_hunk(
+    path: str, original_content: str, edited_content: str, hunk: list
+) -> FileReplacement | None:
+    # Trim leading/trailing 'equal' opcodes
+    while hunk and hunk[0][0] == "equal":
+        hunk.pop(0)
+    while hunk and hunk[-1][0] == "equal":
+        hunk.pop()
+    if not hunk:
+        return None
+    # Determine range of modification
+    i_start = hunk[0][1]
+    i_end = hunk[-1][2]
+    j_start = hunk[0][3]
+    j_end = hunk[-1][4]
+    base_old_text = original_content[i_start:i_end]
+    base_new_text = edited_content[j_start:j_end]
+    if base_old_text == base_new_text:
+        return None
+    # Expand context
+    start, end = _expand_context_for_uniqueness(original_content, i_start, i_end)
+    start, end = _expand_to_word_boundary(original_content, start, end)
+    final_old_text = original_content[start:end]
+    # Reconstruct new text
+    prefix = original_content[start:i_start]
+    suffix = original_content[i_end:end]
+    final_new_text = prefix + base_new_text + suffix
+    if final_old_text == final_new_text:
+        return None
+    return {
+        "path": path,
+        "old_text": final_old_text,
+        "new_text": final_new_text,
+        "count": 1,
+    }
+
+
+def _expand_context_for_uniqueness(
+    content: str, start: int, end: int
+) -> tuple[int, int]:
+    """Expands the range [start, end] until the substring content[start:end] is unique."""
+    while content.count(content[start:end]) > 1:
+        if start == 0 and end == len(content):
+            break
+        if start > 0:
+            start -= 1
+        if end < len(content):
+            end += 1
+    return start, end
+
+
+def _expand_to_word_boundary(content: str, start: int, end: int) -> tuple[int, int]:
+    """Expands the range [start, end] outwards to the nearest whitespace boundaries."""
+
+    def is_boundary(char):
+        return char.isspace()
+
+    while start > 0 and not is_boundary(content[start - 1]):
+        start -= 1
+    while end < len(content) and not is_boundary(content[end]):
+        end += 1
+    return start, end

zrb/task/llm/file_tool_model.py

@@ -0,0 +1,57 @@
+import sys
+from typing import Literal
+
+if sys.version_info >= (3, 12):
+    from typing import NotRequired, TypedDict
+else:
+    from typing_extensions import NotRequired, TypedDict
+
+
+class FileToRead(TypedDict):
+    """
+    Configuration for reading a file or file section.
+
+    Attributes:
+        path (str): Absolute or relative path to the file
+        start_line (int | None): Starting line number (1-based, inclusive).
+            If None, reads from beginning.
+        end_line (int | None): Ending line number (1-based, exclusive). If None, reads to end.
+    """
+
+    path: str
+    start_line: NotRequired[int | None]
+    end_line: NotRequired[int | None]
+
+
+class FileToWrite(TypedDict):
+    """
+    Configuration for writing content to a file.
+
+    Attributes:
+        path (str): Absolute or relative path where file will be written.
+        content (str): Content to write. CRITICAL: For JSON, ensure all special characters
+            in this string are properly escaped.
+        mode (str): Mode for writing:
+            'w' (overwrite, default), 'a' (append), 'x' (create exclusively).
+    """
+
+    path: str
+    content: str
+    mode: NotRequired[Literal["w", "wt", "tw", "a", "at", "ta", "x", "xt", "tx"]]
+
+
+class FileReplacement(TypedDict):
+    """
+    Configuration for a single text replacement operation in a file.
+
+    Attributes:
+        path (str): Absolute or relative path to the file
+        old_text (str): Exact text to find and replace (must match file content exactly)
+        new_text (str): New text to replace with
+        count (int): Optional. Number of occurrences to replace. Defaults to -1 (all).
+    """
+
+    path: str
+    old_text: str
+    new_text: str
+    count: NotRequired[int]

zrb/task/llm/history_processor.py

@@ -0,0 +1,206 @@
+import json
+import sys
+import traceback
+from typing import TYPE_CHECKING, Any, Callable, Coroutine
+
+from zrb.config.llm_config import llm_config
+from zrb.config.llm_rate_limitter import LLMRateLimitter
+from zrb.config.llm_rate_limitter import llm_rate_limitter as default_llm_rate_limitter
+from zrb.context.any_context import AnyContext
+from zrb.task.llm.agent_runner import run_agent_iteration
+from zrb.util.cli.style import stylize_faint
+from zrb.util.markdown import make_markdown_section
+
+if sys.version_info >= (3, 12):
+    from typing import TypedDict
+else:
+    from typing_extensions import TypedDict
+
+
+if TYPE_CHECKING:
+    from pydantic_ai import ModelMessage
+    from pydantic_ai.models import Model
+    from pydantic_ai.settings import ModelSettings
+
+
+class SingleMessage(TypedDict):
+    """
+    SingleConversation
+
+    Attributes:
+        role: Either AI, User, Tool Call, or Tool Result
+        time: yyyy-mm-ddTHH:MM:SSZ:
+        content: The content of the message (summarize if too long)
+    """
+
+    role: str
+    time: str
+    content: str
+
+
+class ConversationSummary(TypedDict):
+    """
+    Conversation history
+
+    Attributes:
+        transcript: Several last transcript of the conversation
+        summary: Descriptive conversation summary
+    """
+
+    transcript: list[SingleMessage]
+    summary: str
+
+
+def save_conversation_summary(conversation_summary: ConversationSummary):
+    """
+    Write conversation summary for main assistant to continue conversation.
+    """
+    return conversation_summary
+
+
+def create_summarize_history_processor(
+    ctx: AnyContext,
+    system_prompt: str,
+    rate_limitter: LLMRateLimitter | None = None,
+    summarization_model: "Model | str | None" = None,
+    summarization_model_settings: "ModelSettings | None" = None,
+    summarization_system_prompt: str | None = None,
+    summarization_token_threshold: int | None = None,
+    summarization_retries: int = 2,
+) -> Callable[[list["ModelMessage"]], Coroutine[None, None, list["ModelMessage"]]]:
+    from pydantic_ai import Agent, ModelMessage, ModelRequest
+    from pydantic_ai.messages import ModelMessagesTypeAdapter, UserPromptPart
+
+    if rate_limitter is None:
+        rate_limitter = default_llm_rate_limitter
+    if summarization_model is None:
+        summarization_model = llm_config.default_small_model
+    if summarization_model_settings is None:
+        summarization_model_settings = llm_config.default_small_model_settings
+    if summarization_system_prompt is None:
+        summarization_system_prompt = llm_config.default_summarization_prompt
+    if summarization_token_threshold is None:
+        summarization_token_threshold = (
+            llm_config.default_history_summarization_token_threshold
+        )
+
+    async def maybe_summarize_history(
+        messages: list[ModelMessage],
+    ) -> list[ModelMessage]:
+        history_list = json.loads(ModelMessagesTypeAdapter.dump_json(messages))
+        history_json_str = json.dumps(history_list)
+        # Estimate token usage
+        # Note: Pydantic ai has run context parameter
+        # (https://ai.pydantic.dev/message-history/#runcontext-parameter)
+        # But we cannot use run_ctx.usage.total_tokens because total token keep increasing
+        # even after summariztion.
+        estimated_token_usage = rate_limitter.count_token(history_json_str)
+        _print_request_info(
+            ctx, estimated_token_usage, summarization_token_threshold, messages
+        )
+        if estimated_token_usage < summarization_token_threshold or len(messages) == 1:
+            return messages
+        history_list_without_instruction = [
+            {
+                key: obj[key]
+                for key in obj
+                if index == len(history_list) - 1 or key != "instructions"
+            }
+            for index, obj in enumerate(history_list)
+        ]
+        history_json_str_without_instruction = json.dumps(
+            history_list_without_instruction
+        )
+        summarization_message = f"Summarize the following conversation: {history_json_str_without_instruction}"
+        summarization_agent = Agent[None, ConversationSummary](
+            model=summarization_model,
+            output_type=save_conversation_summary,
+            instructions=summarization_system_prompt,
+            model_settings=summarization_model_settings,
+            retries=summarization_retries,
+        )
+        try:
+            _print_info(ctx, "📝 Rollup Conversation", 2)
+            summary_run = await run_agent_iteration(
+                ctx=ctx,
+                agent=summarization_agent,
+                user_prompt=summarization_message,
+                attachments=[],
+                history_list=[],
+                rate_limitter=rate_limitter,
+                log_indent_level=2,
+            )
+            if summary_run and summary_run.result and summary_run.result.output:
+                usage = summary_run.result.usage()
+                _print_info(ctx, f"📝 Rollup Conversation Token: {usage}", 2)
+                ctx.print(plain=True)
+                ctx.log_info("History summarized and updated.")
+                condensed_message = make_markdown_section(
+                    header="Past Conversation",
+                    content="\n".join(
+                        [
+                            make_markdown_section(
+                                "Summary", _extract_summary(summary_run.result.output)
+                            ),
+                            make_markdown_section(
+                                "Past Trancript",
+                                _extract_transcript(summary_run.result.output),
+                            ),
+                        ]
+                    ),
+                )
+                return [
+                    ModelRequest(
+                        instructions=system_prompt,
+                        parts=[UserPromptPart(condensed_message)],
+                    )
+                ]
+            ctx.log_warning("History summarization failed or returned no data.")
+        except BaseException as e:
+            ctx.log_warning(f"Error during history summarization: {e}")
+            traceback.print_exc()
+        return messages
+
+    return maybe_summarize_history
+
+
+def _print_request_info(
+    ctx: AnyContext,
+    estimated_token_usage: int,
+    summarization_token_threshold: int,
+    messages: list["ModelMessage"],
+):
+    _print_info(ctx, f"Current request token (estimated): {estimated_token_usage}")
+    _print_info(ctx, f"Summarization token threshold: {summarization_token_threshold}")
+    _print_info(ctx, f"History length: {len(messages)}")
+
+
+def _print_info(ctx: AnyContext, text: str, log_indent_level: int = 0):
+    log_prefix = (2 * (log_indent_level + 1)) * " "
+    ctx.print(stylize_faint(f"{log_prefix}{text}"), plain=True)
+
+
+def _extract_summary(summary_result_output: dict[str, Any] | str) -> str:
+    summary = (
+        summary_result_output.get("summary", "")
+        if isinstance(summary_result_output, dict)
+        else ""
+    )
+    return summary
+
+
+def _extract_transcript(summary_result_output: dict[str, Any] | str) -> str:
+    transcript_list = (
+        summary_result_output.get("transcript", [])
+        if isinstance(summary_result_output, dict)
+        else []
+    )
+    transcript_list = [] if not isinstance(transcript_list, list) else transcript_list
+    return "\n".join(_format_transcript_message(message) for message in transcript_list)
+
+
+def _format_transcript_message(message: dict[str, str]) -> str:
+    role = message.get("role", "Message")
+    time = message.get("time", "<unknown>")
+    content = message.get("content", "<empty>")
+    return f"{role} ({time}): {content}"