zrb 1.9.17__py3-none-any.whl → 1.10.1__py3-none-any.whl

zrb/task/llm/conversation_history.py

@@ -0,0 +1,128 @@
+import json
+from collections.abc import Callable
+from copy import deepcopy
+from typing import Any
+
+from zrb.attr.type import StrAttr
+from zrb.context.any_context import AnyContext
+from zrb.context.any_shared_context import AnySharedContext
+from zrb.task.llm.conversation_history_model import ConversationHistory
+from zrb.task.llm.typing import ListOfDict
+from zrb.util.attr import get_str_attr
+from zrb.util.file import write_file
+from zrb.util.run import run_async
+
+
+def get_history_file(
+    ctx: AnyContext,
+    conversation_history_file_attr: StrAttr | None,
+    render_history_file: bool,
+) -> str:
+    """Gets the path to the conversation history file, rendering if configured."""
+    return get_str_attr(
+        ctx,
+        conversation_history_file_attr,
+        "",
+        auto_render=render_history_file,
+    )
+
+
+async def read_conversation_history(
+    ctx: AnyContext,
+    conversation_history_reader: (
+        Callable[[AnySharedContext], ConversationHistory | dict | list | None] | None
+    ),
+    conversation_history_file_attr: StrAttr | None,
+    render_history_file: bool,
+    conversation_history_attr: (
+        ConversationHistory
+        | Callable[[AnySharedContext], ConversationHistory | dict | list]
+        | dict
+        | list
+    ),
+) -> ConversationHistory:
+    """Reads conversation history from reader, file, or attribute, with validation."""
+    history_file = get_history_file(
+        ctx, conversation_history_file_attr, render_history_file
+    )
+    # Use the class method defined above
+    history_data = await ConversationHistory.read_from_source(
+        ctx=ctx,
+        reader=conversation_history_reader,
+        file_path=history_file,
+    )
+    if history_data:
+        return history_data
+    # Priority 3: Callable or direct conversation_history attribute
+    raw_data_attr: Any = None
+    if callable(conversation_history_attr):
+        try:
+            raw_data_attr = await run_async(conversation_history_attr(ctx))
+        except Exception as e:
+            ctx.log_warning(
+                f"Error executing callable conversation_history attribute: {e}. "
+                "Ignoring."
+            )
+    if raw_data_attr is None:
+        raw_data_attr = conversation_history_attr
+    if raw_data_attr:
+        # Use the class method defined above
+        history_data = ConversationHistory.parse_and_validate(
+            ctx, raw_data_attr, "attribute"
+        )
+        if history_data:
+            return history_data
+    # Fallback: Return default value
+    return ConversationHistory()
+
+
+async def write_conversation_history(
+    ctx: AnyContext,
+    history_data: ConversationHistory,
+    conversation_history_writer: (
+        Callable[[AnySharedContext, ConversationHistory], None] | None
+    ),
+    conversation_history_file_attr: StrAttr | None,
+    render_history_file: bool,
+):
+    """Writes conversation history using the writer or to a file."""
+    if conversation_history_writer is not None:
+        await run_async(conversation_history_writer(ctx, history_data))
+    history_file = get_history_file(
+        ctx, conversation_history_file_attr, render_history_file
+    )
+    if history_file != "":
+        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
+    for history in history_list:
+        if "parts" in history:
+            history_part_len += len(history["parts"])
+        else:
+            history_part_len += 1
+    return history_part_len

zrb/task/llm/conversation_history_model.py

@@ -0,0 +1,438 @@
+import json
+import os
+from collections.abc import Callable
+from typing import Any
+
+from zrb.config.config import CFG
+from zrb.context.any_context import AnyContext
+from zrb.task.llm.typing import ListOfDict
+from zrb.util.file import read_file, read_file_with_line_numbers, write_file
+from zrb.util.run import run_async
+
+
+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,
+    ):
+        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()
+
+    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,
+        }
+
+    def model_dump_json(self, indent: int = 2) -> str:
+        return json.dumps(self.to_dict(), indent=indent)
+
+    @classmethod
+    async def read_from_source(
+        cls,
+        ctx: AnyContext,
+        reader: Callable[[AnyContext], dict[str, Any] | list | None] | None,
+        file_path: str | None,
+    ) -> "ConversationHistory | None":
+        # Priority 1: Reader function
+        if reader:
+            try:
+                raw_data = await run_async(reader(ctx))
+                if raw_data:
+                    instance = cls.parse_and_validate(ctx, raw_data, "reader")
+                    if instance:
+                        return instance
+            except Exception as e:
+                ctx.log_warning(
+                    f"Error executing conversation history reader: {e}. Ignoring."
+                )
+        # Priority 2: History file
+        if file_path and os.path.isfile(file_path):
+            try:
+                content = read_file(file_path)
+                raw_data = json.loads(content)
+                instance = cls.parse_and_validate(ctx, raw_data, f"file '{file_path}'")
+                if instance:
+                    return instance
+            except json.JSONDecodeError:
+                ctx.log_warning(
+                    f"Could not decode JSON from history file '{file_path}'. "
+                    "Ignoring file content."
+                )
+            except Exception as e:
+                ctx.log_warning(
+                    f"Error reading history file '{file_path}': {e}. "
+                    "Ignoring file content."
+                )
+        # Fallback: Return default value
+        return None
+
+    def fetch_newest_notes(self):
+        long_term_note_path = self._get_long_term_note_path()
+        if os.path.isfile(long_term_note_path):
+            self.long_term_note = read_file(long_term_note_path)
+        contextual_note_path = self._get_contextual_note_path()
+        if os.path.isfile(contextual_note_path):
+            self.contextual_note = read_file(contextual_note_path)
+
+    @classmethod
+    def parse_and_validate(
+        cls, ctx: AnyContext, data: Any, source: str
+    ) -> "ConversationHistory":
+        try:
+            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", ""),
+                )
+            elif isinstance(data, list):
+                # Handle very old format (just a list) - wrap it
+                ctx.log_warning(
+                    f"History from {source} contains legacy list format. "
+                    "Wrapping it into the new structure. "
+                    "Consider updating the source format."
+                )
+                return cls(history=data)
+            else:
+                ctx.log_warning(
+                    f"History data from {source} has unexpected format "
+                    f"(type: {type(data)}). Ignoring."
+                )
+        except Exception as e:  # Catch validation errors too
+            ctx.log_warning(
+                f"Error validating/parsing history data from {source}: {e}. Ignoring."
+            )
+        return cls()
+
+    def write_past_conversation_summary(self, past_conversation_summary: str):
+        """
+        Write or update the past conversation summary.
+
+        Use this tool to store or update a summary of previous conversations for
+        future reference. This is useful for providing context to LLMs or other tools
+        that need a concise history.
+
+        Args:
+            past_conversation_summary (str): The summary text to store.
+
+        Returns:
+            None
+
+        Raises:
+            Exception: If the summary cannot be written.
+        """
+        self.past_conversation_summary = past_conversation_summary
+
+    def write_past_conversation_transcript(self, past_conversation_transcript: str):
+        """
+        Write or update the past conversation transcript.
+
+        Use this tool to store or update the full transcript of previous conversations.
+        This is useful for providing detailed context to LLMs or for record-keeping.
+
+        Args:
+            past_conversation_transcript (str): The transcript text to store.
+
+        Returns:
+            None
+
+        Raises:
+            Exception: If the transcript cannot be written.
+        """
+        self.past_conversation_transcript = past_conversation_transcript
+
+    def read_long_term_note(
+        self,
+        start_line: int | None = None,
+        end_line: int | None = None,
+    ) -> str:
+        """
+        Read the content of the long-term note, optionally for a specific line range.
+
+        This tool helps you retrieve knowledge or notes stored for long-term reference.
+        If the note does not exist, you may want to create it using the write tool.
+
+        Args:
+            start_line (int, optional): 1-based line number to start reading from.
+            end_line (int, optional): 1-based line number to stop reading at (inclusive).
+
+        Returns:
+            str: JSON with file path, content (with line numbers), start/end lines,
+                and total lines.
+
+        Raises:
+            Exception: If the note cannot be read.
+                Suggests writing the note if it does not exist.
+        """
+        return self._read_note(
+            self._get_long_term_note_path(),
+            start_line,
+            end_line,
+            note_type="long-term note",
+        )
+
+    def write_long_term_note(self, content: str) -> str:
+        """
+        Write or overwrite the content of the long-term note.
+
+        Use this tool to create a new long-term note or replace its entire content.
+        Always read the note first to avoid accidental data loss, unless you are sure
+        you want to overwrite.
+
+        Args:
+            content (str): The full content to write to the note.
+
+        Returns:
+            str: JSON indicating success and the note path.
+
+        Raises:
+            Exception: If the note cannot be written. Suggests checking permissions or path.
+        """
+        self.long_term_note = content
+        return self._write_note(
+            self._get_long_term_note_path(), content, note_type="long-term note"
+        )
+
+    def replace_in_long_term_note(
+        self,
+        old_string: str,
+        new_string: str,
+    ) -> str:
+        """
+        Replace the first occurrence of a string in the long-term note.
+
+        Use this tool to update a specific part of the long-term note without
+        overwriting the entire content. If the note does not exist, consider writing it
+        first. If the string is not found, check your input or read the note to verify.
+
+        Args:
+            old_string (str): The exact string to search for and replace.
+            new_string (str): The string to replace with.
+
+        Returns:
+            str: JSON indicating success and the note path.
+
+        Raises:
+            Exception: If the note does not exist or the string is not found.
+                Suggests writing or reading the note.
+        """
+        result = self._replace_in_note(
+            self._get_long_term_note_path(),
+            old_string,
+            new_string,
+            note_type="long-term note",
+        )
+        self.long_term_note = new_string
+        return result
+
+    def read_contextual_note(
+        self,
+        start_line: int | None = None,
+        end_line: int | None = None,
+    ) -> str:
+        """
+        Read the content of the contextual note, optionally for a specific line range.
+
+        This tool helps you retrieve project-specific or session-specific notes.
+        If the note does not exist, you may want to create it using the write tool.
+
+        Args:
+            start_line (int, optional): 1-based line number to start reading from.
+            end_line (int, optional): 1-based line number to stop reading at (inclusive).
+
+        Returns:
+            str: JSON with file path, content (with line numbers), start/end lines,
+                and total lines.
+
+        Raises:
+            Exception: If the note cannot be read.
+                Suggests writing the note if it does not exist.
+        """
+        return self._read_note(
+            self._get_contextual_note_path(),
+            start_line,
+            end_line,
+            note_type="contextual note",
+        )
+
+    def write_contextual_note(self, content: str) -> str:
+        """
+        Write or overwrite the content of the contextual note.
+
+        Use this tool to create a new contextual note or replace its entire content.
+        Always read the note first to avoid accidental data loss, unless you are sure
+        you want to overwrite.
+
+        Args:
+            content (str): The full content to write to the note.
+
+        Returns:
+            str: JSON indicating success and the note path.
+
+        Raises:
+            Exception: If the note cannot be written. Suggests checking permissions or path.
+        """
+        self.contextual_note = content
+        return self._write_note(
+            self._get_contextual_note_path(), content, note_type="contextual note"
+        )
+
+    def replace_in_contextual_note(
+        self,
+        old_string: str,
+        new_string: str,
+    ) -> str:
+        """
+        Replace the first occurrence of a string in the contextual note.
+
+        Use this tool to update a specific part of the contextual note without
+        overwriting the entire content. If the note does not exist, consider writing it
+        first. If the string is not found, check your input or read the note to verify.
+
+        Args:
+            old_string (str): The exact string to search for and replace.
+            new_string (str): The string to replace with.
+
+        Returns:
+            str: JSON indicating success and the note path.
+
+        Raises:
+            Exception: If the note does not exist or the string is not found.
+                Suggests writing or reading the note.
+        """
+        result = self._replace_in_note(
+            self._get_contextual_note_path(),
+            old_string,
+            new_string,
+            note_type="contextual note",
+        )
+        self.contextual_note = new_string
+        return result
+
+    def _get_long_term_note_path(self) -> str:
+        return os.path.abspath(os.path.expanduser(CFG.LLM_LONG_TERM_NOTE_PATH))
+
+    def _get_contextual_note_path(self) -> str:
+        return os.path.join(self.project_path, CFG.LLM_CONTEXTUAL_NOTE_FILE)
+
+    def _read_note(
+        self,
+        path: str,
+        start_line: int | None = None,
+        end_line: int | None = None,
+        note_type: str = "note",
+    ) -> str:
+        """
+        Internal helper to read a note file with line numbers and error handling.
+        """
+        if not os.path.exists(path):
+            return json.dumps(
+                {
+                    "path": path,
+                    "content": "",
+                    "start_line": 0,
+                    "end_line": 0,
+                    "total_lines": 0,
+                }
+            )
+        try:
+            content = read_file_with_line_numbers(path)
+            lines = content.splitlines()
+            total_lines = len(lines)
+            start_idx = (start_line - 1) if start_line is not None else 0
+            end_idx = end_line if end_line is not None else total_lines
+            if start_idx < 0:
+                start_idx = 0
+            if end_idx > total_lines:
+                end_idx = total_lines
+            if start_idx > end_idx:
+                start_idx = end_idx
+            selected_lines = lines[start_idx:end_idx]
+            content_result = "\n".join(selected_lines)
+            return json.dumps(
+                {
+                    "path": path,
+                    "content": content_result,
+                    "start_line": start_idx + 1,
+                    "end_line": end_idx,
+                    "total_lines": total_lines,
+                }
+            )
+        except Exception:
+            raise Exception(
+                f"Failed to read the {note_type}. "
+                f"If the {note_type} does not exist, try writing it first."
+            )
+
+    def _write_note(self, path: str, content: str, note_type: str = "note") -> str:
+        """
+        Internal helper to write a note file with error handling.
+        """
+        try:
+            directory = os.path.dirname(path)
+            if directory and not os.path.exists(directory):
+                os.makedirs(directory, exist_ok=True)
+            write_file(path, content)
+            return json.dumps({"success": True, "path": path})
+        except (OSError, IOError):
+            raise Exception(
+                f"Failed to write the {note_type}. "
+                "Please check if the path is correct and you have write permissions."
+            )
+        except Exception:
+            raise Exception(
+                f"Unexpected error while writing the {note_type}. "
+                "Please check your input and try again."
+            )
+
+    def _replace_in_note(
+        self, path: str, old_string: str, new_string: str, note_type: str = "note"
+    ) -> str:
+        """
+        Internal helper to replace a string in a note file with error handling.
+        """
+        if not os.path.exists(path):
+            raise Exception(
+                (
+                    f"{note_type.capitalize()} not found. "
+                    f"Consider writing a new {note_type} first."
+                )
+            )
+        try:
+            content = read_file(path)
+            if old_string not in content:
+                raise Exception(
+                    f"The specified string to replace was not found in the {note_type}. ("
+                    f"Try reading the {note_type} to verify its content or "
+                    f"write a new one if needed)."
+                )
+            new_content = content.replace(old_string, new_string, 1)
+            write_file(path, new_content)
+            return json.dumps({"success": True, "path": path})
+        except Exception:
+            raise Exception(
+                f"Failed to replace content in the {note_type}. ("
+                f"Try reading the {note_type} to verify its content or "
+                "write a new one if needed)."
+            )