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

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

zrb/task/llm/history_summarization.py

@@ -1,36 +1,7 @@
-import json
-import traceback
-from typing import TYPE_CHECKING
-
-from zrb.attr.type import BoolAttr, IntAttr
+from zrb.attr.type import IntAttr
 from zrb.config.llm_config import llm_config
-from zrb.config.llm_rate_limitter import LLMRateLimiter, llm_rate_limitter
 from zrb.context.any_context import AnyContext
-from zrb.task.llm.agent import run_agent_iteration
-from zrb.task.llm.conversation_history import (
-    count_part_in_history_list,
-    inject_conversation_history_notes,
-    replace_system_prompt_in_history,
-)
-from zrb.task.llm.conversation_history_model import ConversationHistory
-from zrb.task.llm.history_summarization_tool import (
-    create_history_summarization_tool,
-)
-from zrb.task.llm.typing import ListOfDict
-from zrb.util.attr import get_bool_attr, get_int_attr
-from zrb.util.cli.style import stylize_faint
-from zrb.util.markdown import make_markdown_section
-from zrb.util.truncate import truncate_str
-
-if TYPE_CHECKING:
-    from pydantic_ai.models import Model
-    from pydantic_ai.settings import ModelSettings
-
-
-def _count_token_in_history(history_list: ListOfDict) -> int:
-    """Counts the total number of tokens in a conversation history list."""
-    text_to_count = json.dumps(history_list)
-    return llm_rate_limitter.count_token(text_to_count)
+from zrb.util.attr import get_int_attr
 
 
 def get_history_summarization_token_threshold(
@@ -52,151 +23,3 @@ def get_history_summarization_token_threshold(
             "Defaulting to -1 (no threshold)."
         )
         return -1
-
-
-def should_summarize_history(
-    ctx: AnyContext,
-    history_list: ListOfDict,
-    should_summarize_history_attr: BoolAttr | None,
-    render_summarize_history: bool,
-    history_summarization_token_threshold_attr: IntAttr | None,
-    render_history_summarization_token_threshold: bool,
-) -> bool:
-    """Determines if history summarization should occur based on token length and config."""
-    history_part_count = count_part_in_history_list(history_list)
-    if history_part_count == 0:
-        return False
-    summarization_token_threshold = get_history_summarization_token_threshold(
-        ctx,
-        history_summarization_token_threshold_attr,
-        render_history_summarization_token_threshold,
-    )
-    history_token_count = _count_token_in_history(history_list)
-    if (
-        summarization_token_threshold == -1
-        or summarization_token_threshold > history_token_count
-    ):
-        return False
-    return get_bool_attr(
-        ctx,
-        should_summarize_history_attr,
-        llm_config.default_summarize_history,
-        auto_render=render_summarize_history,
-    )
-
-
-async def summarize_history(
-    ctx: AnyContext,
-    model: "Model | str | None",
-    settings: "ModelSettings | None",
-    system_prompt: str,
-    conversation_history: ConversationHistory,
-    rate_limitter: LLMRateLimiter | None = None,
-    retries: int = 3,
-) -> ConversationHistory:
-    """Runs an LLM call to update the conversation summary."""
-    from pydantic_ai import Agent
-
-    inject_conversation_history_notes(conversation_history)
-    ctx.log_info("Attempting to summarize conversation history...")
-    # Construct the user prompt for the summarization agent
-    user_prompt = "\n".join(
-        [
-            make_markdown_section(
-                "Past Conversation",
-                "\n".join(
-                    [
-                        make_markdown_section(
-                            "Summary",
-                            conversation_history.past_conversation_summary,
-                            as_code=True,
-                        ),
-                        make_markdown_section(
-                            "Last Transcript",
-                            conversation_history.past_conversation_transcript,
-                            as_code=True,
-                        ),
-                    ]
-                ),
-            ),
-            make_markdown_section(
-                "Recent Conversation (JSON)",
-                json.dumps(truncate_str(conversation_history.history, 1000)),
-                as_code=True,
-            ),
-        ]
-    )
-    summarize = create_history_summarization_tool(conversation_history)
-    summarization_agent = Agent[None, str](
-        model=model,
-        output_type=summarize,
-        system_prompt=system_prompt,
-        model_settings=settings,
-        retries=retries,
-    )
-    try:
-        ctx.print(stylize_faint("  📝 Rollup Conversation"), plain=True)
-        summary_run = await run_agent_iteration(
-            ctx=ctx,
-            agent=summarization_agent,
-            user_prompt=user_prompt,
-            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()
-            ctx.print(
-                stylize_faint(f"  📝 Rollup Conversation Token: {usage}"), plain=True
-            )
-            ctx.print(plain=True)
-            ctx.log_info("History summarized and updated.")
-        else:
-            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 the original summary if summarization fails
-    return conversation_history
-
-
-async def maybe_summarize_history(
-    ctx: AnyContext,
-    conversation_history: ConversationHistory,
-    should_summarize_history_attr: BoolAttr | None,
-    render_summarize_history: bool,
-    history_summarization_token_threshold_attr: IntAttr | None,
-    render_history_summarization_token_threshold: bool,
-    model: "str | Model | None",
-    model_settings: "ModelSettings | None",
-    summarization_prompt: str,
-    rate_limitter: LLMRateLimiter | None = None,
-) -> ConversationHistory:
-    """Summarizes history and updates context if enabled and threshold met."""
-    shorten_history = replace_system_prompt_in_history(conversation_history.history)
-    if should_summarize_history(
-        ctx,
-        shorten_history,
-        should_summarize_history_attr,
-        render_summarize_history,
-        history_summarization_token_threshold_attr,
-        render_history_summarization_token_threshold,
-    ):
-        original_history = conversation_history.history
-        conversation_history.history = shorten_history
-        conversation_history = await summarize_history(
-            ctx=ctx,
-            model=model,
-            settings=model_settings,
-            system_prompt=summarization_prompt,
-            conversation_history=conversation_history,
-            rate_limitter=rate_limitter,
-        )
-        conversation_history.history = original_history
-        if (
-            conversation_history.past_conversation_summary != ""
-            and conversation_history.past_conversation_transcript != ""
-        ):
-            conversation_history.history = []
-    return conversation_history

zrb/task/llm/print_node.py

@@ -2,6 +2,7 @@ import json
 from collections.abc import Callable
 from typing import Any
 
+from zrb.config.config import CFG
 from zrb.util.cli.style import stylize_faint
 
 
@@ -104,12 +105,20 @@ async def print_node(
                         and event.tool_call_id
                     ):
                         call_id = event.tool_call_id
-                        result_content = event.result.content
-                        print_func(
-                            _format_content(
-                                f"{call_id} | {result_content}", log_indent_level
+                        if CFG.LLM_SHOW_TOOL_CALL_RESULT:
+                            result_content = event.result.content
+                            print_func(
+                                _format_content(
+                                    f"{call_id} | Return {result_content}",
+                                    log_indent_level,
+                                )
+                            )
+                        else:
+                            print_func(
+                                _format_content(
+                                    f"{call_id} | Executed", log_indent_level
+                                )
                             )
-                        )
         except UnexpectedModelBehavior as e:
             print_func("")  # ensure newline consistency
             print_func(