zrb 1.13.1__py3-none-any.whl → 1.21.17__py3-none-any.whl

zrb/config/llm_config.py

@@ -2,8 +2,6 @@ import os
 from typing import TYPE_CHECKING, Any, Callable
 
 from zrb.config.config import CFG
-from zrb.config.llm_context.config import llm_context_config
-from zrb.util.llm.prompt import make_prompt_section
 
 if TYPE_CHECKING:
     from pydantic_ai.models import Model
@@ -15,8 +13,11 @@ class LLMConfig:
     def __init__(
         self,
         default_model_name: str | None = None,
-        default_base_url: str | None = None,
-        default_api_key: str | None = None,
+        default_model_base_url: str | None = None,
+        default_model_api_key: str | None = None,
+        default_small_model_name: str | None = None,
+        default_small_model_base_url: str | None = None,
+        default_small_model_api_key: str | None = None,
         default_persona: str | None = None,
         default_system_prompt: str | None = None,
         default_interactive_system_prompt: str | None = None,
@@ -24,15 +25,25 @@ class LLMConfig:
         default_summarization_prompt: str | None = None,
         default_summarize_history: bool | None = None,
         default_history_summarization_token_threshold: int | None = None,
-        default_modes: list[str] | None = None,
+        default_workflows: list[str] | None = None,
         default_model: "Model | None" = None,
         default_model_settings: "ModelSettings | None" = None,
         default_model_provider: "Provider | None" = None,
+        default_small_model: "Model | None" = None,
+        default_small_model_settings: "ModelSettings | None" = None,
+        default_small_model_provider: "Provider | None" = None,
+        default_yolo_mode: bool | list[str] | None = None,
+        default_current_weather_tool: Callable | None = None,
+        default_current_location_tool: Callable | None = None,
+        default_search_internet_tool: Callable | None = None,
     ):
         self.__internal_default_prompt: dict[str, str] = {}
         self._default_model_name = default_model_name
-        self._default_model_base_url = default_base_url
-        self._default_model_api_key = default_api_key
+        self._default_model_base_url = default_model_base_url
+        self._default_model_api_key = default_model_api_key
+        self._default_small_model_name = default_small_model_name
+        self._default_small_model_base_url = default_small_model_base_url
+        self._default_small_model_api_key = default_small_model_api_key
         self._default_persona = default_persona
         self._default_system_prompt = default_system_prompt
         self._default_interactive_system_prompt = default_interactive_system_prompt
@@ -42,10 +53,17 @@ class LLMConfig:
         self._default_history_summarization_token_threshold = (
             default_history_summarization_token_threshold
         )
-        self._default_modes = default_modes
+        self._default_workflows = default_workflows
         self._default_model = default_model
         self._default_model_settings = default_model_settings
         self._default_model_provider = default_model_provider
+        self._default_small_model = default_small_model
+        self._default_small_model_settings = default_small_model_settings
+        self._default_small_model_provider = default_small_model_provider
+        self._default_yolo_mode = default_yolo_mode
+        self._default_current_weather_tool = default_current_weather_tool
+        self._default_current_location_tool = default_current_location_tool
+        self._default_search_internet_tool = default_search_internet_tool
 
     def _get_internal_default_prompt(self, name: str) -> str:
         if name not in self.__internal_default_prompt:
@@ -100,6 +118,54 @@ class LLMConfig:
             base_url=self.default_model_base_url, api_key=self.default_model_api_key
         )
 
+    @property
+    def default_small_model_name(self) -> str | None:
+        return self._get_property(
+            self._default_small_model_name,
+            CFG.LLM_MODEL_SMALL,
+            lambda: self.default_model_name,
+        )
+
+    @property
+    def default_small_model_base_url(self) -> str | None:
+        return self._get_property(
+            self._default_small_model_base_url,
+            CFG.LLM_BASE_URL_SMALL,
+            lambda: self.default_model_base_url,
+        )
+
+    @property
+    def default_small_model_api_key(self) -> str | None:
+        return self._get_property(
+            self._default_small_model_api_key,
+            CFG.LLM_API_KEY_SMALL,
+            lambda: self.default_model_api_key,
+        )
+
+    @property
+    def default_small_model_settings(self) -> "ModelSettings | None":
+        return self._get_property(
+            self._default_small_model_settings,
+            None,
+            lambda: self.default_model_settings,
+        )
+
+    @property
+    def default_small_model_provider(self) -> "Provider | str":
+        if self._default_small_model_provider is not None:
+            return self._default_small_model_provider
+        if (
+            self.default_small_model_base_url is None
+            and self.default_small_model_api_key is None
+        ):
+            return self.default_model_provider
+        from pydantic_ai.providers.openai import OpenAIProvider
+
+        return OpenAIProvider(
+            base_url=self.default_small_model_base_url,
+            api_key=self.default_small_model_api_key,
+        )
+
     @property
     def default_system_prompt(self) -> str:
         return self._get_property(
@@ -125,9 +191,9 @@ class LLMConfig:
         )
 
     @property
-    def default_modes(self) -> list[str]:
+    def default_workflows(self) -> list[str]:
         return self._get_property(
-            self._default_modes, CFG.LLM_MODES, lambda: ["coding"]
+            self._default_workflows, CFG.LLM_WORKFLOWS, lambda: []
         )
 
     @property
@@ -147,19 +213,28 @@ class LLMConfig:
         )
 
     @property
-    def default_model(self) -> "Model | str | None":
+    def default_model(self) -> "Model | str":
         if self._default_model is not None:
             return self._default_model
         model_name = self.default_model_name
         if model_name is None:
-            return None
-        from pydantic_ai.models.openai import OpenAIModel
+            return "openai:gpt-4o"
+        from pydantic_ai.models.openai import OpenAIChatModel
 
-        return OpenAIModel(
+        return OpenAIChatModel(
             model_name=model_name,
             provider=self.default_model_provider,
         )
 
+    @property
+    def default_small_model(self) -> "Model | str":
+        if self._default_small_model is not None:
+            return self._default_small_model
+        model_name = self.default_small_model_name
+        if model_name is None:
+            return "openai:gpt-4o"
+        return self.default_model
+
     @property
     def default_summarize_history(self) -> bool:
         return self._get_property(
@@ -174,6 +249,24 @@ class LLMConfig:
             lambda: 1000,
         )
 
+    @property
+    def default_yolo_mode(self) -> bool | list[str]:
+        return self._get_property(
+            self._default_yolo_mode, CFG.LLM_YOLO_MODE, lambda: False
+        )
+
+    @property
+    def default_current_weather_tool(self) -> Callable | None:
+        return self._default_current_weather_tool
+
+    @property
+    def default_current_location_tool(self) -> Callable | None:
+        return self._default_current_location_tool
+
+    @property
+    def default_search_internet_tool(self) -> Callable | None:
+        return self._default_search_internet_tool
+
     def set_default_persona(self, persona: str):
         self._default_persona = persona
 
@@ -186,18 +279,18 @@ class LLMConfig:
     def set_default_special_instruction_prompt(self, special_instruction_prompt: str):
         self._default_special_instruction_prompt = special_instruction_prompt
 
-    def set_default_modes(self, modes: list[str]):
-        self._default_modes = modes
+    def set_default_workflows(self, workflows: list[str]):
+        self._default_workflows = workflows
 
-    def add_default_mode(self, mode: str):
-        if self._default_modes is None:
-            self._default_modes = []
-        self._default_modes.append(mode)
+    def add_default_workflow(self, workflow: str):
+        if self._default_workflows is None:
+            self._default_workflows = []
+        self._default_workflows.append(workflow)
 
-    def remove_default_mode(self, mode: str):
-        if self._default_modes is None:
-            self._default_modes = []
-        self._default_modes.remove(mode)
+    def remove_default_workflow(self, workflow: str):
+        if self._default_workflows is None:
+            self._default_workflows = []
+        self._default_workflows.remove(workflow)
 
     def set_default_summarization_prompt(self, summarization_prompt: str):
         self._default_summarization_prompt = summarization_prompt
@@ -230,5 +323,17 @@ class LLMConfig:
     def set_default_model_settings(self, model_settings: "ModelSettings"):
         self._default_model_settings = model_settings
 
+    def set_default_yolo_mode(self, yolo_mode: bool | list[str]):
+        self._default_yolo_mode = yolo_mode
+
+    def set_default_current_weather_tool(self, tool: Callable):
+        self._default_current_weather_tool = tool
+
+    def set_default_current_location_tool(self, tool: Callable):
+        self._default_current_location_tool = tool
+
+    def set_default_search_internet_tool(self, tool: Callable):
+        self._default_search_internet_tool = tool
+
 
 llm_config = LLMConfig()

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.replace("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.replace("Workflow:", "").strip()
-                    if workflow_name not in workflows:
-                        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