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

zrb/config/llm_context/config.py

@@ -2,12 +2,124 @@ import os
 
 from zrb.config.config import CFG
 from zrb.config.llm_context.config_parser import markdown_to_dict
-from zrb.util.llm.prompt import demote_markdown_headers
+from zrb.config.llm_context.workflow import LLMWorkflow
+from zrb.util.markdown import demote_markdown_headers
 
 
 class LLMContextConfig:
     """High-level API for interacting with cascaded configurations."""
 
+    def write_note(
+        self,
+        content: str,
+        context_path: str | None = None,
+        cwd: str | None = None,
+    ):
+        """Writes content to a note block in the user's home configuration file."""
+        if cwd is None:
+            cwd = os.getcwd()
+        if context_path is None:
+            context_path = cwd
+        config_file = self._get_home_config_file()
+        sections = {}
+        if os.path.exists(config_file):
+            sections = self._parse_config(config_file)
+        abs_context_path = os.path.abspath(os.path.join(cwd, context_path))
+        found_key = None
+        for key in sections.keys():
+            if not key.startswith("Note:"):
+                continue
+            context_path_str = key[len("Note:") :].strip()
+            abs_key_path = self._normalize_context_path(
+                context_path_str,
+                os.path.dirname(config_file),
+            )
+            if abs_key_path == abs_context_path:
+                found_key = key
+                break
+        if found_key:
+            sections[found_key] = content
+        else:
+            config_dir = os.path.dirname(config_file)
+            formatted_path = self._format_context_path_for_writing(
+                abs_context_path,
+                config_dir,
+            )
+            new_key = f"Note: {formatted_path}"
+            sections[new_key] = content
+        # Serialize back to markdown
+        new_file_content = ""
+        for key, value in sections.items():
+            new_file_content += f"# {key}\n{demote_markdown_headers(value)}\n\n"
+        with open(config_file, "w") as f:
+            f.write(new_file_content)
+
+    def get_notes(self, cwd: str | None = None) -> dict[str, str]:
+        """Gathers all notes for a given path."""
+        if cwd is None:
+            cwd = os.getcwd()
+        config_file = self._get_home_config_file()
+        if not os.path.exists(config_file):
+            return {}
+        config_dir = os.path.dirname(config_file)
+        sections = self._parse_config(config_file)
+        notes: dict[str, str] = {}
+        for key, value in sections.items():
+            if key.lower().startswith("note:"):
+                context_path_str = key[len("note:") :].strip()
+                abs_context_path = self._normalize_context_path(
+                    context_path_str,
+                    config_dir,
+                )
+                # A context is relevant if its path is an ancestor of cwd
+                if os.path.commonpath([cwd, abs_context_path]) == abs_context_path:
+                    notes[abs_context_path] = value
+        return notes
+
+    def get_workflows(self, cwd: str | None = None) -> dict[str, LLMWorkflow]:
+        """Gathers all relevant workflows for a given path."""
+        if cwd is None:
+            cwd = os.getcwd()
+        all_sections = self._get_all_sections(cwd)
+        workflows: dict[str, LLMWorkflow] = {}
+        # Iterate from closest to farthest
+        for config_dir, sections in all_sections:
+            for key, value in sections.items():
+                if key.lower().startswith("workflow:"):
+                    workflow_name = key[len("workflow:") :].strip().lower()
+                    # First one found wins
+                    if workflow_name not in workflows:
+                        workflows[workflow_name] = LLMWorkflow(
+                            name=workflow_name,
+                            content=value,
+                            path=config_dir,
+                        )
+        return workflows
+
+    def _format_context_path_for_writing(
+        self,
+        path_to_write: str,
+        relative_to_dir: str,
+    ) -> str:
+        """Formats a path for writing into a context file key."""
+        home_dir = os.path.expanduser("~")
+        abs_path_to_write = os.path.abspath(
+            os.path.join(relative_to_dir, path_to_write)
+        )
+        abs_relative_to_dir = os.path.abspath(relative_to_dir)
+        # Rule 1: Inside relative_to_dir
+        if abs_path_to_write.startswith(abs_relative_to_dir):
+            if abs_path_to_write == abs_relative_to_dir:
+                return "."
+            return os.path.relpath(abs_path_to_write, abs_relative_to_dir)
+        # Rule 2: Inside Home
+        if abs_path_to_write.startswith(home_dir):
+            if abs_path_to_write == home_dir:
+                return "~"
+            return os.path.join("~", os.path.relpath(abs_path_to_write, home_dir))
+        # Rule 3: Absolute
+        return abs_path_to_write
+
     def _find_config_files(self, cwd: str) -> list[str]:
         configs = []
         current_dir = cwd
@@ -21,6 +133,10 @@ class LLMContextConfig:
             current_dir = os.path.dirname(current_dir)
         return configs
 
+    def _get_home_config_file(self) -> str:
+        home_dir = os.path.expanduser("~")
+        return os.path.join(home_dir, CFG.LLM_CONTEXT_FILE)
+
     def _parse_config(self, file_path: str) -> dict[str, str]:
         with open(file_path, "r") as f:
             content = f.read()
@@ -35,95 +151,16 @@ class LLMContextConfig:
             all_sections.append((config_dir, sections))
         return all_sections
 
-    def get_contexts(self, cwd: str | None = None) -> dict[str, str]:
-        """Gathers all relevant contexts for a given path."""
-        if cwd is None:
-            cwd = os.getcwd()
-        all_sections = self._get_all_sections(cwd)
-        contexts: dict[str, str] = {}
-        for config_dir, sections in reversed(all_sections):
-            for key, value in sections.items():
-                if key.startswith("Context:"):
-                    context_path = key[len("Context:") :].strip()
-                    if context_path == ".":
-                        context_path = config_dir
-                    elif not os.path.isabs(context_path):
-                        context_path = os.path.abspath(
-                            os.path.join(config_dir, context_path)
-                        )
-                    if os.path.isabs(context_path) or cwd.startswith(context_path):
-                        contexts[context_path] = value
-        return contexts
-
-    def get_workflows(self, cwd: str | None = None) -> dict[str, str]:
-        """Gathers all relevant workflows for a given path."""
-        if cwd is None:
-            cwd = os.getcwd()
-        all_sections = self._get_all_sections(cwd)
-        workflows: dict[str, str] = {}
-        for _, sections in reversed(all_sections):
-            for key, value in sections.items():
-                if key.startswith("Workflow:"):
-                    workflow_name = key[len("Workflow:") :].strip()
-                    workflow_name = key.replace("Workflow:", "").lower().strip()
-                    workflows[workflow_name] = value
-        return workflows
-
-    def write_context(
-        self, content: str, context_path: str | None = None, cwd: str | None = None
-    ):
-        """Writes content to a context block in the nearest configuration file."""
-        if cwd is None:
-            cwd = os.getcwd()
-        if context_path is None:
-            context_path = cwd
-
-        config_files = self._find_config_files(cwd)
-        if config_files:
-            config_file = config_files[0]  # Closest config file
-        else:
-            config_file = os.path.join(cwd, CFG.LLM_CONTEXT_FILE)
-
-        sections = {}
-        if os.path.exists(config_file):
-            sections = self._parse_config(config_file)
-
-        # Determine the section key
-        section_key_path = context_path
-        if not os.path.isabs(context_path):
-            config_dir = os.path.dirname(config_file)
-            section_key_path = os.path.abspath(os.path.join(config_dir, context_path))
-
-        # Find existing key
-        found_key = ""
-        for key in sections.keys():
-            if not key.startswith("Context:"):
-                continue
-            key_path = key.replace("Context:", "").strip()
-            if key_path == ".":
-                key_path = os.path.dirname(config_file)
-            elif not os.path.isabs(key_path):
-                key_path = os.path.abspath(
-                    os.path.join(os.path.dirname(config_file), key_path)
-                )
-            if key_path == section_key_path:
-                found_key = key
-                break
-
-        if found_key != "":
-            sections[found_key] = content
-        else:
-            # Add new entry
-            new_key = f"Context: {context_path}"
-            sections[new_key] = content
-
-        # Serialize back to markdown
-        new_file_content = ""
-        for key, value in sections.items():
-            new_file_content += f"# {key}\n{demote_markdown_headers(value)}\n\n"
-
-        with open(config_file, "w") as f:
-            f.write(new_file_content)
+    def _normalize_context_path(
+        self,
+        path_str: str,
+        relative_to_dir: str,
+    ) -> str:
+        """Normalizes a context path string to an absolute path."""
+        expanded_path = os.path.expanduser(path_str)
+        if os.path.isabs(expanded_path):
+            return os.path.abspath(expanded_path)
+        return os.path.abspath(os.path.join(relative_to_dir, expanded_path))
 
 
 llm_context_config = LLMContextConfig()

zrb/config/llm_context/config_parser.py

@@ -1,6 +1,6 @@
 import re
 
-from zrb.util.llm.prompt import promote_markdown_headers
+from zrb.util.markdown import promote_markdown_headers
 
 
 def markdown_to_dict(markdown: str) -> dict[str, str]:
@@ -8,21 +8,17 @@ def markdown_to_dict(markdown: str) -> dict[str, str]:
     current_title = ""
     current_content: list[str] = []
     fence_stack: list[str] = []
-
     fence_pattern = re.compile(r"^([`~]{3,})(.*)$")
     h1_pattern = re.compile(r"^# (.+)$")
-
     for line in markdown.splitlines():
         # Detect code fence open/close
         fence_match = fence_pattern.match(line.strip())
-
         if fence_match:
             fence = fence_match.group(1)
             if fence_stack and fence_stack[-1] == fence:
                 fence_stack.pop()  # close current fence
             else:
                 fence_stack.append(fence)  # open new fence
-
         # Only parse H1 when not inside a code fence
         if not fence_stack:
             h1_match = h1_pattern.match(line)
@@ -34,9 +30,7 @@ def markdown_to_dict(markdown: str) -> dict[str, str]:
                 current_title = h1_match.group(1).strip()
                 current_content = []
                 continue
-
         current_content.append(line)
-
     # Save final section
     if current_title:
         sections[current_title] = "\n".join(current_content).strip()

zrb/config/llm_context/workflow.py

@@ -0,0 +1,81 @@
+class LLMWorkflow:
+    def __init__(
+        self, name: str, path: str, content: str, description: str | None = None
+    ):
+        self._name = name
+        self._path = path
+
+        # Extract YAML metadata and clean content
+        (
+            extracted_description,
+            cleaned_content,
+        ) = self._extract_yaml_metadata_and_clean_content(content)
+        self._content = cleaned_content
+
+        # Use provided description or extracted one
+        self._description = (
+            description if description is not None else extracted_description
+        )
+
+    def _extract_yaml_metadata_and_clean_content(
+        self, content: str
+    ) -> tuple[str | None, str]:
+        """Extract YAML metadata and clean content.
+
+        Looks for YAML metadata between --- lines, extracts the 'description' field,
+        and returns the content without the YAML metadata.
+        """
+        import re
+
+        import yaml
+
+        # Pattern to match YAML metadata between --- delimiters
+        yaml_pattern = r"^---\s*\n(.*?)\n---\s*\n"
+        match = re.search(yaml_pattern, content, re.DOTALL | re.MULTILINE)
+
+        if match:
+            yaml_content = match.group(1)
+            try:
+                metadata = yaml.safe_load(yaml_content)
+                description = (
+                    metadata.get("description") if isinstance(metadata, dict) else None
+                )
+                # Remove the YAML metadata from content
+                cleaned_content = re.sub(
+                    yaml_pattern, "", content, count=1, flags=re.DOTALL | re.MULTILINE
+                )
+                return description, cleaned_content.strip()
+            except yaml.YAMLError:
+                # If YAML parsing fails, return original content
+                pass
+
+        # No YAML metadata found, return original content
+        return None, content
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def path(self) -> str:
+        return self._path
+
+    @property
+    def content(self) -> str:
+        return self._content
+
+    @property
+    def description(self) -> str:
+        if self._description is not None:
+            return self._description
+        if len(self._content) > 1000:
+            non_empty_lines = [
+                line for line in self._content.split("\n") if line.strip() != ""
+            ]
+            first_non_empty_line = (
+                non_empty_lines[0] if len(non_empty_lines) > 0 else ""
+            )
+            if len(first_non_empty_line) > 200:
+                return first_non_empty_line[:200] + "... (more)"
+            return first_non_empty_line
+        return self._content

zrb/config/llm_rate_limitter.py

@@ -1,30 +1,13 @@
 import asyncio
+import json
 import time
 from collections import deque
-from typing import Callable
-
-import tiktoken
+from typing import Any, Callable
 
 from zrb.config.config import CFG
 
 
-def _estimate_token(text: str) -> int:
-    """
-    Estimates the number of tokens in a given text.
-    Tries to use the 'gpt-4o' model's tokenizer for an accurate count.
-    If the tokenizer is unavailable (e.g., due to network issues),
-    it falls back to a heuristic of 4 characters per token.
-    """
-    try:
-        # Primary method: Use tiktoken for an accurate count
-        enc = tiktoken.encoding_for_model("gpt-4o")
-        return len(enc.encode(text))
-    except Exception:
-        # Fallback method: Heuristic (4 characters per token)
-        return len(text) // 4
-
-
-class LLMRateLimiter:
+class LLMRateLimitter:
     """
     Helper class to enforce LLM API rate limits and throttling.
     Tracks requests and tokens in a rolling 60-second window.
@@ -35,14 +18,18 @@ class LLMRateLimiter:
         max_requests_per_minute: int | None = None,
         max_tokens_per_minute: int | None = None,
         max_tokens_per_request: int | None = None,
+        max_tokens_per_tool_call_result: int | None = None,
         throttle_sleep: float | None = None,
-        token_counter_fn: Callable[[str], int] | None = None,
+        use_tiktoken: bool | None = None,
+        tiktoken_encoding_name: str | None = None,
     ):
         self._max_requests_per_minute = max_requests_per_minute
         self._max_tokens_per_minute = max_tokens_per_minute
         self._max_tokens_per_request = max_tokens_per_request
+        self._max_tokens_per_tool_call_result = max_tokens_per_tool_call_result
         self._throttle_sleep = throttle_sleep
-        self._token_counter_fn = token_counter_fn
+        self._use_tiktoken = use_tiktoken
+        self._tiktoken_encoding_name = tiktoken_encoding_name
         self.request_times = deque()
         self.token_times = deque()
 
@@ -64,6 +51,12 @@ class LLMRateLimiter:
             return self._max_tokens_per_request
         return CFG.LLM_MAX_TOKENS_PER_REQUEST
 
+    @property
+    def max_tokens_per_tool_call_result(self) -> int:
+        if self._max_tokens_per_tool_call_result is not None:
+            return self._max_tokens_per_tool_call_result
+        return CFG.LLM_MAX_TOKENS_PER_TOOL_CALL_RESULT
+
     @property
     def throttle_sleep(self) -> float:
         if self._throttle_sleep is not None:
@@ -71,10 +64,16 @@ class LLMRateLimiter:
         return CFG.LLM_THROTTLE_SLEEP
 
     @property
-    def count_token(self) -> Callable[[str], int]:
-        if self._token_counter_fn is not None:
-            return self._token_counter_fn
-        return _estimate_token
+    def use_tiktoken(self) -> bool:
+        if self._use_tiktoken is not None:
+            return self._use_tiktoken
+        return CFG.USE_TIKTOKEN
+
+    @property
+    def tiktoken_encoding_name(self) -> str:
+        if self._tiktoken_encoding_name is not None:
+            return self._tiktoken_encoding_name
+        return CFG.TIKTOKEN_ENCODING_NAME
 
     def set_max_requests_per_minute(self, value: int):
         self._max_requests_per_minute = value
@@ -85,29 +84,56 @@ class LLMRateLimiter:
     def set_max_tokens_per_request(self, value: int):
         self._max_tokens_per_request = value
 
+    def set_max_tokens_per_tool_call_result(self, value: int):
+        self._max_tokens_per_tool_call_result = value
+
     def set_throttle_sleep(self, value: float):
         self._throttle_sleep = value
 
-    def set_token_counter_fn(self, fn: Callable[[str], int]):
-        self._token_counter_fn = fn
-
-    def clip_prompt(self, prompt: str, limit: int) -> str:
-        token_count = self.count_token(prompt)
-        if token_count <= limit:
-            return prompt
-        while token_count > limit:
-            prompt_parts = prompt.split(" ")
-            last_part_index = len(prompt_parts) - 2
-            clipped_prompt = " ".join(prompt_parts[:last_part_index])
-            clipped_prompt += "(Content clipped...)"
-            token_count = self.count_token(clipped_prompt)
-            if token_count < limit:
-                return clipped_prompt
-        return prompt[:limit]
-
-    async def throttle(self, prompt: str):
+    def count_token(self, prompt: Any) -> int:
+        str_prompt = self._prompt_to_str(prompt)
+        if not self.use_tiktoken:
+            return self._fallback_count_token(str_prompt)
+        try:
+            import tiktoken
+
+            enc = tiktoken.get_encoding(self.tiktoken_encoding_name)
+            return len(enc.encode(str_prompt))
+        except Exception:
+            return self._fallback_count_token(str_prompt)
+
+    def _fallback_count_token(self, str_prompt: str) -> int:
+        return len(str_prompt) // 4
+
+    def clip_prompt(self, prompt: Any, limit: int) -> str:
+        str_prompt = self._prompt_to_str(prompt)
+        if not self.use_tiktoken:
+            return self._fallback_clip_prompt(str_prompt, limit)
+        try:
+            import tiktoken
+
+            enc = tiktoken.get_encoding(self.tiktoken_encoding_name)
+            tokens = enc.encode(str_prompt)
+            if len(tokens) <= limit:
+                return str_prompt
+            truncated = tokens[: limit - 3]
+            clipped_text = enc.decode(truncated)
+            return clipped_text + "..."
+        except Exception:
+            return self._fallback_clip_prompt(str_prompt, limit)
+
+    def _fallback_clip_prompt(self, str_prompt: str, limit: int) -> str:
+        char_limit = limit * 4 if limit * 4 <= 10 else limit * 4 - 10
+        return str_prompt[:char_limit] + "..."
+
+    async def throttle(
+        self,
+        prompt: Any,
+        throttle_notif_callback: Callable[[str], Any] | None = None,
+    ):
         now = time.time()
-        tokens = self.count_token(prompt)
+        str_prompt = self._prompt_to_str(prompt)
+        tokens = self.count_token(str_prompt)
         # Clean up old entries
         while self.request_times and now - self.request_times[0] > 60:
             self.request_times.popleft()
@@ -116,13 +142,34 @@ class LLMRateLimiter:
         # Check per-request token limit
         if tokens > self.max_tokens_per_request:
             raise ValueError(
-                f"Request exceeds max_tokens_per_request ({self.max_tokens_per_request})."
+                (
+                    "Request exceeds max_tokens_per_request "
+                    f"({tokens} > {self.max_tokens_per_request})."
+                )
+            )
+        if tokens > self.max_tokens_per_minute:
+            raise ValueError(
+                (
+                    "Request exceeds max_tokens_per_minute "
+                    f"({tokens} > {self.max_tokens_per_minute})."
+                )
             )
         # Wait if over per-minute request or token limit
         while (
             len(self.request_times) >= self.max_requests_per_minute
             or sum(t for _, t in self.token_times) + tokens > self.max_tokens_per_minute
         ):
+            if throttle_notif_callback is not None:
+                if len(self.request_times) >= self.max_requests_per_minute:
+                    rpm = len(self.request_times)
+                    throttle_notif_callback(
+                        f"Max request per minute exceeded: {rpm} of {self.max_requests_per_minute}"
+                    )
+                else:
+                    tpm = sum(t for _, t in self.token_times) + tokens
+                    throttle_notif_callback(
+                        f"Max token per minute exceeded: {tpm} of {self.max_tokens_per_minute}"
+                    )
             await asyncio.sleep(self.throttle_sleep)
             now = time.time()
             while self.request_times and now - self.request_times[0] > 60:
@@ -133,5 +180,11 @@ class LLMRateLimiter:
         self.request_times.append(now)
         self.token_times.append((now, tokens))
 
+    def _prompt_to_str(self, prompt: Any) -> str:
+        try:
+            return json.dumps(prompt)
+        except Exception:
+            return f"{prompt}"
+
 
-llm_rate_limitter = LLMRateLimiter()
+llm_rate_limitter = LLMRateLimitter()

zrb/context/any_shared_context.py

@@ -29,26 +29,32 @@ class AnySharedContext(ABC):
         pass
 
     @property
+    @abstractmethod
     def input(self) -> DotDict:
         pass
 
     @property
+    @abstractmethod
     def env(self) -> DotDict:
         pass
 
     @property
+    @abstractmethod
     def args(self) -> list[Any]:
         pass
 
     @property
-    def xcom(self) -> DotDict[str, Xcom]:
+    @abstractmethod
+    def xcom(self) -> DotDict:
         pass
 
     @property
+    @abstractmethod
     def shared_log(self) -> list[str]:
         pass
 
     @property
+    @abstractmethod
     def session(self) -> any_session.AnySession | None:
         pass
 

zrb/context/context.py

@@ -63,7 +63,7 @@ class Context(AnyContext):
 
     @property
     def session(self) -> AnySession | None:
-        return self._shared_ctx._session
+        return self._shared_ctx.session
 
     def update_task_env(self, task_env: dict[str, str]):
         self._env.update(task_env)
@@ -119,7 +119,13 @@ class Context(AnyContext):
             return
         color = self._color
         icon = self._icon
-        max_name_length = max(len(name) + len(icon) for name in self.session.task_names)
+        # Handle case where session is None (e.g., in tests)
+        if self.session is None:
+            max_name_length = len(self._task_name) + len(icon)
+        else:
+            max_name_length = max(
+                len(name) + len(icon) for name in self.session.task_names
+            )
         styled_task_name = f"{icon} {self._task_name}"
         padded_styled_task_name = styled_task_name.rjust(max_name_length + 1)
         if self._attempt == 0:

zrb/context/shared_context.py

@@ -40,11 +40,7 @@ class SharedContext(AnySharedContext):
 
     def __repr__(self):
         class_name = self.__class__.__name__
-        input = self._input
-        args = self._args
-        env = self._env
-        xcom = self._xcom
-        return f"<{class_name} input={input} args={args} xcom={xcom} env={env}>"
+        return f"<{class_name}>"
 
     @property
     def is_web_mode(self) -> bool:
@@ -70,7 +66,7 @@ class SharedContext(AnySharedContext):
         return self._args
 
     @property
-    def xcom(self) -> DotDict[str, Xcom]:
+    def xcom(self) -> DotDict:
         return self._xcom
 
     @property
@@ -85,7 +81,7 @@ class SharedContext(AnySharedContext):
         self._log.append(message)
         session = self.session
         if session is not None:
-            session_parent: AnySession = session.parent
+            session_parent: AnySession | None = session.parent
             if session_parent is not None:
                 session_parent.shared_ctx.append_to_shared_log(message)
 

zrb/group/any_group.py

@@ -35,11 +35,11 @@ class AnyGroup(ABC):
         pass
 
     @abstractmethod
-    def add_group(self, group: "AnyGroup | str") -> "AnyGroup":
+    def add_group(self, group: "AnyGroup", alias: str | None = None) -> "AnyGroup":
         pass
 
     @abstractmethod
-    def add_task(self, task: AnyTask, alias: str | None = None) -> AnyTask:
+    def add_task(self, task: "AnyTask", alias: str | None = None) -> "AnyTask":
         pass
 
     @abstractmethod
@@ -55,5 +55,5 @@ class AnyGroup(ABC):
         pass
 
     @abstractmethod
-    def get_group_by_alias(self, name: str) -> "AnyGroup | None":
+    def get_group_by_alias(self, alias: str) -> "AnyGroup | None":
         pass