Base Utilities

bruno_llm.base

Base utilities and common functionality for LLM providers.

This module provides shared utilities that all provider implementations can use, including token counting, rate limiting, retry logic, cost tracking, caching, streaming, context management, and middleware.

BaseProvider

Bases: LLMInterface, ABC

Base class for LLM provider implementations.

Provides common functionality that all providers can use: - Retry logic with exponential backoff - Rate limiting - Cost tracking - Error handling patterns

Subclasses must implement the LLMInterface methods: - generate() - stream() - get_token_count() - check_connection() - list_models() - get_model_info() - set_system_prompt() - get_system_prompt()

Example

class MyProvider(BaseProvider): ... async def generate(self, messages, kwargs): ... return await self._with_retry( ... self._generate_impl(messages, kwargs) ... )

Source code in bruno_llm/base/base_provider.py

class BaseProvider(LLMInterface, ABC):
    """
    Base class for LLM provider implementations.

    Provides common functionality that all providers can use:
    - Retry logic with exponential backoff
    - Rate limiting
    - Cost tracking
    - Error handling patterns

    Subclasses must implement the LLMInterface methods:
    - generate()
    - stream()
    - get_token_count()
    - check_connection()
    - list_models()
    - get_model_info()
    - set_system_prompt()
    - get_system_prompt()

    Example:
        >>> class MyProvider(BaseProvider):
        ...     async def generate(self, messages, **kwargs):
        ...         return await self._with_retry(
        ...             self._generate_impl(messages, **kwargs)
        ...         )
    """

    def __init__(
        self,
        provider_name: str,
        max_retries: int = 3,
        timeout: float = 30.0,
        **kwargs: Any,
    ):
        """
        Initialize base provider.

        Args:
            provider_name: Name of the provider (e.g., "ollama", "openai")
            max_retries: Maximum number of retry attempts
            timeout: Request timeout in seconds
            **kwargs: Additional provider-specific configuration
        """
        self.provider_name = provider_name
        self.max_retries = max_retries
        self.timeout = timeout
        self._system_prompt: Optional[str] = None
        self._config = kwargs

    async def _with_retry(
        self,
        coro: Callable[[], T],
        max_retries: Optional[int] = None,
    ) -> T:
        """
        Execute coroutine with retry logic.

        Implements exponential backoff with jitter for retries.

        Args:
            coro: Async function to execute
            max_retries: Override default max_retries

        Returns:
            Result from coroutine

        Raises:
            LLMError: If all retries are exhausted
        """
        retries = max_retries if max_retries is not None else self.max_retries
        last_error = None

        for attempt in range(retries + 1):
            try:
                return await coro()
            except Exception as e:
                last_error = e

                if attempt < retries:
                    # Exponential backoff with jitter
                    delay = (2**attempt) + (asyncio.get_event_loop().time() % 1)
                    await asyncio.sleep(delay)
                    continue

                # All retries exhausted
                if isinstance(e, LLMError):
                    raise
                raise LLMError(
                    f"Request failed after {retries + 1} attempts",
                    provider=self.provider_name,
                    original_error=e,
                ) from e

        # Should never reach here, but for type safety
        raise LLMError(
            "Unexpected retry loop exit",
            provider=self.provider_name,
            original_error=last_error,
        )

    def set_system_prompt(self, prompt: str) -> None:
        """
        Set system prompt for the provider.

        Args:
            prompt: System prompt text
        """
        self._system_prompt = prompt

    def get_system_prompt(self) -> Optional[str]:
        """
        Get current system prompt.

        Returns:
            Current system prompt or None
        """
        return self._system_prompt

    def _add_system_prompt(self, messages: list[Message]) -> list[Message]:
        """
        Add system prompt to messages if set.

        Args:
            messages: Original messages

        Returns:
            Messages with system prompt prepended if set
        """
        if not self._system_prompt:
            return messages

        from bruno_core.models import MessageRole

        # Check if first message is already a system message
        if messages and messages[0].role == MessageRole.SYSTEM:
            return messages

        # Prepend system prompt
        system_message = Message(
            role=MessageRole.SYSTEM,
            content=self._system_prompt,
        )
        return [system_message] + messages

    def get_model_info(self) -> dict[str, Any]:
        """
        Get provider and configuration information.

        Returns:
            Dict with provider information
        """
        return {
            "provider": self.provider_name,
            "max_retries": self.max_retries,
            "timeout": self.timeout,
            "system_prompt": self._system_prompt,
            **self._config,
        }

__init__(provider_name, max_retries=3, timeout=30.0, **kwargs)

Initialize base provider.

Parameters:

Name	Type	Description	Default
provider_name	str	Name of the provider (e.g., "ollama", "openai")	required
max_retries	int	Maximum number of retry attempts	3
timeout	float	Request timeout in seconds	30.0
**kwargs	Any	Additional provider-specific configuration	{}

Source code in bruno_llm/base/base_provider.py

def __init__(
    self,
    provider_name: str,
    max_retries: int = 3,
    timeout: float = 30.0,
    **kwargs: Any,
):
    """
    Initialize base provider.

    Args:
        provider_name: Name of the provider (e.g., "ollama", "openai")
        max_retries: Maximum number of retry attempts
        timeout: Request timeout in seconds
        **kwargs: Additional provider-specific configuration
    """
    self.provider_name = provider_name
    self.max_retries = max_retries
    self.timeout = timeout
    self._system_prompt: Optional[str] = None
    self._config = kwargs

set_system_prompt(prompt)

Set system prompt for the provider.

Parameters:

Name	Type	Description	Default
prompt	str	System prompt text	required

Source code in bruno_llm/base/base_provider.py

def set_system_prompt(self, prompt: str) -> None:
    """
    Set system prompt for the provider.

    Args:
        prompt: System prompt text
    """
    self._system_prompt = prompt

get_system_prompt()

Get current system prompt.

Returns:

Type	Description
Optional[str]	Current system prompt or None

Source code in bruno_llm/base/base_provider.py

def get_system_prompt(self) -> Optional[str]:
    """
    Get current system prompt.

    Returns:
        Current system prompt or None
    """
    return self._system_prompt

get_model_info()

Get provider and configuration information.

Returns:

Type	Description
dict[str, Any]	Dict with provider information

Source code in bruno_llm/base/base_provider.py

def get_model_info(self) -> dict[str, Any]:
    """
    Get provider and configuration information.

    Returns:
        Dict with provider information
    """
    return {
        "provider": self.provider_name,
        "max_retries": self.max_retries,
        "timeout": self.timeout,
        "system_prompt": self._system_prompt,
        **self._config,
    }

CacheEntry dataclass

Cache entry for a response.

Attributes:

Name	Type	Description
response	str	The cached response
timestamp	float	When the entry was cached
hit_count	int	Number of times this entry was accessed
tokens	Optional[int]	Token count for the response (if available)

Source code in bruno_llm/base/cache.py

@dataclass
class CacheEntry:
    """
    Cache entry for a response.

    Attributes:
        response: The cached response
        timestamp: When the entry was cached
        hit_count: Number of times this entry was accessed
        tokens: Token count for the response (if available)
    """

    response: str
    timestamp: float
    hit_count: int = 0
    tokens: Optional[int] = None

ResponseCache

LRU cache for LLM responses with TTL support.

Caches responses to avoid redundant API calls. Uses message content and parameters as cache keys. Includes TTL (time-to-live) to ensure responses don't become stale.

Features: - LRU eviction when max_size is reached - TTL-based expiration - Hit/miss statistics - Thread-safe operations (async-safe)

Parameters:

Name	Type	Description	Default
max_size	int	Maximum number of entries to cache (default: 1000)	1000
ttl	float	Time-to-live in seconds (default: 3600 = 1 hour)	3600

Example

cache = ResponseCache(max_size=100, ttl=300)

Check if response is cached

response = cache.get(messages, temperature=0.7) if response is None: ... response = await provider.generate(messages) ... cache.set(messages, response, temperature=0.7)

Source code in bruno_llm/base/cache.py

class ResponseCache:
    """
    LRU cache for LLM responses with TTL support.

    Caches responses to avoid redundant API calls. Uses message content
    and parameters as cache keys. Includes TTL (time-to-live) to ensure
    responses don't become stale.

    Features:
    - LRU eviction when max_size is reached
    - TTL-based expiration
    - Hit/miss statistics
    - Thread-safe operations (async-safe)

    Args:
        max_size: Maximum number of entries to cache (default: 1000)
        ttl: Time-to-live in seconds (default: 3600 = 1 hour)

    Example:
        >>> cache = ResponseCache(max_size=100, ttl=300)
        >>>
        >>> # Check if response is cached
        >>> response = cache.get(messages, temperature=0.7)
        >>> if response is None:
        ...     response = await provider.generate(messages)
        ...     cache.set(messages, response, temperature=0.7)
    """

    def __init__(self, max_size: int = 1000, ttl: float = 3600):
        """
        Initialize response cache.

        Args:
            max_size: Maximum number of cached entries
            ttl: Time-to-live in seconds for cached entries
        """
        self._cache: OrderedDict[str, CacheEntry] = OrderedDict()
        self._max_size = max_size
        self._ttl = ttl
        self._hits = 0
        self._misses = 0

    def _generate_key(self, messages: list[Message], **kwargs: Any) -> str:
        """
        Generate cache key from messages and parameters.

        Args:
            messages: List of conversation messages
            **kwargs: Additional generation parameters

        Returns:
            Cache key as hex string
        """
        # Convert messages to dict representation
        messages_dict = [{"role": msg.role.value, "content": msg.content} for msg in messages]

        # Create deterministic JSON representation
        key_data = {"messages": messages_dict, "params": dict(sorted(kwargs.items()))}

        key_json = json.dumps(key_data, sort_keys=True)
        return hashlib.sha256(key_json.encode()).hexdigest()

    def get(self, messages: list[Message], **kwargs: Any) -> Optional[str]:
        """
        Get cached response if available and not expired.

        Args:
            messages: List of conversation messages
            **kwargs: Additional generation parameters

        Returns:
            Cached response or None if not found/expired
        """
        key = self._generate_key(messages, **kwargs)

        if key not in self._cache:
            self._misses += 1
            return None

        entry = self._cache[key]
        current_time = time.time()

        # Check if entry has expired
        if current_time - entry.timestamp > self._ttl:
            # Remove expired entry
            del self._cache[key]
            self._misses += 1
            return None

        # Move to end (most recently used)
        self._cache.move_to_end(key)
        entry.hit_count += 1
        self._hits += 1

        return entry.response

    def set(
        self, messages: list[Message], response: str, tokens: Optional[int] = None, **kwargs: Any
    ) -> None:
        """
        Cache a response.

        Args:
            messages: List of conversation messages
            response: The response to cache
            tokens: Token count for the response (optional)
            **kwargs: Additional generation parameters
        """
        key = self._generate_key(messages, **kwargs)

        # Create cache entry
        entry = CacheEntry(
            response=response,
            timestamp=time.time(),
            hit_count=0,
            tokens=tokens,
        )

        # Add to cache
        self._cache[key] = entry
        self._cache.move_to_end(key)

        # Evict oldest entry if max size exceeded
        if len(self._cache) > self._max_size:
            self._cache.popitem(last=False)

    def clear(self) -> None:
        """Clear all cached entries."""
        self._cache.clear()
        self._hits = 0
        self._misses = 0

    def invalidate(self, messages: list[Message], **kwargs: Any) -> bool:
        """
        Invalidate a specific cache entry.

        Args:
            messages: List of conversation messages
            **kwargs: Additional generation parameters

        Returns:
            True if entry was found and removed, False otherwise
        """
        key = self._generate_key(messages, **kwargs)

        if key in self._cache:
            del self._cache[key]
            return True
        return False

    def get_stats(self) -> dict[str, Any]:
        """
        Get cache statistics.

        Returns:
            Dictionary with cache statistics
        """
        total_requests = self._hits + self._misses
        hit_rate = self._hits / total_requests if total_requests > 0 else 0.0

        return {
            "size": len(self._cache),
            "max_size": self._max_size,
            "hits": self._hits,
            "misses": self._misses,
            "hit_rate": hit_rate,
            "ttl": self._ttl,
        }

    def get_size_bytes(self) -> int:
        """
        Estimate cache size in bytes.

        Returns:
            Approximate cache size in bytes
        """
        total_size = 0
        for key, entry in self._cache.items():
            # Key size
            total_size += len(key.encode())
            # Response size
            total_size += len(entry.response.encode())
            # Overhead for entry metadata (~100 bytes)
            total_size += 100

        return total_size

    def cleanup_expired(self) -> int:
        """
        Remove all expired entries.

        Returns:
            Number of entries removed
        """
        current_time = time.time()
        expired_keys = [
            key for key, entry in self._cache.items() if current_time - entry.timestamp > self._ttl
        ]

        for key in expired_keys:
            del self._cache[key]

        return len(expired_keys)

    def get_top_entries(self, n: int = 10) -> list[tuple[str, CacheEntry]]:
        """
        Get top N most frequently accessed entries.

        Args:
            n: Number of entries to return

        Returns:
            List of (key, entry) tuples sorted by hit count
        """
        sorted_entries = sorted(self._cache.items(), key=lambda x: x[1].hit_count, reverse=True)
        return sorted_entries[:n]

__init__(max_size=1000, ttl=3600)

Initialize response cache.

Parameters:

Name	Type	Description	Default
max_size	int	Maximum number of cached entries	1000
ttl	float	Time-to-live in seconds for cached entries	3600

Source code in bruno_llm/base/cache.py

def __init__(self, max_size: int = 1000, ttl: float = 3600):
    """
    Initialize response cache.

    Args:
        max_size: Maximum number of cached entries
        ttl: Time-to-live in seconds for cached entries
    """
    self._cache: OrderedDict[str, CacheEntry] = OrderedDict()
    self._max_size = max_size
    self._ttl = ttl
    self._hits = 0
    self._misses = 0

get(messages, **kwargs)

Get cached response if available and not expired.

Parameters:

Name	Type	Description	Default
messages	list[Message]	List of conversation messages	required
**kwargs	Any	Additional generation parameters	{}

Returns:

Type	Description
Optional[str]	Cached response or None if not found/expired

Source code in bruno_llm/base/cache.py

def get(self, messages: list[Message], **kwargs: Any) -> Optional[str]:
    """
    Get cached response if available and not expired.

    Args:
        messages: List of conversation messages
        **kwargs: Additional generation parameters

    Returns:
        Cached response or None if not found/expired
    """
    key = self._generate_key(messages, **kwargs)

    if key not in self._cache:
        self._misses += 1
        return None

    entry = self._cache[key]
    current_time = time.time()

    # Check if entry has expired
    if current_time - entry.timestamp > self._ttl:
        # Remove expired entry
        del self._cache[key]
        self._misses += 1
        return None

    # Move to end (most recently used)
    self._cache.move_to_end(key)
    entry.hit_count += 1
    self._hits += 1

    return entry.response

set(messages, response, tokens=None, **kwargs)

Cache a response.

Parameters:

Name	Type	Description	Default
messages	list[Message]	List of conversation messages	required
response	str	The response to cache	required
tokens	Optional[int]	Token count for the response (optional)	None
**kwargs	Any	Additional generation parameters	{}

Source code in bruno_llm/base/cache.py

def set(
    self, messages: list[Message], response: str, tokens: Optional[int] = None, **kwargs: Any
) -> None:
    """
    Cache a response.

    Args:
        messages: List of conversation messages
        response: The response to cache
        tokens: Token count for the response (optional)
        **kwargs: Additional generation parameters
    """
    key = self._generate_key(messages, **kwargs)

    # Create cache entry
    entry = CacheEntry(
        response=response,
        timestamp=time.time(),
        hit_count=0,
        tokens=tokens,
    )

    # Add to cache
    self._cache[key] = entry
    self._cache.move_to_end(key)

    # Evict oldest entry if max size exceeded
    if len(self._cache) > self._max_size:
        self._cache.popitem(last=False)

clear()

Clear all cached entries.

Source code in bruno_llm/base/cache.py

def clear(self) -> None:
    """Clear all cached entries."""
    self._cache.clear()
    self._hits = 0
    self._misses = 0

invalidate(messages, **kwargs)

Invalidate a specific cache entry.

Parameters:

Name	Type	Description	Default
messages	list[Message]	List of conversation messages	required
**kwargs	Any	Additional generation parameters	{}

Returns:

Type	Description
bool	True if entry was found and removed, False otherwise

Source code in bruno_llm/base/cache.py

def invalidate(self, messages: list[Message], **kwargs: Any) -> bool:
    """
    Invalidate a specific cache entry.

    Args:
        messages: List of conversation messages
        **kwargs: Additional generation parameters

    Returns:
        True if entry was found and removed, False otherwise
    """
    key = self._generate_key(messages, **kwargs)

    if key in self._cache:
        del self._cache[key]
        return True
    return False

get_stats()

Get cache statistics.

Returns:

Type	Description
dict[str, Any]	Dictionary with cache statistics

Source code in bruno_llm/base/cache.py

def get_stats(self) -> dict[str, Any]:
    """
    Get cache statistics.

    Returns:
        Dictionary with cache statistics
    """
    total_requests = self._hits + self._misses
    hit_rate = self._hits / total_requests if total_requests > 0 else 0.0

    return {
        "size": len(self._cache),
        "max_size": self._max_size,
        "hits": self._hits,
        "misses": self._misses,
        "hit_rate": hit_rate,
        "ttl": self._ttl,
    }

get_size_bytes()

Estimate cache size in bytes.

Returns:

Type	Description
int	Approximate cache size in bytes

Source code in bruno_llm/base/cache.py

def get_size_bytes(self) -> int:
    """
    Estimate cache size in bytes.

    Returns:
        Approximate cache size in bytes
    """
    total_size = 0
    for key, entry in self._cache.items():
        # Key size
        total_size += len(key.encode())
        # Response size
        total_size += len(entry.response.encode())
        # Overhead for entry metadata (~100 bytes)
        total_size += 100

    return total_size

cleanup_expired()

Remove all expired entries.

Returns:

Type	Description
int	Number of entries removed

Source code in bruno_llm/base/cache.py

def cleanup_expired(self) -> int:
    """
    Remove all expired entries.

    Returns:
        Number of entries removed
    """
    current_time = time.time()
    expired_keys = [
        key for key, entry in self._cache.items() if current_time - entry.timestamp > self._ttl
    ]

    for key in expired_keys:
        del self._cache[key]

    return len(expired_keys)

get_top_entries(n=10)

Get top N most frequently accessed entries.

Parameters:

Name	Type	Description	Default
n	int	Number of entries to return	10

Returns:

Type	Description
list[tuple[str, CacheEntry]]	List of (key, entry) tuples sorted by hit count

Source code in bruno_llm/base/cache.py

def get_top_entries(self, n: int = 10) -> list[tuple[str, CacheEntry]]:
    """
    Get top N most frequently accessed entries.

    Args:
        n: Number of entries to return

    Returns:
        List of (key, entry) tuples sorted by hit count
    """
    sorted_entries = sorted(self._cache.items(), key=lambda x: x[1].hit_count, reverse=True)
    return sorted_entries[:n]

ContextLimits dataclass

Context window limits for a model.

Attributes:

Name	Type	Description
max_tokens	int	Maximum total tokens (input + output)
max_input_tokens	Optional[int]	Maximum input tokens
max_output_tokens	Optional[int]	Maximum output tokens
warning_threshold	float	Warn when this % of limit is reached (0.0-1.0)

Source code in bruno_llm/base/context.py

@dataclass
class ContextLimits:
    """
    Context window limits for a model.

    Attributes:
        max_tokens: Maximum total tokens (input + output)
        max_input_tokens: Maximum input tokens
        max_output_tokens: Maximum output tokens
        warning_threshold: Warn when this % of limit is reached (0.0-1.0)
    """

    max_tokens: int
    max_input_tokens: Optional[int] = None
    max_output_tokens: Optional[int] = None
    warning_threshold: float = 0.9

    def __post_init__(self):
        """Validate limits after initialization."""
        if self.max_input_tokens is None:
            self.max_input_tokens = self.max_tokens
        if self.max_output_tokens is None:
            self.max_output_tokens = self.max_tokens // 4  # Default 25% for output

__post_init__()

Validate limits after initialization.

Source code in bruno_llm/base/context.py

def __post_init__(self):
    """Validate limits after initialization."""
    if self.max_input_tokens is None:
        self.max_input_tokens = self.max_tokens
    if self.max_output_tokens is None:
        self.max_output_tokens = self.max_tokens // 4  # Default 25% for output

ContextWindowManager

Manage context windows and message truncation.

Handles: - Token counting for messages - Context limit checking - Automatic message truncation - Warning when approaching limits

Parameters:

Name	Type	Description	Default
model	str	Model name for context limits	required
token_counter	Optional[TokenCounter]	Token counter instance	None
limits	Optional[ContextLimits]	Custom context limits (overrides model defaults)	None
strategy	TruncationStrategy	Truncation strategy to use	SLIDING_WINDOW

Example

manager = ContextWindowManager(model="gpt-4")

Check if messages fit

if manager.check_limit(messages): ... response = await provider.generate(messages) ... else: ... # Truncate messages ... truncated = manager.truncate(messages) ... response = await provider.generate(truncated)

Source code in bruno_llm/base/context.py

class ContextWindowManager:
    """
    Manage context windows and message truncation.

    Handles:
    - Token counting for messages
    - Context limit checking
    - Automatic message truncation
    - Warning when approaching limits

    Args:
        model: Model name for context limits
        token_counter: Token counter instance
        limits: Custom context limits (overrides model defaults)
        strategy: Truncation strategy to use

    Example:
        >>> manager = ContextWindowManager(model="gpt-4")
        >>>
        >>> # Check if messages fit
        >>> if manager.check_limit(messages):
        ...     response = await provider.generate(messages)
        ... else:
        ...     # Truncate messages
        ...     truncated = manager.truncate(messages)
        ...     response = await provider.generate(truncated)
    """

    def __init__(
        self,
        model: str,
        token_counter: Optional[TokenCounter] = None,
        limits: Optional[ContextLimits] = None,
        strategy: TruncationStrategy = TruncationStrategy.SLIDING_WINDOW,
    ):
        """
        Initialize context window manager.

        Args:
            model: Model name
            token_counter: Token counter instance
            limits: Custom context limits
            strategy: Truncation strategy
        """
        self.model = model
        self.token_counter = token_counter or create_token_counter(model)
        self.limits = limits or self._get_model_limits(model)
        self.strategy = strategy
        self._warning_callback: Optional[Callable[[int, int], None]] = None

    def _get_model_limits(self, model: str) -> ContextLimits:
        """
        Get context limits for a model.

        Args:
            model: Model name

        Returns:
            Context limits for the model
        """
        # Try exact match first
        if model in MODEL_LIMITS:
            return MODEL_LIMITS[model]

        # Try partial match
        for model_name, limits in MODEL_LIMITS.items():
            if model.startswith(model_name):
                return limits

        # Default conservative limit
        return ContextLimits(max_tokens=4096)

    def count_tokens(self, messages: list[Message]) -> int:
        """
        Count tokens in messages.

        Args:
            messages: List of messages

        Returns:
            Total token count
        """
        return self.token_counter.count_messages_tokens(messages)

    def check_limit(
        self,
        messages: list[Message],
        max_output_tokens: Optional[int] = None,
    ) -> bool:
        """
        Check if messages fit within context limit.

        Args:
            messages: List of messages
            max_output_tokens: Expected output tokens

        Returns:
            True if messages fit, False otherwise
        """
        input_tokens = self.count_tokens(messages)
        output_tokens = max_output_tokens or self.limits.max_output_tokens
        total_tokens = input_tokens + output_tokens

        # Check warning threshold
        if input_tokens / self.limits.max_input_tokens >= self.limits.warning_threshold:
            if self._warning_callback:
                self._warning_callback(input_tokens, self.limits.max_input_tokens)

        return total_tokens <= self.limits.max_tokens

    def get_available_tokens(self, messages: list[Message]) -> int:
        """
        Get number of tokens available for output.

        Args:
            messages: List of messages

        Returns:
            Available tokens for output
        """
        input_tokens = self.count_tokens(messages)
        return max(0, self.limits.max_tokens - input_tokens)

    def truncate(
        self,
        messages: list[Message],
        max_output_tokens: Optional[int] = None,
    ) -> list[Message]:
        """
        Truncate messages to fit within context limit.

        Args:
            messages: List of messages
            max_output_tokens: Expected output tokens

        Returns:
            Truncated message list

        Raises:
            ContextLengthExceededError: If messages can't be truncated enough
        """
        output_tokens = max_output_tokens or self.limits.max_output_tokens
        target_input_tokens = self.limits.max_tokens - output_tokens

        if target_input_tokens <= 0:
            raise ContextLengthExceededError(
                f"Output tokens ({output_tokens}) exceed total limit ({self.limits.max_tokens})"
            )

        if self.strategy == TruncationStrategy.OLDEST_FIRST:
            return self._truncate_oldest_first(messages, target_input_tokens)
        elif self.strategy == TruncationStrategy.MIDDLE_OUT:
            return self._truncate_middle_out(messages, target_input_tokens)
        elif self.strategy == TruncationStrategy.SLIDING_WINDOW:
            return self._truncate_sliding_window(messages, target_input_tokens)
        elif self.strategy == TruncationStrategy.SMART:
            return self._truncate_smart(messages, target_input_tokens)
        else:
            return self._truncate_oldest_first(messages, target_input_tokens)

    def _truncate_oldest_first(
        self,
        messages: list[Message],
        target_tokens: int,
    ) -> list[Message]:
        """Remove oldest messages first (keep system message)."""
        # Always keep system messages
        system_messages = [m for m in messages if m.role == MessageRole.SYSTEM]
        other_messages = [m for m in messages if m.role != MessageRole.SYSTEM]

        # Start with system messages
        result = system_messages[:]
        current_tokens = self.count_tokens(result)

        # Add messages from newest to oldest
        for message in reversed(other_messages):
            message_tokens = self.token_counter.count_message_tokens(message)
            if current_tokens + message_tokens <= target_tokens:
                result.insert(len(system_messages), message)
                current_tokens += message_tokens
            else:
                break

        # Re-order to maintain chronological order (except system at start)
        return system_messages + list(reversed(result[len(system_messages) :]))

    def _truncate_middle_out(
        self,
        messages: list[Message],
        target_tokens: int,
    ) -> list[Message]:
        """Keep first and last messages, remove middle."""
        if len(messages) <= 2:
            return messages

        # Keep system messages and last message
        system_messages = [m for m in messages if m.role == MessageRole.SYSTEM]
        other_messages = [m for m in messages if m.role != MessageRole.SYSTEM]

        if not other_messages:
            return messages

        result = system_messages + [other_messages[-1]]
        current_tokens = self.count_tokens(result)

        # Add messages from the start
        for message in other_messages[:-1]:
            message_tokens = self.token_counter.count_message_tokens(message)
            if current_tokens + message_tokens <= target_tokens:
                result.insert(len(system_messages), message)
                current_tokens += message_tokens
            else:
                break

        return result

    def _truncate_sliding_window(
        self,
        messages: list[Message],
        target_tokens: int,
    ) -> list[Message]:
        """Keep most recent N messages."""
        # Always keep system messages
        system_messages = [m for m in messages if m.role == MessageRole.SYSTEM]
        other_messages = [m for m in messages if m.role != MessageRole.SYSTEM]

        result = system_messages[:]
        current_tokens = self.count_tokens(result)

        # Add messages from newest to oldest
        for message in reversed(other_messages):
            message_tokens = self.token_counter.count_message_tokens(message)
            if current_tokens + message_tokens <= target_tokens:
                result.append(message)
                current_tokens += message_tokens
            else:
                break

        # Keep system messages at start, reverse others
        return system_messages + list(reversed(result[len(system_messages) :]))

    def _truncate_smart(
        self,
        messages: list[Message],
        target_tokens: int,
    ) -> list[Message]:
        """
        Smart truncation: keep system + important messages + recent.

        Priority:
        1. System messages (always keep)
        2. Last 2 messages (recent context)
        3. Messages with high token count (likely important)
        4. Fill remaining space with recent messages
        """
        system_messages = [m for m in messages if m.role == MessageRole.SYSTEM]
        other_messages = [m for m in messages if m.role != MessageRole.SYSTEM]

        if not other_messages:
            return messages

        # Start with system messages
        result = system_messages[:]
        current_tokens = self.count_tokens(result)

        # Always include last 2 messages (most recent context)
        priority_messages = other_messages[-2:]
        for message in priority_messages:
            message_tokens = self.token_counter.count_message_tokens(message)
            if current_tokens + message_tokens <= target_tokens:
                result.append(message)
                current_tokens += message_tokens

        # Fill remaining space with other messages (newest first)
        remaining = list(other_messages[:-2])
        for message in reversed(remaining):
            message_tokens = self.token_counter.count_message_tokens(message)
            if current_tokens + message_tokens <= target_tokens:
                result.insert(len(system_messages), message)
                current_tokens += message_tokens
            else:
                break

        return result

    def set_warning_callback(self, callback: Callable[[int, int], None]) -> None:
        """
        Set callback for context limit warnings.

        Args:
            callback: Function (current_tokens, max_tokens) -> None
        """
        self._warning_callback = callback

    def get_stats(self, messages: list[Message]) -> dict:
        """
        Get statistics about context usage.

        Args:
            messages: List of messages

        Returns:
            Dictionary with context statistics
        """
        input_tokens = self.count_tokens(messages)
        available_tokens = self.get_available_tokens(messages)
        usage_percent = (input_tokens / self.limits.max_input_tokens) * 100

        return {
            "model": self.model,
            "input_tokens": input_tokens,
            "max_input_tokens": self.limits.max_input_tokens,
            "available_output_tokens": available_tokens,
            "max_output_tokens": self.limits.max_output_tokens,
            "total_limit": self.limits.max_tokens,
            "usage_percent": usage_percent,
            "within_limit": self.check_limit(messages),
            "message_count": len(messages),
        }

__init__(model, token_counter=None, limits=None, strategy=TruncationStrategy.SLIDING_WINDOW)

Initialize context window manager.

Parameters:

Name	Type	Description	Default
model	str	Model name	required
token_counter	Optional[TokenCounter]	Token counter instance	None
limits	Optional[ContextLimits]	Custom context limits	None
strategy	TruncationStrategy	Truncation strategy	SLIDING_WINDOW

Source code in bruno_llm/base/context.py

def __init__(
    self,
    model: str,
    token_counter: Optional[TokenCounter] = None,
    limits: Optional[ContextLimits] = None,
    strategy: TruncationStrategy = TruncationStrategy.SLIDING_WINDOW,
):
    """
    Initialize context window manager.

    Args:
        model: Model name
        token_counter: Token counter instance
        limits: Custom context limits
        strategy: Truncation strategy
    """
    self.model = model
    self.token_counter = token_counter or create_token_counter(model)
    self.limits = limits or self._get_model_limits(model)
    self.strategy = strategy
    self._warning_callback: Optional[Callable[[int, int], None]] = None

count_tokens(messages)

Count tokens in messages.

Parameters:

Name	Type	Description	Default
messages	list[Message]	List of messages	required

Returns:

Type	Description
int	Total token count

Source code in bruno_llm/base/context.py

def count_tokens(self, messages: list[Message]) -> int:
    """
    Count tokens in messages.

    Args:
        messages: List of messages

    Returns:
        Total token count
    """
    return self.token_counter.count_messages_tokens(messages)

check_limit(messages, max_output_tokens=None)

Check if messages fit within context limit.

Parameters:

Name	Type	Description	Default
messages	list[Message]	List of messages	required
max_output_tokens	Optional[int]	Expected output tokens	None

Returns:

Type	Description
bool	True if messages fit, False otherwise

Source code in bruno_llm/base/context.py

def check_limit(
    self,
    messages: list[Message],
    max_output_tokens: Optional[int] = None,
) -> bool:
    """
    Check if messages fit within context limit.

    Args:
        messages: List of messages
        max_output_tokens: Expected output tokens

    Returns:
        True if messages fit, False otherwise
    """
    input_tokens = self.count_tokens(messages)
    output_tokens = max_output_tokens or self.limits.max_output_tokens
    total_tokens = input_tokens + output_tokens

    # Check warning threshold
    if input_tokens / self.limits.max_input_tokens >= self.limits.warning_threshold:
        if self._warning_callback:
            self._warning_callback(input_tokens, self.limits.max_input_tokens)

    return total_tokens <= self.limits.max_tokens

get_available_tokens(messages)

Get number of tokens available for output.

Parameters:

Name	Type	Description	Default
messages	list[Message]	List of messages	required

Returns:

Type	Description
int	Available tokens for output

Source code in bruno_llm/base/context.py

def get_available_tokens(self, messages: list[Message]) -> int:
    """
    Get number of tokens available for output.

    Args:
        messages: List of messages

    Returns:
        Available tokens for output
    """
    input_tokens = self.count_tokens(messages)
    return max(0, self.limits.max_tokens - input_tokens)

truncate(messages, max_output_tokens=None)

Truncate messages to fit within context limit.

Parameters:

Name	Type	Description	Default
messages	list[Message]	List of messages	required
max_output_tokens	Optional[int]	Expected output tokens	None

Returns:

Type	Description
list[Message]	Truncated message list

Raises:

Type	Description
ContextLengthExceededError	If messages can't be truncated enough

Source code in bruno_llm/base/context.py

def truncate(
    self,
    messages: list[Message],
    max_output_tokens: Optional[int] = None,
) -> list[Message]:
    """
    Truncate messages to fit within context limit.

    Args:
        messages: List of messages
        max_output_tokens: Expected output tokens

    Returns:
        Truncated message list

    Raises:
        ContextLengthExceededError: If messages can't be truncated enough
    """
    output_tokens = max_output_tokens or self.limits.max_output_tokens
    target_input_tokens = self.limits.max_tokens - output_tokens

    if target_input_tokens <= 0:
        raise ContextLengthExceededError(
            f"Output tokens ({output_tokens}) exceed total limit ({self.limits.max_tokens})"
        )

    if self.strategy == TruncationStrategy.OLDEST_FIRST:
        return self._truncate_oldest_first(messages, target_input_tokens)
    elif self.strategy == TruncationStrategy.MIDDLE_OUT:
        return self._truncate_middle_out(messages, target_input_tokens)
    elif self.strategy == TruncationStrategy.SLIDING_WINDOW:
        return self._truncate_sliding_window(messages, target_input_tokens)
    elif self.strategy == TruncationStrategy.SMART:
        return self._truncate_smart(messages, target_input_tokens)
    else:
        return self._truncate_oldest_first(messages, target_input_tokens)

set_warning_callback(callback)

Set callback for context limit warnings.

Parameters:

Name	Type	Description	Default
callback	Callable[[int, int], None]	Function (current_tokens, max_tokens) -> None	required

Source code in bruno_llm/base/context.py

def set_warning_callback(self, callback: Callable[[int, int], None]) -> None:
    """
    Set callback for context limit warnings.

    Args:
        callback: Function (current_tokens, max_tokens) -> None
    """
    self._warning_callback = callback

get_stats(messages)

Get statistics about context usage.

Parameters:

Name	Type	Description	Default
messages	list[Message]	List of messages	required

Returns:

Type	Description
dict	Dictionary with context statistics

Source code in bruno_llm/base/context.py

def get_stats(self, messages: list[Message]) -> dict:
    """
    Get statistics about context usage.

    Args:
        messages: List of messages

    Returns:
        Dictionary with context statistics
    """
    input_tokens = self.count_tokens(messages)
    available_tokens = self.get_available_tokens(messages)
    usage_percent = (input_tokens / self.limits.max_input_tokens) * 100

    return {
        "model": self.model,
        "input_tokens": input_tokens,
        "max_input_tokens": self.limits.max_input_tokens,
        "available_output_tokens": available_tokens,
        "max_output_tokens": self.limits.max_output_tokens,
        "total_limit": self.limits.max_tokens,
        "usage_percent": usage_percent,
        "within_limit": self.check_limit(messages),
        "message_count": len(messages),
    }

TruncationStrategy

Bases: Enum

Strategy for truncating messages when context limit is exceeded.

Source code in bruno_llm/base/context.py

class TruncationStrategy(Enum):
    """Strategy for truncating messages when context limit is exceeded."""

    OLDEST_FIRST = "oldest_first"  # Remove oldest messages first
    MIDDLE_OUT = "middle_out"  # Keep first and last, remove middle
    SLIDING_WINDOW = "sliding_window"  # Keep most recent N messages
    SMART = "smart"  # Keep system + important messages + recent

CostTracker

Track API usage costs across requests.

Maintains history of API calls with token usage and costs. Supports multiple models with different pricing.

Example

tracker = CostTracker( ... provider_name="openai", ... pricing={ ... "gpt-4": {"input": 0.03, "output": 0.06}, ... "gpt-3.5-turbo": {"input": 0.001, "output": 0.002}, ... } ... ) tracker.track_request( ... model="gpt-4", ... input_tokens=100, ... output_tokens=50 ... ) print(tracker.get_total_cost()) 6.0 # $0.06 (in cents)

Source code in bruno_llm/base/cost_tracker.py

class CostTracker:
    """
    Track API usage costs across requests.

    Maintains history of API calls with token usage and costs.
    Supports multiple models with different pricing.

    Example:
        >>> tracker = CostTracker(
        ...     provider_name="openai",
        ...     pricing={
        ...         "gpt-4": {"input": 0.03, "output": 0.06},
        ...         "gpt-3.5-turbo": {"input": 0.001, "output": 0.002},
        ...     }
        ... )
        >>> tracker.track_request(
        ...     model="gpt-4",
        ...     input_tokens=100,
        ...     output_tokens=50
        ... )
        >>> print(tracker.get_total_cost())
        6.0  # $0.06 (in cents)
    """

    def __init__(
        self,
        provider_name: str,
        pricing: dict[str, dict[str, float]],
        currency: str = "USD",
    ):
        """
        Initialize cost tracker.

        Args:
            provider_name: Name of the provider
            pricing: Pricing per model (per 1K tokens)
                Format: {"model_name": {"input": price, "output": price}}
            currency: Currency code (default: "USD")
        """
        self.provider_name = provider_name
        self.pricing = pricing
        self.currency = currency
        self.usage_history: list[UsageRecord] = []

    def track_request(
        self,
        model: str,
        input_tokens: int,
        output_tokens: int,
        metadata: Optional[dict[str, str]] = None,
    ) -> UsageRecord:
        """
        Track a single API request.

        Args:
            model: Model name
            input_tokens: Number of input tokens
            output_tokens: Number of output tokens
            metadata: Optional metadata about the request

        Returns:
            UsageRecord with calculated costs
        """
        # Get pricing for model (with fallback to default if available)
        model_pricing = self.pricing.get(model, self.pricing.get("default", {}))

        # Calculate costs (pricing is per 1K tokens)
        input_cost = (input_tokens / 1000.0) * model_pricing.get("input", 0.0)
        output_cost = (output_tokens / 1000.0) * model_pricing.get("output", 0.0)
        total_cost = input_cost + output_cost

        # Create record
        record = UsageRecord(
            timestamp=time.time(),
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            input_cost=input_cost,
            output_cost=output_cost,
            total_cost=total_cost,
            metadata=metadata or {},
        )

        self.usage_history.append(record)
        return record

    def get_total_cost(self, model: Optional[str] = None) -> float:
        """
        Get total cost across all requests.

        Args:
            model: Optional model filter (None = all models)

        Returns:
            Total cost in the configured currency
        """
        total = 0.0
        for record in self.usage_history:
            if model is None or record.model == model:
                total += record.total_cost
        return total

    def get_total_tokens(self, model: Optional[str] = None) -> dict[str, int]:
        """
        Get total tokens used.

        Args:
            model: Optional model filter (None = all models)

        Returns:
            Dict with input, output, and total token counts
        """
        input_tokens = 0
        output_tokens = 0

        for record in self.usage_history:
            if model is None or record.model == model:
                input_tokens += record.input_tokens
                output_tokens += record.output_tokens

        return {
            "input": input_tokens,
            "output": output_tokens,
            "total": input_tokens + output_tokens,
        }

    def get_request_count(self, model: Optional[str] = None) -> int:
        """
        Get number of requests made.

        Args:
            model: Optional model filter (None = all models)

        Returns:
            Number of requests
        """
        if model is None:
            return len(self.usage_history)
        return sum(1 for r in self.usage_history if r.model == model)

    def get_model_breakdown(self) -> dict[str, dict[str, float]]:
        """
        Get cost breakdown by model.

        Returns:
            Dict mapping model names to their usage statistics
        """
        breakdown: dict[str, dict[str, float]] = {}

        for record in self.usage_history:
            if record.model not in breakdown:
                breakdown[record.model] = {
                    "cost": 0.0,
                    "input_tokens": 0,
                    "output_tokens": 0,
                    "requests": 0,
                }

            breakdown[record.model]["cost"] += record.total_cost
            breakdown[record.model]["input_tokens"] += record.input_tokens
            breakdown[record.model]["output_tokens"] += record.output_tokens
            breakdown[record.model]["requests"] += 1

        return breakdown

    def get_usage_report(self) -> dict:
        """
        Get comprehensive usage report.

        Returns:
            Dict with complete usage statistics
        """
        return {
            "provider": self.provider_name,
            "currency": self.currency,
            "total_cost": self.get_total_cost(),
            "total_requests": self.get_request_count(),
            "total_tokens": self.get_total_tokens(),
            "model_breakdown": self.get_model_breakdown(),
            "first_request": (
                self.usage_history[0].datetime.isoformat() if self.usage_history else None
            ),
            "last_request": (
                self.usage_history[-1].datetime.isoformat() if self.usage_history else None
            ),
        }

    def clear_history(self) -> None:
        """Clear all usage history."""
        self.usage_history.clear()

    def export_history(self) -> list[dict]:
        """
        Export usage history as list of dicts.

        Returns:
            List of usage records as dictionaries
        """
        return [
            {
                "timestamp": record.datetime.isoformat(),
                "model": record.model,
                "input_tokens": record.input_tokens,
                "output_tokens": record.output_tokens,
                "input_cost": record.input_cost,
                "output_cost": record.output_cost,
                "total_cost": record.total_cost,
                "metadata": record.metadata,
            }
            for record in self.usage_history
        ]

    def export_to_csv(self, filepath: str) -> None:
        """
        Export usage history to CSV file.

        Args:
            filepath: Path to output CSV file
        """
        import csv

        if not self.usage_history:
            return

        with open(filepath, "w", newline="", encoding="utf-8") as f:
            writer = csv.writer(f)
            writer.writerow(
                [
                    "Timestamp",
                    "Model",
                    "Input Tokens",
                    "Output Tokens",
                    "Total Tokens",
                    "Input Cost",
                    "Output Cost",
                    "Total Cost",
                ]
            )

            for record in self.usage_history:
                writer.writerow(
                    [
                        record.datetime.isoformat(),
                        record.model,
                        record.input_tokens,
                        record.output_tokens,
                        record.total_tokens,
                        f"{record.input_cost:.6f}",
                        f"{record.output_cost:.6f}",
                        f"{record.total_cost:.6f}",
                    ]
                )

    def export_to_json(self, filepath: str) -> None:
        """
        Export usage history to JSON file.

        Args:
            filepath: Path to output JSON file
        """
        import json

        data = {
            "provider": self.provider_name,
            "currency": self.currency,
            "export_date": datetime.now().isoformat(),
            "summary": self.get_usage_report(),
            "history": self.export_history(),
        }

        with open(filepath, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2)

    def get_time_range_report(
        self,
        start_time: Optional[float] = None,
        end_time: Optional[float] = None,
    ) -> dict:
        """
        Get usage report for a specific time range.

        Args:
            start_time: Start timestamp (inclusive)
            end_time: End timestamp (inclusive)

        Returns:
            Usage report for the time range
        """
        filtered_records = self.usage_history

        if start_time:
            filtered_records = [r for r in filtered_records if r.timestamp >= start_time]

        if end_time:
            filtered_records = [r for r in filtered_records if r.timestamp <= end_time]

        if not filtered_records:
            return {
                "total_cost": 0.0,
                "total_tokens": 0,
                "total_requests": 0,
                "model_breakdown": {},
            }

        total_cost = sum(r.total_cost for r in filtered_records)
        total_tokens = sum(r.total_tokens for r in filtered_records)

        # Model breakdown
        breakdown = {}
        for record in filtered_records:
            if record.model not in breakdown:
                breakdown[record.model] = {
                    "cost": 0.0,
                    "input_tokens": 0,
                    "output_tokens": 0,
                    "requests": 0,
                }

            breakdown[record.model]["cost"] += record.total_cost
            breakdown[record.model]["input_tokens"] += record.input_tokens
            breakdown[record.model]["output_tokens"] += record.output_tokens
            breakdown[record.model]["requests"] += 1

        return {
            "total_cost": total_cost,
            "total_tokens": total_tokens,
            "total_requests": len(filtered_records),
            "model_breakdown": breakdown,
            "start_time": (filtered_records[0].datetime.isoformat() if filtered_records else None),
            "end_time": (filtered_records[-1].datetime.isoformat() if filtered_records else None),
        }

    def check_budget(self, budget_limit: float) -> dict:
        """
        Check if spending is within budget.

        Args:
            budget_limit: Budget limit in currency units

        Returns:
            Budget status information
        """
        total_cost = self.get_total_cost()
        remaining = budget_limit - total_cost
        percent_used = (total_cost / budget_limit * 100) if budget_limit > 0 else 0

        return {
            "budget_limit": budget_limit,
            "total_spent": total_cost,
            "remaining": remaining,
            "percent_used": percent_used,
            "within_budget": total_cost <= budget_limit,
            "near_limit": percent_used >= 90,  # Warning at 90%
        }

__init__(provider_name, pricing, currency='USD')

Initialize cost tracker.

Parameters:

Name	Type	Description	Default
provider_name	str	Name of the provider	required
pricing	dict[str, dict[str, float]]	Pricing per model (per 1K tokens) Format: {"model_name": {"input": price, "output": price}}	required
currency	str	Currency code (default: "USD")	'USD'

Source code in bruno_llm/base/cost_tracker.py

def __init__(
    self,
    provider_name: str,
    pricing: dict[str, dict[str, float]],
    currency: str = "USD",
):
    """
    Initialize cost tracker.

    Args:
        provider_name: Name of the provider
        pricing: Pricing per model (per 1K tokens)
            Format: {"model_name": {"input": price, "output": price}}
        currency: Currency code (default: "USD")
    """
    self.provider_name = provider_name
    self.pricing = pricing
    self.currency = currency
    self.usage_history: list[UsageRecord] = []

track_request(model, input_tokens, output_tokens, metadata=None)

Track a single API request.

Parameters:

Name	Type	Description	Default
model	str	Model name	required
input_tokens	int	Number of input tokens	required
output_tokens	int	Number of output tokens	required
metadata	Optional[dict[str, str]]	Optional metadata about the request	None

Returns:

Type	Description
UsageRecord	UsageRecord with calculated costs

Source code in bruno_llm/base/cost_tracker.py

def track_request(
    self,
    model: str,
    input_tokens: int,
    output_tokens: int,
    metadata: Optional[dict[str, str]] = None,
) -> UsageRecord:
    """
    Track a single API request.

    Args:
        model: Model name
        input_tokens: Number of input tokens
        output_tokens: Number of output tokens
        metadata: Optional metadata about the request

    Returns:
        UsageRecord with calculated costs
    """
    # Get pricing for model (with fallback to default if available)
    model_pricing = self.pricing.get(model, self.pricing.get("default", {}))

    # Calculate costs (pricing is per 1K tokens)
    input_cost = (input_tokens / 1000.0) * model_pricing.get("input", 0.0)
    output_cost = (output_tokens / 1000.0) * model_pricing.get("output", 0.0)
    total_cost = input_cost + output_cost

    # Create record
    record = UsageRecord(
        timestamp=time.time(),
        model=model,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        input_cost=input_cost,
        output_cost=output_cost,
        total_cost=total_cost,
        metadata=metadata or {},
    )

    self.usage_history.append(record)
    return record

get_total_cost(model=None)

Get total cost across all requests.

Parameters:

Name	Type	Description	Default
model	Optional[str]	Optional model filter (None = all models)	None

Returns:

Type	Description
float	Total cost in the configured currency

Source code in bruno_llm/base/cost_tracker.py

def get_total_cost(self, model: Optional[str] = None) -> float:
    """
    Get total cost across all requests.

    Args:
        model: Optional model filter (None = all models)

    Returns:
        Total cost in the configured currency
    """
    total = 0.0
    for record in self.usage_history:
        if model is None or record.model == model:
            total += record.total_cost
    return total

get_total_tokens(model=None)

Get total tokens used.

Parameters:

Name	Type	Description	Default
model	Optional[str]	Optional model filter (None = all models)	None

Returns:

Type	Description
dict[str, int]	Dict with input, output, and total token counts

Source code in bruno_llm/base/cost_tracker.py

def get_total_tokens(self, model: Optional[str] = None) -> dict[str, int]:
    """
    Get total tokens used.

    Args:
        model: Optional model filter (None = all models)

    Returns:
        Dict with input, output, and total token counts
    """
    input_tokens = 0
    output_tokens = 0

    for record in self.usage_history:
        if model is None or record.model == model:
            input_tokens += record.input_tokens
            output_tokens += record.output_tokens

    return {
        "input": input_tokens,
        "output": output_tokens,
        "total": input_tokens + output_tokens,
    }

get_request_count(model=None)

Get number of requests made.

Parameters:

Name	Type	Description	Default
model	Optional[str]	Optional model filter (None = all models)	None

Returns:

Type	Description
int	Number of requests

Source code in bruno_llm/base/cost_tracker.py

def get_request_count(self, model: Optional[str] = None) -> int:
    """
    Get number of requests made.

    Args:
        model: Optional model filter (None = all models)

    Returns:
        Number of requests
    """
    if model is None:
        return len(self.usage_history)
    return sum(1 for r in self.usage_history if r.model == model)

get_model_breakdown()

Get cost breakdown by model.

Returns:

Type	Description
dict[str, dict[str, float]]	Dict mapping model names to their usage statistics

Source code in bruno_llm/base/cost_tracker.py

def get_model_breakdown(self) -> dict[str, dict[str, float]]:
    """
    Get cost breakdown by model.

    Returns:
        Dict mapping model names to their usage statistics
    """
    breakdown: dict[str, dict[str, float]] = {}

    for record in self.usage_history:
        if record.model not in breakdown:
            breakdown[record.model] = {
                "cost": 0.0,
                "input_tokens": 0,
                "output_tokens": 0,
                "requests": 0,
            }

        breakdown[record.model]["cost"] += record.total_cost
        breakdown[record.model]["input_tokens"] += record.input_tokens
        breakdown[record.model]["output_tokens"] += record.output_tokens
        breakdown[record.model]["requests"] += 1

    return breakdown

get_usage_report()

Get comprehensive usage report.

Returns:

Type	Description
dict	Dict with complete usage statistics

Source code in bruno_llm/base/cost_tracker.py

def get_usage_report(self) -> dict:
    """
    Get comprehensive usage report.

    Returns:
        Dict with complete usage statistics
    """
    return {
        "provider": self.provider_name,
        "currency": self.currency,
        "total_cost": self.get_total_cost(),
        "total_requests": self.get_request_count(),
        "total_tokens": self.get_total_tokens(),
        "model_breakdown": self.get_model_breakdown(),
        "first_request": (
            self.usage_history[0].datetime.isoformat() if self.usage_history else None
        ),
        "last_request": (
            self.usage_history[-1].datetime.isoformat() if self.usage_history else None
        ),
    }

clear_history()

Clear all usage history.

Source code in bruno_llm/base/cost_tracker.py

def clear_history(self) -> None:
    """Clear all usage history."""
    self.usage_history.clear()

export_history()

Export usage history as list of dicts.

Returns:

Type	Description
list[dict]	List of usage records as dictionaries

Source code in bruno_llm/base/cost_tracker.py

def export_history(self) -> list[dict]:
    """
    Export usage history as list of dicts.

    Returns:
        List of usage records as dictionaries
    """
    return [
        {
            "timestamp": record.datetime.isoformat(),
            "model": record.model,
            "input_tokens": record.input_tokens,
            "output_tokens": record.output_tokens,
            "input_cost": record.input_cost,
            "output_cost": record.output_cost,
            "total_cost": record.total_cost,
            "metadata": record.metadata,
        }
        for record in self.usage_history
    ]

export_to_csv(filepath)

Export usage history to CSV file.

Parameters:

Name	Type	Description	Default
filepath	str	Path to output CSV file	required

Source code in bruno_llm/base/cost_tracker.py

def export_to_csv(self, filepath: str) -> None:
    """
    Export usage history to CSV file.

    Args:
        filepath: Path to output CSV file
    """
    import csv

    if not self.usage_history:
        return

    with open(filepath, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(
            [
                "Timestamp",
                "Model",
                "Input Tokens",
                "Output Tokens",
                "Total Tokens",
                "Input Cost",
                "Output Cost",
                "Total Cost",
            ]
        )

        for record in self.usage_history:
            writer.writerow(
                [
                    record.datetime.isoformat(),
                    record.model,
                    record.input_tokens,
                    record.output_tokens,
                    record.total_tokens,
                    f"{record.input_cost:.6f}",
                    f"{record.output_cost:.6f}",
                    f"{record.total_cost:.6f}",
                ]
            )

export_to_json(filepath)

Export usage history to JSON file.

Parameters:

Name	Type	Description	Default
filepath	str	Path to output JSON file	required

Source code in bruno_llm/base/cost_tracker.py

def export_to_json(self, filepath: str) -> None:
    """
    Export usage history to JSON file.

    Args:
        filepath: Path to output JSON file
    """
    import json

    data = {
        "provider": self.provider_name,
        "currency": self.currency,
        "export_date": datetime.now().isoformat(),
        "summary": self.get_usage_report(),
        "history": self.export_history(),
    }

    with open(filepath, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)

get_time_range_report(start_time=None, end_time=None)

Get usage report for a specific time range.

Parameters:

Name	Type	Description	Default
start_time	Optional[float]	Start timestamp (inclusive)	None
end_time	Optional[float]	End timestamp (inclusive)	None

Returns:

Type	Description
dict	Usage report for the time range

Source code in bruno_llm/base/cost_tracker.py

def get_time_range_report(
    self,
    start_time: Optional[float] = None,
    end_time: Optional[float] = None,
) -> dict:
    """
    Get usage report for a specific time range.

    Args:
        start_time: Start timestamp (inclusive)
        end_time: End timestamp (inclusive)

    Returns:
        Usage report for the time range
    """
    filtered_records = self.usage_history

    if start_time:
        filtered_records = [r for r in filtered_records if r.timestamp >= start_time]

    if end_time:
        filtered_records = [r for r in filtered_records if r.timestamp <= end_time]

    if not filtered_records:
        return {
            "total_cost": 0.0,
            "total_tokens": 0,
            "total_requests": 0,
            "model_breakdown": {},
        }

    total_cost = sum(r.total_cost for r in filtered_records)
    total_tokens = sum(r.total_tokens for r in filtered_records)

    # Model breakdown
    breakdown = {}
    for record in filtered_records:
        if record.model not in breakdown:
            breakdown[record.model] = {
                "cost": 0.0,
                "input_tokens": 0,
                "output_tokens": 0,
                "requests": 0,
            }

        breakdown[record.model]["cost"] += record.total_cost
        breakdown[record.model]["input_tokens"] += record.input_tokens
        breakdown[record.model]["output_tokens"] += record.output_tokens
        breakdown[record.model]["requests"] += 1

    return {
        "total_cost": total_cost,
        "total_tokens": total_tokens,
        "total_requests": len(filtered_records),
        "model_breakdown": breakdown,
        "start_time": (filtered_records[0].datetime.isoformat() if filtered_records else None),
        "end_time": (filtered_records[-1].datetime.isoformat() if filtered_records else None),
    }

check_budget(budget_limit)

Check if spending is within budget.

Parameters:

Name	Type	Description	Default
budget_limit	float	Budget limit in currency units	required

Returns:

Type	Description
dict	Budget status information

Source code in bruno_llm/base/cost_tracker.py

def check_budget(self, budget_limit: float) -> dict:
    """
    Check if spending is within budget.

    Args:
        budget_limit: Budget limit in currency units

    Returns:
        Budget status information
    """
    total_cost = self.get_total_cost()
    remaining = budget_limit - total_cost
    percent_used = (total_cost / budget_limit * 100) if budget_limit > 0 else 0

    return {
        "budget_limit": budget_limit,
        "total_spent": total_cost,
        "remaining": remaining,
        "percent_used": percent_used,
        "within_budget": total_cost <= budget_limit,
        "near_limit": percent_used >= 90,  # Warning at 90%
    }

UsageRecord dataclass

Record of a single API usage event.

Attributes:

Name	Type	Description
timestamp	float	When the request was made
model	str	Model name used
input_tokens	int	Number of input tokens
output_tokens	int	Number of output tokens
input_cost	float	Cost for input tokens
output_cost	float	Cost for output tokens
total_cost	float	Total cost for this request

Source code in bruno_llm/base/cost_tracker.py

@dataclass
class UsageRecord:
    """
    Record of a single API usage event.

    Attributes:
        timestamp: When the request was made
        model: Model name used
        input_tokens: Number of input tokens
        output_tokens: Number of output tokens
        input_cost: Cost for input tokens
        output_cost: Cost for output tokens
        total_cost: Total cost for this request
    """

    timestamp: float
    model: str
    input_tokens: int
    output_tokens: int
    input_cost: float
    output_cost: float
    total_cost: float
    metadata: dict[str, str] = field(default_factory=dict)

    @property
    def datetime(self) -> datetime:
        """Get datetime from timestamp."""
        return datetime.fromtimestamp(self.timestamp)

    @property
    def total_tokens(self) -> int:
        """Get total tokens used."""
        return self.input_tokens + self.output_tokens

datetime property

Get datetime from timestamp.

total_tokens property

Get total tokens used.

BaseEmbeddingProvider

Bases: EmbeddingInterface, ABC

Base class for embedding providers using numpy for vector operations.

This base class implements the EmbeddingInterface from bruno-core and provides: - Efficient vector operations using numpy - Input validation and error handling - Similarity calculation using numpy's optimized functions - Batch processing utilities

All embedding providers should inherit from this class and implement the abstract methods for their specific API or model.

Parameters:

Name	Type	Description	Default
model	str	Model name or identifier	required
timeout	float	Request timeout in seconds	30.0

Examples:

>>> class MyEmbeddingProvider(BaseEmbeddingProvider):
...     async def embed_text(self, text: str) -> List[float]:
...         # Implementation specific to your provider
...         return await self._call_api(text)
...
...     def get_dimension(self) -> int:
...         return 768  # Your model's dimension

Source code in bruno_llm/base/embedding_interface.py

class BaseEmbeddingProvider(EmbeddingInterface, ABC):
    """
    Base class for embedding providers using numpy for vector operations.

    This base class implements the EmbeddingInterface from bruno-core and provides:
    - Efficient vector operations using numpy
    - Input validation and error handling
    - Similarity calculation using numpy's optimized functions
    - Batch processing utilities

    All embedding providers should inherit from this class and implement the
    abstract methods for their specific API or model.

    Args:
        model: Model name or identifier
        timeout: Request timeout in seconds

    Examples:
        >>> class MyEmbeddingProvider(BaseEmbeddingProvider):
        ...     async def embed_text(self, text: str) -> List[float]:
        ...         # Implementation specific to your provider
        ...         return await self._call_api(text)
        ...
        ...     def get_dimension(self) -> int:
        ...         return 768  # Your model's dimension
    """

    def __init__(
        self,
        model: str,
        timeout: float = 30.0,
        **kwargs,
    ):
        self.model = model
        self.timeout = timeout
        self.config = kwargs
        self._dimension_cache: Optional[int] = None

    @abstractmethod
    async def embed_text(self, text: str) -> list[float]:
        """
        Generate embedding for a single text.

        Args:
            text: Input text to embed

        Returns:
            Embedding vector as list of floats

        Raises:
            LLMError: If embedding generation fails
        """
        pass

    async def embed_texts(self, texts: list[str]) -> list[list[float]]:
        """
        Generate embeddings for multiple texts.

        Default implementation calls embed_text for each text.
        Override for providers with native batch support.

        Args:
            texts: List of texts to embed

        Returns:
            List of embedding vectors

        Raises:
            LLMError: If any embedding generation fails
        """
        embeddings = []
        for text in texts:
            embedding = await self.embed_text(text)
            embeddings.append(embedding)
        return embeddings

    async def embed_message(self, message: Message) -> list[float]:
        """
        Generate embedding for a message.

        Args:
            message: Message object to embed

        Returns:
            Embedding vector for the message content
        """
        return await self.embed_text(message.content)

    @abstractmethod
    def get_dimension(self) -> int:
        """
        Get the embedding dimension for this provider.

        Returns:
            Number of dimensions in embeddings
        """
        pass

    def get_model_name(self) -> str:
        """
        Get the model name used by this provider.

        Returns:
            Model name string
        """
        return self.model

    def calculate_similarity(self, embedding1: list[float], embedding2: list[float]) -> float:
        """
        Calculate cosine similarity between two embeddings using numpy.

        Uses numpy's optimized dot product and norm calculations for efficiency.

        Args:
            embedding1: First embedding vector
            embedding2: Second embedding vector

        Returns:
            Cosine similarity (-1.0 to 1.0)

        Raises:
            ValueError: If vectors have different dimensions
        """
        # Convert to numpy arrays for efficient computation
        vec1 = np.array(embedding1, dtype=np.float32)
        vec2 = np.array(embedding2, dtype=np.float32)

        if vec1.shape != vec2.shape:
            raise ValueError(f"Vector dimensions don't match: {vec1.shape} != {vec2.shape}")

        # Handle edge cases
        if vec1.size == 0 or vec2.size == 0:
            return 0.0

        # Calculate cosine similarity using numpy
        dot_product = np.dot(vec1, vec2)
        norm1 = np.linalg.norm(vec1)
        norm2 = np.linalg.norm(vec2)

        # Handle zero vectors
        if norm1 == 0.0 or norm2 == 0.0:
            return 0.0

        similarity = dot_product / (norm1 * norm2)

        # Ensure result is in valid range due to floating point precision
        return float(np.clip(similarity, -1.0, 1.0))

    @abstractmethod
    async def check_connection(self) -> bool:
        """
        Check if the provider is accessible.

        Returns:
            True if provider can be used, False otherwise
        """
        pass

    # Utility methods for common operations

    def validate_embedding(
        self, embedding: list[float], expected_dimension: Optional[int] = None
    ) -> None:
        """
        Validate an embedding vector using numpy.

        Args:
            embedding: Embedding to validate
            expected_dimension: Expected dimension (uses get_dimension() if None)

        Raises:
            ValueError: If embedding is invalid
        """
        if not isinstance(embedding, list):
            raise ValueError(f"Embedding must be a list, got {type(embedding)}")

        if len(embedding) == 0:
            raise ValueError("Embedding cannot be empty")

        # Convert to numpy for validation
        vec = np.array(embedding, dtype=np.float32)

        # Check for invalid values using numpy
        if np.any(np.isnan(vec)):
            raise ValueError("Embedding contains NaN values")

        if np.any(np.isinf(vec)):
            raise ValueError("Embedding contains infinite values")

        # Check dimension
        expected_dim = expected_dimension or self.get_dimension()
        if len(embedding) != expected_dim:
            raise ValueError(
                f"Embedding dimension mismatch: expected {expected_dim}, got {len(embedding)}"
            )

    def batch_cosine_similarity(
        self, embeddings1: list[list[float]], embeddings2: list[list[float]]
    ) -> np.ndarray:
        """
        Calculate pairwise cosine similarities between two sets of embeddings.

        Uses numpy for efficient batch computation.

        Args:
            embeddings1: First set of embeddings
            embeddings2: Second set of embeddings

        Returns:
            Similarity matrix as numpy array (len(embeddings1) x len(embeddings2))
        """
        # Convert to numpy arrays
        matrix1 = np.array(embeddings1, dtype=np.float32)
        matrix2 = np.array(embeddings2, dtype=np.float32)

        # Normalize vectors
        norm1 = np.linalg.norm(matrix1, axis=1, keepdims=True)
        norm2 = np.linalg.norm(matrix2, axis=1, keepdims=True)

        # Handle zero vectors
        norm1 = np.where(norm1 == 0, 1, norm1)
        norm2 = np.where(norm2 == 0, 1, norm2)

        normalized1 = matrix1 / norm1
        normalized2 = matrix2 / norm2

        # Calculate cosine similarity matrix
        similarity_matrix = np.dot(normalized1, normalized2.T)

        return similarity_matrix

    def find_most_similar(
        self,
        query_embedding: list[float],
        candidate_embeddings: list[list[float]],
        top_k: Optional[int] = None,
    ) -> list[tuple[int, float]]:
        """
        Find most similar embeddings to a query using numpy.

        Args:
            query_embedding: Query embedding vector
            candidate_embeddings: List of candidate embeddings
            top_k: Number of top results (None = all)

        Returns:
            List of (index, similarity) tuples sorted by similarity (descending)
        """
        if not candidate_embeddings:
            return []

        # Calculate similarities using numpy
        query_vec = np.array(query_embedding, dtype=np.float32).reshape(1, -1)
        candidates_matrix = np.array(candidate_embeddings, dtype=np.float32)

        # Use batch similarity calculation
        similarities = self.batch_cosine_similarity(query_vec.tolist(), candidates_matrix.tolist())[
            0
        ]  # Get first (and only) row

        # Create index-similarity pairs
        indexed_similarities = list(enumerate(similarities.tolist()))

        # Sort by similarity (descending)
        indexed_similarities.sort(key=lambda x: x[1], reverse=True)

        # Return top-k results
        if top_k is not None:
            indexed_similarities = indexed_similarities[:top_k]

        return indexed_similarities

    def average_embeddings(self, embeddings: list[list[float]]) -> list[float]:
        """
        Calculate average embedding using numpy.

        Args:
            embeddings: List of embedding vectors

        Returns:
            Average embedding vector

        Raises:
            ValueError: If embeddings list is empty
        """
        if not embeddings:
            raise ValueError("Cannot average empty list of embeddings")

        # Use numpy for efficient averaging
        matrix = np.array(embeddings, dtype=np.float32)
        average = np.mean(matrix, axis=0)

        return average.tolist()

    def weighted_average_embeddings(
        self,
        embeddings: list[list[float]],
        weights: list[float],
    ) -> list[float]:
        """
        Calculate weighted average embedding using numpy.

        Args:
            embeddings: List of embedding vectors
            weights: List of weights (must sum to 1.0)

        Returns:
            Weighted average embedding

        Raises:
            ValueError: If inputs are invalid
        """
        if len(embeddings) != len(weights):
            raise ValueError(
                f"Number of embeddings ({len(embeddings)}) must match "
                f"number of weights ({len(weights)})"
            )

        if not embeddings:
            raise ValueError("Cannot average empty list of embeddings")

        # Validate weights sum using numpy
        weights_array = np.array(weights, dtype=np.float32)
        if not np.isclose(weights_array.sum(), 1.0, atol=1e-6):
            raise ValueError(f"Weights must sum to 1.0, got {weights_array.sum()}")

        # Calculate weighted average using numpy
        embeddings_matrix = np.array(embeddings, dtype=np.float32)
        weighted_avg = np.average(embeddings_matrix, axis=0, weights=weights_array)

        return weighted_avg.tolist()

embed_text(text) abstractmethod async

Generate embedding for a single text.

Parameters:

Name	Type	Description	Default
text	str	Input text to embed	required

Returns:

Type	Description
list[float]	Embedding vector as list of floats

Raises:

Type	Description
LLMError	If embedding generation fails

Source code in bruno_llm/base/embedding_interface.py

@abstractmethod
async def embed_text(self, text: str) -> list[float]:
    """
    Generate embedding for a single text.

    Args:
        text: Input text to embed

    Returns:
        Embedding vector as list of floats

    Raises:
        LLMError: If embedding generation fails
    """
    pass

embed_texts(texts) async

Generate embeddings for multiple texts.

Default implementation calls embed_text for each text. Override for providers with native batch support.

Parameters:

Name	Type	Description	Default
texts	list[str]	List of texts to embed	required

Returns:

Type	Description
list[list[float]]	List of embedding vectors

Raises:

Type	Description
LLMError	If any embedding generation fails

Source code in bruno_llm/base/embedding_interface.py

async def embed_texts(self, texts: list[str]) -> list[list[float]]:
    """
    Generate embeddings for multiple texts.

    Default implementation calls embed_text for each text.
    Override for providers with native batch support.

    Args:
        texts: List of texts to embed

    Returns:
        List of embedding vectors

    Raises:
        LLMError: If any embedding generation fails
    """
    embeddings = []
    for text in texts:
        embedding = await self.embed_text(text)
        embeddings.append(embedding)
    return embeddings

embed_message(message) async

Generate embedding for a message.

Parameters:

Name	Type	Description	Default
message	Message	Message object to embed	required

Returns:

Type	Description
list[float]	Embedding vector for the message content

Source code in bruno_llm/base/embedding_interface.py

async def embed_message(self, message: Message) -> list[float]:
    """
    Generate embedding for a message.

    Args:
        message: Message object to embed

    Returns:
        Embedding vector for the message content
    """
    return await self.embed_text(message.content)

get_dimension() abstractmethod

Get the embedding dimension for this provider.

Returns:

Type	Description
int	Number of dimensions in embeddings

Source code in bruno_llm/base/embedding_interface.py

@abstractmethod
def get_dimension(self) -> int:
    """
    Get the embedding dimension for this provider.

    Returns:
        Number of dimensions in embeddings
    """
    pass

get_model_name()

Get the model name used by this provider.

Returns:

Type	Description
str	Model name string

Source code in bruno_llm/base/embedding_interface.py

def get_model_name(self) -> str:
    """
    Get the model name used by this provider.

    Returns:
        Model name string
    """
    return self.model

calculate_similarity(embedding1, embedding2)

Calculate cosine similarity between two embeddings using numpy.

Uses numpy's optimized dot product and norm calculations for efficiency.

Parameters:

Name	Type	Description	Default
embedding1	list[float]	First embedding vector	required
embedding2	list[float]	Second embedding vector	required

Returns:

Type	Description
float	Cosine similarity (-1.0 to 1.0)

Raises:

Type	Description
ValueError	If vectors have different dimensions

Source code in bruno_llm/base/embedding_interface.py

def calculate_similarity(self, embedding1: list[float], embedding2: list[float]) -> float:
    """
    Calculate cosine similarity between two embeddings using numpy.

    Uses numpy's optimized dot product and norm calculations for efficiency.

    Args:
        embedding1: First embedding vector
        embedding2: Second embedding vector

    Returns:
        Cosine similarity (-1.0 to 1.0)

    Raises:
        ValueError: If vectors have different dimensions
    """
    # Convert to numpy arrays for efficient computation
    vec1 = np.array(embedding1, dtype=np.float32)
    vec2 = np.array(embedding2, dtype=np.float32)

    if vec1.shape != vec2.shape:
        raise ValueError(f"Vector dimensions don't match: {vec1.shape} != {vec2.shape}")

    # Handle edge cases
    if vec1.size == 0 or vec2.size == 0:
        return 0.0

    # Calculate cosine similarity using numpy
    dot_product = np.dot(vec1, vec2)
    norm1 = np.linalg.norm(vec1)
    norm2 = np.linalg.norm(vec2)

    # Handle zero vectors
    if norm1 == 0.0 or norm2 == 0.0:
        return 0.0

    similarity = dot_product / (norm1 * norm2)

    # Ensure result is in valid range due to floating point precision
    return float(np.clip(similarity, -1.0, 1.0))

check_connection() abstractmethod async

Check if the provider is accessible.

Returns:

Type	Description
bool	True if provider can be used, False otherwise

Source code in bruno_llm/base/embedding_interface.py

@abstractmethod
async def check_connection(self) -> bool:
    """
    Check if the provider is accessible.

    Returns:
        True if provider can be used, False otherwise
    """
    pass

validate_embedding(embedding, expected_dimension=None)

Validate an embedding vector using numpy.

Parameters:

Name	Type	Description	Default
embedding	list[float]	Embedding to validate	required
expected_dimension	Optional[int]	Expected dimension (uses get_dimension() if None)	None

Raises:

Type	Description
ValueError	If embedding is invalid

Source code in bruno_llm/base/embedding_interface.py

def validate_embedding(
    self, embedding: list[float], expected_dimension: Optional[int] = None
) -> None:
    """
    Validate an embedding vector using numpy.

    Args:
        embedding: Embedding to validate
        expected_dimension: Expected dimension (uses get_dimension() if None)

    Raises:
        ValueError: If embedding is invalid
    """
    if not isinstance(embedding, list):
        raise ValueError(f"Embedding must be a list, got {type(embedding)}")

    if len(embedding) == 0:
        raise ValueError("Embedding cannot be empty")

    # Convert to numpy for validation
    vec = np.array(embedding, dtype=np.float32)

    # Check for invalid values using numpy
    if np.any(np.isnan(vec)):
        raise ValueError("Embedding contains NaN values")

    if np.any(np.isinf(vec)):
        raise ValueError("Embedding contains infinite values")

    # Check dimension
    expected_dim = expected_dimension or self.get_dimension()
    if len(embedding) != expected_dim:
        raise ValueError(
            f"Embedding dimension mismatch: expected {expected_dim}, got {len(embedding)}"
        )

batch_cosine_similarity(embeddings1, embeddings2)

Calculate pairwise cosine similarities between two sets of embeddings.

Uses numpy for efficient batch computation.

Parameters:

Name	Type	Description	Default
embeddings1	list[list[float]]	First set of embeddings	required
embeddings2	list[list[float]]	Second set of embeddings	required

Returns:

Type	Description
ndarray	Similarity matrix as numpy array (len(embeddings1) x len(embeddings2))

Source code in bruno_llm/base/embedding_interface.py

def batch_cosine_similarity(
    self, embeddings1: list[list[float]], embeddings2: list[list[float]]
) -> np.ndarray:
    """
    Calculate pairwise cosine similarities between two sets of embeddings.

    Uses numpy for efficient batch computation.

    Args:
        embeddings1: First set of embeddings
        embeddings2: Second set of embeddings

    Returns:
        Similarity matrix as numpy array (len(embeddings1) x len(embeddings2))
    """
    # Convert to numpy arrays
    matrix1 = np.array(embeddings1, dtype=np.float32)
    matrix2 = np.array(embeddings2, dtype=np.float32)

    # Normalize vectors
    norm1 = np.linalg.norm(matrix1, axis=1, keepdims=True)
    norm2 = np.linalg.norm(matrix2, axis=1, keepdims=True)

    # Handle zero vectors
    norm1 = np.where(norm1 == 0, 1, norm1)
    norm2 = np.where(norm2 == 0, 1, norm2)

    normalized1 = matrix1 / norm1
    normalized2 = matrix2 / norm2

    # Calculate cosine similarity matrix
    similarity_matrix = np.dot(normalized1, normalized2.T)

    return similarity_matrix

find_most_similar(query_embedding, candidate_embeddings, top_k=None)

Find most similar embeddings to a query using numpy.

Parameters:

Name	Type	Description	Default
query_embedding	list[float]	Query embedding vector	required
candidate_embeddings	list[list[float]]	List of candidate embeddings	required
top_k	Optional[int]	Number of top results (None = all)	None

Returns:

Type	Description
list[tuple[int, float]]	List of (index, similarity) tuples sorted by similarity (descending)

Source code in bruno_llm/base/embedding_interface.py

def find_most_similar(
    self,
    query_embedding: list[float],
    candidate_embeddings: list[list[float]],
    top_k: Optional[int] = None,
) -> list[tuple[int, float]]:
    """
    Find most similar embeddings to a query using numpy.

    Args:
        query_embedding: Query embedding vector
        candidate_embeddings: List of candidate embeddings
        top_k: Number of top results (None = all)

    Returns:
        List of (index, similarity) tuples sorted by similarity (descending)
    """
    if not candidate_embeddings:
        return []

    # Calculate similarities using numpy
    query_vec = np.array(query_embedding, dtype=np.float32).reshape(1, -1)
    candidates_matrix = np.array(candidate_embeddings, dtype=np.float32)

    # Use batch similarity calculation
    similarities = self.batch_cosine_similarity(query_vec.tolist(), candidates_matrix.tolist())[
        0
    ]  # Get first (and only) row

    # Create index-similarity pairs
    indexed_similarities = list(enumerate(similarities.tolist()))

    # Sort by similarity (descending)
    indexed_similarities.sort(key=lambda x: x[1], reverse=True)

    # Return top-k results
    if top_k is not None:
        indexed_similarities = indexed_similarities[:top_k]

    return indexed_similarities

average_embeddings(embeddings)

Calculate average embedding using numpy.

Parameters:

Name	Type	Description	Default
embeddings	list[list[float]]	List of embedding vectors	required

Returns:

Type	Description
list[float]	Average embedding vector

Raises:

Type	Description
ValueError	If embeddings list is empty

Source code in bruno_llm/base/embedding_interface.py

def average_embeddings(self, embeddings: list[list[float]]) -> list[float]:
    """
    Calculate average embedding using numpy.

    Args:
        embeddings: List of embedding vectors

    Returns:
        Average embedding vector

    Raises:
        ValueError: If embeddings list is empty
    """
    if not embeddings:
        raise ValueError("Cannot average empty list of embeddings")

    # Use numpy for efficient averaging
    matrix = np.array(embeddings, dtype=np.float32)
    average = np.mean(matrix, axis=0)

    return average.tolist()

weighted_average_embeddings(embeddings, weights)

Calculate weighted average embedding using numpy.

Parameters:

Name	Type	Description	Default
embeddings	list[list[float]]	List of embedding vectors	required
weights	list[float]	List of weights (must sum to 1.0)	required

Returns:

Type	Description
list[float]	Weighted average embedding

Raises:

Type	Description
ValueError	If inputs are invalid

Source code in bruno_llm/base/embedding_interface.py

def weighted_average_embeddings(
    self,
    embeddings: list[list[float]],
    weights: list[float],
) -> list[float]:
    """
    Calculate weighted average embedding using numpy.

    Args:
        embeddings: List of embedding vectors
        weights: List of weights (must sum to 1.0)

    Returns:
        Weighted average embedding

    Raises:
        ValueError: If inputs are invalid
    """
    if len(embeddings) != len(weights):
        raise ValueError(
            f"Number of embeddings ({len(embeddings)}) must match "
            f"number of weights ({len(weights)})"
        )

    if not embeddings:
        raise ValueError("Cannot average empty list of embeddings")

    # Validate weights sum using numpy
    weights_array = np.array(weights, dtype=np.float32)
    if not np.isclose(weights_array.sum(), 1.0, atol=1e-6):
        raise ValueError(f"Weights must sum to 1.0, got {weights_array.sum()}")

    # Calculate weighted average using numpy
    embeddings_matrix = np.array(embeddings, dtype=np.float32)
    weighted_avg = np.average(embeddings_matrix, axis=0, weights=weights_array)

    return weighted_avg.tolist()

CachingMiddleware

Bases: Middleware

Cache responses using ResponseCache.

Parameters:

Name	Type	Description	Default
cache		ResponseCache instance	required
cache_streaming	bool	Whether to cache streaming responses	True

Example

from bruno_llm.base.cache import ResponseCache cache = ResponseCache(max_size=100, ttl=300) middleware = CachingMiddleware(cache)

Source code in bruno_llm/base/middleware.py

class CachingMiddleware(Middleware):
    """
    Cache responses using ResponseCache.

    Args:
        cache: ResponseCache instance
        cache_streaming: Whether to cache streaming responses

    Example:
        >>> from bruno_llm.base.cache import ResponseCache
        >>> cache = ResponseCache(max_size=100, ttl=300)
        >>> middleware = CachingMiddleware(cache)
    """

    def __init__(self, cache, cache_streaming: bool = True):
        """
        Initialize caching middleware.

        Args:
            cache: ResponseCache instance
            cache_streaming: Whether to cache streaming responses
        """
        self.cache = cache
        self.cache_streaming = cache_streaming
        self._current_stream_chunks: Optional[list[str]] = None

    async def before_request(
        self, messages: list[Message], **kwargs: Any
    ) -> tuple[list[Message], dict[str, Any]]:
        """Check cache before request."""
        # Cache lookup is handled externally
        return messages, kwargs

    async def after_response(self, messages: list[Message], response: str, **kwargs: Any) -> str:
        """Cache response after receiving."""
        self.cache.set(messages, response, **kwargs)
        return response

    async def on_stream_chunk(self, chunk: str, **kwargs: Any) -> str:
        """Collect stream chunks for caching."""
        if self.cache_streaming:
            if self._current_stream_chunks is None:
                self._current_stream_chunks = []
            self._current_stream_chunks.append(chunk)

        return chunk

__init__(cache, cache_streaming=True)

Initialize caching middleware.

Parameters:

Name	Type	Description	Default
cache		ResponseCache instance	required
cache_streaming	bool	Whether to cache streaming responses	True

Source code in bruno_llm/base/middleware.py

def __init__(self, cache, cache_streaming: bool = True):
    """
    Initialize caching middleware.

    Args:
        cache: ResponseCache instance
        cache_streaming: Whether to cache streaming responses
    """
    self.cache = cache
    self.cache_streaming = cache_streaming
    self._current_stream_chunks: Optional[list[str]] = None

before_request(messages, **kwargs) async

Check cache before request.

Source code in bruno_llm/base/middleware.py

async def before_request(
    self, messages: list[Message], **kwargs: Any
) -> tuple[list[Message], dict[str, Any]]:
    """Check cache before request."""
    # Cache lookup is handled externally
    return messages, kwargs

after_response(messages, response, **kwargs) async

Cache response after receiving.

Source code in bruno_llm/base/middleware.py

async def after_response(self, messages: list[Message], response: str, **kwargs: Any) -> str:
    """Cache response after receiving."""
    self.cache.set(messages, response, **kwargs)
    return response

on_stream_chunk(chunk, **kwargs) async

Collect stream chunks for caching.

Source code in bruno_llm/base/middleware.py

async def on_stream_chunk(self, chunk: str, **kwargs: Any) -> str:
    """Collect stream chunks for caching."""
    if self.cache_streaming:
        if self._current_stream_chunks is None:
            self._current_stream_chunks = []
        self._current_stream_chunks.append(chunk)

    return chunk

LoggingMiddleware

Bases: Middleware

Log all requests and responses.

Parameters:

Name	Type	Description	Default
logger		Logger instance (defaults to structlog)	None
log_messages	bool	Whether to log full message content	False

Example

middleware = LoggingMiddleware(log_messages=False) provider = MiddlewareProvider(base_provider, [middleware])

Source code in bruno_llm/base/middleware.py

class LoggingMiddleware(Middleware):
    """
    Log all requests and responses.

    Args:
        logger: Logger instance (defaults to structlog)
        log_messages: Whether to log full message content

    Example:
        >>> middleware = LoggingMiddleware(log_messages=False)
        >>> provider = MiddlewareProvider(base_provider, [middleware])
    """

    def __init__(self, logger=None, log_messages: bool = False):
        """
        Initialize logging middleware.

        Args:
            logger: Logger instance
            log_messages: Whether to log message content
        """
        self.log_messages = log_messages

        if logger is None:
            import structlog

            self.logger = structlog.get_logger(__name__)
        else:
            self.logger = logger

    async def before_request(
        self, messages: list[Message], **kwargs: Any
    ) -> tuple[list[Message], dict[str, Any]]:
        """Log before request."""
        log_data = {
            "event": "llm_request",
            "message_count": len(messages),
            "params": {k: v for k, v in kwargs.items() if k not in ["api_key"]},
        }

        if self.log_messages:
            log_data["messages"] = [{"role": m.role.value, "content": m.content} for m in messages]

        self.logger.info("LLM request", **log_data)
        return messages, kwargs

    async def after_response(self, messages: list[Message], response: str, **kwargs: Any) -> str:
        """Log after response."""
        log_data = {
            "event": "llm_response",
            "response_length": len(response),
        }

        if self.log_messages:
            log_data["response"] = response

        self.logger.info("LLM response", **log_data)
        return response

    async def on_error(self, error: Exception, messages: list[Message], **kwargs: Any) -> None:
        """Log errors."""
        self.logger.error(
            "LLM error",
            event="llm_error",
            error=str(error),
            error_type=type(error).__name__,
            message_count=len(messages),
        )

__init__(logger=None, log_messages=False)

Initialize logging middleware.

Parameters:

Name	Type	Description	Default
logger		Logger instance	None
log_messages	bool	Whether to log message content	False

Source code in bruno_llm/base/middleware.py

def __init__(self, logger=None, log_messages: bool = False):
    """
    Initialize logging middleware.

    Args:
        logger: Logger instance
        log_messages: Whether to log message content
    """
    self.log_messages = log_messages

    if logger is None:
        import structlog

        self.logger = structlog.get_logger(__name__)
    else:
        self.logger = logger

before_request(messages, **kwargs) async

Log before request.

Source code in bruno_llm/base/middleware.py

async def before_request(
    self, messages: list[Message], **kwargs: Any
) -> tuple[list[Message], dict[str, Any]]:
    """Log before request."""
    log_data = {
        "event": "llm_request",
        "message_count": len(messages),
        "params": {k: v for k, v in kwargs.items() if k not in ["api_key"]},
    }

    if self.log_messages:
        log_data["messages"] = [{"role": m.role.value, "content": m.content} for m in messages]

    self.logger.info("LLM request", **log_data)
    return messages, kwargs

after_response(messages, response, **kwargs) async

Log after response.

Source code in bruno_llm/base/middleware.py

async def after_response(self, messages: list[Message], response: str, **kwargs: Any) -> str:
    """Log after response."""
    log_data = {
        "event": "llm_response",
        "response_length": len(response),
    }

    if self.log_messages:
        log_data["response"] = response

    self.logger.info("LLM response", **log_data)
    return response

on_error(error, messages, **kwargs) async

Log errors.

Source code in bruno_llm/base/middleware.py

async def on_error(self, error: Exception, messages: list[Message], **kwargs: Any) -> None:
    """Log errors."""
    self.logger.error(
        "LLM error",
        event="llm_error",
        error=str(error),
        error_type=type(error).__name__,
        message_count=len(messages),
    )

Middleware

Bases: ABC

Base class for provider middleware.

Middleware can intercept and modify: - Request messages before sending to provider - Response text after receiving from provider - Streaming chunks as they arrive - Request parameters (temperature, max_tokens, etc.)

Example

class LoggingMiddleware(Middleware): ... async def before_request(self, messages, kwargs): ... print(f"Sending {len(messages)} messages") ... return messages, kwargs ... ... async def after_response(self, messages, response, kwargs): ... print(f"Received {len(response)} chars") ... return response

Source code in bruno_llm/base/middleware.py

class Middleware(ABC):
    """
    Base class for provider middleware.

    Middleware can intercept and modify:
    - Request messages before sending to provider
    - Response text after receiving from provider
    - Streaming chunks as they arrive
    - Request parameters (temperature, max_tokens, etc.)

    Example:
        >>> class LoggingMiddleware(Middleware):
        ...     async def before_request(self, messages, **kwargs):
        ...         print(f"Sending {len(messages)} messages")
        ...         return messages, kwargs
        ...
        ...     async def after_response(self, messages, response, **kwargs):
        ...         print(f"Received {len(response)} chars")
        ...         return response
    """

    @abstractmethod
    async def before_request(
        self, messages: list[Message], **kwargs: Any
    ) -> tuple[list[Message], dict[str, Any]]:
        """
        Process messages and parameters before request.

        Args:
            messages: Input messages
            **kwargs: Request parameters

        Returns:
            Tuple of (modified_messages, modified_kwargs)
        """
        pass

    @abstractmethod
    async def after_response(self, messages: list[Message], response: str, **kwargs: Any) -> str:
        """
        Process response after receiving.

        Args:
            messages: Original input messages
            response: Provider response
            **kwargs: Request parameters used

        Returns:
            Modified response
        """
        pass

    async def on_stream_chunk(self, chunk: str, **kwargs: Any) -> str:
        """
        Process individual stream chunks.

        Args:
            chunk: Stream chunk
            **kwargs: Request parameters

        Returns:
            Modified chunk
        """
        return chunk

    async def on_error(self, error: Exception, messages: list[Message], **kwargs: Any) -> None:  # noqa: B027
        """
        Handle errors during request.

        Args:
            error: The exception that occurred
            messages: Messages that were being processed
            **kwargs: Request parameters
        """
        pass

before_request(messages, **kwargs) abstractmethod async

Process messages and parameters before request.

Parameters:

Name	Type	Description	Default
messages	list[Message]	Input messages	required
**kwargs	Any	Request parameters	{}

Returns:

Type	Description
tuple[list[Message], dict[str, Any]]	Tuple of (modified_messages, modified_kwargs)

Source code in bruno_llm/base/middleware.py

@abstractmethod
async def before_request(
    self, messages: list[Message], **kwargs: Any
) -> tuple[list[Message], dict[str, Any]]:
    """
    Process messages and parameters before request.

    Args:
        messages: Input messages
        **kwargs: Request parameters

    Returns:
        Tuple of (modified_messages, modified_kwargs)
    """
    pass

after_response(messages, response, **kwargs) abstractmethod async

Process response after receiving.

Parameters:

Name	Type	Description	Default
messages	list[Message]	Original input messages	required
response	str	Provider response	required
**kwargs	Any	Request parameters used	{}

Returns:

Type	Description
str	Modified response

Source code in bruno_llm/base/middleware.py

@abstractmethod
async def after_response(self, messages: list[Message], response: str, **kwargs: Any) -> str:
    """
    Process response after receiving.

    Args:
        messages: Original input messages
        response: Provider response
        **kwargs: Request parameters used

    Returns:
        Modified response
    """
    pass

on_stream_chunk(chunk, **kwargs) async

Process individual stream chunks.

Parameters:

Name	Type	Description	Default
chunk	str	Stream chunk	required
**kwargs	Any	Request parameters	{}

Returns:

Type	Description
str	Modified chunk

Source code in bruno_llm/base/middleware.py

async def on_stream_chunk(self, chunk: str, **kwargs: Any) -> str:
    """
    Process individual stream chunks.

    Args:
        chunk: Stream chunk
        **kwargs: Request parameters

    Returns:
        Modified chunk
    """
    return chunk

on_error(error, messages, **kwargs) async

Handle errors during request.

Parameters:

Name	Type	Description	Default
error	Exception	The exception that occurred	required
messages	list[Message]	Messages that were being processed	required
**kwargs	Any	Request parameters	{}

Source code in bruno_llm/base/middleware.py

async def on_error(self, error: Exception, messages: list[Message], **kwargs: Any) -> None:  # noqa: B027
    """
    Handle errors during request.

    Args:
        error: The exception that occurred
        messages: Messages that were being processed
        **kwargs: Request parameters
    """
    pass

MiddlewareChain

Chain multiple middleware together.

Executes middleware in order for before_request, and in reverse order for after_response.

Parameters:

Name	Type	Description	Default
middlewares	list[Middleware]	List of middleware instances	required

Example

chain = MiddlewareChain([ ... LoggingMiddleware(), ... ValidationMiddleware(), ... CachingMiddleware(cache), ... ]) messages, kwargs = await chain.before_request(messages, **kwargs)

Source code in bruno_llm/base/middleware.py

class MiddlewareChain:
    """
    Chain multiple middleware together.

    Executes middleware in order for before_request,
    and in reverse order for after_response.

    Args:
        middlewares: List of middleware instances

    Example:
        >>> chain = MiddlewareChain([
        ...     LoggingMiddleware(),
        ...     ValidationMiddleware(),
        ...     CachingMiddleware(cache),
        ... ])
        >>> messages, kwargs = await chain.before_request(messages, **kwargs)
    """

    def __init__(self, middlewares: list[Middleware]):
        """
        Initialize middleware chain.

        Args:
            middlewares: List of middleware to chain
        """
        self.middlewares = middlewares

    async def before_request(
        self, messages: list[Message], **kwargs: Any
    ) -> tuple[list[Message], dict[str, Any]]:
        """Execute all middleware before_request in order."""
        for middleware in self.middlewares:
            messages, kwargs = await middleware.before_request(messages, **kwargs)
        return messages, kwargs

    async def after_response(self, messages: list[Message], response: str, **kwargs: Any) -> str:
        """Execute all middleware after_response in reverse order."""
        for middleware in reversed(self.middlewares):
            response = await middleware.after_response(messages, response, **kwargs)
        return response

    async def on_stream_chunk(self, chunk: str, **kwargs: Any) -> str:
        """Execute all middleware on_stream_chunk in order."""
        for middleware in self.middlewares:
            chunk = await middleware.on_stream_chunk(chunk, **kwargs)
        return chunk

    async def on_error(self, error: Exception, messages: list[Message], **kwargs: Any) -> None:
        """Execute all middleware on_error."""
        for middleware in self.middlewares:
            await middleware.on_error(error, messages, **kwargs)

__init__(middlewares)

Initialize middleware chain.

Parameters:

Name	Type	Description	Default
middlewares	list[Middleware]	List of middleware to chain	required

Source code in bruno_llm/base/middleware.py

def __init__(self, middlewares: list[Middleware]):
    """
    Initialize middleware chain.

    Args:
        middlewares: List of middleware to chain
    """
    self.middlewares = middlewares

before_request(messages, **kwargs) async

Execute all middleware before_request in order.

Source code in bruno_llm/base/middleware.py

async def before_request(
    self, messages: list[Message], **kwargs: Any
) -> tuple[list[Message], dict[str, Any]]:
    """Execute all middleware before_request in order."""
    for middleware in self.middlewares:
        messages, kwargs = await middleware.before_request(messages, **kwargs)
    return messages, kwargs

after_response(messages, response, **kwargs) async

Execute all middleware after_response in reverse order.

Source code in bruno_llm/base/middleware.py

async def after_response(self, messages: list[Message], response: str, **kwargs: Any) -> str:
    """Execute all middleware after_response in reverse order."""
    for middleware in reversed(self.middlewares):
        response = await middleware.after_response(messages, response, **kwargs)
    return response

on_stream_chunk(chunk, **kwargs) async

Execute all middleware on_stream_chunk in order.

Source code in bruno_llm/base/middleware.py

async def on_stream_chunk(self, chunk: str, **kwargs: Any) -> str:
    """Execute all middleware on_stream_chunk in order."""
    for middleware in self.middlewares:
        chunk = await middleware.on_stream_chunk(chunk, **kwargs)
    return chunk

on_error(error, messages, **kwargs) async

Execute all middleware on_error.

Source code in bruno_llm/base/middleware.py

async def on_error(self, error: Exception, messages: list[Message], **kwargs: Any) -> None:
    """Execute all middleware on_error."""
    for middleware in self.middlewares:
        await middleware.on_error(error, messages, **kwargs)

RetryMiddleware

Bases: Middleware

Add retry logic with exponential backoff.

Note: This is typically handled by BaseProvider's retry logic, but can be used for additional retry layers.

Parameters:

Name	Type	Description	Default
max_retries	int	Maximum number of retries	3
base_delay	float	Base delay in seconds	1.0

Example

middleware = RetryMiddleware(max_retries=3, base_delay=1.0)

Source code in bruno_llm/base/middleware.py

class RetryMiddleware(Middleware):
    """
    Add retry logic with exponential backoff.

    Note: This is typically handled by BaseProvider's retry logic,
    but can be used for additional retry layers.

    Args:
        max_retries: Maximum number of retries
        base_delay: Base delay in seconds

    Example:
        >>> middleware = RetryMiddleware(max_retries=3, base_delay=1.0)
    """

    def __init__(self, max_retries: int = 3, base_delay: float = 1.0):
        """
        Initialize retry middleware.

        Args:
            max_retries: Maximum retry attempts
            base_delay: Base delay for exponential backoff
        """
        self.max_retries = max_retries
        self.base_delay = base_delay
        self.retry_count = 0

    async def before_request(
        self, messages: list[Message], **kwargs: Any
    ) -> tuple[list[Message], dict[str, Any]]:
        """Reset retry count before request."""
        self.retry_count = 0
        return messages, kwargs

    async def after_response(self, messages: list[Message], response: str, **kwargs: Any) -> str:
        """No processing after successful response."""
        return response

    async def on_error(self, error: Exception, messages: list[Message], **kwargs: Any) -> None:
        """Handle retry logic on error."""
        import asyncio

        self.retry_count += 1

        if self.retry_count <= self.max_retries:
            delay = self.base_delay * (2 ** (self.retry_count - 1))
            await asyncio.sleep(delay)

__init__(max_retries=3, base_delay=1.0)

Initialize retry middleware.

Parameters:

Name	Type	Description	Default
max_retries	int	Maximum retry attempts	3
base_delay	float	Base delay for exponential backoff	1.0

Source code in bruno_llm/base/middleware.py

def __init__(self, max_retries: int = 3, base_delay: float = 1.0):
    """
    Initialize retry middleware.

    Args:
        max_retries: Maximum retry attempts
        base_delay: Base delay for exponential backoff
    """
    self.max_retries = max_retries
    self.base_delay = base_delay
    self.retry_count = 0

before_request(messages, **kwargs) async

Reset retry count before request.

Source code in bruno_llm/base/middleware.py

async def before_request(
    self, messages: list[Message], **kwargs: Any
) -> tuple[list[Message], dict[str, Any]]:
    """Reset retry count before request."""
    self.retry_count = 0
    return messages, kwargs

after_response(messages, response, **kwargs) async

No processing after successful response.

Source code in bruno_llm/base/middleware.py

async def after_response(self, messages: list[Message], response: str, **kwargs: Any) -> str:
    """No processing after successful response."""
    return response

on_error(error, messages, **kwargs) async

Handle retry logic on error.

Source code in bruno_llm/base/middleware.py

async def on_error(self, error: Exception, messages: list[Message], **kwargs: Any) -> None:
    """Handle retry logic on error."""
    import asyncio

    self.retry_count += 1

    if self.retry_count <= self.max_retries:
        delay = self.base_delay * (2 ** (self.retry_count - 1))
        await asyncio.sleep(delay)

ValidationMiddleware

Bases: Middleware

Validate messages and parameters.

Parameters:

Name	Type	Description	Default
max_message_length	Optional[int]	Max length for individual messages	None
allowed_roles	Optional[list[str]]	Allowed message roles	None
required_params	Optional[list[str]]	Required parameter names	None

Example

middleware = ValidationMiddleware( ... max_message_length=10000, ... allowed_roles=["user", "assistant", "system"] ... )

Source code in bruno_llm/base/middleware.py

class ValidationMiddleware(Middleware):
    """
    Validate messages and parameters.

    Args:
        max_message_length: Max length for individual messages
        allowed_roles: Allowed message roles
        required_params: Required parameter names

    Example:
        >>> middleware = ValidationMiddleware(
        ...     max_message_length=10000,
        ...     allowed_roles=["user", "assistant", "system"]
        ... )
    """

    def __init__(
        self,
        max_message_length: Optional[int] = None,
        allowed_roles: Optional[list[str]] = None,
        required_params: Optional[list[str]] = None,
    ):
        """
        Initialize validation middleware.

        Args:
            max_message_length: Maximum message length
            allowed_roles: Allowed message roles
            required_params: Required parameters
        """
        self.max_message_length = max_message_length
        self.allowed_roles = allowed_roles
        self.required_params = required_params or []

    async def before_request(
        self, messages: list[Message], **kwargs: Any
    ) -> tuple[list[Message], dict[str, Any]]:
        """Validate before request."""
        # Validate message lengths
        if self.max_message_length:
            for msg in messages:
                if len(msg.content) > self.max_message_length:
                    raise ValueError(
                        f"Message content exceeds max length "
                        f"({len(msg.content)} > {self.max_message_length})"
                    )

        # Validate roles
        if self.allowed_roles:
            for msg in messages:
                if msg.role.value not in self.allowed_roles:
                    raise ValueError(
                        f"Invalid message role: {msg.role.value}. Allowed: {self.allowed_roles}"
                    )

        # Validate required parameters
        for param in self.required_params:
            if param not in kwargs:
                raise ValueError(f"Required parameter missing: {param}")

        return messages, kwargs

    async def after_response(self, messages: list[Message], response: str, **kwargs: Any) -> str:
        """No validation after response."""
        return response

__init__(max_message_length=None, allowed_roles=None, required_params=None)

Initialize validation middleware.

Parameters:

Name	Type	Description	Default
max_message_length	Optional[int]	Maximum message length	None
allowed_roles	Optional[list[str]]	Allowed message roles	None
required_params	Optional[list[str]]	Required parameters	None

Source code in bruno_llm/base/middleware.py

def __init__(
    self,
    max_message_length: Optional[int] = None,
    allowed_roles: Optional[list[str]] = None,
    required_params: Optional[list[str]] = None,
):
    """
    Initialize validation middleware.

    Args:
        max_message_length: Maximum message length
        allowed_roles: Allowed message roles
        required_params: Required parameters
    """
    self.max_message_length = max_message_length
    self.allowed_roles = allowed_roles
    self.required_params = required_params or []

before_request(messages, **kwargs) async

Validate before request.

Source code in bruno_llm/base/middleware.py

async def before_request(
    self, messages: list[Message], **kwargs: Any
) -> tuple[list[Message], dict[str, Any]]:
    """Validate before request."""
    # Validate message lengths
    if self.max_message_length:
        for msg in messages:
            if len(msg.content) > self.max_message_length:
                raise ValueError(
                    f"Message content exceeds max length "
                    f"({len(msg.content)} > {self.max_message_length})"
                )

    # Validate roles
    if self.allowed_roles:
        for msg in messages:
            if msg.role.value not in self.allowed_roles:
                raise ValueError(
                    f"Invalid message role: {msg.role.value}. Allowed: {self.allowed_roles}"
                )

    # Validate required parameters
    for param in self.required_params:
        if param not in kwargs:
            raise ValueError(f"Required parameter missing: {param}")

    return messages, kwargs

after_response(messages, response, **kwargs) async

No validation after response.

Source code in bruno_llm/base/middleware.py

async def after_response(self, messages: list[Message], response: str, **kwargs: Any) -> str:
    """No validation after response."""
    return response

RateLimiter

Async rate limiter using token bucket algorithm.

Controls the rate of API calls to prevent exceeding provider limits. Thread-safe and supports multiple concurrent requests.

Example

limiter = RateLimiter(requests_per_minute=60) async with limiter: ... # Make API call ... response = await api_call()

Source code in bruno_llm/base/rate_limiter.py

class RateLimiter:
    """
    Async rate limiter using token bucket algorithm.

    Controls the rate of API calls to prevent exceeding provider limits.
    Thread-safe and supports multiple concurrent requests.

    Example:
        >>> limiter = RateLimiter(requests_per_minute=60)
        >>> async with limiter:
        ...     # Make API call
        ...     response = await api_call()
    """

    def __init__(
        self,
        requests_per_minute: int = 60,
        tokens_per_minute: Optional[int] = None,
    ):
        """
        Initialize rate limiter.

        Args:
            requests_per_minute: Maximum requests allowed per minute
            tokens_per_minute: Maximum tokens allowed per minute (optional)
        """
        self.requests_per_minute = requests_per_minute
        self.tokens_per_minute = tokens_per_minute

        # Calculate minimum interval between requests
        self.min_interval = 60.0 / requests_per_minute if requests_per_minute > 0 else 0

        # Token bucket for requests
        self._request_tokens = float(requests_per_minute)
        self._max_request_tokens = float(requests_per_minute)
        self._last_update = time.time()

        # Token bucket for API tokens (if specified)
        self._api_tokens = float(tokens_per_minute) if tokens_per_minute else None
        self._max_api_tokens = float(tokens_per_minute) if tokens_per_minute else None

        # Lock for thread safety
        self._lock = asyncio.Lock()

    async def _refill_tokens(self) -> None:
        """Refill token buckets based on elapsed time."""
        now = time.time()
        elapsed = now - self._last_update

        if elapsed <= 0:
            return

        # Refill request tokens
        tokens_to_add = (elapsed * self.requests_per_minute) / 60.0
        self._request_tokens = min(self._max_request_tokens, self._request_tokens + tokens_to_add)

        # Refill API tokens if applicable
        if self._api_tokens is not None and self.tokens_per_minute:
            api_tokens_to_add = (elapsed * self.tokens_per_minute) / 60.0
            self._api_tokens = min(self._max_api_tokens or 0, self._api_tokens + api_tokens_to_add)

        self._last_update = now

    async def acquire(self, api_tokens: int = 0) -> None:
        """
        Acquire permission to make a request.

        Blocks until rate limit allows the request.

        Args:
            api_tokens: Number of API tokens the request will consume
        """
        async with self._lock:
            while True:
                await self._refill_tokens()

                # Check if we have enough request tokens
                if self._request_tokens < 1:
                    # Calculate wait time
                    wait_time = (1 - self._request_tokens) * (60.0 / self.requests_per_minute)
                    await asyncio.sleep(wait_time)
                    continue

                # Check if we have enough API tokens (if applicable)
                if self._api_tokens is not None and api_tokens > 0:
                    if self._api_tokens < api_tokens:
                        wait_time = (api_tokens - self._api_tokens) * (
                            60.0 / (self.tokens_per_minute or 1)
                        )
                        await asyncio.sleep(wait_time)
                        continue

                # Consume tokens
                self._request_tokens -= 1
                if self._api_tokens is not None and api_tokens > 0:
                    self._api_tokens -= api_tokens

                break

    async def __aenter__(self) -> "RateLimiter":
        """Context manager entry."""
        await self.acquire()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Context manager exit."""
        pass

    def get_stats(self) -> dict:
        """
        Get current rate limiter statistics.

        Returns:
            Dict with current token levels and limits
        """
        return {
            "requests_per_minute": self.requests_per_minute,
            "tokens_per_minute": self.tokens_per_minute,
            "available_request_tokens": self._request_tokens,
            "available_api_tokens": self._api_tokens,
            "last_update": self._last_update,
        }

__init__(requests_per_minute=60, tokens_per_minute=None)

Initialize rate limiter.

Parameters:

Name	Type	Description	Default
requests_per_minute	int	Maximum requests allowed per minute	60
tokens_per_minute	Optional[int]	Maximum tokens allowed per minute (optional)	None

Source code in bruno_llm/base/rate_limiter.py

def __init__(
    self,
    requests_per_minute: int = 60,
    tokens_per_minute: Optional[int] = None,
):
    """
    Initialize rate limiter.

    Args:
        requests_per_minute: Maximum requests allowed per minute
        tokens_per_minute: Maximum tokens allowed per minute (optional)
    """
    self.requests_per_minute = requests_per_minute
    self.tokens_per_minute = tokens_per_minute

    # Calculate minimum interval between requests
    self.min_interval = 60.0 / requests_per_minute if requests_per_minute > 0 else 0

    # Token bucket for requests
    self._request_tokens = float(requests_per_minute)
    self._max_request_tokens = float(requests_per_minute)
    self._last_update = time.time()

    # Token bucket for API tokens (if specified)
    self._api_tokens = float(tokens_per_minute) if tokens_per_minute else None
    self._max_api_tokens = float(tokens_per_minute) if tokens_per_minute else None

    # Lock for thread safety
    self._lock = asyncio.Lock()

acquire(api_tokens=0) async

Acquire permission to make a request.

Blocks until rate limit allows the request.

Parameters:

Name	Type	Description	Default
api_tokens	int	Number of API tokens the request will consume	0

Source code in bruno_llm/base/rate_limiter.py

async def acquire(self, api_tokens: int = 0) -> None:
    """
    Acquire permission to make a request.

    Blocks until rate limit allows the request.

    Args:
        api_tokens: Number of API tokens the request will consume
    """
    async with self._lock:
        while True:
            await self._refill_tokens()

            # Check if we have enough request tokens
            if self._request_tokens < 1:
                # Calculate wait time
                wait_time = (1 - self._request_tokens) * (60.0 / self.requests_per_minute)
                await asyncio.sleep(wait_time)
                continue

            # Check if we have enough API tokens (if applicable)
            if self._api_tokens is not None and api_tokens > 0:
                if self._api_tokens < api_tokens:
                    wait_time = (api_tokens - self._api_tokens) * (
                        60.0 / (self.tokens_per_minute or 1)
                    )
                    await asyncio.sleep(wait_time)
                    continue

            # Consume tokens
            self._request_tokens -= 1
            if self._api_tokens is not None and api_tokens > 0:
                self._api_tokens -= api_tokens

            break

__aenter__() async

Context manager entry.

Source code in bruno_llm/base/rate_limiter.py

async def __aenter__(self) -> "RateLimiter":
    """Context manager entry."""
    await self.acquire()
    return self

__aexit__(exc_type, exc_val, exc_tb) async

Context manager exit.

Source code in bruno_llm/base/rate_limiter.py

async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
    """Context manager exit."""
    pass

get_stats()

Get current rate limiter statistics.

Returns:

Type	Description
dict	Dict with current token levels and limits

Source code in bruno_llm/base/rate_limiter.py

def get_stats(self) -> dict:
    """
    Get current rate limiter statistics.

    Returns:
        Dict with current token levels and limits
    """
    return {
        "requests_per_minute": self.requests_per_minute,
        "tokens_per_minute": self.tokens_per_minute,
        "available_request_tokens": self._request_tokens,
        "available_api_tokens": self._api_tokens,
        "last_update": self._last_update,
    }

RetryConfig

Configuration for retry behavior.

Example

config = RetryConfig( ... max_retries=5, ... initial_delay=1.0, ... max_delay=60.0, ... exponential_base=2.0 ... )

Source code in bruno_llm/base/retry.py

class RetryConfig:
    """
    Configuration for retry behavior.

    Example:
        >>> config = RetryConfig(
        ...     max_retries=5,
        ...     initial_delay=1.0,
        ...     max_delay=60.0,
        ...     exponential_base=2.0
        ... )
    """

    def __init__(
        self,
        max_retries: int = 3,
        initial_delay: float = 1.0,
        max_delay: float = 60.0,
        exponential_base: float = 2.0,
        jitter: bool = True,
        retry_on: Optional[tuple[type[Exception], ...]] = None,
    ):
        """
        Initialize retry configuration.

        Args:
            max_retries: Maximum number of retry attempts
            initial_delay: Initial delay between retries in seconds
            max_delay: Maximum delay between retries in seconds
            exponential_base: Base for exponential backoff
            jitter: Whether to add random jitter to delays
            retry_on: Tuple of exception types to retry on (None = all)
        """
        self.max_retries = max_retries
        self.initial_delay = initial_delay
        self.max_delay = max_delay
        self.exponential_base = exponential_base
        self.jitter = jitter
        self.retry_on = retry_on or (Exception,)

    def calculate_delay(self, attempt: int) -> float:
        """
        Calculate delay for given retry attempt.

        Uses exponential backoff with optional jitter.

        Args:
            attempt: Retry attempt number (0-indexed)

        Returns:
            Delay in seconds
        """
        # Exponential backoff
        delay = min(self.initial_delay * (self.exponential_base**attempt), self.max_delay)

        # Add jitter if enabled
        if self.jitter:
            jitter_amount = delay * 0.1  # 10% jitter
            delay += random.uniform(-jitter_amount, jitter_amount)

        return max(0, delay)

    def should_retry(self, exception: Exception, attempt: int) -> bool:
        """
        Determine if retry should be attempted.

        Args:
            exception: Exception that was raised
            attempt: Current attempt number (0-indexed)

        Returns:
            True if should retry, False otherwise
        """
        # Check if we've exhausted retries
        if attempt >= self.max_retries:
            return False

        # Check if exception type is retryable
        if not isinstance(exception, self.retry_on):
            return False

        # Special handling for rate limit errors
        if isinstance(exception, RateLimitError):
            return True

        return True

__init__(max_retries=3, initial_delay=1.0, max_delay=60.0, exponential_base=2.0, jitter=True, retry_on=None)

Initialize retry configuration.

Parameters:

Name	Type	Description	Default
max_retries	int	Maximum number of retry attempts	3
initial_delay	float	Initial delay between retries in seconds	1.0
max_delay	float	Maximum delay between retries in seconds	60.0
exponential_base	float	Base for exponential backoff	2.0
jitter	bool	Whether to add random jitter to delays	True
retry_on	Optional[tuple[type[Exception], ...]]	Tuple of exception types to retry on (None = all)	None

Source code in bruno_llm/base/retry.py

def __init__(
    self,
    max_retries: int = 3,
    initial_delay: float = 1.0,
    max_delay: float = 60.0,
    exponential_base: float = 2.0,
    jitter: bool = True,
    retry_on: Optional[tuple[type[Exception], ...]] = None,
):
    """
    Initialize retry configuration.

    Args:
        max_retries: Maximum number of retry attempts
        initial_delay: Initial delay between retries in seconds
        max_delay: Maximum delay between retries in seconds
        exponential_base: Base for exponential backoff
        jitter: Whether to add random jitter to delays
        retry_on: Tuple of exception types to retry on (None = all)
    """
    self.max_retries = max_retries
    self.initial_delay = initial_delay
    self.max_delay = max_delay
    self.exponential_base = exponential_base
    self.jitter = jitter
    self.retry_on = retry_on or (Exception,)

calculate_delay(attempt)

Calculate delay for given retry attempt.

Uses exponential backoff with optional jitter.

Parameters:

Name	Type	Description	Default
attempt	int	Retry attempt number (0-indexed)	required

Returns:

Type	Description
float	Delay in seconds

Source code in bruno_llm/base/retry.py

def calculate_delay(self, attempt: int) -> float:
    """
    Calculate delay for given retry attempt.

    Uses exponential backoff with optional jitter.

    Args:
        attempt: Retry attempt number (0-indexed)

    Returns:
        Delay in seconds
    """
    # Exponential backoff
    delay = min(self.initial_delay * (self.exponential_base**attempt), self.max_delay)

    # Add jitter if enabled
    if self.jitter:
        jitter_amount = delay * 0.1  # 10% jitter
        delay += random.uniform(-jitter_amount, jitter_amount)

    return max(0, delay)

should_retry(exception, attempt)

Determine if retry should be attempted.

Parameters:

Name	Type	Description	Default
exception	Exception	Exception that was raised	required
attempt	int	Current attempt number (0-indexed)	required

Returns:

Type	Description
bool	True if should retry, False otherwise

Source code in bruno_llm/base/retry.py

def should_retry(self, exception: Exception, attempt: int) -> bool:
    """
    Determine if retry should be attempted.

    Args:
        exception: Exception that was raised
        attempt: Current attempt number (0-indexed)

    Returns:
        True if should retry, False otherwise
    """
    # Check if we've exhausted retries
    if attempt >= self.max_retries:
        return False

    # Check if exception type is retryable
    if not isinstance(exception, self.retry_on):
        return False

    # Special handling for rate limit errors
    if isinstance(exception, RateLimitError):
        return True

    return True

RetryDecorator

Decorator for adding retry logic to async functions.

Example

@RetryDecorator(max_retries=5) ... async def api_call(): ... return await external_api()

Source code in bruno_llm/base/retry.py

class RetryDecorator:
    """
    Decorator for adding retry logic to async functions.

    Example:
        >>> @RetryDecorator(max_retries=5)
        ... async def api_call():
        ...     return await external_api()
    """

    def __init__(
        self,
        max_retries: int = 3,
        initial_delay: float = 1.0,
        max_delay: float = 60.0,
        exponential_base: float = 2.0,
        jitter: bool = True,
    ):
        """
        Initialize retry decorator.

        Args:
            max_retries: Maximum number of retry attempts
            initial_delay: Initial delay between retries
            max_delay: Maximum delay between retries
            exponential_base: Base for exponential backoff
            jitter: Whether to add jitter
        """
        self.config = RetryConfig(
            max_retries=max_retries,
            initial_delay=initial_delay,
            max_delay=max_delay,
            exponential_base=exponential_base,
            jitter=jitter,
        )

    def __call__(self, func: Callable[..., T]) -> Callable[..., T]:
        """
        Wrap function with retry logic.

        Args:
            func: Function to wrap

        Returns:
            Wrapped function
        """

        async def wrapper(*args: Any, **kwargs: Any) -> T:
            return await retry_async(func, *args, config=self.config, **kwargs)

        return wrapper

__init__(max_retries=3, initial_delay=1.0, max_delay=60.0, exponential_base=2.0, jitter=True)

Initialize retry decorator.

Parameters:

Name	Type	Description	Default
max_retries	int	Maximum number of retry attempts	3
initial_delay	float	Initial delay between retries	1.0
max_delay	float	Maximum delay between retries	60.0
exponential_base	float	Base for exponential backoff	2.0
jitter	bool	Whether to add jitter	True

Source code in bruno_llm/base/retry.py

def __init__(
    self,
    max_retries: int = 3,
    initial_delay: float = 1.0,
    max_delay: float = 60.0,
    exponential_base: float = 2.0,
    jitter: bool = True,
):
    """
    Initialize retry decorator.

    Args:
        max_retries: Maximum number of retry attempts
        initial_delay: Initial delay between retries
        max_delay: Maximum delay between retries
        exponential_base: Base for exponential backoff
        jitter: Whether to add jitter
    """
    self.config = RetryConfig(
        max_retries=max_retries,
        initial_delay=initial_delay,
        max_delay=max_delay,
        exponential_base=exponential_base,
        jitter=jitter,
    )

__call__(func)

Wrap function with retry logic.

Parameters:

Name	Type	Description	Default
func	Callable[..., T]	Function to wrap	required

Returns:

Type	Description
Callable[..., T]	Wrapped function

Source code in bruno_llm/base/retry.py

def __call__(self, func: Callable[..., T]) -> Callable[..., T]:
    """
    Wrap function with retry logic.

    Args:
        func: Function to wrap

    Returns:
        Wrapped function
    """

    async def wrapper(*args: Any, **kwargs: Any) -> T:
        return await retry_async(func, *args, config=self.config, **kwargs)

    return wrapper

StreamAggregator

Aggregate streaming chunks with various strategies.

Provides different aggregation strategies for streaming responses: - Word-by-word: Buffer until complete words - Sentence-by-sentence: Buffer until sentence boundaries - Fixed-size: Buffer until fixed character count - Time-based: Buffer for fixed time intervals

Parameters:

Name	Type	Description	Default
strategy	str	Aggregation strategy ('word', 'sentence', 'fixed', 'time')	'word'
size	int	Size parameter (chars for 'fixed', seconds for 'time')	10

Example

aggregator = StreamAggregator(strategy='word') async for chunk in aggregator.aggregate(stream): ... print(chunk) # Prints complete words

Source code in bruno_llm/base/streaming.py

class StreamAggregator:
    """
    Aggregate streaming chunks with various strategies.

    Provides different aggregation strategies for streaming responses:
    - Word-by-word: Buffer until complete words
    - Sentence-by-sentence: Buffer until sentence boundaries
    - Fixed-size: Buffer until fixed character count
    - Time-based: Buffer for fixed time intervals

    Args:
        strategy: Aggregation strategy ('word', 'sentence', 'fixed', 'time')
        size: Size parameter (chars for 'fixed', seconds for 'time')

    Example:
        >>> aggregator = StreamAggregator(strategy='word')
        >>> async for chunk in aggregator.aggregate(stream):
        ...     print(chunk)  # Prints complete words
    """

    def __init__(
        self,
        strategy: str = "word",
        size: int = 10,
    ):
        """
        Initialize stream aggregator.

        Args:
            strategy: Aggregation strategy
            size: Size parameter for aggregation
        """
        self.strategy = strategy
        self.size = size
        self._buffer = ""

    async def aggregate(self, stream: AsyncIterator[str]) -> AsyncIterator[str]:
        """
        Aggregate stream chunks according to strategy.

        Args:
            stream: Input stream to aggregate

        Yields:
            Aggregated chunks
        """
        if self.strategy == "word":
            async for chunk in self._aggregate_words(stream):
                yield chunk
        elif self.strategy == "sentence":
            async for chunk in self._aggregate_sentences(stream):
                yield chunk
        elif self.strategy == "fixed":
            async for chunk in self._aggregate_fixed(stream):
                yield chunk
        elif self.strategy == "time":
            async for chunk in self._aggregate_time(stream):
                yield chunk
        else:
            # No aggregation, pass through
            async for chunk in stream:
                yield chunk

    async def _aggregate_words(self, stream: AsyncIterator[str]) -> AsyncIterator[str]:
        """Aggregate chunks into complete words."""
        async for chunk in stream:
            self._buffer += chunk

            # Split on whitespace while keeping incomplete words
            parts = self._buffer.split()

            if len(parts) > 1:
                # Yield all complete words
                for word in parts[:-1]:
                    yield word + " "

                # Keep the last part as buffer (might be incomplete)
                self._buffer = parts[-1]
            elif self._buffer.endswith((" ", "\n", "\t")):
                # Buffer ends with whitespace, yield it
                yield self._buffer
                self._buffer = ""

        # Flush remaining buffer
        if self._buffer:
            yield self._buffer
            self._buffer = ""

    async def _aggregate_sentences(self, stream: AsyncIterator[str]) -> AsyncIterator[str]:
        """Aggregate chunks into complete sentences."""
        sentence_endings = (".", "!", "?", "\n")

        async for chunk in stream:
            self._buffer += chunk

            # Check if buffer contains sentence ending
            while any(self._buffer.endswith(end) for end in sentence_endings):
                # Find last sentence ending
                last_idx = -1
                for ending in sentence_endings:
                    idx = self._buffer.rfind(ending)
                    if idx > last_idx:
                        last_idx = idx

                if last_idx >= 0:
                    # Yield complete sentence(s)
                    sentence = self._buffer[: last_idx + 1]
                    yield sentence
                    self._buffer = self._buffer[last_idx + 1 :]
                else:
                    break

        # Flush remaining buffer
        if self._buffer:
            yield self._buffer
            self._buffer = ""

    async def _aggregate_fixed(self, stream: AsyncIterator[str]) -> AsyncIterator[str]:
        """Aggregate chunks into fixed-size chunks."""
        async for chunk in stream:
            self._buffer += chunk

            # Yield chunks of fixed size
            while len(self._buffer) >= self.size:
                yield self._buffer[: self.size]
                self._buffer = self._buffer[self.size :]

        # Flush remaining buffer
        if self._buffer:
            yield self._buffer
            self._buffer = ""

    async def _aggregate_time(self, stream: AsyncIterator[str]) -> AsyncIterator[str]:
        """Aggregate chunks based on time intervals."""
        import time

        last_yield = time.time()

        async for chunk in stream:
            self._buffer += chunk

            current_time = time.time()
            if current_time - last_yield >= self.size:
                if self._buffer:
                    yield self._buffer
                    self._buffer = ""
                    last_yield = current_time

        # Flush remaining buffer
        if self._buffer:
            yield self._buffer
            self._buffer = ""

__init__(strategy='word', size=10)

Initialize stream aggregator.

Parameters:

Name	Type	Description	Default
strategy	str	Aggregation strategy	'word'
size	int	Size parameter for aggregation	10

Source code in bruno_llm/base/streaming.py

def __init__(
    self,
    strategy: str = "word",
    size: int = 10,
):
    """
    Initialize stream aggregator.

    Args:
        strategy: Aggregation strategy
        size: Size parameter for aggregation
    """
    self.strategy = strategy
    self.size = size
    self._buffer = ""

aggregate(stream) async

Aggregate stream chunks according to strategy.

Parameters:

Name	Type	Description	Default
stream	AsyncIterator[str]	Input stream to aggregate	required

Yields:

Type	Description
AsyncIterator[str]	Aggregated chunks

Source code in bruno_llm/base/streaming.py

async def aggregate(self, stream: AsyncIterator[str]) -> AsyncIterator[str]:
    """
    Aggregate stream chunks according to strategy.

    Args:
        stream: Input stream to aggregate

    Yields:
        Aggregated chunks
    """
    if self.strategy == "word":
        async for chunk in self._aggregate_words(stream):
            yield chunk
    elif self.strategy == "sentence":
        async for chunk in self._aggregate_sentences(stream):
            yield chunk
    elif self.strategy == "fixed":
        async for chunk in self._aggregate_fixed(stream):
            yield chunk
    elif self.strategy == "time":
        async for chunk in self._aggregate_time(stream):
            yield chunk
    else:
        # No aggregation, pass through
        async for chunk in stream:
            yield chunk

StreamBuffer dataclass

Buffer for managing streaming chunks.

Provides buffering, batching, and aggregation of stream chunks.

Attributes:

Name	Type	Description
buffer	deque[str]	Internal deque for storing chunks
max_size	int	Maximum buffer size in characters
batch_size	int	Number of chunks to batch before yielding
stats	StreamStats	Stream statistics

Source code in bruno_llm/base/streaming.py

@dataclass
class StreamBuffer:
    """
    Buffer for managing streaming chunks.

    Provides buffering, batching, and aggregation of stream chunks.

    Attributes:
        buffer: Internal deque for storing chunks
        max_size: Maximum buffer size in characters
        batch_size: Number of chunks to batch before yielding
        stats: Stream statistics
    """

    buffer: deque[str] = field(default_factory=deque)
    max_size: int = 10000
    batch_size: int = 1
    stats: StreamStats = field(default_factory=StreamStats)

    def add(self, chunk: str) -> None:
        """
        Add a chunk to the buffer.

        Args:
            chunk: Text chunk to buffer

        Raises:
            StreamError: If buffer is full
        """
        current_size = sum(len(c) for c in self.buffer)

        if current_size + len(chunk) > self.max_size:
            raise StreamError(
                f"Stream buffer full ({current_size} chars). "
                f"Consider increasing max_size or consuming buffer faster."
            )

        self.buffer.append(chunk)
        self.stats.chunks_received += 1
        self.stats.total_chars += len(chunk)

    def get_batch(self) -> Optional[str]:
        """
        Get a batch of chunks.

        Returns:
            Concatenated batch or None if not enough chunks
        """
        if len(self.buffer) < self.batch_size:
            return None

        chunks = []
        for _ in range(min(self.batch_size, len(self.buffer))):
            chunks.append(self.buffer.popleft())

        return "".join(chunks)

    def flush(self) -> str:
        """
        Flush all remaining chunks.

        Returns:
            All remaining chunks concatenated
        """
        chunks = list(self.buffer)
        self.buffer.clear()
        return "".join(chunks)

    def is_empty(self) -> bool:
        """Check if buffer is empty."""
        return len(self.buffer) == 0

    def clear(self) -> None:
        """Clear buffer and reset stats."""
        self.buffer.clear()
        self.stats = StreamStats()

add(chunk)

Add a chunk to the buffer.

Parameters:

Name	Type	Description	Default
chunk	str	Text chunk to buffer	required

Raises:

Type	Description
StreamError	If buffer is full

Source code in bruno_llm/base/streaming.py

def add(self, chunk: str) -> None:
    """
    Add a chunk to the buffer.

    Args:
        chunk: Text chunk to buffer

    Raises:
        StreamError: If buffer is full
    """
    current_size = sum(len(c) for c in self.buffer)

    if current_size + len(chunk) > self.max_size:
        raise StreamError(
            f"Stream buffer full ({current_size} chars). "
            f"Consider increasing max_size or consuming buffer faster."
        )

    self.buffer.append(chunk)
    self.stats.chunks_received += 1
    self.stats.total_chars += len(chunk)

get_batch()

Get a batch of chunks.

Returns:

Type	Description
Optional[str]	Concatenated batch or None if not enough chunks

Source code in bruno_llm/base/streaming.py

def get_batch(self) -> Optional[str]:
    """
    Get a batch of chunks.

    Returns:
        Concatenated batch or None if not enough chunks
    """
    if len(self.buffer) < self.batch_size:
        return None

    chunks = []
    for _ in range(min(self.batch_size, len(self.buffer))):
        chunks.append(self.buffer.popleft())

    return "".join(chunks)

flush()

Flush all remaining chunks.

Returns:

Type	Description
str	All remaining chunks concatenated

Source code in bruno_llm/base/streaming.py

def flush(self) -> str:
    """
    Flush all remaining chunks.

    Returns:
        All remaining chunks concatenated
    """
    chunks = list(self.buffer)
    self.buffer.clear()
    return "".join(chunks)

is_empty()

Check if buffer is empty.

Source code in bruno_llm/base/streaming.py

def is_empty(self) -> bool:
    """Check if buffer is empty."""
    return len(self.buffer) == 0

clear()

Clear buffer and reset stats.

Source code in bruno_llm/base/streaming.py

def clear(self) -> None:
    """Clear buffer and reset stats."""
    self.buffer.clear()
    self.stats = StreamStats()

StreamProcessor

Process streaming responses with callbacks and error handling.

Provides a framework for processing streams with: - Progress callbacks - Error recovery - Automatic retry on connection loss - Statistics tracking

Parameters:

Name	Type	Description	Default
on_chunk	Optional[Callable[[str], None]]	Callback for each chunk (chunk: str) -> None	None
on_error	Optional[Callable[[Exception], None]]	Callback for errors (error: Exception) -> None	None
on_complete	Optional[Callable[[StreamStats], None]]	Callback when stream completes (stats: StreamStats) -> None	None
max_retries	int	Maximum number of retries on error	3

Example

processor = StreamProcessor( ... on_chunk=lambda chunk: print(chunk, end=""), ... on_error=lambda e: print(f"Error: {e}"), ... on_complete=lambda stats: print(f"\nReceived {stats.chunks_received} chunks") ... ) await processor.process(stream)

Source code in bruno_llm/base/streaming.py

class StreamProcessor:
    """
    Process streaming responses with callbacks and error handling.

    Provides a framework for processing streams with:
    - Progress callbacks
    - Error recovery
    - Automatic retry on connection loss
    - Statistics tracking

    Args:
        on_chunk: Callback for each chunk (chunk: str) -> None
        on_error: Callback for errors (error: Exception) -> None
        on_complete: Callback when stream completes (stats: StreamStats) -> None
        max_retries: Maximum number of retries on error

    Example:
        >>> processor = StreamProcessor(
        ...     on_chunk=lambda chunk: print(chunk, end=""),
        ...     on_error=lambda e: print(f"Error: {e}"),
        ...     on_complete=lambda stats: print(f"\\nReceived {stats.chunks_received} chunks")
        ... )
        >>> await processor.process(stream)
    """

    def __init__(
        self,
        on_chunk: Optional[Callable[[str], None]] = None,
        on_error: Optional[Callable[[Exception], None]] = None,
        on_complete: Optional[Callable[[StreamStats], None]] = None,
        max_retries: int = 3,
    ):
        """
        Initialize stream processor.

        Args:
            on_chunk: Callback for each chunk
            on_error: Callback for errors
            on_complete: Callback when stream completes
            max_retries: Maximum number of retries on error
        """
        self.on_chunk = on_chunk
        self.on_error = on_error
        self.on_complete = on_complete
        self.max_retries = max_retries
        self.stats = StreamStats()

    async def process(
        self,
        stream: AsyncIterator[str],
        retry_on_error: bool = True,
    ) -> list[str]:
        """
        Process a stream with callbacks and error handling.

        Args:
            stream: Input stream to process
            retry_on_error: Whether to retry on errors

        Returns:
            List of all chunks received

        Raises:
            StreamError: If max retries exceeded
        """
        import time

        chunks = []
        start_time = time.time()
        retries = 0

        try:
            async for chunk in stream:
                chunks.append(chunk)
                self.stats.chunks_received += 1
                self.stats.total_chars += len(chunk)

                if self.on_chunk:
                    self.on_chunk(chunk)

            # Calculate duration after stream completes
            self.stats.duration = time.time() - start_time

            if self.on_complete:
                self.on_complete(self.stats)

            return chunks

        except Exception as e:
            self.stats.errors += 1
            self.stats.duration = time.time() - start_time

            if self.on_error:
                self.on_error(e)

            if retry_on_error and retries < self.max_retries:
                retries += 1
                await asyncio.sleep(2**retries)  # Exponential backoff
                return await self.process(stream, retry_on_error)

            raise StreamError(f"Stream processing failed: {e}") from e

__init__(on_chunk=None, on_error=None, on_complete=None, max_retries=3)

Initialize stream processor.

Parameters:

Name	Type	Description	Default
on_chunk	Optional[Callable[[str], None]]	Callback for each chunk	None
on_error	Optional[Callable[[Exception], None]]	Callback for errors	None
on_complete	Optional[Callable[[StreamStats], None]]	Callback when stream completes	None
max_retries	int	Maximum number of retries on error	3

Source code in bruno_llm/base/streaming.py

def __init__(
    self,
    on_chunk: Optional[Callable[[str], None]] = None,
    on_error: Optional[Callable[[Exception], None]] = None,
    on_complete: Optional[Callable[[StreamStats], None]] = None,
    max_retries: int = 3,
):
    """
    Initialize stream processor.

    Args:
        on_chunk: Callback for each chunk
        on_error: Callback for errors
        on_complete: Callback when stream completes
        max_retries: Maximum number of retries on error
    """
    self.on_chunk = on_chunk
    self.on_error = on_error
    self.on_complete = on_complete
    self.max_retries = max_retries
    self.stats = StreamStats()

process(stream, retry_on_error=True) async

Process a stream with callbacks and error handling.

Parameters:

Name	Type	Description	Default
stream	AsyncIterator[str]	Input stream to process	required
retry_on_error	bool	Whether to retry on errors	True

Returns:

Type	Description
list[str]	List of all chunks received

Raises:

Type	Description
StreamError	If max retries exceeded

Source code in bruno_llm/base/streaming.py

async def process(
    self,
    stream: AsyncIterator[str],
    retry_on_error: bool = True,
) -> list[str]:
    """
    Process a stream with callbacks and error handling.

    Args:
        stream: Input stream to process
        retry_on_error: Whether to retry on errors

    Returns:
        List of all chunks received

    Raises:
        StreamError: If max retries exceeded
    """
    import time

    chunks = []
    start_time = time.time()
    retries = 0

    try:
        async for chunk in stream:
            chunks.append(chunk)
            self.stats.chunks_received += 1
            self.stats.total_chars += len(chunk)

            if self.on_chunk:
                self.on_chunk(chunk)

        # Calculate duration after stream completes
        self.stats.duration = time.time() - start_time

        if self.on_complete:
            self.on_complete(self.stats)

        return chunks

    except Exception as e:
        self.stats.errors += 1
        self.stats.duration = time.time() - start_time

        if self.on_error:
            self.on_error(e)

        if retry_on_error and retries < self.max_retries:
            retries += 1
            await asyncio.sleep(2**retries)  # Exponential backoff
            return await self.process(stream, retry_on_error)

        raise StreamError(f"Stream processing failed: {e}") from e

StreamStats dataclass

Statistics for a streaming session.

Attributes:

Name	Type	Description
chunks_received	int	Number of chunks received
total_chars	int	Total characters streamed
total_tokens	int	Estimated token count
duration	float	Duration of stream in seconds
errors	int	Number of errors encountered

Source code in bruno_llm/base/streaming.py

@dataclass
class StreamStats:
    """
    Statistics for a streaming session.

    Attributes:
        chunks_received: Number of chunks received
        total_chars: Total characters streamed
        total_tokens: Estimated token count
        duration: Duration of stream in seconds
        errors: Number of errors encountered
    """

    chunks_received: int = 0
    total_chars: int = 0
    total_tokens: int = 0
    duration: float = 0.0
    errors: int = 0

SimpleTokenCounter

Bases: TokenCounter

Simple token counter using word splitting.

This is a fallback implementation that approximates token count by counting words. Not as accurate as provider-specific tokenizers but works universally.

Example

counter = SimpleTokenCounter() tokens = counter.count_tokens("Hello world!") print(tokens) # Approximately 2-3

Source code in bruno_llm/base/token_counter.py

class SimpleTokenCounter(TokenCounter):
    """
    Simple token counter using word splitting.

    This is a fallback implementation that approximates token count
    by counting words. Not as accurate as provider-specific tokenizers
    but works universally.

    Example:
        >>> counter = SimpleTokenCounter()
        >>> tokens = counter.count_tokens("Hello world!")
        >>> print(tokens)  # Approximately 2-3
    """

    def __init__(self, chars_per_token: float = 4.0):
        """
        Initialize simple token counter.

        Args:
            chars_per_token: Average characters per token (default: 4)
        """
        self.chars_per_token = chars_per_token

    def count_tokens(self, text: str) -> int:
        """
        Count tokens using character-based estimation.

        Uses the common approximation that 1 token ≈ 4 characters
        in English text.

        Args:
            text: Text to count tokens for

        Returns:
            Estimated token count
        """
        if not text:
            return 0
        return max(1, int(len(text) / self.chars_per_token))

__init__(chars_per_token=4.0)

Initialize simple token counter.

Parameters:

Name	Type	Description	Default
chars_per_token	float	Average characters per token (default: 4)	4.0

Source code in bruno_llm/base/token_counter.py

def __init__(self, chars_per_token: float = 4.0):
    """
    Initialize simple token counter.

    Args:
        chars_per_token: Average characters per token (default: 4)
    """
    self.chars_per_token = chars_per_token

count_tokens(text)

Count tokens using character-based estimation.

Uses the common approximation that 1 token ≈ 4 characters in English text.

Parameters:

Name	Type	Description	Default
text	str	Text to count tokens for	required

Returns:

Type	Description
int	Estimated token count

Source code in bruno_llm/base/token_counter.py

def count_tokens(self, text: str) -> int:
    """
    Count tokens using character-based estimation.

    Uses the common approximation that 1 token ≈ 4 characters
    in English text.

    Args:
        text: Text to count tokens for

    Returns:
        Estimated token count
    """
    if not text:
        return 0
    return max(1, int(len(text) / self.chars_per_token))

TikTokenCounter

Bases: TokenCounter

Token counter using OpenAI's tiktoken library.

Provides accurate token counting for OpenAI models. Falls back to SimpleTokenCounter if tiktoken is not available.

Example

counter = TikTokenCounter(model="gpt-4") tokens = counter.count_tokens("Hello world!") print(tokens)

Source code in bruno_llm/base/token_counter.py

class TikTokenCounter(TokenCounter):
    """
    Token counter using OpenAI's tiktoken library.

    Provides accurate token counting for OpenAI models.
    Falls back to SimpleTokenCounter if tiktoken is not available.

    Example:
        >>> counter = TikTokenCounter(model="gpt-4")
        >>> tokens = counter.count_tokens("Hello world!")
        >>> print(tokens)
    """

    def __init__(self, model: str = "gpt-4"):
        """
        Initialize tiktoken-based counter.

        Args:
            model: Model name for tiktoken encoding
        """
        self.model = model
        self._encoding = None
        self._fallback = SimpleTokenCounter()

        try:
            import tiktoken

            self._encoding = tiktoken.encoding_for_model(model)
        except ImportError:
            # tiktoken not available, will use fallback
            pass
        except Exception:
            # Model not found or other error, use fallback
            pass

    def count_tokens(self, text: str) -> int:
        """
        Count tokens using tiktoken or fallback.

        Args:
            text: Text to count tokens for

        Returns:
            Accurate token count (if tiktoken available) or estimate
        """
        if not text:
            return 0

        if self._encoding is not None:
            try:
                return len(self._encoding.encode(text))
            except Exception:
                pass

        # Fallback to simple counting
        return self._fallback.count_tokens(text)

__init__(model='gpt-4')

Initialize tiktoken-based counter.

Parameters:

Name	Type	Description	Default
model	str	Model name for tiktoken encoding	'gpt-4'

Source code in bruno_llm/base/token_counter.py

def __init__(self, model: str = "gpt-4"):
    """
    Initialize tiktoken-based counter.

    Args:
        model: Model name for tiktoken encoding
    """
    self.model = model
    self._encoding = None
    self._fallback = SimpleTokenCounter()

    try:
        import tiktoken

        self._encoding = tiktoken.encoding_for_model(model)
    except ImportError:
        # tiktoken not available, will use fallback
        pass
    except Exception:
        # Model not found or other error, use fallback
        pass

count_tokens(text)

Count tokens using tiktoken or fallback.

Parameters:

Name	Type	Description	Default
text	str	Text to count tokens for	required

Returns:

Type	Description
int	Accurate token count (if tiktoken available) or estimate

Source code in bruno_llm/base/token_counter.py

def count_tokens(self, text: str) -> int:
    """
    Count tokens using tiktoken or fallback.

    Args:
        text: Text to count tokens for

    Returns:
        Accurate token count (if tiktoken available) or estimate
    """
    if not text:
        return 0

    if self._encoding is not None:
        try:
            return len(self._encoding.encode(text))
        except Exception:
            pass

    # Fallback to simple counting
    return self._fallback.count_tokens(text)

TokenCounter

Bases: ABC

Abstract base class for token counting.

Different providers may have different tokenization methods. Subclasses should implement provider-specific counting logic.

Source code in bruno_llm/base/token_counter.py

class TokenCounter(ABC):
    """
    Abstract base class for token counting.

    Different providers may have different tokenization methods.
    Subclasses should implement provider-specific counting logic.
    """

    @abstractmethod
    def count_tokens(self, text: str) -> int:
        """
        Count tokens in text.

        Args:
            text: Text to count tokens for

        Returns:
            Number of tokens
        """
        pass

    def count_message_tokens(self, message: Message) -> int:
        """
        Count tokens in a message.

        Args:
            message: Message to count tokens for

        Returns:
            Number of tokens
        """
        return self.count_tokens(message.content)

    def count_messages_tokens(self, messages: list[Message]) -> int:
        """
        Count tokens in multiple messages.

        Args:
            messages: List of messages

        Returns:
            Total number of tokens
        """
        total = 0
        for message in messages:
            total += self.count_message_tokens(message)
            # Add overhead for message formatting (role, etc.)
            total += 4  # Approximate overhead per message
        return total

count_tokens(text) abstractmethod

Count tokens in text.

Parameters:

Name	Type	Description	Default
text	str	Text to count tokens for	required

Returns:

Type	Description
int	Number of tokens

Source code in bruno_llm/base/token_counter.py

@abstractmethod
def count_tokens(self, text: str) -> int:
    """
    Count tokens in text.

    Args:
        text: Text to count tokens for

    Returns:
        Number of tokens
    """
    pass

count_message_tokens(message)

Count tokens in a message.

Parameters:

Name	Type	Description	Default
message	Message	Message to count tokens for	required

Returns:

Type	Description
int	Number of tokens

Source code in bruno_llm/base/token_counter.py

def count_message_tokens(self, message: Message) -> int:
    """
    Count tokens in a message.

    Args:
        message: Message to count tokens for

    Returns:
        Number of tokens
    """
    return self.count_tokens(message.content)

count_messages_tokens(messages)

Count tokens in multiple messages.

Parameters:

Name	Type	Description	Default
messages	list[Message]	List of messages	required

Returns:

Type	Description
int	Total number of tokens

Source code in bruno_llm/base/token_counter.py

def count_messages_tokens(self, messages: list[Message]) -> int:
    """
    Count tokens in multiple messages.

    Args:
        messages: List of messages

    Returns:
        Total number of tokens
    """
    total = 0
    for message in messages:
        total += self.count_message_tokens(message)
        # Add overhead for message formatting (role, etc.)
        total += 4  # Approximate overhead per message
    return total

retry_async(func, *args, config=None, **kwargs) async

Execute async function with retry logic.

Parameters:

Name	Type	Description	Default
func	Callable[..., T]	Async function to execute	required
*args	Any	Positional arguments for func	()
config	Optional[RetryConfig]	Retry configuration (uses defaults if None)	None
**kwargs	Any	Keyword arguments for func	{}

Returns:

Type	Description
T	Result from func

Raises:

Type	Description
Exception	Last exception if all retries fail

Example

async def api_call(): ... # May fail transiently ... return await external_api() result = await retry_async(api_call, config=RetryConfig(max_retries=5))

Source code in bruno_llm/base/retry.py

async def retry_async(
    func: Callable[..., T],
    *args: Any,
    config: Optional[RetryConfig] = None,
    **kwargs: Any,
) -> T:
    """
    Execute async function with retry logic.

    Args:
        func: Async function to execute
        *args: Positional arguments for func
        config: Retry configuration (uses defaults if None)
        **kwargs: Keyword arguments for func

    Returns:
        Result from func

    Raises:
        Exception: Last exception if all retries fail

    Example:
        >>> async def api_call():
        ...     # May fail transiently
        ...     return await external_api()
        >>> result = await retry_async(api_call, config=RetryConfig(max_retries=5))
    """
    if config is None:
        config = RetryConfig()

    last_exception = None

    for attempt in range(config.max_retries + 1):
        try:
            return await func(*args, **kwargs)
        except Exception as e:
            last_exception = e

            # Check if we should retry
            if not config.should_retry(e, attempt):
                raise

            # Calculate and wait
            if attempt < config.max_retries:
                delay = config.calculate_delay(attempt)

                # Special handling for rate limit with retry_after
                if isinstance(e, RateLimitError) and e.retry_after:
                    delay = max(delay, e.retry_after)

                await asyncio.sleep(delay)

    # Should not reach here, but for safety
    if last_exception:
        raise last_exception
    raise LLMError("Retry loop exited unexpectedly")

stream_with_timeout(stream, timeout=30.0) async

Wrap a stream with timeout protection.

Raises TimeoutError if no chunk is received within timeout.

Parameters:

Name	Type	Description	Default
stream	AsyncIterator[str]	Input stream to wrap	required
timeout	float	Timeout in seconds for each chunk	30.0

Yields:

Type	Description
AsyncIterator[str]	Stream chunks

Raises:

Type	Description
TimeoutError	If timeout is exceeded

Example

async for chunk in stream_with_timeout(stream, timeout=10.0): ... print(chunk)

Source code in bruno_llm/base/streaming.py

async def stream_with_timeout(
    stream: AsyncIterator[str],
    timeout: float = 30.0,
) -> AsyncIterator[str]:
    """
    Wrap a stream with timeout protection.

    Raises TimeoutError if no chunk is received within timeout.

    Args:
        stream: Input stream to wrap
        timeout: Timeout in seconds for each chunk

    Yields:
        Stream chunks

    Raises:
        TimeoutError: If timeout is exceeded

    Example:
        >>> async for chunk in stream_with_timeout(stream, timeout=10.0):
        ...     print(chunk)
    """
    async for chunk in stream:
        try:
            yield await asyncio.wait_for(_async_identity(chunk), timeout=timeout)
        except asyncio.TimeoutError as e:
            raise TimeoutError(f"No chunk received within {timeout} seconds") from e

create_token_counter(provider='simple', model=None)

Factory function to create appropriate token counter.

Parameters:

Name	Type	Description	Default
provider	str	Provider name ("simple", "openai", "tiktoken")	'simple'
model	Optional[str]	Optional model name for provider-specific counting	None

Returns:

Type	Description
TokenCounter	TokenCounter instance

Example

counter = create_token_counter("openai", model="gpt-4") tokens = counter.count_tokens("Hello!")

Source code in bruno_llm/base/token_counter.py

def create_token_counter(
    provider: str = "simple",
    model: Optional[str] = None,
) -> TokenCounter:
    """
    Factory function to create appropriate token counter.

    Args:
        provider: Provider name ("simple", "openai", "tiktoken")
        model: Optional model name for provider-specific counting

    Returns:
        TokenCounter instance

    Example:
        >>> counter = create_token_counter("openai", model="gpt-4")
        >>> tokens = counter.count_tokens("Hello!")
    """
    if provider in ("openai", "tiktoken"):
        if model:
            return TikTokenCounter(model=model)
        return TikTokenCounter()

    return SimpleTokenCounter()