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

zrb/llm/config/__init__.py

@@ -0,0 +1,4 @@
+from zrb.llm.config.config import LLMConfig, llm_config
+from zrb.llm.config.limiter import LLMLimiter, llm_limiter
+
+__all__ = ["LLMConfig", "llm_config", "LLMLimiter", "llm_limiter"]

zrb/llm/config/config.py

@@ -0,0 +1,122 @@
+from typing import TYPE_CHECKING
+
+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:
+    """
+    Configuration provider for Pollux.
+    Allows runtime configuration while falling back to ZRB global settings.
+    """
+
+    def __init__(self):
+        self._model: "str | Model | None" = None
+        self._model_settings: "ModelSettings | None" = None
+        self._system_prompt: str | None = None
+        self._summarization_prompt: str | None = None
+
+        # Optional overrides for provider resolution
+        self._api_key: str | None = None
+        self._base_url: str | None = None
+        self._provider: "str | Provider | None" = None
+
+    # --- Model ---
+
+    @property
+    def model(self) -> "str | Model":
+        """
+        The LLM model to use. Returns a model string (e.g. 'openai:gpt-4o')
+        or a pydantic_ai Model object.
+        """
+        if self._model is not None:
+            return self._model
+
+        model_name = CFG.LLM_MODEL or "openai:gpt-4o"
+        provider = self.provider
+
+        return self._resolve_model(model_name, provider)
+
+    @model.setter
+    def model(self, value: "str | Model"):
+        self._model = value
+
+    # --- Model Settings ---
+
+    @property
+    def model_settings(self) -> "ModelSettings | None":
+        """Runtime settings for the model (temperature, etc.)."""
+        return self._model_settings
+
+    @model_settings.setter
+    def model_settings(self, value: "ModelSettings"):
+        self._model_settings = value
+
+    # --- Provider Helpers (Advanced) ---
+
+    @property
+    def api_key(self) -> str | None:
+        return self._api_key or getattr(CFG, "LLM_API_KEY", None)
+
+    @api_key.setter
+    def api_key(self, value: str):
+        self._api_key = value
+
+    @property
+    def base_url(self) -> str | None:
+        return self._base_url or getattr(CFG, "LLM_BASE_URL", None)
+
+    @base_url.setter
+    def base_url(self, value: str):
+        self._base_url = value
+
+    @property
+    def provider(self) -> "str | Provider":
+        """Resolves the model provider based on config."""
+        if self._provider is not None:
+            return self._provider
+
+        # If API Key or Base URL is set, we assume OpenAI-compatible provider
+        if self.api_key or self.base_url:
+            from pydantic_ai.providers.openai import OpenAIProvider
+
+            return OpenAIProvider(api_key=self.api_key, base_url=self.base_url)
+
+        return "openai"
+
+    @provider.setter
+    def provider(self, value: "str | Provider"):
+        self._provider = value
+
+    # --- Internal Logic ---
+
+    def _resolve_model(
+        self, model_name: str, provider: "str | Provider"
+    ) -> "str | Model":
+        # Strip existing provider prefix if present
+        clean_model_name = model_name.split(":", 1)[-1]
+
+        # 1. Provider is an Object (e.g. OpenAIProvider created from custom config)
+        # We check specific types we know how to wrap
+        try:
+            from pydantic_ai.models.openai import OpenAIChatModel
+            from pydantic_ai.providers.openai import OpenAIProvider
+
+            if isinstance(provider, OpenAIProvider):
+                return OpenAIChatModel(model_name=clean_model_name, provider=provider)
+        except ImportError:
+            pass
+
+        # 2. Provider is a String
+        if isinstance(provider, str):
+            return f"{provider}:{clean_model_name}"
+
+        # 3. Fallback (Provider is None or unknown object)
+        return model_name
+
+
+llm_config = LLMConfig()

zrb/llm/config/limiter.py

@@ -0,0 +1,247 @@
+import asyncio
+import json
+import time
+from collections import deque
+from typing import Any, Callable
+
+from zrb.config.config import CFG
+
+
+class LLMLimiter:
+    """
+    Manages LLM constraints: Context Window (Pruning) and Rate Limits (Throttling).
+    Designed as a singleton to share limits across tasks.
+    """
+
+    def __init__(self):
+        # Sliding window logs
+        self._request_log: deque[float] = deque()
+        self._token_log: deque[tuple[float, int]] = deque()
+
+        # Internal overrides
+        self._max_requests_per_minute: int | None = None
+        self._max_tokens_per_minute: int | None = None
+        self._max_tokens_per_request: int | None = None
+        self._throttle_check_interval: float | None = None
+
+    # --- Configuration Properties ---
+
+    @property
+    def max_requests_per_minute(self) -> int:
+        if self._max_requests_per_minute is not None:
+            return self._max_requests_per_minute
+        return getattr(CFG, "LLM_MAX_REQUESTS_PER_MINUTE", None) or 60
+
+    @max_requests_per_minute.setter
+    def max_requests_per_minute(self, value: int):
+        self._max_requests_per_minute = value
+
+    @property
+    def max_tokens_per_minute(self) -> int:
+        if self._max_tokens_per_minute is not None:
+            return self._max_tokens_per_minute
+        return getattr(CFG, "LLM_MAX_TOKENS_PER_MINUTE", None) or 100_000
+
+    @max_tokens_per_minute.setter
+    def max_tokens_per_minute(self, value: int):
+        self._max_tokens_per_minute = value
+
+    @property
+    def max_tokens_per_request(self) -> int:
+        if self._max_tokens_per_request is not None:
+            return self._max_tokens_per_request
+        return getattr(CFG, "LLM_MAX_TOKENS_PER_REQUEST", None) or 16_000
+
+    @max_tokens_per_request.setter
+    def max_tokens_per_request(self, value: int):
+        self._max_tokens_per_request = value
+
+    @property
+    def throttle_check_interval(self) -> float:
+        if self._throttle_check_interval is not None:
+            return self._throttle_check_interval
+        return getattr(CFG, "LLM_THROTTLE_SLEEP", None) or 0.1
+
+    @throttle_check_interval.setter
+    def throttle_check_interval(self, value: float):
+        self._throttle_check_interval = value
+
+    @property
+    def use_tiktoken(self) -> bool:
+        return getattr(CFG, "USE_TIKTOKEN", False)
+
+    @property
+    def tiktoken_encoding(self) -> str:
+        return getattr(CFG, "TIKTOKEN_ENCODING_NAME", "cl100k_base")
+
+    # --- Public API ---
+
+    def fit_context_window(self, history: list[Any], new_message: Any) -> list[Any]:
+        """
+        Prunes the history (removing oldest turns) so that 'history + new_message'
+        fits within 'max_tokens_per_request'.
+        Ensures strict tool call pairing by removing full conversation turns.
+        """
+        if not history:
+            return history
+
+        # Import message types locally to avoid circular deps or startup cost
+        try:
+            from pydantic_ai.messages import (
+                ModelRequest,
+                ToolReturnPart,
+                UserPromptPart,
+            )
+        except ImportError:
+            # Fallback if pydantic_ai is not installed (unlikely in context)
+            return []
+
+        def is_turn_start(msg: Any) -> bool:
+            """Identify start of a new user interaction (User Prompt without Tool Return)."""
+            if not isinstance(msg, ModelRequest):
+                return False
+            has_user = any(isinstance(p, UserPromptPart) for p in msg.parts)
+            has_return = any(isinstance(p, ToolReturnPart) for p in msg.parts)
+            return has_user and not has_return
+
+        new_msg_tokens = self._count_tokens(new_message)
+        if new_msg_tokens > self.max_tokens_per_request:
+            return []
+
+        pruned_history = list(history)
+
+        while pruned_history:
+            history_tokens = self._count_tokens(pruned_history)
+            total_tokens = history_tokens + new_msg_tokens
+
+            if total_tokens <= self.max_tokens_per_request:
+                break
+
+            # Pruning Strategy: Find the start of the *next* turn and cut everything before it.
+            # We start searching from index 1 because removing index 0 (current start) is the goal.
+            next_turn_index = -1
+            for i in range(1, len(pruned_history)):
+                if is_turn_start(pruned_history[i]):
+                    next_turn_index = i
+                    break
+
+            if next_turn_index != -1:
+                # Remove everything up to the next turn
+                pruned_history = pruned_history[next_turn_index:]
+            else:
+                # No subsequent turns found.
+                # This implies the history contains only one (potentially long) turn or partial fragments.
+                # To satisfy the limit, we must clear the history entirely.
+                pruned_history = []
+
+        return pruned_history
+
+    async def acquire(self, content: Any, notifier: Callable[[str], Any] | None = None):
+        """
+        Acquires permission to proceed with the given content.
+        Calculates token count internally and waits if rate limits are exceeded.
+        """
+        # Calculate tokens once
+        estimated_tokens = self._count_tokens(content)
+
+        # 1. Prune logs older than 60 seconds
+        self._prune_logs()
+
+        # 2. Check limits loop
+        notified = False
+        while not self._can_proceed(estimated_tokens):
+            wait_time = self._calculate_wait_time(estimated_tokens)
+            reason = self._get_limit_reason(estimated_tokens)
+
+            if notifier:
+                msg = f"Rate Limit Reached: {reason}. Waiting {wait_time:.1f}s..."
+                # Only notify once or if status changes? Simple is better.
+                notifier(msg)
+                notified = True
+
+            await asyncio.sleep(self.throttle_check_interval)
+            self._prune_logs()
+
+        if notified and notifier:
+            notifier("")  # Clear status
+
+        # 3. Record usage
+        now = time.time()
+        self._request_log.append(now)
+        self._token_log.append((now, estimated_tokens))
+
+    def count_tokens(self, content: Any) -> int:
+        """Public alias for internal counter."""
+        return self._count_tokens(content)
+
+    # --- Internal Helpers ---
+
+    def _count_tokens(self, content: Any) -> int:
+        text = self._to_str(content)
+        if self.use_tiktoken:
+            try:
+                import tiktoken
+
+                enc = tiktoken.get_encoding(self.tiktoken_encoding)
+                return len(enc.encode(text))
+            except ImportError:
+                pass
+        # Fallback approximation (char/4)
+        return len(text) // 4
+
+    def _to_str(self, content: Any) -> str:
+        if isinstance(content, str):
+            return content
+        try:
+            return json.dumps(content, default=str)
+        except Exception:
+            return str(content)
+
+    def _prune_logs(self):
+        now = time.time()
+        window_start = now - 60
+
+        while self._request_log and self._request_log[0] < window_start:
+            self._request_log.popleft()
+
+        while self._token_log and self._token_log[0][0] < window_start:
+            self._token_log.popleft()
+
+    def _can_proceed(self, tokens: int) -> bool:
+        requests_ok = len(self._request_log) < self.max_requests_per_minute
+
+        current_tokens = sum(t for _, t in self._token_log)
+        tokens_ok = (current_tokens + tokens) <= self.max_tokens_per_minute
+
+        return requests_ok and tokens_ok
+
+    def _get_limit_reason(self, tokens: int) -> str:
+        if len(self._request_log) >= self.max_requests_per_minute:
+            return f"Max Requests ({self.max_requests_per_minute}/min)"
+        return f"Max Tokens ({self.max_tokens_per_minute}/min)"
+
+    def _calculate_wait_time(self, tokens: int) -> float:
+        now = time.time()
+        # Default wait
+        wait = 1.0
+
+        # If request limit hit, wait until oldest request expires
+        if len(self._request_log) >= self.max_requests_per_minute:
+            oldest = self._request_log[0]
+            wait = max(0.1, 60 - (now - oldest))
+
+        # If token limit hit, wait until enough tokens expire
+        current_tokens = sum(t for _, t in self._token_log)
+        if current_tokens + tokens > self.max_tokens_per_minute:
+            needed = (current_tokens + tokens) - self.max_tokens_per_minute
+            freed = 0
+            for ts, count in self._token_log:
+                freed += count
+                if freed >= needed:
+                    wait = max(wait, 60 - (now - ts))
+                    break
+
+        return wait
+
+
+llm_limiter = LLMLimiter()

zrb/llm/history_manager/__init__.py

@@ -0,0 +1,4 @@
+from zrb.llm.history_manager.any_history_manager import AnyHistoryManager
+from zrb.llm.history_manager.file_history_manager import FileHistoryManager
+
+__all__ = ["AnyHistoryManager", "FileHistoryManager"]

zrb/llm/history_manager/any_history_manager.py

@@ -0,0 +1,23 @@
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pydantic_ai import ModelMessage
+
+
+class AnyHistoryManager(ABC):
+    @abstractmethod
+    def load(self, conversation_name: str) -> "list[ModelMessage]":
+        pass
+
+    @abstractmethod
+    def save(self, conversation_name: str):
+        pass
+
+    @abstractmethod
+    def update(self, conversation_name: str, messages: "list[ModelMessage]"):
+        pass
+
+    @abstractmethod
+    def search(self, keyword: str) -> list[str]:
+        pass

zrb/llm/history_manager/file_history_manager.py

@@ -0,0 +1,91 @@
+import json
+import os
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pydantic_ai import ModelMessage
+
+from zrb.llm.history_manager.any_history_manager import AnyHistoryManager
+from zrb.util.match import fuzzy_match
+
+
+class FileHistoryManager(AnyHistoryManager):
+    def __init__(self, history_dir: str):
+        self._history_dir = os.path.expanduser(history_dir)
+        self._cache: "dict[str, list[ModelMessage]]" = {}
+        if not os.path.exists(self._history_dir):
+            os.makedirs(self._history_dir, exist_ok=True)
+
+    def _get_file_path(self, conversation_name: str) -> str:
+        # Sanitize conversation name to be safe for filename
+        safe_name = "".join(
+            c for c in conversation_name if c.isalnum() or c in (" ", ".", "_", "-")
+        ).strip()
+        if not safe_name:
+            safe_name = "default"
+        return os.path.join(self._history_dir, f"{safe_name}.json")
+
+    def load(self, conversation_name: str) -> "list[ModelMessage]":
+        from pydantic_ai.messages import ModelMessagesTypeAdapter
+
+        if conversation_name in self._cache:
+            return self._cache[conversation_name]
+
+        file_path = self._get_file_path(conversation_name)
+        if not os.path.exists(file_path):
+            return []
+
+        try:
+            with open(file_path, "r", encoding="utf-8") as f:
+                content = f.read()
+                if not content.strip():
+                    return []
+                data = json.loads(content)
+                messages = ModelMessagesTypeAdapter.validate_python(data)
+                self._cache[conversation_name] = messages
+                return messages
+        except (json.JSONDecodeError, OSError) as e:
+            # Log error or warn? For now, return empty list or re-raise.
+            # Returning empty list is safer for UI not to crash.
+            print(f"Warning: Failed to load history for {conversation_name}: {e}")
+            return []
+
+    def update(self, conversation_name: str, messages: "list[ModelMessage]"):
+        self._cache[conversation_name] = messages
+
+    def save(self, conversation_name: str):
+        from pydantic_ai.messages import ModelMessagesTypeAdapter
+
+        if conversation_name not in self._cache:
+            return
+
+        messages = self._cache[conversation_name]
+        file_path = self._get_file_path(conversation_name)
+
+        try:
+            data = ModelMessagesTypeAdapter.dump_python(messages, mode="json")
+            with open(file_path, "w", encoding="utf-8") as f:
+                json.dump(data, f, indent=2)
+        except OSError as e:
+            print(f"Error: Failed to save history for {conversation_name}: {e}")
+
+    def search(self, keyword: str) -> list[str]:
+        if not os.path.exists(self._history_dir):
+            return []
+
+        matches = []
+        for filename in os.listdir(self._history_dir):
+            if not filename.endswith(".json"):
+                continue
+
+            # Remove extension to get session name
+            session_name = filename[:-5]
+
+            is_match, score = fuzzy_match(session_name, keyword)
+            if is_match:
+                matches.append((session_name, score))
+
+        # Sort by score (lower is better)
+        matches.sort(key=lambda x: x[1])
+
+        return [m[0] for m in matches]

zrb/llm/history_processor/summarizer.py

@@ -0,0 +1,108 @@
+from typing import TYPE_CHECKING, Any, Awaitable, Callable
+
+from zrb.config.config import CFG
+from zrb.llm.agent.summarizer import create_summarizer_agent
+from zrb.llm.config.limiter import LLMLimiter
+from zrb.llm.config.limiter import llm_limiter as default_llm_limiter
+from zrb.util.markdown import make_markdown_section
+
+if TYPE_CHECKING:
+    from pydantic_ai.messages import ModelMessage
+else:
+    ModelMessage = Any
+
+
+def is_turn_start(msg: Any) -> bool:
+    """Identify start of a new user interaction (User Prompt without Tool Return)."""
+    from pydantic_ai.messages import ModelRequest, ToolReturnPart, UserPromptPart
+
+    if not isinstance(msg, ModelRequest):
+        return False
+    # In pydantic_ai, ModelRequest parts can be list of various parts
+    has_user = any(isinstance(p, UserPromptPart) for p in msg.parts)
+    has_return = any(isinstance(p, ToolReturnPart) for p in msg.parts)
+    return has_user and not has_return
+
+
+async def summarize_history(
+    messages: "list[ModelMessage]",
+    agent: Any = None,
+    summary_window: int | None = None,
+) -> "list[ModelMessage]":
+    """
+    Summarizes the history, keeping the last `summary_window` messages intact.
+    Returns a new list of messages where older messages are replaced by a summary.
+    """
+    from pydantic_ai.messages import ModelRequest, UserPromptPart
+
+    if summary_window is None:
+        summary_window = CFG.LLM_HISTORY_SUMMARIZATION_WINDOW
+    if len(messages) <= summary_window:
+        return messages
+
+    # Determine split index
+    # We want to keep at least summary_window messages.
+    # So split_idx <= len(messages) - summary_window.
+    # We search backwards from there for a clean turn start.
+    start_search_idx = max(0, len(messages) - summary_window)
+    split_idx = -1
+
+    # Iterate backwards from start_search_idx to 0
+    for i in range(start_search_idx, -1, -1):
+        if is_turn_start(messages[i]):
+            split_idx = i
+            break
+
+    if split_idx <= 0:
+        return messages
+
+    to_summarize = messages[:split_idx]
+    to_keep = messages[split_idx:]
+
+    # Simple text representation for now
+    history_text = "\n".join([str(m) for m in to_summarize])
+
+    summarizer_agent = agent or create_summarizer_agent()
+
+    # Run the summarizer agent
+    result = await summarizer_agent.run(
+        f"Summarize this conversation history:\n{history_text}"
+    )
+    summary_text = result.output
+
+    # Create a summary message injected as user context
+    summary_message = ModelRequest(
+        parts=[
+            UserPromptPart(
+                content=make_markdown_section(
+                    "Previous conversation summary", summary_text
+                )
+            )
+        ]
+    )
+
+    return [summary_message] + to_keep
+
+
+def create_summarizer_history_processor(
+    agent: Any = None,
+    limiter: LLMLimiter | None = None,
+    token_threshold: int | None = None,
+    summary_window: int | None = None,
+) -> "Callable[[list[ModelMessage]], Awaitable[list[ModelMessage]]]":
+    """
+    Creates a history processor that auto-summarizes history when it exceeds `token_threshold`.
+    """
+    llm_limiter = limiter or default_llm_limiter
+    if token_threshold is None:
+        token_threshold = CFG.LLM_HISTORY_SUMMARIZATION_TOKEN_THRESHOLD
+
+    async def process_history(messages: "list[ModelMessage]") -> "list[ModelMessage]":
+        current_tokens = llm_limiter.count_tokens(messages)
+
+        if current_tokens <= token_threshold:
+            return messages
+
+        return await summarize_history(messages, agent, summary_window)
+
+    return process_history

zrb/llm/note/__init__.py

@@ -0,0 +1,3 @@
+from zrb.llm.note.manager import NoteManager
+
+__all__ = ["NoteManager"]