zrb 1.5.17__py3-none-any.whl → 1.6.1__py3-none-any.whl

zrb/task/llm/agent.py

@@ -2,14 +2,12 @@ from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
-    from openai import APIError
     from pydantic_ai import Agent, Tool
     from pydantic_ai.agent import AgentRun
     from pydantic_ai.mcp import MCPServer
     from pydantic_ai.models import Model
     from pydantic_ai.settings import ModelSettings
 else:
-    APIError = Any
     Agent = Any
     Tool = Any
     AgentRun = Any
@@ -130,6 +128,7 @@ async def run_agent_iteration(
     Raises:
         Exception: If any error occurs during agent execution.
     """
+    from openai import APIError
     from pydantic_ai.messages import ModelMessagesTypeAdapter
 
     async with agent.run_mcp_servers():
@@ -141,7 +140,7 @@ async def run_agent_iteration(
                 # Each node represents a step in the agent's execution
                 # Reference: https://ai.pydantic.dev/agents/#streaming
                 try:
-                    await print_node(ctx.print, agent_run, node)
+                    await print_node(_get_plain_printer(ctx), agent_run, node)
                 except APIError as e:
                     # Extract detailed error information from the response
                     error_details = extract_api_error_details(e)
@@ -152,3 +151,12 @@ async def run_agent_iteration(
                     ctx.log_error(f"Error type: {type(e).__name__}")
                     raise
             return agent_run
+
+
+def _get_plain_printer(ctx: AnyContext):
+    def printer(*args, **kwargs):
+        if "plain" not in kwargs:
+            kwargs["plain"] = True
+        return ctx.print(*args, **kwargs)
+
+    return printer

zrb/task/llm/context.py

@@ -73,9 +73,7 @@ def get_conversation_context(
     Retrieves the conversation context.
     If a value in the context dict is callable, it executes it with ctx.
     """
-    raw_context = get_attr(
-        ctx, conversation_context_attr, {}, auto_render=False
-    )  # Context usually shouldn't be rendered
+    raw_context = get_attr(ctx, conversation_context_attr, {}, auto_render=False)
     if not isinstance(raw_context, dict):
         ctx.log_warning(
             f"Conversation context resolved to type {type(raw_context)}, "

zrb/task/llm/context_enrichment.py

@@ -1,6 +1,5 @@
 import json
 import traceback
-from textwrap import dedent
 from typing import TYPE_CHECKING, Any
 
 from pydantic import BaseModel
@@ -9,8 +8,13 @@ from zrb.attr.type import BoolAttr, IntAttr
 from zrb.context.any_context import AnyContext
 from zrb.llm_config import llm_config
 from zrb.task.llm.agent import run_agent_iteration
+from zrb.task.llm.history import (
+    count_part_in_history_list,
+    replace_system_prompt_in_history_list,
+)
 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
 
 if TYPE_CHECKING:
     from pydantic_ai.models import Model
@@ -47,15 +51,13 @@ async def enrich_context(
         context_json = json.dumps(conversation_context)
         history_json = json.dumps(history_list)
         # The user prompt will now contain the dynamic data
-        user_prompt_data = dedent(
-            f"""
-            Analyze the following
-            # Current Context
-            {context_json}
-            # Conversation History
-            {history_json}
-        """
-        ).strip()
+        user_prompt_data = "\n".join(
+            [
+                "Extract context from the following conversation info",
+                f"Existing Context: {context_json}",
+                f"Conversation History: {history_json}",
+            ]
+        )
     except Exception as e:
         ctx.log_warning(f"Error formatting context/history for enrichment: {e}")
         return conversation_context  # Return original context if formatting fails
@@ -72,6 +74,7 @@ async def enrich_context(
     )
 
     try:
+        ctx.print(stylize_faint("[Context Enrichment Triggered]"), plain=True)
         enrichment_run = await run_agent_iteration(
             ctx=ctx,
             agent=enrichment_agent,
@@ -80,6 +83,8 @@ async def enrich_context(
         )
         if enrichment_run and enrichment_run.result.output:
             response = enrichment_run.result.output.response
+            usage = enrichment_run.result.usage()
+            ctx.print(stylize_faint(f"[Token Usage] {usage}"), plain=True)
             if response:
                 conversation_context.update(response)
                 ctx.log_info("Context enriched based on history.")
@@ -127,15 +132,15 @@ def should_enrich_context(
     """
     Determines if context enrichment should occur based on history, threshold, and config.
     """
-    history_len = len(history_list)
-    if history_len == 0:
+    history_part_count = count_part_in_history_list(history_list)
+    if history_part_count == 0:
         return False
     enrichment_threshold = get_context_enrichment_threshold(
         ctx,
         context_enrichment_threshold_attr,
         render_context_enrichment_threshold,
     )
-    if enrichment_threshold == -1 or enrichment_threshold > history_len:
+    if enrichment_threshold == -1 or enrichment_threshold > history_part_count:
         return False
     return get_bool_attr(
         ctx,
@@ -158,9 +163,10 @@ async def maybe_enrich_context(
     context_enrichment_prompt: str,
 ) -> dict[str, Any]:
     """Enriches context based on history if enabled and threshold met."""
+    shorten_history_list = replace_system_prompt_in_history_list(history_list)
     if should_enrich_context(
         ctx,
-        history_list,
+        shorten_history_list,
         should_enrich_context_attr,
         render_enrich_context,
         context_enrichment_threshold_attr,
@@ -174,6 +180,6 @@ async def maybe_enrich_context(
                 prompt=context_enrichment_prompt,
             ),
             conversation_context=conversation_context,
-            history_list=history_list,
+            history_list=shorten_history_list,
         )
     return conversation_context

zrb/task/llm/history.py

@@ -1,6 +1,7 @@
 import json
 import os
 from collections.abc import Callable
+from copy import deepcopy
 from typing import Any, Optional
 
 from pydantic import BaseModel
@@ -177,66 +178,34 @@ async def write_conversation_history(
         write_file(history_file, history_data.model_dump_json(indent=2))
 
 
-async def prepare_initial_state(
-    ctx: AnyContext,
-    conversation_history_reader: (
-        Callable[[AnySharedContext], ConversationHistoryData | dict | list | None]
-        | None
-    ),
-    conversation_history_file_attr: StrAttr | None,
-    render_history_file: bool,
-    conversation_history_attr: (
-        ConversationHistoryData
-        | Callable[[AnySharedContext], ConversationHistoryData | dict | list]
-        | dict
-        | list
-    ),
-    conversation_context_getter: Callable[[AnyContext], dict[str, Any]],
-) -> tuple[ListOfDict, dict[str, Any]]:
-    """Reads history and prepares the initial conversation context."""
-    history_data: ConversationHistoryData = await read_conversation_history(
-        ctx,
-        conversation_history_reader,
-        conversation_history_file_attr,
-        render_history_file,
-        conversation_history_attr,
-    )
-    # Clean the history list to remove context from historical user prompts
-    cleaned_history_list = []
-    for interaction in history_data.history:
-        cleaned_history_list.append(
-            remove_context_from_interaction_history(interaction)
-        )
-    conversation_context = conversation_context_getter(ctx)
-    # Merge history context from loaded data without overwriting existing keys
-    for key, value in history_data.context.items():
-        if key not in conversation_context:
-            conversation_context[key] = value
-    # Return the CLEANED history list
-    return cleaned_history_list, conversation_context
-
-
-def remove_context_from_interaction_history(
-    interaction: dict[str, Any],
-) -> dict[str, Any]:
-    try:
-        cleaned_interaction = json.loads(json.dumps(interaction))
-    except Exception:
-        # Fallback to shallow copy if not JSON serializable (less safe)
-        cleaned_interaction = interaction.copy()
-    if "parts" in cleaned_interaction and isinstance(
-        cleaned_interaction["parts"], list
-    ):
-        for part in cleaned_interaction["parts"]:
-            is_user_prompt = part.get("part_kind") == "user-prompt"
-            has_str_content = isinstance(part.get("content"), str)
-            if is_user_prompt and has_str_content:
-                content = part["content"]
-                user_message_marker = "# User Message\n"
-                marker_index = content.find(user_message_marker)
-                if marker_index != -1:
-                    # Extract message after the marker and strip whitespace
-                    start_index = marker_index + len(user_message_marker)
-                    part["content"] = content[start_index:].strip()
-                # else: If marker not found, leave content as is (old format/error)
-    return cleaned_interaction
+def replace_system_prompt_in_history_list(
+    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/history_summarization.py

@@ -7,8 +7,13 @@ from zrb.attr.type import BoolAttr, IntAttr
 from zrb.context.any_context import AnyContext
 from zrb.llm_config import llm_config
 from zrb.task.llm.agent import run_agent_iteration
+from zrb.task.llm.history import (
+    count_part_in_history_list,
+    replace_system_prompt_in_history_list,
+)
 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
 
 if TYPE_CHECKING:
     from pydantic_ai.models import Model
@@ -18,17 +23,6 @@ else:
     ModelSettings = Any
 
 
-def get_history_part_len(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
-
-
 def get_history_summarization_threshold(
     ctx: AnyContext,
     history_summarization_threshold_attr: IntAttr | None,
@@ -60,15 +54,15 @@ def should_summarize_history(
     render_history_summarization_threshold: bool,
 ) -> bool:
     """Determines if history summarization should occur based on length and config."""
-    history_part_len = get_history_part_len(history_list)
-    if history_part_len == 0:
+    history_part_count = count_part_in_history_list(history_list)
+    if history_part_count == 0:
         return False
     summarization_threshold = get_history_summarization_threshold(
         ctx,
         history_summarization_threshold_attr,
         render_history_summarization_threshold,
     )
-    if summarization_threshold == -1 or summarization_threshold > history_part_len:
+    if summarization_threshold == -1 or summarization_threshold > history_part_count:
         return False
     return get_bool_attr(
         ctx,
@@ -111,15 +105,18 @@ async def summarize_history(
     try:
         context_json = json.dumps(conversation_context)
         history_to_summarize_json = json.dumps(history_list)
-        summarization_user_prompt = (
-            f"# Current Context\n{context_json}\n\n"
-            f"# Conversation History to Summarize\n{history_to_summarize_json}"
+        summarization_user_prompt = "\n".join(
+            [
+                f"Current Context: {context_json}",
+                f"Conversation History to Summarize: {history_to_summarize_json}",
+            ]
         )
     except Exception as e:
         ctx.log_warning(f"Error formatting context/history for summarization: {e}")
         return conversation_context  # Return original context if formatting fails
 
     try:
+        ctx.print(stylize_faint("[Summarization Triggered]"), plain=True)
         summary_run = await run_agent_iteration(
             ctx=ctx,
             agent=summarization_agent,
@@ -128,6 +125,8 @@ async def summarize_history(
         )
         if summary_run and summary_run.result.output:
             summary_text = str(summary_run.result.output)
+            usage = summary_run.result.usage()
+            ctx.print(stylize_faint(f"[Token Usage] {usage}"), plain=True)
             # Update context with the new summary
             conversation_context["history_summary"] = summary_text
             ctx.log_info("History summarized and added/updated in context.")
@@ -153,9 +152,10 @@ async def maybe_summarize_history(
     summarization_prompt: str,
 ) -> tuple[ListOfDict, dict[str, Any]]:
     """Summarizes history and updates context if enabled and threshold met."""
+    shorten_history_list = replace_system_prompt_in_history_list(history_list)
     if should_summarize_history(
         ctx,
-        history_list,
+        shorten_history_list,
         should_summarize_history_attr,
         render_summarize_history,
         history_summarization_threshold_attr,
@@ -170,7 +170,7 @@ async def maybe_summarize_history(
                 prompt=summarization_prompt,
             ),
             conversation_context=conversation_context,
-            history_list=history_list,  # Pass the full list for context
+            history_list=shorten_history_list,  # Pass the full list for context
         )
         # Truncate the history list after summarization
         return [], updated_context

zrb/task/llm/print_node.py

@@ -30,7 +30,7 @@ async def print_node(print_func: Callable, agent_run: Any, node: Any):
             async for event in request_stream:
                 if isinstance(event, PartStartEvent):
                     if is_streaming:
-                        print_func("", plain=True)
+                        print_func("")
                     print_func(
                         stylize_faint(
                             f"[Request] Starting part {event.index}: {event.part!r}"
@@ -42,24 +42,22 @@ async def print_node(print_func: Callable, agent_run: Any, node: Any):
                         print_func(
                             stylize_faint(f"{event.delta.content_delta}"),
                             end="",
-                            plain=is_streaming,
                         )
                     elif isinstance(event.delta, ToolCallPartDelta):
                         print_func(
                             stylize_faint(f"{event.delta.args_delta}"),
                             end="",
-                            plain=is_streaming,
                         )
                     is_streaming = True
                 elif isinstance(event, FinalResultEvent):
                     if is_streaming:
-                        print_func("", plain=True)
+                        print_func("")
                     print_func(
                         stylize_faint(f"[Result] tool_name={event.tool_name}"),
                     )
                     is_streaming = False
             if is_streaming:
-                print_func("", plain=True)
+                print_func("")
     elif Agent.is_call_tools_node(node):
         # A handle-response node => The model returned some data, potentially calls a tool
         print_func(
@@ -95,4 +93,5 @@ async def print_node(print_func: Callable, agent_run: Any, node: Any):
                     )
     elif Agent.is_end_node(node):
         # Once an End node is reached, the agent run is complete
-        print_func(stylize_faint(f"{agent_run.result.data}"))
+        print_func(stylize_faint("[End of Response]"))
+        # print_func(stylize_faint(f"{agent_run.result.data}"))

zrb/task/llm/prompt.py

@@ -1,11 +1,6 @@
-import json
-from textwrap import dedent
-from typing import Any
-
 from zrb.attr.type import StrAttr
 from zrb.context.any_context import AnyContext
 from zrb.llm_config import llm_config as default_llm_config
-from zrb.task.llm.context import extract_default_context
 from zrb.util.attr import get_attr, get_str_attr
 
 
@@ -127,25 +122,3 @@ def get_context_enrichment_prompt(
     if context_enrichment_prompt is not None:
         return context_enrichment_prompt
     return default_llm_config.get_default_context_enrichment_prompt()
-
-
-def build_user_prompt(
-    ctx: AnyContext,
-    message_attr: StrAttr | None,
-    conversation_context: dict[str, Any],
-) -> str:
-    """Constructs the final user prompt including context."""
-    original_user_message = get_user_message(ctx, message_attr)
-    # Combine default context, conversation context (potentially enriched/summarized)
-    modified_user_message, default_context = extract_default_context(
-        original_user_message
-    )
-    enriched_context = {**default_context, **conversation_context}
-    return dedent(
-        f"""
-        # Context
-        {json.dumps(enriched_context)}
-        # User Message
-        {modified_user_message}
-        """
-    ).strip()

zrb/task/llm_task.py

@@ -22,26 +22,24 @@ from zrb.input.any_input import AnyInput
 from zrb.task.any_task import AnyTask
 from zrb.task.base_task import BaseTask
 from zrb.task.llm.agent import get_agent, run_agent_iteration
-
-# No longer need llm_config here
 from zrb.task.llm.config import (
     get_model,
     get_model_settings,
 )
-from zrb.task.llm.context import get_conversation_context
+from zrb.task.llm.context import extract_default_context, get_conversation_context
 from zrb.task.llm.context_enrichment import maybe_enrich_context
 from zrb.task.llm.history import (
     ConversationHistoryData,
     ListOfDict,
-    prepare_initial_state,
+    read_conversation_history,
     write_conversation_history,
 )
 from zrb.task.llm.history_summarization import maybe_summarize_history
 from zrb.task.llm.prompt import (
-    build_user_prompt,
     get_combined_system_prompt,
     get_context_enrichment_prompt,
     get_summarization_prompt,
+    get_user_message,
 )
 from zrb.util.cli.style import stylize_faint
 from zrb.xcom.xcom import Xcom
@@ -240,6 +238,7 @@ class LLMTask(BaseTask):
             summarization_prompt_attr=self._summarization_prompt,
             render_summarization_prompt=self._render_summarization_prompt,
         )
+        user_message = get_user_message(ctx, self._message)
         # Get the combined system prompt using the new getter
         system_prompt = get_combined_system_prompt(
             ctx=ctx,
@@ -250,17 +249,19 @@ class LLMTask(BaseTask):
             special_instruction_prompt_attr=self._special_instruction_prompt,
             render_special_instruction_prompt=self._render_special_instruction_prompt,
         )
-        # 1. Prepare initial state (read history, get initial context)
-        history_list, conversation_context = await prepare_initial_state(
+        # 1. Prepare initial state (read history from previous session)
+        conversation_history = await read_conversation_history(
             ctx=ctx,
             conversation_history_reader=self._conversation_history_reader,
             conversation_history_file_attr=self._conversation_history_file,
             render_history_file=self._render_history_file,
             conversation_history_attr=self._conversation_history,
-            conversation_context_getter=lambda c: get_conversation_context(
-                c, self._conversation_context
-            ),
         )
+        history_list = conversation_history.history
+        conversation_context = {
+            **conversation_history.context,
+            **get_conversation_context(ctx, self._conversation_context),
+        }
         # 2. Enrich context (optional)
         conversation_context = await maybe_enrich_context(
             ctx=ctx,
@@ -289,18 +290,21 @@ class LLMTask(BaseTask):
             model_settings=model_settings,
             summarization_prompt=summarization_prompt,
         )
-        # 4. Build the final user prompt
-        user_prompt = build_user_prompt(
-            ctx=ctx,
-            message_attr=self._message,
-            conversation_context=conversation_context,
+        # 4. Build the final user prompt and system prompt
+        final_user_prompt, default_context = extract_default_context(user_message)
+        final_system_prompt = "\n".join(
+            [
+                system_prompt,
+                "# Context",
+                json.dumps({**default_context, **conversation_context}),
+            ]
         )
         # 5. Get the agent instance
         agent = get_agent(
             ctx=ctx,
             agent_attr=self._agent,
             model=model,
-            system_prompt=system_prompt,
+            system_prompt=final_system_prompt,
             model_settings=model_settings,
             tools_attr=self._tools,
             additional_tools=self._additional_tools,
@@ -309,7 +313,7 @@ class LLMTask(BaseTask):
         )
         # 6. Run the agent iteration and save the results/history
         return await self._run_agent_and_save_history(
-            ctx, agent, user_prompt, history_list, conversation_context
+            ctx, agent, final_user_prompt, history_list, conversation_context
         )
 
     async def _run_agent_and_save_history(
@@ -346,7 +350,7 @@ class LLMTask(BaseTask):
                     ctx.xcom[xcom_usage_key] = Xcom([])
                 usage = agent_run.result.usage()
                 ctx.xcom.get(xcom_usage_key).push(usage)
-                ctx.print(stylize_faint(f"[USAGE] {usage}"))
+                ctx.print(stylize_faint(f"[Token Usage] {usage}"), plain=True)
                 return agent_run.result.output
             else:
                 ctx.log_warning("Agent run did not produce a result.")

{zrb-1.5.17.dist-info → zrb-1.6.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.5.17
+Version: 1.6.1
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later