zrb 1.21.29__py3-none-any.whl → 2.0.0a4__py3-none-any.whl

zrb/task/llm/history_processor.py

@@ -1,206 +0,0 @@
-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,25 +0,0 @@
-from zrb.attr.type import IntAttr
-from zrb.config.llm_config import llm_config
-from zrb.context.any_context import AnyContext
-from zrb.util.attr import get_int_attr
-
-
-def get_history_summarization_token_threshold(
-    ctx: AnyContext,
-    history_summarization_token_threshold_attr: IntAttr | None,
-    render_history_summarization_token_threshold: bool,
-) -> int:
-    """Gets the history summarization token threshold, handling defaults and errors."""
-    try:
-        return get_int_attr(
-            ctx,
-            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_token_threshold to int: {e}. "
-            "Defaulting to -1 (no threshold)."
-        )
-        return -1

zrb/task/llm/print_node.py

@@ -1,221 +0,0 @@
-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, 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,
-        FunctionToolResultEvent,
-        PartDeltaEvent,
-        PartStartEvent,
-        TextPartDelta,
-        ThinkingPartDelta,
-        ToolCallPartDelta,
-    )
-
-    meta = getattr(node, "id", None) or getattr(node, "request_id", None)
-    if Agent.is_user_prompt_node(node):
-        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(_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(
-                            _format_content(
-                                f"Result: tool_name={tool_name}", log_indent_level
-                            )
-                        )
-                        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(_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(_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")