shekkari21/agent-from-scratch

1"""Callback utilities for agent execution."""2 3import inspect4from typing import Optional, Callable5 6from .models import ExecutionContext7from .llm import LlmRequest, LlmResponse8from .memory import count_tokens9 10 11def create_optimizer_callback(12    apply_optimization: Callable,13    threshold: int = 50000,14    model_id: str = "gpt-4"15) -> Callable:16    """Factory function that creates a callback applying optimization strategy.17    18    Args:19        apply_optimization: Function that modifies the LlmRequest in place20        threshold: Token count threshold to trigger optimization21        model_id: Model identifier for token counting22    23    Returns:24        Callback function that can be used as before_llm_callback25    """26    async def callback(27        context: ExecutionContext,28        request: LlmRequest29    ) -> Optional[LlmResponse]:30        token_count = count_tokens(request, model_id=model_id)31 32        if token_count < threshold:33            return None34 35        # Support both sync and async functions36        result = apply_optimization(context, request)37        if inspect.isawaitable(result):38            await result39        return None40 41    return callback42