zrb 1.8.15__py3-none-any.whl → 1.9.1__py3-none-any.whl

zrb/task/llm/context.py

@@ -1,14 +1,9 @@
 import datetime
-import inspect
 import os
 import platform
 import re
-from collections.abc import Callable
 from typing import Any
 
-from zrb.context.any_context import AnyContext
-from zrb.context.any_shared_context import AnySharedContext
-from zrb.util.attr import get_attr
 from zrb.util.file import read_dir, read_file_with_line_numbers
 
 
@@ -61,42 +56,3 @@ def extract_default_context(user_message: str) -> tuple[str, dict[str, Any]]:
     }
 
     return modified_user_message, context
-
-
-def get_conversation_context(
-    ctx: AnyContext,
-    conversation_context_attr: (
-        dict[str, Any] | Callable[[AnySharedContext], dict[str, Any]] | None
-    ),
-) -> dict[str, Any]:
-    """
-    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)
-    if not isinstance(raw_context, dict):
-        ctx.log_warning(
-            f"Conversation context resolved to type {type(raw_context)}, "
-            "expected dict. Returning empty context."
-        )
-        return {}
-    # If conversation_context contains callable value, execute them.
-    processed_context: dict[str, Any] = {}
-    for key, value in raw_context.items():
-        if callable(value):
-            try:
-                # Check if the callable expects 'ctx'
-                sig = inspect.signature(value)
-                if "ctx" in sig.parameters:
-                    processed_context[key] = value(ctx)
-                else:
-                    processed_context[key] = value()
-            except Exception as e:
-                ctx.log_warning(
-                    f"Error executing callable for context key '{key}': {e}. "
-                    "Skipping."
-                )
-                processed_context[key] = None
-        else:
-            processed_context[key] = value
-    return processed_context

zrb/task/llm/context_enrichment.py

@@ -1,13 +1,11 @@
 import json
 import traceback
-from typing import TYPE_CHECKING, Any
-
-from pydantic import BaseModel
+from typing import TYPE_CHECKING
 
 from zrb.attr.type import BoolAttr, IntAttr
 from zrb.context.any_context import AnyContext
 from zrb.llm_config import llm_config
-from zrb.llm_rate_limitter import LLMRateLimiter
+from zrb.llm_rate_limitter import LLMRateLimiter, llm_rate_limitter
 from zrb.task.llm.agent import run_agent_iteration
 from zrb.task.llm.history import (
     count_part_in_history_list,
@@ -20,114 +18,87 @@ 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
-
-
-class EnrichmentConfig(BaseModel):
-    model_config = {"arbitrary_types_allowed": True}
-    model: Model | str | None = None
-    settings: ModelSettings | None = None
-    prompt: str
-    retries: int = 3
 
 
-class EnrichmentResult(BaseModel):
-    response: dict[str, Any]  # or further decompose as needed
+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,
-    config: EnrichmentConfig,
-    conversation_context: dict[str, Any],
+    model: "Model | str | None",
+    settings: "ModelSettings | None",
+    prompt: str,
+    previous_long_term_context: str,
     history_list: ListOfDict,
     rate_limitter: LLMRateLimiter | None = None,
-) -> dict[str, Any]:
-    """Runs an LLM call to extract key info and merge it into the context."""
+    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...")
-    # Prepare context and history for the enrichment prompt
-    history_summary = conversation_context.get("history_summary")
-    try:
-        context_json = json.dumps(conversation_context)
-        history_json = json.dumps(history_list)
-        # The user prompt will now contain the dynamic data
-        user_prompt_data = "\n".join(
-            [
-                "Extract context from the following conversation info.",
-                "Extract only contexts that will be relevant across multiple conversations, like",  # noqa
-                "- user name",
-                "- user hobby",
-                "- user's long life goal",
-                "- standard/SOP",
-                "- etc.",
-                "Always maintain the relevant context and remove the irrelevant ones.",
-                "Restructure the context in a helpful way",
-                "Keep the context small",
-                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
-
+    # 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=config.model,
-        system_prompt=config.prompt,  # Use the main prompt as system prompt
-        model_settings=config.settings,
-        retries=config.retries,
-        output_type=EnrichmentResult,
+        model=model,
+        system_prompt=prompt,
+        model_settings=settings,
+        retries=retries,
     )
 
     try:
-        ctx.print(stylize_faint("[Context Enrichment Triggered]"), plain=True)
+        ctx.print(stylize_faint("💡 Enrich Context"), plain=True)
         enrichment_run = await run_agent_iteration(
             ctx=ctx,
             agent=enrichment_agent,
-            user_prompt=user_prompt_data,  # Pass the formatted data as user prompt
-            history_list=[],  # Enrichment agent doesn't need prior history itself
+            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:
-            response = enrichment_run.result.output.response
+            new_long_term_context = str(enrichment_run.result.output)
             usage = enrichment_run.result.usage()
-            ctx.print(stylize_faint(f"[Token Usage] {usage}"), plain=True)
-            if response:
-                conversation_context = response
-                # Re inject history summary
-                conversation_context["history_summary"] = history_summary
-                ctx.log_info("Context enriched based on history.")
-                ctx.log_info(
-                    f"Updated conversation context: {json.dumps(conversation_context)}"
-                )
+            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")
+            ctx.log_warning("Context enrichment returned no data.")
     except Exception as e:
         ctx.log_warning(f"Error during context enrichment LLM call: {e}")
         traceback.print_exc()
-    return conversation_context
+
+    # Return the original context if enrichment fails
+    return previous_long_term_context
 
 
-def get_context_enrichment_threshold(
+def get_context_enrichment_token_threshold(
     ctx: AnyContext,
-    context_enrichment_threshold_attr: IntAttr | None,
-    render_context_enrichment_threshold: bool,
+    context_enrichment_token_threshold_attr: IntAttr | None,
+    render_context_enrichment_token_threshold: bool,
 ) -> int:
-    """Gets the context enrichment threshold, handling defaults and errors."""
+    """Gets the context enrichment token threshold, handling defaults and errors."""
     try:
         return get_int_attr(
             ctx,
-            context_enrichment_threshold_attr,
-            # Use llm_config default if attribute is None
-            llm_config.default_context_enrichment_threshold,
-            auto_render=render_context_enrichment_threshold,
+            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_threshold to int: {e}. "
+            f"Could not convert context_enrichment_token_threshold to int: {e}. "
             "Defaulting to -1 (no threshold)."
         )
         return -1
@@ -136,23 +107,27 @@ def get_context_enrichment_threshold(
 def should_enrich_context(
     ctx: AnyContext,
     history_list: ListOfDict,
-    should_enrich_context_attr: BoolAttr | None,  # Allow None
+    should_enrich_context_attr: BoolAttr | None,
     render_enrich_context: bool,
-    context_enrichment_threshold_attr: IntAttr | None,
-    render_context_enrichment_threshold: bool,
+    context_enrichment_token_threshold_attr: IntAttr | None,
+    render_context_enrichment_token_threshold: bool,
 ) -> bool:
     """
-    Determines if context enrichment should occur based on history, threshold, and config.
+    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_threshold = get_context_enrichment_threshold(
+    enrichment_token_threshold = get_context_enrichment_token_threshold(
         ctx,
-        context_enrichment_threshold_attr,
-        render_context_enrichment_threshold,
+        context_enrichment_token_threshold_attr,
+        render_context_enrichment_token_threshold,
     )
-    if enrichment_threshold == -1 or enrichment_threshold > history_part_count:
+    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,
@@ -165,35 +140,33 @@ def should_enrich_context(
 async def maybe_enrich_context(
     ctx: AnyContext,
     history_list: ListOfDict,
-    conversation_context: dict[str, Any],
+    long_term_context: str,
     should_enrich_context_attr: BoolAttr | None,
     render_enrich_context: bool,
-    context_enrichment_threshold_attr: IntAttr | None,
-    render_context_enrichment_threshold: bool,
-    model: str | Model | None,
-    model_settings: ModelSettings | None,
+    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,
-) -> dict[str, Any]:
-    """Enriches context based on history if enabled and threshold met."""
+) -> 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_threshold_attr,
-        render_context_enrichment_threshold,
+        context_enrichment_token_threshold_attr,
+        render_context_enrichment_token_threshold,
     ):
         return await enrich_context(
             ctx=ctx,
-            config=EnrichmentConfig(
-                model=model,
-                settings=model_settings,
-                prompt=context_enrichment_prompt,
-            ),
-            conversation_context=conversation_context,
+            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 conversation_context
+    return long_term_context

zrb/task/llm/error.py

@@ -1,12 +1,10 @@
 import json
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Optional
 
 from pydantic import BaseModel
 
 if TYPE_CHECKING:
     from openai import APIError
-else:
-    APIError = Any
 
 
 # Define a structured error model for tool execution failures
@@ -17,7 +15,7 @@ class ToolExecutionError(BaseModel):
     details: Optional[str] = None
 
 
-def extract_api_error_details(error: APIError) -> str:
+def extract_api_error_details(error: "APIError") -> str:
     """Extract detailed error information from an APIError."""
     details = f"{error.message}"
     # Try to parse the error body as JSON

zrb/task/llm/history.py

@@ -4,7 +4,7 @@ from collections.abc import Callable
 from copy import deepcopy
 from typing import Any, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from zrb.attr.type import StrAttr
 from zrb.context.any_context import AnyContext
@@ -17,8 +17,18 @@ from zrb.util.run import run_async
 
 # Define the new ConversationHistoryData model
 class ConversationHistoryData(BaseModel):
-    context: dict[str, Any] = {}
-    history: ListOfDict = []
+    long_term_context: str = Field(
+        default="",
+        description="A markdown-formatted string containing curated, long-term context.",
+    )
+    conversation_summary: str = Field(
+        default="",
+        description="A free-text summary of the conversation history.",
+    )
+    history: ListOfDict = Field(
+        default_factory=list,
+        description="The recent, un-summarized conversation history.",
+    )
 
     @classmethod
     async def read_from_sources(
@@ -69,19 +79,17 @@ class ConversationHistoryData(BaseModel):
         try:
             if isinstance(data, cls):
                 return data  # Already a valid instance
-            if isinstance(data, dict) and "history" in data:
-                # Standard format {'context': ..., 'history': ...}
-                # Ensure context exists, even if empty
-                data.setdefault("context", {})
+            if isinstance(data, dict):
+                # This handles both the new format and the old {'context': ..., 'history': ...}
                 return cls.model_validate(data)
             elif isinstance(data, list):
-                # Handle old format (just a list) - wrap it
+                # Handle very old format (just a list) - wrap it
                 ctx.log_warning(
-                    f"History from {source} contains old list format. "
-                    "Wrapping it into the new structure {'context': {}, 'history': [...]}. "
+                    f"History from {source} contains legacy list format. "
+                    "Wrapping it into the new structure. "
                     "Consider updating the source format."
                 )
-                return cls(history=data, context={})
+                return cls(history=data)
             else:
                 ctx.log_warning(
                     f"History data from {source} has unexpected format "

zrb/task/llm/history_summarization.py

@@ -1,12 +1,11 @@
 import json
-from typing import TYPE_CHECKING, Any
-
-from pydantic import BaseModel
+import traceback
+from typing import TYPE_CHECKING
 
 from zrb.attr.type import BoolAttr, IntAttr
 from zrb.context.any_context import AnyContext
 from zrb.llm_config import llm_config
-from zrb.llm_rate_limitter import LLMRateLimiter
+from zrb.llm_rate_limitter import LLMRateLimiter, llm_rate_limitter
 from zrb.task.llm.agent import run_agent_iteration
 from zrb.task.llm.history import (
     count_part_in_history_list,
@@ -19,28 +18,30 @@ 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
 
 
-def get_history_summarization_threshold(
+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)
+
+
+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
@@ -49,110 +50,99 @@ def get_history_summarization_threshold(
 def should_summarize_history(
     ctx: AnyContext,
     history_list: ListOfDict,
-    should_summarize_history_attr: BoolAttr | None,  # Allow None
+    should_summarize_history_attr: BoolAttr | None,
     render_summarize_history: bool,
-    history_summarization_threshold_attr: IntAttr | None,  # Allow None
-    render_history_summarization_threshold: bool,
+    history_summarization_token_threshold_attr: IntAttr | None,
+    render_history_summarization_token_threshold: bool,
 ) -> bool:
-    """Determines if history summarization should occur based on length and config."""
+    """Determines if history summarization should occur based on token 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(
+    summarization_token_threshold = get_history_summarization_token_threshold(
         ctx,
-        history_summarization_threshold_attr,
-        render_history_summarization_threshold,
+        history_summarization_token_threshold_attr,
+        render_history_summarization_token_threshold,
     )
-    if summarization_threshold == -1 or summarization_threshold > history_part_count:
+    history_token_count = _count_token_in_history(history_list)
+    if (
+        summarization_token_threshold == -1
+        or summarization_token_threshold > history_token_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],
+    model: "Model | str | None",
+    settings: "ModelSettings | None",
+    prompt: str,
+    previous_summary: str,
     history_list: ListOfDict,
     rate_limitter: LLMRateLimiter | None = None,
-) -> dict[str, Any]:
-    """Runs an LLM call to summarize history and update the context."""
+    retries: int = 3,
+) -> str:
+    """Runs an LLM call to update the conversation summary."""
     from pydantic_ai import Agent
 
     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}
+    )
     summarization_agent = Agent(
-        model=config.model,
-        system_prompt=config.prompt,
-        model_settings=config.settings,
-        retries=config.retries,
+        model=model,
+        system_prompt=prompt,
+        model_settings=settings,
+        retries=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)
+        ctx.print(stylize_faint("📝 Summarize"), 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
+            user_prompt=user_prompt,
+            history_list=[],
             rate_limitter=rate_limitter,
         )
-        if summary_run and summary_run.result.output:
-            summary_text = str(summary_run.result.output)
+        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"[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}")
+            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 Exception as e:
         ctx.log_warning(f"Error during history summarization: {e}")
+        traceback.print_exc()
 
-    return conversation_context
+    # Return the original summary if summarization fails
+    return previous_summary
 
 
 async def maybe_summarize_history(
     ctx: AnyContext,
     history_list: ListOfDict,
-    conversation_context: dict[str, Any],
-    should_summarize_history_attr: BoolAttr | None,  # Allow None
+    conversation_summary: str,
+    should_summarize_history_attr: BoolAttr | 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,
+    history_summarization_token_threshold_attr: IntAttr | None,
+    render_history_summarization_token_threshold: bool,
+    model: "str | Model | None",
+    model_settings: "ModelSettings | None",
     summarization_prompt: str,
     rate_limitter: LLMRateLimiter | None = None,
-) -> tuple[ListOfDict, dict[str, Any]]:
+) -> tuple[ListOfDict, str]:
     """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(
@@ -160,21 +150,18 @@ async def maybe_summarize_history(
         shorten_history_list,
         should_summarize_history_attr,
         render_summarize_history,
-        history_summarization_threshold_attr,
-        render_history_summarization_threshold,
+        history_summarization_token_threshold_attr,
+        render_history_summarization_token_threshold,
     ):
-        # Use summarize_history defined above
-        updated_context = await summarize_history(
+        new_summary = 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
+            model=model,
+            settings=model_settings,
+            prompt=summarization_prompt,
+            previous_summary=conversation_summary,
+            history_list=shorten_history_list,
             rate_limitter=rate_limitter,
         )
-        # Truncate the history list after summarization
-        return [], updated_context
-    return history_list, conversation_context
+        # After summarization, the history is cleared and replaced by the new summary
+        return [], new_summary
+    return history_list, conversation_summary