zrb 1.15.3__py3-none-any.whl → 2.0.0a4__py3-none-any.whl

zrb/config/llm_config.py

@@ -1,243 +0,0 @@
-import os
-from typing import TYPE_CHECKING, Any, Callable
-
-from zrb.config.config import CFG
-
-if TYPE_CHECKING:
-    from pydantic_ai.models import Model
-    from pydantic_ai.providers import Provider
-    from pydantic_ai.settings import ModelSettings
-
-
-class LLMConfig:
-    def __init__(
-        self,
-        default_model_name: str | None = None,
-        default_base_url: str | None = None,
-        default_api_key: str | None = None,
-        default_persona: str | None = None,
-        default_system_prompt: str | None = None,
-        default_interactive_system_prompt: str | None = None,
-        default_special_instruction_prompt: str | None = None,
-        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_model: "Model | None" = None,
-        default_model_settings: "ModelSettings | None" = None,
-        default_model_provider: "Provider | None" = None,
-        default_yolo_mode: bool | 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_persona = default_persona
-        self._default_system_prompt = default_system_prompt
-        self._default_interactive_system_prompt = default_interactive_system_prompt
-        self._default_special_instruction_prompt = default_special_instruction_prompt
-        self._default_summarization_prompt = default_summarization_prompt
-        self._default_summarize_history = default_summarize_history
-        self._default_history_summarization_token_threshold = (
-            default_history_summarization_token_threshold
-        )
-        self._default_modes = default_modes
-        self._default_model = default_model
-        self._default_model_settings = default_model_settings
-        self._default_model_provider = default_model_provider
-        self._default_yolo_mode = default_yolo_mode
-
-    def _get_internal_default_prompt(self, name: str) -> str:
-        if name not in self.__internal_default_prompt:
-            file_path = os.path.join(
-                os.path.dirname(__file__), "default_prompt", f"{name}.md"
-            )
-            with open(file_path, "r") as f:
-                self.__internal_default_prompt[name] = f.read().strip()
-        return self.__internal_default_prompt[name]
-
-    def _get_property(
-        self,
-        instance_var: Any,
-        config_var: Any,
-        default_func: Callable[[], Any],
-    ) -> Any:
-        if instance_var is not None:
-            return instance_var
-        if config_var is not None:
-            return config_var
-        return default_func()
-
-    @property
-    def default_model_name(self) -> str | None:
-        return self._get_property(self._default_model_name, CFG.LLM_MODEL, lambda: None)
-
-    @property
-    def default_model_base_url(self) -> str | None:
-        return self._get_property(
-            self._default_model_base_url, CFG.LLM_BASE_URL, lambda: None
-        )
-
-    @property
-    def default_model_api_key(self) -> str | None:
-        return self._get_property(
-            self._default_model_api_key, CFG.LLM_API_KEY, lambda: None
-        )
-
-    @property
-    def default_model_settings(self) -> "ModelSettings | None":
-        return self._get_property(self._default_model_settings, None, lambda: None)
-
-    @property
-    def default_model_provider(self) -> "Provider | str":
-        if self._default_model_provider is not None:
-            return self._default_model_provider
-        if self.default_model_base_url is None and self.default_model_api_key is None:
-            return "openai"
-        from pydantic_ai.providers.openai import OpenAIProvider
-
-        return OpenAIProvider(
-            base_url=self.default_model_base_url, api_key=self.default_model_api_key
-        )
-
-    @property
-    def default_system_prompt(self) -> str:
-        return self._get_property(
-            self._default_system_prompt,
-            CFG.LLM_SYSTEM_PROMPT,
-            lambda: self._get_internal_default_prompt("system_prompt"),
-        )
-
-    @property
-    def default_interactive_system_prompt(self) -> str:
-        return self._get_property(
-            self._default_interactive_system_prompt,
-            CFG.LLM_INTERACTIVE_SYSTEM_PROMPT,
-            lambda: self._get_internal_default_prompt("interactive_system_prompt"),
-        )
-
-    @property
-    def default_persona(self) -> str:
-        return self._get_property(
-            self._default_persona,
-            CFG.LLM_PERSONA,
-            lambda: self._get_internal_default_prompt("persona"),
-        )
-
-    @property
-    def default_modes(self) -> list[str]:
-        return self._get_property(
-            self._default_modes, CFG.LLM_MODES, lambda: ["coding"]
-        )
-
-    @property
-    def default_special_instruction_prompt(self) -> str:
-        return self._get_property(
-            self._default_special_instruction_prompt,
-            CFG.LLM_SPECIAL_INSTRUCTION_PROMPT,
-            lambda: "",
-        )
-
-    @property
-    def default_summarization_prompt(self) -> str:
-        return self._get_property(
-            self._default_summarization_prompt,
-            CFG.LLM_SUMMARIZATION_PROMPT,
-            lambda: self._get_internal_default_prompt("summarization_prompt"),
-        )
-
-    @property
-    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 "openai:gpt-4o"
-        from pydantic_ai.models.openai import OpenAIModel
-
-        return OpenAIModel(
-            model_name=model_name,
-            provider=self.default_model_provider,
-        )
-
-    @property
-    def default_summarize_history(self) -> bool:
-        return self._get_property(
-            self._default_summarize_history, CFG.LLM_SUMMARIZE_HISTORY, lambda: False
-        )
-
-    @property
-    def default_history_summarization_token_threshold(self) -> int:
-        return self._get_property(
-            self._default_history_summarization_token_threshold,
-            CFG.LLM_HISTORY_SUMMARIZATION_TOKEN_THRESHOLD,
-            lambda: 1000,
-        )
-
-    @property
-    def default_yolo_mode(self) -> bool:
-        return self._get_property(
-            self._default_yolo_mode, CFG.LLM_YOLO_MODE, lambda: False
-        )
-
-    def set_default_persona(self, persona: str):
-        self._default_persona = persona
-
-    def set_default_system_prompt(self, system_prompt: str):
-        self._default_system_prompt = system_prompt
-
-    def set_default_interactive_system_prompt(self, interactive_system_prompt: str):
-        self._default_interactive_system_prompt = interactive_system_prompt
-
-    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 add_default_mode(self, mode: str):
-        if self._default_modes is None:
-            self._default_modes = []
-        self._default_modes.append(mode)
-
-    def remove_default_mode(self, mode: str):
-        if self._default_modes is None:
-            self._default_modes = []
-        self._default_modes.remove(mode)
-
-    def set_default_summarization_prompt(self, summarization_prompt: str):
-        self._default_summarization_prompt = summarization_prompt
-
-    def set_default_model_name(self, model_name: str):
-        self._default_model_name = model_name
-
-    def set_default_model_api_key(self, model_api_key: str):
-        self._default_model_api_key = model_api_key
-
-    def set_default_model_base_url(self, model_base_url: str):
-        self._default_model_base_url = model_base_url
-
-    def set_default_model_provider(self, provider: "Provider | str"):
-        self._default_model_provider = provider
-
-    def set_default_model(self, model: "Model | str"):
-        self._default_model = model
-
-    def set_default_summarize_history(self, summarize_history: bool):
-        self._default_summarize_history = summarize_history
-
-    def set_default_history_summarization_token_threshold(
-        self, history_summarization_token_threshold: int
-    ):
-        self._default_history_summarization_token_threshold = (
-            history_summarization_token_threshold
-        )
-
-    def set_default_model_settings(self, model_settings: "ModelSettings"):
-        self._default_model_settings = model_settings
-
-    def set_default_yolo_mode(self, yolo_mode: bool):
-        self._default_yolo_mode = yolo_mode
-
-
-llm_config = LLMConfig()

zrb/config/llm_context/config.py

@@ -1,129 +0,0 @@
-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
-
-
-class LLMContextConfig:
-    """High-level API for interacting with cascaded configurations."""
-
-    def _find_config_files(self, cwd: str) -> list[str]:
-        configs = []
-        current_dir = cwd
-        home_dir = os.path.expanduser("~")
-        while True:
-            config_path = os.path.join(current_dir, CFG.LLM_CONTEXT_FILE)
-            if os.path.exists(config_path):
-                configs.append(config_path)
-            if current_dir == home_dir or current_dir == "/":
-                break
-            current_dir = os.path.dirname(current_dir)
-        return configs
-
-    def _parse_config(self, file_path: str) -> dict[str, str]:
-        with open(file_path, "r") as f:
-            content = f.read()
-        return markdown_to_dict(content)
-
-    def _get_all_sections(self, cwd: str) -> list[tuple[str, dict[str, str]]]:
-        config_files = self._find_config_files(cwd)
-        all_sections = []
-        for config_file in config_files:
-            config_dir = os.path.dirname(config_file)
-            sections = self._parse_config(config_file)
-            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)
-
-
-llm_context_config = LLMContextConfig()

zrb/config/llm_context/config_parser.py

@@ -1,46 +0,0 @@
-import re
-
-from zrb.util.llm.prompt import promote_markdown_headers
-
-
-def markdown_to_dict(markdown: str) -> dict[str, str]:
-    sections: 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)
-            if h1_match:
-                # Save previous section
-                if current_title:
-                    sections[current_title] = "\n".join(current_content).strip()
-                # Start new section
-                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()
-    return {
-        header: promote_markdown_headers(content)
-        for header, content in sections.items()
-    }

zrb/config/llm_rate_limitter.py

@@ -1,137 +0,0 @@
-import asyncio
-import time
-from collections import deque
-from typing import Callable
-
-import tiktoken
-
-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:
-    """
-    Helper class to enforce LLM API rate limits and throttling.
-    Tracks requests and tokens in a rolling 60-second window.
-    """
-
-    def __init__(
-        self,
-        max_requests_per_minute: int | None = None,
-        max_tokens_per_minute: int | None = None,
-        max_tokens_per_request: int | None = None,
-        throttle_sleep: float | None = None,
-        token_counter_fn: Callable[[str], int] | 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._throttle_sleep = throttle_sleep
-        self._token_counter_fn = token_counter_fn
-        self.request_times = deque()
-        self.token_times = deque()
-
-    @property
-    def max_requests_per_minute(self) -> int:
-        if self._max_requests_per_minute is not None:
-            return self._max_requests_per_minute
-        return CFG.LLM_MAX_REQUESTS_PER_MINUTE
-
-    @property
-    def max_tokens_per_minute(self) -> int:
-        if self._max_tokens_per_minute is not None:
-            return self._max_tokens_per_minute
-        return CFG.LLM_MAX_TOKENS_PER_MINUTE
-
-    @property
-    def max_tokens_per_request(self) -> int:
-        if self._max_tokens_per_request is not None:
-            return self._max_tokens_per_request
-        return CFG.LLM_MAX_TOKENS_PER_REQUEST
-
-    @property
-    def throttle_sleep(self) -> float:
-        if self._throttle_sleep is not None:
-            return self._throttle_sleep
-        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 set_max_requests_per_minute(self, value: int):
-        self._max_requests_per_minute = value
-
-    def set_max_tokens_per_minute(self, value: int):
-        self._max_tokens_per_minute = value
-
-    def set_max_tokens_per_request(self, value: int):
-        self._max_tokens_per_request = 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):
-        now = time.time()
-        tokens = self.count_token(prompt)
-        # Clean up old entries
-        while self.request_times and now - self.request_times[0] > 60:
-            self.request_times.popleft()
-        while self.token_times and now - self.token_times[0][0] > 60:
-            self.token_times.popleft()
-        # 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})."
-            )
-        # 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
-        ):
-            await asyncio.sleep(self.throttle_sleep)
-            now = time.time()
-            while self.request_times and now - self.request_times[0] > 60:
-                self.request_times.popleft()
-            while self.token_times and now - self.token_times[0][0] > 60:
-                self.token_times.popleft()
-        # Record this request
-        self.request_times.append(now)
-        self.token_times.append((now, tokens))
-
-
-llm_rate_limitter = LLMRateLimiter()