zrb 1.8.10__py3-none-any.whl → 1.21.29__py3-none-any.whl

zrb/task/llm/history_processor.py

@@ -0,0 +1,206 @@
+import json
+import sys
+import traceback
+from typing import TYPE_CHECKING, Any, Callable, Coroutine
+
+from zrb.config.llm_config import llm_config
+from zrb.config.llm_rate_limitter import LLMRateLimitter
+from zrb.config.llm_rate_limitter import llm_rate_limitter as default_llm_rate_limitter
+from zrb.context.any_context import AnyContext
+from zrb.task.llm.agent_runner import run_agent_iteration
+from zrb.util.cli.style import stylize_faint
+from zrb.util.markdown import make_markdown_section
+
+if sys.version_info >= (3, 12):
+    from typing import TypedDict
+else:
+    from typing_extensions import TypedDict
+
+
+if TYPE_CHECKING:
+    from pydantic_ai import ModelMessage
+    from pydantic_ai.models import Model
+    from pydantic_ai.settings import ModelSettings
+
+
+class SingleMessage(TypedDict):
+    """
+    SingleConversation
+
+    Attributes:
+        role: Either AI, User, Tool Call, or Tool Result
+        time: yyyy-mm-ddTHH:MM:SSZ:
+        content: The content of the message (summarize if too long)
+    """
+
+    role: str
+    time: str
+    content: str
+
+
+class ConversationSummary(TypedDict):
+    """
+    Conversation history
+
+    Attributes:
+        transcript: Several last transcript of the conversation
+        summary: Descriptive conversation summary
+    """
+
+    transcript: list[SingleMessage]
+    summary: str
+
+
+def save_conversation_summary(conversation_summary: ConversationSummary):
+    """
+    Write conversation summary for main assistant to continue conversation.
+    """
+    return conversation_summary
+
+
+def create_summarize_history_processor(
+    ctx: AnyContext,
+    system_prompt: str,
+    rate_limitter: LLMRateLimitter | None = None,
+    summarization_model: "Model | str | None" = None,
+    summarization_model_settings: "ModelSettings | None" = None,
+    summarization_system_prompt: str | None = None,
+    summarization_token_threshold: int | None = None,
+    summarization_retries: int = 2,
+) -> Callable[[list["ModelMessage"]], Coroutine[None, None, list["ModelMessage"]]]:
+    from pydantic_ai import Agent, ModelMessage, ModelRequest
+    from pydantic_ai.messages import ModelMessagesTypeAdapter, UserPromptPart
+
+    if rate_limitter is None:
+        rate_limitter = default_llm_rate_limitter
+    if summarization_model is None:
+        summarization_model = llm_config.default_small_model
+    if summarization_model_settings is None:
+        summarization_model_settings = llm_config.default_small_model_settings
+    if summarization_system_prompt is None:
+        summarization_system_prompt = llm_config.default_summarization_prompt
+    if summarization_token_threshold is None:
+        summarization_token_threshold = (
+            llm_config.default_history_summarization_token_threshold
+        )
+
+    async def maybe_summarize_history(
+        messages: list[ModelMessage],
+    ) -> list[ModelMessage]:
+        history_list = json.loads(ModelMessagesTypeAdapter.dump_json(messages))
+        history_json_str = json.dumps(history_list)
+        # Estimate token usage
+        # Note: Pydantic ai has run context parameter
+        # (https://ai.pydantic.dev/message-history/#runcontext-parameter)
+        # But we cannot use run_ctx.usage.total_tokens because total token keep increasing
+        # even after summariztion.
+        estimated_token_usage = rate_limitter.count_token(history_json_str)
+        _print_request_info(
+            ctx, estimated_token_usage, summarization_token_threshold, messages
+        )
+        if estimated_token_usage < summarization_token_threshold or len(messages) == 1:
+            return messages
+        history_list_without_instruction = [
+            {
+                key: obj[key]
+                for key in obj
+                if index == len(history_list) - 1 or key != "instructions"
+            }
+            for index, obj in enumerate(history_list)
+        ]
+        history_json_str_without_instruction = json.dumps(
+            history_list_without_instruction
+        )
+        summarization_message = f"Summarize the following conversation: {history_json_str_without_instruction}"
+        summarization_agent = Agent[None, ConversationSummary](
+            model=summarization_model,
+            output_type=save_conversation_summary,
+            instructions=summarization_system_prompt,
+            model_settings=summarization_model_settings,
+            retries=summarization_retries,
+        )
+        try:
+            _print_info(ctx, "📝 Rollup Conversation", 2)
+            summary_run = await run_agent_iteration(
+                ctx=ctx,
+                agent=summarization_agent,
+                user_prompt=summarization_message,
+                attachments=[],
+                history_list=[],
+                rate_limitter=rate_limitter,
+                log_indent_level=2,
+            )
+            if summary_run and summary_run.result and summary_run.result.output:
+                usage = summary_run.result.usage()
+                _print_info(ctx, f"📝 Rollup Conversation Token: {usage}", 2)
+                ctx.print(plain=True)
+                ctx.log_info("History summarized and updated.")
+                condensed_message = make_markdown_section(
+                    header="Past Conversation",
+                    content="\n".join(
+                        [
+                            make_markdown_section(
+                                "Summary", _extract_summary(summary_run.result.output)
+                            ),
+                            make_markdown_section(
+                                "Past Trancript",
+                                _extract_transcript(summary_run.result.output),
+                            ),
+                        ]
+                    ),
+                )
+                return [
+                    ModelRequest(
+                        instructions=system_prompt,
+                        parts=[UserPromptPart(condensed_message)],
+                    )
+                ]
+            ctx.log_warning("History summarization failed or returned no data.")
+        except BaseException as e:
+            ctx.log_warning(f"Error during history summarization: {e}")
+            traceback.print_exc()
+        return messages
+
+    return maybe_summarize_history
+
+
+def _print_request_info(
+    ctx: AnyContext,
+    estimated_token_usage: int,
+    summarization_token_threshold: int,
+    messages: list["ModelMessage"],
+):
+    _print_info(ctx, f"Current request token (estimated): {estimated_token_usage}")
+    _print_info(ctx, f"Summarization token threshold: {summarization_token_threshold}")
+    _print_info(ctx, f"History length: {len(messages)}")
+
+
+def _print_info(ctx: AnyContext, text: str, log_indent_level: int = 0):
+    log_prefix = (2 * (log_indent_level + 1)) * " "
+    ctx.print(stylize_faint(f"{log_prefix}{text}"), plain=True)
+
+
+def _extract_summary(summary_result_output: dict[str, Any] | str) -> str:
+    summary = (
+        summary_result_output.get("summary", "")
+        if isinstance(summary_result_output, dict)
+        else ""
+    )
+    return summary
+
+
+def _extract_transcript(summary_result_output: dict[str, Any] | str) -> str:
+    transcript_list = (
+        summary_result_output.get("transcript", [])
+        if isinstance(summary_result_output, dict)
+        else []
+    )
+    transcript_list = [] if not isinstance(transcript_list, list) else transcript_list
+    return "\n".join(_format_transcript_message(message) for message in transcript_list)
+
+
+def _format_transcript_message(message: dict[str, str]) -> str:
+    role = message.get("role", "Message")
+    time = message.get("time", "<unknown>")
+    content = message.get("content", "<empty>")
+    return f"{role} ({time}): {content}"

zrb/task/llm/history_summarization.py

@@ -1,180 +1,25 @@
-import json
-from typing import TYPE_CHECKING, Any
-
-from pydantic import BaseModel
-
-from zrb.attr.type import BoolAttr, IntAttr
+from zrb.attr.type import IntAttr
+from zrb.config.llm_config import llm_config
 from zrb.context.any_context import AnyContext
-from zrb.llm_config import llm_config
-from zrb.llm_rate_limitter import LLMRateLimiter
-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
-    from pydantic_ai.settings import ModelSettings
-else:
-    Model = Any
-    ModelSettings = Any
+from zrb.util.attr import get_int_attr
 
 
-def get_history_summarization_threshold(
+def get_history_summarization_token_threshold(
     ctx: AnyContext,
-    history_summarization_threshold_attr: IntAttr | None,
-    render_history_summarization_threshold: bool,
+    history_summarization_token_threshold_attr: IntAttr | None,
+    render_history_summarization_token_threshold: bool,
 ) -> int:
-    """Gets the history summarization threshold, handling defaults and errors."""
+    """Gets the history summarization token threshold, handling defaults and errors."""
     try:
         return get_int_attr(
             ctx,
-            history_summarization_threshold_attr,
-            # Use llm_config default if attribute is None
-            llm_config.default_history_summarization_threshold,
-            auto_render=render_history_summarization_threshold,
+            history_summarization_token_threshold_attr,
+            llm_config.default_history_summarization_token_threshold,
+            auto_render=render_history_summarization_token_threshold,
         )
     except ValueError as e:
         ctx.log_warning(
-            f"Could not convert history_summarization_threshold to int: {e}. "
+            f"Could not convert history_summarization_token_threshold to int: {e}. "
             "Defaulting to -1 (no threshold)."
         )
         return -1
-
-
-def should_summarize_history(
-    ctx: AnyContext,
-    history_list: ListOfDict,
-    should_summarize_history_attr: BoolAttr | None,  # Allow None
-    render_summarize_history: bool,
-    history_summarization_threshold_attr: IntAttr | None,  # Allow None
-    render_history_summarization_threshold: bool,
-) -> bool:
-    """Determines if history summarization should occur based on length and config."""
-    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_count:
-        return False
-    return get_bool_attr(
-        ctx,
-        should_summarize_history_attr,
-        # Use llm_config default if attribute is None
-        llm_config.default_summarize_history,
-        auto_render=render_summarize_history,
-    )
-
-
-class SummarizationConfig(BaseModel):
-    model_config = {"arbitrary_types_allowed": True}
-    model: Model | str | None = None
-    settings: ModelSettings | None = None
-    prompt: str
-    retries: int = 3
-
-
-async def summarize_history(
-    ctx: AnyContext,
-    config: SummarizationConfig,
-    conversation_context: dict[str, Any],
-    history_list: ListOfDict,
-    rate_limitter: LLMRateLimiter | None = None,
-) -> dict[str, Any]:
-    """Runs an LLM call to summarize history and update the context."""
-    from pydantic_ai import Agent
-
-    ctx.log_info("Attempting to summarize conversation history...")
-
-    summarization_agent = Agent(
-        model=config.model,
-        system_prompt=config.prompt,
-        model_settings=config.settings,
-        retries=config.retries,
-    )
-
-    # Prepare context and history for summarization prompt
-    try:
-        context_json = json.dumps(conversation_context)
-        history_to_summarize_json = json.dumps(history_list)
-        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,
-            user_prompt=summarization_user_prompt,
-            history_list=[],  # Summarization agent doesn't need prior history
-            rate_limitter=rate_limitter,
-        )
-        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.")
-            ctx.log_info(f"Conversation summary: {summary_text}")
-        else:
-            ctx.log_warning("History summarization failed or returned no data.")
-    except Exception as e:
-        ctx.log_warning(f"Error during history summarization: {e}")
-
-    return conversation_context
-
-
-async def maybe_summarize_history(
-    ctx: AnyContext,
-    history_list: ListOfDict,
-    conversation_context: dict[str, Any],
-    should_summarize_history_attr: BoolAttr | None,  # Allow None
-    render_summarize_history: bool,
-    history_summarization_threshold_attr: IntAttr | None,  # Allow None
-    render_history_summarization_threshold: bool,
-    model: str | Model | None,
-    model_settings: ModelSettings | None,
-    summarization_prompt: str,
-    rate_limitter: LLMRateLimiter | None = None,
-) -> 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,
-        shorten_history_list,
-        should_summarize_history_attr,
-        render_summarize_history,
-        history_summarization_threshold_attr,
-        render_history_summarization_threshold,
-    ):
-        # Use summarize_history defined above
-        updated_context = await summarize_history(
-            ctx=ctx,
-            config=SummarizationConfig(
-                model=model,
-                settings=model_settings,
-                prompt=summarization_prompt,
-            ),
-            conversation_context=conversation_context,
-            history_list=shorten_history_list,  # Pass the full list for context
-            rate_limitter=rate_limitter,
-        )
-        # Truncate the history list after summarization
-        return [], updated_context
-    return history_list, conversation_context

zrb/task/llm/print_node.py

@@ -1,12 +1,17 @@
+import json
 from collections.abc import Callable
 from typing import Any
 
+from zrb.config.config import CFG
 from zrb.util.cli.style import stylize_faint
 
 
-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,
@@ -14,84 +19,203 @@ async def print_node(print_func: Callable, agent_run: Any, node: Any):
         PartDeltaEvent,
         PartStartEvent,
         TextPartDelta,
+        ThinkingPartDelta,
         ToolCallPartDelta,
     )
 
+    meta = getattr(node, "id", None) or getattr(node, "request_id", None)
     if Agent.is_user_prompt_node(node):
-        # A user prompt node => The user has provided input
-        print_func(stylize_faint(f">> UserPromptNode: {node.user_prompt}"))
+        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(">> ModelRequestNode: streaming partial request tokens")
-        )
-        async with node.stream(agent_run.ctx) as request_stream:
-            is_streaming = False
-            async for event in request_stream:
-                if isinstance(event, PartStartEvent):
-                    if is_streaming:
-                        print_func("")
-                    print_func(
-                        stylize_faint(
-                            f"[Request] Starting part {event.index}: {event.part!r}"
-                        ),
-                    )
-                    is_streaming = False
-                elif isinstance(event, PartDeltaEvent):
-                    if isinstance(event.delta, TextPartDelta):
+        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.content_delta}"),
-                            end="",
+                            _format_content(
+                                f"Result: tool_name={tool_name}", log_indent_level
+                            )
                         )
-                    elif isinstance(event.delta, ToolCallPartDelta):
-                        print_func(
-                            stylize_faint(f"{event.delta.args_delta}"),
-                            end="",
-                        )
-                    is_streaming = True
-                elif isinstance(event, FinalResultEvent):
-                    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(">> CallToolsNode: streaming partial response & tool usage")
-        )
-        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"[Tools] The LLM calls tool={event.part.tool_name!r} with args={event.part.args} (tool_call_id={event.part.tool_call_id!r})"  # noqa
-                        )
-                    )
-                elif isinstance(event, FunctionToolResultEvent):
-                    print_func(
-                        stylize_faint(
-                            f"[Tools] Tool call {event.tool_call_id!r} returned => {event.result.content}"  # noqa
+        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)
+                        and event.tool_call_id
+                    ):
+                        call_id = event.tool_call_id
+                        if CFG.LLM_SHOW_TOOL_CALL_RESULT:
+                            result_content = event.result.content
+                            print_func(
+                                _format_content(
+                                    f"{call_id} | Return {result_content}",
+                                    log_indent_level,
+                                )
+                            )
+                        else:
+                            print_func(
+                                _format_content(
+                                    f"{call_id} | Executed", log_indent_level
+                                )
+                            )
+        except UnexpectedModelBehavior as e:
+            print_func("")  # ensure newline consistency
+            print_func(
+                _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("[End of Response]"))
-        # print_func(stylize_faint(f"{agent_run.result.data}"))
+        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")