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

zrb/task/llm/context_enrichment.py

@@ -1,172 +0,0 @@
-import json
-import traceback
-from typing import TYPE_CHECKING
-
-from zrb.attr.type import BoolAttr, IntAttr
-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 (
-    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
-
-
-def _count_token_in_history(history_list: ListOfDict) -> int:
-    """Counts the total number of tokens in a conversation history list."""
-    text_to_count = json.dumps(history_list)
-    return llm_rate_limitter.count_token(text_to_count)
-
-
-async def enrich_context(
-    ctx: AnyContext,
-    model: "Model | str | None",
-    settings: "ModelSettings | None",
-    prompt: str,
-    previous_long_term_context: str,
-    history_list: ListOfDict,
-    rate_limitter: LLMRateLimiter | None = None,
-    retries: int = 3,
-) -> str:
-    """Runs an LLM call to update the long-term context and returns the new context string."""
-    from pydantic_ai import Agent
-
-    ctx.log_info("Attempting to enrich conversation context...")
-    # Construct the user prompt according to the new prompt format
-    user_prompt = json.dumps(
-        {
-            "previous_long_term_context": previous_long_term_context,
-            "recent_conversation_history": history_list,
-        }
-    )
-    enrichment_agent = Agent(
-        model=model,
-        system_prompt=prompt,
-        model_settings=settings,
-        retries=retries,
-    )
-
-    try:
-        ctx.print(stylize_faint("💡 Enrich Context"), plain=True)
-        enrichment_run = await run_agent_iteration(
-            ctx=ctx,
-            agent=enrichment_agent,
-            user_prompt=user_prompt,
-            history_list=[],  # Enrichment agent works off the prompt, not history
-            rate_limitter=rate_limitter,
-        )
-        if enrichment_run and enrichment_run.result.output:
-            new_long_term_context = str(enrichment_run.result.output)
-            usage = enrichment_run.result.usage()
-            ctx.print(
-                stylize_faint(f"💡 Context Enrichment Token: {usage}"), plain=True
-            )
-            ctx.print(plain=True)
-            ctx.log_info("Context enriched based on history.")
-            ctx.log_info(f"Updated long-term context:\n{new_long_term_context}")
-            return new_long_term_context
-        else:
-            ctx.log_warning("Context enrichment returned no data.")
-    except BaseException as e:
-        ctx.log_warning(f"Error during context enrichment LLM call: {e}")
-        traceback.print_exc()
-
-    # Return the original context if enrichment fails
-    return previous_long_term_context
-
-
-def get_context_enrichment_token_threshold(
-    ctx: AnyContext,
-    context_enrichment_token_threshold_attr: IntAttr | None,
-    render_context_enrichment_token_threshold: bool,
-) -> int:
-    """Gets the context enrichment token threshold, handling defaults and errors."""
-    try:
-        return get_int_attr(
-            ctx,
-            context_enrichment_token_threshold_attr,
-            llm_config.default_context_enrichment_token_threshold,
-            auto_render=render_context_enrichment_token_threshold,
-        )
-    except ValueError as e:
-        ctx.log_warning(
-            f"Could not convert context_enrichment_token_threshold to int: {e}. "
-            "Defaulting to -1 (no threshold)."
-        )
-        return -1
-
-
-def should_enrich_context(
-    ctx: AnyContext,
-    history_list: ListOfDict,
-    should_enrich_context_attr: BoolAttr | None,
-    render_enrich_context: bool,
-    context_enrichment_token_threshold_attr: IntAttr | None,
-    render_context_enrichment_token_threshold: bool,
-) -> bool:
-    """
-    Determines if context enrichment should occur based on history, token threshold, and config.
-    """
-    history_part_count = count_part_in_history_list(history_list)
-    if history_part_count == 0:
-        return False
-    enrichment_token_threshold = get_context_enrichment_token_threshold(
-        ctx,
-        context_enrichment_token_threshold_attr,
-        render_context_enrichment_token_threshold,
-    )
-    history_token_count = _count_token_in_history(history_list)
-    if (
-        enrichment_token_threshold == -1
-        or enrichment_token_threshold > history_token_count
-    ):
-        return False
-    return get_bool_attr(
-        ctx,
-        should_enrich_context_attr,
-        llm_config.default_enrich_context,
-        auto_render=render_enrich_context,
-    )
-
-
-async def maybe_enrich_context(
-    ctx: AnyContext,
-    history_list: ListOfDict,
-    long_term_context: str,
-    should_enrich_context_attr: BoolAttr | None,
-    render_enrich_context: bool,
-    context_enrichment_token_threshold_attr: IntAttr | None,
-    render_context_enrichment_token_threshold: bool,
-    model: "str | Model | None",
-    model_settings: "ModelSettings | None",
-    context_enrichment_prompt: str,
-    rate_limitter: LLMRateLimiter | None = None,
-) -> str:
-    """Enriches context based on history if enabled and token threshold met."""
-    shorten_history_list = replace_system_prompt_in_history_list(history_list)
-    if should_enrich_context(
-        ctx,
-        shorten_history_list,
-        should_enrich_context_attr,
-        render_enrich_context,
-        context_enrichment_token_threshold_attr,
-        render_context_enrichment_token_threshold,
-    ):
-        return await enrich_context(
-            ctx=ctx,
-            model=model,
-            settings=model_settings,
-            prompt=context_enrichment_prompt,
-            previous_long_term_context=long_term_context,
-            history_list=shorten_history_list,
-            rate_limitter=rate_limitter,
-        )
-    return long_term_context

zrb/task/llm/history.py

@@ -1,233 +0,0 @@
-import json
-import os
-from collections.abc import Callable
-from copy import deepcopy
-from typing import Any, Optional
-
-from zrb.attr.type import StrAttr
-from zrb.context.any_context import AnyContext
-from zrb.context.any_shared_context import AnySharedContext
-from zrb.task.llm.typing import ListOfDict
-from zrb.util.attr import get_str_attr
-from zrb.util.file import read_file, write_file
-from zrb.util.run import run_async
-
-
-# Define the new ConversationHistoryData model
-class ConversationHistoryData:
-    def __init__(
-        self,
-        long_term_context: str = "",
-        conversation_summary: str = "",
-        history: Optional[ListOfDict] = None,
-        messages: Optional[ListOfDict] = None,
-    ):
-        self.long_term_context = long_term_context
-        self.conversation_summary = conversation_summary
-        self.history = (
-            history
-            if history is not None
-            else (messages if messages is not None else [])
-        )
-
-    def to_dict(self) -> dict[str, Any]:
-        return {
-            "long_term_context": self.long_term_context,
-            "conversation_summary": self.conversation_summary,
-            "history": self.history,
-        }
-
-    def model_dump_json(self, indent: int = 2) -> str:
-        return json.dumps(self.to_dict(), indent=indent)
-
-    @classmethod
-    async def read_from_sources(
-        cls,
-        ctx: AnyContext,
-        reader: Callable[[AnyContext], dict[str, Any] | list | None] | None,
-        file_path: str | None,
-    ) -> Optional["ConversationHistoryData"]:
-        """Reads conversation history from various sources with priority."""
-        # 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."
-                )
-        # If neither reader nor file provided valid data
-        return None
-
-    @classmethod
-    def parse_and_validate(
-        cls, ctx: AnyContext, data: Any, source: str
-    ) -> Optional["ConversationHistoryData"]:
-        """Parses raw data into ConversationHistoryData, handling validation & old formats."""
-        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(
-                    long_term_context=data.get("long_term_context", ""),
-                    conversation_summary=data.get("conversation_summary", ""),
-                    history=data.get("history", data.get("messages")),
-                )
-            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."
-                )
-                return None
-        except Exception as e:  # Catch validation errors too
-            ctx.log_warning(
-                f"Error validating/parsing history data from {source}: {e}. Ignoring."
-            )
-            return None
-
-
-def get_history_file(
-    ctx: AnyContext,
-    conversation_history_file_attr: StrAttr | None,
-    render_history_file: bool,
-) -> str:
-    """Gets the path to the conversation history file, rendering if configured."""
-    return get_str_attr(
-        ctx,
-        conversation_history_file_attr,
-        "",
-        auto_render=render_history_file,
-    )
-
-
-async def read_conversation_history(
-    ctx: AnyContext,
-    conversation_history_reader: (
-        Callable[[AnySharedContext], 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
-    ),
-) -> ConversationHistoryData:
-    """Reads conversation history from reader, file, or attribute, with validation."""
-    history_file = get_history_file(
-        ctx, conversation_history_file_attr, render_history_file
-    )
-    # Use the class method defined above
-    history_data = await ConversationHistoryData.read_from_sources(
-        ctx=ctx,
-        reader=conversation_history_reader,
-        file_path=history_file,
-    )
-    if history_data:
-        return history_data
-    # Priority 3: Callable or direct conversation_history attribute
-    raw_data_attr: Any = None
-    if callable(conversation_history_attr):
-        try:
-            raw_data_attr = await run_async(conversation_history_attr(ctx))
-        except Exception as e:
-            ctx.log_warning(
-                f"Error executing callable conversation_history attribute: {e}. "
-                "Ignoring."
-            )
-    if raw_data_attr is None:
-        raw_data_attr = conversation_history_attr
-    if raw_data_attr:
-        # Use the class method defined above
-        history_data = ConversationHistoryData.parse_and_validate(
-            ctx, raw_data_attr, "attribute"
-        )
-        if history_data:
-            return history_data
-    # Fallback: Return default value
-    return ConversationHistoryData()
-
-
-async def write_conversation_history(
-    ctx: AnyContext,
-    history_data: ConversationHistoryData,
-    conversation_history_writer: (
-        Callable[[AnySharedContext, ConversationHistoryData], None] | None
-    ),
-    conversation_history_file_attr: StrAttr | None,
-    render_history_file: bool,
-):
-    """Writes conversation history using the writer or to a file."""
-    if conversation_history_writer is not None:
-        await run_async(conversation_history_writer(ctx, history_data))
-    history_file = get_history_file(
-        ctx, conversation_history_file_attr, render_history_file
-    )
-    if history_file != "":
-        write_file(history_file, json.dumps(history_data.to_dict(), indent=2))
-
-
-def replace_system_prompt_in_history_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