zrb 1.13.1__py3-none-any.whl → 1.21.17__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_summarization.py

@@ -9,13 +9,18 @@ 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.llm.prompt import make_prompt_section
+from zrb.util.markdown import make_markdown_section
+from zrb.util.truncate import truncate_str
 
 if TYPE_CHECKING:
     from pydantic_ai.models import Model
@@ -88,24 +93,25 @@ async def summarize_history(
     conversation_history: ConversationHistory,
     rate_limitter: LLMRateLimiter | None = None,
     retries: int = 3,
-) -> str:
+) -> 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_prompt_section(
+            make_markdown_section(
                 "Past Conversation",
                 "\n".join(
                     [
-                        make_prompt_section(
+                        make_markdown_section(
                             "Summary",
                             conversation_history.past_conversation_summary,
                             as_code=True,
                         ),
-                        make_prompt_section(
+                        make_markdown_section(
                             "Last Transcript",
                             conversation_history.past_conversation_transcript,
                             as_code=True,
@@ -113,56 +119,37 @@ async def summarize_history(
                     ]
                 ),
             ),
-            make_prompt_section(
+            make_markdown_section(
                 "Recent Conversation (JSON)",
-                json.dumps(conversation_history.history),
+                json.dumps(truncate_str(conversation_history.history, 1000)),
                 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,
-                        ),
-                    ]
-                ),
-            ),
         ]
     )
-    summarization_agent = Agent(
+    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,
-        tools=[
-            conversation_history.write_past_conversation_summary,
-            conversation_history.write_past_conversation_transcript,
-            conversation_history.read_long_term_note,
-            conversation_history.write_long_term_note,
-            conversation_history.read_contextual_note,
-            conversation_history.write_contextual_note,
-        ],
     )
     try:
-        ctx.print(stylize_faint("📝 Summarize Conversation >>>"), plain=True)
+        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"📝 Summarization Token: {usage}"), plain=True)
+            ctx.print(
+                stylize_faint(f"  📝 Rollup Conversation Token: {usage}"), plain=True
+            )
             ctx.print(plain=True)
             ctx.log_info("History summarized and updated.")
         else:

zrb/task/llm/history_summarization_tool.py

@@ -0,0 +1,24 @@
+from typing import Callable
+
+from zrb.task.llm.conversation_history_model import ConversationHistory
+
+
+def create_history_summarization_tool(
+    conversation_history: ConversationHistory,
+) -> Callable[[str, str], str]:
+    def update_conversation_memory(
+        past_conversation_summary: str,
+        past_conversation_transcript: str,
+    ) -> str:
+        """
+        Update the conversation memory including summary and transcript.
+        - past_conversation_summary: A concise narrative that integrates the
+          previous summary with the recent conversation.
+        - past_conversation_transcript: MUST be ONLY the last 4 (four) turns
+          of the conversation.
+        """
+        conversation_history.past_conversation_summary = past_conversation_summary
+        conversation_history.past_conversation_transcript = past_conversation_transcript
+        return "Conversation memory updated"
+
+    return update_conversation_memory

zrb/task/llm/print_node.py

@@ -1,12 +1,16 @@
+import json
 from collections.abc import Callable
 from typing import Any
 
 from zrb.util.cli.style import stylize_faint
 
 
-async def print_node(print_func: Callable, agent_run: Any, node: Any):
+async def print_node(
+    print_func: Callable, agent_run: Any, node: Any, log_indent_level: int = 0
+):
     """Prints the details of an agent execution node using a provided print function."""
     from pydantic_ai import Agent
+    from pydantic_ai.exceptions import UnexpectedModelBehavior
     from pydantic_ai.messages import (
         FinalResultEvent,
         FunctionToolCallEvent,
@@ -18,76 +22,191 @@ async def print_node(print_func: Callable, agent_run: Any, node: Any):
         ToolCallPartDelta,
     )
 
+    meta = getattr(node, "id", None) or getattr(node, "request_id", None)
     if Agent.is_user_prompt_node(node):
-        print_func(stylize_faint("🔠 Receiving input..."))
+        print_func(_format_header("🔠 Receiving input...", log_indent_level))
     elif Agent.is_model_request_node(node):
         # A model request node => We can stream tokens from the model's request
-        print_func(stylize_faint("🧠 Processing..."))
-        async with node.stream(agent_run.ctx) as request_stream:
-            is_streaming = False
-            async for event in request_stream:
-                if isinstance(event, PartStartEvent) and event.part:
-                    if is_streaming:
-                        print_func("")
-                    print_func(
-                        stylize_faint(f"    Starting part {event.index}: {event.part}"),
-                    )
-                    is_streaming = False
-                elif isinstance(event, PartDeltaEvent):
-                    if isinstance(event.delta, TextPartDelta) or isinstance(
-                        event.delta, ThinkingPartDelta
-                    ):
-                        print_func(
-                            stylize_faint(f"{event.delta.content_delta}"),
-                            end="",
-                        )
-                    elif isinstance(event.delta, ToolCallPartDelta):
+        print_func(_format_header("🧠 Processing...", log_indent_level))
+        # Reference: https://ai.pydantic.dev/agents/#streaming-all-events-and-output
+        try:
+            async with node.stream(agent_run.ctx) as request_stream:
+                is_streaming = False
+                async for event in request_stream:
+                    if isinstance(event, PartStartEvent) and event.part:
+                        if is_streaming:
+                            print_func("")
+                        content = _get_event_part_content(event)
+                        print_func(_format_content(content, log_indent_level), end="")
+                        is_streaming = True
+                    elif isinstance(event, PartDeltaEvent):
+                        if isinstance(event.delta, TextPartDelta):
+                            content_delta = event.delta.content_delta
+                            print_func(
+                                _format_stream_content(content_delta, log_indent_level),
+                                end="",
+                            )
+                        elif isinstance(event.delta, ThinkingPartDelta):
+                            content_delta = event.delta.content_delta
+                            print_func(
+                                _format_stream_content(content_delta, log_indent_level),
+                                end="",
+                            )
+                        elif isinstance(event.delta, ToolCallPartDelta):
+                            args_delta = event.delta.args_delta
+                            if isinstance(args_delta, dict):
+                                args_delta = json.dumps(args_delta)
+                            print_func(
+                                _format_stream_content(args_delta, log_indent_level),
+                                end="",
+                            )
+                        is_streaming = True
+                    elif isinstance(event, FinalResultEvent) and event.tool_name:
+                        if is_streaming:
+                            print_func("")
+                        tool_name = event.tool_name
                         print_func(
-                            stylize_faint(f"{event.delta.args_delta}"),
-                            end="",
+                            _format_content(
+                                f"Result: tool_name={tool_name}", log_indent_level
+                            )
                         )
-                    is_streaming = True
-                elif isinstance(event, FinalResultEvent) and event.tool_name:
-                    if is_streaming:
-                        print_func("")
-                    print_func(
-                        stylize_faint(f"    Result: tool_name={event.tool_name}"),
-                    )
-                    is_streaming = False
-            if is_streaming:
-                print_func("")
+                        is_streaming = False
+                if is_streaming:
+                    print_func("")
+        except UnexpectedModelBehavior as e:
+            print_func("")  # ensure newline consistency
+            print_func(
+                _format_content(
+                    (
+                        f"🟡 Unexpected Model Behavior: {e}. "
+                        f"Cause: {e.__cause__}. Node.Id: {meta}"
+                    ),
+                    log_indent_level,
+                )
+            )
     elif Agent.is_call_tools_node(node):
         # A handle-response node => The model returned some data, potentially calls a tool
-        print_func(stylize_faint("🧰 Calling Tool..."))
-        async with node.stream(agent_run.ctx) as handle_stream:
-            async for event in handle_stream:
-                if isinstance(event, FunctionToolCallEvent):
-                    # Handle empty arguments across different providers
-                    if event.part.args == "" or event.part.args is None:
-                        event.part.args = {}
-                    elif isinstance(
-                        event.part.args, str
-                    ) and event.part.args.strip() in ["null", "{}"]:
-                        # Some providers might send "null" or "{}" as a string
-                        event.part.args = {}
-                    # Handle dummy property if present (from our schema sanitization)
-                    if (
-                        isinstance(event.part.args, dict)
-                        and "_dummy" in event.part.args
-                    ):
-                        del event.part.args["_dummy"]
-                    print_func(
-                        stylize_faint(
-                            f"    {event.part.tool_call_id} | "
-                            f"Call {event.part.tool_name} {event.part.args}"
+        print_func(_format_header("🧰 Calling Tool...", log_indent_level))
+        try:
+            async with node.stream(agent_run.ctx) as handle_stream:
+                async for event in handle_stream:
+                    if isinstance(event, FunctionToolCallEvent):
+                        args = _get_event_part_args(event)
+                        call_id = event.part.tool_call_id
+                        tool_name = event.part.tool_name
+                        print_func(
+                            _format_content(
+                                f"{call_id} | Call {tool_name} {args}", log_indent_level
+                            )
                         )
-                    )
-                elif isinstance(event, FunctionToolResultEvent):
-                    print_func(
-                        stylize_faint(
-                            f"    {event.tool_call_id} | {event.result.content}"
+                    elif (
+                        isinstance(event, FunctionToolResultEvent)
+                        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
+                            )
                         )
-                    )
+        except UnexpectedModelBehavior as e:
+            print_func("")  # ensure newline consistency
+            print_func(
+                _format_content(
+                    (
+                        f"🟡 Unexpected Model Behavior: {e}. "
+                        f"Cause: {e.__cause__}. Node.Id: {meta}"
+                    ),
+                    log_indent_level,
+                )
+            )
     elif Agent.is_end_node(node):
         # Once an End node is reached, the agent run is complete
-        print_func(stylize_faint("✅ Completed..."))
+        print_func(_format_header("✅ Completed...", log_indent_level))
+
+
+def _format_header(text: str | None, log_indent_level: int = 0) -> str:
+    return _format(
+        text,
+        base_indent=2,
+        first_indent=0,
+        indent=0,
+        log_indent_level=log_indent_level,
+    )
+
+
+def _format_content(text: str | None, log_indent_level: int = 0) -> str:
+    return _format(
+        text,
+        base_indent=2,
+        first_indent=3,
+        indent=3,
+        log_indent_level=log_indent_level,
+    )
+
+
+def _format_stream_content(text: str | None, log_indent_level: int = 0) -> str:
+    return _format(
+        text,
+        base_indent=2,
+        indent=3,
+        log_indent_level=log_indent_level,
+        is_stream=True,
+    )
+
+
+def _format(
+    text: str | None,
+    base_indent: int = 0,
+    first_indent: int = 0,
+    indent: int = 0,
+    log_indent_level: int = 0,
+    is_stream: bool = False,
+) -> str:
+    if text is None:
+        text = ""
+    line_prefix = (base_indent * (log_indent_level + 1) + indent) * " "
+    processed_text = text.replace("\n", f"\n{line_prefix}")
+    if is_stream:
+        return stylize_faint(processed_text)
+    first_line_prefix = (base_indent * (log_indent_level + 1) + first_indent) * " "
+    return stylize_faint(f"{first_line_prefix}{processed_text}")
+
+
+def _get_event_part_args(event: Any) -> Any:
+    # Handle empty arguments across different providers
+    if event.part.args == "" or event.part.args is None:
+        return {}
+    if isinstance(event.part.args, str):
+        # Some providers might send "null" or "{}" as a string
+        if event.part.args.strip() in ["null", "{}"]:
+            return {}
+        try:
+            obj = json.loads(event.part.args)
+            if isinstance(obj, dict):
+                return _truncate_kwargs(obj)
+        except json.JSONDecodeError:
+            pass
+    # Handle dummy property if present (from our schema sanitization)
+    if isinstance(event.part.args, dict):
+        return _truncate_kwargs(event.part.args)
+    return event.part.args
+
+
+def _truncate_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
+    return {key: _truncate_arg(val) for key, val in kwargs.items()}
+
+
+def _truncate_arg(arg: str, length: int = 19) -> str:
+    if isinstance(arg, str) and len(arg) > length:
+        return f"{arg[:length-4]} ..."
+    return arg
+
+
+def _get_event_part_content(event: Any) -> str:
+    if not hasattr(event, "part"):
+        return f"{event}"
+    if not hasattr(event.part, "content"):
+        return f"{event.part}"
+    return getattr(event.part, "content")