zrb 1.5.5__py3-none-any.whl → 1.5.6__py3-none-any.whl

zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/common/error.py

@@ -1,35 +1,33 @@
-from typing import Dict
-
 from fastapi import HTTPException
 
 
 class NotFoundError(HTTPException):
-    def __init__(self, message: str, headers: Dict[str, str] | None = None) -> None:
+    def __init__(self, message: str, headers: dict[str, str] | None = None) -> None:
         super().__init__(404, {"message": message}, headers)
 
 
 class ForbiddenError(HTTPException):
-    def __init__(self, message: str, headers: Dict[str, str] | None = None) -> None:
+    def __init__(self, message: str, headers: dict[str, str] | None = None) -> None:
         super().__init__(403, {"message": message}, headers)
 
 
 class UnauthorizedError(HTTPException):
-    def __init__(self, message: str, headers: Dict[str, str] | None = None) -> None:
+    def __init__(self, message: str, headers: dict[str, str] | None = None) -> None:
         super().__init__(401, {"message": message}, headers)
 
 
 class InvalidValueError(HTTPException):
-    def __init__(self, message: str, headers: Dict[str, str] | None = None) -> None:
+    def __init__(self, message: str, headers: dict[str, str] | None = None) -> None:
         super().__init__(422, {"message": message}, headers)
 
 
 class InternalServerError(HTTPException):
-    def __init__(self, message: str, headers: Dict[str, str] | None = None) -> None:
+    def __init__(self, message: str, headers: dict[str, str] | None = None) -> None:
         super().__init__(500, {"message": message}, headers)
 
 
 class ClientAPIError(HTTPException):
     def __init__(
-        self, status_code: int, message: str, headers: Dict[str, str] | None = None
+        self, status_code: int, message: str, headers: dict[str, str] | None = None
     ) -> None:
         super().__init__(status_code, {"message": message}, headers)

zrb/config.py

@@ -31,6 +31,7 @@ def _get_log_level(level: str) -> int:
     return logging.WARNING
 
 
+LOGGER = logging.getLogger()
 DEFAULT_SHELL = os.getenv("ZRB_SHELL", _get_current_shell())
 DEFAULT_EDITOR = os.getenv("ZRB_EDITOR", "nano")
 INIT_MODULES_STR = os.getenv("ZRB_INIT_MODULES", "")

zrb/llm_config.py

@@ -6,27 +6,61 @@ from pydantic_ai.providers import Provider
 from pydantic_ai.providers.openai import OpenAIProvider
 
 DEFAULT_SYSTEM_PROMPT = """
-You have access to tools.
-Your goal is to provide insightful and accurate information based on user queries.
-Follow these instructions precisely:
-1. ALWAYS use available tools to gather information BEFORE asking the user questions.
-2. For tools that require arguments: provide arguments in valid JSON format.
-3. For tools with no args: call the tool without args. Do NOT pass "" or {}.
-4. NEVER pass arguments to tools that don't accept parameters.
-5. NEVER ask users for information obtainable through tools.
-6. Use tools in a logical sequence until you have sufficient information.
-7. If a tool call fails, check if you're passing arguments in the correct format.
-   Consider alternative strategies if the issue persists.
-8. Only after exhausting relevant tools should you request clarification.
-9. Understand the context of user queries to provide relevant and accurate responses.
-10. Engage with users in a conversational manner once the necessary information is gathered.
-11. Adapt to different query types or scenarios to improve flexibility and effectiveness.
+You are a highly capable AI assistant with access to tools. Your primary goal is to
+provide accurate, reliable, and helpful responses.
+
+Key Instructions:
+1.  **Prioritize Tool Use:** Always attempt to use available tools to find
+    information or perform actions before asking the user.
+2.  **Correct Tool Invocation:** Use tools precisely as defined. Provide arguments
+    in valid JSON where required. Do not pass arguments to tools that don't accept
+    them. Handle tool errors gracefully and retry or adapt your strategy if necessary.
+3.  **Accuracy is Paramount:** Ensure all information, code, or outputs provided are
+    correct and reliable. Verify facts and generate executable code when requested.
+4.  **Clarity and Conciseness:** Respond clearly and directly to the user's query
+    after gathering the necessary information. Avoid unnecessary conversation.
+5.  **Context Awareness:** Understand the user's request fully to provide the most
+    relevant and effective assistance.
 """.strip()
 
 DEFAULT_PERSONA = """
 You are an expert in various fields including technology, science, history, and more.
 """.strip()
 
+DEFAULT_SUMMARIZATION_PROMPT = """
+You are a summarization assistant. Your task is to synthesize the provided 
+conversation history and the existing context (which might contain a
+previous 'history_summary') into a comprehensive, updated summary
+Carefully review the '# Current Context' which includes any previous summary
+('history_summary').
+Then, review the '# Conversation History to Summarize'.
+Combine the information from both the existing context/summary and the new
+history. Create a single, coherent summary that captures ALL essential
+information, including:
+- Key decisions made
+- Actions taken (including tool usage and their results)
+- Important facts or data points mentioned
+- Outcomes of discussions or actions
+- Any unresolved questions or pending items
+The goal is to provide a complete yet concise background so that the main
+assistant can seamlessly continue the conversation without losing critical
+context from the summarized parts.
+Output *only* the updated summary text."
+""".strip()
+
+DEFAULT_CONTEXT_ENRICHMENT_PROMPT = """
+You are an information extraction assistant.
+Analyze the conversation history and current context to extract key facts such as:
+  - user_name
+  - user_roles
+  - preferences
+  - goals
+  - etc
+Return only a JSON object containing a single key "response",
+whose value is another JSON object with these details.
+If nothing is found, output {"response": {}}.
+""".strip()
+
 
 class LLMConfig:
 
@@ -37,6 +71,8 @@ class LLMConfig:
         default_api_key: str | None = None,
         default_persona: str | None = None,
         default_system_prompt: str | None = None,
+        default_summarization_prompt: str | None = None,
+        default_context_enrichment_prompt: str | None = None,
     ):
         self._default_model_name = (
             default_model_name
@@ -63,6 +99,16 @@ class LLMConfig:
             if default_persona is not None
             else os.getenv("ZRB_LLM_PERSONA", None)
         )
+        self._default_summarization_prompt = (
+            default_summarization_prompt
+            if default_summarization_prompt is not None
+            else os.getenv("ZRB_LLM_SUMMARIZATION_PROMPT", None)
+        )
+        self._default_context_enrichment_prompt = (
+            default_context_enrichment_prompt
+            if default_context_enrichment_prompt is not None
+            else os.getenv("ZRB_LLM_CONTEXT_ENRICHMENT_PROMPT", None)
+        )
         self._default_provider = None
         self._default_model = None
 
@@ -93,6 +139,20 @@ class LLMConfig:
             return f"{persona}\n{system_prompt}"
         return system_prompt
 
+    def get_default_summarization_prompt(self) -> str:
+        return (
+            DEFAULT_SUMMARIZATION_PROMPT
+            if self._default_summarization_prompt is None
+            else self._default_summarization_prompt
+        )
+
+    def get_default_context_enrichment_prompt(self) -> str:
+        return (
+            DEFAULT_CONTEXT_ENRICHMENT_PROMPT
+            if self._default_context_enrichment_prompt is None
+            else self._default_context_enrichment_prompt
+        )
+
     def get_default_model(self) -> Model | str | None:
         if self._default_model is not None:
             return self._default_model
@@ -110,6 +170,12 @@ class LLMConfig:
     def set_default_system_prompt(self, system_prompt: str):
         self._default_system_prompt = system_prompt
 
+    def set_default_summarization_prompt(self, summarization_prompt: str):
+        self._default_summarization_prompt = summarization_prompt
+
+    def set_default_context_enrichment_prompt(self, context_enrichment_prompt: str):
+        self._default_context_enrichment_prompt = context_enrichment_prompt
+
     def set_default_model_name(self, model_name: str):
         self._default_model_name = model_name
 

zrb/task/llm/agent_runner.py

@@ -0,0 +1,53 @@
+from typing import Any
+
+from openai import APIError
+from pydantic_ai import Agent
+from pydantic_ai.messages import ModelMessagesTypeAdapter
+
+from zrb.context.any_context import AnyContext
+from zrb.task.llm.error import extract_api_error_details
+from zrb.task.llm.history import ListOfDict
+from zrb.task.llm.print_node import print_node
+
+
+async def run_agent_iteration(
+    ctx: AnyContext,
+    agent: Agent,
+    user_prompt: str,
+    history_list: ListOfDict,
+) -> Any:
+    """
+    Runs a single iteration of the agent execution loop.
+
+    Args:
+        ctx: The task context.
+        agent: The Pydantic AI agent instance.
+        user_prompt: The user's input prompt.
+        history_list: The current conversation history.
+
+    Returns:
+        The agent run result object.
+
+    Raises:
+        Exception: If any error occurs during agent execution.
+    """
+    async with agent.run_mcp_servers():
+        async with agent.iter(
+            user_prompt=user_prompt,
+            message_history=ModelMessagesTypeAdapter.validate_python(history_list),
+        ) as agent_run:
+            async for node in agent_run:
+                # 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)
+                except APIError as e:
+                    # Extract detailed error information from the response
+                    error_details = extract_api_error_details(e)
+                    ctx.log_error(f"API Error: {error_details}")
+                    raise
+                except Exception as e:
+                    ctx.log_error(f"Error processing node: {str(e)}")
+                    ctx.log_error(f"Error type: {type(e).__name__}")
+                    raise
+            return agent_run

zrb/task/llm/context_enricher.py

@@ -0,0 +1,86 @@
+import json
+import traceback
+from textwrap import dedent
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic_ai import Agent
+from pydantic_ai.models import Model
+from pydantic_ai.settings import ModelSettings
+
+from zrb.context.any_context import AnyContext
+from zrb.task.llm.agent_runner import run_agent_iteration
+from zrb.task.llm.history import ListOfDict
+
+
+# Configuration model for context enrichment
+class EnrichmentConfig(BaseModel):
+    model_config = {"arbitrary_types_allowed": True}
+    model: Model | str | None = None
+    settings: ModelSettings | None = None
+    prompt: str
+    retries: int = 1
+
+
+class EnrichmentResult(BaseModel):
+    response: dict[str, Any]  # or further decompose as needed
+
+
+async def enrich_context(
+    ctx: AnyContext,
+    config: EnrichmentConfig,
+    conversation_context: dict[str, Any],
+    history_list: ListOfDict,
+) -> dict[str, Any]:
+    """Runs an LLM call to extract key info and merge it into the context."""
+    ctx.log_info("Attempting to enrich conversation context...")
+    # Prepare context and history for the enrichment prompt
+    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 = dedent(
+            f"""
+            Analyze the following
+            # Current Context
+            {context_json}
+            # Conversation History
+            {history_json}
+        """
+        ).strip()
+    except Exception as e:
+        ctx.log_warning(f"Error formatting context/history for enrichment: {e}")
+        return conversation_context  # Return original context if formatting fails
+
+    enrichment_agent = Agent(
+        model=config.model,
+        # System prompt is part of the user prompt for this specific call
+        system_prompt=config.prompt,  # Use the main prompt as system prompt
+        tools=[],
+        mcp_servers=[],
+        model_settings=config.settings,
+        retries=config.retries,
+        result_type=EnrichmentResult,
+    )
+
+    try:
+        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
+        )
+        if enrichment_run and enrichment_run.result.data:
+            response = enrichment_run.result.data.response
+            if response:
+                conversation_context.update(response)
+                ctx.log_info("History summarized and added/updated in context.")
+                ctx.log_info(
+                    f"New conversation context: {json.dumps(conversation_context)}"
+                )
+        else:
+            ctx.log_warning("Context enrichment return no data")
+    except Exception as e:
+        ctx.log_warning(f"Error during context enrichment LLM call: {e}")
+        traceback.print_exc()
+    return conversation_context

zrb/task/llm/default_context.py

@@ -0,0 +1,44 @@
+import datetime
+import os
+import platform
+import re
+from typing import Any
+
+from zrb.util.file import read_dir, read_file
+
+
+def get_default_context(user_message: str) -> dict[str, Any]:
+    references = re.findall(r"@(\S+)", user_message)
+    current_references = []
+
+    for ref in references:
+        resource_path = os.path.abspath(os.path.expanduser(ref))
+        if os.path.isfile(resource_path):
+            content = read_file(resource_path)
+            current_references.append(
+                {
+                    "reference": ref,
+                    "name": resource_path,
+                    "type": "file",
+                    "content": content,
+                }
+            )
+        elif os.path.isdir(resource_path):
+            content = read_dir(resource_path)
+            current_references.append(
+                {
+                    "reference": ref,
+                    "name": resource_path,
+                    "type": "directory",
+                    "content": content,
+                }
+            )
+
+    return {
+        "current_time": datetime.datetime.now().isoformat(),
+        "current_working_directory": os.getcwd(),
+        "current_os": platform.system(),
+        "os_version": platform.version(),
+        "python_version": platform.python_version(),
+        "current_references": current_references,
+    }

zrb/task/llm/error.py

@@ -0,0 +1,77 @@
+import json
+from typing import Optional
+
+from openai import APIError
+from pydantic import BaseModel
+
+
+# Define a structured error model for tool execution failures
+class ToolExecutionError(BaseModel):
+    tool_name: str
+    error_type: str
+    message: str
+    details: Optional[str] = None
+
+
+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
+    if error.body:
+        try:
+            if isinstance(error.body, str):
+                body_json = json.loads(error.body)
+            elif isinstance(error.body, bytes):
+                body_json = json.loads(error.body.decode("utf-8"))
+            else:
+                body_json = error.body
+            # Extract error details from the JSON structure
+            if isinstance(body_json, dict):
+                if "error" in body_json:
+                    error_obj = body_json["error"]
+                    if isinstance(error_obj, dict):
+                        if "message" in error_obj:
+                            details += f"\nProvider message: {error_obj['message']}"
+                        if "code" in error_obj:
+                            details += f"\nError code: {error_obj['code']}"
+                        if "status" in error_obj:
+                            details += f"\nStatus: {error_obj['status']}"
+                # Check for metadata that might contain provider-specific information
+                if "metadata" in body_json and isinstance(body_json["metadata"], dict):
+                    metadata = body_json["metadata"]
+                    if "provider_name" in metadata:
+                        details += f"\nProvider: {metadata['provider_name']}"
+                    if "raw" in metadata:
+                        try:
+                            raw_json = json.loads(metadata["raw"])
+                            if "error" in raw_json and isinstance(
+                                raw_json["error"], dict
+                            ):
+                                raw_error = raw_json["error"]
+                                if "message" in raw_error:
+                                    details += (
+                                        f"\nRaw error message: {raw_error['message']}"
+                                    )
+                        except (KeyError, TypeError, ValueError):
+                            # If we can't parse the raw JSON, just include it as is
+                            details += f"\nRaw error data: {metadata['raw']}"
+        except json.JSONDecodeError:
+            # If we can't parse the JSON, include the raw body
+            details += f"\nRaw error body: {error.body}"
+        except Exception as e:
+            # Catch any other exceptions during parsing
+            details += f"\nError parsing error body: {str(e)}"
+    # Include request information if available
+    if hasattr(error, "request") and error.request:
+        if hasattr(error.request, "method") and hasattr(error.request, "url"):
+            details += f"\nRequest: {error.request.method} {error.request.url}"
+        # Include a truncated version of the request content if available
+        if hasattr(error.request, "content") and error.request.content:
+            content = error.request.content
+            if isinstance(content, bytes):
+                try:
+                    content = content.decode("utf-8")
+                except UnicodeDecodeError:
+                    content = str(content)
+            details += f"\nRequest content: {content}"
+    return details

zrb/task/llm/history.py

@@ -0,0 +1,92 @@
+import json
+import os
+from collections.abc import Callable
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+from zrb.context.any_context import AnyContext
+from zrb.util.file import read_file
+from zrb.util.run import run_async
+
+ListOfDict = list[dict[str, Any]]
+
+
+# Define the new ConversationHistoryData model
+class ConversationHistoryData(BaseModel):
+    context: dict[str, Any] = {}
+    history: ListOfDict = []
+
+    @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) and "history" in data:
+                # Standard format {'context': ..., 'history': ...}
+                # Ensure context exists, even if empty
+                data.setdefault("context", {})
+                return cls.model_validate(data)
+            elif isinstance(data, list):
+                # Handle 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': [...]}. "
+                    "Consider updating the source format."
+                )
+                return cls(history=data, context={})
+            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

zrb/task/llm/history_summarizer.py

@@ -0,0 +1,71 @@
+import json
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic_ai import Agent
+from pydantic_ai.models import Model
+from pydantic_ai.settings import ModelSettings
+
+from zrb.context.any_context import AnyContext
+from zrb.task.llm.agent_runner import run_agent_iteration
+from zrb.task.llm.history import ListOfDict
+
+
+# Configuration model for history summarization
+class SummarizationConfig(BaseModel):
+    model_config = {"arbitrary_types_allowed": True}
+    model: Model | str | None = None
+    settings: ModelSettings | None = None
+    prompt: str
+    retries: int = 1
+
+
+async def summarize_history(
+    ctx: AnyContext,
+    config: SummarizationConfig,
+    conversation_context: dict[str, Any],
+    history_list: ListOfDict,
+) -> dict[str, Any]:
+    """Runs an LLM call to summarize history and update the context."""
+    ctx.log_info("Attempting to summarize conversation history...")
+
+    summarization_agent = Agent(
+        model=config.model,
+        system_prompt=config.prompt,
+        tools=[],  # No tools needed for summarization
+        mcp_servers=[],
+        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 = (
+            f"# Current Context\n{context_json}\n\n"
+            f"# Conversation History to Summarize\n{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:
+        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
+        )
+        if summary_run and summary_run.result.data:
+            summary_text = str(summary_run.result.data)
+            # 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"Conversaion 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