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

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)."
+            )

zrb/task/llm/history_summarization.py

@@ -7,13 +7,15 @@ 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.history import (
+from zrb.task.llm.conversation_history import (
     count_part_in_history_list,
-    replace_system_prompt_in_history_list,
+    replace_system_prompt_in_history,
 )
+from zrb.task.llm.conversation_history_model import ConversationHistory
 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
 
 if TYPE_CHECKING:
     from pydantic_ai.models import Model
@@ -82,9 +84,8 @@ async def summarize_history(
     ctx: AnyContext,
     model: "Model | str | None",
     settings: "ModelSettings | None",
-    prompt: str,
-    previous_summary: str,
-    history_list: ListOfDict,
+    system_prompt: str,
+    conversation_history: ConversationHistory,
     rate_limitter: LLMRateLimiter | None = None,
     retries: int = 3,
 ) -> str:
@@ -93,16 +94,65 @@ async def summarize_history(
 
     ctx.log_info("Attempting to summarize conversation history...")
     # Construct the user prompt for the summarization agent
-    user_prompt = json.dumps(
-        {"previous_summary": previous_summary, "recent_history": history_list}
+    user_prompt = "\n".join(
+        [
+            make_prompt_section(
+                "Past Conversation",
+                "\n".join(
+                    [
+                        make_prompt_section(
+                            "Summary",
+                            conversation_history.past_conversation_summary,
+                            as_code=True,
+                        ),
+                        make_prompt_section(
+                            "Last Transcript",
+                            conversation_history.past_conversation_transcript,
+                            as_code=True,
+                        ),
+                    ]
+                ),
+            ),
+            make_prompt_section(
+                "Recent Conversation (JSON)",
+                json.dumps(conversation_history.history),
+                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(
         model=model,
-        system_prompt=prompt,
+        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_contextual_note,
+            conversation_history.write_contextual_note,
+            conversation_history.replace_in_contextual_note,
+            conversation_history.read_long_term_note,
+            conversation_history.write_long_term_note,
+            conversation_history.replace_in_long_term_note,
+        ],
     )
-
     try:
         ctx.print(stylize_faint("📝 Summarize"), plain=True)
         summary_run = await run_agent_iteration(
@@ -113,27 +163,22 @@ async def summarize_history(
             rate_limitter=rate_limitter,
         )
         if summary_run and summary_run.result and summary_run.result.output:
-            new_summary = str(summary_run.result.output)
             usage = summary_run.result.usage()
             ctx.print(stylize_faint(f"📝 Summarization Token: {usage}"), plain=True)
             ctx.print(plain=True)
             ctx.log_info("History summarized and updated.")
-            ctx.log_info(f"New conversation summary:\n{new_summary}")
-            return new_summary
         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 previous_summary
+    return conversation_history
 
 
 async def maybe_summarize_history(
     ctx: AnyContext,
-    history_list: ListOfDict,
-    conversation_summary: str,
+    conversation_history: ConversationHistory,
     should_summarize_history_attr: BoolAttr | None,
     render_summarize_history: bool,
     history_summarization_token_threshold_attr: IntAttr | None,
@@ -142,26 +187,31 @@ async def maybe_summarize_history(
     model_settings: "ModelSettings | None",
     summarization_prompt: str,
     rate_limitter: LLMRateLimiter | None = None,
-) -> tuple[ListOfDict, str]:
+) -> ConversationHistory:
     """Summarizes history and updates context if enabled and threshold met."""
-    shorten_history_list = replace_system_prompt_in_history_list(history_list)
+    shorten_history = replace_system_prompt_in_history(conversation_history.history)
     if should_summarize_history(
         ctx,
-        shorten_history_list,
+        shorten_history,
         should_summarize_history_attr,
         render_summarize_history,
         history_summarization_token_threshold_attr,
         render_history_summarization_token_threshold,
     ):
-        new_summary = await summarize_history(
+        original_history = conversation_history.history
+        conversation_history.history = shorten_history
+        conversation_history = await summarize_history(
             ctx=ctx,
             model=model,
             settings=model_settings,
-            prompt=summarization_prompt,
-            previous_summary=conversation_summary,
-            history_list=shorten_history_list,
+            system_prompt=summarization_prompt,
+            conversation_history=conversation_history,
             rate_limitter=rate_limitter,
         )
-        # After summarization, the history is cleared and replaced by the new summary
-        return [], new_summary
-    return history_list, conversation_summary
+        conversation_history.history = original_history
+        if (
+            conversation_history.past_conversation_summary != ""
+            and conversation_history.past_conversation_transcript != ""
+        ):
+            conversation_history.history = []
+    return conversation_history